Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 import torch
-from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
 from PIL import Image, ImageDraw, ImageFont
 import cv2
 import numpy as np
 import pandas as pd
 from pathlib import Path
 import tempfile
+from huggingface_hub import hf_hub_download
+import os
 
 # Class labels for the model
 CLASS_NAMES = {
@@ -30,35 +31,46 @@ class SoccerDetector:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {self.device}")
         
-        # Load the model and processor
-        model_name = "julianzu9612/RFDETR-Soccernet"
-        self.processor = RTDetrImageProcessor.from_pretrained(model_name)
-        self.model = RTDetrForObjectDetection.from_pretrained(model_name).to(self.device)
-        self.model.eval()
-        print("Model loaded successfully!")
+        try:
+            # Try to download the model file from Hugging Face
+            print("Downloading model from Hugging Face...")
+            model_path = hf_hub_download(
+                repo_id="julianzu9612/RFDETR-Soccernet",
+                filename="best.pt"  # or "model.pt" - we'll need to check
+            )
+            print(f"Model downloaded to: {model_path}")
+            
+            # Load with Ultralytics YOLO (RF-DETR is YOLO-based)
+            from ultralytics import RTDETR
+            self.model = RTDETR(model_path)
+            print("Model loaded successfully!")
+            
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            print("\nTrying alternative loading method...")
+            
+            # Alternative: Try loading directly from hub
+            from ultralytics import RTDETR
+            try:
+                self.model = RTDETR("julianzu9612/RFDETR-Soccernet")
+                print("Model loaded via direct hub access!")
+            except Exception as e2:
+                print(f"Alternative method failed: {e2}")
+                raise Exception("Could not load model. Please check the model repository structure.")
     
     def process_image(self, image, confidence_threshold=0.5):
         """Process a single image and return detections"""
         # Convert to PIL if needed
         if isinstance(image, np.ndarray):
+            image_array = image
             image = Image.fromarray(image)
-        
-        # Prepare image
-        inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+        else:
+            image_array = np.array(image)
         
         # Run inference
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        
-        # Post-process
-        target_sizes = torch.tensor([image.size[::-1]]).to(self.device)
-        results = self.processor.post_process_object_detection(
-            outputs, 
-            target_sizes=target_sizes, 
-            threshold=confidence_threshold
-        )[0]
+        results = self.model(image_array, conf=confidence_threshold, verbose=False)
         
-        return results, image
+        return results[0], image
     
     def draw_detections(self, image, results):
         """Draw bounding boxes on image"""
@@ -69,19 +81,28 @@ class SoccerDetector:
         except:
             font = ImageFont.load_default()
         
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            box = box.cpu().numpy()
-            x1, y1, x2, y2 = box
+        # Get boxes, scores, and classes
+        boxes = results.boxes
+        
+        if boxes is None or len(boxes) == 0:
+            return image
+        
+        for box in boxes:
+            # Get coordinates
+            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+            
+            # Get class and confidence
+            cls = int(box.cls[0].cpu().numpy())
+            conf = float(box.conf[0].cpu().numpy())
             
-            label_id = label.item()
-            class_name = CLASS_NAMES.get(label_id, f"class_{label_id}")
-            color = CLASS_COLORS.get(label_id, (255, 255, 255))
+            class_name = CLASS_NAMES.get(cls, f"class_{cls}")
+            color = CLASS_COLORS.get(cls, (255, 255, 255))
             
             # Draw box
             draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
             
             # Draw label
-            text = f"{class_name}: {score:.2f}"
+            text = f"{class_name}: {conf:.2f}"
             
             # Draw text background
             bbox = draw.textbbox((x1, y1), text, font=font)
@@ -90,21 +111,29 @@ class SoccerDetector:
         
         return image
     
-    def create_detections_dataframe(self, results, image_size):
+    def create_detections_dataframe(self, results):
         """Create a pandas DataFrame from detection results"""
         data = []
         
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            box = box.cpu().numpy()
-            x1, y1, x2, y2 = box
+        boxes = results.boxes
+        
+        if boxes is None or len(boxes) == 0:
+            return pd.DataFrame()
+        
+        for box in boxes:
+            # Get coordinates
+            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+            
+            # Get class and confidence
+            cls = int(box.cls[0].cpu().numpy())
+            conf = float(box.conf[0].cpu().numpy())
             
-            label_id = label.item()
-            class_name = CLASS_NAMES.get(label_id, f"class_{label_id}")
+            class_name = CLASS_NAMES.get(cls, f"class_{cls}")
             
             data.append({
                 'class_name': class_name,
-                'class_id': label_id,
-                'confidence': float(score),
+                'class_id': cls,
+                'confidence': conf,
                 'x1': float(x1),
                 'y1': float(y1),
                 'x2': float(x2),
@@ -153,7 +182,7 @@ class SoccerDetector:
                 rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 
                 # Run detection
-                results, _ = self.process_image(rgb_frame, confidence_threshold)
+                results = self.model(rgb_frame, conf=confidence_threshold, verbose=False)[0]
                 
                 # Draw detections
                 pil_image = Image.fromarray(rgb_frame)
@@ -161,28 +190,31 @@ class SoccerDetector:
                 annotated_frame = cv2.cvtColor(np.array(annotated_image), cv2.COLOR_RGB2BGR)
                 
                 # Save detections to list
-                for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-                    box = box.cpu().numpy()
-                    x1, y1, x2, y2 = box
-                    label_id = label.item()
-                    class_name = CLASS_NAMES.get(label_id, f"class_{label_id}")
-                    
-                    all_detections.append({
-                        'frame': frame_num,
-                        'timestamp': frame_num / fps,
-                        'class_name': class_name,
-                        'class_id': label_id,
-                        'confidence': float(score),
-                        'x1': float(x1),
-                        'y1': float(y1),
-                        'x2': float(x2),
-                        'y2': float(y2),
-                        'width': float(x2 - x1),
-                        'height': float(y2 - y1),
-                        'center_x': float((x1 + x2) / 2),
-                        'center_y': float((y1 + y2) / 2),
-                        'area': float((x2 - x1) * (y2 - y1))
-                    })
+                boxes = results.boxes
+                if boxes is not None and len(boxes) > 0:
+                    for box in boxes:
+                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+                        cls = int(box.cls[0].cpu().numpy())
+                        conf = float(box.conf[0].cpu().numpy())
+                        
+                        class_name = CLASS_NAMES.get(cls, f"class_{cls}")
+                        
+                        all_detections.append({
+                            'frame': frame_num,
+                            'timestamp': frame_num / fps,
+                            'class_name': class_name,
+                            'class_id': cls,
+                            'confidence': conf,
+                            'x1': float(x1),
+                            'y1': float(y1),
+                            'x2': float(x2),
+                            'y2': float(y2),
+                            'width': float(x2 - x1),
+                            'height': float(y2 - y1),
+                            'center_x': float((x1 + x2) / 2),
+                            'center_y': float((y1 + y2) / 2),
+                            'area': float((x2 - x1) * (y2 - y1))
+                        })
                 
                 out.write(annotated_frame)
             else:
@@ -214,7 +246,7 @@ def process_image_interface(image, confidence_threshold):
     
     results, original_image = detector.process_image(image, confidence_threshold)
     annotated_image = detector.draw_detections(original_image.copy(), results)
-    df = detector.create_detections_dataframe(results, original_image.size)
+    df = detector.create_detections_dataframe(results)
     
     return annotated_image, df