Spaces:
Sleeping
Sleeping
File size: 1,999 Bytes
b61f3f8 4560e58 b61f3f8 08cfba0 b61f3f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import os
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
from networks import GMM, UnetGenerator, load_checkpoint
def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output_dir):
os.makedirs(output_dir, exist_ok=True)
# Preprocessing
im_h, im_w = 256, 192
tf = transforms.Compose([
transforms.Resize((im_h, im_w)),
transforms.ToTensor()
])
dress_img = Image.open(dress_path).convert("RGB")
design_img = Image.open(design_path).convert("RGB")
dress_tensor = tf(dress_img).unsqueeze(0).cpu()
design_tensor = tf(design_img).unsqueeze(0).cpu()
design_mask = torch.ones_like(design_tensor[:, :1, :, :]) # full white mask
# Fake agnostic: use the dress image itself
agnostic = dress_tensor.clone()
# ----- GMM -----
gmm = GMM(opt=None)
load_checkpoint(gmm, gmm_ckpt)
gmm.cuda().eval()
with torch.no_grad():
grid, _ = gmm(agnostic, design_mask)
warped_design = F.grid_sample(design_tensor, grid, padding_mode='border')
warped_mask = F.grid_sample(design_mask, grid, padding_mode='zeros')
# ----- TOM -----
tom = UnetGenerator(26, 4, 6, ngf=64, norm_layer=torch.nn.InstanceNorm2d)
load_checkpoint(tom, tom_ckpt)
tom.cpu().eval()
with torch.no_grad():
tom_input = torch.cat([agnostic, warped_design, warped_mask], 1)
output = tom(tom_input)
p_rendered, m_composite = torch.split(output, 3, 1)
p_rendered = torch.tanh(p_rendered)
m_composite = torch.sigmoid(m_composite)
tryon = warped_design * m_composite + p_rendered * (1 - m_composite)
# Save output
out_img = tryon.squeeze().permute(1, 2, 0).cpu().numpy()
out_img = (out_img * 255).astype("uint8")
out_pil = Image.fromarray(out_img)
output_path = os.path.join(output_dir, "tryon.jpg")
out_pil.save(output_path)
return output_path |