Spaces:
Sleeping
Sleeping
File size: 2,393 Bytes
123489f |
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 |
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
|