0135

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+detectron2/_C.cpython-39-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+playground.py
+__pycache__

README.md

@@ -0,0 +1,13 @@
+---
+title: CatVTON
+emoji: 🐈
+colorFrom: indigo
+colorTo: blue
+sdk: gradio
+sdk_version: 4.40.0
+app_file: app.py
+pinned: false
+license: cc-by-nc-sa-4.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

api.py

@@ -0,0 +1,1014 @@
+"""
+api.py는 fastapi를 사용하고 있지만,
+gradio를 사용한 웹 데모를 확인하고 싶다면 gradio 폴더 안에 있는 app-final.py를 사용하면 된다
+
+api.py의 마스킹 부분을 제대로 수정한게,
+api2.py임
+
+
+"""
+
+
+import os
+os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
+from datetime import datetime
+
+from pydantic import BaseModel
+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
+
+from fastapi import FastAPI, File, Form, UploadFile
+from typing import List
+from typing import Optional
+import shutil
+
+from fastapi.responses import JSONResponse
+import uuid
+import base64
+from io import BytesIO
+
+from fastapi.middleware.cors import CORSMiddleware
+
+app = FastAPI()
+origins = [
+    "http://localhost",
+    "http://localhost:8080",
+    "http://localhost:3000",
+    "http://127.0.0.1:8080",
+    "http://127.0.0.1:3000",
+]
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+print ('starting app')
+# api 연결하면서 추가한 코드 
+def pil_to_base64(img: Image.Image) -> str:
+    buffer = BytesIO()
+    img.save(buffer, format="PNG")  # PNG 형식으로 저장
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+# GPU에서 현재 할당된 메모리 확인 (GPU 0번 기준)
+#allocated_memory = torch.cuda.memory_allocated(0) 
+#print(f"GPU 0에서 할당된 메모리: {allocated_memory / (1024 ** 2)} MB")  # MB로 변환하여 출력
+
+
+# 설정값을 환경 변수로 정의
+BASE_MODEL_PATH = os.getenv("BASE_MODEL_PATH", "booksforcharlie/stable-diffusion-inpainting")
+RESUME_PATH = os.getenv("RESUME_PATH", "zhengchong/CatVTON")
+OUTPUT_DIR = os.getenv("OUTPUT_DIR", "resource/demo/output")
+WIDTH = int(os.getenv("WIDTH", 768))
+HEIGHT = int(os.getenv("HEIGHT", 1024))
+
+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
+
+
+repo_path = snapshot_download(repo_id=RESUME_PATH)
+print ('repo_path')
+# Pipeline
+pipeline = CatVTONPipeline(
+    base_ckpt=BASE_MODEL_PATH,
+    attn_ckpt=repo_path,
+    attn_ckpt_version="mix",
+    weight_dtype=init_weight_dtype("no"),
+    use_tf32=True,
+    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', 
+)
+
+@spaces.GPU(duration=120)
+# 매개변수로 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)
+    mask = None
+    tmp_folder = "resource/demo/output"
+    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, (768, 1024))
+    cloth_image = resize_and_padding(cloth_image, (768, 1024))
+    
+    #예외처리
+    #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, (768, 1024))
+
+    compare_image_mlvl1 = Image.open("./resource/demo/example/person/men/m_lvl1.png").convert("RGB")
+    compare_image_mlvl1 = resize_and_crop(compare_image_mlvl1, (768, 1024))
+
+    compare_image_mlvl2 = Image.open("./resource/demo/example/person/men/m_lvl2.png").convert("RGB")
+    compare_image_mlvl2 = resize_and_crop(compare_image_mlvl2, (768, 1024))
+
+    compare_image_mlvl3 = Image.open("./resource/demo/example/person/men/m_lvl3.png").convert("RGB")
+    compare_image_mlvl3 = resize_and_crop(compare_image_mlvl3, (768, 1024))
+
+    #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, (768, 1024))
+    
+    compare_image_wlvl1 = Image.open("./resource/demo/example/person/women/w_lvl1.png").convert("RGB")
+    compare_image_wlvl1 = resize_and_crop(compare_image_wlvl1, (768, 1024))
+    
+    compare_image_wlvl2 = Image.open("./resource/demo/example/person/women/w_lvl2.png").convert("RGB")
+    compare_image_wlvl2 = resize_and_crop(compare_image_wlvl2, (768, 1024))
+    
+    compare_image_wlvl3 = Image.open("./resource/demo/example/person/women/w_lvl3.png").convert("RGB")
+    compare_image_wlvl3 = resize_and_crop(compare_image_wlvl3, (768, 1024))
+
+    # Process mask
+    if mask is not None:
+        mask = resize_and_crop(mask, (768, 1024))
+    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, (768, 1024))
+            mask = automasker(
+                person_image2,
+                cloth_type
+            )['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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+
+            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 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, (768, 1024))
+            mask = automasker(
+                person_image2,
+                cloth_type
+            )['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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+            
+
+        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, (768, 1024))
+            mask = automasker(
+                person_image2,
+                cloth_type
+            )['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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+            
+
+
+        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, (768, 1024))
+            mask = automasker(
+                person_image2,
+                cloth_type
+            )['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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+
+
+
+        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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+
+        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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+
+        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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+
+            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 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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+
+            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)
+
+            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)
+
+        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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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, (768, 1024))
+            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": result_image, "masked_person": masked_person}
+    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
+
+
+###### Fastapi의 api
+
+
+
+
+# FastAPI 함수 정의
+@app.post("/process-image")
+async def process_image(
+    cloth_type: str = Form(...),
+    fitting_type: str = Form(...),
+    person_image: UploadFile = File(...),
+    cloth_image: UploadFile = File(...)
+):
+    try:
+        # 고유한 파일 이름 생성
+        person_filename = f"received_{uuid.uuid4().hex}_{person_image.filename}"
+        cloth_filename = f"received_{uuid.uuid4().hex}_{cloth_image.filename}"
+
+        print ('person_filename: ', person_filename)
+        print ('cloth_filename: ', cloth_filename)
+
+        # 이미지 저장 디렉토리 생성
+        os.makedirs("uploads", exist_ok=True)
+
+        # 업로드된 이미지 저장
+        person_path = os.path.join("uploads", person_filename)
+        cloth_path = os.path.join("uploads", cloth_filename)
+
+        with open(person_path, "wb") as buffer:
+            shutil.copyfileobj(person_image.file, buffer)
+
+        with open(cloth_path, "wb") as buffer:
+            shutil.copyfileobj(cloth_image.file, buffer)
+
+        # 이미지 처리 함수 호출
+        result = submit_function(
+            person_image=person_path,
+            cloth_image=cloth_path,
+            cloth_type=cloth_type,
+            fitting_type=fitting_type,
+            num_inference_steps=25,  
+            guidance_scale=2.5,       
+            seed=42,                  
+            show_type='result only'   
+        )
+
+        print ('processing done')
+        # 반환된 이미지 추출
+        result_image = result['result_image']
+        masked_person = result['masked_person']
+        result_image.save('results/result.png')
+        # 이미지를 Base64로 인코딩
+        result_image_b64 = pil_to_base64(result_image)
+        masked_person_b64 = pil_to_base64(masked_person)
+
+        # 임시 파일 삭제 (필요 시)
+        os.remove(person_path)
+        os.remove(cloth_path)
+
+        return {
+            "message": "이미지가 처리되었습니다",
+            "result_image": result_image_b64,
+            "masked_person": masked_person_b64
+        }
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"message": f"오류 발생: {str(e)}"})
+
+
+
+
+@app.post("/send-to-ssh")
+async def send_to_ssh(
+    cloth_type: str = Form(...),
+    fitting_type: str = Form(...),
+    person_image: UploadFile = File(...),  # 이미지 파일 업로드로 처리
+    cloth_image: UploadFile = File(...)
+):
+    # 받은 데이터를 처리하거나 저장하는 로직
+    return {"message": "데이터가 성공적으로 처리되었습니다."}
+
+@app.get('/test')
+async def test():
+    return JSONResponse(status_code=200, content={"message": "hello"})

api2.py

@@ -0,0 +1,1283 @@
+"""
+api.py는 fastapi를 사용하고 있지만,
+gradio를 사용한 웹 데모를 확인하고 싶다면 gradio 폴더 안에 있는 app-final.py를 사용하면 된다
+
+api.py의 마스킹 부분을 제대로 수정한게,
+api2.py임 
+
+"""
+
+
+import os
+os.environ['CUDA_HOME'] = '/usr/local/cuda'
+os.environ['PATH'] = os.environ['PATH'] + ':/usr/local/cuda/bin'
+from datetime import datetime
+
+from pydantic import BaseModel
+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, morph_open2, extend_mask_downward, extend_mask_downward2, image_equal
+import cv2
+
+from fastapi import FastAPI, File, Form, UploadFile, WebSocket, WebSocketDisconnect
+from typing import List
+from typing import Optional
+import shutil
+
+from fastapi.responses import JSONResponse
+import uuid
+import base64
+from io import BytesIO
+
+from fastapi.middleware.cors import CORSMiddleware
+import asyncio
+
+app = FastAPI()
+origins = [
+    "http://localhost",
+    "http://localhost:8080",
+    "http://localhost:3000",
+    "http://127.0.0.1:8080",
+    "http://127.0.0.1:3000",
+]
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # You can set this to specific origins instead of '*' for security
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+# app.add_middleware(
+#     CORSMiddleware,
+#     allow_origins=origins,
+#     allow_credentials=True,
+#     allow_methods=["*"],
+#     allow_headers=["*"],
+# )
+print ('starting app')
+# api 연결하면서 추가한 코드 
+def pil_to_base64(img: Image.Image) -> str:
+    buffer = BytesIO()
+    img.save(buffer, format="PNG")  # PNG 형식으로 저장
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+# GPU에서 현재 할당된 메모리 확인 (GPU 0번 기준)
+#allocated_memory = torch.cuda.memory_allocated(0) 
+#print(f"GPU 0에서 할당된 메모리: {allocated_memory / (1024 ** 2)} MB")  # MB로 변환하여 출력
+
+
+# 설정값을 환경 변수로 정의
+BASE_MODEL_PATH = os.getenv("BASE_MODEL_PATH", "booksforcharlie/stable-diffusion-inpainting")
+RESUME_PATH = os.getenv("RESUME_PATH", "zhengchong/CatVTON")
+OUTPUT_DIR = os.getenv("OUTPUT_DIR", "resource/demo/output")
+WIDTH = int(os.getenv("WIDTH", 768))
+HEIGHT = int(os.getenv("HEIGHT", 1024))
+
+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
+
+
+repo_path = snapshot_download(repo_id=RESUME_PATH)
+print ('repo_path')
+# Pipeline
+pipeline = CatVTONPipeline(
+    base_ckpt=BASE_MODEL_PATH,
+    attn_ckpt=repo_path,
+    attn_ckpt_version="mix",
+    weight_dtype=init_weight_dtype("no"),
+    use_tf32=True,
+    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', 
+)
+
+
+@spaces.GPU(duration=120)
+# 매개변수로 fitting_type 추가해야 함. cloth_type 밑에.
+
+async 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)
+    mask = None
+    tmp_folder = "resource/demo/output"
+    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, (768, 1024))
+    cloth_image = resize_and_padding(cloth_image, (768, 1024))
+    
+    #예외처리
+    #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, (768, 1024))
+
+    compare_image_mlvl1 = Image.open("./resource/demo/example/person/men/m_lvl1.png").convert("RGB")
+    compare_image_mlvl1 = resize_and_crop(compare_image_mlvl1, (768, 1024))
+
+    compare_image_mlvl2 = Image.open("./resource/demo/example/person/men/m_lvl2.png").convert("RGB")
+    compare_image_mlvl2 = resize_and_crop(compare_image_mlvl2, (768, 1024))
+
+    compare_image_mlvl3 = Image.open("./resource/demo/example/person/men/m_lvl3.png").convert("RGB")
+    compare_image_mlvl3 = resize_and_crop(compare_image_mlvl3, (768, 1024))
+
+    #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, (768, 1024))
+    
+    compare_image_wlvl1 = Image.open("./resource/demo/example/person/women/w_lvl1.png").convert("RGB")
+    compare_image_wlvl1 = resize_and_crop(compare_image_wlvl1, (768, 1024))
+    
+    compare_image_wlvl2 = Image.open("./resource/demo/example/person/women/w_lvl2.png").convert("RGB")
+    compare_image_wlvl2 = resize_and_crop(compare_image_wlvl2, (768, 1024))
+    
+    compare_image_wlvl3 = Image.open("./resource/demo/example/person/women/w_lvl3.png").convert("RGB")
+    compare_image_wlvl3 = resize_and_crop(compare_image_wlvl3, (768, 1024))
+
+# Process mask
+    if mask is not None:
+        mask = resize_and_crop(mask, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, 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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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
+
+        # 그 외 디폴트
+        # flag2
+        else:
+            if cloth_type == "upper":
+                opened_mask = morph_open(mask)
+                extended_mask = extend_mask_downward2(np.array(mask), pixels=50)
+                mask = extended_mask
+            else:
+                opened_mask = morph_open(mask)
+                extended_mask = extend_mask_downward(np.array(mask), pixels=70)
+                final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask))
+                final_mask = 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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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":
+            # flag
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1110)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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)) 원래 주석 아닌데 주석처리함. 1110
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_upper_np = np.array(sam_mask_upper)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_upper_np = cv2.dilate(sam_mask_upper_np, kernel, iterations=1)
+
+            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, (WIDTH, HEIGHT))
+            sam_mask_lower_np = np.array(sam_mask_lower)
+
+            # 추가한 부분 (1113)
+            kernel = np.ones((40, 40), np.uint8)
+            sam_mask_lower_np = cv2.dilate(sam_mask_lower_np, kernel, iterations=1)
+
+            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:
+            if cloth_type == "upper":
+                opened_mask = morph_open(mask)
+                extended_mask = extend_mask_downward2(np.array(mask), pixels=70)
+                mask = extended_mask
+            else:
+                opened_mask = morph_open(mask)
+                extended_mask = extend_mask_downward(np.array(mask), pixels=100)
+                final_mask = Image.fromarray(np.array(opened_mask) | np.array(extended_mask))
+                #morph_open 감싸던 morph_close 없앰.
+                final_mask = 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]
+    # print(pipeline.progress)
+    # except Exception as e:
+    #     raise gr.Error(
+    #         "An error occurred. Please try again later: {}".format(e)
+    #     )
+
+    #inside of submit_function
+    pipeline.progress = 0.02
+    pipeline_future = asyncio.to_thread(
+        pipeline,
+        image=person_image,
+        condition_image=cloth_image,
+        mask=mask,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    )
+    
+    # Immediately print progress after starting pipeline()
+    #print(pipeline.progress)
+    await asyncio.sleep(1)
+    async def print_progress():
+        print('starting printing progress: ', pipeline_future)
+        while pipeline.progress < 0.9:
+            print(f"Progress: {pipeline.progress}")
+            await asyncio.sleep(1)
+        # Optionally, print the final progress after completion
+        print(f"Final Progress: {pipeline.progress}")
+
+    progress_task = asyncio.create_task(print_progress())
+    #await print_progress()
+    #progress_task
+    result = await pipeline_future
+
+    # Wait for the result_image after the print statement
+    result_image = result[0]
+    await progress_task
+
+    # 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": result_image, "masked_person": masked_person}
+    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
+
+
+
+#add websocket related function inside process_image or submit_function so that it sends the progress every 1 second 
+
+# FastAPI 함수 정의
+@app.post("/process-image")
+async def process_image(
+    cloth_type: str = Form(...),
+    fitting_type: str = Form(...),
+    person_image: UploadFile = File(...),
+    cloth_image: UploadFile = File(...)
+):
+    try:
+        # 고유한 파일 이름 생성
+        person_filename = f"received_{uuid.uuid4().hex}_{person_image.filename}"
+        cloth_filename = f"received_{uuid.uuid4().hex}_{cloth_image.filename}"
+
+        print ('person_filename: ', person_filename)
+        print ('cloth_filename: ', cloth_filename)
+
+        # 이미지 저장 디렉토리 생성
+        os.makedirs("uploads", exist_ok=True)
+
+        # 업로드된 이미지 저장
+        person_path = os.path.join("uploads", person_filename)
+        cloth_path = os.path.join("uploads", cloth_filename)
+
+        with open(person_path, "wb") as buffer:
+            shutil.copyfileobj(person_image.file, buffer)
+
+        with open(cloth_path, "wb") as buffer:
+            shutil.copyfileobj(cloth_image.file, buffer)
+
+
+        pipeline.progress = 0.02
+
+        # 이미지 처리 함수 호출
+        # num inference : 50보다 25가 나은듯
+        result = await submit_function(
+            person_image=person_path,
+            cloth_image=cloth_path,
+            cloth_type=cloth_type,
+            fitting_type=fitting_type,
+            num_inference_steps=25,  
+            guidance_scale=2.5,       
+            seed=42,                  
+            show_type='result only'
+        )
+
+        print ('processing done')
+        # 반환된 이미지 추출
+        result_image = result['result_image']
+        masked_person = result['masked_person']
+        result_image.save('results/result.png')
+        # 이미지를 Base64로 인코딩
+        result_image_b64 = pil_to_base64(result_image)
+        masked_person_b64 = pil_to_base64(masked_person)
+
+        # 임시 파일 삭제 (필요 시)
+        os.remove(person_path)
+        os.remove(cloth_path)
+
+        return {
+            "message": "이미지가 처리되었습니다",
+            "result_image": result_image_b64,
+            "masked_person": masked_person_b64
+        }
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"message": f"오류 발생: {str(e)}"})
+#blocking language
+#syncronous language
+@app.websocket("/ws/")
+async def websocket_endpoint(websocket: WebSocket):
+    print('awaiting websocket connection')
+    await websocket.accept()
+    print ('websocket accepted')
+    try:
+        while True:
+            if pipeline.progress >= 0.94:
+                await asyncio.sleep(0.1);
+            if pipeline.progress >= 0.94:
+                break
+            if pipeline.progress < 0.02:
+                pipeline.progress = 0.02
+            await websocket.send_text(f"{pipeline.progress}")
+            await asyncio.sleep(0.5)
+        await websocket.send_text("Processing complete")
+    except WebSocketDisconnect:
+        print(f"Client disconnected")
+    finally:
+        await websocket.close()
+
+
+@app.post("/send-to-ssh")
+async def send_to_ssh(
+    cloth_type: str = Form(...),
+    fitting_type: str = Form(...),
+    person_image: UploadFile = File(...),  # 이미지 파일 업로드로 처리
+    cloth_image: UploadFile = File(...)
+):
+    # 받은 데이터를 처리하거나 저장하는 로직
+    return {"message": "데이터가 성공적으로 처리되었습니다."}
+
+@app.get('/test')
+async def test():
+    return JSONResponse(status_code=200, content={"message": "hello"})
+
+
+def person_example_fn(image_path):
+    return image_path
+

app-final.py

@@ -0,0 +1,1187 @@
+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', 
+)
+
+@spaces.GPU(duration=120)
+# 매개변수로 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()

densepose/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+from .data.datasets import builtin  # just to register data
+from .converters import builtin as builtin_converters  # register converters
+from .config import (
+    add_densepose_config,
+    add_densepose_head_config,
+    add_hrnet_config,
+    add_dataset_category_config,
+    add_bootstrap_config,
+    load_bootstrap_config,
+)
+from .structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
+from .evaluation import DensePoseCOCOEvaluator
+from .modeling.roi_heads import DensePoseROIHeads
+from .modeling.test_time_augmentation import (
+    DensePoseGeneralizedRCNNWithTTA,
+    DensePoseDatasetMapperTTA,
+)
+from .utils.transform import load_from_cfg
+from .modeling.hrfpn import build_hrfpn_backbone

densepose/config.py

@@ -0,0 +1,277 @@
+# -*- coding = utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+# pyre-ignore-all-errors
+
+from detectron2.config import CfgNode as CN
+
+
+def add_dataset_category_config(cfg: CN) -> None:
+    """
+    Add config for additional category-related dataset options
+     - category whitelisting
+     - category mapping
+    """
+    _C = cfg
+    _C.DATASETS.CATEGORY_MAPS = CN(new_allowed=True)
+    _C.DATASETS.WHITELISTED_CATEGORIES = CN(new_allowed=True)
+    # class to mesh mapping
+    _C.DATASETS.CLASS_TO_MESH_NAME_MAPPING = CN(new_allowed=True)
+
+
+def add_evaluation_config(cfg: CN) -> None:
+    _C = cfg
+    _C.DENSEPOSE_EVALUATION = CN()
+    # evaluator type, possible values:
+    #  - "iou": evaluator for models that produce iou data
+    #  - "cse": evaluator for models that produce cse data
+    _C.DENSEPOSE_EVALUATION.TYPE = "iou"
+    # storage for DensePose results, possible values:
+    #  - "none": no explicit storage, all the results are stored in the
+    #            dictionary with predictions, memory intensive;
+    #            historically the default storage type
+    #  - "ram": RAM storage, uses per-process RAM storage, which is
+    #           reduced to a single process storage on later stages,
+    #           less memory intensive
+    #  - "file": file storage, uses per-process file-based storage,
+    #            the least memory intensive, but may create bottlenecks
+    #            on file system accesses
+    _C.DENSEPOSE_EVALUATION.STORAGE = "none"
+    # minimum threshold for IOU values: the lower its values is,
+    # the more matches are produced (and the higher the AP score)
+    _C.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD = 0.5
+    # Non-distributed inference is slower (at inference time) but can avoid RAM OOM
+    _C.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE = True
+    # evaluate mesh alignment based on vertex embeddings, only makes sense in CSE context
+    _C.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT = False
+    # meshes to compute mesh alignment for
+    _C.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES = []
+
+
+def add_bootstrap_config(cfg: CN) -> None:
+    """ """
+    _C = cfg
+    _C.BOOTSTRAP_DATASETS = []
+    _C.BOOTSTRAP_MODEL = CN()
+    _C.BOOTSTRAP_MODEL.WEIGHTS = ""
+    _C.BOOTSTRAP_MODEL.DEVICE = "cuda"
+
+
+def get_bootstrap_dataset_config() -> CN:
+    _C = CN()
+    _C.DATASET = ""
+    # ratio used to mix data loaders
+    _C.RATIO = 0.1
+    # image loader
+    _C.IMAGE_LOADER = CN(new_allowed=True)
+    _C.IMAGE_LOADER.TYPE = ""
+    _C.IMAGE_LOADER.BATCH_SIZE = 4
+    _C.IMAGE_LOADER.NUM_WORKERS = 4
+    _C.IMAGE_LOADER.CATEGORIES = []
+    _C.IMAGE_LOADER.MAX_COUNT_PER_CATEGORY = 1_000_000
+    _C.IMAGE_LOADER.CATEGORY_TO_CLASS_MAPPING = CN(new_allowed=True)
+    # inference
+    _C.INFERENCE = CN()
+    # batch size for model inputs
+    _C.INFERENCE.INPUT_BATCH_SIZE = 4
+    # batch size to group model outputs
+    _C.INFERENCE.OUTPUT_BATCH_SIZE = 2
+    # sampled data
+    _C.DATA_SAMPLER = CN(new_allowed=True)
+    _C.DATA_SAMPLER.TYPE = ""
+    _C.DATA_SAMPLER.USE_GROUND_TRUTH_CATEGORIES = False
+    # filter
+    _C.FILTER = CN(new_allowed=True)
+    _C.FILTER.TYPE = ""
+    return _C
+
+
+def load_bootstrap_config(cfg: CN) -> None:
+    """
+    Bootstrap datasets are given as a list of `dict` that are not automatically
+    converted into CfgNode. This method processes all bootstrap dataset entries
+    and ensures that they are in CfgNode format and comply with the specification
+    """
+    if not cfg.BOOTSTRAP_DATASETS:
+        return
+
+    bootstrap_datasets_cfgnodes = []
+    for dataset_cfg in cfg.BOOTSTRAP_DATASETS:
+        _C = get_bootstrap_dataset_config().clone()
+        _C.merge_from_other_cfg(CN(dataset_cfg))
+        bootstrap_datasets_cfgnodes.append(_C)
+    cfg.BOOTSTRAP_DATASETS = bootstrap_datasets_cfgnodes
+
+
+def add_densepose_head_cse_config(cfg: CN) -> None:
+    """
+    Add configuration options for Continuous Surface Embeddings (CSE)
+    """
+    _C = cfg
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE = CN()
+    # Dimensionality D of the embedding space
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE = 16
+    # Embedder specifications for various mesh IDs
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS = CN(new_allowed=True)
+    # normalization coefficient for embedding distances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_DIST_GAUSS_SIGMA = 0.01
+    # normalization coefficient for geodesic distances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.GEODESIC_DIST_GAUSS_SIGMA = 0.01
+    # embedding loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_WEIGHT = 0.6
+    # embedding loss name, currently the following options are supported:
+    # - EmbeddingLoss: cross-entropy on vertex labels
+    # - SoftEmbeddingLoss: cross-entropy on vertex label combined with
+    #    Gaussian penalty on distance between vertices
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_NAME = "EmbeddingLoss"
+    # optimizer hyperparameters
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR = 1.0
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR = 1.0
+    # Shape to shape cycle consistency loss parameters:
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
+    # shape to shape cycle consistency loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.025
+    # norm type used for loss computation
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
+    # normalization term for embedding similarity matrices
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.TEMPERATURE = 0.05
+    # maximum number of vertices to include into shape to shape cycle loss
+    # if negative or zero, all vertices are considered
+    # if positive, random subset of vertices of given size is considered
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.MAX_NUM_VERTICES = 4936
+    # Pixel to shape cycle consistency loss parameters:
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
+    # pixel to shape cycle consistency loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.0001
+    # norm type used for loss computation
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
+    # map images to all meshes and back (if false, use only gt meshes from the batch)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.USE_ALL_MESHES_NOT_GT_ONLY = False
+    # Randomly select at most this number of pixels from every instance
+    # if negative or zero, all vertices are considered
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NUM_PIXELS_TO_SAMPLE = 100
+    # normalization factor for pixel to pixel distances (higher value = smoother distribution)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.PIXEL_SIGMA = 5.0
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_PIXEL_TO_VERTEX = 0.05
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_VERTEX_TO_PIXEL = 0.05
+
+
+def add_densepose_head_config(cfg: CN) -> None:
+    """
+    Add config for densepose head.
+    """
+    _C = cfg
+
+    _C.MODEL.DENSEPOSE_ON = True
+
+    _C.MODEL.ROI_DENSEPOSE_HEAD = CN()
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NAME = ""
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS = 8
+    # Number of parts used for point labels
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES = 24
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL = 4
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM = 512
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL = 3
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UP_SCALE = 2
+    _C.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE = 112
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE = "ROIAlignV2"
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION = 28
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO = 2
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS = 2  # 15 or 2
+    # Overlap threshold for an RoI to be considered foreground (if >= FG_IOU_THRESHOLD)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.FG_IOU_THRESHOLD = 0.7
+    # Loss weights for annotation masks.(14 Parts)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.INDEX_WEIGHTS = 5.0
+    # Loss weights for surface parts. (24 Parts)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.PART_WEIGHTS = 1.0
+    # Loss weights for UV regression.
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POINT_REGRESSION_WEIGHTS = 0.01
+    # Coarse segmentation is trained using instance segmentation task data
+    _C.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS = False
+    # For Decoder
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_ON = True
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NUM_CLASSES = 256
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_CONV_DIMS = 256
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NORM = ""
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_COMMON_STRIDE = 4
+    # For DeepLab head
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB = CN()
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NORM = "GN"
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NONLOCAL_ON = 0
+    # Predictor class name, must be registered in DENSEPOSE_PREDICTOR_REGISTRY
+    # Some registered predictors:
+    #   "DensePoseChartPredictor": predicts segmentation and UV coordinates for predefined charts
+    #   "DensePoseChartWithConfidencePredictor": predicts segmentation, UV coordinates
+    #       and associated confidences for predefined charts (default)
+    #   "DensePoseEmbeddingWithConfidencePredictor": predicts segmentation, embeddings
+    #       and associated confidences for CSE
+    _C.MODEL.ROI_DENSEPOSE_HEAD.PREDICTOR_NAME = "DensePoseChartWithConfidencePredictor"
+    # Loss class name, must be registered in DENSEPOSE_LOSS_REGISTRY
+    # Some registered losses:
+    #   "DensePoseChartLoss": loss for chart-based models that estimate
+    #      segmentation and UV coordinates
+    #   "DensePoseChartWithConfidenceLoss": loss for chart-based models that estimate
+    #      segmentation, UV coordinates and the corresponding confidences (default)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.LOSS_NAME = "DensePoseChartWithConfidenceLoss"
+    # Confidences
+    # Enable learning UV confidences (variances) along with the actual values
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE = CN({"ENABLED": False})
+    # UV confidence lower bound
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.EPSILON = 0.01
+    # Enable learning segmentation confidences (variances) along with the actual values
+    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE = CN({"ENABLED": False})
+    # Segmentation confidence lower bound
+    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE.EPSILON = 0.01
+    # Statistical model type for confidence learning, possible values:
+    # - "iid_iso": statistically independent identically distributed residuals
+    #    with isotropic covariance
+    # - "indep_aniso": statistically independent residuals with anisotropic
+    #    covariances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.TYPE = "iid_iso"
+    # List of angles for rotation in data augmentation during training
+    _C.INPUT.ROTATION_ANGLES = [0]
+    _C.TEST.AUG.ROTATION_ANGLES = ()  # Rotation TTA
+
+    add_densepose_head_cse_config(cfg)
+
+
+def add_hrnet_config(cfg: CN) -> None:
+    """
+    Add config for HRNet backbone.
+    """
+    _C = cfg
+
+    # For HigherHRNet w32
+    _C.MODEL.HRNET = CN()
+    _C.MODEL.HRNET.STEM_INPLANES = 64
+    _C.MODEL.HRNET.STAGE2 = CN()
+    _C.MODEL.HRNET.STAGE2.NUM_MODULES = 1
+    _C.MODEL.HRNET.STAGE2.NUM_BRANCHES = 2
+    _C.MODEL.HRNET.STAGE2.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE2.NUM_BLOCKS = [4, 4]
+    _C.MODEL.HRNET.STAGE2.NUM_CHANNELS = [32, 64]
+    _C.MODEL.HRNET.STAGE2.FUSE_METHOD = "SUM"
+    _C.MODEL.HRNET.STAGE3 = CN()
+    _C.MODEL.HRNET.STAGE3.NUM_MODULES = 4
+    _C.MODEL.HRNET.STAGE3.NUM_BRANCHES = 3
+    _C.MODEL.HRNET.STAGE3.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE3.NUM_BLOCKS = [4, 4, 4]
+    _C.MODEL.HRNET.STAGE3.NUM_CHANNELS = [32, 64, 128]
+    _C.MODEL.HRNET.STAGE3.FUSE_METHOD = "SUM"
+    _C.MODEL.HRNET.STAGE4 = CN()
+    _C.MODEL.HRNET.STAGE4.NUM_MODULES = 3
+    _C.MODEL.HRNET.STAGE4.NUM_BRANCHES = 4
+    _C.MODEL.HRNET.STAGE4.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+    _C.MODEL.HRNET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
+    _C.MODEL.HRNET.STAGE4.FUSE_METHOD = "SUM"
+
+    _C.MODEL.HRNET.HRFPN = CN()
+    _C.MODEL.HRNET.HRFPN.OUT_CHANNELS = 256
+
+
+def add_densepose_config(cfg: CN) -> None:
+    add_densepose_head_config(cfg)
+    add_hrnet_config(cfg)
+    add_bootstrap_config(cfg)
+    add_dataset_category_config(cfg)
+    add_evaluation_config(cfg)

densepose/converters/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .hflip import HFlipConverter
+from .to_mask import ToMaskConverter
+from .to_chart_result import ToChartResultConverter, ToChartResultConverterWithConfidences
+from .segm_to_mask import (
+    predictor_output_with_fine_and_coarse_segm_to_mask,
+    predictor_output_with_coarse_segm_to_mask,
+    resample_fine_and_coarse_segm_to_bbox,
+)
+from .chart_output_to_chart_result import (
+    densepose_chart_predictor_output_to_result,
+    densepose_chart_predictor_output_to_result_with_confidences,
+)
+from .chart_output_hflip import densepose_chart_predictor_output_hflip

densepose/converters/base.py

@@ -0,0 +1,95 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Tuple, Type
+import torch
+
+
+class BaseConverter:
+    """
+    Converter base class to be reused by various converters.
+    Converter allows one to convert data from various source types to a particular
+    destination type. Each source type needs to register its converter. The
+    registration for each source type is valid for all descendants of that type.
+    """
+
+    @classmethod
+    def register(cls, from_type: Type, converter: Any = None):
+        """
+        Registers a converter for the specified type.
+        Can be used as a decorator (if converter is None), or called as a method.
+
+        Args:
+            from_type (type): type to register the converter for;
+                all instances of this type will use the same converter
+            converter (callable): converter to be registered for the given
+                type; if None, this method is assumed to be a decorator for the converter
+        """
+
+        if converter is not None:
+            cls._do_register(from_type, converter)
+
+        def wrapper(converter: Any) -> Any:
+            cls._do_register(from_type, converter)
+            return converter
+
+        return wrapper
+
+    @classmethod
+    def _do_register(cls, from_type: Type, converter: Any):
+        cls.registry[from_type] = converter  # pyre-ignore[16]
+
+    @classmethod
+    def _lookup_converter(cls, from_type: Type) -> Any:
+        """
+        Perform recursive lookup for the given type
+        to find registered converter. If a converter was found for some base
+        class, it gets registered for this class to save on further lookups.
+
+        Args:
+            from_type: type for which to find a converter
+        Return:
+            callable or None - registered converter or None
+                if no suitable entry was found in the registry
+        """
+        if from_type in cls.registry:  # pyre-ignore[16]
+            return cls.registry[from_type]
+        for base in from_type.__bases__:
+            converter = cls._lookup_converter(base)
+            if converter is not None:
+                cls._do_register(from_type, converter)
+                return converter
+        return None
+
+    @classmethod
+    def convert(cls, instance: Any, *args, **kwargs):
+        """
+        Convert an instance to the destination type using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            instance: source instance to convert to the destination type
+        Return:
+            An instance of the destination type obtained from the source instance
+            Raises KeyError, if no suitable converter found
+        """
+        instance_type = type(instance)
+        converter = cls._lookup_converter(instance_type)
+        if converter is None:
+            if cls.dst_type is None:  # pyre-ignore[16]
+                output_type_str = "itself"
+            else:
+                output_type_str = cls.dst_type
+            raise KeyError(f"Could not find converter from {instance_type} to {output_type_str}")
+        return converter(instance, *args, **kwargs)
+
+
+IntTupleBox = Tuple[int, int, int, int]
+
+
+def make_int_box(box: torch.Tensor) -> IntTupleBox:
+    int_box = [0, 0, 0, 0]
+    int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
+    return int_box[0], int_box[1], int_box[2], int_box[3]

densepose/converters/builtin.py

@@ -0,0 +1,33 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from ..structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
+from . import (
+    HFlipConverter,
+    ToChartResultConverter,
+    ToChartResultConverterWithConfidences,
+    ToMaskConverter,
+    densepose_chart_predictor_output_hflip,
+    densepose_chart_predictor_output_to_result,
+    densepose_chart_predictor_output_to_result_with_confidences,
+    predictor_output_with_coarse_segm_to_mask,
+    predictor_output_with_fine_and_coarse_segm_to_mask,
+)
+
+ToMaskConverter.register(
+    DensePoseChartPredictorOutput, predictor_output_with_fine_and_coarse_segm_to_mask
+)
+ToMaskConverter.register(
+    DensePoseEmbeddingPredictorOutput, predictor_output_with_coarse_segm_to_mask
+)
+
+ToChartResultConverter.register(
+    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result
+)
+
+ToChartResultConverterWithConfidences.register(
+    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result_with_confidences
+)
+
+HFlipConverter.register(DensePoseChartPredictorOutput, densepose_chart_predictor_output_hflip)

densepose/converters/chart_output_hflip.py

@@ -0,0 +1,73 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+from dataclasses import fields
+import torch
+
+from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
+
+
+def densepose_chart_predictor_output_hflip(
+    densepose_predictor_output: DensePoseChartPredictorOutput,
+    transform_data: DensePoseTransformData,
+) -> DensePoseChartPredictorOutput:
+    """
+    Change  to take into account a Horizontal flip.
+    """
+    if len(densepose_predictor_output) > 0:
+
+        PredictorOutput = type(densepose_predictor_output)
+        output_dict = {}
+
+        for field in fields(densepose_predictor_output):
+            field_value = getattr(densepose_predictor_output, field.name)
+            # flip tensors
+            if isinstance(field_value, torch.Tensor):
+                setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
+
+        densepose_predictor_output = _flip_iuv_semantics_tensor(
+            densepose_predictor_output, transform_data
+        )
+        densepose_predictor_output = _flip_segm_semantics_tensor(
+            densepose_predictor_output, transform_data
+        )
+
+        for field in fields(densepose_predictor_output):
+            output_dict[field.name] = getattr(densepose_predictor_output, field.name)
+
+        return PredictorOutput(**output_dict)
+    else:
+        return densepose_predictor_output
+
+
+def _flip_iuv_semantics_tensor(
+    densepose_predictor_output: DensePoseChartPredictorOutput,
+    dp_transform_data: DensePoseTransformData,
+) -> DensePoseChartPredictorOutput:
+    point_label_symmetries = dp_transform_data.point_label_symmetries
+    uv_symmetries = dp_transform_data.uv_symmetries
+
+    N, C, H, W = densepose_predictor_output.u.shape
+    u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
+    v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
+    Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
+        None, :, None, None
+    ].expand(N, C - 1, H, W)
+    densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
+    densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
+
+    for el in ["fine_segm", "u", "v"]:
+        densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
+            :, point_label_symmetries, :, :
+        ]
+    return densepose_predictor_output
+
+
+def _flip_segm_semantics_tensor(
+    densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
+):
+    if densepose_predictor_output.coarse_segm.shape[1] > 2:
+        densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
+            :, dp_transform_data.mask_label_symmetries, :, :
+        ]
+    return densepose_predictor_output

densepose/converters/chart_output_to_chart_result.py

@@ -0,0 +1,190 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Dict
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures.boxes import Boxes, BoxMode
+
+from ..structures import (
+    DensePoseChartPredictorOutput,
+    DensePoseChartResult,
+    DensePoseChartResultWithConfidences,
+)
+from . import resample_fine_and_coarse_segm_to_bbox
+from .base import IntTupleBox, make_int_box
+
+
+def resample_uv_tensors_to_bbox(
+    u: torch.Tensor,
+    v: torch.Tensor,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        u (tensor [1, C, H, W] of float): U coordinates
+        v (tensor [1, C, H, W] of float): V coordinates
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
+    v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
+    uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
+    for part_id in range(1, u_bbox.size(1)):
+        uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
+        uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
+    return uv
+
+
+def resample_uv_to_bbox(
+    predictor_output: DensePoseChartPredictorOutput,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be resampled
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    return resample_uv_tensors_to_bbox(
+        predictor_output.u,
+        predictor_output.v,
+        labels,
+        box_xywh_abs,
+    )
+
+
+def densepose_chart_predictor_output_to_result(
+    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
+) -> DensePoseChartResult:
+    """
+    Convert densepose chart predictor outputs to results
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be converted to results, must contain only 1 output
+        boxes (Boxes): bounding box that corresponds to the predictor output,
+            must contain only 1 bounding box
+    Return:
+       DensePose chart-based result (DensePoseChartResult)
+    """
+    assert len(predictor_output) == 1 and len(boxes) == 1, (
+        f"Predictor output to result conversion can operate only single outputs"
+        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
+    )
+
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
+    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
+    return DensePoseChartResult(labels=labels, uv=uv)
+
+
+def resample_confidences_to_bbox(
+    predictor_output: DensePoseChartPredictorOutput,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> Dict[str, torch.Tensor]:
+    """
+    Resamples confidences for the given bounding box
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be resampled
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled confidences - a dict of [H, W] tensors of float
+    """
+
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+
+    confidence_names = [
+        "sigma_1",
+        "sigma_2",
+        "kappa_u",
+        "kappa_v",
+        "fine_segm_confidence",
+        "coarse_segm_confidence",
+    ]
+    confidence_results = {key: None for key in confidence_names}
+    confidence_names = [
+        key for key in confidence_names if getattr(predictor_output, key) is not None
+    ]
+    confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device)
+
+    # assign data from channels that correspond to the labels
+    for key in confidence_names:
+        resampled_confidence = F.interpolate(
+            getattr(predictor_output, key),
+            (h, w),
+            mode="bilinear",
+            align_corners=False,
+        )
+        result = confidence_base.clone()
+        for part_id in range(1, predictor_output.u.size(1)):
+            if resampled_confidence.size(1) != predictor_output.u.size(1):
+                # confidence is not part-based, don't try to fill it part by part
+                continue
+            result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id]
+
+        if resampled_confidence.size(1) != predictor_output.u.size(1):
+            # confidence is not part-based, fill the data with the first channel
+            # (targeted for segmentation confidences that have only 1 channel)
+            result = resampled_confidence[0, 0]
+
+        confidence_results[key] = result
+
+    return confidence_results  # pyre-ignore[7]
+
+
+def densepose_chart_predictor_output_to_result_with_confidences(
+    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
+) -> DensePoseChartResultWithConfidences:
+    """
+    Convert densepose chart predictor outputs to results
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output with confidences to be converted to results, must contain only 1 output
+        boxes (Boxes): bounding box that corresponds to the predictor output,
+            must contain only 1 bounding box
+    Return:
+       DensePose chart-based result with confidences (DensePoseChartResultWithConfidences)
+    """
+    assert len(predictor_output) == 1 and len(boxes) == 1, (
+        f"Predictor output to result conversion can operate only single outputs"
+        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
+    )
+
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
+    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
+    confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh)
+    return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences)

densepose/converters/hflip.py

@@ -0,0 +1,36 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any
+
+from .base import BaseConverter
+
+
+class HFlipConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = None
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(cls, predictor_outputs: Any, transform_data: Any, *args, **kwargs):
+        """
+        Performs an horizontal flip on DensePose predictor outputs.
+        Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            predictor_outputs: DensePose predictor output to be converted to BitMasks
+            transform_data: Anything useful for the flip
+        Return:
+            An instance of the same type as predictor_outputs
+        """
+        return super(HFlipConverter, cls).convert(
+            predictor_outputs, transform_data, *args, **kwargs
+        )

densepose/converters/segm_to_mask.py

@@ -0,0 +1,152 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures import BitMasks, Boxes, BoxMode
+
+from .base import IntTupleBox, make_int_box
+from .to_mask import ImageSizeType
+
+
+def resample_coarse_segm_tensor_to_bbox(coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox):
+    """
+    Resample coarse segmentation tensor to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    labels = F.interpolate(coarse_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+    return labels
+
+
+def resample_fine_and_coarse_segm_tensors_to_bbox(
+    fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
+):
+    """
+    Resample fine and coarse segmentation tensors to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        fine_segm: float tensor of shape [1, C, Hout, Wout]
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    # coarse segmentation
+    coarse_segm_bbox = F.interpolate(
+        coarse_segm,
+        (h, w),
+        mode="bilinear",
+        align_corners=False,
+    ).argmax(dim=1)
+    # combined coarse and fine segmentation
+    labels = (
+        F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+        * (coarse_segm_bbox > 0).long()
+    )
+    return labels
+
+
+def resample_fine_and_coarse_segm_to_bbox(predictor_output: Any, box_xywh_abs: IntTupleBox):
+    """
+    Resample fine and coarse segmentation outputs from a predictor to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        predictor_output: DensePose predictor output that contains segmentation
+            results to be resampled
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    return resample_fine_and_coarse_segm_tensors_to_bbox(
+        predictor_output.fine_segm,
+        predictor_output.coarse_segm,
+        box_xywh_abs,
+    )
+
+
+def predictor_output_with_coarse_segm_to_mask(
+    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
+) -> BitMasks:
+    """
+    Convert predictor output with coarse and fine segmentation to a mask.
+    Assumes that predictor output has the following attributes:
+     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
+         unnormalized scores for N instances; D is the number of coarse
+         segmentation labels, H and W is the resolution of the estimate
+
+    Args:
+        predictor_output: DensePose predictor output to be converted to mask
+        boxes (Boxes): bounding boxes that correspond to the DensePose
+            predictor outputs
+        image_size_hw (tuple [int, int]): image height Himg and width Wimg
+    Return:
+        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
+        a mask of the size of the image for each instance
+    """
+    H, W = image_size_hw
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    N = len(boxes_xywh_abs)
+    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
+    for i in range(len(boxes_xywh_abs)):
+        box_xywh = make_int_box(boxes_xywh_abs[i])
+        box_mask = resample_coarse_segm_tensor_to_bbox(predictor_output[i].coarse_segm, box_xywh)
+        x, y, w, h = box_xywh
+        masks[i, y : y + h, x : x + w] = box_mask
+
+    return BitMasks(masks)
+
+
+def predictor_output_with_fine_and_coarse_segm_to_mask(
+    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
+) -> BitMasks:
+    """
+    Convert predictor output with coarse and fine segmentation to a mask.
+    Assumes that predictor output has the following attributes:
+     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
+         unnormalized scores for N instances; D is the number of coarse
+         segmentation labels, H and W is the resolution of the estimate
+     - fine_segm (tensor of size [N, C, H, W]): fine segmentation
+         unnormalized scores for N instances; C is the number of fine
+         segmentation labels, H and W is the resolution of the estimate
+
+    Args:
+        predictor_output: DensePose predictor output to be converted to mask
+        boxes (Boxes): bounding boxes that correspond to the DensePose
+            predictor outputs
+        image_size_hw (tuple [int, int]): image height Himg and width Wimg
+    Return:
+        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
+        a mask of the size of the image for each instance
+    """
+    H, W = image_size_hw
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    N = len(boxes_xywh_abs)
+    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
+    for i in range(len(boxes_xywh_abs)):
+        box_xywh = make_int_box(boxes_xywh_abs[i])
+        labels_i = resample_fine_and_coarse_segm_to_bbox(predictor_output[i], box_xywh)
+        x, y, w, h = box_xywh
+        masks[i, y : y + h, x : x + w] = labels_i > 0
+    return BitMasks(masks)

densepose/converters/to_chart_result.py

@@ -0,0 +1,72 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any
+
+from detectron2.structures import Boxes
+
+from ..structures import DensePoseChartResult, DensePoseChartResultWithConfidences
+from .base import BaseConverter
+
+
+class ToChartResultConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = DensePoseChartResult
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs) -> DensePoseChartResult:
+        """
+        Convert DensePose predictor outputs to DensePoseResult using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output to be
+                converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+        Return:
+            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
+        """
+        return super(ToChartResultConverter, cls).convert(predictor_outputs, boxes, *args, **kwargs)
+
+
+class ToChartResultConverterWithConfidences(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = DensePoseChartResultWithConfidences
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(
+        cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs
+    ) -> DensePoseChartResultWithConfidences:
+        """
+        Convert DensePose predictor outputs to DensePoseResult with confidences
+        using some registered converter. Does recursive lookup for base classes,
+        so there's no need for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output with confidences
+                to be converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+        Return:
+            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
+        """
+        return super(ToChartResultConverterWithConfidences, cls).convert(
+            predictor_outputs, boxes, *args, **kwargs
+        )

densepose/converters/to_mask.py

@@ -0,0 +1,51 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Tuple
+
+from detectron2.structures import BitMasks, Boxes
+
+from .base import BaseConverter
+
+ImageSizeType = Tuple[int, int]
+
+
+class ToMaskConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to masks
+    in bit mask format (see `BitMasks`). Each DensePose predictor output type
+    has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = BitMasks
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(
+        cls,
+        densepose_predictor_outputs: Any,
+        boxes: Boxes,
+        image_size_hw: ImageSizeType,
+        *args,
+        **kwargs
+    ) -> BitMasks:
+        """
+        Convert DensePose predictor outputs to BitMasks using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output to be
+                converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+            image_size_hw (tuple [int, int]): image height and width
+        Return:
+            An instance of `BitMasks`. If no suitable converter was found, raises KeyError
+        """
+        return super(ToMaskConverter, cls).convert(
+            densepose_predictor_outputs, boxes, image_size_hw, *args, **kwargs
+        )

densepose/data/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .meshes import builtin
+from .build import (
+    build_detection_test_loader,
+    build_detection_train_loader,
+    build_combined_loader,
+    build_frame_selector,
+    build_inference_based_loaders,
+    has_inference_based_loaders,
+    BootstrapDatasetFactoryCatalog,
+)
+from .combined_loader import CombinedDataLoader
+from .dataset_mapper import DatasetMapper
+from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
+from .image_list_dataset import ImageListDataset
+from .utils import is_relative_local_path, maybe_prepend_base_path
+
+# ensure the builtin datasets are registered
+from . import datasets
+
+# ensure the bootstrap datasets builders are registered
+from . import build
+
+__all__ = [k for k in globals().keys() if not k.startswith("_")]

densepose/data/build.py

@@ -0,0 +1,738 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import itertools
+import logging
+import numpy as np
+from collections import UserDict, defaultdict
+from dataclasses import dataclass
+from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Sequence, Tuple
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.config import CfgNode
+from detectron2.data.build import build_detection_test_loader as d2_build_detection_test_loader
+from detectron2.data.build import build_detection_train_loader as d2_build_detection_train_loader
+from detectron2.data.build import (
+    load_proposals_into_dataset,
+    print_instances_class_histogram,
+    trivial_batch_collator,
+    worker_init_reset_seed,
+)
+from detectron2.data.catalog import DatasetCatalog, Metadata, MetadataCatalog
+from detectron2.data.samplers import TrainingSampler
+from detectron2.utils.comm import get_world_size
+
+from densepose.config import get_bootstrap_dataset_config
+from densepose.modeling import build_densepose_embedder
+
+from .combined_loader import CombinedDataLoader, Loader
+from .dataset_mapper import DatasetMapper
+from .datasets.coco import DENSEPOSE_CSE_KEYS_WITHOUT_MASK, DENSEPOSE_IUV_KEYS_WITHOUT_MASK
+from .datasets.dataset_type import DatasetType
+from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
+from .samplers import (
+    DensePoseConfidenceBasedSampler,
+    DensePoseCSEConfidenceBasedSampler,
+    DensePoseCSEUniformSampler,
+    DensePoseUniformSampler,
+    MaskFromDensePoseSampler,
+    PredictionToGroundTruthSampler,
+)
+from .transform import ImageResizeTransform
+from .utils import get_category_to_class_mapping, get_class_to_mesh_name_mapping
+from .video import (
+    FirstKFramesSelector,
+    FrameSelectionStrategy,
+    LastKFramesSelector,
+    RandomKFramesSelector,
+    VideoKeyframeDataset,
+    video_list_from_file,
+)
+
+__all__ = ["build_detection_train_loader", "build_detection_test_loader"]
+
+
+Instance = Dict[str, Any]
+InstancePredicate = Callable[[Instance], bool]
+
+
+def _compute_num_images_per_worker(cfg: CfgNode) -> int:
+    num_workers = get_world_size()
+    images_per_batch = cfg.SOLVER.IMS_PER_BATCH
+    assert (
+        images_per_batch % num_workers == 0
+    ), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format(
+        images_per_batch, num_workers
+    )
+    assert (
+        images_per_batch >= num_workers
+    ), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format(
+        images_per_batch, num_workers
+    )
+    images_per_worker = images_per_batch // num_workers
+    return images_per_worker
+
+
+def _map_category_id_to_contiguous_id(dataset_name: str, dataset_dicts: Iterable[Instance]) -> None:
+    meta = MetadataCatalog.get(dataset_name)
+    for dataset_dict in dataset_dicts:
+        for ann in dataset_dict["annotations"]:
+            ann["category_id"] = meta.thing_dataset_id_to_contiguous_id[ann["category_id"]]
+
+
+@dataclass
+class _DatasetCategory:
+    """
+    Class representing category data in a dataset:
+     - id: category ID, as specified in the dataset annotations file
+     - name: category name, as specified in the dataset annotations file
+     - mapped_id: category ID after applying category maps (DATASETS.CATEGORY_MAPS config option)
+     - mapped_name: category name after applying category maps
+     - dataset_name: dataset in which the category is defined
+
+    For example, when training models in a class-agnostic manner, one could take LVIS 1.0
+    dataset and map the animal categories to the same category as human data from COCO:
+     id = 225
+     name = "cat"
+     mapped_id = 1
+     mapped_name = "person"
+     dataset_name = "lvis_v1_animals_dp_train"
+    """
+
+    id: int
+    name: str
+    mapped_id: int
+    mapped_name: str
+    dataset_name: str
+
+
+_MergedCategoriesT = Dict[int, List[_DatasetCategory]]
+
+
+def _add_category_id_to_contiguous_id_maps_to_metadata(
+    merged_categories: _MergedCategoriesT,
+) -> None:
+    merged_categories_per_dataset = {}
+    for contiguous_cat_id, cat_id in enumerate(sorted(merged_categories.keys())):
+        for cat in merged_categories[cat_id]:
+            if cat.dataset_name not in merged_categories_per_dataset:
+                merged_categories_per_dataset[cat.dataset_name] = defaultdict(list)
+            merged_categories_per_dataset[cat.dataset_name][cat_id].append(
+                (
+                    contiguous_cat_id,
+                    cat,
+                )
+            )
+
+    logger = logging.getLogger(__name__)
+    for dataset_name, merged_categories in merged_categories_per_dataset.items():
+        meta = MetadataCatalog.get(dataset_name)
+        if not hasattr(meta, "thing_classes"):
+            meta.thing_classes = []
+            meta.thing_dataset_id_to_contiguous_id = {}
+            meta.thing_dataset_id_to_merged_id = {}
+        else:
+            meta.thing_classes.clear()
+            meta.thing_dataset_id_to_contiguous_id.clear()
+            meta.thing_dataset_id_to_merged_id.clear()
+        logger.info(f"Dataset {dataset_name}: category ID to contiguous ID mapping:")
+        for _cat_id, categories in sorted(merged_categories.items()):
+            added_to_thing_classes = False
+            for contiguous_cat_id, cat in categories:
+                if not added_to_thing_classes:
+                    meta.thing_classes.append(cat.mapped_name)
+                    added_to_thing_classes = True
+                meta.thing_dataset_id_to_contiguous_id[cat.id] = contiguous_cat_id
+                meta.thing_dataset_id_to_merged_id[cat.id] = cat.mapped_id
+                logger.info(f"{cat.id} ({cat.name}) -> {contiguous_cat_id}")
+
+
+def _maybe_create_general_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    def has_annotations(instance: Instance) -> bool:
+        return "annotations" in instance
+
+    def has_only_crowd_anotations(instance: Instance) -> bool:
+        for ann in instance["annotations"]:
+            if ann.get("is_crowd", 0) == 0:
+                return False
+        return True
+
+    def general_keep_instance_predicate(instance: Instance) -> bool:
+        return has_annotations(instance) and not has_only_crowd_anotations(instance)
+
+    if not cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS:
+        return None
+    return general_keep_instance_predicate
+
+
+def _maybe_create_keypoints_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+
+    min_num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+
+    def has_sufficient_num_keypoints(instance: Instance) -> bool:
+        num_kpts = sum(
+            (np.array(ann["keypoints"][2::3]) > 0).sum()
+            for ann in instance["annotations"]
+            if "keypoints" in ann
+        )
+        return num_kpts >= min_num_keypoints
+
+    if cfg.MODEL.KEYPOINT_ON and (min_num_keypoints > 0):
+        return has_sufficient_num_keypoints
+    return None
+
+
+def _maybe_create_mask_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    if not cfg.MODEL.MASK_ON:
+        return None
+
+    def has_mask_annotations(instance: Instance) -> bool:
+        return any("segmentation" in ann for ann in instance["annotations"])
+
+    return has_mask_annotations
+
+
+def _maybe_create_densepose_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    if not cfg.MODEL.DENSEPOSE_ON:
+        return None
+
+    use_masks = cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
+
+    def has_densepose_annotations(instance: Instance) -> bool:
+        for ann in instance["annotations"]:
+            if all(key in ann for key in DENSEPOSE_IUV_KEYS_WITHOUT_MASK) or all(
+                key in ann for key in DENSEPOSE_CSE_KEYS_WITHOUT_MASK
+            ):
+                return True
+            if use_masks and "segmentation" in ann:
+                return True
+        return False
+
+    return has_densepose_annotations
+
+
+def _maybe_create_specific_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    specific_predicate_creators = [
+        _maybe_create_keypoints_keep_instance_predicate,
+        _maybe_create_mask_keep_instance_predicate,
+        _maybe_create_densepose_keep_instance_predicate,
+    ]
+    predicates = [creator(cfg) for creator in specific_predicate_creators]
+    predicates = [p for p in predicates if p is not None]
+    if not predicates:
+        return None
+
+    def combined_predicate(instance: Instance) -> bool:
+        return any(p(instance) for p in predicates)
+
+    return combined_predicate
+
+
+def _get_train_keep_instance_predicate(cfg: CfgNode):
+    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
+    combined_specific_keep_predicate = _maybe_create_specific_keep_instance_predicate(cfg)
+
+    def combined_general_specific_keep_predicate(instance: Instance) -> bool:
+        return general_keep_predicate(instance) and combined_specific_keep_predicate(instance)
+
+    if (general_keep_predicate is None) and (combined_specific_keep_predicate is None):
+        return None
+    if general_keep_predicate is None:
+        return combined_specific_keep_predicate
+    if combined_specific_keep_predicate is None:
+        return general_keep_predicate
+    return combined_general_specific_keep_predicate
+
+
+def _get_test_keep_instance_predicate(cfg: CfgNode):
+    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
+    return general_keep_predicate
+
+
+def _maybe_filter_and_map_categories(
+    dataset_name: str, dataset_dicts: List[Instance]
+) -> List[Instance]:
+    meta = MetadataCatalog.get(dataset_name)
+    category_id_map = meta.thing_dataset_id_to_contiguous_id
+    filtered_dataset_dicts = []
+    for dataset_dict in dataset_dicts:
+        anns = []
+        for ann in dataset_dict["annotations"]:
+            cat_id = ann["category_id"]
+            if cat_id not in category_id_map:
+                continue
+            ann["category_id"] = category_id_map[cat_id]
+            anns.append(ann)
+        dataset_dict["annotations"] = anns
+        filtered_dataset_dicts.append(dataset_dict)
+    return filtered_dataset_dicts
+
+
+def _add_category_whitelists_to_metadata(cfg: CfgNode) -> None:
+    for dataset_name, whitelisted_cat_ids in cfg.DATASETS.WHITELISTED_CATEGORIES.items():
+        meta = MetadataCatalog.get(dataset_name)
+        meta.whitelisted_categories = whitelisted_cat_ids
+        logger = logging.getLogger(__name__)
+        logger.info(
+            "Whitelisted categories for dataset {}: {}".format(
+                dataset_name, meta.whitelisted_categories
+            )
+        )
+
+
+def _add_category_maps_to_metadata(cfg: CfgNode) -> None:
+    for dataset_name, category_map in cfg.DATASETS.CATEGORY_MAPS.items():
+        category_map = {
+            int(cat_id_src): int(cat_id_dst) for cat_id_src, cat_id_dst in category_map.items()
+        }
+        meta = MetadataCatalog.get(dataset_name)
+        meta.category_map = category_map
+        logger = logging.getLogger(__name__)
+        logger.info("Category maps for dataset {}: {}".format(dataset_name, meta.category_map))
+
+
+def _add_category_info_to_bootstrapping_metadata(dataset_name: str, dataset_cfg: CfgNode) -> None:
+    meta = MetadataCatalog.get(dataset_name)
+    meta.category_to_class_mapping = get_category_to_class_mapping(dataset_cfg)
+    meta.categories = dataset_cfg.CATEGORIES
+    meta.max_count_per_category = dataset_cfg.MAX_COUNT_PER_CATEGORY
+    logger = logging.getLogger(__name__)
+    logger.info(
+        "Category to class mapping for dataset {}: {}".format(
+            dataset_name, meta.category_to_class_mapping
+        )
+    )
+
+
+def _maybe_add_class_to_mesh_name_map_to_metadata(dataset_names: List[str], cfg: CfgNode) -> None:
+    for dataset_name in dataset_names:
+        meta = MetadataCatalog.get(dataset_name)
+        if not hasattr(meta, "class_to_mesh_name"):
+            meta.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
+
+
+def _merge_categories(dataset_names: Collection[str]) -> _MergedCategoriesT:
+    merged_categories = defaultdict(list)
+    category_names = {}
+    for dataset_name in dataset_names:
+        meta = MetadataCatalog.get(dataset_name)
+        whitelisted_categories = meta.get("whitelisted_categories")
+        category_map = meta.get("category_map", {})
+        cat_ids = (
+            whitelisted_categories if whitelisted_categories is not None else meta.categories.keys()
+        )
+        for cat_id in cat_ids:
+            cat_name = meta.categories[cat_id]
+            cat_id_mapped = category_map.get(cat_id, cat_id)
+            if cat_id_mapped == cat_id or cat_id_mapped in cat_ids:
+                category_names[cat_id] = cat_name
+            else:
+                category_names[cat_id] = str(cat_id_mapped)
+            # assign temporary mapped category name, this name can be changed
+            # during the second pass, since mapped ID can correspond to a category
+            # from a different dataset
+            cat_name_mapped = meta.categories[cat_id_mapped]
+            merged_categories[cat_id_mapped].append(
+                _DatasetCategory(
+                    id=cat_id,
+                    name=cat_name,
+                    mapped_id=cat_id_mapped,
+                    mapped_name=cat_name_mapped,
+                    dataset_name=dataset_name,
+                )
+            )
+    # second pass to assign proper mapped category names
+    for cat_id, categories in merged_categories.items():
+        for cat in categories:
+            if cat_id in category_names and cat.mapped_name != category_names[cat_id]:
+                cat.mapped_name = category_names[cat_id]
+
+    return merged_categories
+
+
+def _warn_if_merged_different_categories(merged_categories: _MergedCategoriesT) -> None:
+    logger = logging.getLogger(__name__)
+    for cat_id in merged_categories:
+        merged_categories_i = merged_categories[cat_id]
+        first_cat_name = merged_categories_i[0].name
+        if len(merged_categories_i) > 1 and not all(
+            cat.name == first_cat_name for cat in merged_categories_i[1:]
+        ):
+            cat_summary_str = ", ".join(
+                [f"{cat.id} ({cat.name}) from {cat.dataset_name}" for cat in merged_categories_i]
+            )
+            logger.warning(
+                f"Merged category {cat_id} corresponds to the following categories: "
+                f"{cat_summary_str}"
+            )
+
+
+def combine_detection_dataset_dicts(
+    dataset_names: Collection[str],
+    keep_instance_predicate: Optional[InstancePredicate] = None,
+    proposal_files: Optional[Collection[str]] = None,
+) -> List[Instance]:
+    """
+    Load and prepare dataset dicts for training / testing
+
+    Args:
+        dataset_names (Collection[str]): a list of dataset names
+        keep_instance_predicate (Callable: Dict[str, Any] -> bool): predicate
+            applied to instance dicts which defines whether to keep the instance
+        proposal_files (Collection[str]): if given, a list of object proposal files
+            that match each dataset in `dataset_names`.
+    """
+    assert len(dataset_names)
+    if proposal_files is None:
+        proposal_files = [None] * len(dataset_names)
+    assert len(dataset_names) == len(proposal_files)
+    # load datasets and metadata
+    dataset_name_to_dicts = {}
+    for dataset_name in dataset_names:
+        dataset_name_to_dicts[dataset_name] = DatasetCatalog.get(dataset_name)
+        assert len(dataset_name_to_dicts), f"Dataset '{dataset_name}' is empty!"
+    # merge categories, requires category metadata to be loaded
+    # cat_id -> [(orig_cat_id, cat_name, dataset_name)]
+    merged_categories = _merge_categories(dataset_names)
+    _warn_if_merged_different_categories(merged_categories)
+    merged_category_names = [
+        merged_categories[cat_id][0].mapped_name for cat_id in sorted(merged_categories)
+    ]
+    # map to contiguous category IDs
+    _add_category_id_to_contiguous_id_maps_to_metadata(merged_categories)
+    # load annotations and dataset metadata
+    for dataset_name, proposal_file in zip(dataset_names, proposal_files):
+        dataset_dicts = dataset_name_to_dicts[dataset_name]
+        assert len(dataset_dicts), f"Dataset '{dataset_name}' is empty!"
+        if proposal_file is not None:
+            dataset_dicts = load_proposals_into_dataset(dataset_dicts, proposal_file)
+        dataset_dicts = _maybe_filter_and_map_categories(dataset_name, dataset_dicts)
+        print_instances_class_histogram(dataset_dicts, merged_category_names)
+        dataset_name_to_dicts[dataset_name] = dataset_dicts
+
+    if keep_instance_predicate is not None:
+        all_datasets_dicts_plain = [
+            d
+            for d in itertools.chain.from_iterable(dataset_name_to_dicts.values())
+            if keep_instance_predicate(d)
+        ]
+    else:
+        all_datasets_dicts_plain = list(
+            itertools.chain.from_iterable(dataset_name_to_dicts.values())
+        )
+    return all_datasets_dicts_plain
+
+
+def build_detection_train_loader(cfg: CfgNode, mapper=None):
+    """
+    A data loader is created in a way similar to that of Detectron2.
+    The main differences are:
+     - it allows to combine datasets with different but compatible object category sets
+
+    The data loader is created by the following steps:
+    1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
+    2. Start workers to work on the dicts. Each worker will:
+        * Map each metadata dict into another format to be consumed by the model.
+        * Batch them by simply putting dicts into a list.
+    The batched ``list[mapped_dict]`` is what this dataloader will return.
+
+    Args:
+        cfg (CfgNode): the config
+        mapper (callable): a callable which takes a sample (dict) from dataset and
+            returns the format to be consumed by the model.
+            By default it will be `DatasetMapper(cfg, True)`.
+
+    Returns:
+        an infinite iterator of training data
+    """
+
+    _add_category_whitelists_to_metadata(cfg)
+    _add_category_maps_to_metadata(cfg)
+    _maybe_add_class_to_mesh_name_map_to_metadata(cfg.DATASETS.TRAIN, cfg)
+    dataset_dicts = combine_detection_dataset_dicts(
+        cfg.DATASETS.TRAIN,
+        keep_instance_predicate=_get_train_keep_instance_predicate(cfg),
+        proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+    )
+    if mapper is None:
+        mapper = DatasetMapper(cfg, True)
+    return d2_build_detection_train_loader(cfg, dataset=dataset_dicts, mapper=mapper)
+
+
+def build_detection_test_loader(cfg, dataset_name, mapper=None):
+    """
+    Similar to `build_detection_train_loader`.
+    But this function uses the given `dataset_name` argument (instead of the names in cfg),
+    and uses batch size 1.
+
+    Args:
+        cfg: a detectron2 CfgNode
+        dataset_name (str): a name of the dataset that's available in the DatasetCatalog
+        mapper (callable): a callable which takes a sample (dict) from dataset
+            and returns the format to be consumed by the model.
+            By default it will be `DatasetMapper(cfg, False)`.
+
+    Returns:
+        DataLoader: a torch DataLoader, that loads the given detection
+            dataset, with test-time transformation and batching.
+    """
+    _add_category_whitelists_to_metadata(cfg)
+    _add_category_maps_to_metadata(cfg)
+    _maybe_add_class_to_mesh_name_map_to_metadata([dataset_name], cfg)
+    dataset_dicts = combine_detection_dataset_dicts(
+        [dataset_name],
+        keep_instance_predicate=_get_test_keep_instance_predicate(cfg),
+        proposal_files=(
+            [cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]]
+            if cfg.MODEL.LOAD_PROPOSALS
+            else None
+        ),
+    )
+    sampler = None
+    if not cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE:
+        sampler = torch.utils.data.SequentialSampler(dataset_dicts)
+    if mapper is None:
+        mapper = DatasetMapper(cfg, False)
+    return d2_build_detection_test_loader(
+        dataset_dicts, mapper=mapper, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler
+    )
+
+
+def build_frame_selector(cfg: CfgNode):
+    strategy = FrameSelectionStrategy(cfg.STRATEGY)
+    if strategy == FrameSelectionStrategy.RANDOM_K:
+        frame_selector = RandomKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.FIRST_K:
+        frame_selector = FirstKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.LAST_K:
+        frame_selector = LastKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.ALL:
+        frame_selector = None
+    # pyre-fixme[61]: `frame_selector` may not be initialized here.
+    return frame_selector
+
+
+def build_transform(cfg: CfgNode, data_type: str):
+    if cfg.TYPE == "resize":
+        if data_type == "image":
+            return ImageResizeTransform(cfg.MIN_SIZE, cfg.MAX_SIZE)
+    raise ValueError(f"Unknown transform {cfg.TYPE} for data type {data_type}")
+
+
+def build_combined_loader(cfg: CfgNode, loaders: Collection[Loader], ratios: Sequence[float]):
+    images_per_worker = _compute_num_images_per_worker(cfg)
+    return CombinedDataLoader(loaders, images_per_worker, ratios)
+
+
+def build_bootstrap_dataset(dataset_name: str, cfg: CfgNode) -> Sequence[torch.Tensor]:
+    """
+    Build dataset that provides data to bootstrap on
+
+    Args:
+        dataset_name (str): Name of the dataset, needs to have associated metadata
+            to load the data
+        cfg (CfgNode): bootstrapping config
+    Returns:
+        Sequence[Tensor] - dataset that provides image batches, Tensors of size
+            [N, C, H, W] of type float32
+    """
+    logger = logging.getLogger(__name__)
+    _add_category_info_to_bootstrapping_metadata(dataset_name, cfg)
+    meta = MetadataCatalog.get(dataset_name)
+    factory = BootstrapDatasetFactoryCatalog.get(meta.dataset_type)
+    dataset = None
+    if factory is not None:
+        dataset = factory(meta, cfg)
+    if dataset is None:
+        logger.warning(f"Failed to create dataset {dataset_name} of type {meta.dataset_type}")
+    return dataset
+
+
+def build_data_sampler(cfg: CfgNode, sampler_cfg: CfgNode, embedder: Optional[torch.nn.Module]):
+    if sampler_cfg.TYPE == "densepose_uniform":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseUniformSampler(count_per_class=sampler_cfg.COUNT_PER_CLASS),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_UV_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="sigma_2",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_fine_segm_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="fine_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_coarse_segm_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="coarse_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_cse_uniform":
+        assert embedder is not None
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseCSEUniformSampler(
+                cfg=cfg,
+                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
+                embedder=embedder,
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_cse_coarse_segm_confidence":
+        assert embedder is not None
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseCSEConfidenceBasedSampler(
+                cfg=cfg,
+                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
+                embedder=embedder,
+                confidence_channel="coarse_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+
+    raise ValueError(f"Unknown data sampler type {sampler_cfg.TYPE}")
+
+
+def build_data_filter(cfg: CfgNode):
+    if cfg.TYPE == "detection_score":
+        min_score = cfg.MIN_VALUE
+        return ScoreBasedFilter(min_score=min_score)
+    raise ValueError(f"Unknown data filter type {cfg.TYPE}")
+
+
+def build_inference_based_loader(
+    cfg: CfgNode,
+    dataset_cfg: CfgNode,
+    model: torch.nn.Module,
+    embedder: Optional[torch.nn.Module] = None,
+) -> InferenceBasedLoader:
+    """
+    Constructs data loader based on inference results of a model.
+    """
+    dataset = build_bootstrap_dataset(dataset_cfg.DATASET, dataset_cfg.IMAGE_LOADER)
+    meta = MetadataCatalog.get(dataset_cfg.DATASET)
+    training_sampler = TrainingSampler(len(dataset))
+    data_loader = torch.utils.data.DataLoader(
+        dataset,  # pyre-ignore[6]
+        batch_size=dataset_cfg.IMAGE_LOADER.BATCH_SIZE,
+        sampler=training_sampler,
+        num_workers=dataset_cfg.IMAGE_LOADER.NUM_WORKERS,
+        collate_fn=trivial_batch_collator,
+        worker_init_fn=worker_init_reset_seed,
+    )
+    return InferenceBasedLoader(
+        model,
+        data_loader=data_loader,
+        data_sampler=build_data_sampler(cfg, dataset_cfg.DATA_SAMPLER, embedder),
+        data_filter=build_data_filter(dataset_cfg.FILTER),
+        shuffle=True,
+        batch_size=dataset_cfg.INFERENCE.OUTPUT_BATCH_SIZE,
+        inference_batch_size=dataset_cfg.INFERENCE.INPUT_BATCH_SIZE,
+        category_to_class_mapping=meta.category_to_class_mapping,
+    )
+
+
+def has_inference_based_loaders(cfg: CfgNode) -> bool:
+    """
+    Returns True, if at least one inferense-based loader must
+    be instantiated for training
+    """
+    return len(cfg.BOOTSTRAP_DATASETS) > 0
+
+
+def build_inference_based_loaders(
+    cfg: CfgNode, model: torch.nn.Module
+) -> Tuple[List[InferenceBasedLoader], List[float]]:
+    loaders = []
+    ratios = []
+    embedder = build_densepose_embedder(cfg).to(device=model.device)  # pyre-ignore[16]
+    for dataset_spec in cfg.BOOTSTRAP_DATASETS:
+        dataset_cfg = get_bootstrap_dataset_config().clone()
+        dataset_cfg.merge_from_other_cfg(CfgNode(dataset_spec))
+        loader = build_inference_based_loader(cfg, dataset_cfg, model, embedder)
+        loaders.append(loader)
+        ratios.append(dataset_cfg.RATIO)
+    return loaders, ratios
+
+
+def build_video_list_dataset(meta: Metadata, cfg: CfgNode):
+    video_list_fpath = meta.video_list_fpath
+    video_base_path = meta.video_base_path
+    category = meta.category
+    if cfg.TYPE == "video_keyframe":
+        frame_selector = build_frame_selector(cfg.SELECT)
+        transform = build_transform(cfg.TRANSFORM, data_type="image")
+        video_list = video_list_from_file(video_list_fpath, video_base_path)
+        keyframe_helper_fpath = getattr(cfg, "KEYFRAME_HELPER", None)
+        return VideoKeyframeDataset(
+            video_list, category, frame_selector, transform, keyframe_helper_fpath
+        )
+
+
+class _BootstrapDatasetFactoryCatalog(UserDict):
+    """
+    A global dictionary that stores information about bootstrapped datasets creation functions
+    from metadata and config, for diverse DatasetType
+    """
+
+    def register(self, dataset_type: DatasetType, factory: Callable[[Metadata, CfgNode], Dataset]):
+        """
+        Args:
+            dataset_type (DatasetType): a DatasetType e.g. DatasetType.VIDEO_LIST
+            factory (Callable[Metadata, CfgNode]): a callable which takes Metadata and cfg
+            arguments and returns a dataset object.
+        """
+        assert dataset_type not in self, "Dataset '{}' is already registered!".format(dataset_type)
+        self[dataset_type] = factory
+
+
+BootstrapDatasetFactoryCatalog = _BootstrapDatasetFactoryCatalog()
+BootstrapDatasetFactoryCatalog.register(DatasetType.VIDEO_LIST, build_video_list_dataset)

densepose/data/combined_loader.py

@@ -0,0 +1,46 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from collections import deque
+from typing import Any, Collection, Deque, Iterable, Iterator, List, Sequence
+
+Loader = Iterable[Any]
+
+
+def _pooled_next(iterator: Iterator[Any], pool: Deque[Any]):
+    if not pool:
+        pool.extend(next(iterator))
+    return pool.popleft()
+
+
+class CombinedDataLoader:
+    """
+    Combines data loaders using the provided sampling ratios
+    """
+
+    BATCH_COUNT = 100
+
+    def __init__(self, loaders: Collection[Loader], batch_size: int, ratios: Sequence[float]):
+        self.loaders = loaders
+        self.batch_size = batch_size
+        self.ratios = ratios
+
+    def __iter__(self) -> Iterator[List[Any]]:
+        iters = [iter(loader) for loader in self.loaders]
+        indices = []
+        pool = [deque()] * len(iters)
+        # infinite iterator, as in D2
+        while True:
+            if not indices:
+                # just a buffer of indices, its size doesn't matter
+                # as long as it's a multiple of batch_size
+                k = self.batch_size * self.BATCH_COUNT
+                indices = random.choices(range(len(self.loaders)), self.ratios, k=k)
+            try:
+                batch = [_pooled_next(iters[i], pool[i]) for i in indices[: self.batch_size]]
+            except StopIteration:
+                break
+            indices = indices[self.batch_size :]
+            yield batch

densepose/data/dataset_mapper.py

@@ -0,0 +1,170 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import copy
+import logging
+from typing import Any, Dict, List, Tuple
+import torch
+
+from detectron2.data import MetadataCatalog
+from detectron2.data import detection_utils as utils
+from detectron2.data import transforms as T
+from detectron2.layers import ROIAlign
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from densepose.structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
+
+
+def build_augmentation(cfg, is_train):
+    logger = logging.getLogger(__name__)
+    result = utils.build_augmentation(cfg, is_train)
+    if is_train:
+        random_rotation = T.RandomRotation(
+            cfg.INPUT.ROTATION_ANGLES, expand=False, sample_style="choice"
+        )
+        result.append(random_rotation)
+        logger.info("DensePose-specific augmentation used in training: " + str(random_rotation))
+    return result
+
+
+class DatasetMapper:
+    """
+    A customized version of `detectron2.data.DatasetMapper`
+    """
+
+    def __init__(self, cfg, is_train=True):
+        self.augmentation = build_augmentation(cfg, is_train)
+
+        # fmt: off
+        self.img_format     = cfg.INPUT.FORMAT
+        self.mask_on        = (
+            cfg.MODEL.MASK_ON or (
+                cfg.MODEL.DENSEPOSE_ON
+                and cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS)
+        )
+        self.keypoint_on    = cfg.MODEL.KEYPOINT_ON
+        self.densepose_on   = cfg.MODEL.DENSEPOSE_ON
+        assert not cfg.MODEL.LOAD_PROPOSALS, "not supported yet"
+        # fmt: on
+        if self.keypoint_on and is_train:
+            # Flip only makes sense in training
+            self.keypoint_hflip_indices = utils.create_keypoint_hflip_indices(cfg.DATASETS.TRAIN)
+        else:
+            self.keypoint_hflip_indices = None
+
+        if self.densepose_on:
+            densepose_transform_srcs = [
+                MetadataCatalog.get(ds).densepose_transform_src
+                for ds in cfg.DATASETS.TRAIN + cfg.DATASETS.TEST
+            ]
+            assert len(densepose_transform_srcs) > 0
+            # TODO: check that DensePose transformation data is the same for
+            # all the datasets. Otherwise one would have to pass DB ID with
+            # each entry to select proper transformation data. For now, since
+            # all DensePose annotated data uses the same data semantics, we
+            # omit this check.
+            densepose_transform_data_fpath = PathManager.get_local_path(densepose_transform_srcs[0])
+            self.densepose_transform_data = DensePoseTransformData.load(
+                densepose_transform_data_fpath
+            )
+
+        self.is_train = is_train
+
+    def __call__(self, dataset_dict):
+        """
+        Args:
+            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
+
+        Returns:
+            dict: a format that builtin models in detectron2 accept
+        """
+        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
+        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
+        utils.check_image_size(dataset_dict, image)
+
+        image, transforms = T.apply_transform_gens(self.augmentation, image)
+        image_shape = image.shape[:2]  # h, w
+        dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
+
+        if not self.is_train:
+            dataset_dict.pop("annotations", None)
+            return dataset_dict
+
+        for anno in dataset_dict["annotations"]:
+            if not self.mask_on:
+                anno.pop("segmentation", None)
+            if not self.keypoint_on:
+                anno.pop("keypoints", None)
+
+        # USER: Implement additional transformations if you have other types of data
+        # USER: Don't call transpose_densepose if you don't need
+        annos = [
+            self._transform_densepose(
+                utils.transform_instance_annotations(
+                    obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices
+                ),
+                transforms,
+            )
+            for obj in dataset_dict.pop("annotations")
+            if obj.get("iscrowd", 0) == 0
+        ]
+
+        if self.mask_on:
+            self._add_densepose_masks_as_segmentation(annos, image_shape)
+
+        instances = utils.annotations_to_instances(annos, image_shape, mask_format="bitmask")
+        densepose_annotations = [obj.get("densepose") for obj in annos]
+        if densepose_annotations and not all(v is None for v in densepose_annotations):
+            instances.gt_densepose = DensePoseList(
+                densepose_annotations, instances.gt_boxes, image_shape
+            )
+
+        dataset_dict["instances"] = instances[instances.gt_boxes.nonempty()]
+        return dataset_dict
+
+    def _transform_densepose(self, annotation, transforms):
+        if not self.densepose_on:
+            return annotation
+
+        # Handle densepose annotations
+        is_valid, reason_not_valid = DensePoseDataRelative.validate_annotation(annotation)
+        if is_valid:
+            densepose_data = DensePoseDataRelative(annotation, cleanup=True)
+            densepose_data.apply_transform(transforms, self.densepose_transform_data)
+            annotation["densepose"] = densepose_data
+        else:
+            # logger = logging.getLogger(__name__)
+            # logger.debug("Could not load DensePose annotation: {}".format(reason_not_valid))
+            DensePoseDataRelative.cleanup_annotation(annotation)
+            # NOTE: annotations for certain instances may be unavailable.
+            # 'None' is accepted by the DensePostList data structure.
+            annotation["densepose"] = None
+        return annotation
+
+    def _add_densepose_masks_as_segmentation(
+        self, annotations: List[Dict[str, Any]], image_shape_hw: Tuple[int, int]
+    ):
+        for obj in annotations:
+            if ("densepose" not in obj) or ("segmentation" in obj):
+                continue
+            # DP segmentation: torch.Tensor [S, S] of float32, S=256
+            segm_dp = torch.zeros_like(obj["densepose"].segm)
+            segm_dp[obj["densepose"].segm > 0] = 1
+            segm_h, segm_w = segm_dp.shape
+            bbox_segm_dp = torch.tensor((0, 0, segm_h - 1, segm_w - 1), dtype=torch.float32)
+            # image bbox
+            x0, y0, x1, y1 = (
+                v.item() for v in BoxMode.convert(obj["bbox"], obj["bbox_mode"], BoxMode.XYXY_ABS)
+            )
+            segm_aligned = (
+                ROIAlign((y1 - y0, x1 - x0), 1.0, 0, aligned=True)
+                .forward(segm_dp.view(1, 1, *segm_dp.shape), bbox_segm_dp)
+                .squeeze()
+            )
+            image_mask = torch.zeros(*image_shape_hw, dtype=torch.float32)
+            image_mask[y0:y1, x0:x1] = segm_aligned
+            # segmentation for BitMask: np.array [H, W] of bool
+            obj["segmentation"] = image_mask >= 0.5

densepose/data/datasets/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from . import builtin  # ensure the builtin datasets are registered
+
+__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]

densepose/data/datasets/builtin.py

@@ -0,0 +1,18 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+from .chimpnsee import register_dataset as register_chimpnsee_dataset
+from .coco import BASE_DATASETS as BASE_COCO_DATASETS
+from .coco import DATASETS as COCO_DATASETS
+from .coco import register_datasets as register_coco_datasets
+from .lvis import DATASETS as LVIS_DATASETS
+from .lvis import register_datasets as register_lvis_datasets
+
+DEFAULT_DATASETS_ROOT = "datasets"
+
+
+register_coco_datasets(COCO_DATASETS, DEFAULT_DATASETS_ROOT)
+register_coco_datasets(BASE_COCO_DATASETS, DEFAULT_DATASETS_ROOT)
+register_lvis_datasets(LVIS_DATASETS, DEFAULT_DATASETS_ROOT)
+
+register_chimpnsee_dataset(DEFAULT_DATASETS_ROOT)  # pyre-ignore[19]

densepose/data/datasets/chimpnsee.py

@@ -0,0 +1,31 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Optional
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+
+from ..utils import maybe_prepend_base_path
+from .dataset_type import DatasetType
+
+CHIMPNSEE_DATASET_NAME = "chimpnsee"
+
+
+def register_dataset(datasets_root: Optional[str] = None) -> None:
+    def empty_load_callback():
+        pass
+
+    video_list_fpath = maybe_prepend_base_path(
+        datasets_root,
+        "chimpnsee/cdna.eva.mpg.de/video_list.txt",
+    )
+    video_base_path = maybe_prepend_base_path(datasets_root, "chimpnsee/cdna.eva.mpg.de")
+
+    DatasetCatalog.register(CHIMPNSEE_DATASET_NAME, empty_load_callback)
+    MetadataCatalog.get(CHIMPNSEE_DATASET_NAME).set(
+        dataset_type=DatasetType.VIDEO_LIST,
+        video_list_fpath=video_list_fpath,
+        video_base_path=video_base_path,
+        category="chimpanzee",
+    )

densepose/data/datasets/coco.py

@@ -0,0 +1,434 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+import contextlib
+import io
+import logging
+import os
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Optional
+from fvcore.common.timer import Timer
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+
+DENSEPOSE_MASK_KEY = "dp_masks"
+DENSEPOSE_IUV_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_I", "dp_U", "dp_V"]
+DENSEPOSE_CSE_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_vertex", "ref_model"]
+DENSEPOSE_ALL_POSSIBLE_KEYS = set(
+    DENSEPOSE_IUV_KEYS_WITHOUT_MASK + DENSEPOSE_CSE_KEYS_WITHOUT_MASK + [DENSEPOSE_MASK_KEY]
+)
+DENSEPOSE_METADATA_URL_PREFIX = "https://dl.fbaipublicfiles.com/densepose/data/"
+
+
+@dataclass
+class CocoDatasetInfo:
+    name: str
+    images_root: str
+    annotations_fpath: str
+
+
+DATASETS = [
+    CocoDatasetInfo(
+        name="densepose_coco_2014_train",
+        images_root="coco/train2014",
+        annotations_fpath="coco/annotations/densepose_train2014.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival",
+        images_root="coco/val2014",
+        annotations_fpath="coco/annotations/densepose_minival2014.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival_100",
+        images_root="coco/val2014",
+        annotations_fpath="coco/annotations/densepose_minival2014_100.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_valminusminival",
+        images_root="coco/val2014",
+        annotations_fpath="coco/annotations/densepose_valminusminival2014.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_train_cse",
+        images_root="coco/train2014",
+        annotations_fpath="coco_cse/densepose_train2014_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival_cse",
+        images_root="coco/val2014",
+        annotations_fpath="coco_cse/densepose_minival2014_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_minival_100_cse",
+        images_root="coco/val2014",
+        annotations_fpath="coco_cse/densepose_minival2014_100_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_coco_2014_valminusminival_cse",
+        images_root="coco/val2014",
+        annotations_fpath="coco_cse/densepose_valminusminival2014_cse.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_chimps",
+        images_root="densepose_chimps/images",
+        annotations_fpath="densepose_chimps/densepose_chimps_densepose.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_chimps_cse_train",
+        images_root="densepose_chimps/images",
+        annotations_fpath="densepose_chimps/densepose_chimps_cse_train.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_chimps_cse_val",
+        images_root="densepose_chimps/images",
+        annotations_fpath="densepose_chimps/densepose_chimps_cse_val.json",
+    ),
+    CocoDatasetInfo(
+        name="posetrack2017_train",
+        images_root="posetrack2017/posetrack_data_2017",
+        annotations_fpath="posetrack2017/densepose_posetrack_train2017.json",
+    ),
+    CocoDatasetInfo(
+        name="posetrack2017_val",
+        images_root="posetrack2017/posetrack_data_2017",
+        annotations_fpath="posetrack2017/densepose_posetrack_val2017.json",
+    ),
+    CocoDatasetInfo(
+        name="lvis_v05_train",
+        images_root="coco/train2017",
+        annotations_fpath="lvis/lvis_v0.5_plus_dp_train.json",
+    ),
+    CocoDatasetInfo(
+        name="lvis_v05_val",
+        images_root="coco/val2017",
+        annotations_fpath="lvis/lvis_v0.5_plus_dp_val.json",
+    ),
+]
+
+
+BASE_DATASETS = [
+    CocoDatasetInfo(
+        name="base_coco_2017_train",
+        images_root="coco/train2017",
+        annotations_fpath="coco/annotations/instances_train2017.json",
+    ),
+    CocoDatasetInfo(
+        name="base_coco_2017_val",
+        images_root="coco/val2017",
+        annotations_fpath="coco/annotations/instances_val2017.json",
+    ),
+    CocoDatasetInfo(
+        name="base_coco_2017_val_100",
+        images_root="coco/val2017",
+        annotations_fpath="coco/annotations/instances_val2017_100.json",
+    ),
+]
+
+
+def get_metadata(base_path: Optional[str]) -> Dict[str, Any]:
+    """
+    Returns metadata associated with COCO DensePose datasets
+
+    Args:
+    base_path: Optional[str]
+        Base path used to load metadata from
+
+    Returns:
+    Dict[str, Any]
+        Metadata in the form of a dictionary
+    """
+    meta = {
+        "densepose_transform_src": maybe_prepend_base_path(base_path, "UV_symmetry_transforms.mat"),
+        "densepose_smpl_subdiv": maybe_prepend_base_path(base_path, "SMPL_subdiv.mat"),
+        "densepose_smpl_subdiv_transform": maybe_prepend_base_path(
+            base_path,
+            "SMPL_SUBDIV_TRANSFORM.mat",
+        ),
+    }
+    return meta
+
+
+def _load_coco_annotations(json_file: str):
+    """
+    Load COCO annotations from a JSON file
+
+    Args:
+        json_file: str
+            Path to the file to load annotations from
+    Returns:
+        Instance of `pycocotools.coco.COCO` that provides access to annotations
+        data
+    """
+    from pycocotools.coco import COCO
+
+    logger = logging.getLogger(__name__)
+    timer = Timer()
+    with contextlib.redirect_stdout(io.StringIO()):
+        coco_api = COCO(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
+    return coco_api
+
+
+def _add_categories_metadata(dataset_name: str, categories: List[Dict[str, Any]]):
+    meta = MetadataCatalog.get(dataset_name)
+    meta.categories = {c["id"]: c["name"] for c in categories}
+    logger = logging.getLogger(__name__)
+    logger.info("Dataset {} categories: {}".format(dataset_name, meta.categories))
+
+
+def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]):
+    if "minival" in json_file:
+        # Skip validation on COCO2014 valminusminival and minival annotations
+        # The ratio of buggy annotations there is tiny and does not affect accuracy
+        # Therefore we explicitly white-list them
+        return
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
+        json_file
+    )
+
+
+def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    if "bbox" not in ann_dict:
+        return
+    obj["bbox"] = ann_dict["bbox"]
+    obj["bbox_mode"] = BoxMode.XYWH_ABS
+
+
+def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    if "segmentation" not in ann_dict:
+        return
+    segm = ann_dict["segmentation"]
+    if not isinstance(segm, dict):
+        # filter out invalid polygons (< 3 points)
+        segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
+        if len(segm) == 0:
+            return
+    obj["segmentation"] = segm
+
+
+def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    if "keypoints" not in ann_dict:
+        return
+    keypts = ann_dict["keypoints"]  # list[int]
+    for idx, v in enumerate(keypts):
+        if idx % 3 != 2:
+            # COCO's segmentation coordinates are floating points in [0, H or W],
+            # but keypoint coordinates are integers in [0, H-1 or W-1]
+            # Therefore we assume the coordinates are "pixel indices" and
+            # add 0.5 to convert to floating point coordinates.
+            keypts[idx] = v + 0.5
+    obj["keypoints"] = keypts
+
+
+def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
+    for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
+        if key in ann_dict:
+            obj[key] = ann_dict[key]
+
+
+def _combine_images_with_annotations(
+    dataset_name: str,
+    image_root: str,
+    img_datas: Iterable[Dict[str, Any]],
+    ann_datas: Iterable[Iterable[Dict[str, Any]]],
+):
+
+    ann_keys = ["iscrowd", "category_id"]
+    dataset_dicts = []
+    contains_video_frame_info = False
+
+    for img_dict, ann_dicts in zip(img_datas, ann_datas):
+        record = {}
+        record["file_name"] = os.path.join(image_root, img_dict["file_name"])
+        record["height"] = img_dict["height"]
+        record["width"] = img_dict["width"]
+        record["image_id"] = img_dict["id"]
+        record["dataset"] = dataset_name
+        if "frame_id" in img_dict:
+            record["frame_id"] = img_dict["frame_id"]
+            record["video_id"] = img_dict.get("vid_id", None)
+            contains_video_frame_info = True
+        objs = []
+        for ann_dict in ann_dicts:
+            assert ann_dict["image_id"] == record["image_id"]
+            assert ann_dict.get("ignore", 0) == 0
+            obj = {key: ann_dict[key] for key in ann_keys if key in ann_dict}
+            _maybe_add_bbox(obj, ann_dict)
+            _maybe_add_segm(obj, ann_dict)
+            _maybe_add_keypoints(obj, ann_dict)
+            _maybe_add_densepose(obj, ann_dict)
+            objs.append(obj)
+        record["annotations"] = objs
+        dataset_dicts.append(record)
+    if contains_video_frame_info:
+        create_video_frame_mapping(dataset_name, dataset_dicts)
+    return dataset_dicts
+
+
+def get_contiguous_id_to_category_id_map(metadata):
+    cat_id_2_cont_id = metadata.thing_dataset_id_to_contiguous_id
+    cont_id_2_cat_id = {}
+    for cat_id, cont_id in cat_id_2_cont_id.items():
+        if cont_id in cont_id_2_cat_id:
+            continue
+        cont_id_2_cat_id[cont_id] = cat_id
+    return cont_id_2_cat_id
+
+
+def maybe_filter_categories_cocoapi(dataset_name, coco_api):
+    meta = MetadataCatalog.get(dataset_name)
+    cont_id_2_cat_id = get_contiguous_id_to_category_id_map(meta)
+    cat_id_2_cont_id = meta.thing_dataset_id_to_contiguous_id
+    # filter categories
+    cats = []
+    for cat in coco_api.dataset["categories"]:
+        cat_id = cat["id"]
+        if cat_id not in cat_id_2_cont_id:
+            continue
+        cont_id = cat_id_2_cont_id[cat_id]
+        if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
+            cats.append(cat)
+    coco_api.dataset["categories"] = cats
+    # filter annotations, if multiple categories are mapped to a single
+    # contiguous ID, use only one category ID and map all annotations to that category ID
+    anns = []
+    for ann in coco_api.dataset["annotations"]:
+        cat_id = ann["category_id"]
+        if cat_id not in cat_id_2_cont_id:
+            continue
+        cont_id = cat_id_2_cont_id[cat_id]
+        ann["category_id"] = cont_id_2_cat_id[cont_id]
+        anns.append(ann)
+    coco_api.dataset["annotations"] = anns
+    # recreate index
+    coco_api.createIndex()
+
+
+def maybe_filter_and_map_categories_cocoapi(dataset_name, coco_api):
+    meta = MetadataCatalog.get(dataset_name)
+    category_id_map = meta.thing_dataset_id_to_contiguous_id
+    # map categories
+    cats = []
+    for cat in coco_api.dataset["categories"]:
+        cat_id = cat["id"]
+        if cat_id not in category_id_map:
+            continue
+        cat["id"] = category_id_map[cat_id]
+        cats.append(cat)
+    coco_api.dataset["categories"] = cats
+    # map annotation categories
+    anns = []
+    for ann in coco_api.dataset["annotations"]:
+        cat_id = ann["category_id"]
+        if cat_id not in category_id_map:
+            continue
+        ann["category_id"] = category_id_map[cat_id]
+        anns.append(ann)
+    coco_api.dataset["annotations"] = anns
+    # recreate index
+    coco_api.createIndex()
+
+
+def create_video_frame_mapping(dataset_name, dataset_dicts):
+    mapping = defaultdict(dict)
+    for d in dataset_dicts:
+        video_id = d.get("video_id")
+        if video_id is None:
+            continue
+        mapping[video_id].update({d["frame_id"]: d["file_name"]})
+    MetadataCatalog.get(dataset_name).set(video_frame_mapping=mapping)
+
+
+def load_coco_json(annotations_json_file: str, image_root: str, dataset_name: str):
+    """
+    Loads a JSON file with annotations in COCO instances format.
+    Replaces `detectron2.data.datasets.coco.load_coco_json` to handle metadata
+    in a more flexible way. Postpones category mapping to a later stage to be
+    able to combine several datasets with different (but coherent) sets of
+    categories.
+
+    Args:
+
+    annotations_json_file: str
+        Path to the JSON file with annotations in COCO instances format.
+    image_root: str
+        directory that contains all the images
+    dataset_name: str
+        the name that identifies a dataset, e.g. "densepose_coco_2014_train"
+    extra_annotation_keys: Optional[List[str]]
+        If provided, these keys are used to extract additional data from
+        the annotations.
+    """
+    coco_api = _load_coco_annotations(PathManager.get_local_path(annotations_json_file))
+    _add_categories_metadata(dataset_name, coco_api.loadCats(coco_api.getCatIds()))
+    # sort indices for reproducible results
+    img_ids = sorted(coco_api.imgs.keys())
+    # imgs is a list of dicts, each looks something like:
+    # {'license': 4,
+    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
+    #  'file_name': 'COCO_val2014_000000001268.jpg',
+    #  'height': 427,
+    #  'width': 640,
+    #  'date_captured': '2013-11-17 05:57:24',
+    #  'id': 1268}
+    imgs = coco_api.loadImgs(img_ids)
+    logger = logging.getLogger(__name__)
+    logger.info("Loaded {} images in COCO format from {}".format(len(imgs), annotations_json_file))
+    # anns is a list[list[dict]], where each dict is an annotation
+    # record for an object. The inner list enumerates the objects in an image
+    # and the outer list enumerates over images.
+    anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
+    _verify_annotations_have_unique_ids(annotations_json_file, anns)
+    dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
+    return dataset_records
+
+
+def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None):
+    """
+    Registers provided COCO DensePose dataset
+
+    Args:
+    dataset_data: CocoDatasetInfo
+        Dataset data
+    datasets_root: Optional[str]
+        Datasets root folder (default: None)
+    """
+    annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
+    images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
+
+    def load_annotations():
+        return load_coco_json(
+            annotations_json_file=annotations_fpath,
+            image_root=images_root,
+            dataset_name=dataset_data.name,
+        )
+
+    DatasetCatalog.register(dataset_data.name, load_annotations)
+    MetadataCatalog.get(dataset_data.name).set(
+        json_file=annotations_fpath,
+        image_root=images_root,
+        **get_metadata(DENSEPOSE_METADATA_URL_PREFIX)
+    )
+
+
+def register_datasets(
+    datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
+):
+    """
+    Registers provided COCO DensePose datasets
+
+    Args:
+    datasets_data: Iterable[CocoDatasetInfo]
+        An iterable of dataset datas
+    datasets_root: Optional[str]
+        Datasets root folder (default: None)
+    """
+    for dataset_data in datasets_data:
+        register_dataset(dataset_data, datasets_root)

densepose/data/datasets/dataset_type.py

@@ -0,0 +1,13 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from enum import Enum
+
+
+class DatasetType(Enum):
+    """
+    Dataset type, mostly used for datasets that contain data to bootstrap models on
+    """
+
+    VIDEO_LIST = "video_list"

densepose/data/datasets/lvis.py

@@ -0,0 +1,259 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+import logging
+import os
+from typing import Any, Dict, Iterable, List, Optional
+from fvcore.common.timer import Timer
+
+from detectron2.data import DatasetCatalog, MetadataCatalog
+from detectron2.data.datasets.lvis import get_lvis_instances_meta
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+from .coco import (
+    DENSEPOSE_ALL_POSSIBLE_KEYS,
+    DENSEPOSE_METADATA_URL_PREFIX,
+    CocoDatasetInfo,
+    get_metadata,
+)
+
+DATASETS = [
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds1_train_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds1_train_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds1_val_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds1_val_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds2_train_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds2_train_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds2_val_v1",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_ds2_val_v1.json",
+    ),
+    CocoDatasetInfo(
+        name="densepose_lvis_v1_ds1_val_animals_100",
+        images_root="coco_",
+        annotations_fpath="lvis/densepose_lvis_v1_val_animals_100_v2.json",
+    ),
+]
+
+
+def _load_lvis_annotations(json_file: str):
+    """
+    Load COCO annotations from a JSON file
+
+    Args:
+        json_file: str
+            Path to the file to load annotations from
+    Returns:
+        Instance of `pycocotools.coco.COCO` that provides access to annotations
+        data
+    """
+    from lvis import LVIS
+
+    json_file = PathManager.get_local_path(json_file)
+    logger = logging.getLogger(__name__)
+    timer = Timer()
+    lvis_api = LVIS(json_file)
+    if timer.seconds() > 1:
+        logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
+    return lvis_api
+
+
+def _add_categories_metadata(dataset_name: str) -> None:
+    metadict = get_lvis_instances_meta(dataset_name)
+    categories = metadict["thing_classes"]
+    metadata = MetadataCatalog.get(dataset_name)
+    metadata.categories = {i + 1: categories[i] for i in range(len(categories))}
+    logger = logging.getLogger(__name__)
+    logger.info(f"Dataset {dataset_name} has {len(categories)} categories")
+
+
+def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]) -> None:
+    ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
+    assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
+        json_file
+    )
+
+
+def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    if "bbox" not in ann_dict:
+        return
+    obj["bbox"] = ann_dict["bbox"]
+    obj["bbox_mode"] = BoxMode.XYWH_ABS
+
+
+def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    if "segmentation" not in ann_dict:
+        return
+    segm = ann_dict["segmentation"]
+    if not isinstance(segm, dict):
+        # filter out invalid polygons (< 3 points)
+        segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
+        if len(segm) == 0:
+            return
+    obj["segmentation"] = segm
+
+
+def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    if "keypoints" not in ann_dict:
+        return
+    keypts = ann_dict["keypoints"]  # list[int]
+    for idx, v in enumerate(keypts):
+        if idx % 3 != 2:
+            # COCO's segmentation coordinates are floating points in [0, H or W],
+            # but keypoint coordinates are integers in [0, H-1 or W-1]
+            # Therefore we assume the coordinates are "pixel indices" and
+            # add 0.5 to convert to floating point coordinates.
+            keypts[idx] = v + 0.5
+    obj["keypoints"] = keypts
+
+
+def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
+    for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
+        if key in ann_dict:
+            obj[key] = ann_dict[key]
+
+
+def _combine_images_with_annotations(
+    dataset_name: str,
+    image_root: str,
+    img_datas: Iterable[Dict[str, Any]],
+    ann_datas: Iterable[Iterable[Dict[str, Any]]],
+):
+
+    dataset_dicts = []
+
+    def get_file_name(img_root, img_dict):
+        # Determine the path including the split folder ("train2017", "val2017", "test2017") from
+        # the coco_url field. Example:
+        #   'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg'
+        split_folder, file_name = img_dict["coco_url"].split("/")[-2:]
+        return os.path.join(img_root + split_folder, file_name)
+
+    for img_dict, ann_dicts in zip(img_datas, ann_datas):
+        record = {}
+        record["file_name"] = get_file_name(image_root, img_dict)
+        record["height"] = img_dict["height"]
+        record["width"] = img_dict["width"]
+        record["not_exhaustive_category_ids"] = img_dict.get("not_exhaustive_category_ids", [])
+        record["neg_category_ids"] = img_dict.get("neg_category_ids", [])
+        record["image_id"] = img_dict["id"]
+        record["dataset"] = dataset_name
+
+        objs = []
+        for ann_dict in ann_dicts:
+            assert ann_dict["image_id"] == record["image_id"]
+            obj = {}
+            _maybe_add_bbox(obj, ann_dict)
+            obj["iscrowd"] = ann_dict.get("iscrowd", 0)
+            obj["category_id"] = ann_dict["category_id"]
+            _maybe_add_segm(obj, ann_dict)
+            _maybe_add_keypoints(obj, ann_dict)
+            _maybe_add_densepose(obj, ann_dict)
+            objs.append(obj)
+        record["annotations"] = objs
+        dataset_dicts.append(record)
+    return dataset_dicts
+
+
+def load_lvis_json(annotations_json_file: str, image_root: str, dataset_name: str):
+    """
+    Loads a JSON file with annotations in LVIS instances format.
+    Replaces `detectron2.data.datasets.coco.load_lvis_json` to handle metadata
+    in a more flexible way. Postpones category mapping to a later stage to be
+    able to combine several datasets with different (but coherent) sets of
+    categories.
+
+    Args:
+
+    annotations_json_file: str
+        Path to the JSON file with annotations in COCO instances format.
+    image_root: str
+        directory that contains all the images
+    dataset_name: str
+        the name that identifies a dataset, e.g. "densepose_coco_2014_train"
+    extra_annotation_keys: Optional[List[str]]
+        If provided, these keys are used to extract additional data from
+        the annotations.
+    """
+    lvis_api = _load_lvis_annotations(PathManager.get_local_path(annotations_json_file))
+
+    _add_categories_metadata(dataset_name)
+
+    # sort indices for reproducible results
+    img_ids = sorted(lvis_api.imgs.keys())
+    # imgs is a list of dicts, each looks something like:
+    # {'license': 4,
+    #  'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
+    #  'file_name': 'COCO_val2014_000000001268.jpg',
+    #  'height': 427,
+    #  'width': 640,
+    #  'date_captured': '2013-11-17 05:57:24',
+    #  'id': 1268}
+    imgs = lvis_api.load_imgs(img_ids)
+    logger = logging.getLogger(__name__)
+    logger.info("Loaded {} images in LVIS format from {}".format(len(imgs), annotations_json_file))
+    # anns is a list[list[dict]], where each dict is an annotation
+    # record for an object. The inner list enumerates the objects in an image
+    # and the outer list enumerates over images.
+    anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
+
+    _verify_annotations_have_unique_ids(annotations_json_file, anns)
+    dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
+    return dataset_records
+
+
+def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None) -> None:
+    """
+    Registers provided LVIS DensePose dataset
+
+    Args:
+        dataset_data: CocoDatasetInfo
+            Dataset data
+        datasets_root: Optional[str]
+            Datasets root folder (default: None)
+    """
+    annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
+    images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
+
+    def load_annotations():
+        return load_lvis_json(
+            annotations_json_file=annotations_fpath,
+            image_root=images_root,
+            dataset_name=dataset_data.name,
+        )
+
+    DatasetCatalog.register(dataset_data.name, load_annotations)
+    MetadataCatalog.get(dataset_data.name).set(
+        json_file=annotations_fpath,
+        image_root=images_root,
+        evaluator_type="lvis",
+        **get_metadata(DENSEPOSE_METADATA_URL_PREFIX),
+    )
+
+
+def register_datasets(
+    datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
+) -> None:
+    """
+    Registers provided LVIS DensePose datasets
+
+    Args:
+        datasets_data: Iterable[CocoDatasetInfo]
+            An iterable of dataset datas
+        datasets_root: Optional[str]
+            Datasets root folder (default: None)
+    """
+    for dataset_data in datasets_data:
+        register_dataset(dataset_data, datasets_root)

densepose/data/image_list_dataset.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import logging
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Union
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.data.detection_utils import read_image
+
+ImageTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+class ImageListDataset(Dataset):
+    """
+    Dataset that provides images from a list.
+    """
+
+    _EMPTY_IMAGE = torch.empty((0, 3, 1, 1))
+
+    def __init__(
+        self,
+        image_list: List[str],
+        category_list: Union[str, List[str], None] = None,
+        transform: Optional[ImageTransform] = None,
+    ):
+        """
+        Args:
+            image_list (List[str]): list of paths to image files
+            category_list (Union[str, List[str], None]): list of animal categories for
+                each image. If it is a string, or None, this applies to all images
+        """
+        if type(category_list) is list:
+            self.category_list = category_list
+        else:
+            self.category_list = [category_list] * len(image_list)
+        assert len(image_list) == len(
+            self.category_list
+        ), "length of image and category lists must be equal"
+        self.image_list = image_list
+        self.transform = transform
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Gets selected images from the list
+
+        Args:
+            idx (int): video index in the video list file
+        Returns:
+            A dictionary containing two keys:
+                images (torch.Tensor): tensor of size [N, 3, H, W] (N = 1, or 0 for _EMPTY_IMAGE)
+                categories (List[str]): categories of the frames
+        """
+        categories = [self.category_list[idx]]
+        fpath = self.image_list[idx]
+        transform = self.transform
+
+        try:
+            image = torch.from_numpy(np.ascontiguousarray(read_image(fpath, format="BGR")))
+            image = image.permute(2, 0, 1).unsqueeze(0).float()  # HWC -> NCHW
+            if transform is not None:
+                image = transform(image)
+            return {"images": image, "categories": categories}
+        except (OSError, RuntimeError) as e:
+            logger = logging.getLogger(__name__)
+            logger.warning(f"Error opening image file container {fpath}: {e}")
+
+        return {"images": self._EMPTY_IMAGE, "categories": []}
+
+    def __len__(self):
+        return len(self.image_list)

densepose/data/inference_based_loader.py

@@ -0,0 +1,174 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple
+import torch
+from torch import nn
+
+SampledData = Any
+ModelOutput = Any
+
+
+def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]:
+    """
+    Group elements of an iterable by chunks of size `n`, e.g.
+    grouper(range(9), 4) ->
+        (0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None)
+    """
+    it = iter(iterable)
+    while True:
+        values = []
+        for _ in range(n):
+            try:
+                value = next(it)
+            except StopIteration:
+                if values:
+                    values.extend([fillvalue] * (n - len(values)))
+                    yield tuple(values)
+                return
+            values.append(value)
+        yield tuple(values)
+
+
+class ScoreBasedFilter:
+    """
+    Filters entries in model output based on their scores
+    Discards all entries with score less than the specified minimum
+    """
+
+    def __init__(self, min_score: float = 0.8):
+        self.min_score = min_score
+
+    def __call__(self, model_output: ModelOutput) -> ModelOutput:
+        for model_output_i in model_output:
+            instances = model_output_i["instances"]
+            if not instances.has("scores"):
+                continue
+            instances_filtered = instances[instances.scores >= self.min_score]
+            model_output_i["instances"] = instances_filtered
+        return model_output
+
+
+class InferenceBasedLoader:
+    """
+    Data loader based on results inferred by a model. Consists of:
+     - a data loader that provides batches of images
+     - a model that is used to infer the results
+     - a data sampler that converts inferred results to annotations
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        data_loader: Iterable[List[Dict[str, Any]]],
+        data_sampler: Optional[Callable[[ModelOutput], List[SampledData]]] = None,
+        data_filter: Optional[Callable[[ModelOutput], ModelOutput]] = None,
+        shuffle: bool = True,
+        batch_size: int = 4,
+        inference_batch_size: int = 4,
+        drop_last: bool = False,
+        category_to_class_mapping: Optional[dict] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          model (torch.nn.Module): model used to produce data
+          data_loader (Iterable[List[Dict[str, Any]]]): iterable that provides
+            dictionaries with "images" and "categories" fields to perform inference on
+          data_sampler (Callable: ModelOutput -> SampledData): functor
+              that produces annotation data from inference results;
+              (optional, default: None)
+          data_filter (Callable: ModelOutput -> ModelOutput): filter
+              that selects model outputs for further processing
+              (optional, default: None)
+          shuffle (bool): if True, the input images get shuffled
+          batch_size (int): batch size for the produced annotation data
+          inference_batch_size (int): batch size for input images
+          drop_last (bool): if True, drop the last batch if it is undersized
+          category_to_class_mapping (dict): category to class mapping
+        """
+        self.model = model
+        self.model.eval()
+        self.data_loader = data_loader
+        self.data_sampler = data_sampler
+        self.data_filter = data_filter
+        self.shuffle = shuffle
+        self.batch_size = batch_size
+        self.inference_batch_size = inference_batch_size
+        self.drop_last = drop_last
+        if category_to_class_mapping is not None:
+            self.category_to_class_mapping = category_to_class_mapping
+        else:
+            self.category_to_class_mapping = {}
+
+    def __iter__(self) -> Iterator[List[SampledData]]:
+        for batch in self.data_loader:
+            # batch : List[Dict[str: Tensor[N, C, H, W], str: Optional[str]]]
+            # images_batch : Tensor[N, C, H, W]
+            # image : Tensor[C, H, W]
+            images_and_categories = [
+                {"image": image, "category": category}
+                for element in batch
+                for image, category in zip(element["images"], element["categories"])
+            ]
+            if not images_and_categories:
+                continue
+            if self.shuffle:
+                random.shuffle(images_and_categories)
+            yield from self._produce_data(images_and_categories)  # pyre-ignore[6]
+
+    def _produce_data(
+        self, images_and_categories: List[Tuple[torch.Tensor, Optional[str]]]
+    ) -> Iterator[List[SampledData]]:
+        """
+        Produce batches of data from images
+
+        Args:
+          images_and_categories (List[Tuple[torch.Tensor, Optional[str]]]):
+            list of images and corresponding categories to process
+
+        Returns:
+          Iterator over batches of data sampled from model outputs
+        """
+        data_batches: List[SampledData] = []
+        category_to_class_mapping = self.category_to_class_mapping
+        batched_images_and_categories = _grouper(images_and_categories, self.inference_batch_size)
+        for batch in batched_images_and_categories:
+            batch = [
+                {
+                    "image": image_and_category["image"].to(self.model.device),
+                    "category": image_and_category["category"],
+                }
+                for image_and_category in batch
+                if image_and_category is not None
+            ]
+            if not batch:
+                continue
+            with torch.no_grad():
+                model_output = self.model(batch)
+            for model_output_i, batch_i in zip(model_output, batch):
+                assert len(batch_i["image"].shape) == 3
+                model_output_i["image"] = batch_i["image"]
+                instance_class = category_to_class_mapping.get(batch_i["category"], 0)
+                model_output_i["instances"].dataset_classes = torch.tensor(
+                    [instance_class] * len(model_output_i["instances"])
+                )
+            model_output_filtered = (
+                model_output if self.data_filter is None else self.data_filter(model_output)
+            )
+            data = (
+                model_output_filtered
+                if self.data_sampler is None
+                else self.data_sampler(model_output_filtered)
+            )
+            for data_i in data:
+                if len(data_i["instances"]):
+                    data_batches.append(data_i)
+            if len(data_batches) >= self.batch_size:
+                yield data_batches[: self.batch_size]
+                data_batches = data_batches[self.batch_size :]
+        if not self.drop_last and data_batches:
+            yield data_batches

densepose/data/meshes/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+from . import builtin
+
+__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]

densepose/data/meshes/builtin.py

@@ -0,0 +1,103 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+from .catalog import MeshInfo, register_meshes
+
+DENSEPOSE_MESHES_DIR = "https://dl.fbaipublicfiles.com/densepose/meshes/"
+
+MESHES = [
+    MeshInfo(
+        name="smpl_27554",
+        data="smpl_27554.pkl",
+        geodists="geodists/geodists_smpl_27554.pkl",
+        symmetry="symmetry/symmetry_smpl_27554.pkl",
+        texcoords="texcoords/texcoords_smpl_27554.pkl",
+    ),
+    MeshInfo(
+        name="chimp_5029",
+        data="chimp_5029.pkl",
+        geodists="geodists/geodists_chimp_5029.pkl",
+        symmetry="symmetry/symmetry_chimp_5029.pkl",
+        texcoords="texcoords/texcoords_chimp_5029.pkl",
+    ),
+    MeshInfo(
+        name="cat_5001",
+        data="cat_5001.pkl",
+        geodists="geodists/geodists_cat_5001.pkl",
+        symmetry="symmetry/symmetry_cat_5001.pkl",
+        texcoords="texcoords/texcoords_cat_5001.pkl",
+    ),
+    MeshInfo(
+        name="cat_7466",
+        data="cat_7466.pkl",
+        geodists="geodists/geodists_cat_7466.pkl",
+        symmetry="symmetry/symmetry_cat_7466.pkl",
+        texcoords="texcoords/texcoords_cat_7466.pkl",
+    ),
+    MeshInfo(
+        name="sheep_5004",
+        data="sheep_5004.pkl",
+        geodists="geodists/geodists_sheep_5004.pkl",
+        symmetry="symmetry/symmetry_sheep_5004.pkl",
+        texcoords="texcoords/texcoords_sheep_5004.pkl",
+    ),
+    MeshInfo(
+        name="zebra_5002",
+        data="zebra_5002.pkl",
+        geodists="geodists/geodists_zebra_5002.pkl",
+        symmetry="symmetry/symmetry_zebra_5002.pkl",
+        texcoords="texcoords/texcoords_zebra_5002.pkl",
+    ),
+    MeshInfo(
+        name="horse_5004",
+        data="horse_5004.pkl",
+        geodists="geodists/geodists_horse_5004.pkl",
+        symmetry="symmetry/symmetry_horse_5004.pkl",
+        texcoords="texcoords/texcoords_zebra_5002.pkl",
+    ),
+    MeshInfo(
+        name="giraffe_5002",
+        data="giraffe_5002.pkl",
+        geodists="geodists/geodists_giraffe_5002.pkl",
+        symmetry="symmetry/symmetry_giraffe_5002.pkl",
+        texcoords="texcoords/texcoords_giraffe_5002.pkl",
+    ),
+    MeshInfo(
+        name="elephant_5002",
+        data="elephant_5002.pkl",
+        geodists="geodists/geodists_elephant_5002.pkl",
+        symmetry="symmetry/symmetry_elephant_5002.pkl",
+        texcoords="texcoords/texcoords_elephant_5002.pkl",
+    ),
+    MeshInfo(
+        name="dog_5002",
+        data="dog_5002.pkl",
+        geodists="geodists/geodists_dog_5002.pkl",
+        symmetry="symmetry/symmetry_dog_5002.pkl",
+        texcoords="texcoords/texcoords_dog_5002.pkl",
+    ),
+    MeshInfo(
+        name="dog_7466",
+        data="dog_7466.pkl",
+        geodists="geodists/geodists_dog_7466.pkl",
+        symmetry="symmetry/symmetry_dog_7466.pkl",
+        texcoords="texcoords/texcoords_dog_7466.pkl",
+    ),
+    MeshInfo(
+        name="cow_5002",
+        data="cow_5002.pkl",
+        geodists="geodists/geodists_cow_5002.pkl",
+        symmetry="symmetry/symmetry_cow_5002.pkl",
+        texcoords="texcoords/texcoords_cow_5002.pkl",
+    ),
+    MeshInfo(
+        name="bear_4936",
+        data="bear_4936.pkl",
+        geodists="geodists/geodists_bear_4936.pkl",
+        symmetry="symmetry/symmetry_bear_4936.pkl",
+        texcoords="texcoords/texcoords_bear_4936.pkl",
+    ),
+]
+
+register_meshes(MESHES, DENSEPOSE_MESHES_DIR)

densepose/data/meshes/catalog.py

@@ -0,0 +1,73 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+import logging
+from collections import UserDict
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+from ..utils import maybe_prepend_base_path
+
+
+@dataclass
+class MeshInfo:
+    name: str
+    data: str
+    geodists: Optional[str] = None
+    symmetry: Optional[str] = None
+    texcoords: Optional[str] = None
+
+
+class _MeshCatalog(UserDict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mesh_ids = {}
+        self.mesh_names = {}
+        self.max_mesh_id = -1
+
+    def __setitem__(self, key, value):
+        if key in self:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                f"Overwriting mesh catalog entry '{key}': old value {self[key]}"
+                f", new value {value}"
+            )
+            mesh_id = self.mesh_ids[key]
+        else:
+            self.max_mesh_id += 1
+            mesh_id = self.max_mesh_id
+        super().__setitem__(key, value)
+        self.mesh_ids[key] = mesh_id
+        self.mesh_names[mesh_id] = key
+
+    def get_mesh_id(self, shape_name: str) -> int:
+        return self.mesh_ids[shape_name]
+
+    def get_mesh_name(self, mesh_id: int) -> str:
+        return self.mesh_names[mesh_id]
+
+
+MeshCatalog = _MeshCatalog()
+
+
+def register_mesh(mesh_info: MeshInfo, base_path: Optional[str]) -> None:
+    geodists, symmetry, texcoords = mesh_info.geodists, mesh_info.symmetry, mesh_info.texcoords
+    if geodists:
+        geodists = maybe_prepend_base_path(base_path, geodists)
+    if symmetry:
+        symmetry = maybe_prepend_base_path(base_path, symmetry)
+    if texcoords:
+        texcoords = maybe_prepend_base_path(base_path, texcoords)
+    MeshCatalog[mesh_info.name] = MeshInfo(
+        name=mesh_info.name,
+        data=maybe_prepend_base_path(base_path, mesh_info.data),
+        geodists=geodists,
+        symmetry=symmetry,
+        texcoords=texcoords,
+    )
+
+
+def register_meshes(mesh_infos: Iterable[MeshInfo], base_path: Optional[str]) -> None:
+    for mesh_info in mesh_infos:
+        register_mesh(mesh_info, base_path)

densepose/data/samplers/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .densepose_uniform import DensePoseUniformSampler
+from .densepose_confidence_based import DensePoseConfidenceBasedSampler
+from .densepose_cse_uniform import DensePoseCSEUniformSampler
+from .densepose_cse_confidence_based import DensePoseCSEConfidenceBasedSampler
+from .mask_from_densepose import MaskFromDensePoseSampler
+from .prediction_to_gt import PredictionToGroundTruthSampler

densepose/data/samplers/densepose_base.py

@@ -0,0 +1,205 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Dict, List, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures import BoxMode, Instances
+
+from densepose.converters import ToChartResultConverter
+from densepose.converters.base import IntTupleBox, make_int_box
+from densepose.structures import DensePoseDataRelative, DensePoseList
+
+
+class DensePoseBaseSampler:
+    """
+    Base DensePose sampler to produce DensePose data from DensePose predictions.
+    Samples for each class are drawn according to some distribution over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(self, count_per_class: int = 8):
+        """
+        Constructor
+
+        Args:
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        self.count_per_class = count_per_class
+
+    def __call__(self, instances: Instances) -> DensePoseList:
+        """
+        Convert DensePose predictions (an instance of `DensePoseChartPredictorOutput`)
+        into DensePose annotations data (an instance of `DensePoseList`)
+        """
+        boxes_xyxy_abs = instances.pred_boxes.tensor.clone().cpu()
+        boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+        dp_datas = []
+        for i in range(len(boxes_xywh_abs)):
+            annotation_i = self._sample(instances[i], make_int_box(boxes_xywh_abs[i]))
+            annotation_i[DensePoseDataRelative.S_KEY] = self._resample_mask(  # pyre-ignore[6]
+                instances[i].pred_densepose
+            )
+            dp_datas.append(DensePoseDataRelative(annotation_i))
+        # create densepose annotations on CPU
+        dp_list = DensePoseList(dp_datas, boxes_xyxy_abs, instances.image_size)
+        return dp_list
+
+    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
+        """
+        Sample DensPoseDataRelative from estimation results
+        """
+        labels, dp_result = self._produce_labels_and_results(instance)
+        annotation = {
+            DensePoseDataRelative.X_KEY: [],
+            DensePoseDataRelative.Y_KEY: [],
+            DensePoseDataRelative.U_KEY: [],
+            DensePoseDataRelative.V_KEY: [],
+            DensePoseDataRelative.I_KEY: [],
+        }
+        n, h, w = dp_result.shape
+        for part_id in range(1, DensePoseDataRelative.N_PART_LABELS + 1):
+            # indices - tuple of 3 1D tensors of size k
+            # 0: index along the first dimension N
+            # 1: index along H dimension
+            # 2: index along W dimension
+            indices = torch.nonzero(labels.expand(n, h, w) == part_id, as_tuple=True)
+            # values - an array of size [n, k]
+            # n: number of channels (U, V, confidences)
+            # k: number of points labeled with part_id
+            values = dp_result[indices].view(n, -1)
+            k = values.shape[1]
+            count = min(self.count_per_class, k)
+            if count <= 0:
+                continue
+            index_sample = self._produce_index_sample(values, count)
+            sampled_values = values[:, index_sample]
+            sampled_y = indices[1][index_sample] + 0.5
+            sampled_x = indices[2][index_sample] + 0.5
+            # prepare / normalize data
+            x = (sampled_x / w * 256.0).cpu().tolist()
+            y = (sampled_y / h * 256.0).cpu().tolist()
+            u = sampled_values[0].clamp(0, 1).cpu().tolist()
+            v = sampled_values[1].clamp(0, 1).cpu().tolist()
+            fine_segm_labels = [part_id] * count
+            # extend annotations
+            annotation[DensePoseDataRelative.X_KEY].extend(x)
+            annotation[DensePoseDataRelative.Y_KEY].extend(y)
+            annotation[DensePoseDataRelative.U_KEY].extend(u)
+            annotation[DensePoseDataRelative.V_KEY].extend(v)
+            annotation[DensePoseDataRelative.I_KEY].extend(fine_segm_labels)
+        return annotation
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Abstract method to produce a sample of indices to select data
+        To be implemented in descendants
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        raise NotImplementedError
+
+    def _produce_labels_and_results(self, instance: Instances) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of `DensePoseChartPredictorOutput`
+
+        Return:
+            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
+            dp_result (torch.Tensor): shape [2, H, W], stacked DensePose results u and v
+        """
+        converter = ToChartResultConverter
+        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
+        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
+        return labels, dp_result
+
+    def _resample_mask(self, output: Any) -> torch.Tensor:
+        """
+        Convert DensePose predictor output to segmentation annotation - tensors of size
+        (256, 256) and type `int64`.
+
+        Args:
+            output: DensePose predictor output with the following attributes:
+             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
+               segmentation scores
+             - fine_segm: tensor of size [N, C, H, W] with unnormalized fine
+               segmentation scores
+        Return:
+            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
+            where S = DensePoseDataRelative.MASK_SIZE
+        """
+        sz = DensePoseDataRelative.MASK_SIZE
+        S = (
+            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
+            .argmax(dim=1)
+            .long()
+        )
+        I = (
+            (
+                F.interpolate(
+                    output.fine_segm,
+                    (sz, sz),
+                    mode="bilinear",
+                    align_corners=False,
+                ).argmax(dim=1)
+                * (S > 0).long()
+            )
+            .squeeze()
+            .cpu()
+        )
+        # Map fine segmentation results to coarse segmentation ground truth
+        # TODO: extract this into separate classes
+        # coarse segmentation: 1 = Torso, 2 = Right Hand, 3 = Left Hand,
+        # 4 = Left Foot, 5 = Right Foot, 6 = Upper Leg Right, 7 = Upper Leg Left,
+        # 8 = Lower Leg Right, 9 = Lower Leg Left, 10 = Upper Arm Left,
+        # 11 = Upper Arm Right, 12 = Lower Arm Left, 13 = Lower Arm Right,
+        # 14 = Head
+        # fine segmentation: 1, 2 = Torso, 3 = Right Hand, 4 = Left Hand,
+        # 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right,
+        # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right,
+        # 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left,
+        # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left,
+        # 20, 22 = Lower Arm Right, 23, 24 = Head
+        FINE_TO_COARSE_SEGMENTATION = {
+            1: 1,
+            2: 1,
+            3: 2,
+            4: 3,
+            5: 4,
+            6: 5,
+            7: 6,
+            8: 7,
+            9: 6,
+            10: 7,
+            11: 8,
+            12: 9,
+            13: 8,
+            14: 9,
+            15: 10,
+            16: 11,
+            17: 10,
+            18: 11,
+            19: 12,
+            20: 13,
+            21: 12,
+            22: 13,
+            23: 14,
+            24: 14,
+        }
+        mask = torch.zeros((sz, sz), dtype=torch.int64, device=torch.device("cpu"))
+        for i in range(DensePoseDataRelative.N_PART_LABELS):
+            mask[I == i + 1] = FINE_TO_COARSE_SEGMENTATION[i + 1]
+        return mask

densepose/data/samplers/densepose_confidence_based.py

@@ -0,0 +1,110 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from typing import Optional, Tuple
+import torch
+
+from densepose.converters import ToChartResultConverterWithConfidences
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseConfidenceBasedSampler(DensePoseBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn using confidence value estimates.
+    """
+
+    def __init__(
+        self,
+        confidence_channel: str,
+        count_per_class: int = 8,
+        search_count_multiplier: Optional[float] = None,
+        search_proportion: Optional[float] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          confidence_channel (str): confidence channel to use for sampling;
+            possible values:
+              "sigma_2": confidences for UV values
+              "fine_segm_confidence": confidences for fine segmentation
+              "coarse_segm_confidence": confidences for coarse segmentation
+            (default: "sigma_2")
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category (default: 8)
+          search_count_multiplier (float or None): if not None, the total number
+              of the most confident estimates of a given class to consider is
+              defined as `min(search_count_multiplier * count_per_class, N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_proportion` (default: None)
+          search_proportion (float or None): if not None, the total number of the
+              of the most confident estimates of a given class to consider is
+              defined as `min(max(search_proportion * N, count_per_class), N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_count_multiplier` (default: None)
+        """
+        super().__init__(count_per_class)
+        self.confidence_channel = confidence_channel
+        self.search_count_multiplier = search_count_multiplier
+        self.search_proportion = search_proportion
+        assert (search_count_multiplier is None) or (search_proportion is None), (
+            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
+            f"and search_proportion (={search_proportion})"
+        )
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a sample of indices to select data based on confidences
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        if k == count:
+            index_sample = list(range(k))
+        else:
+            # take the best count * search_count_multiplier pixels,
+            # sample from them uniformly
+            # (here best = smallest variance)
+            _, sorted_confidence_indices = torch.sort(values[2])
+            if self.search_count_multiplier is not None:
+                search_count = min(int(count * self.search_count_multiplier), k)
+            elif self.search_proportion is not None:
+                search_count = min(max(int(k * self.search_proportion), count), k)
+            else:
+                search_count = min(count, k)
+            sample_from_top = random.sample(range(search_count), count)
+            index_sample = sorted_confidence_indices[:search_count][sample_from_top]
+        return index_sample
+
+    def _produce_labels_and_results(self, instance) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance, with confidences
+
+        Args:
+            instance (Instances): an instance of `DensePoseChartPredictorOutputWithConfidences`
+
+        Return:
+            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
+            dp_result (torch.Tensor): shape [3, H, W], DensePose results u and v
+                stacked with the confidence channel
+        """
+        converter = ToChartResultConverterWithConfidences
+        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
+        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
+        dp_result = torch.cat(
+            (dp_result, getattr(chart_result, self.confidence_channel)[None].cpu())
+        )
+
+        return labels, dp_result

densepose/data/samplers/densepose_cse_base.py

@@ -0,0 +1,141 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Dict, List, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.config import CfgNode
+from detectron2.structures import Instances
+
+from densepose.converters.base import IntTupleBox
+from densepose.data.utils import get_class_to_mesh_name_mapping
+from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
+from densepose.structures import DensePoseDataRelative
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseCSEBaseSampler(DensePoseBaseSampler):
+    """
+    Base DensePose sampler to produce DensePose data from DensePose predictions.
+    Samples for each class are drawn according to some distribution over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(
+        self,
+        cfg: CfgNode,
+        use_gt_categories: bool,
+        embedder: torch.nn.Module,
+        count_per_class: int = 8,
+    ):
+        """
+        Constructor
+
+        Args:
+          cfg (CfgNode): the config of the model
+          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        super().__init__(count_per_class)
+        self.embedder = embedder
+        self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
+        self.use_gt_categories = use_gt_categories
+
+    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
+        """
+        Sample DensPoseDataRelative from estimation results
+        """
+        if self.use_gt_categories:
+            instance_class = instance.dataset_classes.tolist()[0]
+        else:
+            instance_class = instance.pred_classes.tolist()[0]
+        mesh_name = self.class_to_mesh_name[instance_class]
+
+        annotation = {
+            DensePoseDataRelative.X_KEY: [],
+            DensePoseDataRelative.Y_KEY: [],
+            DensePoseDataRelative.VERTEX_IDS_KEY: [],
+            DensePoseDataRelative.MESH_NAME_KEY: mesh_name,
+        }
+
+        mask, embeddings, other_values = self._produce_mask_and_results(instance, bbox_xywh)
+        indices = torch.nonzero(mask, as_tuple=True)
+        selected_embeddings = embeddings.permute(1, 2, 0)[indices].cpu()
+        values = other_values[:, indices[0], indices[1]]
+        k = values.shape[1]
+
+        count = min(self.count_per_class, k)
+        if count <= 0:
+            return annotation
+
+        index_sample = self._produce_index_sample(values, count)
+        closest_vertices = squared_euclidean_distance_matrix(
+            selected_embeddings[index_sample], self.embedder(mesh_name)
+        )
+        closest_vertices = torch.argmin(closest_vertices, dim=1)
+
+        sampled_y = indices[0][index_sample] + 0.5
+        sampled_x = indices[1][index_sample] + 0.5
+        # prepare / normalize data
+        _, _, w, h = bbox_xywh
+        x = (sampled_x / w * 256.0).cpu().tolist()
+        y = (sampled_y / h * 256.0).cpu().tolist()
+        # extend annotations
+        annotation[DensePoseDataRelative.X_KEY].extend(x)
+        annotation[DensePoseDataRelative.Y_KEY].extend(y)
+        annotation[DensePoseDataRelative.VERTEX_IDS_KEY].extend(closest_vertices.cpu().tolist())
+        return annotation
+
+    def _produce_mask_and_results(
+        self, instance: Instances, bbox_xywh: IntTupleBox
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of `DensePoseEmbeddingPredictorOutput`
+            bbox_xywh (IntTupleBox): the corresponding bounding box
+
+        Return:
+            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
+            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W],
+                DensePose CSE Embeddings
+            other_values (Tuple[torch.Tensor]): a tensor of shape [0, H, W],
+                for potential other values
+        """
+        densepose_output = instance.pred_densepose
+        S = densepose_output.coarse_segm
+        E = densepose_output.embedding
+        _, _, w, h = bbox_xywh
+        embeddings = F.interpolate(E, size=(h, w), mode="bilinear")[0]
+        coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0]
+        mask = coarse_segm_resized.argmax(0) > 0
+        other_values = torch.empty((0, h, w), device=E.device)
+        return mask, embeddings, other_values
+
+    def _resample_mask(self, output: Any) -> torch.Tensor:
+        """
+        Convert DensePose predictor output to segmentation annotation - tensors of size
+        (256, 256) and type `int64`.
+
+        Args:
+            output: DensePose predictor output with the following attributes:
+             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
+               segmentation scores
+        Return:
+            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
+            where S = DensePoseDataRelative.MASK_SIZE
+        """
+        sz = DensePoseDataRelative.MASK_SIZE
+        mask = (
+            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
+            .argmax(dim=1)
+            .long()
+            .squeeze()
+            .cpu()
+        )
+        return mask

densepose/data/samplers/densepose_cse_confidence_based.py

@@ -0,0 +1,121 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from typing import Optional, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.config import CfgNode
+from detectron2.structures import Instances
+
+from densepose.converters.base import IntTupleBox
+
+from .densepose_cse_base import DensePoseCSEBaseSampler
+
+
+class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn using confidence value estimates.
+    """
+
+    def __init__(
+        self,
+        cfg: CfgNode,
+        use_gt_categories: bool,
+        embedder: torch.nn.Module,
+        confidence_channel: str,
+        count_per_class: int = 8,
+        search_count_multiplier: Optional[float] = None,
+        search_proportion: Optional[float] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          cfg (CfgNode): the config of the model
+          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
+          confidence_channel (str): confidence channel to use for sampling;
+            possible values:
+              "coarse_segm_confidence": confidences for coarse segmentation
+            (default: "coarse_segm_confidence")
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category (default: 8)
+          search_count_multiplier (float or None): if not None, the total number
+              of the most confident estimates of a given class to consider is
+              defined as `min(search_count_multiplier * count_per_class, N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_proportion` (default: None)
+          search_proportion (float or None): if not None, the total number of the
+              of the most confident estimates of a given class to consider is
+              defined as `min(max(search_proportion * N, count_per_class), N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_count_multiplier` (default: None)
+        """
+        super().__init__(cfg, use_gt_categories, embedder, count_per_class)
+        self.confidence_channel = confidence_channel
+        self.search_count_multiplier = search_count_multiplier
+        self.search_proportion = search_proportion
+        assert (search_count_multiplier is None) or (search_proportion is None), (
+            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
+            f"and search_proportion (={search_proportion})"
+        )
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a sample of indices to select data based on confidences
+
+        Args:
+            values (torch.Tensor): a tensor of length k that contains confidences
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        if k == count:
+            index_sample = list(range(k))
+        else:
+            # take the best count * search_count_multiplier pixels,
+            # sample from them uniformly
+            # (here best = smallest variance)
+            _, sorted_confidence_indices = torch.sort(values[0])
+            if self.search_count_multiplier is not None:
+                search_count = min(int(count * self.search_count_multiplier), k)
+            elif self.search_proportion is not None:
+                search_count = min(max(int(k * self.search_proportion), count), k)
+            else:
+                search_count = min(count, k)
+            sample_from_top = random.sample(range(search_count), count)
+            index_sample = sorted_confidence_indices[-search_count:][sample_from_top]
+        return index_sample
+
+    def _produce_mask_and_results(
+        self, instance: Instances, bbox_xywh: IntTupleBox
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of
+                `DensePoseEmbeddingPredictorOutputWithConfidences`
+            bbox_xywh (IntTupleBox): the corresponding bounding box
+
+        Return:
+            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
+            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W]
+                DensePose CSE Embeddings
+            other_values: a tensor of shape [1, H, W], DensePose CSE confidence
+        """
+        _, _, w, h = bbox_xywh
+        densepose_output = instance.pred_densepose
+        mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh)
+        other_values = F.interpolate(
+            getattr(densepose_output, self.confidence_channel),
+            size=(h, w),
+            mode="bilinear",
+        )[0].cpu()
+        return mask, embeddings, other_values

densepose/data/samplers/densepose_cse_uniform.py

@@ -0,0 +1,14 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .densepose_cse_base import DensePoseCSEBaseSampler
+from .densepose_uniform import DensePoseUniformSampler
+
+
+class DensePoseCSEUniformSampler(DensePoseCSEBaseSampler, DensePoseUniformSampler):
+    """
+    Uniform Sampler for CSE
+    """
+
+    pass

densepose/data/samplers/densepose_uniform.py

@@ -0,0 +1,43 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+import torch
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseUniformSampler(DensePoseBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn uniformly over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(self, count_per_class: int = 8):
+        """
+        Constructor
+
+        Args:
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        super().__init__(count_per_class)
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a uniform sample of indices to select data
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        return random.sample(range(k), count)

densepose/data/samplers/mask_from_densepose.py

@@ -0,0 +1,30 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from detectron2.structures import BitMasks, Instances
+
+from densepose.converters import ToMaskConverter
+
+
+class MaskFromDensePoseSampler:
+    """
+    Produce mask GT from DensePose predictions
+    This sampler simply converts DensePose predictions to BitMasks
+    that a contain a bool tensor of the size of the input image
+    """
+
+    def __call__(self, instances: Instances) -> BitMasks:
+        """
+        Converts predicted data from `instances` into the GT mask data
+
+        Args:
+            instances (Instances): predicted results, expected to have `pred_densepose` field
+
+        Returns:
+            Boolean Tensor of the size of the input image that has non-zero
+            values at pixels that are estimated to belong to the detected object
+        """
+        return ToMaskConverter.convert(
+            instances.pred_densepose, instances.pred_boxes, instances.image_size
+        )

densepose/data/samplers/prediction_to_gt.py

@@ -0,0 +1,100 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from detectron2.structures import Instances
+
+ModelOutput = Dict[str, Any]
+SampledData = Dict[str, Any]
+
+
+@dataclass
+class _Sampler:
+    """
+    Sampler registry entry that contains:
+     - src (str): source field to sample from (deleted after sampling)
+     - dst (Optional[str]): destination field to sample to, if not None
+     - func (Optional[Callable: Any -> Any]): function that performs sampling,
+         if None, reference copy is performed
+    """
+
+    src: str
+    dst: Optional[str]
+    func: Optional[Callable[[Any], Any]]
+
+
+class PredictionToGroundTruthSampler:
+    """
+    Sampler implementation that converts predictions to GT using registered
+    samplers for different fields of `Instances`.
+    """
+
+    def __init__(self, dataset_name: str = ""):
+        self.dataset_name = dataset_name
+        self._samplers = {}
+        self.register_sampler("pred_boxes", "gt_boxes", None)
+        self.register_sampler("pred_classes", "gt_classes", None)
+        # delete scores
+        self.register_sampler("scores")
+
+    def __call__(self, model_output: List[ModelOutput]) -> List[SampledData]:
+        """
+        Transform model output into ground truth data through sampling
+
+        Args:
+          model_output (Dict[str, Any]): model output
+        Returns:
+          Dict[str, Any]: sampled data
+        """
+        for model_output_i in model_output:
+            instances: Instances = model_output_i["instances"]
+            # transform data in each field
+            for _, sampler in self._samplers.items():
+                if not instances.has(sampler.src) or sampler.dst is None:
+                    continue
+                if sampler.func is None:
+                    instances.set(sampler.dst, instances.get(sampler.src))
+                else:
+                    instances.set(sampler.dst, sampler.func(instances))
+            # delete model output data that was transformed
+            for _, sampler in self._samplers.items():
+                if sampler.src != sampler.dst and instances.has(sampler.src):
+                    instances.remove(sampler.src)
+            model_output_i["dataset"] = self.dataset_name
+        return model_output
+
+    def register_sampler(
+        self,
+        prediction_attr: str,
+        gt_attr: Optional[str] = None,
+        func: Optional[Callable[[Any], Any]] = None,
+    ):
+        """
+        Register sampler for a field
+
+        Args:
+          prediction_attr (str): field to replace with a sampled value
+          gt_attr (Optional[str]): field to store the sampled value to, if not None
+          func (Optional[Callable: Any -> Any]): sampler function
+        """
+        self._samplers[(prediction_attr, gt_attr)] = _Sampler(
+            src=prediction_attr, dst=gt_attr, func=func
+        )
+
+    def remove_sampler(
+        self,
+        prediction_attr: str,
+        gt_attr: Optional[str] = None,
+    ):
+        """
+        Remove sampler for a field
+
+        Args:
+          prediction_attr (str): field to replace with a sampled value
+          gt_attr (Optional[str]): field to store the sampled value to, if not None
+        """
+        assert (prediction_attr, gt_attr) in self._samplers
+        del self._samplers[(prediction_attr, gt_attr)]

densepose/data/transform/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .image import ImageResizeTransform

densepose/data/transform/image.py

@@ -0,0 +1,41 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import torch
+
+
+class ImageResizeTransform:
+    """
+    Transform that resizes images loaded from a dataset
+    (BGR data in NCHW channel order, typically uint8) to a format ready to be
+    consumed by DensePose training (BGR float32 data in NCHW channel order)
+    """
+
+    def __init__(self, min_size: int = 800, max_size: int = 1333):
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            images (torch.Tensor): tensor of size [N, 3, H, W] that contains
+                BGR data (typically in uint8)
+        Returns:
+            images (torch.Tensor): tensor of size [N, 3, H1, W1] where
+                H1 and W1 are chosen to respect the specified min and max sizes
+                and preserve the original aspect ratio, the data channels
+                follow BGR order and the data type is `torch.float32`
+        """
+        # resize with min size
+        images = images.float()
+        min_size = min(images.shape[-2:])
+        max_size = max(images.shape[-2:])
+        scale = min(self.min_size / min_size, self.max_size / max_size)
+        images = torch.nn.functional.interpolate(
+            images,
+            scale_factor=scale,
+            mode="bilinear",
+            align_corners=False,
+        )
+        return images

densepose/data/utils.py

@@ -0,0 +1,40 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import os
+from typing import Dict, Optional
+
+from detectron2.config import CfgNode
+
+
+def is_relative_local_path(path: str) -> bool:
+    path_str = os.fsdecode(path)
+    return ("://" not in path_str) and not os.path.isabs(path)
+
+
+def maybe_prepend_base_path(base_path: Optional[str], path: str):
+    """
+    Prepends the provided path with a base path prefix if:
+    1) base path is not None;
+    2) path is a local path
+    """
+    if base_path is None:
+        return path
+    if is_relative_local_path(path):
+        return os.path.join(base_path, path)
+    return path
+
+
+def get_class_to_mesh_name_mapping(cfg: CfgNode) -> Dict[int, str]:
+    return {
+        int(class_id): mesh_name
+        for class_id, mesh_name in cfg.DATASETS.CLASS_TO_MESH_NAME_MAPPING.items()
+    }
+
+
+def get_category_to_class_mapping(dataset_cfg: CfgNode) -> Dict[str, int]:
+    return {
+        category: int(class_id)
+        for category, class_id in dataset_cfg.CATEGORY_TO_CLASS_MAPPING.items()
+    }

densepose/data/video/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .frame_selector import (
+    FrameSelectionStrategy,
+    RandomKFramesSelector,
+    FirstKFramesSelector,
+    LastKFramesSelector,
+    FrameTsList,
+    FrameSelector,
+)
+
+from .video_keyframe_dataset import (
+    VideoKeyframeDataset,
+    video_list_from_file,
+    list_keyframes,
+    read_keyframes,
+)

densepose/data/video/frame_selector.py

@@ -0,0 +1,89 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from collections.abc import Callable
+from enum import Enum
+from typing import Callable as TCallable
+from typing import List
+
+FrameTsList = List[int]
+FrameSelector = TCallable[[FrameTsList], FrameTsList]
+
+
+class FrameSelectionStrategy(Enum):
+    """
+    Frame selection strategy used with videos:
+     - "random_k": select k random frames
+     - "first_k": select k first frames
+     - "last_k": select k last frames
+     - "all": select all frames
+    """
+
+    # fmt: off
+    RANDOM_K = "random_k"
+    FIRST_K  = "first_k"
+    LAST_K   = "last_k"
+    ALL      = "all"
+    # fmt: on
+
+
+class RandomKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` random frames
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` random frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return random.sample(frame_tss, min(self.k, len(frame_tss)))
+
+
+class FirstKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` first frames
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` first frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return frame_tss[: self.k]
+
+
+class LastKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` last frames from video data
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` last frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return frame_tss[-self.k :]

densepose/data/video/video_keyframe_dataset.py

@@ -0,0 +1,304 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import csv
+import logging
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Union
+import av
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+from .frame_selector import FrameSelector, FrameTsList
+
+FrameList = List[av.frame.Frame]  # pyre-ignore[16]
+FrameTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+def list_keyframes(video_fpath: str, video_stream_idx: int = 0) -> FrameTsList:
+    """
+    Traverses all keyframes of a video file. Returns a list of keyframe
+    timestamps. Timestamps are counts in timebase units.
+
+    Args:
+       video_fpath (str): Video file path
+       video_stream_idx (int): Video stream index (default: 0)
+    Returns:
+       List[int]: list of keyframe timestaps (timestamp is a count in timebase
+           units)
+    """
+    try:
+        with PathManager.open(video_fpath, "rb") as io:
+            # pyre-fixme[16]: Module `av` has no attribute `open`.
+            container = av.open(io, mode="r")
+            stream = container.streams.video[video_stream_idx]
+            keyframes = []
+            pts = -1
+            # Note: even though we request forward seeks for keyframes, sometimes
+            # a keyframe in backwards direction is returned. We introduce tolerance
+            # as a max count of ignored backward seeks
+            tolerance_backward_seeks = 2
+            while True:
+                try:
+                    container.seek(pts + 1, backward=False, any_frame=False, stream=stream)
+                except av.AVError as e:
+                    # the exception occurs when the video length is exceeded,
+                    # we then return whatever data we've already collected
+                    logger = logging.getLogger(__name__)
+                    logger.debug(
+                        f"List keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts + 1}, AV error: {e}"
+                    )
+                    return keyframes
+                except OSError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"List keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts + 1}, OS error: {e}"
+                    )
+                    return []
+                packet = next(container.demux(video=video_stream_idx))
+                if packet.pts is not None and packet.pts <= pts:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Video file {video_fpath}, stream {video_stream_idx}: "
+                        f"bad seek for packet {pts + 1} (got packet {packet.pts}), "
+                        f"tolerance {tolerance_backward_seeks}."
+                    )
+                    tolerance_backward_seeks -= 1
+                    if tolerance_backward_seeks == 0:
+                        return []
+                    pts += 1
+                    continue
+                tolerance_backward_seeks = 2
+                pts = packet.pts
+                if pts is None:
+                    return keyframes
+                if packet.is_keyframe:
+                    keyframes.append(pts)
+            return keyframes
+    except OSError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"List keyframes: Error opening video file container {video_fpath}, " f"OS error: {e}"
+        )
+    except RuntimeError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"List keyframes: Error opening video file container {video_fpath}, "
+            f"Runtime error: {e}"
+        )
+    return []
+
+
+def read_keyframes(
+    video_fpath: str, keyframes: FrameTsList, video_stream_idx: int = 0
+) -> FrameList:  # pyre-ignore[11]
+    """
+    Reads keyframe data from a video file.
+
+    Args:
+        video_fpath (str): Video file path
+        keyframes (List[int]): List of keyframe timestamps (as counts in
+            timebase units to be used in container seek operations)
+        video_stream_idx (int): Video stream index (default: 0)
+    Returns:
+        List[Frame]: list of frames that correspond to the specified timestamps
+    """
+    try:
+        with PathManager.open(video_fpath, "rb") as io:
+            # pyre-fixme[16]: Module `av` has no attribute `open`.
+            container = av.open(io)
+            stream = container.streams.video[video_stream_idx]
+            frames = []
+            for pts in keyframes:
+                try:
+                    container.seek(pts, any_frame=False, stream=stream)
+                    frame = next(container.decode(video=0))
+                    frames.append(frame)
+                except av.AVError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}, AV error: {e}"
+                    )
+                    container.close()
+                    return frames
+                except OSError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}, OS error: {e}"
+                    )
+                    container.close()
+                    return frames
+                except StopIteration:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error decoding frame from {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}"
+                    )
+                    container.close()
+                    return frames
+
+            container.close()
+            return frames
+    except OSError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Read keyframes: Error opening video file container {video_fpath}, OS error: {e}"
+        )
+    except RuntimeError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Read keyframes: Error opening video file container {video_fpath}, Runtime error: {e}"
+        )
+    return []
+
+
+def video_list_from_file(video_list_fpath: str, base_path: Optional[str] = None):
+    """
+    Create a list of paths to video files from a text file.
+
+    Args:
+        video_list_fpath (str): path to a plain text file with the list of videos
+        base_path (str): base path for entries from the video list (default: None)
+    """
+    video_list = []
+    with PathManager.open(video_list_fpath, "r") as io:
+        for line in io:
+            video_list.append(maybe_prepend_base_path(base_path, str(line.strip())))
+    return video_list
+
+
+def read_keyframe_helper_data(fpath: str):
+    """
+    Read keyframe data from a file in CSV format: the header should contain
+    "video_id" and "keyframes" fields. Value specifications are:
+      video_id: int
+      keyframes: list(int)
+    Example of contents:
+      video_id,keyframes
+      2,"[1,11,21,31,41,51,61,71,81]"
+
+    Args:
+        fpath (str): File containing keyframe data
+
+    Return:
+        video_id_to_keyframes (dict: int -> list(int)): for a given video ID it
+          contains a list of keyframes for that video
+    """
+    video_id_to_keyframes = {}
+    try:
+        with PathManager.open(fpath, "r") as io:
+            csv_reader = csv.reader(io)
+            header = next(csv_reader)
+            video_id_idx = header.index("video_id")
+            keyframes_idx = header.index("keyframes")
+            for row in csv_reader:
+                video_id = int(row[video_id_idx])
+                assert (
+                    video_id not in video_id_to_keyframes
+                ), f"Duplicate keyframes entry for video {fpath}"
+                video_id_to_keyframes[video_id] = (
+                    [int(v) for v in row[keyframes_idx][1:-1].split(",")]
+                    if len(row[keyframes_idx]) > 2
+                    else []
+                )
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(f"Error reading keyframe helper data from {fpath}: {e}")
+    return video_id_to_keyframes
+
+
+class VideoKeyframeDataset(Dataset):
+    """
+    Dataset that provides keyframes for a set of videos.
+    """
+
+    _EMPTY_FRAMES = torch.empty((0, 3, 1, 1))
+
+    def __init__(
+        self,
+        video_list: List[str],
+        category_list: Union[str, List[str], None] = None,
+        frame_selector: Optional[FrameSelector] = None,
+        transform: Optional[FrameTransform] = None,
+        keyframe_helper_fpath: Optional[str] = None,
+    ):
+        """
+        Dataset constructor
+
+        Args:
+            video_list (List[str]): list of paths to video files
+            category_list (Union[str, List[str], None]): list of animal categories for each
+                video file. If it is a string, or None, this applies to all videos
+            frame_selector (Callable: KeyFrameList -> KeyFrameList):
+                selects keyframes to process, keyframes are given by
+                packet timestamps in timebase counts. If None, all keyframes
+                are selected (default: None)
+            transform (Callable: torch.Tensor -> torch.Tensor):
+                transforms a batch of RGB images (tensors of size [B, 3, H, W]),
+                returns a tensor of the same size. If None, no transform is
+                applied (default: None)
+
+        """
+        if type(category_list) is list:
+            self.category_list = category_list
+        else:
+            self.category_list = [category_list] * len(video_list)
+        assert len(video_list) == len(
+            self.category_list
+        ), "length of video and category lists must be equal"
+        self.video_list = video_list
+        self.frame_selector = frame_selector
+        self.transform = transform
+        self.keyframe_helper_data = (
+            read_keyframe_helper_data(keyframe_helper_fpath)
+            if keyframe_helper_fpath is not None
+            else None
+        )
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Gets selected keyframes from a given video
+
+        Args:
+            idx (int): video index in the video list file
+        Returns:
+            A dictionary containing two keys:
+                images (torch.Tensor): tensor of size [N, H, W, 3] or of size
+                    defined by the transform that contains keyframes data
+                categories (List[str]): categories of the frames
+        """
+        categories = [self.category_list[idx]]
+        fpath = self.video_list[idx]
+        keyframes = (
+            list_keyframes(fpath)
+            if self.keyframe_helper_data is None or idx not in self.keyframe_helper_data
+            else self.keyframe_helper_data[idx]
+        )
+        transform = self.transform
+        frame_selector = self.frame_selector
+        if not keyframes:
+            return {"images": self._EMPTY_FRAMES, "categories": []}
+        if frame_selector is not None:
+            keyframes = frame_selector(keyframes)
+        frames = read_keyframes(fpath, keyframes)
+        if not frames:
+            return {"images": self._EMPTY_FRAMES, "categories": []}
+        frames = np.stack([frame.to_rgb().to_ndarray() for frame in frames])
+        frames = torch.as_tensor(frames, device=torch.device("cpu"))
+        frames = frames[..., [2, 1, 0]]  # RGB -> BGR
+        frames = frames.permute(0, 3, 1, 2).float()  # NHWC -> NCHW
+        if transform is not None:
+            frames = transform(frames)
+        return {"images": frames, "categories": categories}
+
+    def __len__(self):
+        return len(self.video_list)

densepose/engine/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .trainer import Trainer

densepose/engine/trainer.py

@@ -0,0 +1,260 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+import logging
+import os
+from collections import OrderedDict
+from typing import List, Optional, Union
+import torch
+from torch import nn
+
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import CfgNode
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import (
+    DatasetEvaluator,
+    DatasetEvaluators,
+    inference_on_dataset,
+    print_csv_format,
+)
+from detectron2.solver.build import get_default_optimizer_params, maybe_add_gradient_clipping
+from detectron2.utils import comm
+from detectron2.utils.events import EventWriter, get_event_storage
+
+from densepose import DensePoseDatasetMapperTTA, DensePoseGeneralizedRCNNWithTTA, load_from_cfg
+from densepose.data import (
+    DatasetMapper,
+    build_combined_loader,
+    build_detection_test_loader,
+    build_detection_train_loader,
+    build_inference_based_loaders,
+    has_inference_based_loaders,
+)
+from densepose.evaluation.d2_evaluator_adapter import Detectron2COCOEvaluatorAdapter
+from densepose.evaluation.evaluator import DensePoseCOCOEvaluator, build_densepose_evaluator_storage
+from densepose.modeling.cse import Embedder
+
+
+class SampleCountingLoader:
+    def __init__(self, loader):
+        self.loader = loader
+
+    def __iter__(self):
+        it = iter(self.loader)
+        storage = get_event_storage()
+        while True:
+            try:
+                batch = next(it)
+                num_inst_per_dataset = {}
+                for data in batch:
+                    dataset_name = data["dataset"]
+                    if dataset_name not in num_inst_per_dataset:
+                        num_inst_per_dataset[dataset_name] = 0
+                    num_inst = len(data["instances"])
+                    num_inst_per_dataset[dataset_name] += num_inst
+                for dataset_name in num_inst_per_dataset:
+                    storage.put_scalar(f"batch/{dataset_name}", num_inst_per_dataset[dataset_name])
+                yield batch
+            except StopIteration:
+                break
+
+
+class SampleCountMetricPrinter(EventWriter):
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+
+    def write(self):
+        storage = get_event_storage()
+        batch_stats_strs = []
+        for key, buf in storage.histories().items():
+            if key.startswith("batch/"):
+                batch_stats_strs.append(f"{key} {buf.avg(20)}")
+        self.logger.info(", ".join(batch_stats_strs))
+
+
+class Trainer(DefaultTrainer):
+    @classmethod
+    def extract_embedder_from_model(cls, model: nn.Module) -> Optional[Embedder]:
+        if isinstance(model, nn.parallel.DistributedDataParallel):
+            model = model.module
+        if hasattr(model, "roi_heads") and hasattr(model.roi_heads, "embedder"):
+            return model.roi_heads.embedder
+        return None
+
+    # TODO: the only reason to copy the base class code here is to pass the embedder from
+    # the model to the evaluator; that should be refactored to avoid unnecessary copy-pasting
+    @classmethod
+    def test(
+        cls,
+        cfg: CfgNode,
+        model: nn.Module,
+        evaluators: Optional[Union[DatasetEvaluator, List[DatasetEvaluator]]] = None,
+    ):
+        """
+        Args:
+            cfg (CfgNode):
+            model (nn.Module):
+            evaluators (DatasetEvaluator, list[DatasetEvaluator] or None): if None, will call
+                :meth:`build_evaluator`. Otherwise, must have the same length as
+                ``cfg.DATASETS.TEST``.
+
+        Returns:
+            dict: a dict of result metrics
+        """
+        logger = logging.getLogger(__name__)
+        if isinstance(evaluators, DatasetEvaluator):
+            evaluators = [evaluators]
+        if evaluators is not None:
+            assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
+                len(cfg.DATASETS.TEST), len(evaluators)
+            )
+
+        results = OrderedDict()
+        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
+            data_loader = cls.build_test_loader(cfg, dataset_name)
+            # When evaluators are passed in as arguments,
+            # implicitly assume that evaluators can be created before data_loader.
+            if evaluators is not None:
+                evaluator = evaluators[idx]
+            else:
+                try:
+                    embedder = cls.extract_embedder_from_model(model)
+                    evaluator = cls.build_evaluator(cfg, dataset_name, embedder=embedder)
+                except NotImplementedError:
+                    logger.warn(
+                        "No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
+                        "or implement its `build_evaluator` method."
+                    )
+                    results[dataset_name] = {}
+                    continue
+            if cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE or comm.is_main_process():
+                results_i = inference_on_dataset(model, data_loader, evaluator)
+            else:
+                results_i = {}
+            results[dataset_name] = results_i
+            if comm.is_main_process():
+                assert isinstance(
+                    results_i, dict
+                ), "Evaluator must return a dict on the main process. Got {} instead.".format(
+                    results_i
+                )
+                logger.info("Evaluation results for {} in csv format:".format(dataset_name))
+                print_csv_format(results_i)
+
+        if len(results) == 1:
+            results = list(results.values())[0]
+        return results
+
+    @classmethod
+    def build_evaluator(
+        cls,
+        cfg: CfgNode,
+        dataset_name: str,
+        output_folder: Optional[str] = None,
+        embedder: Optional[Embedder] = None,
+    ) -> DatasetEvaluators:
+        if output_folder is None:
+            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
+        evaluators = []
+        distributed = cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE
+        # Note: we currently use COCO evaluator for both COCO and LVIS datasets
+        # to have compatible metrics. LVIS bbox evaluator could also be used
+        # with an adapter to properly handle filtered / mapped categories
+        # evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
+        # if evaluator_type == "coco":
+        #     evaluators.append(COCOEvaluator(dataset_name, output_dir=output_folder))
+        # elif evaluator_type == "lvis":
+        #     evaluators.append(LVISEvaluator(dataset_name, output_dir=output_folder))
+        evaluators.append(
+            Detectron2COCOEvaluatorAdapter(
+                dataset_name, output_dir=output_folder, distributed=distributed
+            )
+        )
+        if cfg.MODEL.DENSEPOSE_ON:
+            storage = build_densepose_evaluator_storage(cfg, output_folder)
+            evaluators.append(
+                DensePoseCOCOEvaluator(
+                    dataset_name,
+                    distributed,
+                    output_folder,
+                    evaluator_type=cfg.DENSEPOSE_EVALUATION.TYPE,
+                    min_iou_threshold=cfg.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD,
+                    storage=storage,
+                    embedder=embedder,
+                    should_evaluate_mesh_alignment=cfg.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT,
+                    mesh_alignment_mesh_names=cfg.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES,
+                )
+            )
+        return DatasetEvaluators(evaluators)
+
+    @classmethod
+    def build_optimizer(cls, cfg: CfgNode, model: nn.Module):
+        params = get_default_optimizer_params(
+            model,
+            base_lr=cfg.SOLVER.BASE_LR,
+            weight_decay_norm=cfg.SOLVER.WEIGHT_DECAY_NORM,
+            bias_lr_factor=cfg.SOLVER.BIAS_LR_FACTOR,
+            weight_decay_bias=cfg.SOLVER.WEIGHT_DECAY_BIAS,
+            overrides={
+                "features": {
+                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR,
+                },
+                "embeddings": {
+                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR,
+                },
+            },
+        )
+        optimizer = torch.optim.SGD(
+            params,
+            cfg.SOLVER.BASE_LR,
+            momentum=cfg.SOLVER.MOMENTUM,
+            nesterov=cfg.SOLVER.NESTEROV,
+            weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+        )
+        # pyre-fixme[6]: For 2nd param expected `Type[Optimizer]` but got `SGD`.
+        return maybe_add_gradient_clipping(cfg, optimizer)
+
+    @classmethod
+    def build_test_loader(cls, cfg: CfgNode, dataset_name):
+        return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
+
+    @classmethod
+    def build_train_loader(cls, cfg: CfgNode):
+        data_loader = build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
+        if not has_inference_based_loaders(cfg):
+            return data_loader
+        model = cls.build_model(cfg)
+        model.to(cfg.BOOTSTRAP_MODEL.DEVICE)
+        DetectionCheckpointer(model).resume_or_load(cfg.BOOTSTRAP_MODEL.WEIGHTS, resume=False)
+        inference_based_loaders, ratios = build_inference_based_loaders(cfg, model)
+        loaders = [data_loader] + inference_based_loaders
+        ratios = [1.0] + ratios
+        combined_data_loader = build_combined_loader(cfg, loaders, ratios)
+        sample_counting_loader = SampleCountingLoader(combined_data_loader)
+        return sample_counting_loader
+
+    def build_writers(self):
+        writers = super().build_writers()
+        writers.append(SampleCountMetricPrinter())
+        return writers
+
+    @classmethod
+    def test_with_TTA(cls, cfg: CfgNode, model):
+        logger = logging.getLogger("detectron2.trainer")
+        # In the end of training, run an evaluation with TTA
+        # Only support some R-CNN models.
+        logger.info("Running inference with test-time augmentation ...")
+        transform_data = load_from_cfg(cfg)
+        model = DensePoseGeneralizedRCNNWithTTA(
+            cfg, model, transform_data, DensePoseDatasetMapperTTA(cfg)
+        )
+        evaluators = [
+            cls.build_evaluator(
+                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
+            )
+            for name in cfg.DATASETS.TEST
+        ]
+        res = cls.test(cfg, model, evaluators)  # pyre-ignore[6]
+        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
+        return res

densepose/evaluation/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .evaluator import DensePoseCOCOEvaluator