tennis_analysis

The main.py script serves as the primary entrypoint for the tennis analysis pipeline. It orchestrates the following workflow:

1. Video Ingestion: Reads video frames using read_video.
2. Detection & Tracking: Uses PlayerTracker and BallTracker to detect and track players and the ball. It supports loading detections from stubs (e.g., .pkl files) to skip detection steps during development.
3. Court Analysis: Uses CourtLineDetector to identify court keypoints.
4. Player Filtering: Filters detected players based on court keypoints using player_tracker.choose_and_filter_players.
5. Coordinate Transformation: Uses MiniCourt to convert bounding box detections into normalized mini-court coordinates.
6. Shot & Speed Analysis: Identifies ball shots, calculates ball speed (km/h), and determines which player hit the ball.
7. Visualization: Draws bounding boxes, court keypoints, mini-court overlays, and player statistics onto the video frames.
8. Output: Saves the processed video using save_video.

# To run the full pipeline, ensure you have the required models in your models/ directory
# and input videos in input_videos/

from main import main

if __name__ == "__main__":
    main()

A ball hit is identified by detecting rapid changes in the vertical direction (sign changes in delta_y) that persist over a specific number of frames.

1. Calculate delta_y using the difference of the smoothed mid_y_rolling_mean.
2. Iterate through the frames to find points where the direction of vertical movement flips (positive to negative or vice versa).
3. Verify if these direction changes occur frequently enough within a window defined by minimum_change_frames_for_hit to qualify as a hit.
4. Mark the hit frame in the ball_hit column with a value of 1.

# Calculate vertical velocity/change
df_ball_positions['delta_y'] = df_ball_positions['mid_y_rolling_mean'].diff()

# Initialize hit column
df_ball_positions['ball_hit'] = 0

# Detection logic
minimum_change_frames_for_hit = 25
for i in range(1, len(df_ball_positions) - int(minimum_change_frames_for_hit * 1.2)):
    negative_position_change = df_ball_positions['delta_y'].iloc[i] > 0 and df_ball_positions['delta_y'].iloc[i+1] < 0
    positive_position_change = df_ball_positions['delta_y'].iloc[i] < 0 and df_ball_positions['delta_y'].iloc[i+1] > 0

    if negative_position_change or positive_position_change:
        change_count = 0 
        for change_frame in range(i+1, i + int(minimum_change_frames_for_hit * 1.2) + 1):
            negative_position_change_following_frame = df_ball_positions['delta_y'].iloc[i] > 0 and df_ball_positions['delta_y'].iloc[change_frame] < 0
            positive_position_change_following_frame = df_ball_positions['delta_y'].iloc[i] < 0 and df_ball_positions['delta_y'].iloc[change_frame] > 0

            if negative_position_change and negative_position_change_following_frame:
                change_count += 1
            elif positive_position_change and positive_position_change_following_frame:
                change_count += 1

        if change_count > minimum_change_frames_for_hit - 1:
            df_ball_positions.iloc[i, df_ball_positions.columns.get_loc('ball_hit')] = 1

# Retrieve indices of detected hits
frame_nums_with_ball_hits = df_ball_positions[df_ball_positions['ball_hit'] == 1].index.tolist()

Training is performed using Mean Squared Error (MSELoss) and the Adam optimizer. The training loop iterates through the DataLoader, performing forward passes, calculating loss, and updating weights via backpropagation.

After training, the model weights can be saved using torch.save.

import torch

# Setup
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
epochs = 20

# Training Loop
for epoch in range(epochs):
    for i, (imgs, kps) in enumerate(train_loader):
        imgs = imgs.to(device)
        kps = kps.to(device)

        optimizer.zero_grad()
        outputs = model(imgs)
        loss = criterion(outputs, kps)
        loss.backward()
        optimizer.step()

        if i % 10 == 0:
            print(f"Epoch {epoch}, iter {i}, loss: {loss.item()}")

# Save the trained model
torch.save(model.state_dict(), "keypoints_model.pth")