feat(idle_model): write code for getting inference from yolor model

Author: eugenos-programos

idle_models/IdleObjectDetectionModel.py

@@ -1,22 +1,34 @@
 import torch
-from ultralytics import YOLO
+from yolor.model import get_model
+import numpy as np
+from yolor.utils.datasets import letterbox
+from yolor.utils.general import non_max_suppression, scale_coords
 
 
 class IdleObjectDetectionModel:
     def __init__(self, path: str, conf_thresh, iou_thresh, classes) -> None:
-        self.model = YOLO(path)
+        self.model, self.device = get_model(path, "yolor/yolor_csp_x.cfg")
         self.conf_thresh = conf_thresh
         self.iou_thresh = iou_thresh
         self.classes = classes
 
+    def __preprocess_image__(self, img: np.array) -> np.array:
+        self.img_shape = img.shape
+        img = letterbox(img.copy(), new_shape=1280, auto_size=64)[0]
+        img = img[:, :, ::-1].transpose(2, 0, 1)
+        img = np.ascontiguousarray(img)
+        img = torch.from_numpy(img).to(self.device)
+        img = img.float()
+        img /= 255.0
+        img = img.unsqueeze(0)
+        return img
+
     @torch.no_grad()
-    def __call__(self, img) -> list:
-        results = self.model(
-            source=img,
-            conf=self.conf_thresh,
-            iou=self.iou_thresh,
-            max_det=600,
-            classes=self.classes,
-            verbose=False
-        )[0].boxes
-        return results.xyxy, results.conf
+    def __call__(self, img: np.array) -> list:
+        img = self.__preprocess_image__(img)
+        pred = self.model(img, augment=False)[0]
+        pred = non_max_suppression(
+            pred, 0.45, 0.5, classes=[67], agnostic=False)[0]
+        pred[:, :4] = scale_coords(
+            img.shape[2:], pred[:, :4], self.img_shape).round()
+        return pred[:, :4], pred[:, 4]

idle_models/requirements.txt

@@ -7,5 +7,4 @@ pydantic==1.10.2
 python-dotenv==1.0.0
 PyYAML==6.0
 requests==2.27.1
-ultralytics==8.0.112
 Flask==2.2.2

idle_models/yolor/coco.names

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorcycle
+airplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+couch
+potted plant
+bed
+dining table
+toilet
+tv
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

idle_models/yolor/model.py

@@ -0,0 +1,12 @@
+import torch
+from yolor.utils.torch_utils import select_device
+from yolor.models.models import Darknet
+
+
+def get_model(weights, cfg):
+    imgsz = 1280
+    device = select_device('cpu')
+    model = Darknet(cfg, imgsz)
+    model.load_state_dict(torch.load(weights, map_location=device)['model'])
+    model.to(device).eval()
+    return model, device

send_request.py

@@ -18,8 +18,8 @@ def predict(img: np.array, server_url: str, logger: Logger):
         confidences = np.array(response.json().get("confidences"))
     else:
         logger.warning(
-            "Status code = {}\n JSON = {}".format(
-                response.status_code, response.json())
+            "Status code = {}\n response = {}".format(
+                response.status_code, response)
         )
         coordinates = confidences = None
     return [coordinates, confidences]