Spaces:
Paused
Paused
File size: 2,603 Bytes
a422282 43be61b a422282 bf56017 a422282 bf56017 a422282 43be61b a422282 43be61b a422282 0723bf0 bf56017 a422282 0723bf0 a422282 0723bf0 a422282 0723bf0 a422282 0723bf0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
"""
Video processing and object detection module
"""
import cv2
import platform
import numpy as np
from tqdm import tqdm
from ultralytics import YOLO
class Detection:
"""Class to store detection results"""
def __init__(self, frame_idx, class_id, class_name, bbox, confidence):
self.frame_idx = frame_idx
self.class_id = class_id
self.class_name = class_name
self.bbox = bbox # [x1, y1, x2, y2]
self.confidence = confidence
def process_video(video_path, sample_rate=1):
"""
Process video and detect golfer, club, and ball
Args:
video_path (str): Path to the video file
sample_rate (int): Process every nth frame (default: 1 for all frames)
Returns:
tuple: (frames, detections)
- frames: List of processed frames
- detections: List of Detection objects
"""
model = YOLO("yolov8n.pt")
class_names = model.names
# On macOS ("Darwin"), the AVFoundation backend is often more reliable.
# For other systems, FFMPEG is a good choice.
backend = cv2.CAP_AVFOUNDATION if platform.system() == "Darwin" else cv2.CAP_FFMPEG
cap = cv2.VideoCapture(video_path, backend)
if not cap.isOpened():
backend_name = "AVFoundation" if platform.system() == "Darwin" else "FFMPEG"
print(f"Warning: Could not open video with {backend_name} backend. Trying default.")
cap = cv2.VideoCapture(video_path) # Fallback to default
if not cap.isOpened():
raise ValueError("Error opening video file with any available backend.")
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# Process all frames by default
frames = []
detections = []
for frame_idx in tqdm(range(0, frame_count, sample_rate), desc="Processing frames"):
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
ret, frame = cap.read()
if not ret:
print(f"Warning: Could not read frame {frame_idx}")
continue
# Store original frame
frames.append(frame)
# Run YOLOv8 detection
results = model(frame)
# Process detection results
for result in results:
boxes = result.boxes
for box in boxes:
class_id = int(box.cls.item())
class_name = class_names[class_id]
bbox = box.xyxy[0].tolist()
confidence = box.conf.item()
detections.append(Detection(frame_idx, class_id, class_name, bbox, confidence))
cap.release()
return frames, detections |