Update app.py
Browse files
app.py
CHANGED
@@ -31,7 +31,7 @@ BATCH_SIZE = 4 # Process 4 frames at a time
|
|
31 |
|
32 |
# Load model
|
33 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
34 |
-
model = attempt_load("yolov5s.pt") # Load yolov5s.pt
|
35 |
model.to(device).eval() # Move model to device and set to evaluation mode
|
36 |
|
37 |
# Function to process video and detect ball
|
@@ -240,14 +240,18 @@ def drs_analysis(video):
|
|
240 |
# Overlay pitch map in top-right corner
|
241 |
if pitch_x is not None and pitch_y is not None:
|
242 |
map_width = 200
|
243 |
-
map_height
|
|
|
244 |
pitch_map = np.zeros((map_height, map_width, 3), dtype=np.uint8)
|
245 |
pitch_map[:] = (0, 255, 0) # Green pitch
|
246 |
cv2.rectangle(pitch_map, (0, map_height-10), (map_width, map_height), (0, 51, 51), -1) # Brown stumps
|
247 |
bounce_x = int((pitch_x + PITCH_WIDTH/2) / PITCH_WIDTH * map_width)
|
248 |
bounce_y = int((1 - pitch_y / PITCH_LENGTH) * map_height)
|
249 |
cv2.circle(pitch_map, (bounce_x, bounce_y), 5, (0, 0, 255), -1) # Red bounce point
|
250 |
-
|
|
|
|
|
|
|
251 |
|
252 |
# Add text annotations
|
253 |
text = f"LBW: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h"
|
|
|
31 |
|
32 |
# Load model
|
33 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
34 |
+
model = attempt_load("yolov5s.pt") # Load yolov5s.pt
|
35 |
model.to(device).eval() # Move model to device and set to evaluation mode
|
36 |
|
37 |
# Function to process video and detect ball
|
|
|
240 |
# Overlay pitch map in top-right corner
|
241 |
if pitch_x is not None and pitch_y is not None:
|
242 |
map_width = 200
|
243 |
+
# Cap map_height to 25% of frame height to ensure it fits
|
244 |
+
map_height = min(int(map_width * PITCH_LENGTH / PITCH_WIDTH), frame_height // 4)
|
245 |
pitch_map = np.zeros((map_height, map_width, 3), dtype=np.uint8)
|
246 |
pitch_map[:] = (0, 255, 0) # Green pitch
|
247 |
cv2.rectangle(pitch_map, (0, map_height-10), (map_width, map_height), (0, 51, 51), -1) # Brown stumps
|
248 |
bounce_x = int((pitch_x + PITCH_WIDTH/2) / PITCH_WIDTH * map_width)
|
249 |
bounce_y = int((1 - pitch_y / PITCH_LENGTH) * map_height)
|
250 |
cv2.circle(pitch_map, (bounce_x, bounce_y), 5, (0, 0, 255), -1) # Red bounce point
|
251 |
+
# Ensure overlay fits within frame
|
252 |
+
overlay_region = frame[0:map_height, frame_width-map_width:frame_width]
|
253 |
+
if overlay_region.shape[0] >= map_height and overlay_region.shape[1] >= map_width:
|
254 |
+
frame[0:map_height, frame_width-map_width:frame_width] = cv2.resize(pitch_map, (map_width, map_height))
|
255 |
|
256 |
# Add text annotations
|
257 |
text = f"LBW: {lbw_decision}\nSpeed: {speed_kmh:.2f} km/h"
|