Spaces:
Runtime error
Runtime error
File size: 2,073 Bytes
ac7cda5 |
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 |
import cv2
import numpy as np
from ..utils.blend import blend_images_cy
from ..utils.get_mask import get_mask
class PutBackNumpy:
def __init__(
self,
mask_template_path=None,
):
if mask_template_path is None:
mask = get_mask(512, 512, 0.9, 0.9)
self.mask_ori_float = np.concatenate([mask] * 3, 2)
else:
mask = cv2.imread(mask_template_path, cv2.IMREAD_COLOR)
self.mask_ori_float = mask.astype(np.float32) / 255.0
def __call__(self, frame_rgb, render_image, M_c2o):
h, w = frame_rgb.shape[:2]
mask_warped = cv2.warpAffine(
self.mask_ori_float, M_c2o[:2, :], dsize=(w, h), flags=cv2.INTER_LINEAR
).clip(0, 1)
frame_warped = cv2.warpAffine(
render_image, M_c2o[:2, :], dsize=(w, h), flags=cv2.INTER_LINEAR
)
result = mask_warped * frame_warped + (1 - mask_warped) * frame_rgb
result = np.clip(result, 0, 255)
result = result.astype(np.uint8)
return result
class PutBack:
def __init__(
self,
mask_template_path=None,
):
if mask_template_path is None:
mask = get_mask(512, 512, 0.9, 0.9)
mask = np.concatenate([mask] * 3, 2)
else:
mask = cv2.imread(mask_template_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
self.mask_ori_float = np.ascontiguousarray(mask)[:,:,0]
self.result_buffer = None
def __call__(self, frame_rgb, render_image, M_c2o):
h, w = frame_rgb.shape[:2]
mask_warped = cv2.warpAffine(
self.mask_ori_float, M_c2o[:2, :], dsize=(w, h), flags=cv2.INTER_LINEAR
).clip(0, 1)
frame_warped = cv2.warpAffine(
render_image, M_c2o[:2, :], dsize=(w, h), flags=cv2.INTER_LINEAR
)
self.result_buffer = np.empty((h, w, 3), dtype=np.uint8)
# Use Cython implementation for blending
blend_images_cy(mask_warped, frame_warped, frame_rgb, self.result_buffer)
return self.result_buffer |