Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from yolov5 import YOLOv5 | |
| # Load the trained YOLOv5 model (best.pt) using the YOLOv5 library | |
| model = YOLOv5("best.pt") # This will use YOLOv5's internal loading mechanism | |
| # Function to process the video and calculate ball trajectory, speed, and visualize the pitch | |
| def process_video(video_file): | |
| # Load video file using OpenCV | |
| video = cv2.VideoCapture(video_file.name) | |
| ball_positions = [] | |
| speed_data = [] | |
| frame_count = 0 | |
| last_position = None | |
| while video.isOpened(): | |
| ret, frame = video.read() | |
| if not ret: | |
| break | |
| frame_count += 1 | |
| # Run YOLOv5 model on the frame to detect ball | |
| results = model(frame) # Using the pre-trained YOLOv5 model | |
| # Extract the ball position (assuming class 0 = ball) | |
| ball_detections = results.pandas().xywh | |
| ball = ball_detections[ball_detections['class'] == 0] # class 0 is ball, adjust as needed | |
| if not ball.empty: | |
| ball_x = ball.iloc[0]['xmin'] + (ball.iloc[0]['xmax'] - ball.iloc[0]['xmin']) / 2 | |
| ball_y = ball.iloc[0]['ymin'] + (ball.iloc[0]['ymax'] - ball.iloc[0]['ymin']) / 2 | |
| ball_positions.append((frame_count, ball_x, ball_y)) # Track position in each frame | |
| if last_position is not None: | |
| # Calculate speed based on pixel displacement between frames | |
| distance = np.sqrt((ball_x - last_position[1]) ** 2 + (ball_y - last_position[2]) ** 2) | |
| fps = video.get(cv2.CAP_PROP_FPS) # Frames per second of the video | |
| speed = distance * fps # Speed = distance / time (time between frames is 1/fps) | |
| speed_data.append(speed) | |
| last_position = (frame_count, ball_x, ball_y) # Update last position | |
| video.release() | |
| # Ball trajectory plot | |
| plot_trajectory(ball_positions) | |
| # Return results | |
| avg_speed = np.mean(speed_data) if speed_data else 0 | |
| return f"Average Ball Speed: {avg_speed:.2f} pixels per second" | |
| # Function to plot ball trajectory using matplotlib | |
| def plot_trajectory(ball_positions): | |
| x_positions = [pos[1] for pos in ball_positions] | |
| y_positions = [pos[2] for pos in ball_positions] | |
| plt.figure(figsize=(10, 6)) | |
| plt.plot(x_positions, y_positions, label="Ball Trajectory", color='b') | |
| plt.title("Ball Trajectory on Pitch") | |
| plt.xlabel("X Position (pitch width)") | |
| plt.ylabel("Y Position (pitch length)") | |
| plt.grid(True) | |
| plt.legend() | |
| plt.show() | |
| # Gradio interface for the app | |
| iface = gr.Interface( | |
| fn=process_video, # Function to call when video is uploaded | |
| inputs=gr.inputs.File(label="Upload a Video File"), # File input (video) | |
| outputs="text", # Output the result as text | |
| live=True # Keep the interface live | |
| ) | |
| iface.launch(debug=True) | |