AjaykumarPilla commited on
Commit
3320b0d
·
verified ·
1 Parent(s): 3c5c2d2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -5
app.py CHANGED
@@ -3,10 +3,17 @@ import torch
3
  import cv2
4
  import numpy as np
5
  import matplotlib.pyplot as plt
6
- from yolov5 import YOLOv5
7
 
8
- # Load the trained YOLOv5 model (best.pt) using the YOLOv5 library
9
- model = YOLOv5("best.pt") # This will use YOLOv5's internal loading mechanism
 
 
 
 
 
 
 
10
 
11
  # Function to process the video and calculate ball trajectory, speed, and visualize the pitch
12
  def process_video(video_file):
@@ -25,12 +32,12 @@ def process_video(video_file):
25
 
26
  frame_count += 1
27
 
28
- # Run YOLOv5 model on the frame to detect ball
29
  results = model(frame) # Using the pre-trained YOLOv5 model
30
 
31
  # Extract the ball position (assuming class 0 = ball)
32
  ball_detections = results.pandas().xywh
33
- ball = ball_detections[ball_detections['class'] == 0] # class 0 is ball, adjust as needed
34
 
35
  if not ball.empty:
36
  ball_x = ball.iloc[0]['xmin'] + (ball.iloc[0]['xmax'] - ball.iloc[0]['xmin']) / 2
 
3
  import cv2
4
  import numpy as np
5
  import matplotlib.pyplot as plt
6
+ from pathlib import Path
7
 
8
+ # Load the trained YOLOv5 model from a local path (directly using PyTorch)
9
+ def load_model(model_path):
10
+ # Load model using torch.load, ensuring no downloading from Hugging Face
11
+ model = torch.load(model_path, map_location="cpu") # Load model from the local path
12
+ model.eval() # Set the model to evaluation mode
13
+ return model
14
+
15
+ # Load the model directly from the local file
16
+ model = load_model("best.pt") # Ensure 'best.pt' is in the correct path
17
 
18
  # Function to process the video and calculate ball trajectory, speed, and visualize the pitch
19
  def process_video(video_file):
 
32
 
33
  frame_count += 1
34
 
35
+ # Run the model on the frame to detect the ball
36
  results = model(frame) # Using the pre-trained YOLOv5 model
37
 
38
  # Extract the ball position (assuming class 0 = ball)
39
  ball_detections = results.pandas().xywh
40
+ ball = ball_detections[ball_detections['class'] == 0] # Class 0 for ball (adjust if needed)
41
 
42
  if not ball.empty:
43
  ball_x = ball.iloc[0]['xmin'] + (ball.iloc[0]['xmax'] - ball.iloc[0]['xmin']) / 2