Update app.py

app.py

@@ -15,22 +15,22 @@ with open(LABELS_PATH, "r") as f:
 
 img_transforms = transforms.ToTensor()
 
-# inter_class_nms implemented by GPT
-def inter_class_nms(boxes, scores, iou_threshold=0.5):
+def inter_class_nms(boxes, scores, labels, iou_threshold=0.5):
     # Perform non-maximum suppression
     keep = nms(boxes, scores, iou_threshold)
 
     # Filter boxes and scores
     new_boxes = boxes[keep]
     new_scores = scores[keep]
+    new_labels = labels[keep]
 
     # Return the result in a dictionary
-    return {'boxes': new_boxes, 'scores': new_scores}
+    return {'boxes': new_boxes, 'scores': new_scores, 'labels': new_labels}
 
 def predict(img, conf_thresh=0.4):
     img_input = [img_transforms(img)]
     _, pred = model(img_input)
-    pred = [inter_class_nms(pred[0]['boxes'], pred[0]['scores'])]
+    pred = [inter_class_nms(pred[0]['boxes'], pred[0]['scores'], pred[0]['labels'])]
     out_img = img.copy()
     draw = ImageDraw.Draw(out_img)
     font = ImageFont.truetype("res/Tuffy_Bold.ttf", 25)