Update app.py

app.py

@@ -4,71 +4,71 @@ import cv2
 import numpy as np
 from torchvision import transforms
 from PIL import Image
-from transformers import DPTForDepthEstimation, DPTFeatureExtractor
+from transformers import DPTForDepthEstimation, DPTFeatureExtractor, MidasForDepthEstimation, MidasImageProcessor
 
-# Load depth estimation model
-model_name = "Intel/dpt-large"
-feature_extractor = DPTFeatureExtractor.from_pretrained(model_name)
-depth_model = DPTForDepthEstimation.from_pretrained(model_name)
+# Load depth estimation model (MiDaS v3 for better accuracy)
+model_name = "Intel/midas-v3"  # Upgraded model
+processor = MidasImageProcessor.from_pretrained(model_name)
+depth_model = MidasForDepthEstimation.from_pretrained(model_name)
 depth_model.eval()
 
 def estimate_depth(image):
-    """Estimate depth map from image."""
+    """Estimate depth map from image using MiDaS v3."""
     image = image.convert("RGB")
-    image = image.resize((384, 384))  # Resize for model input
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = depth_model(**inputs)
         depth = outputs.predicted_depth.squeeze().cpu().numpy()
-    depth = cv2.resize(depth, (image.width, image.height))  # Resize back to original
+    depth = cv2.resize(depth, (image.width, image.height))
     depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255
     return depth.astype(np.uint8)
 
+def apply_tps_warping(design, depth):
+    """Apply Thin Plate Spline (TPS) warping based on depth."""
+    h, w = depth.shape
+    grid_x, grid_y = np.meshgrid(np.arange(w), np.arange(h))
+    displacement_x = cv2.Sobel(depth, cv2.CV_32F, 1, 0, ksize=5)
+    displacement_y = cv2.Sobel(depth, cv2.CV_32F, 0, 1, ksize=5)
+    displacement_x = cv2.normalize(displacement_x, None, -10, 10, cv2.NORM_MINMAX)
+    displacement_y = cv2.normalize(displacement_y, None, -10, 10, cv2.NORM_MINMAX)
+    map_x = np.clip(grid_x + displacement_x, 0, w - 1).astype(np.float32)
+    map_y = np.clip(grid_y + displacement_y, 0, h - 1).astype(np.float32)
+    warped_design = cv2.remap(design, map_x, map_y, interpolation=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REFLECT)
+    return warped_design
+
 def blend_design(cloth_img, design_img):
-    """Blend design onto clothing naturally with fold adaptation."""
+    """Blend design onto clothing naturally with fold adaptation using TPS warping."""
     cloth_img = cloth_img.convert("RGB")
     design_img = design_img.convert("RGBA")
     cloth_np = np.array(cloth_img)
     design_np = np.array(design_img)
     
-    # Resize design to fit within clothing
+    # Resize design
     h, w, _ = cloth_np.shape
     dh, dw, _ = design_np.shape
-    scale_factor = min(w / dw, h / dh) * 0.4  # Scale to 40% of clothing area
+    scale_factor = min(w / dw, h / dh) * 0.4  
     new_w, new_h = int(dw * scale_factor), int(dh * scale_factor)
     design_np = cv2.resize(design_np, (new_w, new_h), interpolation=cv2.INTER_AREA)
     
-    # Extract alpha channel for transparency
+    # Extract alpha channel
     alpha_channel = design_np[:, :, 3] / 255.0
     design_np = design_np[:, :, :3]
     
-    # Create a blank canvas and paste the resized design at the center
+    # Create placement area
     x_offset = (w - new_w) // 2
-    y_offset = int(h * 0.35)  # Move slightly upward for a natural position
-    
-    for c in range(3):
-        cloth_np[y_offset:y_offset+new_h, x_offset:x_offset+new_w, c] = (
-            cloth_np[y_offset:y_offset+new_h, x_offset:x_offset+new_w, c] * (1 - alpha_channel) + 
-            design_np[:, :, c] * alpha_channel
-        )
+    y_offset = int(h * 0.35)
+    design_canvas = np.zeros_like(cloth_np)
+    design_canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = design_np
     
-    # Estimate depth for fold detection
+    # Estimate depth and apply TPS warping
     depth_map = estimate_depth(cloth_img)
-    depth_map = cv2.resize(depth_map, (w, h))
+    warped_design = apply_tps_warping(design_canvas, depth_map)
     
-    # Generate displacement map based on depth
-    displacement_x = cv2.Sobel(depth_map, cv2.CV_32F, 1, 0, ksize=5)
-    displacement_y = cv2.Sobel(depth_map, cv2.CV_32F, 0, 1, ksize=5)
-    displacement_x = cv2.normalize(displacement_x, None, -5, 5, cv2.NORM_MINMAX)
-    displacement_y = cv2.normalize(displacement_y, None, -5, 5, cv2.NORM_MINMAX)
-    
-    # Warp design using displacement map
-    map_x, map_y = np.meshgrid(np.arange(w), np.arange(h))
-    map_x = np.clip(np.float32(map_x + displacement_x), 0, w - 1)
-    map_y = np.clip(np.float32(map_y + displacement_y), 0, h - 1)
-    warped_cloth = cv2.remap(cloth_np, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
+    # Blend design onto cloth
+    for c in range(3):
+        cloth_np[:, :, c] = (cloth_np[:, :, c] * (1 - alpha_channel) + warped_design[:, :, c] * alpha_channel)
     
-    return Image.fromarray(warped_cloth)
+    return Image.fromarray(cloth_np)
 
 def main(cloth, design):
     return blend_design(cloth, design)
@@ -82,4 +82,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)