Fix

app.py

@@ -6,7 +6,33 @@ import cv2
 
 # Configuration
 MODEL_ONNX_PATH = "model.onnx"
-CLASS_NAMES = ["Coffee_Franch_Original_Box_250g", "Coffee_Franch_Instant_Crema_80g"]
+CLASS_NAMES = [
+    'Butter_Dukat_Maslac_Stick_250g',
+    'Butter_Zbregov_Maslac_Stick_250g',
+    'Butter_Zdenka_Maslac_Stick_250g',
+    'Cheese_President_Gouda_Cube_250g',
+    'Chicken_Cekin_Pileca_Prsa_500g',
+    'Coffee_Franch_Crema_Bag_175g',
+    'Coffee_Franch_Crema_Box_250g',
+    'Coffee_Franch_Instant_Crema_80g',
+    'Coffee_Franch_Intense_Box_250g',
+    'Coffee_Franch_Original_Box_250g',
+    'Coffee_Franch_Sensual_Box_250g',
+    'Drink_CocaCola_Original_Bottle_1l',
+    'Flour_Mlineta_Brasno_Ostro_1kg',
+    'Juice_Vindi_Naranca_Nektar_1l',
+    'Ketchup_Zvijezda_Mild_Bottle_500g',
+    'Mayonnaise_Zvijezda_Delicate_Bottle_400g',
+    'Milk_Zbregov_Trajno_28_1l',
+    'Oil_Dijamant_Suncokretovo_Bottle_1l',
+    'Oil_Zvijezda_Suncokretovo_Ulje_1l',
+    'Pasta_Barilla_Fusilli_Box_500g',
+    'Rice_Gallo_Long_Grain_900g',
+    'Rice_Kplus_Arborio_BijeliDugi_1kg',
+    'Salt_SolanaPag_Sitna_Box_1kg',
+    'Spaghetti_PastaZara_Spaghettini_Bag_500g',
+    'Tuna_RioMare_Tonno_Oliva'
+]
 INPUT_SIZE = 640
 
 # Initialize detector

inference.py

@@ -62,49 +62,65 @@ class ObjectDetector:
         blob = canvas.astype(np.float32) / 255.0
         return blob.transpose(2, 0, 1)[None, ...], scale, (pw, ph)
 
-    def postprocess(
-        self,
-        predictions: np.ndarray,
-        original_shape: Tuple[int, int],
-        scale: float,
-        padding: Tuple[int, int],
-        conf_threshold: float = 0.5,
-        iou_threshold: float = 0.45
-    ) -> List[Dict]:
-        predictions = np.squeeze(predictions).T
-        scores = np.max(predictions[:, 4:], axis=1)
-        valid = scores > conf_threshold
-        predictions = predictions[valid]
-        
-        if predictions.size == 0:
-            return []
+def postprocess(
+    self,
+    predictions: np.ndarray,
+    original_shape: Tuple[int, int],
+    scale: float,
+    padding: Tuple[int, int],
+    conf_threshold: float = 0.5,
+    iou_threshold: float = 0.45
+) -> List[Dict]:
+    # Reshape and transpose predictions to (8400, 29)
+    predictions = np.squeeze(predictions).T
+    
+    # Split into box coordinates, objectness score, and class probabilities
+    boxes = predictions[:, :4]
+    obj_scores = predictions[:, 4]
+    cls_scores = predictions[:, 5:5 + len(self.class_names)]  # Dynamic class count
+    
+    # Combine objectness and class scores
+    final_scores = obj_scores[:, None] * cls_scores  # (8400, num_classes)
+    best_class_ids = np.argmax(final_scores, axis=1)
+    best_scores = np.max(final_scores, axis=1)
+    
+    # Filter by confidence threshold
+    valid = best_scores > conf_threshold
+    boxes = boxes[valid]
+    best_scores = best_scores[valid]
+    best_class_ids = best_class_ids[valid]
+    
+    if boxes.size == 0:
+        return []
 
-        boxes = predictions[:, :4]
-        boxes[:, [0, 1]] = boxes[:, [0, 1]] - boxes[:, [2, 3]] / 2
-        boxes[:, [2, 3]] = boxes[:, [0, 1]] + boxes[:, [2, 3]]
-        
-        pad_w, pad_h = padding
-        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
-        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
-        
-        h, w = original_shape
-        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
-        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
-        
-        class_ids = np.argmax(predictions[:, 4:], axis=1)
-        indices = self.nms(boxes, scores[valid], iou_threshold)
-        
-        return [{
-            "class": self.class_names[int(class_ids[i])],
-            "confidence": float(scores[valid][i]),
-            "bbox": boxes[i].tolist(),
-            "bbox_normalized": [
-                float((boxes[i][0] + boxes[i][2])/2 / w),
-                float((boxes[i][1] + boxes[i][3])/2 / h),
-                float((boxes[i][2] - boxes[i][0]) / w),
-                float((boxes[i][3] - boxes[i][1]) / h)
-            ]
-        } for i in indices]
+    # Convert from center_x, center_y, width, height to xyxy format
+    boxes[:, [0, 1]] = boxes[:, [0, 1]] - boxes[:, [2, 3]] / 2
+    boxes[:, [2, 3]] = boxes[:, [0, 1]] + boxes[:, [2, 3]]
+    
+    # Adjust for padding and scale
+    pad_w, pad_h = padding
+    boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / scale
+    boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / scale
+    
+    # Clip boxes to image boundaries
+    h, w = original_shape
+    boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w)
+    boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h)
+    
+    # Perform NMS
+    indices = self.nms(boxes, best_scores, iou_threshold)
+    
+    return [{
+        "class": self.class_names[int(best_class_ids[i])],
+        "confidence": float(best_scores[i]),
+        "bbox": boxes[i].tolist(),
+        "bbox_normalized": [
+            float((boxes[i][0] + boxes[i][2])/2 / w),
+            float((boxes[i][1] + boxes[i][3])/2 / h),
+            float((boxes[i][2] - boxes[i][0]) / w),
+            float((boxes[i][3] - boxes[i][1]) / h)
+        ]
+    } for i in indices]
 
     def predict(self, image: np.ndarray) -> List[Dict]:
         """Main prediction method"""

model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:39582e24d58748bb0304397f40f5d88bff8f7c7a1584bbec11e9f9490fb12690
-size 37908587
+oid sha256:47fb25ba1e7fb96e8577581ecf4811788f36a0117ea9390859635bd89e4d439d
+size 103663101