Create app.py

app.py

@@ -0,0 +1,205 @@
+import os
+from collections import deque
+from typing import Optional, List
+from io import BytesIO
+import base64
+
+import cv2
+import numpy as np
+from PIL import Image
+
+import torch
+from tqdm import tqdm
+
+import supervision as sv
+from inference_sdk import get_model # Modified import
+from sports.common.team import TeamClassifier
+from sports.common.view import ViewTransformer
+from sports.annotators.soccer import draw_pitch, draw_points_on_pitch
+from sports.configs.soccer import SoccerPitchConfiguration
+
+import gradio as gr
+import plotly.graph_objects as go
+from transformers import AutoProcessor, SiglipVisionModel
+from more_itertools import chunked
+from sklearn.cluster import KMeans
+import umap
+
+# ==============================================
+# Environment variables
+# ==============================================
+HF_TOKEN = os.environ.get("HF_TOKEN")
+ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY")
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# ==============================================
+# Load Detection Models
+# ==============================================
+PLAYER_DETECTION_MODEL_ID = "football-players-detection-3zvbc/11"
+PLAYER_DETECTION_MODEL = get_model(model_id=PLAYER_DETECTION_MODEL_ID, api_key=ROBOFLOW_API_KEY)
+
+FIELD_DETECTION_MODEL_ID = "football-field-detection-f07vi/14"
+FIELD_DETECTION_MODEL = get_model(model_id=FIELD_DETECTION_MODEL_ID, api_key=ROBOFLOW_API_KEY)
+
+team_classifier = TeamClassifier(device=DEVICE)
+CONFIG = SoccerPitchConfiguration()
+
+# ==============================================
+# Load SigLIP Model
+# ==============================================
+SIGLIP_MODEL_PATH = 'google/siglip-base-patch16-224'
+EMBEDDINGS_MODEL = SiglipVisionModel.from_pretrained(SIGLIP_MODEL_PATH).to(DEVICE)
+EMBEDDINGS_PROCESSOR = AutoProcessor.from_pretrained(SIGLIP_MODEL_PATH)
+
+# ==============================================
+# Helper Functions
+# ==============================================
+
+def resolve_goalkeepers_team_id(players: sv.Detections, goalkeepers: sv.Detections) -> np.ndarray:
+    goalkeepers_xy = goalkeepers.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+    players_xy = players.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+    team_0_centroid = players_xy[players.class_id == 0].mean(axis=0)
+    team_1_centroid = players_xy[players.class_id == 1].mean(axis=0)
+    return np.array([0 if np.linalg.norm(gk - team_0_centroid) < np.linalg.norm(gk - team_1_centroid) else 1 for gk in goalkeepers_xy])
+
+def pil_image_to_data_uri(image: Image.Image) -> str:
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{img_str}"
+
+def create_umap_3d_plot(crops: List[Image.Image]):
+    BATCH_SIZE = 32
+    crops = [sv.cv2_to_pillow(crop) for crop in crops]
+    batches = list(chunked(crops, BATCH_SIZE))
+    data = []
+    with torch.no_grad():
+        for batch in tqdm(batches, desc='embedding extraction'):
+            inputs = EMBEDDINGS_PROCESSOR(images=batch, return_tensors="pt").to(DEVICE)
+            outputs = EMBEDDINGS_MODEL(**inputs)
+            embeddings = torch.mean(outputs.last_hidden_state, dim=1).cpu().numpy()
+            data.append(embeddings)
+    data = np.concatenate(data)
+
+    # UMAP and clustering
+    REDUCER = umap.UMAP(n_components=3)
+    CLUSTERING_MODEL = KMeans(n_clusters=2)
+    projections = REDUCER.fit_transform(data)
+    clusters = CLUSTERING_MODEL.fit_predict(projections)
+
+    # Prepare image data URIs
+    image_data_uris = {f"image_{i}": pil_image_to_data_uri(image) for i, image in enumerate(crops)}
+    image_ids = np.array([f"image_{i}" for i in range(len(crops))])
+
+    # Plotly 3D scatter
+    traces = []
+    for lbl in np.unique(clusters):
+        mask = clusters == lbl
+        trace = go.Scatter3d(
+            x=projections[mask][:,0],
+            y=projections[mask][:,1],
+            z=projections[mask][:,2],
+            mode='markers+text',
+            text=clusters[mask],
+            customdata=image_ids[mask],
+            name=str(lbl),
+            marker=dict(size=6),
+            hovertemplate="<b>Cluster: %{text}</b><br>Image ID: %{customdata}<extra></extra>"
+        )
+        traces.append(trace)
+
+    fig = go.Figure(data=traces)
+    fig.update_layout(width=800, height=800)
+    return fig, image_data_uris
+
+# ==============================================
+# Main Video Processing
+# ==============================================
+def analyze_football_video(video_path: str):
+    BALL_ID, GOALKEEPER_ID, PLAYER_ID, REFEREE_ID = 0,1,2,3
+    MAXLEN = 5
+    M = deque(maxlen=MAXLEN)
+    path_raw = []
+
+    # Annotators
+    ellipse_annotator = sv.EllipseAnnotator(color=sv.ColorPalette.from_hex(['#00BFFF','#FF1493','#FFD700']), thickness=2)
+    label_annotator = sv.LabelAnnotator(color=sv.ColorPalette.from_hex(['#00BFFF','#FF1493','#FFD700']), text_color=sv.Color.from_hex('#000000'))
+    triangle_annotator = sv.TriangleAnnotator(color=sv.Color.from_hex('#FFD700'), base=20, height=17)
+    tracker = sv.ByteTrack()
+    tracker.reset()
+
+    cap = cv2.VideoCapture(video_path)
+    width, height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter("/tmp/annotated_video.mp4", fourcc, fps, (width, height))
+
+    player_crops = []
+
+    while True:
+        ret, frame = cap.read()
+        if not ret: break
+
+        # Player/ball/referee detection
+        result = PLAYER_DETECTION_MODEL.infer(frame, confidence=0.3)[0]
+        detections = sv.Detections.from_inference(result)
+        ball_detections = detections[detections.class_id==BALL_ID]
+        ball_detections.xyxy = sv.pad_boxes(ball_detections.xyxy, px=10)
+        all_detections = detections[detections.class_id != BALL_ID].with_nms(threshold=0.5, class_agnostic=True)
+        all_detections = tracker.update_with_detections(all_detections)
+
+        goalkeepers_detections = all_detections[all_detections.class_id==GOALKEEPER_ID]
+        players_detections = all_detections[all_detections.class_id==PLAYER_ID]
+        referees_detections = all_detections[all_detections.class_id==REFEREE_ID]
+
+        # Team classification
+        if len(players_detections.xyxy) > 0:
+            crops = [sv.crop_image(frame, xyxy) for xyxy in players_detections.xyxy]
+            players_detections.class_id = team_classifier.predict(crops)
+            player_crops += crops
+
+        goalkeepers_detections.class_id = resolve_goalkeepers_team_id(players_detections, goalkeepers_detections)
+        referees_detections.class_id -= 1
+        all_detections = sv.Detections.merge([players_detections, goalkeepers_detections, referees_detections])
+        labels = [f"#{tid}" for tid in all_detections.tracker_id]
+        all_detections.class_id = all_detections.class_id.astype(int)
+
+        # Annotate frame
+        annotated_frame = frame.copy()
+        annotated_frame = ellipse_annotator.annotate(annotated_frame, all_detections)
+        annotated_frame = label_annotator.annotate(annotated_frame, all_detections, labels=labels)
+        annotated_frame = triangle_annotator.annotate(annotated_frame, ball_detections)
+        out.write(annotated_frame)
+
+        # Field detection and projection
+        result_field = FIELD_DETECTION_MODEL.infer(frame, confidence=0.3)[0]
+        key_points = sv.KeyPoints.from_inference(result_field)
+        filter = key_points.confidence[0] > 0.5
+        frame_ref_points = key_points.xy[0][filter]
+        pitch_ref_points = np.array(CONFIG.vertices)[filter]
+        transformer = ViewTransformer(source=frame_ref_points, target=pitch_ref_points)
+        M.append(transformer.m)
+        transformer.m = np.mean(np.array(M), axis=0)
+
+        # Ball & players projected
+        pitch_ball_xy = transformer.transform_points(ball_detections.get_anchors_coordinates(sv.Position.BOTTOM_CENTER))
+        pitch_players_xy = transformer.transform_points(players_detections.get_anchors_coordinates(sv.Position.BOTTOM_CENTER))
+
+    cap.release()
+    out.release()
+
+    # Create UMAP 3D plot
+    umap_fig, image_data_uris = create_umap_3d_plot(player_crops)
+
+    return "/tmp/annotated_video.mp4", umap_fig
+
+# ==============================================
+# Gradio Interface
+# ==============================================
+iface = gr.Interface(
+    fn=analyze_football_video,
+    inputs=gr.Video(label="Upload Football Video"),
+    outputs=[gr.Video(label="Annotated Video"), gr.Plot(label="3D Player Embeddings")],
+    title="Football Video Analyzer with SigLIP Player Embeddings"
+)