sidharthism's picture
Added sudo to openpose
e2d6896
# -*- 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/
BASE_DIR = "/home/user/app/fashion-eye-try-on"
import os
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
import os
import subprocess
from os.path import exists, join, basename, splitext
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"sudo 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("sudo 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("sudo 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"sudo 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'
import gdown
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}/")
import sys
# 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")
from cloth_segmentation.networks import U2NET
import torchvision.transforms as transforms
import torch.nn.functional as F
import os
from PIL import Image
from collections import OrderedDict
import torch
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}")
from PIL import Image
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 gradio as gr
# import cv2
from PIL import Image
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'
from threading import Thread
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)
# !pip freeze > /content/requirements_final.txt