Update app.py

app.py

@@ -1,3 +1,81 @@
+# import gradio as gr
+# import io
+# from ultralytics import YOLO
+# import cv2
+# import numpy as np
+# from PIL import Image
+# import json
+
+# # Load your custom YOLO model
+# model = YOLO("fentanyl_oft.pt")
+# # model = YOLO("avatar_ckpt.pt")
+
+# def detect_keypoints(image):
+#     """
+#     Run YOLO inference and return keypoints data
+#     """
+#     try:
+#         # Convert PIL Image to numpy array
+#         if isinstance(image, Image.Image):
+#             image_np = np.array(image)
+#             image_cv2 = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+#         else:
+#             image_cv2 = image
+        
+#         # Run inference
+#         results = model.predict(
+#             source=image_cv2,
+#             conf=0.05,
+#             iou=0.7,
+#             max_det=1,
+#             imgsz=1440,
+#             device='cpu',
+#             verbose=False
+#         )
+        
+#         keypoints_data = []
+#         if results and len(results) > 0:
+#             result = results[0]
+#             if result.keypoints is not None:
+#                 kpts = result.keypoints.xy.cpu().numpy()
+#                 conf = result.keypoints.conf.cpu().numpy()
+                
+#                 for i in range(kpts.shape[1]):
+#                     if i < len(kpts[0]):
+#                         x, y = kpts[0][i]
+#                         confidence = conf[0][i] if i < len(conf[0]) else 0.0
+#                         keypoints_data.append({
+#                             "id": i,
+#                             "x": float(x),
+#                             "y": float(y),
+#                             "confidence": float(confidence)
+#                         })
+        
+#         return {
+#             "success": True,
+#             "keypoints": keypoints_data,
+#             "image_width": image_cv2.shape[1],
+#             "image_height": image_cv2.shape[0],
+#             "num_keypoints": len(keypoints_data)
+#         }
+        
+#     except Exception as e:
+#         return {"success": False, "error": str(e)}
+
+# # Create Gradio interface with API access enabled
+# iface = gr.Interface(
+#     fn=detect_keypoints,
+#     inputs=gr.Image(type="pil"),
+#     outputs=gr.JSON(),
+#     title="YOLO Keypoint Detection",
+#     description="Upload an image to detect keypoints using custom YOLO model",
+#     api_name="predict"  # This enables API access at /api/predict
+# )
+
+# # Launch with API enabled
+# if __name__ == "__main__":
+#     iface.launch(share=False)
+
 import gradio as gr
 import io
 from ultralytics import YOLO
@@ -6,15 +84,28 @@ import numpy as np
 from PIL import Image
 import json
 
-# Load your custom YOLO model
-model = YOLO("fentanyl_oft.pt")
-# model = YOLO("avatar_ckpt.pt")
+# Load both models
+single_animal_model = YOLO("fentanyl_oft.pt")  # Single animal model
+multi_animal_model = YOLO("avatar_ckpt.pt")   # Multi-animal model
 
-def detect_keypoints(image):
+def detect_keypoints(image, mode="single"):
     """
     Run YOLO inference and return keypoints data
+    Args:
+        image: PIL Image
+        mode: "single" or "multi" to determine which model to use
     """
     try:
+        # Select model and parameters based on mode
+        if mode == "multi":
+            model = multi_animal_model
+            imgsz = 1504
+            max_det = 5
+        else:  # default to single
+            model = single_animal_model
+            imgsz = 1440
+            max_det = 1
+        
         # Convert PIL Image to numpy array
         if isinstance(image, Image.Image):
             image_np = np.array(image)
@@ -22,13 +113,13 @@ def detect_keypoints(image):
         else:
             image_cv2 = image
         
-        # Run inference
+        # Run inference with mode-specific parameters
         results = model.predict(
             source=image_cv2,
             conf=0.05,
             iou=0.7,
-            max_det=1,
-            imgsz=1440,
+            max_det=max_det,
+            imgsz=imgsz,
             device='cpu',
             verbose=False
         )
@@ -40,35 +131,54 @@ def detect_keypoints(image):
                 kpts = result.keypoints.xy.cpu().numpy()
                 conf = result.keypoints.conf.cpu().numpy()
                 
-                for i in range(kpts.shape[1]):
-                    if i < len(kpts[0]):
-                        x, y = kpts[0][i]
-                        confidence = conf[0][i] if i < len(conf[0]) else 0.0
+                # Handle multiple detections (for multi-animal mode)
+                for detection_idx in range(kpts.shape[0]):
+                    detection_keypoints = []
+                    for i in range(kpts.shape[1]):
+                        if i < len(kpts[detection_idx]):
+                            x, y = kpts[detection_idx][i]
+                            confidence = conf[detection_idx][i] if i < len(conf[detection_idx]) else 0.0
+                            detection_keypoints.append({
+                                "id": i,
+                                "x": float(x),
+                                "y": float(y),
+                                "confidence": float(confidence)
+                            })
+                    
+                    # For single animal mode, flatten the structure
+                    if mode == "single":
+                        keypoints_data = detection_keypoints
+                        break  # Only take first detection
+                    else:
+                        # For multi-animal mode, keep detection structure
                         keypoints_data.append({
-                            "id": i,
-                            "x": float(x),
-                            "y": float(y),
-                            "confidence": float(confidence)
+                            "detection_id": detection_idx,
+                            "keypoints": detection_keypoints
                         })
         
         return {
             "success": True,
+            "mode": mode,
             "keypoints": keypoints_data,
             "image_width": image_cv2.shape[1],
             "image_height": image_cv2.shape[0],
-            "num_keypoints": len(keypoints_data)
+            "num_detections": len(keypoints_data) if mode == "multi" else (1 if keypoints_data else 0),
+            "num_keypoints": len(keypoints_data) if mode == "single" else sum(len(det["keypoints"]) for det in keypoints_data) if mode == "multi" else 0
         }
         
     except Exception as e:
-        return {"success": False, "error": str(e)}
+        return {"success": False, "error": str(e), "mode": mode}
 
-# Create Gradio interface with API access enabled
+# Create Gradio interface with mode parameter
 iface = gr.Interface(
     fn=detect_keypoints,
-    inputs=gr.Image(type="pil"),
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Dropdown(choices=["single", "multi"], value="single", label="Detection Mode")
+    ],
     outputs=gr.JSON(),
     title="YOLO Keypoint Detection",
-    description="Upload an image to detect keypoints using custom YOLO model",
+    description="Upload an image to detect keypoints using custom YOLO model. Choose single or multi-animal mode.",
     api_name="predict"  # This enables API access at /api/predict
 )