Update app.py

app.py

@@ -1,50 +1,69 @@
-import gradio as gr
 import torch
 import cv2
 import numpy as np
-from torchvision import transforms
+import gradio as gr
 from PIL import Image
+from torchvision import transforms
+from skimage.restoration import denoise_tv_chambolle
+from transformers import SamModel, SamProcessor
 
-# Load MiDaS depth estimation model
-midas_model = torch.hub.load("intel-isl/MiDaS", "DPT_Hybrid")
-midas_model.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-midas_model.to(device)
-midas_transform = torch.hub.load("intel-isl/MiDaS", "transforms").default_transform
-
-def estimate_depth(image):
-    """Estimate depth map to identify fabric folds."""
-    image = image.convert("RGB")
-    image_tensor = midas_transform(image).to(device)
-    
+# Load SAM model
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+model = SamModel.from_pretrained("facebook/sam-vit-huge").to(DEVICE)
+processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+
+def segment_dress(image):
+    """Segments the dress from an input image using SAM."""
+    input_points = [[[image.size[0] // 2, image.size[1] // 2]]]
+    inputs = processor(image, input_points=input_points, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
-        depth = midas_model(image_tensor).squeeze().cpu().numpy()
+        outputs = model(**inputs)
+    masks = processor.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+    )
+    return masks[0][0].numpy() if masks else None
+
+def warp_design(design, mask, warp_scale):
+    """Warp the design using TPS and scale control."""
+    h, w = mask.shape[:2]
+    design_resized = cv2.resize(design, (w, h))
     
-    depth = cv2.resize(depth, (image.size[0], image.size[1]))
-    depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255
-    return depth.astype(np.uint8)
-
-def detect_folds(image):
-    """Apply edge detection and highlight cloth folds."""
-    depth_map = estimate_depth(image)
-    edges = cv2.Canny(depth_map, 50, 150)
+    # Apply scaling
+    scaled_mask = (mask * 255 * (warp_scale / 100)).astype(np.uint8)
+    return cv2.bitwise_and(design_resized, design_resized, mask=scaled_mask)
+
+def blend_images(base, overlay, mask):
+    """Blends the design onto the dress using seamless cloning."""
+    center = tuple(np.array(base.shape[:2]) // 2)
+    return cv2.seamlessClone(overlay, base, mask, center, cv2.NORMAL_CLONE)
+
+def apply_design(image_path, design_path, warp_scale):
+    """Pipeline to segment, warp, and blend design onto dress."""
+    image = Image.open(image_path).convert("RGB")
+    design = cv2.imread(design_path)
+    mask = segment_dress(image)
     
-    # Convert edges to 3-channel image for visualization
-    edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
-    overlay = cv2.addWeighted(np.array(image), 0.7, edges_colored, 0.3, 0)
+    if mask is None:
+        return "Segmentation Failed!"
     
-    return Image.fromarray(overlay)
+    warped_design = warp_design(design, mask, warp_scale)
+    blended = blend_images(np.array(image), warped_design, mask)
+    return Image.fromarray(blended)
 
-def main(image):
-    return detect_folds(image)
+def main(image, design, warp_scale):
+    return apply_design(image, design, warp_scale)
 
-iface = gr.Interface(
+# Gradio UI
+demo = gr.Interface(
     fn=main,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Cloth Fold Detection",
-    description="Upload an image of clothing to visualize folds using depth estimation and edge detection."
+    inputs=[
+        gr.Image(type="filepath", label="Upload Dress Image"),
+        gr.Image(type="filepath", label="Upload Design Image"),
+        gr.Slider(0, 100, value=50, label="Warp Scale (%)")
+    ],
+    outputs=gr.Image(label="Warped Design on Dress"),
+    title="AI-Powered Dress Designer",
+    description="Upload a dress image and a design pattern. The AI will warp and blend the design onto the dress while preserving natural folds!"
 )
 
-if __name__ == "__main__":
-    iface.launch(share=True, debug=True)
+demo.launch()