Minor fixes

.gitignore

@@ -0,0 +1,7 @@
+.venv/
+out/*
+!out/README.md
+__pycache__/
+
+# ignore model weight in github, they will be re-downloaded
+*.pth

example.py

@@ -15,6 +15,40 @@ from pathlib import Path
 # Import our RF-DETR SoccerNet class
 from inference import RFDETRSoccerNet
 
+# for drawing result
+from PIL import Image, ImageDraw, ImageFont
+
+def draw_detections_on_image(image, df):
+    """Draw bounding boxes on PIL image"""
+    draw = ImageDraw.Draw(image)
+    
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 16)
+    except:
+        font = ImageFont.load_default()
+    
+    colors = {
+        'ball': (255, 0, 0),
+        'player': (0, 255, 0),
+        'referee': (255, 255, 0),
+        'goalkeeper': (0, 0, 255)
+    }
+    
+    for _, row in df.iterrows():
+        x1, y1, x2, y2 = row['x1'], row['y1'], row['x2'], row['y2']
+        class_name = row['class_name']
+        conf = row['confidence']
+        color = colors.get(class_name, (255, 255, 255))
+        
+        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
+        
+        text = f"{class_name}: {conf:.2f}"
+        bbox = draw.textbbox((x1, y1-20), text, font=font)
+        draw.rectangle([bbox[0]-2, bbox[1]-2, bbox[2]+2, bbox[3]+2], fill=color)
+        draw.text((x1, y1-20), text, fill=(0, 0, 0), font=font)
+    
+    return image
+
 
 def basic_image_analysis():
     """Basic image processing example using included sample."""
@@ -44,6 +78,12 @@ def basic_image_analysis():
     # Display results
     print(f"\n📊 ANALYSIS RESULTS")
     print(f"Total detections: {len(df):,}")
+
+    # save result image with drawn boxes
+    image = Image.open(image_path).convert("RGB")
+    result_image = draw_detections_on_image(image, df)
+    result_image.save("out/image_analysis_result.jpg")
+
     
     if len(df) > 0:
         # Show detections by class
@@ -55,8 +95,8 @@ def basic_image_analysis():
         
         # Save results
         print(f"\n💾 Saving results...")
-        model.save_results(df, "image_analysis.csv", format="csv")
-        model.save_results(df, "image_analysis.json", format="json")
+        model.save_results(df, "out/image_analysis.csv", format="csv")
+        model.save_results(df, "out/image_analysis.json", format="json")
     else:
         print("No detections found. Try lowering the confidence threshold.")
     
@@ -136,7 +176,7 @@ def ball_possession_example():
         print(f"  Frames with possession: {possession_df['frame'].nunique()}")
         
         # Save possession analysis
-        model.save_results(possession_df, "ball_possession_analysis.csv")
+        model.save_results(possession_df, "out/ball_possession_analysis.csv")
         print(f"💾 Possession analysis saved to: ball_possession_analysis.csv")
     else:
         print("No possession events found. Try increasing the distance threshold.")

inference.py

@@ -13,7 +13,7 @@ import cv2
 import pandas as pd
 import numpy as np
 import torch
-from rfdetr import RFDETRBase
+from rfdetr import RFDETRLarge
 from PIL import Image
 from typing import Union, Optional, List, Dict, Tuple
 import os
@@ -84,7 +84,7 @@ class RFDETRSoccerNet:
             print(f"📦 Loading model from {self.model_path}...")
             
             # Initialize base model
-            self.model = RFDETRBase()
+            self.model = RFDETRLarge()
             
             # Reinitialize detection head for 4 classes (critical for compatibility)
             print(f"🔧 Reinitializing detection head for {self.num_classes} classes...")

out/README.md

@@ -0,0 +1 @@
+This directory is for output files.