import os os.environ["OMP_NUM_THREADS"] = "1" import random import glob import cv2 import tqdm import numpy as np from typing import Union from utils.commons.tensor_utils import convert_to_np from utils.commons.os_utils import multiprocess_glob import pickle import traceback import multiprocessing from utils.commons.multiprocess_utils import multiprocess_run_tqdm from scipy.ndimage import binary_erosion, binary_dilation from sklearn.neighbors import NearestNeighbors from mediapipe.tasks.python import vision from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter, encode_segmap_mask_to_image, decode_segmap_mask_from_image, job_cal_seg_map_for_image seg_model = None segmenter = None mat_model = None lama_model = None lama_config = None from data_gen.utils.process_video.split_video_to_imgs import extract_img_job BG_NAME_MAP = { "knn": "", } FRAME_SELECT_INTERVAL = 5 SIM_METHOD = "mse" SIM_THRESHOLD = 3 def save_file(name, content): with open(name, "wb") as f: pickle.dump(content, f) def load_file(name): with open(name, "rb") as f: content = pickle.load(f) return content def save_rgb_alpha_image_to_path(img, alpha, img_path): try: os.makedirs(os.path.dirname(img_path), exist_ok=True) except: pass cv2.imwrite(img_path, np.concatenate([cv2.cvtColor(img, cv2.COLOR_RGB2BGR), alpha], axis=-1)) def save_rgb_image_to_path(img, img_path): try: os.makedirs(os.path.dirname(img_path), exist_ok=True) except: pass cv2.imwrite(img_path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) def load_rgb_image_to_path(img_path): return cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB) def image_similarity(x: np.ndarray, y: np.ndarray, method="mse"): if method == "mse": return np.mean((x - y) ** 2) else: raise NotImplementedError def extract_background(img_lst, segmap_mask_lst=None, method="knn", device='cpu', mix_bg=True): """ img_lst: list of rgb ndarray method: "knn" """ global segmenter global seg_model global mat_model global lama_model global lama_config assert len(img_lst) > 0 if segmap_mask_lst is not None: assert len(segmap_mask_lst) == len(img_lst) else: del segmenter del seg_model seg_model = MediapipeSegmenter() segmenter = vision.ImageSegmenter.create_from_options(seg_model.video_options) def get_segmap_mask(img_lst, segmap_mask_lst, index): if segmap_mask_lst is not None: segmap = refresh_segment_mask(segmap_mask_lst[index]) else: segmap = seg_model._cal_seg_map(refresh_image(img_lst[index]), segmenter=segmenter) return segmap if method == "knn": num_frames = len(img_lst) if num_frames < 100: FRAME_SELECT_INTERVAL = 5 elif num_frames < 10000: FRAME_SELECT_INTERVAL = 20 else: FRAME_SELECT_INTERVAL = num_frames // 500 img_lst = img_lst[::FRAME_SELECT_INTERVAL] if num_frames > FRAME_SELECT_INTERVAL else img_lst[0:1] if segmap_mask_lst is not None: segmap_mask_lst = segmap_mask_lst[::FRAME_SELECT_INTERVAL] if num_frames > FRAME_SELECT_INTERVAL else segmap_mask_lst[0:1] assert len(img_lst) == len(segmap_mask_lst) # get H/W h, w = refresh_image(img_lst[0]).shape[:2] # nearest neighbors all_xys = np.mgrid[0:h, 0:w].reshape(2, -1).transpose() # [512*512, 2] coordinate grid distss = [] for idx, img in tqdm.tqdm(enumerate(img_lst), desc='combining backgrounds...', total=len(img_lst)): segmap = get_segmap_mask(img_lst=img_lst, segmap_mask_lst=segmap_mask_lst, index=idx) bg = (segmap[0]).astype(bool) # [h,w] bool mask fg_xys = np.stack(np.nonzero(~bg)).transpose(1, 0) # [N_nonbg,2] coordinate of non-bg pixels nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(fg_xys) dists, _ = nbrs.kneighbors(all_xys) # [512*512, 1] distance to nearest non-bg pixel distss.append(dists) distss = np.stack(distss) # [B, 512*512, 1] max_dist = np.max(distss, 0) # [512*512, 1] max_id = np.argmax(distss, 0) # id of frame bc_pixs = max_dist > 10 # 在各个frame有一个出现过是bg的pixel，bg标准是离最近的non-bg pixel距离大于10 bc_pixs_id = np.nonzero(bc_pixs) bc_ids = max_id[bc_pixs] # TODO: maybe we should reimplement here to avoid memory costs? # though there is upper limits of images here num_pixs = distss.shape[1] bg_img = np.zeros((h*w, 3), dtype=np.uint8) img_lst = [refresh_image(img) for img in img_lst] imgs = np.stack(img_lst).reshape(-1, num_pixs, 3) bg_img[bc_pixs_id, :] = imgs[bc_ids, bc_pixs_id, :] # 对那些铁bg的pixel，直接去对应的image里面采样 bg_img = bg_img.reshape(h, w, 3) max_dist = max_dist.reshape(h, w) bc_pixs = max_dist > 10 # 5 bg_xys = np.stack(np.nonzero(~bc_pixs)).transpose() fg_xys = np.stack(np.nonzero(bc_pixs)).transpose() nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(fg_xys) distances, indices = nbrs.kneighbors(bg_xys) # 对non-bg img，用KNN找最近的bg pixel bg_fg_xys = fg_xys[indices[:, 0]] bg_img[bg_xys[:, 0], bg_xys[:, 1], :] = bg_img[bg_fg_xys[:, 0], bg_fg_xys[:, 1], :] else: raise NotImplementedError # deperated return bg_img def inpaint_torso_job(gt_img, segmap): bg_part = (segmap[0]).astype(bool) head_part = (segmap[1] + segmap[3] + segmap[5]).astype(bool) neck_part = (segmap[2]).astype(bool) torso_part = (segmap[4]).astype(bool) img = gt_img.copy() img[head_part] = 0 # torso part "vertical" in-painting... L = 8 + 1 torso_coords = np.stack(np.nonzero(torso_part), axis=-1) # [M, 2] # lexsort: sort 2D coords first by y then by x, # ref: https://stackoverflow.com/questions/2706605/sorting-a-2d-numpy-array-by-multiple-axes inds = np.lexsort((torso_coords[:, 0], torso_coords[:, 1])) torso_coords = torso_coords[inds] # choose the top pixel for each column u, uid, ucnt = np.unique(torso_coords[:, 1], return_index=True, return_counts=True) top_torso_coords = torso_coords[uid] # [m, 2] # only keep top-is-head pixels top_torso_coords_up = top_torso_coords.copy() - np.array([1, 0]) # [N, 2] mask = head_part[tuple(top_torso_coords_up.T)] if mask.any(): top_torso_coords = top_torso_coords[mask] # get the color top_torso_colors = gt_img[tuple(top_torso_coords.T)] # [m, 3] # construct inpaint coords (vertically up, or minus in x) inpaint_torso_coords = top_torso_coords[None].repeat(L, 0) # [L, m, 2] inpaint_offsets = np.stack([-np.arange(L), np.zeros(L, dtype=np.int32)], axis=-1)[:, None] # [L, 1, 2] inpaint_torso_coords += inpaint_offsets inpaint_torso_coords = inpaint_torso_coords.reshape(-1, 2) # [Lm, 2] inpaint_torso_colors = top_torso_colors[None].repeat(L, 0) # [L, m, 3] darken_scaler = 0.98 ** np.arange(L).reshape(L, 1, 1) # [L, 1, 1] inpaint_torso_colors = (inpaint_torso_colors * darken_scaler).reshape(-1, 3) # [Lm, 3] # set color img[tuple(inpaint_torso_coords.T)] = inpaint_torso_colors inpaint_torso_mask = np.zeros_like(img[..., 0]).astype(bool) inpaint_torso_mask[tuple(inpaint_torso_coords.T)] = True else: inpaint_torso_mask = None # neck part "vertical" in-painting... push_down = 4 L = 48 + push_down + 1 neck_part = binary_dilation(neck_part, structure=np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=bool), iterations=3) neck_coords = np.stack(np.nonzero(neck_part), axis=-1) # [M, 2] # lexsort: sort 2D coords first by y then by x, # ref: https://stackoverflow.com/questions/2706605/sorting-a-2d-numpy-array-by-multiple-axes inds = np.lexsort((neck_coords[:, 0], neck_coords[:, 1])) neck_coords = neck_coords[inds] # choose the top pixel for each column u, uid, ucnt = np.unique(neck_coords[:, 1], return_index=True, return_counts=True) top_neck_coords = neck_coords[uid] # [m, 2] # only keep top-is-head pixels top_neck_coords_up = top_neck_coords.copy() - np.array([1, 0]) mask = head_part[tuple(top_neck_coords_up.T)] top_neck_coords = top_neck_coords[mask] # push these top down for 4 pixels to make the neck inpainting more natural... offset_down = np.minimum(ucnt[mask] - 1, push_down) top_neck_coords += np.stack([offset_down, np.zeros_like(offset_down)], axis=-1) # get the color top_neck_colors = gt_img[tuple(top_neck_coords.T)] # [m, 3] # construct inpaint coords (vertically up, or minus in x) inpaint_neck_coords = top_neck_coords[None].repeat(L, 0) # [L, m, 2] inpaint_offsets = np.stack([-np.arange(L), np.zeros(L, dtype=np.int32)], axis=-1)[:, None] # [L, 1, 2] inpaint_neck_coords += inpaint_offsets inpaint_neck_coords = inpaint_neck_coords.reshape(-1, 2) # [Lm, 2] inpaint_neck_colors = top_neck_colors[None].repeat(L, 0) # [L, m, 3] darken_scaler = 0.98 ** np.arange(L).reshape(L, 1, 1) # [L, 1, 1] inpaint_neck_colors = (inpaint_neck_colors * darken_scaler).reshape(-1, 3) # [Lm, 3] # set color img[tuple(inpaint_neck_coords.T)] = inpaint_neck_colors # apply blurring to the inpaint area to avoid vertical-line artifects... inpaint_mask = np.zeros_like(img[..., 0]).astype(bool) inpaint_mask[tuple(inpaint_neck_coords.T)] = True blur_img = img.copy() blur_img = cv2.GaussianBlur(blur_img, (5, 5), cv2.BORDER_DEFAULT) img[inpaint_mask] = blur_img[inpaint_mask] # set mask torso_img_mask = (neck_part | torso_part | inpaint_mask) torso_with_bg_img_mask = (bg_part | neck_part | torso_part | inpaint_mask) if inpaint_torso_mask is not None: torso_img_mask = torso_img_mask | inpaint_torso_mask torso_with_bg_img_mask = torso_with_bg_img_mask | inpaint_torso_mask torso_img = img.copy() torso_img[~torso_img_mask] = 0 torso_with_bg_img = img.copy() torso_img[~torso_with_bg_img_mask] = 0 return torso_img, torso_img_mask, torso_with_bg_img, torso_with_bg_img_mask def load_segment_mask_from_file(filename: str): encoded_segmap = load_rgb_image_to_path(filename) segmap_mask = decode_segmap_mask_from_image(encoded_segmap) return segmap_mask # load segment mask to memory if not loaded yet def refresh_segment_mask(segmap_mask: Union[str, np.ndarray]): if isinstance(segmap_mask, str): segmap_mask = load_segment_mask_from_file(segmap_mask) return segmap_mask # load segment mask to memory if not loaded yet def refresh_image(image: Union[str, np.ndarray]): if isinstance(image, str): image = load_rgb_image_to_path(image) return image def generate_segment_imgs_job(img_name, segmap, img): out_img_name = segmap_name = img_name.replace("/gt_imgs/", "/segmaps/").replace(".jpg", ".png") # 存成jpg的话，pixel value会有误差 try: os.makedirs(os.path.dirname(out_img_name), exist_ok=True) except: pass encoded_segmap = encode_segmap_mask_to_image(segmap) save_rgb_image_to_path(encoded_segmap, out_img_name) for mode in ['head', 'torso', 'person', 'bg']: out_img, mask = seg_model._seg_out_img_with_segmap(img, segmap, mode=mode) img_alpha = 255 * np.ones((img.shape[0], img.shape[1], 1), dtype=np.uint8) # alpha mask = mask[0][..., None] img_alpha[~mask] = 0 out_img_name = img_name.replace("/gt_imgs/", f"/{mode}_imgs/").replace(".jpg", ".png") save_rgb_alpha_image_to_path(out_img, img_alpha, out_img_name) inpaint_torso_img, inpaint_torso_img_mask, inpaint_torso_with_bg_img, inpaint_torso_with_bg_img_mask = inpaint_torso_job(img, segmap) img_alpha = 255 * np.ones((img.shape[0], img.shape[1], 1), dtype=np.uint8) # alpha img_alpha[~inpaint_torso_img_mask[..., None]] = 0 out_img_name = img_name.replace("/gt_imgs/", f"/inpaint_torso_imgs/").replace(".jpg", ".png") save_rgb_alpha_image_to_path(inpaint_torso_img, img_alpha, out_img_name) return segmap_name def segment_and_generate_for_image_job(img_name, img, segmenter_options=None, segmenter=None, store_in_memory=False): img = refresh_image(img) segmap_mask, segmap_image = job_cal_seg_map_for_image(img, segmenter_options=segmenter_options, segmenter=segmenter) segmap_name = generate_segment_imgs_job(img_name=img_name, segmap=segmap_mask, img=img) if store_in_memory: return segmap_mask else: return segmap_name def extract_segment_job( video_name, nerf=False, background_method='knn', device="cpu", total_gpus=0, mix_bg=True, store_in_memory=False, # set to True to speed up a bit of preprocess, but leads to HUGE memory costs (100GB for 5-min video) force_single_process=False, # turn this on if you find multi-process does not work on your environment ): global segmenter global seg_model del segmenter del seg_model seg_model = MediapipeSegmenter() segmenter = vision.ImageSegmenter.create_from_options(seg_model.options) # nerf means that we extract only one video, so can enable multi-process acceleration multiprocess_enable = nerf and not force_single_process try: if "cuda" in device: # determine which cuda index from subprocess id pname = multiprocessing.current_process().name pid = int(pname.rsplit("-", 1)[-1]) - 1 cuda_id = pid % total_gpus device = f"cuda:{cuda_id}" if nerf: # single video raw_img_dir = video_name.replace(".mp4", "/gt_imgs/").replace("/raw/","/processed/") else: # whole dataset raw_img_dir = video_name.replace(".mp4", "").replace("/video/", "/gt_imgs/") if not os.path.exists(raw_img_dir): extract_img_job(video_name, raw_img_dir) # use ffmpeg to split video into imgs img_names = glob.glob(os.path.join(raw_img_dir, "*.jpg")) img_lst = [] for img_name in img_names: if store_in_memory: img = load_rgb_image_to_path(img_name) else: img = img_name img_lst.append(img) print("| Extracting Segmaps && Saving...") args = [] segmap_mask_lst = [] # preparing parameters for segment for i in range(len(img_lst)): img_name = img_names[i] img = img_lst[i] if multiprocess_enable: # create seg_model in subprocesses here options = seg_model.options segmenter_arg = None else: # use seg_model of this process options = None segmenter_arg = segmenter arg = (img_name, img, options, segmenter_arg, store_in_memory) args.append(arg) if multiprocess_enable: for (_, res) in multiprocess_run_tqdm(segment_and_generate_for_image_job, args=args, num_workers=16, desc='generating segment images in multi-processes...'): segmap_mask = res segmap_mask_lst.append(segmap_mask) else: for index in tqdm.tqdm(range(len(img_lst)), desc="generating segment images in single-process..."): segmap_mask = segment_and_generate_for_image_job(*args[index]) segmap_mask_lst.append(segmap_mask) print("| Extracted Segmaps Done.") print("| Extracting background...") bg_prefix_name = f"bg{BG_NAME_MAP[background_method]}" bg_img = extract_background(img_lst, segmap_mask_lst, method=background_method, device=device, mix_bg=mix_bg) if nerf: out_img_name = video_name.replace("/raw/", "/processed/").replace(".mp4", f"/{bg_prefix_name}.jpg") else: out_img_name = video_name.replace("/video/", f"/{bg_prefix_name}_img/").replace(".mp4", ".jpg") save_rgb_image_to_path(bg_img, out_img_name) print("| Extracted background done.") print("| Extracting com_imgs...") com_prefix_name = f"com{BG_NAME_MAP[background_method]}" for i in tqdm.trange(len(img_names), desc='extracting com_imgs'): img_name = img_names[i] com_img = refresh_image(img_lst[i]).copy() segmap = refresh_segment_mask(segmap_mask_lst[i]) bg_part = segmap[0].astype(bool)[..., None].repeat(3,axis=-1) com_img[bg_part] = bg_img[bg_part] out_img_name = img_name.replace("/gt_imgs/", f"/{com_prefix_name}_imgs/") save_rgb_image_to_path(com_img, out_img_name) print("| Extracted com_imgs done.") return 0 except Exception as e: print(str(type(e)), e) traceback.print_exc(e) return 1 def out_exist_job(vid_name, background_method='knn'): com_prefix_name = f"com{BG_NAME_MAP[background_method]}" img_dir = vid_name.replace("/video/", "/gt_imgs/").replace(".mp4", "") out_dir1 = img_dir.replace("/gt_imgs/", "/head_imgs/") out_dir2 = img_dir.replace("/gt_imgs/", f"/{com_prefix_name}_imgs/") if os.path.exists(img_dir) and os.path.exists(out_dir1) and os.path.exists(out_dir1) and os.path.exists(out_dir2) : num_frames = len(os.listdir(img_dir)) if len(os.listdir(out_dir1)) == num_frames and len(os.listdir(out_dir2)) == num_frames: return None else: return vid_name else: return vid_name def get_todo_vid_names(vid_names, background_method='knn'): if len(vid_names) == 1: # nerf return vid_names todo_vid_names = [] fn_args = [(vid_name, background_method) for vid_name in vid_names] for i, res in multiprocess_run_tqdm(out_exist_job, fn_args, num_workers=16, desc="checking todo videos..."): if res is not None: todo_vid_names.append(res) return todo_vid_names if __name__ == '__main__': import argparse, glob, tqdm, random parser = argparse.ArgumentParser() parser.add_argument("--vid_dir", default='/home/tiger/datasets/raw/TH1KH_512/video') parser.add_argument("--ds_name", default='TH1KH_512') parser.add_argument("--num_workers", default=48, type=int) parser.add_argument("--seed", default=0, type=int) parser.add_argument("--process_id", default=0, type=int) parser.add_argument("--total_process", default=1, type=int) parser.add_argument("--reset", action='store_true') parser.add_argument("--load_names", action="store_true") parser.add_argument("--background_method", choices=['knn', 'mat', 'ddnm', 'lama'], type=str, default='knn') parser.add_argument("--total_gpus", default=0, type=int) # zero gpus means utilizing cpu parser.add_argument("--no_mix_bg", action="store_true") parser.add_argument("--store_in_memory", action="store_true") # set to True to speed up preprocess, but leads to high memory costs parser.add_argument("--force_single_process", action="store_true") # turn this on if you find multi-process does not work on your environment args = parser.parse_args() vid_dir = args.vid_dir ds_name = args.ds_name load_names = args.load_names background_method = args.background_method total_gpus = args.total_gpus mix_bg = not args.no_mix_bg store_in_memory = args.store_in_memory force_single_process = args.force_single_process devices = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(",") for d in devices[:total_gpus]: os.system(f'pkill -f "voidgpu{d}"') if ds_name.lower() == 'nerf': # 处理单个视频 vid_names = [vid_dir] out_names = [video_name.replace("/raw/", "/processed/").replace(".mp4","_lms.npy") for video_name in vid_names] else: # 处理整个数据集 if ds_name in ['lrs3_trainval']: vid_name_pattern = os.path.join(vid_dir, "*/*.mp4") elif ds_name in ['TH1KH_512', 'CelebV-HQ']: vid_name_pattern = os.path.join(vid_dir, "*.mp4") elif ds_name in ['lrs2', 'lrs3', 'voxceleb2']: vid_name_pattern = os.path.join(vid_dir, "*/*/*.mp4") elif ds_name in ["RAVDESS", 'VFHQ']: vid_name_pattern = os.path.join(vid_dir, "*/*/*/*.mp4") else: raise NotImplementedError() vid_names_path = os.path.join(vid_dir, "vid_names.pkl") if os.path.exists(vid_names_path) and load_names: print(f"loading vid names from {vid_names_path}") vid_names = load_file(vid_names_path) else: vid_names = multiprocess_glob(vid_name_pattern) vid_names = sorted(vid_names) print(f"saving vid names to {vid_names_path}") save_file(vid_names_path, vid_names) vid_names = sorted(vid_names) random.seed(args.seed) random.shuffle(vid_names) process_id = args.process_id total_process = args.total_process if total_process > 1: assert process_id <= total_process -1 num_samples_per_process = len(vid_names) // total_process if process_id == total_process: vid_names = vid_names[process_id * num_samples_per_process : ] else: vid_names = vid_names[process_id * num_samples_per_process : (process_id+1) * num_samples_per_process] if not args.reset: vid_names = get_todo_vid_names(vid_names, background_method) print(f"todo videos number: {len(vid_names)}") device = "cuda" if total_gpus > 0 else "cpu" extract_job = extract_segment_job fn_args = [(vid_name, ds_name=='nerf', background_method, device, total_gpus, mix_bg, store_in_memory, force_single_process) for i, vid_name in enumerate(vid_names)] if ds_name == 'nerf': # 处理单个视频 extract_job(*fn_args[0]) else: for vid_name in multiprocess_run_tqdm(extract_job, fn_args, desc=f"Root process {args.process_id}: segment images", num_workers=args.num_workers): pass