TheNetherWatcher's picture
Upload folder using huggingface_hub
d0ffe9c verified
raw
history blame
2.19 kB
import glob
import logging
import os
from pathlib import Path
import cv2
import numpy as np
import torch
from PIL import Image
from rembg import new_session, remove
from tqdm.rich import tqdm
logger = logging.getLogger(__name__)
def rembg_create_fg(frame_dir, output_dir, output_mask_dir, masked_area_list,
bg_color=(0,255,0),
mask_padding=0,
):
frame_list = sorted(glob.glob( os.path.join(frame_dir, "[0-9]*.png"), recursive=False))
if mask_padding != 0:
kernel = np.ones((abs(mask_padding),abs(mask_padding)),np.uint8)
kernel2 = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
session = new_session(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
for i, frame in tqdm(enumerate(frame_list),total=len(frame_list), desc=f"creating mask"):
frame = Path(frame)
file_name = frame.name
cur_frame_no = int(frame.stem)
img = Image.open(frame)
img_array = np.asarray(img)
mask_array = remove(img_array, only_mask=True, session=session)
#mask_array = mask_array[None,...]
if mask_padding < 0:
mask_array = cv2.erode(mask_array.astype(np.uint8),kernel,iterations = 1)
elif mask_padding > 0:
mask_array = cv2.dilate(mask_array.astype(np.uint8),kernel,iterations = 1)
mask_array = cv2.morphologyEx(mask_array, cv2.MORPH_OPEN, kernel2)
mask_array = cv2.GaussianBlur(mask_array, (7, 7), sigmaX=3, sigmaY=3, borderType=cv2.BORDER_DEFAULT)
if masked_area_list[cur_frame_no] is not None:
masked_area_list[cur_frame_no] = np.where(masked_area_list[cur_frame_no] > mask_array[None,...], masked_area_list[cur_frame_no], mask_array[None,...])
else:
masked_area_list[cur_frame_no] = mask_array[None,...]
if output_mask_dir:
Image.fromarray(mask_array).save( output_mask_dir / file_name )
img_array = np.asarray(img).copy()
if bg_color is not None:
img_array[mask_array == 0] = bg_color
img = Image.fromarray(img_array)
img.save( output_dir / file_name )
return masked_area_list