Create app.py

app.py

@@ -0,0 +1,180 @@
+import gradio as gr
+import os
+from ultralytics import YOLO
+import numpy as np
+import json
+from PIL import Image, ImageDraw
+
+# Define keypoints we need for rigging
+KEYPOINTS = {
+    0: {"name": "chin (nose)"},
+    7: {"name": "left_elbow"},
+    8: {"name": "right_elbow"},
+    9: {"name": "left_wrist"},
+    10: {"name": "right_wrist"},
+    13: {"name": "left_knee"},
+    14: {"name": "right_knee"}
+}
+
+# Initialize model
+model = None
+
+def load_model():
+    """Load the YOLO pose estimation model"""
+    global model
+    if model is None:
+        model_path = 'yolov8s-pose.pt'
+        if os.path.exists(model_path):
+            try:
+                model = YOLO(model_path)
+                print("Model loaded successfully")
+            except Exception as e:
+                print(f"Error loading model: {e}")
+                model = None
+        else:
+            print(f"Model file not found: {model_path}")
+    return model
+
+def process_image(input_image):
+    """
+    Process an image for pose estimation and return keypoint coordinates
+    
+    Args:
+        input_image: Input image (PIL Image or numpy array)
+    
+    Returns:
+        Tuple of (visualization image, JSON results string)
+    """
+    # Load model if not already loaded
+    if load_model() is None:
+        return None, json.dumps({"error": "Model not available"})
+    
+    try:
+        # Convert to PIL if needed
+        if not isinstance(input_image, np.ndarray):
+            input_image = np.array(input_image)
+        
+        # Run inference
+        results = model.predict(input_image, verbose=False)
+        
+        # Process keypoint data
+        keypoint_data = {}
+        
+        if not results or len(results) == 0:
+            return input_image, json.dumps({"error": "No pose detection results found"})
+            
+        result = results[0]
+        
+        if not hasattr(result, "keypoints") or result.keypoints is None:
+            return input_image, json.dumps({"error": "No keypoints detected in the image"})
+            
+        try:
+            keypoints = result.keypoints.data.cpu().numpy()
+        except AttributeError:
+            return input_image, json.dumps({"error": "Error accessing keypoints data"})
+            
+        if len(keypoints) == 0:
+            return input_image, json.dumps({"error": "No people detected in the image"})
+            
+        # Get first person's keypoints
+        kp = keypoints[0]
+        
+        # Extract keypoints
+        for idx, keypoint_info in KEYPOINTS.items():
+            if idx < len(kp) and kp[idx][2] > 0.5:  # Confidence threshold
+                x, y, conf = kp[idx]
+                keypoint_data[keypoint_info["name"]] = {
+                    "x": int(x),
+                    "y": int(y),
+                    "confidence": float(conf)
+                }
+        
+        # Add groin point (midpoint between points 11 and 12)
+        if len(kp) > 12 and kp[11][2] > 0.5 and kp[12][2] > 0.5:
+            groin_x = int((kp[11][0] + kp[12][0]) / 2)
+            groin_y = int((kp[11][1] + kp[12][1]) / 2)
+            groin_conf = (float(kp[11][2]) + float(kp[12][2])) / 2
+            keypoint_data["groin"] = {
+                "x": groin_x,
+                "y": groin_y,
+                "confidence": groin_conf
+            }
+        
+        # Create visualization image
+        vis_image = Image.fromarray(input_image.copy())
+        draw = ImageDraw.Draw(vis_image)
+        
+        # Draw keypoints
+        for point_name, point_data in keypoint_data.items():
+            x, y = point_data["x"], point_data["y"]
+            # Draw a circle at each keypoint
+            radius = 5
+            draw.ellipse(
+                [(x - radius, y - radius), (x + radius, y + radius)],
+                fill="red"
+            )
+            # Add text label
+            draw.text((x + 10, y), point_name, fill="black")
+        
+        return np.array(vis_image), json.dumps({"keypoints": keypoint_data}, indent=2)
+    
+    except Exception as e:
+        return input_image, json.dumps({"error": f"Error processing image: {str(e)}"})
+
+# Create Gradio interface
+def create_gradio_app():
+    with gr.Blocks() as demo:
+        gr.Markdown("# YOLO Pose Estimation API")
+        gr.Markdown("Upload an image to detect pose keypoints")
+        
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="numpy", label="Input Image")
+                submit_btn = gr.Button("Process Image")
+            
+            with gr.Column():
+                output_image = gr.Image(label="Visualization")
+                output_json = gr.JSON(label="Keypoint Data")
+        
+        submit_btn.click(
+            fn=process_image,
+            inputs=[input_image],
+            outputs=[output_image, output_json]
+        )
+        
+        # Add API documentation
+        gr.Markdown("""
+        ## API Usage
+        
+        This Gradio app also provides a REST API endpoint at `/api/predict`.
+        
+        Example usage:
+        ```python
+        import requests
+        
+        # Send a POST request to the API endpoint
+        response = requests.post(
+            "YOUR_HUGGINGFACE_SPACE_URL/api/predict",
+            files={"input_image": open("image.jpg", "rb")}
+        )
+        
+        # Process results
+        if response.status_code == 200:
+            results = response.json()
+            keypoints = results.get("keypoints", {})
+            print(keypoints)
+        else:
+            print(f"Error: {response.text}")
+        ```
+        """)
+    
+    return demo
+
+demo = create_gradio_app()
+
+# Launch app
+if __name__ == "__main__":
+    demo.launch()
+else:
+    # For Hugging Face Spaces
+    demo.launch(share=False)