app.py

# -*- 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