Added app.py

app.py

@@ -0,0 +1,482 @@
+# -*- coding: utf-8 -*-
+"""With os FASHION-EYE_VITON-HD Integrated Full Model Final.ipynb
+
+Automatically generated by Colaboratory.
+"""
+
+# !rm -rf sample_data
+# !rm -rf fashion-eye-try-on/
+
+import sys
+from threading import Thread
+import gradio as gr
+import torch
+from collections import OrderedDict
+from PIL import Image
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from cloth_segmentation.networks import U2NET
+import gdown
+from os.path import exists, join, basename, splitext
+import subprocess
+import os
+BASE_DIR = "/home/user/app/fashion-eye-try-on"
+
+os.system(
+    f"git clone https://huggingface.co/spaces/sidharthism/fashion-eye-try-on {BASE_DIR}")
+
+# !pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
+# !pip install -r /content/fashion-eye-try-on/requirements.txt
+os.system("pip install torch>=1.6.0 torchvision -f https://download.pytorch.org/whl/cu92/torch_stable.html")
+os.system("pip install opencv-python torchgeometry gdown Pillow")
+
+os.system(f"cd {BASE_DIR}")
+
+# Download and save checkpoints for cloth mask generation
+os.system(f"rm -rf {BASE_DIR}/cloth_segmentation/checkpoints/")
+os.system(
+    f"gdown --id 1mhF3yqd7R-Uje092eypktNl-RoZNuiCJ -O {BASE_DIR}/cloth_segmentation/checkpoints/")
+
+os.system(
+    f"git clone https://github.com/shadow2496/VITON-HD {BASE_DIR}/VITON-HD")
+
+# checkpoints
+os.system(
+    f"gdown 1RM4OthSM6V4r7kWCu8SbPIPY14Oz8B2u -O {BASE_DIR}/VITON-HD/checkpoints/alias_final.pth")
+os.system(
+    f"gdown 1MBHBddaAs7sy8W40jzLmNL83AUh035F1 -O {BASE_DIR}/VITON-HD/checkpoints/gmm_final.pth")
+os.system(
+    f"gdown 1MBHBddaAs7sy8W40jzLmNL83AUh035F1 -O {BASE_DIR}/VITON-HD/checkpoints/gmm_final.pth")
+os.system(
+    f"gdown 17U1sooR3mVIbe8a7rZuFIF3kukPchHfZ -O {BASE_DIR}/VITON-HD/checkpoints/seg_final.pth")
+# test data
+os.system(
+    f"gdown 1ncEHn_6liOot8sgt3A2DOFJBffvx8tW8 -O {BASE_DIR}/VITON-HD/datasets/test_pairs.txt")
+os.system(
+    f"gdown 1ZA2C8yMOprwc0TV4hvrt0X-ljZugrClq -O {BASE_DIR}/VITON-HD/datasets/test.zip")
+
+os.system(
+    f"unzip {BASE_DIR}/VITON-HD/datasets/test.zip -d {BASE_DIR}/VITON-HD/datasets/")
+
+# @title To clear all the already existing test data
+# !rm -rf /content/fashion-eye-try-on/VITON-HD/datasets/test/image
+# !rm -rf /content/fashion-eye-try-on/VITON-HD/datasets/test/image-parse
+# !rm -rf /content/fashion-eye-try-on/VITON-HD/datasets/test/cloth
+# !rm -rf /content/fashion-eye-try-on/VITON-HD/datasets/test/cloth-mask
+# !rm -rf /content/fashion-eye-try-on/VITON-HD/datasets/test/openpose-img
+# !rm -rf /content/fashion-eye-try-on/VITON-HD/datasets/test/openpose-json
+
+"""Paddle"""
+
+os.system(
+    f"git clone https://huggingface.co/spaces/sidharthism/pipeline_paddle {BASE_DIR}/pipeline_paddle")
+
+# Required for paddle and gradio (Jinja2 dependency)
+os.system("pip install paddlepaddle-gpu pymatting")
+os.system(f"pip install -r {BASE_DIR}/pipeline_paddle/requirements.txt")
+
+os.system(f"rm -rf {BASE_DIR}/pipeline_paddle/models")
+if not os.path.exists(f"{BASE_DIR}/pipeline_paddle/models/ppmatting-hrnet_w18-human_1024.pdparams"):
+    if not os.path.exists(f"{BASE_DIR}/pipeline_paddle/models"):
+        os.mkdir(f"{BASE_DIR}/pipeline_paddle/models")
+    os.system(
+        f"wget https://paddleseg.bj.bcebos.com/matting/models/ppmatting-hrnet_w18-human_1024.pdparams -O {BASE_DIR}/pipeline_paddle/models/ppmatting-hrnet_w18-human_1024.pdparams")
+    # !wget "https://bj.bcebos.com/paddleseg/dygraph/hrnet_w18_ssld.tar.gz" -O "/content/fashion-eye-try-on/pipeline_paddle/models/hrnet_w18_ssld.tar.gz"
+
+"""Initialization
+
+Pose estimator - open pose
+"""
+
+# Clone openpose model repo
+# os.system(f"git clone https://github.com/CMU-Perceptual-Computing-Lab/openpose.git {BASE_DIR}/openpose")
+
+
+# @ Building and Installation of openpose model
+
+
+project_name = f"{BASE_DIR}/openpose"
+print(project_name)
+if not exists(project_name):
+    # see: https://github.com/CMU-Perceptual-Computing-Lab/openpose/issues/949
+    # install new CMake becaue of CUDA10
+    os.system(
+        f"wget -q https://cmake.org/files/v3.13/cmake-3.13.0-Linux-x86_64.tar.gz")
+    os.system(
+        f"tar xfz cmake-3.13.0-Linux-x86_64.tar.gz --strip-components=1 -C /usr/local")
+    # clone openpose
+    os.system(
+        f"cd {BASE_DIR} && git clone -q --depth 1 https://github.com/CMU-Perceptual-Computing-Lab/openpose.git")
+    os.system(
+        "sed -i 's/execute_process(COMMAND git checkout master WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}\/3rdparty\/caffe)/execute_process(COMMAND git checkout f019d0dfe86f49d1140961f8c7dec22130c83154 WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}\/3rdparty\/caffe)/g' %s/openpose/CMakeLists.txt" % (BASE_DIR, ))
+    # install system dependencies
+    os.system("apt-get -qq install -y libatlas-base-dev libprotobuf-dev libleveldb-dev libsnappy-dev libhdf5-serial-dev protobuf-compiler libgflags-dev libgoogle-glog-dev liblmdb-dev opencl-headers ocl-icd-opencl-dev libviennacl-dev")
+    # build openpose
+    print("Building openpose ... May take nearly 15 mins to build ...")
+    os.system(f"cd {BASE_DIR}/openpose && rm -rf {BASE_DIR}/openpose/build || true && mkdir {BASE_DIR}/openpose/build && cd {BASE_DIR}/openpose/build && cmake .. && make -j`nproc`")
+    print("Openpose successfully build and installed.")
+    # subprocess.Popen(f"cd {BASE_DIR}/openpose && rm -rf {BASE_DIR}/openpose/build || true && mkdir {BASE_DIR}/openpose/build && cd {BASE_DIR}/openpose/build && cmake .. && make -j`nproc`")
+    # subprocess.call(["cd", f"{BASE_DIR}/openpose"])
+    # subprocess.check_output(["rm", "-rf", f"{BASE_DIR}/openpose/build || true"])
+    # subprocess.check_output(["mkdir", f"{BASE_DIR}/openpose/build"])
+    # subprocess.check_output(["cd", f"{BASE_DIR}/openpose/build"])
+    # subprocess.check_output(["cmake", ".."])
+    # subprocess.check_output(["make","-j`nproc`"])
+
+# !cd {BASE_DIR}/openpose && rm -rf {BASE_DIR}/openpose/build || true && mkdir {BASE_DIR}/openpose/build && cd {BASE_DIR}/openpose/build && cmake .. && make -j`nproc`
+
+"""Self correction human parsing"""
+
+os.system(
+    f"git clone https://github.com/PeikeLi/Self-Correction-Human-Parsing.git {BASE_DIR}/human_parse")
+
+os.system(f"cd {BASE_DIR}/human_parse")
+os.system(f"mkdir {BASE_DIR}/human_parse/checkpoints")
+# !mkdir inputs
+# !mkdir outputs
+
+dataset = 'lip'
+
+
+dataset_url = 'https://drive.google.com/uc?id=1k4dllHpu0bdx38J7H28rVVLpU-kOHmnH'
+output = f'{BASE_DIR}/human_parse/checkpoints/final.pth'
+gdown.download(dataset_url, output, quiet=False)
+
+# For human parse
+os.system("pip install ninja")
+
+"""Preprocessing"""
+
+# png to jpg
+
+
+def convert_to_jpg(path):
+    from PIL import Image
+    import os
+    if os.path.exists(path):
+        cl = Image.open(path)
+        jpg_path = path[:-4] + ".jpg"
+        cl.save(jpg_path)
+
+
+def resize_img(path):
+    from PIL import Image
+    print(path)
+    im = Image.open(path)
+    im = im.resize((768, 1024), Image.BICUBIC)
+    im.save(path)
+
+
+def remove_ipynb_checkpoints():
+    import os
+    os.system(
+        f"rm -rf {BASE_DIR}/VITON-HD/datasets/test/image/.ipynb_checkpoints")
+    os.system(
+        f"rm -rf {BASE_DIR}/VITON-HD/datasets/test/cloth/.ipynb_checkpoints")
+    os.system(
+        f"rm -rf {BASE_DIR}/VITON-HD/datasets/test/cloth-mask/.ipynb_checkpoints")
+
+# os.chdir('/content/fashion-eye-try-on')
+
+
+def preprocess():
+    remove_ipynb_checkpoints()
+    for path in os.listdir(f'{BASE_DIR}/VITON-HD/datasets/test/image/'):
+        resize_img(f'{BASE_DIR}/VITON-HD/datasets/test/image/{path}')
+    for path in os.listdir(f'{BASE_DIR}/VITON-HD/datasets/test/cloth/'):
+        resize_img(f'{BASE_DIR}/VITON-HD/datasets/test/cloth/{path}')
+    # for path in os.listdir('/content/fashion-eye-try-on/VITON-HD/datasets/test/cloth-mask/'):
+    #   resize_img(f'/content/fashion-eye-try-on/VITON-HD/datasets/test/cloth-mask/{path}')
+
+
+"""Paddle - removing background"""
+
+# PPMatting hrnet 1024
+# --fg_estimate True - for higher quality output but slower prediction
+
+
+def upload_remove_background_and_save_person_image(person_img):
+    # !export CUDA_VISIBLE_DEVICES=0
+    person_img = person_img.resize((768, 1024), Image.BICUBIC)
+    if os.path.exists(f"{BASE_DIR}/pipeline_paddle/image/person.jpg"):
+        os.remove(f"{BASE_DIR}/pipeline_paddle/image/person.jpg")
+    person_img.save(f"{BASE_DIR}/pipeline_paddle/image/person.jpg")
+    # resize_img(f'/content/fashion-eye-try-on/pipeline_paddle/image/person.jpg')
+    os.system(f"cd {BASE_DIR}/pipeline_paddle/")
+    os.system(f"python {BASE_DIR}/pipeline_paddle/bg_replace.py \
+        --config {BASE_DIR}/pipeline_paddle/configs/ppmatting/ppmatting-hrnet_w18-human_1024.yml \
+        --model_path {BASE_DIR}/pipeline_paddle/models/ppmatting-hrnet_w18-human_1024.pdparams \
+        --image_path {BASE_DIR}/pipeline_paddle/image/person.jpg \
+        --background 'w' \
+        --save_dir {BASE_DIR}/VITON-HD/datasets/test/image \
+        --fg_estimate True")
+    # --save_dir /content/fashion-eye-try-on/pipeline_paddle/output \
+    try:
+        convert_to_jpg(f"{BASE_DIR}/VITON-HD/datasets/test/image/person.png")
+        # os.remove("/content/fashion-eye-try-on/pipeline_paddle/output/person_alpha.png")
+        os.remove(f"{BASE_DIR}/VITON-HD/datasets/test/image/person_alpha.png")
+        # os.remove("/content/fashion-eye-try-on/pipeline_paddle/output/person_rgba.png")
+        os.remove(f"{BASE_DIR}/VITON-HD/datasets/test/image/person_rgba.png")
+        os.system(f"cd {BASE_DIR}")
+    except Exception as e:
+        print(e)
+        os.system(f"cd {BASE_DIR}")
+
+
+# @title If multiple GPU available,uncomment and try this code
+os.system("export CUDA_VISIBLE_DEVICES=0")
+
+# Openpose pose estimation
+# Ubuntu and Mac
+
+
+def estimate_pose():
+    os.system(f"cd {BASE_DIR}/openpose && ./build/examples/openpose/openpose.bin --image_dir {BASE_DIR}/VITON-HD/datasets/test/image --write_json {BASE_DIR}/VITON-HD/datasets/test/openpose-json/ --display 0 --face --hand --render_pose 0")
+    os.system(f"cd {BASE_DIR}/openpose && ./build/examples/openpose/openpose.bin --image_dir {BASE_DIR}/VITON-HD/datasets/test/image --write_images {BASE_DIR}/VITON-HD/datasets/test/openpose-img/ --display 0 --hand --render_pose 1 --disable_blending true")
+    os.system(f"cd {BASE_DIR}")
+    # !cd /content/fashion-eye-try-on/openpose && ./build/examples/openpose/openpose.bin --image_dir /content/fashion-eye-try-on/pipeline_paddle/output/ --write_images /content/fashion-eye-try-on/openpose_img/ --display 0 --hand --render_pose 1 --disable_blending true
+
+# Run self correction human parser
+# !python3 /content/fashion-eye-try-on/human_parse/simple_extractor.py --dataset 'lip' --model-restore '/content/fashion-eye-try-on/human_parse/checkpoints/final.pth' --input-dir '/content/fashion-eye-try-on/image' --output-dir '/content/fashion-eye-try-on/VITON-HD/datasets/test/image-parse'
+
+
+def generate_human_segmentation_map():
+    # remove_ipynb_checkpoints()
+    os.system(f"python3 {BASE_DIR}/human_parse/simple_extractor.py --dataset 'lip' --model-restore '{BASE_DIR}/human_parse/checkpoints/final.pth' --input-dir '{BASE_DIR}/VITON-HD/datasets/test/image' --output-dir '{BASE_DIR}/VITON-HD/datasets/test/image-parse'")
+
+# model_image = os.listdir('/content/fashion-eye-try-on/VITON-HD/datasets/test/image')
+# cloth_image = os.listdir('/content/fashion-eye-try-on/VITON-HD/datasets/test/cloth')
+# pairs = zip(model_image, cloth_image)
+
+# with open('/content/fashion-eye-try-on/VITON-HD/datasets/test_pairs.txt', 'w') as file:
+#     for model, cloth in pairs:
+#         file.write(f"{model} {cloth}\n")
+
+
+def generate_test_pairs_txt():
+    with open(f"{BASE_DIR}/VITON-HD/datasets/test_pairs.txt", 'w') as file:
+        file.write(f"person.jpg cloth.jpg\n")
+
+# VITON-HD
+# Transfer the cloth to the model
+
+
+def generate_viton_hd():
+    os.system(f"python {BASE_DIR}/VITON-HD/test.py --name output --dataset_list {BASE_DIR}/VITON-HD/datasets/test_pairs.txt  --dataset_dir {BASE_DIR}/VITON-HD/datasets/ --checkpoint_dir {BASE_DIR}/VITON-HD/checkpoints --save_dir {BASE_DIR}/")
+
+
+# To resolve ModuleNotFoundError during imports
+if BASE_DIR not in sys.path:
+    sys.path.append(BASE_DIR)
+    sys.path.append(f"{BASE_DIR}/cloth_segmentation")
+
+
+device = 'cuda' if torch.cuda.is_available() else "cpu"
+
+if device == 'cuda':
+    torch.cuda.empty_cache()
+
+# for hugging face
+# BASE_DIR = "/home/path/app"
+
+image_dir = 'cloth'
+result_dir = 'cloth_mask'
+checkpoint_path = 'cloth_segmentation/checkpoints/cloth_segm_u2net_latest.pth'
+
+
+def load_checkpoint_mgpu(model, checkpoint_path):
+    if not os.path.exists(checkpoint_path):
+        print("----No checkpoints at given path----")
+        return
+    model_state_dict = torch.load(
+        checkpoint_path, map_location=torch.device("cpu"))
+    new_state_dict = OrderedDict()
+    for k, v in model_state_dict.items():
+        name = k[7:]  # remove `module.`
+        new_state_dict[name] = v
+
+    model.load_state_dict(new_state_dict)
+    print("----checkpoints loaded from path: {}----".format(checkpoint_path))
+    return model
+
+
+class Normalize_image(object):
+    """Normalize given tensor into given mean and standard dev
+    Args:
+        mean (float): Desired mean to substract from tensors
+        std (float): Desired std to divide from tensors
+    """
+
+    def __init__(self, mean, std):
+        assert isinstance(mean, (float))
+        if isinstance(mean, float):
+            self.mean = mean
+
+        if isinstance(std, float):
+            self.std = std
+
+        self.normalize_1 = transforms.Normalize(self.mean, self.std)
+        self.normalize_3 = transforms.Normalize(
+            [self.mean] * 3, [self.std] * 3)
+        self.normalize_18 = transforms.Normalize(
+            [self.mean] * 18, [self.std] * 18)
+
+    def __call__(self, image_tensor):
+        if image_tensor.shape[0] == 1:
+            return self.normalize_1(image_tensor)
+
+        elif image_tensor.shape[0] == 3:
+            return self.normalize_3(image_tensor)
+
+        elif image_tensor.shape[0] == 18:
+            return self.normalize_18(image_tensor)
+
+        else:
+            assert "Please set proper channels! Normlization implemented only for 1, 3 and 18"
+
+
+def get_palette(num_cls):
+    """ Returns the color map for visualizing the segmentation mask.
+    Args:
+        num_cls: Number of classes
+    Returns:
+        The color map
+    """
+    n = num_cls
+    palette = [0] * (n * 3)
+    for j in range(0, n):
+        lab = j
+        palette[j * 3 + 0] = 0
+        palette[j * 3 + 1] = 0
+        palette[j * 3 + 2] = 0
+        i = 0
+        while lab:
+            palette[j * 3 + 0] = 255
+            palette[j * 3 + 1] = 255
+            palette[j * 3 + 2] = 255
+            # palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
+            # palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
+            # palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
+            i += 1
+            lab >>= 3
+    return palette
+
+
+def generate_cloth_mask(img_dir, output_dir, chkpt_dir):
+    global image_dir
+    global result_dir
+    global checkpoint_path
+    image_dir = img_dir
+    result_dir = output_dir
+    checkpoint_path = chkpt_dir
+    transforms_list = []
+    transforms_list += [transforms.ToTensor()]
+    transforms_list += [Normalize_image(0.5, 0.5)]
+    transform_rgb = transforms.Compose(transforms_list)
+
+    net = U2NET(in_ch=3, out_ch=4)
+    with torch.no_grad():
+        net = load_checkpoint_mgpu(net, checkpoint_path)
+        net = net.to(device)
+        net = net.eval()
+
+        palette = get_palette(4)
+
+        images_list = sorted(os.listdir(image_dir))
+        for image_name in images_list:
+            img = Image.open(os.path.join(
+                image_dir, image_name)).convert('RGB')
+            img_size = img.size
+            img = img.resize((768, 768), Image.BICUBIC)
+            image_tensor = transform_rgb(img)
+            image_tensor = torch.unsqueeze(image_tensor, 0)
+
+            output_tensor = net(image_tensor.to(device))
+            output_tensor = F.log_softmax(output_tensor[0], dim=1)
+            output_tensor = torch.max(output_tensor, dim=1, keepdim=True)[1]
+            output_tensor = torch.squeeze(output_tensor, dim=0)
+            output_tensor = torch.squeeze(output_tensor, dim=0)
+            output_arr = output_tensor.cpu().numpy()
+
+            output_img = Image.fromarray(output_arr.astype('uint8'), mode='L')
+            output_img = output_img.resize(img_size, Image.BICUBIC)
+
+            output_img.putpalette(palette)
+            output_img = output_img.convert('L')
+            output_img.save(os.path.join(result_dir, image_name[:-4]+'.jpg'))
+
+
+os.system(f"cd {BASE_DIR}")
+
+
+def upload_resize_generate_cloth_mask_and_move_to_viton_hd_test_inputs(cloth_img):
+    os.system(f"cd {BASE_DIR}")
+    cloth_img = cloth_img.resize((768, 1024), Image.BICUBIC)
+    cloth_img.save(f"{BASE_DIR}/cloth/cloth.jpg")
+    cloth_img.save(f"{BASE_DIR}/VITON-HD/datasets/test/cloth/cloth.jpg")
+    try:
+        generate_cloth_mask(f"{BASE_DIR}/cloth", f"{BASE_DIR}/cloth_mask",
+                            f"{BASE_DIR}/cloth_segmentation/checkpoints/cloth_segm_u2net_latest.pth")
+        cloth_mask_img = Image.open(f"{BASE_DIR}/cloth_mask/cloth.jpg")
+        cloth_mask_img.save(
+            f"{BASE_DIR}/VITON-HD/datasets/test/cloth-mask/cloth.jpg")
+    except Exception as e:
+        print(e)
+
+
+# Gradio
+os.system("pip install gradio")
+
+# import cv2
+IMAGEPATH = '/content/fashion-eye-try-on/VITON-HD/datasets/test/image'
+CLOTHPATH = '/content/fashion-eye-try-on/VITON-HD/datasets/test/cloth'
+CLOTHMASKPATH = '/content/fashion-eye-try-on/VITON-HD/datasets/test/image'
+
+
+def fashion_eye_tryon(person_img, cloth_img):
+    result_img = person_img
+    # img.save(IMAGEPATH + "person.jpg")
+    # dress.save(CLOTHPATH + "cloth.jpg")
+
+    # txt = open("/content/VITON-HD/datasets/test_pairs.txt", "a")
+    # txt.write("person_img.jpg dress_img.jpg\n")
+    # txt.close()
+    # # result
+    # print(person_img.info, cloth_img.info)
+    # p_t1 = Thread(target=upload_remove_background_and_save_person_image, args=(person_img, ))
+    # c_t2 = Thread(target=upload_resize_generate_cloth_mask_and_move_to_viton_hd_test_inputs, args=(cloth_img, ))
+    # p_t1.start()
+    # c_t2.start()
+    # p_t1.join()
+    # c_t2.join()
+    # Estimate pose
+    try:
+        upload_resize_generate_cloth_mask_and_move_to_viton_hd_test_inputs(
+            cloth_img)
+        upload_remove_background_and_save_person_image(person_img)
+        remove_ipynb_checkpoints()
+        estimate_pose()
+        # Generate human parse
+        remove_ipynb_checkpoints()
+        generate_human_segmentation_map()
+        generate_test_pairs_txt()
+        remove_ipynb_checkpoints()
+        generate_viton_hd()
+        for p in ["/content/fashion-eye-try-on/output/person_cloth.jpg", "/content/fashion-eye-try-on/output/person.jpg_cloth.jpg"]:
+            if os.path.exists(p):
+                result_img = Image.open(p)
+    except Exception as e:
+        print(e)
+        return
+    return result_img
+
+
+# res = fashion_eye_tryon("", "")
+# res.show()
+gr.Interface(fn=fashion_eye_tryon,
+             inputs=[gr.Image(type="pil", label="Your image"),
+                     gr.Image(type="pil", label="Dress")],
+             outputs="image"
+             ).launch(debug=True, inbrowser=True, share=True)
+
+# !pip freeze > /content/requirements_final.txt