Feat/Improve Yolo-v4 Model

.gitignore

@@ -131,4 +131,6 @@ dmypy.json
 node_modules
 .serverless
 .env.*
-**/test_images/*.txt
+**/test_images/*.txt
+**/local_prediction/*.cfg
+**/local_prediction/*.weights

lambda_backend/predict_microservice/base_inference.py

@@ -0,0 +1,148 @@
+import cv2
+import numpy as np
+from configparser import RawConfigParser
+
+CONFIDENCE_THRESHOLD = 0.3
+NMS_IOU_THRESHOLD = 0.4
+SCALE_FACTOR = 0.00392
+
+class BaseInference:
+    """
+    Base Inference class for predictions
+
+    ...
+
+    Attributes
+    ----------
+    weight_path : str
+        path to the .weights file
+    config_path : str
+        path to the .cfg file
+    classes : list
+        names of classes detected
+    score_threshold : float
+        threshold to classify object as detected
+    nms_threshold : float
+        threshold for non-max suppression
+
+    Methods
+    -------
+    run()
+        Obtains predicted boxes 
+    """
+
+    def __init__(self, weight_path, config_path, classes, score_threshold=None, nms_thresh=None):
+        self.weight_path = weight_path
+        self.config_path = config_path
+        self.classes = classes
+        self.net = None
+        self.score_threshold = score_threshold if score_threshold is not None else CONFIDENCE_THRESHOLD
+        self.nms_thresh = nms_thresh if nms_thresh is not None else NMS_IOU_THRESHOLD
+
+        self._initialize_model()
+        self._read_config()
+
+    def _initialize_model(self):
+        # Load Yolo
+        self.net = cv2.dnn.readNet(
+            self.weight_path,
+            self.config_path
+        )
+        layer_names = self.net.getLayerNames()
+        # Gets the indexes of layers with unconnected outputs, 
+        # then stores the associated names into output_layers
+        self.output_layers = [layer_names[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
+
+    def _read_config(self):
+        cfg = RawConfigParser(strict=False)
+        cfg.read(self.config_path)
+
+        assert 'net' in cfg, 'No net section in config'
+
+        net_dict = dict(cfg.items('net'))
+
+        assert 'height' in net_dict and 'weight' in net_dict, 'No height and/or weight in config'
+        self.train_height_width = (int(net_dict['height']), int(net_dict['width']))
+
+    def run(self, img, height=None, width=None):
+        """
+        Parameters
+        ----------
+        img : cv2.Mat
+            Image as a matrix
+        height : int, optional
+            Height of img (default is None)
+        width : int, optional
+            Width of img (default is None)
+
+        Returns
+        -------
+        class_ids : list(int)
+            Class IDs of boxes
+        box_dims : list(list(int))
+            Dimensions of boxes
+        box_confidences : list(float)
+            Confidence scores of boxes
+        box_dims_norm : list(list(float))
+            Normalised dimensions of boxes
+        indexes : list(int)
+            Indexes of boxes that passed NMS
+        """        
+
+        # If run is called without height or width given
+        if height is None or width is None:
+            height, width, channels = img.shape
+
+        # Detecting objects
+        blob = cv2.dnn.blobFromImage(
+            image       = img, 
+            scalefactor = SCALE_FACTOR, 
+            size        = self.train_height_width, 
+            mean        = (0,0,0), 
+            swapRB      = True, 
+            crop        = False
+        )
+
+        self.net.setInput(blob)
+        outs = self.net.forward(self.output_layers)
+
+        class_ids, box_dims, box_confidences, box_dims_norm, indexes = self._get_filtered_boxes(outs, height, width)
+
+        return class_ids, box_dims, box_confidences, box_dims_norm, indexes
+
+    def _get_filtered_boxes(self, output, height, width):
+        # Showing informations on the screen
+        class_ids = []
+        box_confidences = []
+        box_dims = []
+        # Saving to txt
+        box_dims_norm = []
+
+        for out in output:
+            for detection in out:
+
+                scores = detection[5:]
+                class_id = np.argmax(scores)
+                confidence = scores[class_id]
+                if confidence > self.score_threshold:
+                    # Object detected
+                    center_x = int(detection[0] * width)
+                    center_y = int(detection[1] * height)
+                    w = int(detection[2] * width)
+                    h = int(detection[3] * height)
+
+                    # Rectangle coordinates
+                    x = int(center_x - w / 2)
+                    y = int(center_y - h / 2)
+
+                    box_dims.append([x, y, w, h])
+                    box_confidences.append(float(confidence))
+                    class_ids.append(class_id)
+
+                    # Save normalised format
+                    box_dims_norm.append(detection[:4])
+
+        indexes = cv2.dnn.NMSBoxes(box_dims, box_confidences, self.score_threshold, self.nms_thresh)
+        indexes = [int(i) for i in indexes]
+
+        return class_ids, box_dims, box_confidences, box_dims_norm, indexes

lambda_backend/predict_microservice/handler.py

@@ -1,29 +1,17 @@
-import numpy as np
-import cv2
-
-from os.path import join
 import os
-import glob
 import json
 import base64
 
 from utils import exception_handler, retrieve_numpy_image, parse_multipart_data, get_response_headers
+from service_inference import ServiceInference
 
-ALLOWED_TYPES = ["image/jpeg"]
-
-# Load Yolo
-net = cv2.dnn.readNet(
-    join("weights", "yolov4-tiny-obj.weights"),
-    join("weights", "yolov4-tiny-obj.cfg")
-)
-
-# Names of custom objects
-# Refer to colab notebook for index to class mappings
-classes = ["hold"]
+DEFAULT_WEIGHTS = os.path.join("weights", "yolov4-tiny-obj.weights")
+DEFAULT_CONFIG = os.path.join("weights", "yolov4-tiny-obj.cfg")
+DEFAULT_CLASSES = ["hold"]
 
-layer_names = net.getLayerNames()
-output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+ALLOWED_TYPES = ["image/jpeg"]
 
+inference = ServiceInference(DEFAULT_WEIGHTS, DEFAULT_CONFIG, DEFAULT_CLASSES)
 
 @exception_handler
 def predict(event, context):
@@ -57,68 +45,11 @@ def predict(event, context):
     assert image_dict["type"] in ALLOWED_TYPES, "Unallowed file type"
 
     img = retrieve_numpy_image(image_dict["content"])
-    height, width, channels = img.shape
 
     scaled_width = int(width_dict["content"].decode("utf-8"))
 
-    # If given width is 0, do not scale
-    scaled_width = scaled_width if scaled_width != 0 else width
-    scaled_height = int((scaled_width / width) * height)
-
-    # Image Blob
-    blob = cv2.dnn.blobFromImage(
-        img,
-        0.00392,
-        (416, 416),
-        (0, 0, 0),
-        True,
-        crop=False
-    )
-
-    net.setInput(blob)
-    outs = net.forward(output_layers)
-
-    box_dimensions = []
-    box_confidences = []
-    class_ids = []
-
-    for out in outs:
-        for detection in out:
-            scores = detection[5:]
-            class_id = np.argmax(scores)
-            confidence = scores[class_id]
-            if confidence > 0.3:
-                # Object detected
-                center_x = int(detection[0] * scaled_width)
-                center_y = int(detection[1] * scaled_height)
-                w = int(detection[2] * scaled_width)
-                h = int(detection[3] * scaled_height)
-
-                # Rectangle coordinates
-                x = int(center_x - w / 2)
-                y = int(center_y - h / 2)
-
-                box_dimensions.append([x, y, w, h])
-                box_confidences.append(float(confidence))
-                class_ids.append(class_id)
-
-    boxes = []
-    # Non Maximum Suppression
-    indexes = cv2.dnn.NMSBoxes(box_dimensions, box_confidences, 0.5, 0.4)
-    for i in indexes:
-        i = int(i)
-        x, y, w, h = box_dimensions[i]
-        boxes.append({
-            "x": x,
-            "y": y,
-            "w": w,
-            "h": h,
-            "confidence": float(box_confidences[i]),
-            "class": str(classes[class_ids[i]])
-        })
-
-    # Sort boxes in descending sizes
-    boxes = sorted(boxes, key=lambda box: box["w"] * box["h"], reverse=True)
+    # Run inference on image
+    scaled_height, scaled_width, boxes = inference.run(img, scaled_width)
 
     return {
         "statusCode": "200",

lambda_backend/predict_microservice/service_inference.py

@@ -0,0 +1,78 @@
+from base_inference import BaseInference
+
+class ServiceInference(BaseInference):
+    """
+    Inference class for predictions on predict_microservice
+
+    ...
+
+    Methods
+    -------
+    run(img, scaled_width)
+        Obtains predicted boxes for predict_microservice
+    """
+
+    def run(self, img, scaled_width):
+        """
+        Parameters
+        ----------
+        img : cv2.Mat
+            Image as a matrix
+        scaled_width : int
+            Scaled width of images
+
+        Returns
+        -------
+        scaled_height : int
+            Scaled height of img
+        scaled_width : int
+            Scaled width of img
+        boxes : list
+            List of predicted boxes in JSON format
+        """
+
+        height, width = img.shape
+        # If given width is 0, do not scale
+        scaled_width = scaled_width if scaled_width != 0 else width
+        scaled_height = int((scaled_width / width) * height)
+
+        class_ids, box_dims, box_confidences, _, indexes = super().run(img, scaled_height, scaled_width)
+
+        boxes = self._get_boxes_dict(box_dims, box_confidences, class_ids, indexes)
+        return scaled_height, scaled_width, boxes
+
+    def _get_boxes_dict(self, box_dims, box_confidences, class_ids, indexes):
+        """
+        Parameters
+        ----------
+        box_dims : list
+            Dimensions of predicted boxes
+        box_confidences : list
+            Confidence scores of predicted boxes
+        class_ids : list
+            Class IDs of predicted boxes
+        indexes : list
+            Indexes of predicted boxes after NMS
+
+        Returns
+        -------
+        boxes : list
+            List of predicted boxes in JSON format
+        """
+
+        boxes = []
+        for i in indexes:
+            x, y, w, h = box_dims[i]
+            boxes.append({
+                "x": x,
+                "y": y,
+                "w": w,
+                "h": h,
+                "confidence": float(box_confidences[i]),
+                "class": str(self.classes[class_ids[i]])
+            })
+
+        # Sort boxes in descending sizes
+        boxes = sorted(boxes, key=lambda box: box["w"] * box["h"], reverse=True)
+
+        return boxes

lambda_backend/predict_microservice/weights/yolov4-tiny-obj.cfg

@@ -1,10 +1,10 @@
 [net]
 # Testing
-batch=8
-subdivisions=16
+#batch=64
+#subdivisions=16
 # Training
-# batch=64
-# subdivisions=16
+batch=64
+subdivisions=16
 width=640
 height=640
 channels=3
@@ -226,7 +226,7 @@ iou_normalizer=0.07
 iou_loss=ciou
 ignore_thresh = .7
 truth_thresh = 1
-random=0
+random=1
 resize=1.5
 nms_kind=greedynms
 beta_nms=0.6
@@ -275,7 +275,7 @@ iou_normalizer=0.07
 iou_loss=ciou
 ignore_thresh = .7
 truth_thresh = 1
-random=0
+random=1
 resize=1.5
 nms_kind=greedynms
 beta_nms=0.6