Spaces:
Sleeping
Sleeping
File size: 2,287 Bytes
a63ce99 |
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 |
import os
import cv2
import numpy as np
def compose_final_video(mask_path, frames_folder, extracted_frames_folder, output_path, fps=24):
"""
Composes the final heatmap video by placing each heatmap frame
on top of its corresponding original frame using the table region defined by the mask.
"""
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
if mask is None:
raise ValueError(f"Mask not found at {mask_path}")
_, binary_mask = cv2.threshold(mask, 200, 255, cv2.THRESH_BINARY)
height, width = binary_mask.shape
contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
raise ValueError("No white area found in mask")
table_contour = max(contours, key=cv2.contourArea)
x, y, w, h = cv2.boundingRect(table_contour)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
heatmap_files = sorted([f for f in os.listdir(frames_folder) if f.endswith(('.bmp', '.jpg', '.png'))])
original_files = sorted([f for f in os.listdir(extracted_frames_folder) if f.endswith(('.bmp', '.jpg', '.png'))])
if len(heatmap_files) != len(original_files):
raise ValueError("Mismatch between number of heatmap frames and original frames")
for heatmap_file, original_file in zip(heatmap_files, original_files):
heatmap = cv2.imread(os.path.join(frames_folder, heatmap_file))
original = cv2.imread(os.path.join(extracted_frames_folder, original_file))
if heatmap is None or original is None:
continue
ph, pw = heatmap.shape[:2]
if pw > w or ph > h:
scale = min(w / pw, h / ph)
heatmap = cv2.resize(heatmap, (int(pw * scale), int(ph * scale)))
pw, ph = heatmap.shape[1], heatmap.shape[0]
x_offset = x + (w - pw) // 2
y_offset = y + (h - ph) // 2
result = original.copy()
result[y_offset:y_offset+ph, x_offset:x_offset+pw] = heatmap
cv2.drawContours(result, [table_contour], -1, (0, 255, 0), 2)
video_writer.write(result)
video_writer.release()
print("[Video] Final video saved to:", output_path)
|