Spaces:
Paused
Paused
File size: 2,603 Bytes
05fcd0f |
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 |
import os
import cv2
import numpy as np
import math
from modules.video_queue import JobStatus
def assemble_grid_video(grid_job, child_jobs, settings):
"""
Assembles a grid video from the results of child jobs.
"""
print(f"Starting grid assembly for job {grid_job.id}")
output_dir = settings.get("output_dir", "outputs")
os.makedirs(output_dir, exist_ok=True)
video_paths = [child.result for child in child_jobs if child.status == JobStatus.COMPLETED and child.result and os.path.exists(child.result)]
if not video_paths:
print(f"No valid video paths found for grid job {grid_job.id}")
return None
print(f"Found {len(video_paths)} videos for grid assembly.")
# Determine grid size (e.g., 2x2, 3x3)
num_videos = len(video_paths)
grid_size = math.ceil(math.sqrt(num_videos))
# Get video properties from the first video
try:
cap = cv2.VideoCapture(video_paths[0])
if not cap.isOpened():
raise IOError(f"Cannot open video file: {video_paths[0]}")
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)
cap.release()
except Exception as e:
print(f"Error getting video properties from {video_paths[0]}: {e}")
return None
output_filename = os.path.join(output_dir, f"grid_{grid_job.id}.mp4")
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video_writer = cv2.VideoWriter(output_filename, fourcc, fps, (width * grid_size, height * grid_size))
caps = [cv2.VideoCapture(p) for p in video_paths]
while True:
frames = []
all_frames_read = True
for cap in caps:
ret, frame = cap.read()
if ret:
frames.append(frame)
else:
# If one video ends, stop processing
all_frames_read = False
break
if not all_frames_read or not frames:
break
# Create a blank canvas for the grid
grid_frame = np.zeros((height * grid_size, width * grid_size, 3), dtype=np.uint8)
# Place frames into the grid
for i, frame in enumerate(frames):
row = i // grid_size
col = i % grid_size
grid_frame[row*height:(row+1)*height, col*width:(col+1)*width] = frame
video_writer.write(grid_frame)
for cap in caps:
cap.release()
video_writer.release()
print(f"Grid video saved to {output_filename}")
return output_filename
|