wunderscout 0.1.2__py3-none-any.whl

wunderscout/__init__.py

@@ -0,0 +1,6 @@
+from .vision import VisionEngine
+from .geometry import PitchMapper
+from .teams import TeamClassifier
+from .core import ScoutingPipeline
+
+__all__ = ["VisionEngine", "PitchMapper", "TeamClassifier", "ScoutingPipeline"]

wunderscout/core.py

@@ -0,0 +1,164 @@
+import cv2
+import supervision as sv
+import numpy as np
+from pathlib import Path
+from .vision import VisionEngine
+from .geometry import PitchMapper
+from .teams import TeamClassifier
+from .exporters import DataExporter
+
+
+class ScoutingPipeline:
+    def __init__(self, player_weights, field_weights):
+        self.engine = VisionEngine(player_weights, field_weights)
+        self.mapper = PitchMapper()
+        self.classifier = TeamClassifier()
+
+    def run(self, video_path, output_video_path):
+        # 1. Warm-up (Calibration)
+        print("WORKER: Calibrating teams...")
+        crops = self.engine.get_calibration_crops(video_path)
+        if len(crops) > 0:
+            embeddings = self.engine.get_embeddings(crops)
+            self.classifier.fit(embeddings)
+        else:
+            print("WARNING: No player crops found for calibration.")
+
+        # 2. Setup Video I/O
+        output_path_obj = Path(output_video_path)
+        output_path_obj.parent.mkdir(parents=True, exist_ok=True)
+        # ---------------------------------------------------------------------
+
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+        out = cv2.VideoWriter(
+            output_video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
+        )
+
+        if not out.isOpened():
+            print(f"ERROR: Could not create video file at {output_video_path}")
+            return
+
+        tracker = sv.ByteTrack()
+        tracking_results = {}
+
+        # ID Constants
+        BALL_ID = 0
+        GOALKEEPER_ID = 1
+        PLAYER_ID = 2
+        REFEREE_ID = 3
+
+        # 3. Main Processing Loop
+        print(f"WORKER: Starting processing: {video_path}")
+        frame_generator = sv.get_video_frames_generator(video_path)
+
+        for frame_idx, frame in enumerate(frame_generator):
+            if frame_idx % 100 == 0:
+                print(f"WORKER: Processing frame {frame_idx}")
+
+            # --- A. DETECTION ---
+            all_dets = self.engine.detect_players(frame)
+            f_res = self.engine.detect_field(frame)
+
+            # --- B. FIELD HOMOGRAPHY ---
+            H = None
+            if f_res.keypoints is not None and len(f_res.keypoints.xy) > 0:
+                H = self.mapper.get_matrix(
+                    f_res.keypoints.xy[0].cpu().numpy(),
+                    f_res.keypoints.conf[0].cpu().numpy(),
+                )
+            else:
+                H = self.mapper.last_h
+
+            # --- C. SEPARATE BALL & OTHERS ---
+            ball_detections = all_dets[all_dets.class_id == BALL_ID]
+            ball_detections.xyxy = sv.pad_boxes(xyxy=ball_detections.xyxy, px=10)
+
+            other_detections = all_dets[all_dets.class_id != BALL_ID]
+            other_detections = other_detections.with_nms(threshold=0.5)
+
+            # --- D. TRACKING ---
+            tracked_objects = tracker.update_with_detections(other_detections)
+
+            # Split tracked objects
+            tracked_players = tracked_objects[tracked_objects.class_id == PLAYER_ID]
+            tracked_gks = tracked_objects[tracked_objects.class_id == GOALKEEPER_ID]
+            tracked_refs = tracked_objects[tracked_objects.class_id == REFEREE_ID]
+
+            # --- E. TEAM CLASSIFICATION ---
+
+            # 1. Players
+            if len(tracked_players) > 0:
+                p_crops = [sv.crop_image(frame, xyxy) for xyxy in tracked_players.xyxy]
+                p_pil = [sv.cv2_to_pillow(c) for c in p_crops]
+                p_embeddings = self.engine.get_embeddings(p_pil)
+
+                final_team_ids = []
+                for i, tid in enumerate(tracked_players.tracker_id):
+                    team_id = self.classifier.get_consensus_team(tid, p_embeddings[i])
+                    final_team_ids.append(team_id)
+
+                tracked_players.class_id = np.array(final_team_ids)
+
+            # 2. Goalkeepers
+            if len(tracked_gks) > 0 and len(tracked_players) > 0:
+                tracked_gks.class_id = self.classifier.resolve_goalkeepers_team_id(
+                    tracked_players, tracked_gks
+                )
+
+            # 3. Referees (Shift ID 3 -> 2)
+            if len(tracked_refs) > 0:
+                tracked_refs.class_id -= 1
+
+            # --- F. DATA STORAGE ---
+            tracking_results[frame_idx] = {"players": {}, "ball": None}
+            data_targets = sv.Detections.merge([tracked_players, tracked_gks])
+
+            if H is not None:
+                if len(data_targets) > 0:
+                    feet_coords = data_targets.get_anchors_coordinates(
+                        sv.Position.BOTTOM_CENTER
+                    )
+                    transformed_feet = self.mapper.transform(feet_coords, H)
+
+                    for i, tid in enumerate(data_targets.tracker_id):
+                        px, py = transformed_feet[i]
+                        tracking_results[frame_idx]["players"][tid] = (
+                            max(0.0, min(1.0, px)),
+                            max(0.0, min(1.0, py)),
+                        )
+
+                if len(ball_detections) > 0:
+                    ball_coords = ball_detections.get_anchors_coordinates(
+                        sv.Position.CENTER
+                    )
+                    transformed_ball = self.mapper.transform([ball_coords[0]], H)
+                    bx, by = transformed_ball[0]
+                    tracking_results[frame_idx]["ball"] = (
+                        max(0.0, min(1.0, bx)),
+                        max(0.0, min(1.0, by)),
+                    )
+
+            # --- G. DRAW & WRITE VIDEO ---
+            all_tracked = sv.Detections.merge(
+                [tracked_players, tracked_gks, tracked_refs]
+            )
+            annotated_frame = self.engine.draw_annotations(
+                frame, all_tracked, ball_detections
+            )
+            out.write(annotated_frame)
+
+        # 4. Cleanup
+        out.release()
+        cap.release()
+        print(f"WORKER: Video saved to {output_video_path}")
+
+        # Save CSVs
+        final_assignments = self.classifier.get_final_assignments()
+        csv_path = output_video_path.replace(".mp4", ".csv")
+        DataExporter.save_csvs(
+            tracking_results, final_assignments, frame_idx, fps, csv_path
+        )

wunderscout/exporters.py

@@ -0,0 +1,43 @@
+import csv
+from pathlib import Path
+
+
+class DataExporter:
+    @staticmethod
+    def save_csvs(tracking_data, team_assignments, total_frames, fps, output_path):
+        """
+        tracking_data: {frame_idx: {"ball": (x,y), "players": {id: (x,y)}}}
+        team_assignments: {tracker_id: team_id}
+        """
+        path_obj = Path(output_path)
+        path_obj.parent.mkdir(parents=True, exist_ok=True)
+        base_name = str(path_obj.with_suffix(""))
+        home_ids = [tid for tid, team in team_assignments.items() if team == 0]
+        away_ids = [tid for tid, team in team_assignments.items() if team == 1]
+
+        def write_file(filename, team_name, ids):
+            with open(filename, "w", newline="") as f:
+                writer = csv.writer(f)
+                writer.writerow(
+                    ["", "", ""] + [team_name for _ in ids for _ in (0, 1)] + ["", ""]
+                )
+                writer.writerow(
+                    ["", "", ""] + [str(pid) for pid in ids for _ in (0, 1)] + ["", ""]
+                )
+                writer.writerow(
+                    ["Period", "Frame", "Time [s]"]
+                    + [f"Player{pid}_{axis}" for pid in ids for axis in ("X", "Y")]
+                    + ["Ball_X", "Ball_Y"]
+                )
+
+                for f_idx in range(total_frames):
+                    data = tracking_data.get(f_idx, {"ball": None, "players": {}})
+                    row = [1, f_idx, f"{f_idx / fps:.2f}"]
+                    for tid in ids:
+                        coords = data["players"].get(tid, ("NaN", "NaN"))
+                        row.extend(coords)
+                    row.extend(data["ball"] if data["ball"] else ("NaN", "NaN"))
+                    writer.writerow(row)
+
+        write_file(f"{base_name}_Home.csv", "Home", sorted(home_ids))
+        write_file(f"{base_name}_Away.csv", "Away", sorted(away_ids))

wunderscout/geometry.py

@@ -0,0 +1,74 @@
+import cv2
+import numpy as np
+
+PITCH_CONFIG = {
+    # --- LEFT GOAL LINE ---
+    0: (0.000, 0.000),  # Top-Left Corner
+    1: (0.000, 0.204),  # Top Edge of Penalty Box
+    2: (0.000, 0.365),  # Top Edge of Goal Area
+    3: (0.000, 0.635),  # Bottom Edge of Goal Area
+    4: (0.000, 0.796),  # Bottom Edge of Penalty Box
+    5: (0.000, 1.000),  # Bottom-Left Corner
+    # --- LEFT PENALTY AREA ---
+    6: (0.052, 0.365),
+    7: (0.052, 0.635),
+    8: (0.105, 0.500),  # Penalty Spot (Left)
+    9: (0.157, 0.204),
+    10: (0.157, 0.392),
+    11: (0.157, 0.608),
+    12: (0.157, 0.796),
+    # --- MIDFIELD ---
+    13: (0.413, 0.500),
+    14: (0.500, 0.000),
+    15: (0.500, 0.365),
+    16: (0.500, 0.635),
+    17: (0.500, 1.000),
+    18: (0.587, 0.500),
+    # --- RIGHT PENALTY AREA ---
+    19: (0.843, 0.204),
+    20: (0.843, 0.392),
+    21: (0.843, 0.608),
+    22: (0.843, 0.796),
+    23: (0.895, 0.500),  # Penalty Spot (Right)
+    24: (0.948, 0.365),
+    25: (0.948, 0.635),
+    # --- RIGHT GOAL LINE ---
+    26: (1.000, 0.000),
+    27: (1.000, 0.204),
+    28: (1.000, 0.365),
+    29: (1.000, 0.635),
+    30: (1.000, 0.796),
+    31: (1.000, 1.000),
+}
+
+
+class PitchMapper:
+    def __init__(self, pitch_config=PITCH_CONFIG):
+        self.pitch_config = pitch_config
+        self.last_h = None
+
+    def get_matrix(self, keypoints_xy, keypoints_conf):
+        src_points = []
+        dst_points = []
+
+        for i, (xy, conf) in enumerate(zip(keypoints_xy, keypoints_conf)):
+            if conf > 0.5 and i in self.pitch_config:
+                src_points.append(xy)
+                dst_points.append(self.pitch_config[i])
+
+        if len(src_points) >= 4:
+            H, _ = cv2.findHomography(
+                np.array(src_points), np.array(dst_points), cv2.RANSAC
+            )
+            self.last_h = H
+
+        return self.last_h
+
+    def transform(self, points, H=None):
+        target_h = H if H is not None else self.last_h
+        if target_h is None or len(points) == 0:
+            return []
+
+        points_reshaped = np.array(points).reshape(-1, 1, 2).astype(np.float32)
+        projected = cv2.perspectiveTransform(points_reshaped, target_h)
+        return projected.reshape(-1, 2)

wunderscout/heatmap.py

@@ -0,0 +1,115 @@
+import pandas as pd
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+import json
+from scipy.stats import gaussian_kde
+
+# === Load the raw tracking data CSV ===
+# NOTE: "header=2" → skip the first two rows (team/labels) and use row 3 as header
+df = pd.read_csv(
+    "./data/Sample_Game_1_RawTrackingData_Away_Team.csv",
+    header=2,
+)
+
+# === Clean column names so each player has _X and _Y ===
+cleaned_colums = []
+colnames = df.columns.tolist()
+i = 0
+while i < len(colnames):
+    col = colnames[i]
+    if col.startswith("Player") or col.startswith("Ball"):
+        cleaned_colums.append(f"{col}_X")
+        cleaned_colums.append(f"{col}_Y")
+        i += 2
+    else:
+        cleaned_colums.append(col)
+        i += 1
+df.columns = cleaned_colums
+
+print("Columns cleaned. First few rows:")
+print(df.head())
+
+# === Extract Player17 (drop NaN values where tracking failed) ===
+player17 = df[["Player17_X", "Player17_Y"]].dropna()
+x = player17["Player17_X"].to_numpy()
+y = player17["Player17_Y"].to_numpy()
+
+# === Detect scale (normalized [0,1] or real meters) ===
+if x.max() <= 1.5 and y.max() <= 1.5:
+    print("Scaling Player17 data from normalized [0,1] to meters...")
+    x = x * 105  # pitch length in meters
+    y = y * 68  # pitch width in meters
+else:
+    print("Data appears to already be in meters, leaving as is.")
+
+print("First 10 points:", list(zip(x[:10], y[:10])))
+
+# =============================================================================
+# 1. Scatter Plot (sanity check, raw positions)
+# =============================================================================
+fig, ax = plt.subplots(figsize=(10, 7))
+# Pitch outline
+ax.plot([0, 105, 105, 0, 0], [0, 0, 68, 68, 0], color="black")
+ax.plot([52.5, 52.5], [0, 68], color="black")  # halfway line
+# Player positions
+ax.scatter(x, y, s=1, alpha=0.3, color="blue")
+ax.set_xlim(0, 105)
+ax.set_ylim(0, 68)
+ax.set_title("Player17 Movement Scatter (raw positions)")
+plt.savefig("./heatmap/player17_scatter.png", dpi=150, bbox_inches="tight")
+
+# =============================================================================
+# 2. Histogram Heatmap (occupancy grid)
+# =============================================================================
+heatmap, xedges, yedges = np.histogram2d(x, y, bins=(50, 34), range=[[0, 105], [0, 68]])
+
+fig, ax = plt.subplots(figsize=(10, 7))
+extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
+im = ax.imshow(
+    heatmap.T, origin="lower", extent=extent, cmap="Blues", alpha=0.7, aspect="auto"
+)
+ax.plot([0, 105, 105, 0, 0], [0, 0, 68, 68, 0], color="black")
+ax.plot([52.5, 52.5], [0, 68], color="black")
+fig.colorbar(im, ax=ax, label="Frames")
+ax.set_title("Player17 Heatmap (Histogram)")
+plt.savefig("./heatmap/player17_histogram.png", dpi=150, bbox_inches="tight")
+
+# === Export histogram data as JSON for three.js ===
+heatmap_data = {
+    "xedges": xedges.tolist(),
+    "yedges": yedges.tolist(),
+    "values": heatmap.T.tolist(),  # transpose so rows correspond to y-axis correctly
+}
+with open("./heatmap/player17_histogram.json", "w") as f:
+    json.dump(heatmap_data, f)
+
+# =============================================================================
+# 3. KDE Heatmap (smoothed density field)
+# =============================================================================
+values = np.vstack([x, y])
+kde = gaussian_kde(values)
+
+# Define mesh grid
+X, Y = np.meshgrid(np.linspace(0, 105, 100), np.linspace(0, 68, 68))
+Z = kde(np.vstack([X.ravel(), Y.ravel()])).reshape(X.shape)
+
+fig, ax = plt.subplots(figsize=(10, 7))
+sns.kdeplot(x=x, y=y, fill=True, cmap="Blues", alpha=0.7, thresh=0.05, levels=50, ax=ax)
+ax.plot([0, 105, 105, 0, 0], [0, 0, 68, 68, 0], color="black")
+ax.plot([52.5, 52.5], [0, 68], color="black")
+ax.set_xlim(0, 105)
+ax.set_ylim(0, 68)
+ax.set_title("Player17 Heatmap (KDE Smoothed)")
+plt.savefig("./heatmap/player17_kde.png", dpi=150, bbox_inches="tight")
+
+# === Export KDE density field for three.js ===
+kde_data = {
+    "x": X[0].tolist(),  # x grid coordinates
+    "y": Y[:, 0].tolist(),  # y grid coordinates
+    "values": Z.tolist(),  # density values
+}
+with open("./heatmap/player17_kde.json", "w") as f:
+    json.dump(kde_data, f)
+
+print("Outputs saved: scatter, histogram PNG+JSON, KDE PNG+JSON for Player17")

wunderscout/main.py

@@ -0,0 +1,598 @@
+import os
+from dotenv import load_dotenv
+from roboflow import Roboflow
+import torch
+from ultralytics import YOLO
+import supervision as sv
+import cv2
+from tqdm import tqdm
+from transformers import AutoProcessor, SiglipVisionModel
+from more_itertools import chunked
+import numpy as np
+import umap
+from sklearn.cluster import KMeans
+import plotly.graph_objects as go
+import base64
+from io import BytesIO
+
+load_dotenv()
+
+SIGLIP_MODEL_PATH = "google/siglip-base-patch16-224"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+EMBEDDINGS_MODEL = SiglipVisionModel.from_pretrained(SIGLIP_MODEL_PATH).to(DEVICE)
+EMBEDDINGS_PROCESSOR = AutoProcessor.from_pretrained(SIGLIP_MODEL_PATH)
+ROBOFLOW_API_KEY = os.getenv("ROBOFLOW_API_KEY")
+
+FIELD_DETECTION_MODEL_ID = "football-field-detection-f07vi/14"
+
+
+def train():
+    print("Beginning training ...")
+    # This is to get the roboflow dataset and train.
+    rf = Roboflow(api_key=ROBOFLOW_API_KEY)
+
+    project = rf.workspace("roboflow-jvuqo").project("football-players-detection-3zvbc")
+    version = project.version(20)
+    dataset = version.download("yolov11")
+
+    model_base_s = YOLO("yolo11s.pt")
+
+    results = model_base_s.train(
+        data="./football-players-detection-20/data.yaml",
+        save=True,
+        epochs=50,
+        imgsz=1280,
+        plots=True,
+        device=0,
+        batch=6,
+        project="runs/detect",
+    )
+
+
+def inference(video_path):
+    print("Inference 3 ...")
+    # We do another inference with different annotations
+    BALL_ID = 0
+
+    model_trained = YOLO("./runs/detect/train/weights/best.pt")
+    ellipse_annotator = sv.EllipseAnnotator(
+        color=sv.ColorPalette.from_hex(["#00BFFF", "#FF1493", "FFD700"]), thickness=2
+    )
+
+    triangle_annotator = sv.TriangleAnnotator(
+        color=sv.Color.from_hex("#FFD700"), base=25, height=21, outline_thickness=1
+    )
+
+    frame_generator_2 = sv.get_video_frames_generator(video_path)
+    frame_2 = next(frame_generator_2)
+
+    result_3 = model_trained.predict(
+        frame_2, save=True, project="runs", name="inference"
+    )[0]
+    detections_2 = sv.Detections.from_ultralytics(result_3)
+
+    ball_detections = detections_2[detections_2.class_id == BALL_ID]
+    ball_detections.xyxy = sv.pad_boxes(xyxy=ball_detections.xyxy, px=10)
+
+    all_detections = detections_2[detections_2.class_id != BALL_ID]
+    all_detections = all_detections.with_nms(threshold=0.5, class_agnostic=True)
+    all_detections.class_id -= 1
+
+    annotated_frame_2 = frame_2.copy()
+    annotated_frame_2 = ellipse_annotator.annotate(
+        scene=annotated_frame_2, detections=all_detections
+    )
+    annotated_frame_2 = triangle_annotator.annotate(
+        scene=annotated_frame_2, detections=ball_detections
+    )
+
+    cv2.imwrite("./runs/annotated_frame.jpg", annotated_frame_2)
+
+
+def inference_with_player_tracking(video_path):
+    print("Inference with tracking ...")
+    # We will do inference again but this time with the video.
+    BALL_ID = 0
+
+    model_trained = YOLO("./runs/detect/train/weights/best.pt")
+    ellipse_annotator_2 = sv.EllipseAnnotator(
+        color=sv.ColorPalette.from_hex(["#00BFFF", "#FF1493", "#FFD700"]), thickness=2
+    )
+
+    label_annotator_2 = sv.LabelAnnotator(
+        color=sv.ColorPalette.from_hex(["#00BFFF", "#FF1493", "#FFD700"]),
+        text_color=sv.Color.from_hex("#000000"),
+        text_position=sv.Position.BOTTOM_CENTER,
+    )
+
+    triangle_annotator_2 = sv.TriangleAnnotator(
+        color=sv.Color.from_hex("#FFD700"),
+        base=25,
+        height=21,
+        outline_thickness=1,
+    )
+
+    tracker = sv.ByteTrack()
+
+    frame_generator_3 = sv.get_video_frames_generator(video_path)
+
+    cap = cv2.VideoCapture(video_path)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(
+        "./runs/video_tracked.mp4", fourcc, fps, (frame_width, frame_height)
+    )
+
+    for frame in frame_generator_3:
+        result_tracking = model_trained.predict(frame, conf=0.3)[0]
+
+        detections = sv.Detections.from_ultralytics(result_tracking)
+
+        ball_detections = detections[detections.class_id == BALL_ID]
+        ball_detections.xyxy = sv.pad_boxes(xyxy=ball_detections.xyxy, px=10)
+
+        all_detections = detections[detections.class_id != BALL_ID]
+        all_detections = all_detections.with_nms(threshold=0.5, class_agnostic=True)
+        all_detections.class_id -= 1
+        all_detections = tracker.update_with_detections(detections=all_detections)
+
+        labels = [f"#{tracker_id}" for tracker_id in all_detections.tracker_id]
+
+        annotated_frame = frame.copy()
+        annotated_frame = ellipse_annotator_2.annotate(
+            scene=annotated_frame, detections=all_detections
+        )
+        annotated_frame = label_annotator_2.annotate(
+            scene=annotated_frame, detections=all_detections, labels=labels
+        )
+        annotated_frame = triangle_annotator_2.annotate(
+            scene=annotated_frame, detections=ball_detections
+        )
+
+        out.write(annotated_frame)
+    out.release()
+
+
+def create_player_crops(video_path):
+    print("Generating player crops ...")
+    # Getting training data for cluster model
+    PLAYER_ID = 2
+    STRIDE = 30
+
+    model_trained = YOLO("./runs/detect/train/weights/best.pt")
+    frame_generator_4 = sv.get_video_frames_generator(
+        source_path=video_path, stride=STRIDE
+    )
+
+    crops = []
+    for frame in tqdm(frame_generator_4, desc="collecting crops"):
+        result = model_trained.predict(frame, conf=0.3)[0]
+        detections = sv.Detections.from_ultralytics(result)
+        detections = detections.with_nms(threshold=0.5, class_agnostic=True)
+        detections = detections[detections.class_id == PLAYER_ID]
+        players_crops = [sv.crop_image(frame, xyxy) for xyxy in detections.xyxy]
+        crops += players_crops
+
+    print(f"Total crops collected: {len(crops)}.")
+
+    # Save to a file
+    # out_dir = "./runs/player_crops"
+    # os.makedirs(out_dir, exist_ok=True)
+
+    # for i, crop in enumerate(crops):
+    #     filename = os.path.join(out_dir, f"crop_{i:04d}.jpg")
+    #     cv2.imwrite(filename, crop)
+
+    # Convert to PIL
+    pil_crops = [sv.cv2_to_pillow(c) for c in crops]
+    return pil_crops
+
+
+def extract_embeddings(crops):
+    BATCH_SIZE = 32
+
+    batches = 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)
+    return data
+
+
+def cluster_players_by_team(embeddings):
+    REDUCER = umap.UMAP(n_components=3)
+    CLUSTERING_MODEL = KMeans(n_clusters=2, n_init=10, random_state=42)
+
+    projections = REDUCER.fit_transform(embeddings)
+    clusters = CLUSTERING_MODEL.fit_predict(projections)
+
+    return projections, clusters
+
+
+def save_projection_plot_html(projections, clusters, crops):
+    # inline helper: convert a crop (PIL image) to base64 string
+    def pil_image_to_data_uri(image):
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{img_str}"
+
+    # convert crops to base64-encoded URIs
+    image_data_uris = {
+        f"image_{i}": pil_image_to_data_uri(crop) for i, crop in enumerate(crops)
+    }
+    image_ids = np.array([f"image_{i}" for i in range(len(crops))])
+
+    traces = []
+    unique_clusters = np.unique(clusters)
+    for unique_cluster in unique_clusters:
+        mask = clusters == unique_cluster
+        trace = go.Scatter3d(
+            x=projections[mask][:, 0],
+            y=projections[mask][:, 1],
+            z=projections[mask][:, 2],
+            mode="markers+text",  # hard-coded: markers+text
+            text=clusters[mask],
+            customdata=image_ids[mask],
+            name=str(unique_cluster),
+            marker=dict(size=6),
+            hovertemplate="<b>class: %{text}</b><br>image ID: %{customdata}<extra></extra>",
+        )
+        traces.append(trace)
+
+    # make cube axis range
+    min_val = np.min(projections)
+    max_val = np.max(projections)
+    padding = (max_val - min_val) * 0.05
+    axis_range = [min_val - padding, max_val + padding]
+
+    fig = go.Figure(data=traces)
+    fig.update_layout(
+        scene=dict(
+            xaxis=dict(title="X", range=axis_range),
+            yaxis=dict(title="Y", range=axis_range),
+            zaxis=dict(title="Z", range=axis_range),
+            aspectmode="cube",
+        ),
+        width=1000,
+        height=1000,
+        showlegend=True,
+    )
+
+    # embed chart HTML with custom JS for crop preview
+    plotly_div = fig.to_html(
+        full_html=False, include_plotlyjs=True, div_id="scatter-plot-3d"
+    )
+    javascript_code = f"""
+    <script>
+        function displayImage(imageId) {{
+            var imageElement = document.getElementById('image-display');
+            var placeholderText = document.getElementById('placeholder-text');
+            var imageDataURIs = {image_data_uris};
+            imageElement.src = imageDataURIs[imageId];
+            imageElement.style.display = 'block';
+            placeholderText.style.display = 'none';
+        }}
+
+        var chartElement = document.getElementById('scatter-plot-3d');
+        chartElement.on('plotly_click', function(data) {{
+            var customdata = data.points[0].customdata;
+            displayImage(customdata);
+        }});
+    </script>
+    """
+
+    html_template = f"""
+    <!DOCTYPE html>
+    <html>
+        <head>
+            <script src="https://cdn.plot.ly/plotly-latest.min.js"></script>
+            <style>
+                #image-container {{
+                    position: fixed;
+                    top: 0;
+                    left: 0;
+                    width: 200px;
+                    height: 200px;
+                    padding: 5px;
+                    border: 1px solid #ccc;
+                    background-color: white;
+                    z-index: 1000;
+                    box-sizing: border-box;
+                    display: flex;
+                    align-items: center;
+                    justify-content: center;
+                    text-align: center;
+                }}
+                #image-display {{
+                    width: 100%;
+                    height: 100%;
+                    object-fit: contain;
+                }}
+            </style>
+        </head>
+        <body>
+            {plotly_div}
+            <div id="image-container">
+                <img id="image-display" src="" alt="Selected image" style="display: none;" />
+                <p id="placeholder-text">Click a data point to display the image</p>
+            </div>
+            {javascript_code}
+        </body>
+    </html>
+    """
+
+    out_path = "./runs/player_clusters.html"
+    with open(out_path, "w") as f:
+        f.write(html_template)
+
+    print(f"Interactive plot saved to {out_path}. Open it in your browser.")
+
+
+def resolve_goalkeepers_team_id(players, goalkeepers):
+    players_xy = players.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+    goalkeepers_xy = goalkeepers.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)
+    goalkeepers_team_id = []
+
+    for gk_xy in goalkeepers_xy:
+        dist_0 = np.linalg.norm(gk_xy - team_0_centroid)
+        dist_1 = np.linalg.norm(gk_xy - team_1_centroid)
+        goalkeepers_team_id.append(0 if dist_0 < dist_1 else 1)
+
+    return np.array(goalkeepers_team_id)
+
+
+def inference_with_goalkeepers(video_path):
+    BALL_ID = 0
+    GOALKEEPER_ID = 1
+    PLAYER_ID = 2
+    REFEREE_ID = 3
+
+    crops = create_player_crops(video_path)
+    embeddings = extract_embeddings(crops)
+
+    REDUCER = umap.UMAP(n_components=3)
+    CLUSTERING_MODEL = KMeans(n_clusters=2, n_init=10, random_state=42)
+
+    projections = REDUCER.fit_transform(embeddings)
+    clustering_model = CLUSTERING_MODEL.fit(projections)
+
+    model_trained = YOLO("./runs/detect/train/weights/best.pt")
+
+    frame_generator = sv.get_video_frames_generator(video_path)
+    frame = next(frame_generator)
+
+    tracker = sv.ByteTrack()
+
+    cap = cv2.VideoCapture(video_path)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(
+        "./runs/video_tracked_2.mp4", fourcc, fps, (frame_width, frame_height)
+    )
+
+    for frame in frame_generator:
+        result = model_trained.predict(frame, conf=0.3)[0]
+
+        detections = sv.Detections.from_ultralytics(result)
+
+        ball_detections = detections[detections.class_id == BALL_ID]
+        ball_detections.xyxy = sv.pad_boxes(xyxy=ball_detections.xyxy, px=10)
+
+        all_detections = detections[detections.class_id != BALL_ID]
+        all_detections = all_detections.with_nms(threshold=0.5, class_agnostic=True)
+        all_detections = tracker.update_with_detections(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]
+
+        players_crops = [sv.crop_image(frame, xyxy) for xyxy in players_detections.xyxy]
+        player_embeddings = extract_embeddings(players_crops)
+        player_projection = REDUCER.transform(player_embeddings)
+        players_detections.class_id = clustering_model.predict(player_projection)
+
+        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"#{tracker_id}" for tracker_id in all_detections.tracker_id]
+
+        all_detections.class_id = all_detections.class_id.astype(int)
+
+        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"),
+            text_position=sv.Position.BOTTOM_CENTER,
+        )
+        triangle_annotator = sv.TriangleAnnotator(
+            color=sv.Color.from_hex("#FFD700"), base=25, height=21, outline_thickness=1
+        )
+
+        annotated_frame = frame.copy()
+        annotated_frame = ellipse_annotator.annotate(
+            scene=annotated_frame, detections=all_detections
+        )
+        annotated_frame = label_annotator.annotate(
+            scene=annotated_frame, detections=all_detections, labels=labels
+        )
+        annotated_frame = triangle_annotator.annotate(
+            scene=annotated_frame, detections=ball_detections
+        )
+
+        out.write(annotated_frame)
+    out.release()
+
+
+def keypoint_detection_train():
+    print("Beginning training ...")
+    # This is to get the roboflow dataset and train.
+    rf = Roboflow(api_key=ROBOFLOW_API_KEY)
+    project = rf.workspace("roboflow-jvuqo").project("football-field-detection-f07vi")
+    version = project.version(15)
+    dataset = version.download("yolov8")
+
+    model_base_s = YOLO("yolo11m-pose.pt")
+    results = model_base_s.train(
+        data="./football-field-detection-15/data.yaml",
+        save=True,
+        epochs=300,
+        plots=True,
+        imgsz=1080,
+        device=0,
+        batch=2,
+        project="runs/keypoint_detect/",
+    )
+
+
+def keypoint_detection_inference(video_path: str) -> None:
+    model = YOLO("./runs/keypoint_detect/train4/weights/best.pt")
+    vertex_annotator = sv.VertexAnnotator(
+        color=sv.Color.from_hex("#FF1493"),
+        radius=8,
+    )
+
+    frame = next(sv.get_video_frames_generator(video_path))
+    res = model.predict(frame, conf=0.3, iou=0.7, imgsz=1088, verbose=False)[0]
+
+    if res.boxes is None or len(res.boxes) == 0:
+        cv2.imwrite("./runs/annotated_frame_keypoint.jpg", frame)
+        return
+
+    best_i = int(np.argmax(res.boxes.conf.cpu().numpy()))
+
+    kxy = res.keypoints.xy[best_i].cpu().numpy()
+    kconf = res.keypoints.conf[best_i].cpu().numpy()
+
+    keep = kconf > 0.5
+    kxy = kxy[keep]
+
+    kp = sv.KeyPoints(xy=kxy[np.newaxis, ...])
+
+    out = vertex_annotator.annotate(scene=frame.copy(), key_points=kp)
+    cv2.imwrite("./runs/annotated_frame_keypoint.jpg", out)
+
+
+def keypoint_detection_indices(video_path, output_path="keypoint_debug.mp4"):
+    print(f"DEBUG: Starting keypoint visualization for {video_path}")
+
+    model_trained = YOLO("./runs/keypoint_detect/train4/weights/best.pt")
+
+    cap = cv2.VideoCapture(video_path)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    frame_count = 0
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        frame_count += 1
+        if frame_count % 30 == 0:
+            print(f"Processing frame {frame_count}...")
+
+        result = model_trained.predict(frame, conf=0.5, verbose=False)[0]
+
+        if result.keypoints is not None and result.keypoints.xy.shape[0] > 0:
+            kpts_xy = result.keypoints.xy[0].cpu().numpy()
+            kpts_conf = result.keypoints.conf[0].cpu().numpy()
+
+            for i, (x, y) in enumerate(kpts_xy):
+                conf = kpts_conf[i]
+
+                if conf < 0.5 or (x == 0 and y == 0):
+                    continue
+
+                x, y = int(x), int(y)
+                cv2.circle(frame, (x, y), 5, (0, 0, 255), -1)
+                cv2.putText(
+                    frame,
+                    str(i + 1),
+                    (x + 10, y),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.6,
+                    (0, 255, 255),
+                    2,
+                )
+        else:
+            cv2.putText(
+                frame,
+                "NO PITCH DETECTED",
+                (50, 50),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                1,
+                (0, 0, 255),
+                2,
+            )
+        out.write(frame)
+
+    cap.release()
+    out.release()
+    print("Video generation complete.")
+
+
+def keypoint_detection_val_test(model_path, data_yaml):
+    model = YOLO(model_path)
+    results = model.val(data=data_yaml, split="test", plots=True)
+
+    print(results)
+
+
+def main():
+    print("CUDA available:", torch.cuda.is_available())
+    if torch.cuda.is_available():
+        print("GPU name:", torch.cuda.get_device_name(0))
+
+    video_path = "/home/lucas/Documents/dev/local/yolo_torch/video.mp4"
+
+    # train()
+    # inference(video_path)
+    # inference_with_player_tracking(video_path)
+    # crops = create_player_crops(video_path)
+    # embeddings = extract_embeddings(crops)
+    # projections, clusters = cluster_players_by_team(embeddings)
+    # save_projection_plot_html(projections, clusters, crops)
+    # inference_with_goalkeepers(video_path)
+    # keypoint_detection_train()
+    # keypoint_detection_inference(video_path)
+    keypoint_detection_indices(video_path)
+    # val_test(
+    #     "./runs/keypoint_detect/train4/weights/best.pt",
+    #     "./football-field-detection-15/data.yaml",
+    # )
+
+
+if __name__ == "__main__":
+    main()

wunderscout/pass_network.py

@@ -0,0 +1,103 @@
+import json
+import networkx as nx
+import matplotlib.pyplot as plt
+from collections import defaultdict
+import os
+from pathlib import Path
+
+# Load match events JSON (replace with your actual file path)
+with open(
+    "./data/3825818.json",
+    "r",
+) as f:
+    events = json.load(f)
+
+# Build mapping from player ID -> player name from the Starting XI event
+player_id_to_name = {}
+
+for ev in events:
+    if ev["type"]["name"] == "Starting XI":
+        for lineup in ev["tactics"]["lineup"]:
+            pid = lineup["player"]["id"]
+            name = lineup["player"]["name"]
+            player_id_to_name[pid] = name
+
+# Data structures for passes and positions
+edges = defaultdict(int)  # (passer, recipient) -> count of passes
+player_positions = defaultdict(list)  # player_id -> list of [x, y] positions
+
+TEAM_NAME = "Real Sociedad"
+
+# Extract completed passes
+for ev in events:
+    if ev["type"]["name"] == "Pass" and ev["team"]["name"] == TEAM_NAME:
+        passer = ev["player"]["id"]
+        recipient = ev.get("pass", {}).get("recipient", {}).get("id")
+        outcome = ev.get("pass", {}).get("outcome", {"name": "Complete"})["name"]
+
+        if outcome == "Complete" and recipient is not None:
+            edges[(passer, recipient)] += 1
+            start = ev["location"]
+            end = ev["pass"]["end_location"]
+            player_positions[passer].append(start)
+            player_positions[recipient].append(end)
+
+# Calculate avg positions
+avg_positions = {}
+for player_id, coords in player_positions.items():
+    xs = [pt[0] for pt in coords]
+    ys = [pt[1] for pt in coords]
+    avg_positions[player_id] = [sum(xs) / len(xs), sum(ys) / len(ys)]
+
+# Build a JSON-friendly structure for export (nodes + links)
+nodes = [{"id": pid, "x": pos[0], "y": pos[1]} for pid, pos in avg_positions.items()]
+links = [
+    {"source": src, "target": tgt, "value": count}
+    for (src, tgt), count in edges.items()
+]
+
+network = {"nodes": nodes, "links": links}
+os.makedirs("pass_network", exist_ok=True)
+with open("./pass_network/pass_network.json", "w") as f:
+    json.dump(network, f, indent=2)
+
+# Build NetworkX graph
+G = nx.DiGraph()
+
+# Add nodes with positions
+for pid, pos in avg_positions.items():
+    G.add_node(pid, pos=(pos[0], pos[1]))
+
+# Add edges with weights
+for (src, tgt), count in edges.items():
+    G.add_edge(src, tgt, weight=count)
+
+# Draw graph
+pos = nx.get_node_attributes(G, "pos")
+labels = {pid: player_id_to_name.get(pid, str(pid)) for pid in G.nodes()}
+
+fig, ax = plt.subplots(figsize=(10, 7))
+
+# Draw pitch outline
+ax.set_xlim(0, 120)
+ax.set_ylim(0, 80)
+ax.plot([0, 120, 120, 0, 0], [0, 0, 80, 80, 0], color="black")
+
+# Draw nodes
+nx.draw_networkx_nodes(G, pos, ax=ax, node_color="skyblue", node_size=500)
+
+# Draw edges
+nx.draw_networkx_edges(
+    G,
+    pos,
+    ax=ax,
+    width=[d["weight"] * 0.2 for _, _, d in G.edges(data=True)],
+    alpha=0.7,
+    arrowsize=10,
+)
+
+# Draw player names
+nx.draw_networkx_labels(G, pos, labels=labels, ax=ax, font_size=8)
+
+plt.title("Team Pass Network")
+plt.savefig("./pass_network/pass_network_viz.png", dpi=150, bbox_inches="tight")

wunderscout/teams.py

@@ -0,0 +1,76 @@
+import numpy as np
+import umap
+from sklearn.cluster import KMeans
+import supervision as sv
+
+
+class TeamClassifier:
+    def __init__(self):
+        self.reducer = umap.UMAP(n_components=3)
+        self.clusterer = KMeans(n_clusters=2, n_init=10, random_state=42)
+        self.history = {}
+
+    def fit(self, embeddings):
+        projections = self.reducer.fit_transform(embeddings)
+        self.clusterer.fit(projections)
+
+    def get_consensus_team(self, tracker_id, embedding):
+        proj = self.reducer.transform(embedding.reshape(1, -1))
+        pred = self.clusterer.predict(proj)[0]
+
+        if tracker_id not in self.history:
+            self.history[tracker_id] = []
+        self.history[tracker_id].append(pred)
+        if len(self.history[tracker_id]) > 50:
+            self.history[tracker_id].pop(0)
+
+        return (
+            1
+            if (sum(self.history[tracker_id]) / len(self.history[tracker_id])) > 0.5
+            else 0
+        )
+
+    def resolve_goalkeepers_team_id(self, players, goalkeepers):
+        """
+        Assigns goalkeepers to the team whose centroid is closest.
+        players: sv.Detections (already classified with class_id 0 or 1)
+        goalkeepers: sv.Detections
+        """
+        if len(players) == 0 or len(goalkeepers) == 0:
+            return np.array([0] * len(goalkeepers))
+
+        players_xy = players.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+        goalkeepers_xy = goalkeepers.get_anchors_coordinates(sv.Position.BOTTOM_CENTER)
+
+        # Calculate centroids for Team 0 and Team 1
+        team_0_mask = players.class_id == 0
+        team_1_mask = players.class_id == 1
+
+        # Handle cases where one team might not be detected yet
+        if np.any(team_0_mask):
+            team_0_centroid = players_xy[team_0_mask].mean(axis=0)
+        else:
+            team_0_centroid = np.array([0, 0])
+
+        if np.any(team_1_mask):
+            team_1_centroid = players_xy[team_1_mask].mean(axis=0)
+        else:
+            team_1_centroid = np.array([10000, 10000])  # Far away
+
+        goalkeepers_team_id = []
+
+        for gk_xy in goalkeepers_xy:
+            dist_0 = np.linalg.norm(gk_xy - team_0_centroid)
+            dist_1 = np.linalg.norm(gk_xy - team_1_centroid)
+            goalkeepers_team_id.append(0 if dist_0 < dist_1 else 1)
+
+        return np.array(goalkeepers_team_id)
+
+    def get_final_assignments(self):
+        assignments = {}
+        for tid, votes in self.history.items():
+            if len(votes) > 0:
+                avg = sum(votes) / len(votes)
+                assignments[tid] = 1 if avg > 0.5 else 0
+        return assignments
+

wunderscout/vision.py

@@ -0,0 +1,155 @@
+import torch
+from ultralytics import YOLO
+import supervision as sv
+from transformers import AutoProcessor, SiglipVisionModel, data
+from roboflow import Roboflow
+from tqdm import tqdm
+from more_itertools import chunked
+import numpy as np
+from pathlib import Path
+
+
+class VisionEngine:
+    def __init__(self, player_weights, field_weights, device=None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.player_model = YOLO(player_weights)
+        self.field_model = YOLO(field_weights)
+
+        # Siglip for embeddings
+        siglip_path = "google/siglip-base-patch16-224"
+        self.siglip_model = SiglipVisionModel.from_pretrained(siglip_path).to(
+            self.device
+        )
+        self.siglip_processor = AutoProcessor.from_pretrained(siglip_path)
+
+        # --- Annotators ---
+        # Palette: 0=Blue, 1=Pink, 2=Yellow (Referee)
+        self.palette = sv.ColorPalette.from_hex(["#00BFFF", "#FF1493", "#FFD700"])
+
+        self.ellipse_annotator = sv.EllipseAnnotator(
+            color=self.palette,
+            thickness=2,
+        )
+        self.label_annotator = sv.LabelAnnotator(
+            color=self.palette,
+            text_color=sv.Color.from_hex("#000000"),
+            text_position=sv.Position.BOTTOM_CENTER,
+        )
+        self.triangle_annotator = sv.TriangleAnnotator(
+            color=sv.Color.from_hex("#FFD700"), base=25, height=21, outline_thickness=1
+        )
+
+    def get_calibration_crops(self, video_path, stride=30):
+        PLAYER_ID = 2
+        frame_generator = sv.get_video_frames_generator(
+            source_path=video_path, stride=stride
+        )
+
+        crops = []
+        for frame in frame_generator:
+            detections = self.detect_players(frame)
+            # Filter for players only for calibration
+            players = detections[detections.class_id == PLAYER_ID]
+            frame_crops = [sv.crop_image(frame, xyxy) for xyxy in players.xyxy]
+            crops += [sv.cv2_to_pillow(c) for c in frame_crops]
+
+        print(f"VisionEngine: Collected {len(crops)} calibration crops.")
+        return crops
+
+    def get_embeddings(self, pil_crops, batch_size=32):
+        batches = chunked(pil_crops, batch_size)
+        data_list = []
+
+        with torch.no_grad():
+            for batch in batches:
+                inputs = self.siglip_processor(images=batch, return_tensors="pt").to(
+                    self.device
+                )
+                outputs = self.siglip_model(**inputs)
+                embeddings = torch.mean(outputs.last_hidden_state, dim=1).cpu().numpy()
+                data_list.append(embeddings)
+
+        return np.concatenate(data_list) if data_list else np.array([])
+
+    def detect_players(self, frame, conf=0.3):
+        result = self.player_model.predict(frame, conf=conf, verbose=False)[0]
+        return sv.Detections.from_ultralytics(result)
+
+    def detect_field(self, frame, conf=0.3):
+        result = self.field_model.predict(frame, conf=conf, verbose=False)[0]
+        return result
+
+    def draw_annotations(self, frame, all_detections, ball_detections):
+        annotated_frame = frame.copy()
+
+        # 1. Draw Ball
+        annotated_frame = self.triangle_annotator.annotate(
+            scene=annotated_frame, detections=ball_detections
+        )
+
+        # 2. Draw People (Players, GKs, Refs)
+        if len(all_detections) > 0:
+            # Ensure class_id is int for color mapping
+            all_detections.class_id = all_detections.class_id.astype(int)
+
+            labels = [f"#{tracker_id}" for tracker_id in all_detections.tracker_id]
+
+            annotated_frame = self.ellipse_annotator.annotate(
+                scene=annotated_frame, detections=all_detections
+            )
+            annotated_frame = self.label_annotator.annotate(
+                scene=annotated_frame, detections=all_detections, labels=labels
+            )
+
+        return annotated_frame
+
+
+class ScoutingTrainer:
+    def __init__(self, api_key):
+        self.rf = Roboflow(api_key=api_key)
+
+    def train_players(
+        self,
+        workspace,
+        project,
+        version,
+        epochs=300,
+        output_dir="../runs/training/player",
+    ):
+        project = self.rf.workspace(workspace).project(project)
+        dataset = project.version(version).download("yolov11")
+        model = YOLO("../data/base_models/yolo11m.pt")
+
+        return model.train(
+            data=f"{dataset.location}/data.yaml",
+            epochs=epochs,
+            imgsz=1280,
+            plots=True,
+            device=0,
+            batch=2,
+            project=output_dir,
+        )
+
+    def train_field(
+        self,
+        workspace,
+        project,
+        version,
+        epochs=300,
+        output_dir="../runs/training/field",
+    ):
+        project = self.rf.workspace(workspace).project(project)
+        version = project.version(15)
+        dataset = version.download("yolov8", location="../data/data_sets/")
+        model = YOLO("yolo11m-pose.pt")
+
+        return model.train(
+            data=f"{dataset.location}/data.yaml",
+            save=True,
+            epochs=epochs,
+            plots=True,
+            imgsz=1080,
+            device=0,
+            batch=2,
+            project=output_dir,
+        )

wunderscout-0.1.2.dist-info/METADATA

@@ -0,0 +1,27 @@
+Metadata-Version: 2.4
+Name: wunderscout
+Version: 0.1.2
+Summary: Add your description here
+Requires-Python: >=3.11
+Requires-Dist: ipython>=9.5.0
+Requires-Dist: matplotlib>=3.10.6
+Requires-Dist: memory-profiler>=0.61.0
+Requires-Dist: more-itertools>=10.8.0
+Requires-Dist: networkx>=3.5
+Requires-Dist: numba>=0.58
+Requires-Dist: numpy>=2.3.2
+Requires-Dist: opencv-python>=4.11.0.86
+Requires-Dist: pandas>=2.3.2
+Requires-Dist: plotly>=6.3.0
+Requires-Dist: protobuf>=6.32.1
+Requires-Dist: psutil>=7.0.0
+Requires-Dist: python-dotenv>=1.1.1
+Requires-Dist: roboflow>=1.2.7
+Requires-Dist: scikit-learn>=1.7.2
+Requires-Dist: seaborn>=0.13.2
+Requires-Dist: sentencepiece>=0.2.1
+Requires-Dist: supervision>=0.26.1
+Requires-Dist: tqdm>=4.67.1
+Requires-Dist: transformers>=4.56.1
+Requires-Dist: ultralytics>=8.3.193
+Requires-Dist: umap-learn>=0.5.9.post2

wunderscout-0.1.2.dist-info/RECORD

@@ -0,0 +1,12 @@
+wunderscout/__init__.py,sha256=UwZMG43lwIAxHCHPg8SYS9A2Azkem9vcZw6tkwFG9kU,217
+wunderscout/core.py,sha256=suVWCsiVmwsevte8tsX_53GEdZ_cQhKcFTknY20BWEw,6345
+wunderscout/exporters.py,sha256=CRczFAYUS6EOedL84wq9WaFWTPKJyGxL7kbZQvLzLOA,1875
+wunderscout/geometry.py,sha256=I5lt00O9jOiEoVpPGy5iVglzA7cgaAdUvzfuFBcJbRA,2197
+wunderscout/heatmap.py,sha256=5R8Zw5Bnk-8eHTcWudo-1a3Mt83WrSbqYX9MkUoht_8,4268
+wunderscout/main.py,sha256=mkiZcWgAuAAc8bUmKQXLkw--PX8dwrRR6qKHZqmyi5M,20008
+wunderscout/pass_network.py,sha256=QC859Pi5VKSgHu1qrE3Zybvu97lNorsb03UAf1IrSbs,3099
+wunderscout/teams.py,sha256=y0IclDACo3F8buVdpqqMCSmZJeWx2uqMkGNbZ6YToVc,2628
+wunderscout/vision.py,sha256=fVX3wtwCwe6AiiGxZjH8u4q2gk3t4gCb4MNvmQf7Lhs,5257
+wunderscout-0.1.2.dist-info/METADATA,sha256=LUmQrjy8MaxHjTmCFILvR95dSQ8kWTYfpxpJPWudjx4,841
+wunderscout-0.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+wunderscout-0.1.2.dist-info/RECORD,,

wunderscout-0.1.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any