Create app.py

app.py

@@ -0,0 +1,190 @@
+import gradio as gr
+import cv2
+import numpy as np
+import base64
+import json
+import math
+from PIL import Image
+import io
+import time
+from collections import deque
+
+# Import your pool detector class (you can copy the class definition here)
+# from pool_detector import PoolBallDetector
+
+class PoolBallDetector:
+    # ... (copy the entire PoolBallDetector class from the previous artifact)
+    def __init__(self):
+        self.ball_history = deque(maxlen=10)
+        self.cue_history = deque(maxlen=5)
+        self.table_bounds = None
+        self.setup_detection_params()
+    
+    def setup_detection_params(self):
+        self.ball_colors = {
+            'cue': {'lower': np.array([0, 0, 200]), 'upper': np.array([180, 30, 255])},
+            'black': {'lower': np.array([0, 0, 0]), 'upper': np.array([180, 255, 50])},
+            'solid': {'lower': np.array([0, 50, 50]), 'upper': np.array([10, 255, 255])},
+            'stripe': {'lower': np.array([20, 50, 50]), 'upper': np.array([30, 255, 255])}
+        }
+        self.cue_color = {'lower': np.array([10, 50, 20]), 'upper': np.array([20, 255, 200])}
+    
+    # ... (include all other methods from PoolBallDetector)
+
+# Global detector instance
+detector = PoolBallDetector()
+
+def process_pool_image(image_data):
+    """Process image and return predictions"""
+    try:
+        # Decode base64 image if it's a string
+        if isinstance(image_data, str):
+            image_data = base64.b64decode(image_data)
+            image = Image.open(io.BytesIO(image_data))
+        else:
+            image = image_data
+        
+        # Convert to OpenCV format
+        frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        
+        # Detect components
+        balls = detector.detect_balls(frame)
+        cue_data = detector.detect_cue_stick(frame)
+        
+        # Calculate trajectory
+        trajectory = []
+        if cue_data.get('detected'):
+            trajectory = detector.calculate_trajectory(cue_data, balls)
+        
+        # Calculate metrics
+        shot_angle = cue_data.get('angle', 0) if cue_data.get('detected') else 0
+        shot_power = 0
+        
+        if cue_data.get('detected') and balls:
+            cue_ball = next((ball for ball in balls if ball['type'] == 'cue'), None)
+            if cue_ball:
+                shot_power = detector.calculate_shot_power(cue_data, cue_ball)
+        
+        # Create visualization
+        result_frame = frame.copy()
+        
+        # Draw balls
+        for ball in balls:
+            color = (255, 255, 255) if ball['type'] == 'cue' else (0, 255, 0)
+            cv2.circle(result_frame, (int(ball['x']), int(ball['y'])), int(ball['radius']), color, 2)
+            cv2.putText(result_frame, ball['type'], (int(ball['x']-20), int(ball['y']-30)), 
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
+        
+        # Draw cue stick
+        if cue_data.get('detected'):
+            cv2.line(result_frame, 
+                    (int(cue_data['start_x']), int(cue_data['start_y'])),
+                    (int(cue_data['end_x']), int(cue_data['end_y'])), 
+                    (0, 255, 255), 3)
+        
+        # Draw trajectory
+        if trajectory:
+            trajectory_points = [(int(p['x']), int(p['y'])) for p in trajectory]
+            for i in range(len(trajectory_points) - 1):
+                cv2.line(result_frame, trajectory_points[i], trajectory_points[i+1], (0, 0, 255), 2)
+        
+        # Convert back to PIL for Gradio
+        result_image = Image.fromarray(cv2.cvtColor(result_frame, cv2.COLOR_BGR2RGB))
+        
+        # Prepare text output
+        info_text = f"""
+Detection Results:
+- Cue Detected: {cue_data.get('detected', False)}
+- Balls Found: {len(balls)}
+- Shot Power: {shot_power:.2f}
+- Shot Angle: {shot_angle:.1f}°
+- Trajectory Points: {len(trajectory)}
+        """
+        
+        return result_image, info_text
+        
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def predict_api(image_b64):
+    """API endpoint for mobile app"""
+    try:
+        # Process the image
+        image_data = base64.b64decode(image_b64)
+        image = Image.open(io.BytesIO(image_data))
+        frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+        
+        # Detect components
+        balls = detector.detect_balls(frame)
+        cue_data = detector.detect_cue_stick(frame)
+        trajectory = []
+        
+        if cue_data.get('detected'):
+            trajectory = detector.calculate_trajectory(cue_data, balls)
+        
+        # Calculate metrics
+        shot_angle = cue_data.get('angle', 0) if cue_data.get('detected') else 0
+        shot_power = 0
+        
+        if cue_data.get('detected') and balls:
+            cue_ball = next((ball for ball in balls if ball['type'] == 'cue'), None)
+            if cue_ball:
+                shot_power = detector.calculate_shot_power(cue_data, cue_ball)
+        
+        response = {
+            'timestamp': int(time.time() * 1000),
+            'cue_detected': cue_data.get('detected', False),
+            'balls': balls,
+            'trajectory': trajectory,
+            'power': shot_power,
+            'angle': shot_angle,
+            'table_bounds': detector.table_bounds
+        }
+        
+        if cue_data.get('detected'):
+            response['cue_line'] = {
+                'start_x': cue_data['start_x'],
+                'start_y': cue_data['start_y'],
+                'end_x': cue_data['end_x'],
+                'end_y': cue_data['end_y'],
+                'center_x': cue_data['center_x'],
+                'center_y': cue_data['center_y'],
+                'length': cue_data['length']
+            }
+        
+        return json.dumps(response)
+        
+    except Exception as e:
+        return json.dumps({'error': str(e)})
+
+# Create Gradio interface
+with gr.Blocks(title="8-Ball Pool Trajectory Predictor") as demo:
+    gr.Markdown("# 🎱 8-Ball Pool Trajectory Predictor")
+    gr.Markdown("Upload a screenshot of your 8-ball pool game to get trajectory predictions!")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Pool Table Screenshot")
+            predict_btn = gr.Button("Predict Trajectory", variant="primary")
+        
+        with gr.Column():
+            output_image = gr.Image(label="Prediction Results")
+            output_text = gr.Textbox(label="Detection Info", lines=8)
+    
+    predict_btn.click(
+        fn=process_pool_image,
+        inputs=[input_image],
+        outputs=[output_image, output_text]
+    )
+    
+    # API endpoint for mobile app
+    gr.Interface(
+        fn=predict_api,
+        inputs=gr.Textbox(label="Base64 Image Data"),
+        outputs=gr.Textbox(label="JSON Response"),
+        title="API Endpoint",
+        description="For mobile app integration - send base64 encoded image"
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)