Update warp_design_on_dress.py

warp_design_on_dress.py

@@ -3,13 +3,15 @@ import torch
 import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image
-from networks import GMM, TOM, load_checkpoint, Options  # Updated imports
+import numpy as np
+from networks import GMM, TOM, 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)
-
-    # Preprocessing with enhanced normalization
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Preprocessing
     im_h, im_w = 256, 192
     tf = transforms.Compose([
         transforms.Resize((im_h, im_w)),
@@ -17,39 +19,28 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ])
 
-    # Load and prepare images with error handling
-    try:
-        dress_img = Image.open(dress_path).convert("RGB")
-        design_img = Image.open(design_path).convert("RGB")
-    except Exception as e:
-        raise ValueError(f"Error loading images: {str(e)}")
-
-    # Convert to tensors
-    dress_tensor = tf(dress_img).unsqueeze(0).cpu()
-    design_tensor = tf(design_img).unsqueeze(0).cpu()
-    design_mask = torch.ones_like(design_tensor[:, :1, :, :])
-
-    # Prepare agnostic (dress image)
-    agnostic = dress_tensor.clone()
+    # Load images
+    dress_img = Image.open(dress_path).convert("RGB")
+    design_img = Image.open(design_path).convert("RGB")
     
-    # Initialize models with proper device handling
-    opt = Options()
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    # Convert to tensors
+    dress_tensor = tf(dress_img).unsqueeze(0).to(device)
+    design_tensor = tf(design_img).unsqueeze(0).to(device)
     
+    # Create design mask (assuming design covers entire area)
+    design_mask = torch.ones_like(design_tensor[:, :1, :, :]).to(device)
+
+    # Prepare agnostic input (dress image)
+    agnostic = pad_to_22_channels(dress_tensor).to(device)
+
     # GMM Processing
+    opt = Options()
     gmm = GMM(opt).to(device)
     load_checkpoint(gmm, gmm_ckpt, strict=False)
     gmm.eval()
 
-    # Convert to required channels and move to device
-    agnostic_22ch = pad_to_22_channels(agnostic).contiguous().to(device)
-    design_mask_22ch = pad_to_22_channels(design_mask).contiguous().to(device)
-    design_tensor = design_tensor.to(device)
-    design_mask = design_mask.to(device)
-
     with torch.no_grad():
-        # Process through GMM with align_corners
-        grid, _ = gmm(agnostic_22ch, design_mask_22ch)
+        grid, _ = gmm(agnostic, design_mask)
         warped_design = F.grid_sample(
             design_tensor, 
             grid, 
@@ -64,43 +55,40 @@ def run_design_warp_on_dress(dress_path, design_path, gmm_ckpt, tom_ckpt, output
         )
 
     # TOM Processing
-    tom = TOM(opt).to(device)  # Using the new TOM class
+    tom = TOM(opt).to(device)
     load_checkpoint(tom, tom_ckpt, strict=False)
     tom.eval()
 
     with torch.no_grad():
-        # Prepare proper 26-channel input
-        # Generate additional features (replace with actual feature extraction if available)
-        gray = agnostic.mean(dim=1, keepdim=True)
-        edges_x = torch.abs(F.conv2d(gray, 
-                           torch.tensor([[[[1,0,-1],[2,0,-2],[1,0,-1]]]], device=device).float()))
-        edges_y = torch.abs(F.conv2d(gray,
-                           torch.tensor([[[[1,2,1],[0,0,0],[-1,-2,-1]]]], device=device).float()))
+        # Create simplified feature inputs
+        gray = dress_tensor.mean(dim=1, keepdim=True)
+        edges = torch.abs(F.conv2d(
+            gray, 
+            torch.tensor([[[[-1,-1,-1],[-1,8,-1],[-1,-1,-1]]], dtype=torch.float32, device=device),
+            padding=1
+        ))
         
-        # Combine all features (3+3+1+19=26)
+        # Combine inputs (agnostic + warped design + warped mask + features)
         tom_input = torch.cat([
-            agnostic,          # 3 channels
-            warped_design,     # 3 channels
-            warped_mask,       # 1 channel
-            gray,              # 1 channel
-            edges_x,           # 1 channel
-            edges_y,           # 1 channel
-            torch.zeros_like(agnostic)[:, :16]  # 16 dummy channels (replace with real features)
-        ], dim=1)
+            dress_tensor,          # 3 channels
+            warped_design,          # 3 channels
+            warped_mask,            # 1 channel
+            gray,                   # 1 channel
+            edges,                  # 1 channel
+            torch.zeros_like(dress_tensor)[:, :17]  # 17 dummy channels
+        ], dim=1)  # Total: 3+3+1+1+1+17 = 26 channels
 
-        # Process through TOM
+        # Generate try-on result
         p_rendered, m_composite = tom(tom_input)
         tryon = warped_design * m_composite + p_rendered * (1 - m_composite)
 
-        # Save output with proper denormalization
-        tryon = tryon.clamp(-1, 1)  # Ensure valid range
-        out_img = (tryon.squeeze().permute(1, 2, 0).cpu().numpy() + 1) * 127.5
-        out_img = out_img.clip(0, 255).astype("uint8")
-        
-        try:
-            out_pil = Image.fromarray(out_img)
-            output_path = os.path.join(output_dir, "tryon.jpg")
-            out_pil.save(output_path)
-            return output_path
-        except Exception as e:
-            raise ValueError(f"Error saving output image: {str(e)}")
+    # Convert to PIL image
+    tryon = tryon.squeeze().detach().cpu()
+    tryon = (tryon.permute(1, 2, 0).numpy() + 1) * 127.5
+    tryon = np.clip(tryon, 0, 255).astype("uint8")
+    out_pil = Image.fromarray(tryon)
+    
+    # Save output
+    output_path = os.path.join(output_dir, "tryon.jpg")
+    out_pil.save(output_path)
+    return output_path