Create warp_design_on_dress.py

warp_design_on_dress.py

@@ -0,0 +1,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).cuda()
+    design_tensor = tf(design_img).unsqueeze(0).cuda()
+    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.cuda().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