|
import os |
|
os.environ["OMP_NUM_THREADS"] = "1" |
|
|
|
import glob |
|
import cv2 |
|
import tqdm |
|
import numpy as np |
|
import PIL |
|
from utils.commons.tensor_utils import convert_to_np |
|
import torch |
|
import mediapipe as mp |
|
from utils.commons.multiprocess_utils import multiprocess_run_tqdm |
|
from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter |
|
from data_gen.utils.process_video.extract_segment_imgs import inpaint_torso_job, extract_background, save_rgb_image_to_path |
|
seg_model = MediapipeSegmenter() |
|
|
|
|
|
def extract_segment_job(img_name): |
|
try: |
|
img = cv2.imread(img_name) |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
|
|
segmap = seg_model._cal_seg_map(img) |
|
bg_img = extract_background([img], [segmap]) |
|
out_img_name = img_name.replace("/images_512/",f"/bg_img/").replace(".mp4", ".jpg") |
|
save_rgb_image_to_path(bg_img, out_img_name) |
|
|
|
com_img = img.copy() |
|
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("/images_512/",f"/com_imgs/") |
|
save_rgb_image_to_path(com_img, out_img_name) |
|
|
|
for mode in ['head', 'torso', 'person', 'torso_with_bg', 'bg']: |
|
out_img, _ = seg_model._seg_out_img_with_segmap(img, segmap, mode=mode) |
|
out_img_name = img_name.replace("/images_512/",f"/{mode}_imgs/") |
|
out_img = cv2.cvtColor(out_img, cv2.COLOR_RGB2BGR) |
|
try: os.makedirs(os.path.dirname(out_img_name), exist_ok=True) |
|
except: pass |
|
cv2.imwrite(out_img_name, out_img) |
|
|
|
inpaint_torso_img, inpaint_torso_with_bg_img, _, _ = inpaint_torso_job(img, segmap) |
|
out_img_name = img_name.replace("/images_512/",f"/inpaint_torso_imgs/") |
|
save_rgb_image_to_path(inpaint_torso_img, out_img_name) |
|
inpaint_torso_with_bg_img[bg_part] = bg_img[bg_part] |
|
out_img_name = img_name.replace("/images_512/",f"/inpaint_torso_with_com_bg_imgs/") |
|
save_rgb_image_to_path(inpaint_torso_with_bg_img, out_img_name) |
|
return 0 |
|
except Exception as e: |
|
print(e) |
|
return 1 |
|
|
|
def out_exist_job(img_name): |
|
out_name1 = img_name.replace("/images_512/", "/head_imgs/") |
|
out_name2 = img_name.replace("/images_512/", "/com_imgs/") |
|
out_name3 = img_name.replace("/images_512/", "/inpaint_torso_with_com_bg_imgs/") |
|
|
|
if os.path.exists(out_name1) and os.path.exists(out_name2) and os.path.exists(out_name3): |
|
return None |
|
else: |
|
return img_name |
|
|
|
def get_todo_img_names(img_names): |
|
todo_img_names = [] |
|
for i, res in multiprocess_run_tqdm(out_exist_job, img_names, num_workers=64): |
|
if res is not None: |
|
todo_img_names.append(res) |
|
return todo_img_names |
|
|
|
|
|
if __name__ == '__main__': |
|
import argparse, glob, tqdm, random |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--img_dir", default='./images_512') |
|
|
|
parser.add_argument("--ds_name", default='FFHQ') |
|
parser.add_argument("--num_workers", default=1, 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') |
|
|
|
args = parser.parse_args() |
|
img_dir = args.img_dir |
|
if args.ds_name == 'FFHQ_MV': |
|
img_name_pattern1 = os.path.join(img_dir, "ref_imgs/*.png") |
|
img_names1 = glob.glob(img_name_pattern1) |
|
img_name_pattern2 = os.path.join(img_dir, "mv_imgs/*.png") |
|
img_names2 = glob.glob(img_name_pattern2) |
|
img_names = img_names1 + img_names2 |
|
elif args.ds_name == 'FFHQ': |
|
img_name_pattern = os.path.join(img_dir, "*.png") |
|
img_names = glob.glob(img_name_pattern) |
|
|
|
img_names = sorted(img_names) |
|
random.seed(args.seed) |
|
random.shuffle(img_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(img_names) // total_process |
|
if process_id == total_process: |
|
img_names = img_names[process_id * num_samples_per_process : ] |
|
else: |
|
img_names = img_names[process_id * num_samples_per_process : (process_id+1) * num_samples_per_process] |
|
|
|
if not args.reset: |
|
img_names = get_todo_img_names(img_names) |
|
print(f"todo images number: {len(img_names)}") |
|
|
|
for vid_name in multiprocess_run_tqdm(extract_segment_job ,img_names, desc=f"Root process {args.process_id}: extracting segment images", num_workers=args.num_workers): |
|
pass |
