Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 from torchvision import transforms
 from PIL import Image
 from transformers import DPTForDepthEstimation, DPTFeatureExtractor
+import torchvision.transforms.functional as F
 
 # Load depth estimation model
 model_name = "Intel/dpt-large"
@@ -30,39 +31,18 @@ def warp_design(cloth_img, design_img):
     cloth_np = np.array(cloth_img)
     design_np = np.array(design_img)
     h, w, _ = cloth_np.shape
-    dh, dw, _ = design_np.shape
-    
-    # Resize design to fit within 70% of the clothing area
-    scale_factor = min(w / dw, h / dh) * 0.7
-    new_w, new_h = int(dw * scale_factor), int(dh * scale_factor)
-    design_np = cv2.resize(design_np, (new_w, new_h))
-    
-    # Create blank canvas with transparent background
-    design_canvas = np.zeros_like(cloth_np, dtype=np.uint8)
-    x_offset = (w - new_w) // 2
-    y_offset = (h - new_h) // 2
-    design_canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = design_np
     
     # Estimate depth map
     depth_map = estimate_depth(cloth_img)
     depth_map = cv2.resize(depth_map, (w, h))
     
-    # 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, -3, 3, cv2.NORM_MINMAX)
-    displacement_y = cv2.normalize(displacement_y, None, -3, 3, cv2.NORM_MINMAX)
-    
-    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_design = cv2.remap(design_canvas, map_x, map_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
+    # Compute optical flow for warping
+    flow = cv2.calcOpticalFlowFarneback(depth_map, depth_map, None, 0.5, 3, 15, 3, 5, 1.2, 0)
+    flow_map = np.column_stack((flow[..., 0] + np.arange(w), flow[..., 1] + np.arange(h)[:, None]))
+    warped_design = cv2.remap(design_np, flow_map, None, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT)
     
-    # Blend images without excessive transparency
-    mask = np.any(warped_design > 0, axis=-1).astype(np.uint8) * 255
-    blended = cloth_np.copy()
-    np.copyto(blended, warped_design, where=(mask[..., None] > 0))
+    # Blending
+    blended = cv2.addWeighted(cloth_np, 0.7, warped_design, 0.3, 0)
     
     return Image.fromarray(blended)
 
@@ -78,4 +58,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch(share=True)