Re-wrote YOLOv5 Plugin using torch instead of OpenCV

object_detection/config/params.yaml

@@ -5,8 +5,8 @@ object_detection:
     output_img_topic: object_detection/img
     publish_output_img: 1
     model_params:
-      detector_type: YOLOv5
-      model_dir_path: /root/percep_ws/models/yolov5
-      weight_file_name: yolov5.onnx
+      detector_type: YOLOv8
+      model_dir_path: /root/percep_ws/models/yolov8
+      weight_file_name: yolov8s.pt
       confidence_threshold : 0.5
       show_fps : 1

object_detection/object_detection/Detectors/YOLOv5.py

@@ -1,148 +1,53 @@
+import torch
 import os
-import cv2
-import numpy as np
 
 from ..DetectorBase import DetectorBase
 
 
 class YOLOv5(DetectorBase):
-    def __init__(self, conf_threshold = 0.7, score_threshold = 0.4, nms_threshold = 0.25, is_cuda = 1):
+    def __init__(self, conf_threshold = 0.7):
 
         super().__init__()
 
-        # opencv img input
-        self.frame = None
-        self.net = None
-        self.INPUT_WIDTH = 640
-        self.INPUT_HEIGHT = 640
-        self.CONFIDENCE_THRESHOLD = conf_threshold
+        self.conf_threshold = conf_threshold
 
-        self.is_cuda = is_cuda   
-
-
-    # load model and prepare its backend to either run on GPU or CPU, see if it can be added in constructor
     def build_model(self, model_dir_path, weight_file_name):
-        model_path = os.path.join(model_dir_path, weight_file_name)
-
-        try:
-            self.net = cv2.dnn.readNet(model_path)
-        except:
-            raise Exception("Error loading given model from path: {}. Maybe the file doesn't exist?".format(model_path))
-
-        if self.is_cuda:
-            print("is_cuda was set to True. Attempting to use CUDA")
-            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
-            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)
-        else:
-            print("is_cuda was set to False. Running on CPU")
-            self.net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
-            self.net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
 
+        try : 
+            model_path = os.path.join(model_dir_path, weight_file_name)
+            self.model = torch.hub.load('ultralytics/yolov5:v6.0', 'custom', path = model_path, force_reload = True)
 
-    # load classes.txt that contains mapping of model with labels
-    # TODO: add try/except to raise exception that tells the use to check the name if it is classes.txt
+        except :
+            raise Exception("Error loading given model from path: {}. Maybe the file doesn't exist?".format(model_path))
+
     def load_classes(self, model_dir_path):
+
         self.class_list = []
-        with open(model_dir_path + "/classes.txt", "r") as f:
+
+        with open(os.path.join(model_dir_path, 'classes.txt')) as f :
             self.class_list = [cname.strip() for cname in f.readlines()]
-        return self.class_list
 
+        return self.class_list
 
-    def detect(self, image):
-        # convert image to 640x640 
-        blob = cv2.dnn.blobFromImage(image, 1/255.0, (self.INPUT_WIDTH, self.INPUT_HEIGHT), swapRB=True, crop=False)
-        self.net.setInput(blob)
-        preds = self.net.forward()
-        return preds
-
-
-    # extract bounding box, class IDs and confidences of detected objects
-    # YOLOv5 returns a 3D tensor of dimension 25200*(5 + n_classes)
-    def wrap_detection(self, input_image, output_data):
-        class_ids = []
-        confidences = []
-        boxes = []
-
-        rows = output_data.shape[0]
-
-        image_width, image_height, _ = input_image.shape
-
-        x_factor = image_width / self.INPUT_WIDTH
-        y_factor =  image_height / self.INPUT_HEIGHT
-
-        # Iterate through all the 25200 vectors
-        for r in range(rows):
-            row = output_data[r]
-
-            # Continue only if Pc > conf_threshold
-            confidence = row[4]
-            if confidence >= self.CONFIDENCE_THRESHOLD:
-
-                # One-hot encoded vector representing class of object
-                classes_scores = row[5:]
-
-                # Returns min and max values in a array alongwith their indices
-                _, _, _, max_indx = cv2.minMaxLoc(classes_scores)
-
-                # Extract the column index of the maximum values in classes_scores
-                class_id = max_indx[1]
-
-                # Continue of the class score is greater than a threshold
-                # class_score represents the probability of an object belonging to that class
-                if (classes_scores[class_id] > .25):
-
-                    confidences.append(confidence)
-
-                    class_ids.append(class_id)
-
-                    x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item() 
-                    left = int((x - 0.5 * w) * x_factor)
-                    top = int((y - 0.5 * h) * y_factor)
-                    width = int(w * x_factor)
-                    height = int(h * y_factor)
-                    box = np.array([left, top, width, height])
-                    boxes.append(box)
-
-        # removing intersecting bounding boxes
-        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.25, 0.45) 
-
-        result_class_ids = []
-        result_confidences = []
-        result_boxes = []
-
-        for i in indexes:
-            result_confidences.append(confidences[i])
-            result_class_ids.append(class_ids[i])
-            result_boxes.append(boxes[i])
-
-        return result_class_ids, result_confidences, result_boxes
-
-
-    # makes image square with dimension max(h, w)
-    def format_yolov5(self):
-        row, col, _ = self.frame.shape
-        _max = max(col, row)
-        result = np.zeros((_max, _max, 3), np.uint8)
-        result[0:row, 0:col] = self.frame
-        return result
-
-
     def get_predictions(self, cv_image):
-        #Clear list
-        self.predictions = []
 
-        if cv_image is None:
+        if cv_image is None :
             # TODO: show warning message (different color, maybe)
-            return None,None
-        else:
+            return None, None
+
+        else :
             self.frame = cv_image
+            class_id = []
+            confidence = []
+            boxes = []
 
-            # make image square
-            inputImage = self.format_yolov5()
+            results = self.model(self.frame)
+
+            for *xyxy, conf, id in results.xyxy[0]:
+                class_id.append(int(id))
+                confidence.append(conf.item())
+                boxes.append([int(xy) for xy in xyxy])
 
-            outs = self.detect(inputImage)
-            class_ids, confidences, boxes = self.wrap_detection(inputImage, outs[0])
+            super().create_predictions_list(class_id, confidence, boxes)
 
-            super().create_predictions_list(class_ids, confidences, boxes)  
-
-            return self.predictions
+            return self.predictions

object_detection/object_detection/Detectors/YOLOv8.py

@@ -5,51 +5,51 @@
 
 
 class YOLOv8(DetectorBase):
-  def __init__(self, conf_threshold = 0.7):
-
-    super().__init__()
-
-    self.conf_threshold = conf_threshold
-
-  def build_model(self, model_dir_path, weight_file_name) :
-
-    try :
-      model_path = os.path.join(model_dir_path, weight_file_name)
-      self.model = YOLO(model_path)
+    def __init__(self, conf_threshold = 0.7):
+
+        super().__init__()
+
+        self.conf_threshold = conf_threshold
+
+    def build_model(self, model_dir_path, weight_file_name) :
+
+        try :
+            model_path = os.path.join(model_dir_path, weight_file_name)
+            self.model = YOLO(model_path)
+
+        except :
+            raise Exception("Error loading given model from path: {}. Maybe the file doesn't exist?".format(model_path))
 
-    except :
-      raise Exception("Error loading given model from path: {}. Maybe the file doesn't exist?".format(model_path))
 
+    def load_classes(self, model_dir_path):
+
+        self.class_list = []
+
+        with open(model_dir_path + "/classes.txt", "r") as f:
+            self.class_list = [cname.strip() for cname in f.readlines()]
+
+        return self.class_list
+
+    def get_predictions(self, cv_image):
+
+        if cv_image is None:
+            # TODO: show warning message (different color, maybe)
+            return None
+
+        else :     
+            self.frame = cv_image   
+            class_id = []
+            confidence = []
+            boxes = []
 
-  def load_classes(self, model_dir_path):
-
-    self.class_list = []
-
-    with open(model_dir_path + "/classes.txt", "r") as f:
-        self.class_list = [cname.strip() for cname in f.readlines()]
-
-    return self.class_list
-
-  def get_predictions(self, cv_image):
-
-    if cv_image is None:
-      # TODO: show warning message (different color, maybe)
-      return None,None
-
-    else :     
-      self.frame = cv_image   
-      class_id = []
-      confidence = []
-      boxes = []
+            result = self.model.predict(self.frame, conf = self.conf_threshold, verbose = False) # Perform object detection on image
+            row = result[0].boxes.cpu()
 
-      result = self.model.predict(self.frame, conf = self.conf_threshold, verbose = False) # Perform object detection on image
-      row = result[0].boxes.cpu()
+            for box in row:
+                class_id.append(box.cls.numpy().tolist()[0])
+                confidence.append(box.conf.numpy().tolist()[0])
+                boxes.append(box.xyxy.numpy().tolist()[0])
 
-      for box in row:
-        class_id.append(box.cls.numpy().tolist()[0])
-        confidence.append(box.conf.numpy().tolist()[0])
-        boxes.append(box.xyxy.numpy().tolist()[0])
+            super().create_predictions_list(class_id, confidence, boxes)
 
-      super().create_predictions_list(class_id, confidence, boxes)
-
-      return self.predictions
+            return self.predictions