Add NMS and fix conversion

app.py

@@ -79,11 +79,13 @@ class KeypointDetector:
             target_sizes=torch.tensor([(image.height, image.width)]),
             threshold=0.3
         )
+        
+        dets = sv.Detections.from_transformers(results[0]).with_nms(0.5)
 
         # Get boxes and scores for human class (index 0 in COCO dataset)
-        boxes = results[0]["boxes"][results[0]["labels"] == 0]
-        scores = results[0]["scores"][results[0]["labels"] == 0]
-        return boxes.cpu().numpy(), scores.cpu().numpy()
+        boxes = dets.xyxy[dets.class_id == 0]
+        scores = dets.confidence[dets.class_id == 0]
+        return boxes, scores
 
     def detect_keypoints(self, image: Image.Image):
         """Detect keypoints in the image"""
@@ -106,7 +108,7 @@ class KeypointDetector:
 
         # Setup detections for bounding boxes
         detections = sv.Detections(
-            xyxy=self.coco_to_xyxy(boxes),
+            xyxy=boxes,
             confidence=scores,
             class_id=np.array([0]*len(scores))
         )