Update app.py

app.py

@@ -4,14 +4,13 @@ import cv2
 import numpy as np
 from PIL import Image
 import json
+from fastapi import FastAPI
 
 # Load your custom YOLO model
 model = YOLO("fentanyl_oft.pt")
 
 def detect_keypoints(image):
-    """
-    Run YOLO inference and return keypoints data
-    """
+    # Your existing detection code...
     try:
         # Convert PIL Image to numpy array
         if isinstance(image, Image.Image):
@@ -20,7 +19,7 @@ def detect_keypoints(image):
         else:
             image_cv2 = image
         
-        # Run inference with your exact parameters
+        # Run inference
         results = model.predict(
             source=image_cv2,
             conf=0.05,
@@ -32,21 +31,21 @@ def detect_keypoints(image):
         )
         
         keypoints_data = []
-        
-        for result in results:
-            if result.keypoints is not None and len(result.keypoints.data) > 0:
-                # Extract keypoints from the first detection
-                kpts = result.keypoints.data[0]  # Shape: [num_keypoints, 3] (x, y, confidence)
-                
-                h, w = image_cv2.shape[:2]
+        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, (x, y, conf) in enumerate(kpts):
-                    if conf > 0.3:  # Confidence threshold for individual keypoints
+                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": int(i),
+                            "id": i,
                             "x": float(x),
-                            "y": float(y), 
-                            "confidence": float(conf)
+                            "y": float(y),
+                            "confidence": float(confidence)
                         })
         
         return {
@@ -58,20 +57,28 @@ def detect_keypoints(image):
         }
         
     except Exception as e:
-        return {
-            "success": False,
-            "error": str(e),
-            "keypoints": []
-        }
+        return {"success": False, "error": str(e)}
+
+# Create FastAPI app for API endpoints
+app = FastAPI()
+
+@app.post("/api/detect")
+async def api_detect_keypoints(file: bytes):
+    try:
+        # Convert bytes to PIL Image
+        image = Image.open(io.BytesIO(file))
+        result = detect_keypoints(image)
+        return result
+    except Exception as e:
+        return {"success": False, "error": str(e)}
 
-# Create simple Gradio interface
-interface = gr.Interface(
+# Create Gradio interface
+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"
+    title="YOLO Keypoint Detection"
 )
 
-if __name__ == "__main__":
-    interface.launch()
+# Mount Gradio on FastAPI
+app = gr.mount_gradio_app(app, iface, path="/")