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,144 @@
+import cv2
+import numpy as np
+from configparser import RawConfigParser
+
+DEF_SCORE = 0.3
+DEF_NMS = 0.4
+
+class BaseInference:
+    """
+    Base Inference class for predictions
+
+    ...
+
+    Attributes
+    ----------
+    weight_path : str
+        path to the .weight file
+    config_path : str
+        path to the .cfg file
+    classes : list
+        names of classes detected
+    score_thresh : float
+        threshold to classify object as detected
+    nms_threshold : float
+        threshold for non-max suppression
+
+    Methods
+    -------
+    initialize()
+        Initializes neural network using given weight and config
+    read_config()
+        Reads config for trained height and width
+    run()
+        Obtains predicted boxes 
+    """
+
+    def __init__(self, weight_path, config_path, classes, score_thresh=None, nms_thresh=None):
+        self.weight_path = weight_path
+        self.config_path = config_path
+        self.classes = classes
+        self.net = None
+        self.score_thresh = score_thresh if score_thresh is not None else DEF_SCORE
+        self.nms_thresh = nms_thresh if nms_thresh is not None else DEF_NMS
+
+    def initialize(self):
+        # Load Yolo
+        self.net = cv2.dnn.readNet(
+            self.weight_path,
+            self.config_path
+        )
+        layer_names = self.net.getLayerNames()
+        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)
+
+        net_dict = dict(cfg.items('net'))
+        self.train_height_width = (int(net_dict['height']), int(net_dict['width']))
+
+    def run(self, img):
+        """
+        Parameters
+        ----------
+        img : cv2.Mat
+            Image as a matrix
+
+        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
+        dets : list(list(float))
+            Normalised dimensions of boxes
+        indexes : list(int)
+            Indexes of boxes that passed NMS
+        """        
+        height, width = self.get_height_width_from_img(img)
+
+        return self.get_filtered_boxes(img, height, width)
+
+    def get_filtered_boxes(self, img, height, width):
+        outs = self.run_single(img)
+        class_ids, box_dims, box_confidences, dets = self.get_boxes(outs, height, width)
+        indexes = self.filter_boxes(box_dims, box_confidences)
+
+        return class_ids, box_dims, box_confidences, dets, indexes
+
+    def get_height_width_from_img(self, img):
+        height, width, channels = img.shape
+        return height, width
+
+    def run_single(self, img):
+        # Detecting objects
+        blob = cv2.dnn.blobFromImage(img, 0.00392, self.train_height_width, (0, 0, 0), True, crop=False)
+
+        self.net.setInput(blob)
+        outs = self.net.forward(self.output_layers)
+
+        return outs
+
+    def get_boxes(self, output, height, width, test=True):
+
+        # Showing informations on the screen
+        class_ids = []
+        box_confidences = []
+        box_dims = []
+        # Saving to txt
+        dets = []
+
+        for out in output:
+            for detection in out:
+
+                scores = detection[5:]
+                class_id = np.argmax(scores)
+                confidence = scores[class_id]
+                if confidence > self.score_thresh:
+                    # 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
+                    dets.append(detection[:4])
+
+        return class_ids, box_dims, box_confidences, dets
+
+    def filter_boxes(self, box_dims, box_confidences):
+        indexes = cv2.dnn.NMSBoxes(box_dims, box_confidences, self.score_thresh, self.nms_thresh)
+        indexes = [int(i) for i in indexes]
+        return indexes
+

lambda_backend/predict_microservice/handler.py

@@ -1,29 +1,23 @@
-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")
-)
+DEF_WEIGHTS = join("weights", "yolov4-tiny-obj.weights")
+DEF_CONFIG = join("weights", "yolov4-tiny-obj.cfg")
+DEF_CLASSES = ["hold"]
 
-# Names of custom objects
-# Refer to colab notebook for index to class mappings
-classes = ["hold"]
+ALLOWED_TYPES = ["image/jpeg"]
 
-layer_names = net.getLayerNames()
-output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
+def setup():
+    inference = ServiceInference(DEF_WEIGHTS, DEF_CONFIG, DEF_CLASSES)
+    inference.initialize()
+    inference.read_config()
+    return inference
 
+inference = setup()
 
 @exception_handler
 def predict(event, context):
@@ -57,68 +51,9 @@ 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, width_dict)
 
     return {
         "statusCode": "200",

lambda_backend/predict_microservice/service_inference.py

@@ -0,0 +1,82 @@
+from base_inference import BaseInference
+
+class ServiceInference(BaseInference):
+    """
+    Inference class for predictions on predict_microservice
+
+    ...
+
+    Methods
+    -------
+    run(img, width_dict)
+        Obtains predicted boxes for predict_microservice
+    """
+
+    def run(self, img, width_dict):
+        """
+        Parameters
+        ----------
+        img : cv2.Mat
+            Image as a matrix
+
+        Returns
+        -------
+        scaled_height : int
+            Scaled height of img
+        scaled_width : int
+            Scaled width of img
+        boxes : list
+            List of predicted boxes in JSON format
+        """
+        scaled_height, scaled_width = self.get_height_width_from_img(img, width_dict)
+        class_ids, box_dims, box_confidences, _, indexes = super().get_filtered_boxes(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_height_width_from_img(self, img, width_dict):
+        height, width = super().get_height_width_from_img(img)
+
+        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)
+
+        return scaled_height, scaled_width
+
+    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:
+            i = int(i)
+            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