feat(model): use yolov6.pt

requirements.txt

@@ -13,3 +13,5 @@ ultralytics==8.3.68
 ultralytics-thop==2.0.14
 #opencv-python==4.11.0.86
 python-dotenv==1.0.0
+onnxruntime==1.19.2
+matplotlib==3.8.1

tasks/image.py

@@ -112,25 +112,123 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # Update the code below to replace the random baseline with your model inference
     #--------------------------------------------------------------------------------------------   
     
+    import cv2
+    import onnxruntime
+    import matplotlib.pyplot as plt
+
+    #PATH_TO_MODEL = 'models/best_YOLOv11n_1280.onnx'
+    PATH_TO_MODEL = 'models/best_yolov6n_1280.pt'
+    INFERENCE_ENGINE_TYPE = 'pt'
+    INPUT_SIZE = 1280
+
+    def preprocessor(frame):
+        #frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # Only when read from file
+        x = cv2.resize(frame, (INPUT_SIZE, INPUT_SIZE))
+        image_data = np.array(x).astype(np.float32) / 255.0  # Normalize to [0, 1] range
+        image_data = np.transpose(image_data, (2, 0, 1))  # (H, W, C) -> (C, H, W)
+        image_data = np.expand_dims(image_data, axis=0)  # Add batch dimension
+        return image_data
+
+    class Inference:
+        def __init__(self, model, image):
+            self.session = onnxruntime.InferenceSession(model,#, providers=["CPUExecutionProvider"]
+                                                        providers=["CUDAExecutionProvider"]
+                                                        )
+            model_inputs = self.session.get_inputs()
+            input_shape = model_inputs[0].shape
+            self.image = image
+            self.input_width = input_shape[2]
+            self.input_height = input_shape[3]
+            self.classes = {0: 'smoke'}
+
+        def detector(self, image_data):
+            ort = onnxruntime.OrtValue.ortvalue_from_numpy(image_data)
+            return self.session.run(["output0"], {"images": ort})
+
+        def postprocessor(self, results, frame, confidence, iou):
+            img_height, img_width = frame.shape[:2]
+            outputs = np.transpose(np.squeeze(results[0]))
+            rows = outputs.shape[0]
+            boxes = []
+            final_boxes = []
+            final_scores = []
+            scores = []
+            class_ids = []
+            x_factor = img_width / self.input_width  
+            y_factor = img_height / self.input_height  
+
+            max_of_max_scores = 0
+            for i in range(rows):
+                classes_scores = outputs[i][4:]
+                max_score = np.amax(classes_scores)
+                if max_score >= confidence:
+                    class_id = np.argmax(classes_scores)
+                    x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
+                    # Calculate the scaled coordinates of the bounding box
+                    left = int(x *  x_factor) / img_width
+                    top = int(y * y_factor) / img_height
+                    width = int(w * x_factor) / img_width
+                    height = int(h * y_factor) / img_height
+                    class_ids.append(class_id)
+                    scores.append(max_score)
+                    boxes.append([left, top, width, height])
+                max_of_max_scores = max(max_of_max_scores, max_score)
+            # Apply non-maximum suppression to filter out overlapping bounding boxes
+            indices = cv2.dnn.NMSBoxes(boxes, scores, confidence, iou)
+            for i in indices:
+                box = boxes[i]
+                score = scores[i]
+                class_id = class_ids[i]
+                final_boxes.append(box)
+                final_scores.append(score)
+            return frame, final_boxes, final_scores
+
+        def pipeline(self):
+            if isinstance(self.image, str):
+                frame = cv2.imread(self.image)
+            else:
+                frame = np.array(self.image)
+            preprocessed = preprocessor(frame)
+            detected = self.detector(preprocessed)
+            frame, boxes, scores = self.postprocessor(detected, frame, 0.20,0.20)
+            return frame, boxes, scores
     
-    PATH_TO_MODEL = f"models/best_YOLOv11n_1280.pt"
-    model = load_model(PATH_TO_MODEL)
-    
-    print(f"Model info: {model.info()}")
+    def predict(inference_engine_type, image, path_to_model):
+        if inference_engine_type == 'pt':
+            print("INFO - Using pytorch model")
+            inference_engine = YOLO(path_to_model)
+            boxes = inference_engine.predict(image)[0].boxes.xywhn.tolist()
+            confidences = inference_engine.predict(image)[0].boxes.conf.tolist()
+            return boxes, confidences
+        
+        elif inference_engine_type == 'onnx':
+            print("INFO -Using onnx model")
+            inference_engine = Inference(path_to_model, image)
+            _, boxes, scores = inference_engine.pipeline()
+            return boxes, scores
+        else:
+            raise ValueError(f"Invalid inference engine type: {inference_engine_type}")
+        
+    print("Starting inference")
     predictions = []
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
     
     n_examples = len(test_dataset)
+    n_boxes = []
     for i, example in enumerate(test_dataset):
         print(f"Running {i+1} of {n_examples}")
+
         # Parse true annotation (YOLO format: class_id x_center y_center width height)
         annotation = example.get("annotations", "").strip()
+        n_annotations = len(annotation.split("\n"))
+        n_boxes.append(n_annotations)
+
         has_smoke = len(annotation) > 0
         true_labels.append(int(has_smoke))
         
-        model_preds = model(example['image'])[0]
+        model_preds, scores = predict(INFERENCE_ENGINE_TYPE,  example['image'], PATH_TO_MODEL)
         pred_has_smoke = len(model_preds) > 0
         predictions.append(int(pred_has_smoke))
         
@@ -142,12 +240,20 @@ async def evaluate_image(request: ImageEvaluationRequest):
             true_boxes_list.append(image_true_boxes)
             
             try:
-                pred_box_list = get_boxes_list(model_preds)[0] # With one bbox to start with (as in the random baseline)
+                print("pred_box_list", model_preds) # With one bbox to start with (as in the random baseline)
+                model_preds = model_preds[0]
+                if len(model_preds) < 1:
+                    model_preds = [0, 0, 0, 0]
             except:
                 print("No boxes found")
-                pred_box_list = [0, 0, 0, 0] # Hacky way to make sure that compute_max_iou doesn't fail
-            pred_boxes.append(pred_box_list)
+                model_preds = [0, 0, 0, 0] # Hacky way to make sure that compute_max_iou doesn't fail
+            pred_boxes.append(model_preds)
 
+        if i == 100000:
+            from collections import Counter
+            n_box_distr = Counter(n_boxes)
+            print(n_box_distr)
+            break
 
     
     #--------------------------------------------------------------------------------------------
@@ -158,6 +264,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
     emissions_data = tracker.stop_task()
     
     # Calculate classification metrics
+    print("true_labels", true_labels)
+    print("predictions", predictions)
     classification_accuracy = accuracy_score(true_labels, predictions)
     classification_precision = precision_score(true_labels, predictions)
     classification_recall = recall_score(true_labels, predictions)