Upload 2 files

app.py

@@ -0,0 +1,152 @@
+# app.py
+import gradio as gr
+import torch
+import torch.nn as nn
+import timm
+import cv2
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import os
+
+# ===============================
+# SIMPLE DPT MODEL (DEPTH ESTIMATION)
+# ===============================
+class SimpleDPT(nn.Module):
+    def __init__(self, backbone_name='vit_base_patch16_384'):
+        super(SimpleDPT, self).__init__()
+        self.backbone = timm.create_model(backbone_name, pretrained=True, features_only=True)
+        feature_info = self.backbone.feature_info
+        channels = [f['num_chs'] for f in feature_info]
+
+        self.decoder = nn.Sequential(
+            nn.Conv2d(channels[-1], 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 1, kernel_size=1)
+        )
+
+    def forward(self, x, target_size):
+        features = self.backbone(x)
+        x = features[-1]
+        depth = self.decoder(x)
+        depth = nn.functional.interpolate(depth, size=target_size, mode='bilinear', align_corners=False)
+        return depth
+
+# ===============================
+# DEPTH → NORMAL MAP
+# ===============================
+def depth_to_normal(depth):
+    dy, dx = np.gradient(depth)
+    normal = np.dstack((-dx, -dy, np.ones_like(depth)))
+    n = np.linalg.norm(normal, axis=2, keepdims=True)
+    normal /= (n + 1e-8)
+    normal = (normal + 1) / 2
+    return normal
+
+# ===============================
+# MAIN PROCESSING FUNCTION
+# ===============================
+def process_saree(base_image: Image.Image, pattern_image: Image.Image):
+    # Convert base to numpy
+    img_pil = base_image.convert("RGB")
+    img_np = np.array(img_pil)
+
+    # Prepare tensor
+    img_resized = img_pil.resize((384, 384))
+    img_tensor = torch.from_numpy(np.array(img_resized)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
+    mean = torch.as_tensor([0.5, 0.5, 0.5], device=img_tensor.device).view(1, 3, 1, 1)
+    std = torch.as_tensor([0.5, 0.5, 0.5], device=img_tensor.device).view(1, 3, 1, 1)
+    img_tensor = (img_tensor - mean) / std
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = SimpleDPT(backbone_name='vit_base_patch16_384').to(device)
+    model.eval()
+
+    # Depth inference
+    with torch.no_grad():
+        target_size = img_pil.size[::-1]
+        depth_map = model(img_tensor.to(device), target_size=target_size)
+        depth_map = depth_map.squeeze().cpu().numpy()
+
+    # Normalize depth
+    depth_vis = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
+
+    # Normal map
+    normal_map = depth_to_normal(depth_vis)
+
+    # Shading map (CLAHE)
+    img_lab = cv2.cvtColor(img_np, cv2.COLOR_RGB2LAB)
+    l_channel, _, _ = cv2.split(img_lab)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    l_clahe = clahe.apply(l_channel)
+    shading_map = l_clahe / 255.0
+
+    # GrabCut mask
+    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
+    grabcut_mask = np.zeros(img_bgr.shape[:2], np.uint8)
+    height, width = img_bgr.shape[:2]
+    margin = int(min(width, height) * 0.05)
+    rect = (margin, margin, width - 2 * margin, height - 2 * margin)
+    bgdModel = np.zeros((1, 65), np.float64)
+    fgdModel = np.zeros((1, 65), np.float64)
+    cv2.grabCut(img_bgr, grabcut_mask, rect, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
+    mask = np.where((grabcut_mask == cv2.GC_FGD) | (grabcut_mask == cv2.GC_PR_FGD), 255, 0).astype(np.uint8)
+
+    # Tile pattern
+    pattern_np = np.array(pattern_image.convert("RGB"))
+    target_h, target_w = img_np.shape[:2]
+    pattern_h, pattern_w = pattern_np.shape[:2]
+    pattern_tiled = np.zeros((target_h, target_w, 3), dtype=np.uint8)
+    for y in range(0, target_h, pattern_h):
+        for x in range(0, target_w, pattern_w):
+            end_y = min(y + pattern_h, target_h)
+            end_x = min(x + pattern_w, target_w)
+            pattern_tiled[y:end_y, x:end_x] = pattern_np[0:(end_y - y), 0:(end_x - x)]
+
+    # Blend pattern
+    normal_map_loaded = normal_map.astype(np.float32)
+    shading_map_loaded = np.stack([shading_map] * 3, axis=-1)
+
+    alpha = 0.7
+    blended_shading = alpha * shading_map_loaded + (1 - alpha)
+
+    pattern_folded = pattern_tiled.astype(np.float32) / 255.0 * blended_shading
+    normal_boost = 0.5 + 0.5 * normal_map_loaded[..., 2:3]
+    pattern_folded *= normal_boost
+    pattern_folded = np.clip(pattern_folded, 0, 1)
+
+    # Clean mask and feather edges
+    mask_float = mask.astype(np.float32) / 255.0
+    kernel = np.ones((3, 3), np.uint8)
+    mask_clean = cv2.morphologyEx((mask_float * 255).astype(np.uint8), cv2.MORPH_OPEN, kernel)
+    mask_clean = cv2.morphologyEx(mask_clean, cv2.MORPH_CLOSE, kernel)
+    mask_clean = cv2.dilate(mask_clean, kernel, iterations=1)
+    mask_blurred = cv2.GaussianBlur(mask_clean, (15, 15), sigmaX=5, sigmaY=5)
+    mask_blurred[mask_blurred < 25] = 0
+    mask_blurred = mask_blurred.astype(np.float32) / 255.0
+
+    # Final RGBA
+    mask_stack = np.stack([mask_blurred] * 3, axis=-1)
+    pattern_final = pattern_folded * mask_stack
+    pattern_rgb = (pattern_final * 255).astype(np.uint8)
+    alpha_channel = (mask_blurred * 255).astype(np.uint8)
+    pattern_rgba = np.dstack((pattern_rgb, alpha_channel))
+
+    return Image.fromarray(pattern_rgba, mode="RGBA")
+
+# ===============================
+# GRADIO INTERFACE
+# ===============================
+iface = gr.Interface(
+    fn=process_saree,
+    inputs=[gr.Image(type="pil", label="Base Saree Image"),
+            gr.Image(type="pil", label="Pattern Image")],
+    outputs=gr.Image(type="pil", label="Final Saree Output"),
+    title="Saree Depth + Pattern Draping",
+    description="Upload base saree & pattern images to get depth-aware draped output (transparent edges, no black outline)."
+)
+
+if __name__ == "__main__":
+    iface.launch()

requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchvision
+timm
+opencv-python
+Pillow
+matplotlib
+tqdm
+gradio