Update warp_design_on_dress.py

warp_design_on_dress.py

@@ -4,7 +4,7 @@ import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
 from networks import GMM, UnetGenerator, load_checkpoint, Options
-
+from preprocessing import pad_to_22_channels
 
 def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output_dir):
     os.makedirs(output_dir, exist_ok=True)
@@ -13,7 +13,8 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
     im_h, im_w = 256, 192
     tf = transforms.Compose([
         transforms.Resize((im_h, im_w)),
-        transforms.ToTensor()
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Added normalization
     ])
 
     dress_img = Image.open(dress_path).convert("RGB")
@@ -28,11 +29,15 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
     
     opt = Options()
     gmm = GMM(opt)
-    load_checkpoint(gmm, gmm_ckpt, strict = False)
+    load_checkpoint(gmm, gmm_ckpt, strict=False)
     gmm.cpu().eval()
 
+    # Convert agnostic to 22 channels before passing to GMM
+    agnostic_22ch = pad_to_22_channels(agnostic)
+    design_mask_22ch = pad_to_22_channels(design_mask)
+
     with torch.no_grad():
-        grid, _ = gmm(agnostic, design_mask)
+        grid, _ = gmm(agnostic_22ch, design_mask_22ch)  # Use padded inputs
         warped_design = F.grid_sample(design_tensor, grid, padding_mode='border')
         warped_mask = F.grid_sample(design_mask, grid, padding_mode='zeros')
 
@@ -52,8 +57,8 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
         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_img = (tryon.squeeze().permute(1, 2, 0).cpu().numpy() + 1) * 127.5  # Denormalize
+        out_img = out_img.clip(0, 255).astype("uint8")
         out_pil = Image.fromarray(out_img)
 
         output_path = os.path.join(output_dir, "tryon.jpg")