BoxMOT tracking backend and YOLO training module

.gitignore

@@ -22,6 +22,16 @@ sh_scripts/abci
 outputs/abci_logs
 *.pt
 
+# Working media and scratch results dropped at the repo root during experiments.
+# Root-anchored so tracked assets under docs/assets/ and outputs/graph/ are unaffected.
+/*.mp4
+/*.png
+/output/
+notebooks/test/
+
+# Vendored upstream clone with its own .git — install boxmot from PyPI instead
+/boxmot/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

configs/default_config.yaml

@@ -153,15 +153,23 @@ create_ground_truth_mot:
   pitch_length: ${pitch_length}
   pitch_width: ${pitch_width}
 
+# YOLO dataset creation
 create_yolo_dataset:
-  video_path: "/data/share/teamtrack/teamtrack-mot/soccer_side/full/combined.mp4"
-  mot_path: "/data/share/teamtrack/teamtrack-mot/soccer_side/full/gt.txt"
-  output_dir: "/home/atom/SoccerTrack-v2/data/v1"
-  frame_interval: 1000
-  train_split: 0.8
-  val_split: 0.1
-  test_split: 0.1
-  overwrite: true
+  match_id: null  # Required
+  half: null      # Required: 1st or 2nd
+  frame_interval: 1
+  train_ratio: 0.8  # Ratio of images to use for training
+  base_dir: null  # Optional: defaults to data/interim/{match_id}
+
+# YOLO model training
+train_yolo_model:
+  match_id: null  # Required
+  half: null      # Required: 1st or 2nd
+  model_type: "yolov8n.pt"
+  epochs: 100
+  batch_size: 16
+  imgsz: 640
+  base_dir: null  # Optional: defaults to data/interim/{match_id}
 
 log_string:
   string: "Hello, world!"

pyproject.toml

@@ -5,6 +5,7 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
+    "boxmot>=11.0.8",
     "deffcode>=0.2.6",
     "huggingface-hub>=0.25.0",
     "joblib>=1.4.2",
@@ -19,6 +20,7 @@ dependencies = [
     "opencv-python>=4.10.0.84",
     "openstarlab-preprocessing>=0.1.11",
     "parse>=1.20.2",
+    "pip>=25.0.1",
     "pydrive2>=1.21.1",
     "pyexiftool>=0.5.6",
     "python-dotenv>=1.0.1",

scripts/convert_coordinates_to_bboxes.sh

@@ -59,7 +59,9 @@ for HALF in "${HALVES[@]}"; do
             --detections_path "$DETECTIONS_PATH" \
             --output_path "$BBOX_MODELS_PATH" \
             --match_id "${MATCH_ID}" \
-            --conf_threshold 0.3
+            --conf_threshold 0.6 \
+            --grid_size 5 \
+            --interp_factor 50
 
         # Then create ground truth MOT file using the regression models
         if [ $? -eq 0 ]; then

scripts/generate_detections.sh

@@ -11,7 +11,8 @@ MATCH_ID=$1
 
 # Set up paths
 BASE_DIR="data/interim/$MATCH_ID"
-WEIGHTS_PATH="models/model=yolov8m-imgsz=2048.pt"
+WEIGHTS_PATH="models/yolov8m.pt"
+# WEIGHTS_PATH="models/model=yolov8m-imgsz=2048.pt"
 
 # Check if required files exist
 if [ ! -f "$WEIGHTS_PATH" ]; then
@@ -57,7 +58,7 @@ for HALF in "${HALVES[@]}"; do
             detect_objects.video_path="$VIDEO_PATH" \
             detect_objects.output_path="$OUTPUT_PATH" \
             detect_objects.weights_path="$WEIGHTS_PATH" \
-            detect_objects.vid_stride=200
+            detect_objects.vid_stride=30
 
         # Check if detection was successful
         if [ $? -ne 0 ]; then

scripts/plot_pitch_coordinates.py

@@ -0,0 +1,158 @@
+"""Plot pitch plane coordinates on a football field visualization."""
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pandas as pd
+from loguru import logger
+from mplsoccer import Pitch
+
+
+def plot_pitch_coordinates(
+    coordinates_path: Path | str,
+    output_path: Path | str,
+    frame_number: int | None = None,
+    pitch_length: float = 105.0,
+    pitch_width: float = 68.0,
+) -> None:
+    """
+    Plot pitch plane coordinates on a football field visualization.
+
+    Args:
+        coordinates_path: Path to the CSV file containing pitch plane coordinates
+        output_path: Path to save the output image
+        frame_number: Specific frame to plot (if None, plots first frame)
+        pitch_length: Length of the pitch in meters
+        pitch_width: Width of the pitch in meters
+    """
+    logger.info(f"Loading coordinates from: {coordinates_path}")
+
+    # Load coordinates
+    coordinates_df = pd.read_csv(coordinates_path)
+    logger.info(f"Loaded {len(coordinates_df)} coordinate entries")
+
+    # Filter for specific frame if requested
+    if frame_number is not None:
+        coordinates_df = coordinates_df[coordinates_df["frame"] == frame_number]
+    else:
+        # Use the first available frame
+        frame_number = coordinates_df["frame"].min()
+        coordinates_df = coordinates_df[coordinates_df["frame"] == frame_number]
+
+    logger.info(f"Plotting frame {frame_number} with {len(coordinates_df)} points")
+
+    # Create pitch
+    pitch = Pitch(
+        pitch_type='custom',
+        pitch_color='grass',
+        line_color='white',
+        pitch_width=pitch_width,
+        pitch_length=pitch_length,
+        goal_type='box',
+        linewidth=2
+    )
+
+    fig, ax = pitch.draw(figsize=(16, 10))
+
+    # Define colors for teams and ball
+    team_colors = {
+        9701: '#0080FF',  # Blue team
+        9834: '#FF8000',  # Orange team
+        'ball': '#FFFF00'  # Yellow for ball
+    }
+
+    # Plot players by team
+    for team_id in coordinates_df['teamId'].unique():
+        if pd.notna(team_id):  # Skip NaN values
+            team_data = coordinates_df[coordinates_df['teamId'] == team_id]
+            # Convert normalized coordinates to pitch coordinates
+            x_coords = team_data['x'].values * pitch_length
+            y_coords = team_data['y'].values * pitch_width
+
+            color = team_colors.get(int(team_id), '#FFFFFF')  # Default white
+            pitch.scatter(
+                x_coords, y_coords,
+                ax=ax,
+                color=color,
+                s=200,
+                edgecolors='black',
+                linewidth=1,
+                label=f'Team {int(team_id)}'
+            )
+
+    # Plot ball separately if it exists
+    ball_data = coordinates_df[coordinates_df['id'] == 'ball']
+    if not ball_data.empty:
+        x_ball = ball_data['x'].values[0] * pitch_length
+        y_ball = ball_data['y'].values[0] * pitch_width
+        pitch.scatter(
+            x_ball, y_ball,
+            ax=ax,
+            color=team_colors['ball'],
+            s=150,
+            edgecolors='black',
+            linewidth=2,
+            marker='o',
+            label='Ball'
+        )
+
+    # Add title and legend
+    plt.title(f'Pitch Coordinates - Frame {frame_number}', fontsize=16, pad=20)
+    plt.legend(loc='upper right', fontsize=12)
+
+    # Save the figure
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    logger.info(f"Saved plot to: {output_path}")
+    plt.close()
+
+
+def main():
+    """Main function to parse arguments and execute plotting."""
+    parser = argparse.ArgumentParser(description="Plot pitch plane coordinates on a football field")
+    parser.add_argument(
+        "--coordinates_path",
+        type=str,
+        required=True,
+        help="Path to the pitch plane coordinates CSV file"
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        default="output/pitch_coordinates_plot.png",
+        help="Path to save the output image"
+    )
+    parser.add_argument(
+        "--frame_number",
+        type=int,
+        default=None,
+        help="Specific frame number to plot (default: first frame)"
+    )
+    parser.add_argument(
+        "--pitch_length",
+        type=float,
+        default=105.0,
+        help="Pitch length in meters"
+    )
+    parser.add_argument(
+        "--pitch_width",
+        type=float,
+        default=68.0,
+        help="Pitch width in meters"
+    )
+
+    args = parser.parse_args()
+
+    plot_pitch_coordinates(
+        coordinates_path=args.coordinates_path,
+        output_path=args.output_path,
+        frame_number=args.frame_number,
+        pitch_length=args.pitch_length,
+        pitch_width=args.pitch_width,
+    )
+
+
+if __name__ == "__main__":
+    main() 

scripts/prepare_and_train_yolo.sh

@@ -0,0 +1,133 @@
+#!/bin/bash
+
+# Print usage
+print_usage() {
+    echo "Usage: $0 <match_id> [options]"
+    echo "Example: $0 117093 --skip-dataset --frame-interval 5"
+    echo ""
+    echo "Options:"
+    echo "  --skip-dataset      Skip dataset creation"
+    echo "  --skip-training     Skip model training" 
+    echo "  --frame-interval N  Extract every Nth frame (default: 1)"
+    echo "  --help             Show this help message"
+    echo "  --device           Device to use for training (default: cuda)"
+}
+
+# Check if help is requested
+if [[ "$1" == "--help" ]]; then
+    print_usage
+    exit 0
+fi
+
+# Check if match_id is provided
+if [ $# -eq 0 ]; then
+    print_usage
+    exit 1
+fi
+
+MATCH_ID=$1
+shift  # Remove match_id from arguments
+
+# Parse flags
+SKIP_DATASET=false
+SKIP_TRAINING=false
+FRAME_INTERVAL=1
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --skip-dataset)
+            SKIP_DATASET=true
+            shift
+            ;;
+        --skip-training)
+            SKIP_TRAINING=true
+            shift
+            ;;
+        --frame-interval)
+            FRAME_INTERVAL="$2"
+            shift 2
+            ;;
+        --device)
+            DEVICE="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            print_usage
+            exit 1
+            ;;
+    esac
+done
+
+# Function to process one half
+process_half() {
+    local half=$1
+
+    echo "Processing ${half} half..."
+
+    # Create YOLO format dataset
+    if [ "$SKIP_DATASET" = false ]; then
+        echo "Creating YOLO format dataset..."
+        uv run python -m src.main \
+            command=create_yolo_dataset \
+            create_yolo_dataset.match_id="$MATCH_ID" \
+            create_yolo_dataset.half="$half" \
+            create_yolo_dataset.frame_interval="$FRAME_INTERVAL"
+
+        if [ $? -ne 0 ]; then
+            echo "Error: Failed to create dataset for ${half} half"
+            exit 1
+        fi
+    else
+        echo "Skipping dataset creation..."
+    fi
+
+    # Train YOLO model
+    if [ "$SKIP_TRAINING" = false ]; then
+        echo "Training YOLO model..."
+        uv run python -m src.main \
+            command=train_yolo_model \
+            train_yolo_model.match_id="$MATCH_ID" \
+            train_yolo_model.half="$half" \
+            train_yolo_model.model_type="yolov8m.pt" \
+            train_yolo_model.epochs=50 \
+            train_yolo_model.batch_size=16 \
+            train_yolo_model.imgsz=1024 \
+            train_yolo_model.name="${MATCH_ID}_${half}_half_yolo" \
+            train_yolo_model.device="$DEVICE"
+
+        if [ $? -ne 0 ]; then
+            echo "Error: Failed to train model for ${half} half"
+            exit 1
+        fi
+    else
+        echo "Skipping model training..."
+    fi
+
+    echo "Completed processing ${half} half"
+    echo "----------------------------------------"
+}
+
+# Process both halves
+process_half "1st"
+process_half "2nd"
+
+# Print completion message
+echo "
+Processing completed successfully for match ${MATCH_ID}!"
+
+# Only show relevant output paths based on what was processed
+if [ "$SKIP_DATASET" = false ]; then
+    echo "
+1. YOLO format datasets:
+   - First half:  data/interim/${MATCH_ID}/ultralytics_format_1st_half_distorted/
+   - Second half: data/interim/${MATCH_ID}/ultralytics_format_2nd_half_distorted/"
+fi
+
+if [ "$SKIP_TRAINING" = false ]; then
+    echo "
+2. Trained models:
+   - Under models/soccer_player_detection/
+   - Best weights will be in 'best.pt'
+   - Last weights will be in 'last.pt'"
+fi