Yolo tflite serving

edge_interface/src/components/Inference.vue

@@ -1,6 +1,6 @@
 <template>
   <div class="mr-4 container">
-    <v-btn color="blue-grey" class="ma-2 white--text" @click="trigger">
+    <v-btn color="blue-grey" class="ma-2 white--text" @click="trigger" @call-trigger="trigger" id="trigger-button">
       Trigger
       <v-icon right dark>
         mdi-cloud-upload
@@ -23,27 +23,32 @@
       </p>
       <div v-for="(object, index) in predictedItem" :key="index">
         <h3>{{ object.cameraId }}</h3>
-        <div>
-          <img class="img-responsive" :src="object.image_url" />
+        <div class="inference-image">
+          <img class="img-responsive" ref="image" :src="object.image_url" @load="on_image_loaded" />
           <div v-for="(inference, model_id) in object.inferences" :key="model_id">
-            <div v-for="(result, object_id) in inference" :key="object_id">
-              <div v-if="'location' in result">
-                <Box
-                  :x-min="result['location'][0]"
-                  :y-min="result['location'][1]"
-                  :x-max="result['location'][2]"
-                  :y-max="result['location'][3]"
-                />
+            <div v-if="inference !== 'NO_DECISION'">
+              <div v-for="(result, object_id) in inference" :key="object_id">
+                <div v-if="'location' in result">
+                  <Box
+                    v-if="imgLoaded"
+                    v-bind:x-min="xoffset + result['location'][0] * width"
+                    v-bind:y-min="yoffset + result['location'][1] * height"
+                    v-bind:x-max="xoffset + result['location'][2] * width"
+                    v-bind:y-max="yoffset + result['location'][3] * height"
+                  />
+                </div>
               </div>
             </div>
           </div>
         </div>
         <div v-for="(inference, model_id) in object.inferences" :key="model_id">
           <h4>{{ model_id }}</h4>
-          <div v-for="(result, object_id) in inference" :key="object_id">
-            <span>{{ object_id }}</span>
-            <div v-for="(value, key) in result" :key="key">
-              <span>{{ key }}: {{ value }}</span>
+          <div v-if="inference !== 'NO_DECISION'">
+            <div v-for="(result, object_id) in inference" :key="object_id">
+              <span>{{ object_id }}</span>
+              <div v-for="(value, key) in result" :key="key">
+                <span>{{ key }}: {{ value }}</span>
+              </div>
             </div>
           </div>
         </div>
@@ -68,9 +73,23 @@ export default {
     itemId: null,
     statusList: null,
     state: undefined,
-    decision: undefined
+    decision: undefined,
+    imgLoaded: false,
+    height: null,
+    width: null,
+    xoffset: null,
+    yoffset: null
   }),
   methods: {
+    on_image_loaded() {
+      const img = this.$refs.image[0]
+      this.height = img.height
+      this.width = img.width
+      this.xoffset = img.offsetLeft
+      this.yoffset = img.offsetTop
+      console.log('Image size : ', this.height, this.width)
+      this.imgLoaded = true
+    },
     getColor(status) {
       if (this.statusList[status] > this.statusList[this.state]) {
         return 'red'
@@ -139,8 +158,7 @@ export default {
 <style lang="scss" scoped>
 .result {
   display: inline-block;
-  vertical-align: top;
-  padding: 0 5rem 0 5rem;
+  position: relative;
 }
 
 .container {

edge_interface/src/views/UploadView.vue

@@ -50,6 +50,9 @@
             <v-btn v-if="!isStart" color="error" class="mr-4" @click="onStart">
               Start Camera
             </v-btn>
+            <v-btn v-if="isStart" color="error" class="mr-4" @click="onStartVid">
+              Start Video Capture
+            </v-btn>
           </div>
         </div>
       </div>
@@ -151,6 +154,18 @@ export default {
       this.isStart = true
       this.$refs.webcam.start()
     },
+    onStartVid() {
+      const execTime = 60 * 1000 // 1 minute
+      const intervalID = setInterval(this.captureTrigger, 2000)
+      setTimeout(() => {
+        clearTimeout(intervalID)
+      }, execTime)
+    },
+    captureTrigger() {
+      this.onCapture()
+      const button = document.getElementById('trigger-button')
+      button.click()
+    },
     checkDeviceId(device) {
       return device.deviceId === this.deviceId
     },

edge_model_serving/tflite_serving/src/tflite_serving/api_routes.py

@@ -3,6 +3,11 @@
 
 import numpy as np
 from fastapi import APIRouter, HTTPException, Request
+from tflite_serving.utils.yolo_postprocessing import (
+    yolo_extract_boxes_information,
+    non_max_suppression,
+    compute_severities,
+)
 
 JSONObject = Dict[AnyStr, Any]
 JSONArray = List[Any]
@@ -44,8 +49,13 @@ async def predict(
         input_data = payload[b"inputs"]
         input_array = np.array(input_data, dtype=input_dtype)
 
-        interpreter.set_tensor(input_details[0]["index"], input_array)
+        model_type = None
+        if b"model_type" in payload.keys():
+            model_type = payload[b"model_type"]
+            if model_type == "yolo":
+                input_array /= 255
 
+        interpreter.set_tensor(input_details[0]["index"], input_array)
         interpreter.invoke()
         # Process image and get predictions
         prediction = {}
@@ -67,8 +77,32 @@ async def predict(
                     "detection_boxes": boxes.tolist(),
                     "detection_classes": classes.tolist(),
                     "detection_scores": scores.tolist(),
+                    "severities": [None],
+                }
+            }
+        elif model_type == "yolo":
+            outputs = interpreter.get_tensor(output_details[0]["index"])[0]
+
+            # Rotate the tensor
+            temp_output = []
+            for i in range(len(outputs[0]), 0, -1):
+                temp_output.append(list(map(lambda x: x[i - 1], outputs)))
+            outputs = np.array(temp_output)
+
+            # Extracting the boxes information to select only the most relevant ones
+            boxes, scores, class_ids = yolo_extract_boxes_information(outputs)
+            boxes, scores, class_ids = non_max_suppression(boxes, scores, class_ids)
+            severities = compute_severities(input_array[0], boxes)
+
+            prediction = {
+                "outputs": {
+                    "detection_boxes": [boxes],
+                    "detection_classes": [class_ids],
+                    "detection_scores": [scores],
+                    "severities": [severities],
                 }
             }
+
         elif len(output_details) == 1:
             scores = interpreter.get_tensor(output_details[0]["index"])
             logging.warning(

edge_model_serving/tflite_serving/src/tflite_serving/utils/yolo_postprocessing.py

@@ -0,0 +1,144 @@
+import numpy as np
+from typing import List
+
+
+def yolo_extract_boxes_information(outputs):
+    rows = outputs.shape[0]
+
+    boxes = []
+    scores = []
+    class_ids = []
+
+    for i in range(rows):
+        classes_scores = outputs[i][4:]
+        max_score, max_class_id = max((v, i) for i, v in enumerate(classes_scores))
+
+        box = [
+            float(outputs[i][0] - (0.5 * outputs[i][2])),
+            float(outputs[i][1] - (0.5 * outputs[i][3])),
+            float(outputs[i][0] + (0.5 * outputs[i][2])),
+            float(outputs[i][1] + (0.5 * outputs[i][3])),
+        ]
+        boxes.append(box)
+        scores.append(float(max_score))
+        class_ids.append(float(max_class_id))
+
+    return boxes, scores, class_ids
+
+
+def compute_severities(image: np.array, boxes: List):
+    severities = []
+    for box in boxes:
+        severities.append(compute_box_severity(image, box))
+    return severities
+
+
+def compute_box_severity(image: np.array, box: List):
+    x1_pixel_index = int(box[0] * len(image))
+    y1_pixel_index = int(box[1] * len(image[0]))
+    x2_pixel_index = int(box[2] * len(image))
+    y2_pixel_index = int(box[3] * len(image[0]))
+
+    # Reshape to only the pixels in the detection box & as a list of pixels instead of a 2D array of them
+    image_detection = image[
+        x1_pixel_index:x2_pixel_index, y1_pixel_index:y2_pixel_index, :
+    ]
+    image_detection = image_detection.reshape(-1, 3)
+
+    # Filtering out light pixels
+    mask_dark_pixels = np.all(image_detection < 0.5, axis=1)
+    # Looking at severity as mean darkness
+    dark_colors = image_detection[mask_dark_pixels].flatten()
+    if dark_colors.size == 0:
+        return 0.1
+    else:
+        severity = round((0.5 - dark_colors.mean()) * 2, 2)
+        return severity
+
+
+def non_max_suppression(
+    boxes, scores, class_ids, score_threshold=0.4, iou_threshold=0.45
+):
+    non_max_suppression_parameters_checks(score_threshold, iou_threshold)
+
+    nms_result_boxes = []
+    nms_result_scores = []
+    nms_result_classes = []
+
+    # Cut values with low confidence
+    scores = np.array(scores)
+    mask_low_confidence = scores > score_threshold
+    scores = scores[mask_low_confidence].tolist()
+    boxes = np.array(boxes)[mask_low_confidence].tolist()
+    class_ids = np.array(class_ids)[mask_low_confidence].tolist()
+
+    # Performing the Non-max suppression loop
+    while len(boxes) != 0:
+        delete_index_list = []
+
+        # Locating & Saving max confidence box
+        highest_confidence_index = scores.index(max(scores))
+        highest_confidence_box = boxes[highest_confidence_index]
+
+        nms_result_boxes.append(highest_confidence_box)
+        nms_result_scores.append(scores[highest_confidence_index])
+        nms_result_classes.append(class_ids[highest_confidence_index])
+        delete_index_list.append(highest_confidence_index)
+
+        # Iterating to analyse the iou scores
+        for index_box, box in enumerate(boxes):
+            iou = compute_iou(highest_confidence_box, box)
+            if iou > iou_threshold:
+                delete_index_list.append(index_box)
+
+        # Rebuild the box list to remove the boxes that were close to the last max score
+        boxes = [
+            box
+            for index_box, box in enumerate(boxes)
+            if index_box not in delete_index_list
+        ]
+        scores = [
+            score
+            for index_score, score in enumerate(scores)
+            if index_score not in delete_index_list
+        ]
+        class_ids = [
+            class_id
+            for index_class, class_id in enumerate(class_ids)
+            if index_class not in delete_index_list
+        ]
+
+    return nms_result_boxes, nms_result_scores, nms_result_classes
+
+
+def non_max_suppression_parameters_checks(conf_thres, iou_thres):
+    assert (
+        0 <= conf_thres <= 1
+    ), f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
+    assert (
+        0 <= iou_thres <= 1
+    ), f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"
+
+
+def compute_iou(box1, box2):
+    box1_width = box1[2] - box1[0]
+    box1_height = box1[3] - box1[1]
+    box2_width = box2[2] - box2[0]
+    box2_height = box2[3] - box2[1]
+
+    # Check if centers of boxes are close enough to be intersected
+    if (
+        abs((box1[0] + box1_width / 2) - (box2[0] + box2_width / 2))
+        < 0.5 * (box2_width + box1_width)
+    ) & (
+        abs((box1[1] + box1_height / 2) - (box2[1] + box2_height / 2))
+        < 0.5 * (box2_height + box1_height)
+    ):
+        intersection_area = (max(box1[2], box2[2]) - min(box1[0], box2[0])) * (
+            max(box1[3], box2[3]) - min(box1[1], box2[1])
+        )
+        union_area = (
+            box1_width * box1_height + box2_width * box2_height - intersection_area
+        )
+        return intersection_area / union_area
+    return 0