File size: 4,111 Bytes
c3c6f00 | 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import cv2
import os
import time
import subprocess
import shutil
from src.tracker import PedestrianTracker
from src.analytics import AnalyticsEngine
from src.visualization import draw_annotations, overlay_heatmap
def get_real_fps(video_path):
cap = cv2.VideoCapture(video_path)
fps = cap.get(cv2.CAP_PROP_FPS)
cap.release()
if fps <= 0 or fps > 120:
return 30.0
return fps
def convert_input_if_needed(video_path):
converted_path = video_path + ".ready.mp4"
try:
# Standardize input to 720p/30fps to ensure smooth processing and playback
subprocess.run([
'ffmpeg', '-y', '-i', video_path,
'-vf', 'scale=-2:720',
'-r', '30',
'-c:v', 'libopenh264',
'-pix_fmt', 'yuv420p',
'-an',
converted_path
], capture_output=True, timeout=120)
if os.path.exists(converted_path):
return converted_path
except:
pass
return video_path
def convert_to_browser_mp4(temp_path, final_path, fps=30):
try:
# Use libopenh264 which we verified works on your system
subprocess.run([
'ffmpeg', '-y', '-i', temp_path,
'-c:v', 'libopenh264',
'-pix_fmt', 'yuv420p',
'-movflags', '+faststart',
'-r', str(fps),
final_path
], capture_output=True, timeout=300)
if os.path.exists(final_path):
if os.path.exists(temp_path): os.remove(temp_path)
return True
except:
pass
if os.path.exists(temp_path):
shutil.move(temp_path, final_path)
return False
def process_video(video_path, output_path=None, model_path='yolov8m.pt', progress_callback=None):
print(f"DEBUG: Starting process_video for {video_path}")
if progress_callback: progress_callback(0.02) # Early signal
working_path = convert_input_if_needed(video_path)
print(f"DEBUG: Input converted to {working_path}")
cap = cv2.VideoCapture(working_path)
if not cap.isOpened():
print(f"ERROR: Cannot open video {working_path}")
return None
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(f"DEBUG: Video info: {width}x{height} @ {fps}fps, Total: {total_frames} frames")
if fps <= 0: fps = 30.0
tracker = PedestrianTracker(model_path=model_path)
analytics = AnalyticsEngine(fps=fps, width=width, height=height)
temp_output = output_path + ".raw.avi" if output_path else None
writer = None
if temp_output:
# Using MJPG in AVI is very stable for raw output before conversion
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
writer = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
frame_count = 0
start_time = time.time()
while cap.isOpened():
ret, frame = cap.read()
if not ret: break
frame_count += 1
results = tracker.track(frame)
analytics.update_tracks(results, frame_count)
if writer:
# THIS IS THE CRITICAL PART: Bake rectangles into the frame
annotated_frame = draw_annotations(frame, results, analytics.tracks)
writer.write(annotated_frame)
if progress_callback:
progress_callback(frame_count / max(total_frames, 1))
cap.release()
if writer: writer.release()
if working_path != video_path and os.path.exists(working_path):
os.remove(working_path)
if temp_output and os.path.exists(temp_output):
convert_to_browser_mp4(temp_output, output_path, fps=fps)
end_time = time.time()
return {
"analytics": analytics,
"processing_time": end_time - start_time,
"fps": frame_count / (end_time - start_time + 0.001)
}
|