Spaces:
No application file
No application file
import argparse | |
import os | |
os.environ['CUDA_HOME'] = '/usr/local/cuda' | |
os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin' | |
from datetime import datetime | |
import gradio as gr | |
import spaces | |
import numpy as np | |
import torch | |
from diffusers.image_processor import VaeImageProcessor | |
from huggingface_hub import snapshot_download | |
from PIL import Image | |
torch.jit.script = lambda f: f | |
from model.cloth_masker import AutoMasker, vis_mask | |
from model.pipeline import CatVTONPipeline | |
from utils import init_weight_dtype, resize_and_crop, resize_and_padding | |
from test import morph_close, morph_open, extend_mask_downward, image_equal | |
import cv2 | |
# GPU์์ ํ์ฌ ํ ๋น๋ ๋ฉ๋ชจ๋ฆฌ ํ์ธ (GPU 0๋ฒ ๊ธฐ์ค) | |
#allocated_memory = torch.cuda.memory_allocated(0) # 0๋ฒ GPU์์ ํ ๋น๋ ๋ฉ๋ชจ๋ฆฌ ์์ ๋ฐํ | |
#print(f"GPU 0์์ ํ ๋น๋ ๋ฉ๋ชจ๋ฆฌ: {allocated_memory / (1024 ** 2)} MB") # MB๋ก ๋ณํํ์ฌ ์ถ๋ ฅ | |
# to chck | |
def parse_args(): | |
parser = argparse.ArgumentParser(description="Simple example of a training script.") | |
parser.add_argument( | |
"--base_model_path", | |
type=str, | |
# default="Kwai-Kolors/Kolors-Inpainting", | |
default="booksforcharlie/stable-diffusion-inpainting", | |
# default="stabilityai/stable-diffusion-2-inpainting", | |
# default="runwayml/stable-diffusion-inpainting", | |
help=( | |
"The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub." | |
), | |
) | |
parser.add_argument( | |
"--resume_path", | |
type=str, | |
default="zhengchong/CatVTON", | |
help=( | |
"The Path to the checkpoint of trained tryon model." | |
), | |
) | |
parser.add_argument( | |
"--output_dir", | |
type=str, | |
default="resource/demo/output", | |
help="The output directory where the model predictions will be written.", | |
) | |
parser.add_argument( | |
"--width", | |
type=int, | |
default=768, | |
help=( | |
"The resolution for input images, all the images in the train/validation dataset will be resized to this" | |
" resolution" | |
), | |
) | |
parser.add_argument( | |
"--height", | |
type=int, | |
default=1024, | |
help=( | |
"The resolution for input images, all the images in the train/validation dataset will be resized to this" | |
" resolution" | |
), | |
) | |
parser.add_argument( | |
"--repaint", | |
action="store_true", | |
help="Whether to repaint the result image with the original background." | |
) | |
parser.add_argument( | |
"--allow_tf32", | |
action="store_true", | |
default=True, | |
help=( | |
"Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see" | |
" https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" | |
), | |
) | |
parser.add_argument( | |
"--mixed_precision", | |
type=str, | |
default="no", | |
choices=["no", "fp16", "bf16"], | |
help=( | |
"Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >=" | |
" 1.10.and an Nvidia Ampere GPU. Default to the value of accelerate config of the current system or the" | |
" flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config." | |
), | |
) | |
args = parser.parse_args() | |
env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) | |
if env_local_rank != -1 and env_local_rank != args.local_rank: | |
args.local_rank = env_local_rank | |
return args | |
def image_grid(imgs, rows, cols): | |
assert len(imgs) == rows * cols | |
w, h = imgs[0].size | |
grid = Image.new("RGB", size=(cols * w, rows * h)) | |
for i, img in enumerate(imgs): | |
grid.paste(img, box=(i % cols * w, i // cols * h)) | |
return grid | |
args = parse_args() | |
repo_path = snapshot_download(repo_id=args.resume_path) | |
# Pipeline | |
pipeline = CatVTONPipeline( | |
base_ckpt=args.base_model_path, | |
attn_ckpt=repo_path, | |
attn_ckpt_version="mix", | |
weight_dtype=init_weight_dtype(args.mixed_precision), | |
use_tf32=args.allow_tf32, | |
device='cuda' | |
) | |
# AutoMasker | |
mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True) | |
automasker = AutoMasker( | |
densepose_ckpt=os.path.join(repo_path, "DensePose"), | |
schp_ckpt=os.path.join(repo_path, "SCHP"), | |
device='cuda', | |
) | |
# ๋งค๊ฐ๋ณ์๋ก fitting_type ์ถ๊ฐํด์ผ ํจ. cloth_type ๋ฐ์. | |
def submit_function( | |
person_image, | |
cloth_image, | |
cloth_type, | |
fitting_type, | |
num_inference_steps, | |
guidance_scale, | |
seed, | |
show_type | |
): | |
person_image, mask = person_image["background"], person_image["layers"][0] # person_image["layers"][0]์ด ์ ์ ๊ฐ ๊ทธ๋ฆฐ ๋ง์คํฌ ๋ ์ด์ด์. | |
mask = Image.open(mask).convert("L") | |
if len(np.unique(np.array(mask))) == 1: | |
mask = None # ์ฌ์ฉ์๊ฐ ๋ง์คํฌ๋ฅผ ๊ทธ๋ฆฌ์ง ์์ ๊ฒฝ์ฐ. | |
else: | |
mask = np.array(mask) | |
mask[mask > 0] = 255 # ๋ฐฐ๊ฒฝ์ด ๊ฒ์์. | |
mask = Image.fromarray(mask) | |
tmp_folder = args.output_dir | |
date_str = datetime.now().strftime("%Y%m%d%H%M%S") | |
result_save_path = os.path.join(tmp_folder, date_str[:8], date_str[8:] + ".png") | |
if not os.path.exists(os.path.join(tmp_folder, date_str[:8])): | |
os.makedirs(os.path.join(tmp_folder, date_str[:8])) | |
generator = None | |
if seed != -1: | |
generator = torch.Generator(device='cuda').manual_seed(seed) | |
person_image = Image.open(person_image).convert("RGB") | |
cloth_image = Image.open(cloth_image).convert("RGB") | |
person_image = resize_and_crop(person_image, (args.width, args.height)) | |
cloth_image = resize_and_padding(cloth_image, (args.width, args.height)) | |
#์์ธ์ฒ๋ฆฌ | |
#man | |
compare_image_mlvl0 = Image.open("./resource/demo/example/person/men/m_lvl0.png").convert("RGB") | |
compare_image_mlvl0 = resize_and_crop(compare_image_mlvl0, (args.width, args.height)) | |
compare_image_mlvl1 = Image.open("./resource/demo/example/person/men/m_lvl1.png").convert("RGB") | |
compare_image_mlvl1 = resize_and_crop(compare_image_mlvl1, (args.width, args.height)) | |
compare_image_mlvl2 = Image.open("./resource/demo/example/person/men/m_lvl2.png").convert("RGB") | |
compare_image_mlvl2 = resize_and_crop(compare_image_mlvl2, (args.width, args.height)) | |
compare_image_mlvl3 = Image.open("./resource/demo/example/person/men/m_lvl3.png").convert("RGB") | |
compare_image_mlvl3 = resize_and_crop(compare_image_mlvl3, (args.width, args.height)) | |
#womam | |
compare_image_wlvl0 = Image.open("./resource/demo/example/person/women/w_lvl0.png").convert("RGB") | |
compare_image_wlvl0 = resize_and_crop(compare_image_wlvl0, (args.width, args.height)) | |
compare_image_wlvl1 = Image.open("./resource/demo/example/person/women/w_lvl1.png").convert("RGB") | |
compare_image_wlvl1 = resize_and_crop(compare_image_wlvl1, (args.width, args.height)) | |
compare_image_wlvl2 = Image.open("./resource/demo/example/person/women/w_lvl2.png").convert("RGB") | |
compare_image_wlvl2 = resize_and_crop(compare_image_wlvl2, (args.width, args.height)) | |
compare_image_wlvl3 = Image.open("./resource/demo/example/person/women/w_lvl3.png").convert("RGB") | |
compare_image_wlvl3 = resize_and_crop(compare_image_wlvl3, (args.width, args.height)) | |
# Process mask | |
if mask is not None: | |
mask = resize_and_crop(mask, (args.width, args.height)) | |
else: | |
if image_equal(person_image, compare_image_mlvl3): | |
person_image2 = Image.open("./resource/demo/example/person/men/m_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl3_lower_sam_v2.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl3_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_wlvl3): | |
person_image2 = Image.open("./resource/demo/example/person/women/w_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
# ์ดํ ์ฒ๋ฆฌ | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl3_lower_sam_v2.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl3_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_mlvl2): | |
person_image2 = Image.open("./resource/demo/example/person/men/m_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl2_lower_sam_v2.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl2_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_wlvl2): | |
person_image2 = Image.open("./resource/demo/example/person/women/w_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
# ์ดํ ์ฒ๋ฆฌ | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl2_lower_sam_v2.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl2_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_mlvl1): | |
person_image2 = Image.open("./resource/demo/example/person/men/m_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl1_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl1_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_wlvl1): | |
person_image2 = Image.open("./resource/demo/example/person/women/w_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
# ์ดํ ์ฒ๋ฆฌ | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl1_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl1_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_mlvl0): | |
person_image2 = Image.open("./resource/demo/example/person/men/m_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl0_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl0_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
elif image_equal(person_image, compare_image_wlvl0): | |
person_image2 = Image.open("./resource/demo/example/person/women/w_lvl0.png").convert("RGB") | |
person_image2 = resize_and_crop(person_image2, (args.width, args.height)) | |
mask = automasker( | |
person_image2, | |
cloth_type | |
)['mask'] | |
# ์ดํ ์ฒ๋ฆฌ | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl0_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl0_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
mask_np = np.array(mask) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
if cloth_type == "upper": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_lower_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_upper_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
elif cloth_type == "lower": | |
kernel = np.ones((10, 10), np.uint8) | |
sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1) | |
result_np = np.where(sam_mask_upper_np== 255, 0, mask_np) | |
result_np = np.where(sam_mask_lower_np== 255, 255, result_np) | |
mask = Image.fromarray(result_np) | |
else: | |
mask = Image.fromarray(mask_np) | |
else: | |
mask = automasker( | |
person_image, | |
cloth_type | |
)['mask'] | |
# mask.save("./app_mask_created.png") | |
# ๊ฐ๋ bmi์ง์ ๋์ ์๋ฐํ์ ๊ฒฝ์ฐ, upper mask๋ฅผ ์ ํํ ์์ฑํด๋ด์ง ๋ชปํ๋ ๊ฒฝ์ฐ๊ฐ ์์ด ์๋์ผ๋ก ํ ๋ฒ ๋ ์ฒ๋ฆฌํด์ค. | |
# ํ์ด๋์จ ๋ถ๋ถ ๋ฐ์ด๋ฒ๋ฆฌ๊ธฐ (๋ ์ฌ์ฉ์๊ฐ ๊ทธ๋ฆฐ mask์ ๋ํด์๋ ์ํ๋๋ฉด ์๋๋ฏ๋ก, else๋ฌธ ์์ ๋ฃ์ด๋๊ธฐ) | |
#if cloth_type == "upper": | |
# height = (np.array(mask)).shape[0] | |
# y_threshold = int(height * 0.7) # ์ด๋ฏธ์ง ๋์ด์ 50ํผ์ผํธ ์ดํ. 50ํผ์ผํธ๊ฐ ๋ฑ ์ ๋นํจ. | |
# ๋ฐ๋ถ๋ถ ์ ๊ฑฐ | |
# mask = remove_bottom_part(np.array(mask), y_threshold) | |
# ์ ๋ฐฉ๋ฒ์ผ๋ก ํด๊ฒฐ ๋ถ๊ฐ์. ํ์ด๋์จ ๋ถ๋ถ | |
# input ๋ target ์ด๋ฏธ์ง๋ง๋ค, ์์ฑ๋๋ mask ์์ญ์ ํฌ๊ธฐ๊ฐ ๋ค๋ฅด๊ธฐ ๋๋ฌธ. mask ํ์ผ ์์ฒด์ ํฌ๊ธฐ๋ ๊ฐ์ ์ง์ธ์ . | |
# ์ถ๊ฐ๋ก Fitting Type์ ๋ฐ๋ผ ๋ง์คํฌ ์ฒ๋ฆฌ (else๋ฌธ ๋ด๋ถ) | |
if fitting_type == "standard": | |
# mlvl3์ ๋ํ upper lower ๊ฐ๊ฐ. | |
if image_equal(person_image, compare_image_mlvl3) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl3_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl3) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl3_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# mlvl2์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_mlvl2) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl2_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl2) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl2_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# mlvl1์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_mlvl1) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl1_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl1) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl1_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# mlvl0์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_mlvl0) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl0_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl0) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl0_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl3์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl3) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl3_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl3) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl3_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl2์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl2) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl2_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl2) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl2_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl1์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl1) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl1_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl1) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl1_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl0์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl0) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl0_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl0) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl0_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# ๊ทธ ์ธ ๋ํดํธ | |
else: | |
opened_mask = morph_open(mask) | |
extended_mask = extend_mask_downward(np.array(mask), pixels=100) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif fitting_type == "loose" : | |
# mlvl3์ ๋ํ upper lower ๊ฐ๊ฐ. | |
if image_equal(person_image, compare_image_mlvl3) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl3_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl3) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl3_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# mlvl2์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_mlvl2) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl2_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl2) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl2_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# mlvl1์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_mlvl1) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl1_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl1) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl1_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# mlvl0์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_mlvl0) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/m_lvl0_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_mlvl0) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/m_lvl0_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl3์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl3) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl3_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl3) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl3_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl2์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl2) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl2_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl2) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl2_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl1์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl1) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl1_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl1) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl1_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# wlvl0์ ๋ํ upper lower ๊ฐ๊ฐ. | |
elif image_equal(person_image, compare_image_wlvl0) and cloth_type == "upper": | |
opened_mask = morph_open(mask) | |
sam_mask_upper = Image.open("./resource/demo/example/person/sam/w_lvl0_upper_sam.png").convert("L") | |
sam_mask_upper = resize_and_crop(sam_mask_upper, (args.width, args.height)) | |
sam_mask_upper_np = np.array(sam_mask_upper) | |
extended_mask = extend_mask_downward(sam_mask_upper_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
elif image_equal(person_image, compare_image_wlvl0) and cloth_type == "lower": | |
opened_mask = morph_open(mask) | |
sam_mask_lower = Image.open("./resource/demo/example/person/sam/w_lvl0_lower_sam.png").convert("L") | |
sam_mask_lower = resize_and_crop(sam_mask_lower, (args.width, args.height)) | |
sam_mask_lower_np = np.array(sam_mask_lower) | |
extended_mask = extend_mask_downward(sam_mask_lower_np, pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# ๊ทธ ์ธ ๋ํดํธ | |
else: | |
opened_mask = morph_open(mask) | |
extended_mask = extend_mask_downward(np.array(mask), pixels=200) | |
#์ต์ข ๋ง์คํฌ ์ฒ๋ฆฌ (test.py ์ค๋ช ์ฐธ๊ณ ) | |
final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) | |
final_mask = morph_close(morph_open(final_mask)) | |
mask = final_mask | |
# ๋ธ๋ฌ์ฒ๋ฆฌ | |
mask = mask_processor.blur(mask, blur_factor=9) | |
# Inference | |
# try: | |
result_image = pipeline( | |
image=person_image, | |
condition_image=cloth_image, | |
mask=mask, | |
num_inference_steps=num_inference_steps, | |
guidance_scale=guidance_scale, | |
generator=generator | |
)[0] | |
# except Exception as e: | |
# raise gr.Error( | |
# "An error occurred. Please try again later: {}".format(e) | |
# ) | |
# Post-process | |
masked_person = vis_mask(person_image, mask) | |
save_result_image = image_grid([person_image, masked_person, cloth_image, result_image], 1, 4) | |
save_result_image.save(result_save_path) | |
if show_type == "result only": | |
return result_image | |
else: | |
width, height = person_image.size | |
if show_type == "input & result": | |
condition_width = width // 2 | |
conditions = image_grid([person_image, cloth_image], 2, 1) | |
else: | |
condition_width = width // 3 | |
conditions = image_grid([person_image, masked_person , cloth_image], 3, 1) | |
conditions = conditions.resize((condition_width, height), Image.NEAREST) | |
new_result_image = Image.new("RGB", (width + condition_width + 5, height)) | |
new_result_image.paste(conditions, (0, 0)) | |
new_result_image.paste(result_image, (condition_width + 5, 0)) | |
return new_result_image | |
def person_example_fn(image_path): | |
return image_path | |
HEADER = """ | |
""" | |
def app_gradio(): | |
with gr.Blocks(title="CatVTON") as demo: | |
gr.Markdown(HEADER) | |
with gr.Row(): | |
with gr.Column(scale=1, min_width=350): | |
with gr.Row(): | |
image_path = gr.Image( | |
type="filepath", | |
interactive=True, | |
visible=False, | |
) | |
person_image = gr.ImageEditor( | |
interactive=True, label="Person Image", type="filepath" | |
) | |
with gr.Row(): | |
with gr.Column(scale=1, min_width=230): | |
cloth_image = gr.Image( | |
interactive=True, label="Condition Image", type="filepath" | |
) | |
with gr.Column(scale=1, min_width=120): | |
gr.Markdown( | |
'<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `๐๏ธ` above to draw the Mask (higher priority)<br>2. Select the `Try-On Cloth Type` to generate automatically </span>' | |
) | |
cloth_type = gr.Radio( | |
label="Try-On Cloth Type", | |
choices=["upper", "lower", "overall"], | |
value="upper", | |
) | |
with gr.Column(scale=1, min_width=120): | |
gr.Markdown( | |
'<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `๐๏ธ` above to draw the Mask (higher priority)<br>2. Select the `Fitting Type` to generate automatically </span>' | |
) | |
fitting_type = gr.Radio( | |
label="Try-On Fitting Type", | |
choices=["fit", "standard", "loose"], | |
value="fit", # default | |
) | |
submit = gr.Button("Submit") | |
gr.Markdown( | |
'<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>' | |
) | |
gr.Markdown( | |
'<span style="color: #808080; font-size: small;">Advanced options can adjust details:<br>1. `Inference Step` may enhance details;<br>2. `CFG` is highly correlated with saturation;<br>3. `Random seed` may improve pseudo-shadow.</span>' | |
) | |
with gr.Accordion("Advanced Options", open=False): | |
num_inference_steps = gr.Slider( | |
label="Inference Step", minimum=10, maximum=100, step=5, value=50 | |
) | |
# Guidence Scale | |
guidance_scale = gr.Slider( | |
label="CFG Strenth", minimum=0.0, maximum=7.5, step=0.5, value=2.5 | |
) | |
# Random Seed | |
seed = gr.Slider( | |
label="Seed", minimum=-1, maximum=10000, step=1, value=42 | |
) | |
show_type = gr.Radio( | |
label="Show Type", | |
choices=["result only", "input & result", "input & mask & result"], | |
value="input & mask & result", | |
) | |
with gr.Column(scale=2, min_width=500): | |
result_image = gr.Image(interactive=False, label="Result") | |
with gr.Row(): | |
# Photo Examples | |
root_path = "resource/demo/example" | |
with gr.Column(): | |
men_exm = gr.Examples( | |
examples=[ | |
os.path.join(root_path, "person", "men", _) | |
for _ in os.listdir(os.path.join(root_path, "person", "men")) | |
], | |
examples_per_page=4, | |
inputs=image_path, | |
label="Person Examples โ ", | |
) | |
women_exm = gr.Examples( | |
examples=[ | |
os.path.join(root_path, "person", "women", _) | |
for _ in os.listdir(os.path.join(root_path, "person", "women")) | |
], | |
examples_per_page=4, | |
inputs=image_path, | |
label="Person Examples โก", | |
) | |
gr.Markdown( | |
'<span style="color: #808080; font-size: small;">*Person examples come from the demos of <a href="https://huggingface.co/spaces/levihsu/OOTDiffusion">OOTDiffusion</a> and <a href="https://www.outfitanyone.org">OutfitAnyone</a>. </span>' | |
) | |
with gr.Column(): | |
condition_upper_exm = gr.Examples( | |
examples=[ | |
os.path.join(root_path, "condition", "upper", _) | |
for _ in os.listdir(os.path.join(root_path, "condition", "upper")) | |
], | |
examples_per_page=4, | |
inputs=cloth_image, | |
label="Condition Upper Examples", | |
) | |
condition_overall_exm = gr.Examples( | |
examples=[ | |
os.path.join(root_path, "condition", "overall", _) | |
for _ in os.listdir(os.path.join(root_path, "condition", "overall")) | |
], | |
examples_per_page=4, | |
inputs=cloth_image, | |
label="Condition Overall Examples", | |
) | |
condition_person_exm = gr.Examples( | |
examples=[ | |
os.path.join(root_path, "condition", "person", _) | |
for _ in os.listdir(os.path.join(root_path, "condition", "person")) | |
], | |
examples_per_page=4, | |
inputs=cloth_image, | |
label="Condition Reference Person Examples", | |
) | |
condition_person_exm = gr.Examples( | |
examples=[ | |
os.path.join(root_path, "condition", "lower", _) | |
for _ in os.listdir(os.path.join(root_path, "condition", "lower")) | |
], | |
examples_per_page=4, | |
inputs=cloth_image, | |
label="Condition Reference lower Examples", | |
) | |
gr.Markdown( | |
'<span style="color: #808080; font-size: small;">*Condition examples come from the Internet. </span>' | |
) | |
image_path.change( | |
person_example_fn, inputs=image_path, outputs=person_image | |
) | |
#์ฌ๊ธฐ๋ ๋งค๊ฐ๋ณ์ fitting_type ์ถ๊ฐํด์ผ ํจ. | |
submit.click( | |
submit_function, | |
[ | |
person_image, | |
cloth_image, | |
cloth_type, | |
fitting_type, | |
num_inference_steps, | |
guidance_scale, | |
seed, | |
show_type, | |
], | |
result_image, | |
) | |
demo.queue().launch(share=True, show_error=True) | |
if __name__ == "__main__": | |
app_gradio() |