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 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 # 패키지 추가 import cv2 ''' def parse_args(): parser = argparse.ArgumentParser(description="Simple example of a training script.") parser.add_argument( "--base_model_path", type=str, default="booksforcharlie/stable-diffusion-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 args = parse_args() ''' RESUME_PATH = os.getenv("RESUME_PATH", "zhengchong/CatVTON") repo_path = snapshot_download(repo_id=RESUME_PATH) # 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', ) person_image = Image.open("./resource/demo/example/person/men/model_7.png").convert("RGB") mask = automasker( person_image, 'upper' )['mask'] # 여기서 리턴되는 mask는 PIL 이미지임.(cloth_masker.py 참조) # 참고로 ['densepose']로 densepose도 확인가능. ### 여기서 mask modify에 사용된 코드를 app.py에 체크해놓은 부분에 추가하면 된다! def remove_bottom_part(mask: np.ndarray, y_threshold: int): """ 이미지의 y_threshold 아래에 있는 부분을 삭제. :param mask: 입력 마스크 (numpy 배열) :param y_threshold: 제거할 Y 좌표 값 :return: 수정된 마스크 (numpy 배열) """ # y_threshold 아래의 모든 픽셀을 0으로 설정 mask[y_threshold:, :] = 0 return Image.fromarray(mask) # closing 연산 / fitting_mode가 standard 나 loose 일때만 사용하기 def morph_close(mask): mask_np = np.array(mask) kernel = np.ones((30, 30), np.uint8) # 커질수록 잘 연결됨 closed_mask = cv2.morphologyEx(mask_np, cv2.MORPH_CLOSE, kernel) return Image.fromarray(closed_mask) # opening 연산 / fitting_mode가 standard 나 loose 일때만 사용하기 def morph_open(mask): mask_np = np.array(mask) kernel = np.ones((30, 30), np.uint8) # 커질수록 잘 사라짐 #closed_mask = cv2.morphologyEx(mask_np, cv2.MORPH_CLOSE, kernel) opened_mask = cv2.morphologyEx(mask_np, cv2.MORPH_OPEN, kernel) #opened_mask는 numpy 연산 결과 이므로 PIL 이미지로 변환 필요 return Image.fromarray(opened_mask) def morph_open2(mask): mask_np = np.array(mask) kernel = np.ones((10, 10), np.uint8) # 커질수록 잘 사라짐 #closed_mask = cv2.morphologyEx(mask_np, cv2.MORPH_CLOSE, kernel) opened_mask = cv2.morphologyEx(mask_np, cv2.MORPH_OPEN, kernel) #opened_mask는 numpy 연산 결과 이므로 PIL 이미지로 변환 필요 return Image.fromarray(opened_mask) ## opened_mask = morph_open(mask) ## opened_mask.save('./opened_mask.png') #opened_mask는 PIL 이미지 형태로 반환되었으므로 (Image.fromarray()사용해서) .save를 바로 사용할 수 있다. #opened_mask2 = morph_open2(mask) #kernel = np.ones((50, 50), np.uint8) #opened_mask2 = cv2.dilate(np.array(opened_mask2), kernel, iterations=1) #opened_mask2 = Image.fromarray(opened_mask2) #opened_mask2 = mask_processor.blur(opened_mask2, blur_factor=9) #opened_mask2.save('./opened_mask2.png') # mask = mask_processor.blur(mask, blur_factor=9) ## mask.save("./test_mask.png") # 마스크를 PNG 파일로 저장 ## masked_person = vis_mask(person_image, mask) # app.py에서도 blur 처리 한 다음에 vis_mask 메서드 호출함. ## masked_person.save("./test_masked_person.png") # 마스크와 target img가 합쳐진 사진을 PNG 파일로 저장 # mask의 y축 음의 방향 이동 def extend_mask_downward(mask_image: np.ndarray, pixels: int) -> np.ndarray: """ y축 음의 방향으로 (아래로) 마스크 이미지를 확장하는 함수. :param mask_image: 마스크 이미지 (numpy 배열) :param pixels: 확장할 픽셀 수 :return: 확장된 마스크 이미지 (numpy 배열) """ # 이진화된 마스크를 만듦 mask = cv2.threshold(mask_image, 127, 255, cv2.THRESH_BINARY)[1] # 확장을 위한 커널. y축으로만 확장하기 위해 세로 길이를 크게 설정함 kernel = np.zeros((pixels, 1), np.uint8) # y축으로만 길어진 커널 # y축 음의 방향으로만 확장 (cv2.dilate 사용) extended_mask = cv2.dilate(mask, kernel, iterations=1) return Image.fromarray(extended_mask) def image_equal(img1, img2): return np.array_equal(np.array(img1), np.array(img2)) # automasker에 대한 fitting 정도를 default 처리하기 위한 코드 추가. def extend_mask_downward2(mask_image: np.ndarray, pixels: int) -> Image.Image: # 이진화된 마스크를 만듦 mask = cv2.threshold(mask_image, 127, 255, cv2.THRESH_BINARY)[1] height, width = mask.shape[:2] # 마스크가 시작하는 y 좌표와 끝나는 y 좌표 찾기 non_zero_rows = np.where(mask.any(axis=1))[0] if len(non_zero_rows) == 0: # 마스크가 비어있는 경우, 원본 이미지를 반환 return Image.fromarray(mask.copy()) mask_start_row = non_zero_rows[0] mask_end_row = non_zero_rows[-1] # 늘어난 마스크의 끝 y 좌표 계산 (이미지 높이를 초과하지 않도록) new_mask_end_row = min(mask_end_row + pixels, height - 1) # 원본 마스크 영역과 늘어난 마스크 영역의 길이 계산 original_mask_length = mask_end_row - mask_start_row + 1 stretched_mask_length = new_mask_end_row - mask_start_row + 1 # y 좌표 매핑 배열 생성 map_y = np.arange(height, dtype=np.float32) # 마스크 영역에 대한 y 좌표 재매핑 (선형적으로 늘리기) if stretched_mask_length > 1: map_y[mask_start_row:new_mask_end_row + 1] = np.linspace( mask_start_row, mask_end_row, stretched_mask_length ) else: map_y[mask_start_row:new_mask_end_row + 1] = mask_start_row # 마스크 아래 영역은 마지막 마스크 행으로 매핑 if new_mask_end_row + 1 < height: map_y[new_mask_end_row + 1:] = mask_end_row # map_y의 크기를 (height, width)로 확장 map_y = np.repeat(map_y[:, np.newaxis], width, axis=1) # x 좌표는 그대로 유지 map_x = np.tile(np.arange(width, dtype=np.float32), (height, 1)) # 이미지 리매핑 extended_mask = cv2.remap( mask, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0 ) return Image.fromarray(extended_mask) # 마스크를 y축 음의 방향으로 50픽셀 확장 ## extended_mask = extend_mask_downward(np.array(mask), pixels=100) # 확장된 마스크 저장 ## extended_mask.save('extended_mask_image.png') # 최종 마스크 저장 # fitting 정도에 따라, extended_mask 함수 호출 변수인 pixels를 조절하면 된다. # 정확도를 위해 그냥 dilation 하지 않고, y좌표가 약간 다른 두 마스크를 합쳤다. ## final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask)) ## final_mask = morph_close(morph_open(final_mask)) #불필요한 동떨어진 부분 삭제 -> 연결되지 않은 부분 연결 ## final_mask.save('final_mask_image.png') ## masked_person2 = vis_mask(person_image, final_mask) # app.py에서도 blur 처리 한 다음에 vis_mask 메서드 호출함. ## masked_person2.save("./test_masked_person2.png") # 마스크와 target img가 합쳐진 사진을 PNG 파일로 저장 #person_image = Image.open("path_to_image").convert("RGB") #standard_image = Image.open("./resource/demo/example/person/men/m_lvl3.png").convert("RGB") """ 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)) person_image2 = Image.open("./resource/demo/example/person/men/m_lvl0.png").convert("RGB") # 이걸 어느 bmi 레벨을 기준으로 쓸지는 뭐.. 실험해보면서 제일 좋은 거 정하면 됨. person_image2 = resize_and_crop(person_image2, (args.width, args.height)) mask = automasker( person_image2, "upper" )['mask'] mask.save("./first_mask.png") # 이후 처리 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_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) kernel = np.ones((10, 10), np.uint8) sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1) sam_mask_lower_np = cv2.dilate(sam_mask_lower_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) mask.save("./last_mask2.png") """