Spaces:
Sleeping
Sleeping
import cv2 | |
from Prediction import predict_fight # Assuming you have this | |
# Process video: read, predict, and output labeled video | |
def process_video(video_path): | |
cap = cv2.VideoCapture(video_path) | |
sequence_length = 40 # Number of frames for one prediction | |
all_frames = [] | |
predictions = [] | |
# Step 1: Read all frames from the video | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
all_frames.append(frame) | |
cap.release() | |
# Step 2: Process frames in chunks of 40 to make predictions | |
for i in range(0, len(all_frames), sequence_length): | |
frames_buffer = all_frames[i:i + sequence_length] # Get a batch of 40 frames | |
# If the number of frames is less than 40 at the end, pad it with the last frame | |
if len(frames_buffer) < sequence_length: | |
frames_buffer += [frames_buffer[-1]] * (sequence_length - len(frames_buffer)) | |
# Perform the prediction on the current batch of frames | |
fight_detected = predict_fight(frames_buffer) | |
# Store the prediction for this batch | |
predictions.append(fight_detected) | |
# Step 3: Create output video with labels | |
output_video_path = "output_labeled.mp4" | |
height, width, _ = all_frames[0].shape | |
fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
out = cv2.VideoWriter(output_video_path, fourcc, 30, (width, height)) # Adjust frame rate if needed | |
frame_idx = 0 | |
for pred in predictions: | |
label = "Violence Detected!" if pred else "No Violence" | |
color = (0, 0, 255) if pred else (0, 255, 0) | |
# For the next 40 frames, show the same label | |
for _ in range(sequence_length): | |
if frame_idx >= len(all_frames): | |
break | |
frame = all_frames[frame_idx] | |
cv2.putText(frame, label, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2) | |
out.write(frame) | |
frame_idx += 1 | |
out.release() | |
return output_video_path | |