init

.gitignore

@@ -0,0 +1 @@
+*.pyc

README copy.md

@@ -0,0 +1,14 @@
+---
+title: Virtual Try On
+emoji: 👕👔👚
+colorFrom: yellow
+colorTo: red
+sdk: gradio
+sdk_version: 4.24.0
+app_file: app.py
+pinned: false
+license: cc-by-nc-sa-4.0
+short_description: High-fidelity Virtual Try-on
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,312 @@
+import gradio as gr
+import spaces
+from PIL import Image
+from src.tryon_pipeline import StableDiffusionXLInpaintPipeline as TryonPipeline
+from src.unet_hacked_garmnet import UNet2DConditionModel as UNet2DConditionModel_ref
+from src.unet_hacked_tryon import UNet2DConditionModel
+from transformers import (
+    CLIPImageProcessor,
+    CLIPVisionModelWithProjection,
+    CLIPTextModel,
+    CLIPTextModelWithProjection,
+)
+from diffusers import DDPMScheduler,AutoencoderKL
+from typing import List
+
+import torch
+import os
+from transformers import AutoTokenizer
+import numpy as np
+from utils_mask import get_mask_location
+from torchvision import transforms
+import apply_net
+from preprocess.humanparsing.run_parsing import Parsing
+from preprocess.openpose.run_openpose import OpenPose
+from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation
+from torchvision.transforms.functional import to_pil_image
+
+
+def pil_to_binary_mask(pil_image, threshold=0):
+    np_image = np.array(pil_image)
+    grayscale_image = Image.fromarray(np_image).convert("L")
+    binary_mask = np.array(grayscale_image) > threshold
+    mask = np.zeros(binary_mask.shape, dtype=np.uint8)
+    for i in range(binary_mask.shape[0]):
+        for j in range(binary_mask.shape[1]):
+            if binary_mask[i,j] == True :
+                mask[i,j] = 1
+    mask = (mask*255).astype(np.uint8)
+    output_mask = Image.fromarray(mask)
+    return output_mask
+
+
+base_path = 'yisol/IDM-VTON'
+example_path = os.path.join(os.path.dirname(__file__), 'example')
+
+unet = UNet2DConditionModel.from_pretrained(
+    base_path,
+    subfolder="unet",
+    torch_dtype=torch.float16,
+)
+unet.requires_grad_(False)
+tokenizer_one = AutoTokenizer.from_pretrained(
+    base_path,
+    subfolder="tokenizer",
+    revision=None,
+    use_fast=False,
+)
+tokenizer_two = AutoTokenizer.from_pretrained(
+    base_path,
+    subfolder="tokenizer_2",
+    revision=None,
+    use_fast=False,
+)
+noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")
+
+text_encoder_one = CLIPTextModel.from_pretrained(
+    base_path,
+    subfolder="text_encoder",
+    torch_dtype=torch.float16,
+)
+text_encoder_two = CLIPTextModelWithProjection.from_pretrained(
+    base_path,
+    subfolder="text_encoder_2",
+    torch_dtype=torch.float16,
+)
+image_encoder = CLIPVisionModelWithProjection.from_pretrained(
+    base_path,
+    subfolder="image_encoder",
+    torch_dtype=torch.float16,
+    )
+vae = AutoencoderKL.from_pretrained(base_path,
+                                    subfolder="vae",
+                                    torch_dtype=torch.float16,
+)
+
+# "stabilityai/stable-diffusion-xl-base-1.0",
+UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(
+    base_path,
+    subfolder="unet_encoder",
+    torch_dtype=torch.float16,
+)
+
+parsing_model = Parsing(0)
+openpose_model = OpenPose(0)
+
+UNet_Encoder.requires_grad_(False)
+image_encoder.requires_grad_(False)
+vae.requires_grad_(False)
+unet.requires_grad_(False)
+text_encoder_one.requires_grad_(False)
+text_encoder_two.requires_grad_(False)
+tensor_transfrom = transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+    )
+
+pipe = TryonPipeline.from_pretrained(
+        base_path,
+        unet=unet,
+        vae=vae,
+        feature_extractor= CLIPImageProcessor(),
+        text_encoder = text_encoder_one,
+        text_encoder_2 = text_encoder_two,
+        tokenizer = tokenizer_one,
+        tokenizer_2 = tokenizer_two,
+        scheduler = noise_scheduler,
+        image_encoder=image_encoder,
+        torch_dtype=torch.float16,
+)
+pipe.unet_encoder = UNet_Encoder
+
+@spaces.GPU
+def start_tryon(dict,garm_img,garment_des,is_checked,is_checked_crop,denoise_steps,seed):
+    device = "cuda"
+    
+    openpose_model.preprocessor.body_estimation.model.to(device)
+    pipe.to(device)
+    pipe.unet_encoder.to(device)
+
+    garm_img= garm_img.convert("RGB").resize((768,1024))
+    human_img_orig = dict["background"].convert("RGB")    
+    
+    if is_checked_crop:
+        width, height = human_img_orig.size
+        target_width = int(min(width, height * (3 / 4)))
+        target_height = int(min(height, width * (4 / 3)))
+        left = (width - target_width) / 2
+        top = (height - target_height) / 2
+        right = (width + target_width) / 2
+        bottom = (height + target_height) / 2
+        cropped_img = human_img_orig.crop((left, top, right, bottom))
+        crop_size = cropped_img.size
+        human_img = cropped_img.resize((768,1024))
+    else:
+        human_img = human_img_orig.resize((768,1024))
+
+
+    if is_checked:
+        keypoints = openpose_model(human_img.resize((384,512)))
+        model_parse, _ = parsing_model(human_img.resize((384,512)))
+        mask, mask_gray = get_mask_location('hd', "upper_body", model_parse, keypoints)
+        mask = mask.resize((768,1024))
+    else:
+        mask = pil_to_binary_mask(dict['layers'][0].convert("RGB").resize((768, 1024)))
+        # mask = transforms.ToTensor()(mask)
+        # mask = mask.unsqueeze(0)
+    mask_gray = (1-transforms.ToTensor()(mask)) * tensor_transfrom(human_img)
+    mask_gray = to_pil_image((mask_gray+1.0)/2.0)
+
+
+    human_img_arg = _apply_exif_orientation(human_img.resize((384,512)))
+    human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")
+     
+    
+
+    args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
+    # verbosity = getattr(args, "verbosity", None)
+    pose_img = args.func(args,human_img_arg)    
+    pose_img = pose_img[:,:,::-1]    
+    pose_img = Image.fromarray(pose_img).resize((768,1024))
+    
+    with torch.no_grad():
+        # Extract the images
+        with torch.cuda.amp.autocast():
+            with torch.no_grad():
+                prompt = "model is wearing " + garment_des
+                negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
+                with torch.inference_mode():
+                    (
+                        prompt_embeds,
+                        negative_prompt_embeds,
+                        pooled_prompt_embeds,
+                        negative_pooled_prompt_embeds,
+                    ) = pipe.encode_prompt(
+                        prompt,
+                        num_images_per_prompt=1,
+                        do_classifier_free_guidance=True,
+                        negative_prompt=negative_prompt,
+                    )
+                                    
+                    prompt = "a photo of " + garment_des
+                    negative_prompt = "monochrome, lowres, bad anatomy, worst quality, low quality"
+                    if not isinstance(prompt, List):
+                        prompt = [prompt] * 1
+                    if not isinstance(negative_prompt, List):
+                        negative_prompt = [negative_prompt] * 1
+                    with torch.inference_mode():
+                        (
+                            prompt_embeds_c,
+                            _,
+                            _,
+                            _,
+                        ) = pipe.encode_prompt(
+                            prompt,
+                            num_images_per_prompt=1,
+                            do_classifier_free_guidance=False,
+                            negative_prompt=negative_prompt,
+                        )
+
+
+
+                    pose_img =  tensor_transfrom(pose_img).unsqueeze(0).to(device,torch.float16)
+                    garm_tensor =  tensor_transfrom(garm_img).unsqueeze(0).to(device,torch.float16)
+                    generator = torch.Generator(device).manual_seed(seed) if seed is not None else None
+                    images = pipe(
+                        prompt_embeds=prompt_embeds.to(device,torch.float16),
+                        negative_prompt_embeds=negative_prompt_embeds.to(device,torch.float16),
+                        pooled_prompt_embeds=pooled_prompt_embeds.to(device,torch.float16),
+                        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds.to(device,torch.float16),
+                        num_inference_steps=denoise_steps,
+                        generator=generator,
+                        strength = 1.0,
+                        pose_img = pose_img.to(device,torch.float16),
+                        text_embeds_cloth=prompt_embeds_c.to(device,torch.float16),
+                        cloth = garm_tensor.to(device,torch.float16),
+                        mask_image=mask,
+                        image=human_img, 
+                        height=1024,
+                        width=768,
+                        ip_adapter_image = garm_img.resize((768,1024)),
+                        guidance_scale=2.0,
+                    )[0]
+
+    if is_checked_crop:
+        out_img = images[0].resize(crop_size)        
+        human_img_orig.paste(out_img, (int(left), int(top)))    
+        return human_img_orig, mask_gray
+    else:
+        return images[0], mask_gray
+    # return images[0], mask_gray
+
+garm_list = os.listdir(os.path.join(example_path,"cloth"))
+garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]
+
+human_list = os.listdir(os.path.join(example_path,"human"))
+human_list_path = [os.path.join(example_path,"human",human) for human in human_list]
+
+human_ex_list = []
+for ex_human in human_list_path:
+    ex_dict= {}
+    ex_dict['background'] = ex_human
+    ex_dict['layers'] = None
+    ex_dict['composite'] = None
+    human_ex_list.append(ex_dict)
+
+##default human
+
+
+image_blocks = gr.Blocks(theme="Nymbo/Alyx_Theme").queue()
+with image_blocks as demo:
+    gr.HTML("<center><h1>Virtual Try-On</h1></center>")
+    gr.HTML("<center><p>Upload an image of a person and an image of a garment ✨</p></center>")
+    with gr.Row():
+        with gr.Column():
+            imgs = gr.ImageEditor(sources='upload', type="pil", label='Human. Mask with pen or use auto-masking', interactive=True)
+            with gr.Row():
+                is_checked = gr.Checkbox(label="Yes", info="Use auto-generated mask (Takes 5 seconds)",value=True)
+            with gr.Row():
+                is_checked_crop = gr.Checkbox(label="Yes", info="Use auto-crop & resizing",value=False)
+
+            example = gr.Examples(
+                inputs=imgs,
+                examples_per_page=10,
+                examples=human_ex_list
+            )
+
+        with gr.Column():
+            garm_img = gr.Image(label="Garment", sources='upload', type="pil")
+            with gr.Row(elem_id="prompt-container"):
+                with gr.Row():
+                    prompt = gr.Textbox(placeholder="Description of garment ex) Short Sleeve Round Neck T-shirts", show_label=False, elem_id="prompt")
+            example = gr.Examples(
+                inputs=garm_img,
+                examples_per_page=8,
+                examples=garm_list_path)
+        with gr.Column():
+            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+            masked_img = gr.Image(label="Masked image output", elem_id="masked-img",show_share_button=False)
+        with gr.Column():
+            # image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+            image_out = gr.Image(label="Output", elem_id="output-img",show_share_button=False)
+
+
+
+
+    with gr.Column():
+        try_button = gr.Button(value="Try-on")
+        with gr.Accordion(label="Advanced Settings", open=False):
+            with gr.Row():
+                denoise_steps = gr.Number(label="Denoising Steps", minimum=20, maximum=40, value=30, step=1)
+                seed = gr.Number(label="Seed", minimum=-1, maximum=2147483647, step=1, value=42)
+
+
+
+    try_button.click(fn=start_tryon, inputs=[imgs, garm_img, prompt, is_checked,is_checked_crop, denoise_steps, seed], outputs=[image_out,masked_img], api_name='tryon')
+
+            
+
+
+image_blocks.launch()

apply_net.py

@@ -0,0 +1,359 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+import argparse
+import glob
+import logging
+import os
+import sys
+from typing import Any, ClassVar, Dict, List
+import torch
+
+from detectron2.config import CfgNode, get_cfg
+from detectron2.data.detection_utils import read_image
+from detectron2.engine.defaults import DefaultPredictor
+from detectron2.structures.instances import Instances
+from detectron2.utils.logger import setup_logger
+
+from densepose import add_densepose_config
+from densepose.structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
+from densepose.utils.logger import verbosity_to_level
+from densepose.vis.base import CompoundVisualizer
+from densepose.vis.bounding_box import ScoredBoundingBoxVisualizer
+from densepose.vis.densepose_outputs_vertex import (
+    DensePoseOutputsTextureVisualizer,
+    DensePoseOutputsVertexVisualizer,
+    get_texture_atlases,
+)
+from densepose.vis.densepose_results import (
+    DensePoseResultsContourVisualizer,
+    DensePoseResultsFineSegmentationVisualizer,
+    DensePoseResultsUVisualizer,
+    DensePoseResultsVVisualizer,
+)
+from densepose.vis.densepose_results_textures import (
+    DensePoseResultsVisualizerWithTexture,
+    get_texture_atlas,
+)
+from densepose.vis.extractor import (
+    CompoundExtractor,
+    DensePoseOutputsExtractor,
+    DensePoseResultExtractor,
+    create_extractor,
+)
+
+DOC = """Apply Net - a tool to print / visualize DensePose results
+"""
+
+LOGGER_NAME = "apply_net"
+logger = logging.getLogger(LOGGER_NAME)
+
+_ACTION_REGISTRY: Dict[str, "Action"] = {}
+
+
+class Action:
+    @classmethod
+    def add_arguments(cls: type, parser: argparse.ArgumentParser):
+        parser.add_argument(
+            "-v",
+            "--verbosity",
+            action="count",
+            help="Verbose mode. Multiple -v options increase the verbosity.",
+        )
+
+
+def register_action(cls: type):
+    """
+    Decorator for action classes to automate action registration
+    """
+    global _ACTION_REGISTRY
+    _ACTION_REGISTRY[cls.COMMAND] = cls
+    return cls
+
+
+class InferenceAction(Action):
+    @classmethod
+    def add_arguments(cls: type, parser: argparse.ArgumentParser):
+        super(InferenceAction, cls).add_arguments(parser)
+        parser.add_argument("cfg", metavar="<config>", help="Config file")
+        parser.add_argument("model", metavar="<model>", help="Model file")
+        parser.add_argument(
+            "--opts",
+            help="Modify config options using the command-line 'KEY VALUE' pairs",
+            default=[],
+            nargs=argparse.REMAINDER,
+        )
+
+    @classmethod
+    def execute(cls: type, args: argparse.Namespace, human_img):
+        logger.info(f"Loading config from {args.cfg}")
+        opts = []
+        cfg = cls.setup_config(args.cfg, args.model, args, opts)
+        logger.info(f"Loading model from {args.model}")
+        predictor = DefaultPredictor(cfg)
+        # logger.info(f"Loading data from {args.input}")
+        # file_list = cls._get_input_file_list(args.input)
+        # if len(file_list) == 0:
+        #     logger.warning(f"No input images for {args.input}")
+        #     return
+        context = cls.create_context(args, cfg)
+        # for file_name in file_list:
+        #     img = read_image(file_name, format="BGR")  # predictor expects BGR image.
+        with torch.no_grad():
+            outputs = predictor(human_img)["instances"]
+            out_pose = cls.execute_on_outputs(context, {"image": human_img}, outputs)
+        cls.postexecute(context)
+        return out_pose
+
+    @classmethod
+    def setup_config(
+        cls: type, config_fpath: str, model_fpath: str, args: argparse.Namespace, opts: List[str]
+    ):
+        cfg = get_cfg()
+        add_densepose_config(cfg)
+        cfg.merge_from_file(config_fpath)
+        cfg.merge_from_list(args.opts)
+        if opts:
+            cfg.merge_from_list(opts)
+        cfg.MODEL.WEIGHTS = model_fpath
+        cfg.freeze()
+        return cfg
+
+    @classmethod
+    def _get_input_file_list(cls: type, input_spec: str):
+        if os.path.isdir(input_spec):
+            file_list = [
+                os.path.join(input_spec, fname)
+                for fname in os.listdir(input_spec)
+                if os.path.isfile(os.path.join(input_spec, fname))
+            ]
+        elif os.path.isfile(input_spec):
+            file_list = [input_spec]
+        else:
+            file_list = glob.glob(input_spec)
+        return file_list
+
+
+@register_action
+class DumpAction(InferenceAction):
+    """
+    Dump action that outputs results to a pickle file
+    """
+
+    COMMAND: ClassVar[str] = "dump"
+
+    @classmethod
+    def add_parser(cls: type, subparsers: argparse._SubParsersAction):
+        parser = subparsers.add_parser(cls.COMMAND, help="Dump model outputs to a file.")
+        cls.add_arguments(parser)
+        parser.set_defaults(func=cls.execute)
+
+    @classmethod
+    def add_arguments(cls: type, parser: argparse.ArgumentParser):
+        super(DumpAction, cls).add_arguments(parser)
+        parser.add_argument(
+            "--output",
+            metavar="<dump_file>",
+            default="results.pkl",
+            help="File name to save dump to",
+        )
+
+    @classmethod
+    def execute_on_outputs(
+        cls: type, context: Dict[str, Any], entry: Dict[str, Any], outputs: Instances
+    ):
+        image_fpath = entry["file_name"]
+        logger.info(f"Processing {image_fpath}")
+        result = {"file_name": image_fpath}
+        if outputs.has("scores"):
+            result["scores"] = outputs.get("scores").cpu()
+        if outputs.has("pred_boxes"):
+            result["pred_boxes_XYXY"] = outputs.get("pred_boxes").tensor.cpu()
+            if outputs.has("pred_densepose"):
+                if isinstance(outputs.pred_densepose, DensePoseChartPredictorOutput):
+                    extractor = DensePoseResultExtractor()
+                elif isinstance(outputs.pred_densepose, DensePoseEmbeddingPredictorOutput):
+                    extractor = DensePoseOutputsExtractor()
+                result["pred_densepose"] = extractor(outputs)[0]
+        context["results"].append(result)
+
+    @classmethod
+    def create_context(cls: type, args: argparse.Namespace, cfg: CfgNode):
+        context = {"results": [], "out_fname": args.output}
+        return context
+
+    @classmethod
+    def postexecute(cls: type, context: Dict[str, Any]):
+        out_fname = context["out_fname"]
+        out_dir = os.path.dirname(out_fname)
+        if len(out_dir) > 0 and not os.path.exists(out_dir):
+            os.makedirs(out_dir)
+        with open(out_fname, "wb") as hFile:
+            torch.save(context["results"], hFile)
+            logger.info(f"Output saved to {out_fname}")
+
+
+@register_action
+class ShowAction(InferenceAction):
+    """
+    Show action that visualizes selected entries on an image
+    """
+
+    COMMAND: ClassVar[str] = "show"
+    VISUALIZERS: ClassVar[Dict[str, object]] = {
+        "dp_contour": DensePoseResultsContourVisualizer,
+        "dp_segm": DensePoseResultsFineSegmentationVisualizer,
+        "dp_u": DensePoseResultsUVisualizer,
+        "dp_v": DensePoseResultsVVisualizer,
+        "dp_iuv_texture": DensePoseResultsVisualizerWithTexture,
+        "dp_cse_texture": DensePoseOutputsTextureVisualizer,
+        "dp_vertex": DensePoseOutputsVertexVisualizer,
+        "bbox": ScoredBoundingBoxVisualizer,
+    }
+
+    @classmethod
+    def add_parser(cls: type, subparsers: argparse._SubParsersAction):
+        parser = subparsers.add_parser(cls.COMMAND, help="Visualize selected entries")
+        cls.add_arguments(parser)
+        parser.set_defaults(func=cls.execute)
+
+    @classmethod
+    def add_arguments(cls: type, parser: argparse.ArgumentParser):
+        super(ShowAction, cls).add_arguments(parser)
+        parser.add_argument(
+            "visualizations",
+            metavar="<visualizations>",
+            help="Comma separated list of visualizations, possible values: "
+            "[{}]".format(",".join(sorted(cls.VISUALIZERS.keys()))),
+        )
+        parser.add_argument(
+            "--min_score",
+            metavar="<score>",
+            default=0.8,
+            type=float,
+            help="Minimum detection score to visualize",
+        )
+        parser.add_argument(
+            "--nms_thresh", metavar="<threshold>", default=None, type=float, help="NMS threshold"
+        )
+        parser.add_argument(
+            "--texture_atlas",
+            metavar="<texture_atlas>",
+            default=None,
+            help="Texture atlas file (for IUV texture transfer)",
+        )
+        parser.add_argument(
+            "--texture_atlases_map",
+            metavar="<texture_atlases_map>",
+            default=None,
+            help="JSON string of a dict containing texture atlas files for each mesh",
+        )
+        parser.add_argument(
+            "--output",
+            metavar="<image_file>",
+            default="outputres.png",
+            help="File name to save output to",
+        )
+
+    @classmethod
+    def setup_config(
+        cls: type, config_fpath: str, model_fpath: str, args: argparse.Namespace, opts: List[str]
+    ):
+        opts.append("MODEL.ROI_HEADS.SCORE_THRESH_TEST")
+        opts.append(str(args.min_score))
+        if args.nms_thresh is not None:
+            opts.append("MODEL.ROI_HEADS.NMS_THRESH_TEST")
+            opts.append(str(args.nms_thresh))
+        cfg = super(ShowAction, cls).setup_config(config_fpath, model_fpath, args, opts)
+        return cfg
+
+    @classmethod
+    def execute_on_outputs(
+        cls: type, context: Dict[str, Any], entry: Dict[str, Any], outputs: Instances
+    ):
+        import cv2
+        import numpy as np
+        visualizer = context["visualizer"]
+        extractor = context["extractor"]
+        # image_fpath = entry["file_name"]
+        # logger.info(f"Processing {image_fpath}")
+        image = cv2.cvtColor(entry["image"], cv2.COLOR_BGR2GRAY)
+        image = np.tile(image[:, :, np.newaxis], [1, 1, 3])
+        data = extractor(outputs)
+        image_vis = visualizer.visualize(image, data)
+
+        return image_vis
+        entry_idx = context["entry_idx"] + 1
+        out_fname = './image-densepose/' + image_fpath.split('/')[-1]
+        out_dir = './image-densepose'
+        out_dir = os.path.dirname(out_fname)
+        if len(out_dir) > 0 and not os.path.exists(out_dir):
+            os.makedirs(out_dir)
+        cv2.imwrite(out_fname, image_vis)
+        logger.info(f"Output saved to {out_fname}")
+        context["entry_idx"] += 1
+
+    @classmethod
+    def postexecute(cls: type, context: Dict[str, Any]):
+        pass
+# python ./apply_net.py show ./configs/densepose_rcnn_R_50_FPN_s1x.yaml https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl /home/alin0222/DressCode/upper_body/images dp_segm -v --opts MODEL.DEVICE cpu
+
+    @classmethod
+    def _get_out_fname(cls: type, entry_idx: int, fname_base: str):
+        base, ext = os.path.splitext(fname_base)
+        return base + ".{0:04d}".format(entry_idx) + ext
+
+    @classmethod
+    def create_context(cls: type, args: argparse.Namespace, cfg: CfgNode) -> Dict[str, Any]:
+        vis_specs = args.visualizations.split(",")
+        visualizers = []
+        extractors = []
+        for vis_spec in vis_specs:
+            texture_atlas = get_texture_atlas(args.texture_atlas)
+            texture_atlases_dict = get_texture_atlases(args.texture_atlases_map)
+            vis = cls.VISUALIZERS[vis_spec](
+                cfg=cfg,
+                texture_atlas=texture_atlas,
+                texture_atlases_dict=texture_atlases_dict,
+            )
+            visualizers.append(vis)
+            extractor = create_extractor(vis)
+            extractors.append(extractor)
+        visualizer = CompoundVisualizer(visualizers)
+        extractor = CompoundExtractor(extractors)
+        context = {
+            "extractor": extractor,
+            "visualizer": visualizer,
+            "out_fname": args.output,
+            "entry_idx": 0,
+        }
+        return context
+
+
+def create_argument_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description=DOC,
+        formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=120),
+    )
+    parser.set_defaults(func=lambda _: parser.print_help(sys.stdout))
+    subparsers = parser.add_subparsers(title="Actions")
+    for _, action in _ACTION_REGISTRY.items():
+        action.add_parser(subparsers)
+    return parser
+
+
+def main():
+    parser = create_argument_parser()
+    args = parser.parse_args()
+    verbosity = getattr(args, "verbosity", None)
+    global logger
+    logger = setup_logger(name=LOGGER_NAME)
+    logger.setLevel(verbosity_to_level(verbosity))
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()
+
+
+# python ./apply_net.py show ./configs/densepose_rcnn_R_50_FPN_s1x.yaml https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_s1x/165712039/model_final_162be9.pkl /home/alin0222/Dresscode/dresses/humanonly dp_segm -v --opts MODEL.DEVICE cuda

requirements.txt

@@ -0,0 +1,23 @@
+transformers==4.36.2
+torch==2.0.1
+torchvision==0.15.2
+torchaudio==2.0.2
+numpy==1.24.4
+scipy==1.10.1
+scikit-image==0.21.0
+opencv-python==4.7.0.72
+pillow==9.4.0
+diffusers==0.25.0
+transformers==4.36.2
+accelerate==0.26.1
+matplotlib==3.7.4
+tqdm==4.64.1
+config==0.5.1
+einops==0.7.0
+onnxruntime==1.16.2
+basicsr
+av
+fvcore
+cloudpickle
+omegaconf
+pycocotools

utils_mask.py

@@ -0,0 +1,167 @@
+import numpy as np
+import cv2
+from PIL import Image, ImageDraw
+
+label_map = {
+    "background": 0,
+    "hat": 1,
+    "hair": 2,
+    "sunglasses": 3,
+    "upper_clothes": 4,
+    "skirt": 5,
+    "pants": 6,
+    "dress": 7,
+    "belt": 8,
+    "left_shoe": 9,
+    "right_shoe": 10,
+    "head": 11,
+    "left_leg": 12,
+    "right_leg": 13,
+    "left_arm": 14,
+    "right_arm": 15,
+    "bag": 16,
+    "scarf": 17,
+}
+
+def extend_arm_mask(wrist, elbow, scale):
+  wrist = elbow + scale * (wrist - elbow)
+  return wrist
+
+def hole_fill(img):
+    img = np.pad(img[1:-1, 1:-1], pad_width = 1, mode = 'constant', constant_values=0)
+    img_copy = img.copy()
+    mask = np.zeros((img.shape[0] + 2, img.shape[1] + 2), dtype=np.uint8)
+
+    cv2.floodFill(img, mask, (0, 0), 255)
+    img_inverse = cv2.bitwise_not(img)
+    dst = cv2.bitwise_or(img_copy, img_inverse)
+    return dst
+
+def refine_mask(mask):
+    contours, hierarchy = cv2.findContours(mask.astype(np.uint8),
+                                           cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
+    area = []
+    for j in range(len(contours)):
+        a_d = cv2.contourArea(contours[j], True)
+        area.append(abs(a_d))
+    refine_mask = np.zeros_like(mask).astype(np.uint8)
+    if len(area) != 0:
+        i = area.index(max(area))
+        cv2.drawContours(refine_mask, contours, i, color=255, thickness=-1)
+
+    return refine_mask
+
+def get_mask_location(model_type, category, model_parse: Image.Image, keypoint: dict, width=384,height=512):
+    im_parse = model_parse.resize((width, height), Image.NEAREST)
+    parse_array = np.array(im_parse)
+
+    if model_type == 'hd':
+        arm_width = 60
+    elif model_type == 'dc':
+        arm_width = 45
+    else:
+        raise ValueError("model_type must be \'hd\' or \'dc\'!")
+
+    parse_head = (parse_array == 1).astype(np.float32) + \
+                 (parse_array == 3).astype(np.float32) + \
+                 (parse_array == 11).astype(np.float32)
+
+    parser_mask_fixed = (parse_array == label_map["left_shoe"]).astype(np.float32) + \
+                        (parse_array == label_map["right_shoe"]).astype(np.float32) + \
+                        (parse_array == label_map["hat"]).astype(np.float32) + \
+                        (parse_array == label_map["sunglasses"]).astype(np.float32) + \
+                        (parse_array == label_map["bag"]).astype(np.float32)
+
+    parser_mask_changeable = (parse_array == label_map["background"]).astype(np.float32)
+
+    arms_left = (parse_array == 14).astype(np.float32)
+    arms_right = (parse_array == 15).astype(np.float32)
+
+    if category == 'dresses':
+        parse_mask = (parse_array == 7).astype(np.float32) + \
+                     (parse_array == 4).astype(np.float32) + \
+                     (parse_array == 5).astype(np.float32) + \
+                     (parse_array == 6).astype(np.float32)
+
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+
+    elif category == 'upper_body':
+        parse_mask = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
+        parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \
+                                        (parse_array == label_map["pants"]).astype(np.float32)
+        parser_mask_fixed += parser_mask_fixed_lower_cloth
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+    elif category == 'lower_body':
+        parse_mask = (parse_array == 6).astype(np.float32) + \
+                     (parse_array == 12).astype(np.float32) + \
+                     (parse_array == 13).astype(np.float32) + \
+                     (parse_array == 5).astype(np.float32)
+        parser_mask_fixed += (parse_array == label_map["upper_clothes"]).astype(np.float32) + \
+                             (parse_array == 14).astype(np.float32) + \
+                             (parse_array == 15).astype(np.float32)
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+    else:
+        raise NotImplementedError
+
+    # Load pose points
+    pose_data = keypoint["pose_keypoints_2d"]
+    pose_data = np.array(pose_data)
+    pose_data = pose_data.reshape((-1, 2))
+
+    im_arms_left = Image.new('L', (width, height))
+    im_arms_right = Image.new('L', (width, height))
+    arms_draw_left = ImageDraw.Draw(im_arms_left)
+    arms_draw_right = ImageDraw.Draw(im_arms_right)
+    if category == 'dresses' or category == 'upper_body':
+        shoulder_right = np.multiply(tuple(pose_data[2][:2]), height / 512.0)
+        shoulder_left = np.multiply(tuple(pose_data[5][:2]), height / 512.0)
+        elbow_right = np.multiply(tuple(pose_data[3][:2]), height / 512.0)
+        elbow_left = np.multiply(tuple(pose_data[6][:2]), height / 512.0)
+        wrist_right = np.multiply(tuple(pose_data[4][:2]), height / 512.0)
+        wrist_left = np.multiply(tuple(pose_data[7][:2]), height / 512.0)
+        ARM_LINE_WIDTH = int(arm_width / 512 * height)
+        size_left = [shoulder_left[0] - ARM_LINE_WIDTH // 2, shoulder_left[1] - ARM_LINE_WIDTH // 2, shoulder_left[0] + ARM_LINE_WIDTH // 2, shoulder_left[1] + ARM_LINE_WIDTH // 2]
+        size_right = [shoulder_right[0] - ARM_LINE_WIDTH // 2, shoulder_right[1] - ARM_LINE_WIDTH // 2, shoulder_right[0] + ARM_LINE_WIDTH // 2,
+                      shoulder_right[1] + ARM_LINE_WIDTH // 2]
+        
+
+        if wrist_right[0] <= 1. and wrist_right[1] <= 1.:
+            im_arms_right = arms_right
+        else:
+            wrist_right = extend_arm_mask(wrist_right, elbow_right, 1.2)
+            arms_draw_right.line(np.concatenate((shoulder_right, elbow_right, wrist_right)).astype(np.uint16).tolist(), 'white', ARM_LINE_WIDTH, 'curve')
+            arms_draw_right.arc(size_right, 0, 360, 'white', ARM_LINE_WIDTH // 2)
+
+        if wrist_left[0] <= 1. and wrist_left[1] <= 1.:
+            im_arms_left = arms_left
+        else:
+            wrist_left = extend_arm_mask(wrist_left, elbow_left, 1.2)
+            arms_draw_left.line(np.concatenate((wrist_left, elbow_left, shoulder_left)).astype(np.uint16).tolist(), 'white', ARM_LINE_WIDTH, 'curve')
+            arms_draw_left.arc(size_left, 0, 360, 'white', ARM_LINE_WIDTH // 2)
+
+        hands_left = np.logical_and(np.logical_not(im_arms_left), arms_left)
+        hands_right = np.logical_and(np.logical_not(im_arms_right), arms_right)
+        parser_mask_fixed += hands_left + hands_right
+
+    parser_mask_fixed = np.logical_or(parser_mask_fixed, parse_head)
+    parse_mask = cv2.dilate(parse_mask, np.ones((5, 5), np.uint16), iterations=5)
+    if category == 'dresses' or category == 'upper_body':
+        neck_mask = (parse_array == 18).astype(np.float32)
+        neck_mask = cv2.dilate(neck_mask, np.ones((5, 5), np.uint16), iterations=1)
+        neck_mask = np.logical_and(neck_mask, np.logical_not(parse_head))
+        parse_mask = np.logical_or(parse_mask, neck_mask)
+        arm_mask = cv2.dilate(np.logical_or(im_arms_left, im_arms_right).astype('float32'), np.ones((5, 5), np.uint16), iterations=4)
+        parse_mask += np.logical_or(parse_mask, arm_mask)
+
+    parse_mask = np.logical_and(parser_mask_changeable, np.logical_not(parse_mask))
+
+    parse_mask_total = np.logical_or(parse_mask, parser_mask_fixed)
+    inpaint_mask = 1 - parse_mask_total
+    img = np.where(inpaint_mask, 255, 0)
+    dst = hole_fill(img.astype(np.uint8))
+    dst = refine_mask(dst)
+    inpaint_mask = dst / 255 * 1
+    mask = Image.fromarray(inpaint_mask.astype(np.uint8) * 255)
+    mask_gray = Image.fromarray(inpaint_mask.astype(np.uint8) * 127)
+
+    return mask, mask_gray