Update app.py

app.py

@@ -5,76 +5,114 @@ import torch
 from transformers import DetrForObjectDetection, DetrImageProcessor
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Load pre-trained model for ball detection
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", cache_dir="/home/user/app/cache")
-model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", cache_dir="/home/user/app/cache")
+try:
+    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", cache_dir="/home/user/app/cache")
+    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", cache_dir="/home/user/app/cache")
+    logger.info("Model and processor loaded successfully")
+except Exception as e:
+    logger.error(f"Error loading model: {str(e)}")
+    raise
 
 def process_drs_video(video_path):
-    # Extract frames from video
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frames.append(frame)
-    cap.release()
-
-    # Detect ball in frames
-    ball_positions = []
-    for frame in frames:
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        inputs = processor(images=frame_rgb, return_tensors="pt")
-        outputs = model(**inputs)
-        target_sizes = torch.tensor([frame_rgb.shape[:2]])
-        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+    try:
+        # Extract frames from video
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            logger.error("Failed to open video file")
+            return "Error: Could not open video file", None
         
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            if score > 0.9:  # High confidence threshold
-                x_center = (box[0] + box[2]) / 2
-                y_center = (box[1] + box[3]) / 2
-                # Simplified depth estimation (z-coordinate, placeholder)
-                z = 100  # Replace with actual depth model or pitch mapping
-                ball_positions.append([x_center.item(), y_center.item(), z])
+        frames = []
+        frame_count = 0
+        while cap.isOpened() and frame_count < 100:  # Limit to 100 frames to avoid memory issues
+            ret, frame = cap.read()
+            if not ret:
                 break
-        else:
-            ball_positions.append(None)
-    
-    ball_positions = [pos for pos in ball_positions if pos is not None]
-    if not ball_positions:
-        return "Error: No ball detected in video", None
+            # Resize frame to reduce memory usage
+            frame = cv2.resize(frame, (640, 360))
+            frames.append(frame)
+            frame_count += 1
+        cap.release()
+        logger.info(f"Extracted {len(frames)} frames from video")
 
-    trajectory = np.array(ball_positions)
+        if not frames:
+            logger.error("No frames extracted from video")
+            return "Error: No frames extracted from video", None
+
+        # Detect ball in frames
+        ball_positions = []
+        for i, frame in enumerate(frames[::2]):  # Process every 2nd frame to reduce load
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            inputs = processor(images=frame_rgb, return_tensors="pt")
+            with torch.no_grad():  # Disable gradients for inference
+                outputs = model(**inputs)
+            
+            target_sizes = torch.tensor([frame_rgb.shape[:2]])
+            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.7)[0]  # Lowered threshold
+            
+            ball_found = False
+            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                if score > 0.7:  # Adjusted confidence threshold
+                    x_center = (box[0] + box[2]) / 2
+                    y_center = (box[1] + box[3]) / 2
+                    z = 100  # Placeholder for depth (replace with depth model if available)
+                    ball_positions.append([x_center.item(), y_center.item(), z])
+                    ball_found = True
+                    logger.info(f"Ball detected in frame {i}: score={score:.2f}, box={box.tolist()}")
+                    break
+            if not ball_found:
+                ball_positions.append(None)
+                logger.warning(f"No ball detected in frame {i}")
+        
+        ball_positions = [pos for pos in ball_positions if pos is not None]
+        if not ball_positions:
+            logger.error("No ball detected in any frame")
+            return "Error: No ball detected in video", None
 
-    # Simplified LBW decision logic (without edge detection)
-    # Assumptions: Ball pitches in line, no bat contact (due to no audio), check stump impact
-    pitching_in_line = trajectory[0][0] > 200 and trajectory[0][0] < 520  # Placeholder: Adjust based on pitch dimensions
-    impact_in_line = trajectory[-1][0] > 200 and trajectory[-1][0] < 520  # Check impact near stumps
-    hits_stumps = trajectory[-1][1] < 300  # Simplified: Ball low enough to hit stumps
-    decision = "Out" if pitching_in_line and impact_in_line and hits_stumps else "Not Out"
+        trajectory = np.array(ball_positions)
+        logger.info(f"Trajectory shape: {trajectory.shape}")
 
-    # 3D Trajectory Plot
-    fig = plt.figure(figsize=(10, 5))
-    ax = fig.add_subplot(121, projection='3d')
-    ax.plot(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2], 'r-')
-    ax.set_xlabel("X (Pitch Width)")
-    ax.set_ylabel("Y (Pitch Length)")
-    ax.set_zlabel("Z (Height)")
-    ax.set_title("3D Ball Trajectory")
+        # Simplified LBW decision logic
+        pitching_in_line = trajectory[0][0] > 100 and trajectory[0][0] < 540  # Adjusted for 640x360 resolution
+        impact_in_line = trajectory[-1][0] > 100 and trajectory[-1][0] < 540
+        hits_stumps = trajectory[-1][1] < 200  # Adjusted for lower resolution
+        decision = "Out" if pitching_in_line and impact_in_line and hits_stumps else "Not Out"
+        logger.info(f"LBW Decision: {decision}, Pitching: {pitching_in_line}, Impact: {impact_in_line}, Stumps: {hits_stumps}")
+
+        # 3D Trajectory Plot
+        fig = plt.figure(figsize=(10, 5))
+        ax = fig.add_subplot(121, projection='3d')
+        ax.plot(trajectory[:, 0], trajectory[:, 1], trajectory[:, 2], 'r-')
+        ax.set_xlabel("X (Pitch Width)")
+        ax.set_ylabel("Y (Pitch Length)")
+        ax.set_zlabel("Z (Height)")
+        ax.set_title("3D Ball Trajectory")
+
+        # Pitch Map Plot
+        ax2 = fig.add_subplot(122)
+        ax2.scatter(trajectory[:, 0], trajectory[:, 1], c='blue', marker='o')
+        ax2.set_xlim(0, 640)  # Match video resolution
+        ax2.set_ylim(0, 360)
+        ax2.set_xlabel("Pitch Width")
+        ax2.set_ylabel("Pitch Length")
+        ax2.set_title("Pitch Map")
+        
+        output_path = "drs_output.png"
+        plt.savefig(output_path)
+        plt.close()
+        logger.info(f"Output saved to {output_path}")
 
-    # Pitch Map Plot
-    ax2 = fig.add_subplot(122)
-    ax2.scatter(trajectory[:, 0], trajectory[:, 1], c='blue', marker='o')
-    ax2.set_xlim(0, 720)  # Adjust based on video resolution or pitch dimensions
-    ax2.set_ylim(0, 1280)
-    ax2.set_xlabel("Pitch Width")
-    ax2.set_ylabel("Pitch Length")
-    ax2.set_title("Pitch Map")
-    plt.savefig("drs_output.png")
-    plt.close()
+        return decision, output_path
 
-    return decision, "drs_output.png"
+    except Exception as e:
+        logger.error(f"Error processing video: {str(e)}")
+        return f"Error: {str(e)}", None
 
 # Gradio interface
 iface = gr.Interface(