Spaces:
Running
Running
| import os | |
| import cv2 | |
| import numpy as np | |
| # resize frames | |
| def resize_frames(frames, size=None): | |
| """ | |
| size: (w, h) | |
| """ | |
| if size is not None: | |
| frames = [cv2.resize(f, size) for f in frames] | |
| frames = np.stack(frames, 0) | |
| return frames | |
| # resize frames | |
| def resize_masks(masks, size=None): | |
| """ | |
| size: (w, h) | |
| """ | |
| if size is not None: | |
| masks = [np.expand_dims(cv2.resize(m, size), 2) for m in masks] | |
| masks = np.stack(masks, 0) | |
| return masks | |
| # apply gaussian blur to mask with defined strength | |
| def apply_blur(frame, strength): | |
| blurred = cv2.GaussianBlur(frame, (strength, strength), 0) | |
| return blurred | |
| # blur frames | |
| def blur_frames_and_write( | |
| frames, masks, ratio, strength, dilate_radius=15, fps=30, output_path="blurred.mp4" | |
| ): | |
| assert frames.shape[:3] == masks.shape, "different size between frames and masks" | |
| assert ratio > 0 and ratio <= 1, "ratio must in (0, 1]" | |
| # -------------------- | |
| # pre-processing | |
| # -------------------- | |
| masks = masks.copy() | |
| masks = np.clip(masks, 0, 1) | |
| kernel = cv2.getStructuringElement(2, (dilate_radius, dilate_radius)) | |
| masks = np.stack([cv2.dilate(mask, kernel) for mask in masks], 0) | |
| T, H, W = masks.shape | |
| masks = np.expand_dims(masks, axis=3) # expand to T, H, W, 1 | |
| # size: (w, h) | |
| if ratio == 1: | |
| size = (W, H) | |
| binary_masks = masks | |
| else: | |
| size = [int(W * ratio), int(H * ratio)] | |
| size = [ | |
| si + 1 if si % 2 > 0 else si for si in size | |
| ] # only consider even values | |
| # shortest side should be larger than 50 | |
| if min(size) < 50: | |
| ratio = 50.0 / min(H, W) | |
| size = [int(W * ratio), int(H * ratio)] | |
| binary_masks = resize_masks(masks, tuple(size)) | |
| frames = resize_frames(frames, tuple(size)) # T, H, W, 3 | |
| if not os.path.exists(os.path.dirname(output_path)): | |
| os.makedirs(os.path.dirname(output_path)) | |
| writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, size) | |
| for frame, mask in zip(frames, binary_masks): | |
| blurred_frame = apply_blur(frame, strength) | |
| masked = cv2.bitwise_or(blurred_frame, blurred_frame, mask=mask) | |
| processed = np.where(masked == (0, 0, 0), frame, masked) | |
| writer.write(processed[:, :, ::-1]) | |
| writer.release() | |
| return output_path | |