Update app.py

app.py

@@ -1,67 +1,71 @@
 import gradio as gr
 import numpy as np
-from PIL import Image, ImageOps
+from PIL import Image
+import torch
 import cv2
+from torchvision import transforms
+from u2net import U2NET  
+from tps_warp import apply_tps_warp  
+
+model = U2NET()
+model.load_state_dict(torch.load("u2net.pth", map_location=torch.device('cpu')))
+model.eval()
+
+def get_dress_mask(image_np):
+    """Segment the dress using U^2-Net."""
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Resize((320, 320))
+    ])
+    
+    image = Image.fromarray(image_np).convert("RGB")
+    input_tensor = transform(image).unsqueeze(0)
+    
+    with torch.no_grad():
+        output = model(input_tensor)[0][0].squeeze().cpu().numpy()
+    
+    mask = (output > 0.5).astype(np.uint8)
+    return mask
 
 def detect_dress_and_warp(hoodie_img, design_img, warp_strength, scale):
     if hoodie_img is None or design_img is None:
         return None
     
-    # Open images
     hoodie = Image.open(hoodie_img).convert("RGBA")
     design = Image.open(design_img).convert("RGBA")
+    hoodie_np = np.array(hoodie)
     
-    # Convert hoodie image to grayscale for dress detection
-    hoodie_cv = cv2.cvtColor(np.array(hoodie), cv2.COLOR_RGBA2GRAY)
-    _, thresh = cv2.threshold(hoodie_cv, 200, 255, cv2.THRESH_BINARY_INV)
-    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    
-    if contours:
-        largest_contour = max(contours, key=cv2.contourArea)
-        x, y, w, h = cv2.boundingRect(largest_contour)
-        center_x, center_y = x + w // 2, y + h // 2
-    else:
-        center_x, center_y = hoodie.size[0] // 2, hoodie.size[1] // 2  # Default center placement
+    dress_mask = get_dress_mask(hoodie_np)
+    if dress_mask is None:
+        return hoodie  
+
+    y, x = np.where(dress_mask > 0)
+    if len(x) == 0 or len(y) == 0:
+        return hoodie  
+
+    x_min, x_max, y_min, y_max = x.min(), x.max(), y.min(), y.max()
+    center_x, center_y = (x_min + x_max) // 2, (y_min + y_max) // 2
+    w, h = x_max - x_min, y_max - y_min
     
-    # Resize design to fit within the detected dress
     new_w, new_h = int(w * scale), int(h * scale)
     design = design.resize((new_w, new_h))
-    design_cv = cv2.cvtColor(np.array(design), cv2.COLOR_RGBA2BGRA)
+    design_np = np.array(design)
     
-    # Generate normal map for fabric folds
-    hoodie_blur = cv2.GaussianBlur(hoodie_cv, (5, 5), 0)
-    grad_x = cv2.Sobel(hoodie_blur, cv2.CV_32F, 1, 0, ksize=5)
-    grad_y = cv2.Sobel(hoodie_blur, cv2.CV_32F, 0, 1, ksize=5)
-    normal_map = np.dstack((grad_x, grad_y, np.ones_like(grad_x)))
-    normal_map /= np.linalg.norm(normal_map, axis=2, keepdims=True)
+    warped_design = apply_tps_warp(design_np, dress_mask, warp_strength)
+    design = Image.fromarray(warped_design)
     
-    # Apply warping effect based on fabric folds
-    flow_map = normal_map[:, :, :2] * warp_strength
-    rows, cols = design_cv.shape[:2]
-    map_x, map_y = np.meshgrid(np.arange(cols), np.arange(rows))
-    map_x = (map_x + flow_map[:, :, 0]).astype(np.float32)
-    map_y = (map_y + flow_map[:, :, 1]).astype(np.float32)
-    warped_design = cv2.remap(design_cv, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
-    design = Image.fromarray(cv2.cvtColor(warped_design, cv2.COLOR_BGRA2RGBA))
-    
-    # Calculate placement at the center of the detected dress
     paste_x = center_x - new_w // 2
     paste_y = center_y - new_h // 2
     
-    # Create a blank image with same size as hoodie
     temp_layer = Image.new("RGBA", hoodie.size, (0, 0, 0, 0))
-    
-    # Paste design onto detected dress center
     temp_layer.paste(design, (paste_x, paste_y), design)
     
-    # Blend the layers
     final_image = Image.alpha_composite(hoodie, temp_layer)
     return final_image
 
 def generate_mockup(hoodie, design, warp_strength=5.0, scale=1.0):
     return detect_dress_and_warp(hoodie, design, warp_strength, scale)
 
-# Gradio UI
 demo = gr.Interface(
     fn=generate_mockup,
     inputs=[
@@ -72,7 +76,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Image(type="pil", label="Generated Mockup"),
     title="Rookus Hoodie Mockup Generator",
-    description="Upload a hoodie mockup and your design to generate a realistic preview with automatic dress detection and warp effect following fabric folds."
+    description="Upload a hoodie mockup and your design to generate a realistic preview with automatic dress detection and warping using TPS transformation."
 )
 
 demo.launch()