DEWARP

app.py

@@ -1,148 +1,76 @@
 import gradio as gr
-from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
-from transformers import AutoImageProcessor, AutoModel
-from huggingface_hub import snapshot_download
-from PIL import Image, ImageDraw
-import torch, numpy as np, cv2, zipfile, io, os, tempfile
-from sklearn.cluster import KMeans
-
-# -----------------------------------------------------
-# 1️⃣  Model Initialization
-# -----------------------------------------------------
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# --- Download SAM checkpoint if missing ---
-if not os.path.exists("sam_vit_b_01ec64.pth"):
-    os.system("wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth")
-
-# --- Load SAM ---
-sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b_01ec64.pth").to(device)
-mask_generator = SamAutomaticMaskGenerator(sam)
-
-# --- Preload DINOv2 ---
-snapshot_download("facebook/dinov2-base")
-processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base")
-dinov2 = AutoModel.from_pretrained("facebook/dinov2-base").to(device)
-
-# -----------------------------------------------------
-# 2️⃣  Utility Functions
-# -----------------------------------------------------
-def get_embeddings(img):
-    """Extract DINOv2 feature embeddings."""
-    inputs = processor(images=img, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = dinov2(**inputs)
-    feat = outputs.last_hidden_state[0].cpu().numpy()
-    return feat.mean(axis=0)
-
-def remove_background(image):
-    """Use largest SAM mask to isolate saree from background."""
-    masks = mask_generator.generate(image)
-    if not masks:
-        return image
-    main_mask = max(masks, key=lambda x: x['area'])['segmentation']
-    image[~main_mask] = 255  # white out background
-    return image
-
-def get_centroid(mask):
-    coords = np.column_stack(np.where(mask))
-    if len(coords) == 0:
-        return (0, 0)
-    y, x = coords.mean(axis=0)
-    return int(x), int(y)
-
-def make_transparent(img, mask):
-    rgba = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
-    rgba[..., 3] = np.where(mask, 255, 0).astype(np.uint8)
-    return rgba
-
-# -----------------------------------------------------
-# 3️⃣  Main Segmentation Function
-# -----------------------------------------------------
-def segment_saree(image):
-    try:
-        image = np.array(image.convert("RGB"))
-        image = remove_background(image)
-        masks = mask_generator.generate(image)
-        if not masks:
-            raise ValueError("No masks generated")
-
-        regions = []
-        for m in masks:
-            mask = m["segmentation"]
-            region_img = Image.fromarray(np.uint8(image) * mask[..., None])
-            emb = get_embeddings(region_img)
-            regions.append((mask, emb))
-
-        if len(regions) < 3:
-            raise ValueError("Insufficient distinct regions")
-
-        features = np.array([r[1] for r in regions])
-        kmeans = KMeans(n_clusters=3, random_state=42).fit(features)
-        labels = kmeans.labels_
-
-        colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0)]
-        names = ["Body", "Border", "Pallu"]
-        seg_out = np.zeros_like(image)
-        layers = [np.zeros_like(image, dtype=np.uint8) for _ in range(3)]
-
-        for i, (mask, _) in enumerate(regions):
-            seg_out[mask] = colors[labels[i]]
-            layers[labels[i]][mask] = image[mask]
-
-        seg_img = Image.fromarray(seg_out)
-        draw = ImageDraw.Draw(seg_img)
-        for (mask, _), lbl in zip(regions, labels):
-            x, y = get_centroid(mask)
-            draw.text((x, y), names[lbl], fill=(255, 255, 255))
-
-        # Transparent layers
-        transparent_imgs = [Image.fromarray(make_transparent(l, l.any(axis=2))) for l in layers]
-
-        # Write ZIP to a temp file (Gradio expects a real path)
-        tmpdir = tempfile.mkdtemp()
-        zip_path = os.path.join(tmpdir, "saree_layers.zip")
-        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
-            for n, t in zip(names, transparent_imgs):
-                tmp_img = os.path.join(tmpdir, f"{n}.png")
-                t.save(tmp_img)
-                zf.write(tmp_img, arcname=f"{n}.png")
-
-        return seg_img, transparent_imgs[0], transparent_imgs[1], transparent_imgs[2], zip_path
-
-    except Exception as e:
-        print("Error:", e)
-        blank = Image.new("RGB", (512, 512), color=(30, 30, 30))
-        return blank, blank, blank, blank, None
-
-# -----------------------------------------------------
-# 4️⃣  Gradio UI
-# -----------------------------------------------------
-description = """
-### 🧶 Saree AI — Intelligent Segmentation & Layer Export  
-Upload a **flat or draped saree image**, and this tool will:
-- ✂️ Remove background  
-- 🧠 Segment into **Body**, **Border**, **Pallu** using SAM + DINOv2  
-- 🪞 Provide transparent PNGs  
-- 📦 Download all masks as a single ZIP  
-
-Built for saree recoloring, catalog automation, and AI draping pipelines.
-"""
-
-demo = gr.Interface(
-    fn=segment_saree,
-    inputs=gr.Image(type="pil", label="Upload Saree Image"),
-    outputs=[
-        gr.Image(type="pil", label="Overlay Mask with Labels"),
-        gr.Image(type="pil", label="Body (Transparent)"),
-        gr.Image(type="pil", label="Border (Transparent)"),
-        gr.Image(type="pil", label="Pallu (Transparent)"),
-        gr.File(label="📦 Download All (ZIP)"),
-    ],
-    title="🧵 Saree AI — SAM + DINOv2 Smart Segmentation",
-    description=description,
-    flagging_mode="never",
-)
-
-if __name__ == "__main__":
-    demo.launch()
+import cv2
+import numpy as np
+from PIL import Image
+
+def flatten_image(img, points):
+    """
+    img: PIL.Image
+    points: list of (x, y) tuples in order [TL, TR, BR, BL]
+    """
+    if img is None or not points or len(points) != 4:
+        return None, "Please click exactly 4 points (TL, TR, BR, BL)."
+
+    # Convert to numpy array
+    image_np = np.array(img)
+    h, w = image_np.shape[:2]
+
+    # Convert input points to float32 numpy array
+    src_pts = np.array(points, dtype=np.float32)
+
+    # Compute output rectangle size using distances
+    width_top = np.linalg.norm(src_pts[0] - src_pts[1])
+    width_bottom = np.linalg.norm(src_pts[3] - src_pts[2])
+    height_left = np.linalg.norm(src_pts[0] - src_pts[3])
+    height_right = np.linalg.norm(src_pts[1] - src_pts[2])
+
+    max_width = int(max(width_top, width_bottom))
+    max_height = int(max(height_left, height_right))
+
+    dst_pts = np.array([
+        [0, 0],
+        [max_width - 1, 0],
+        [max_width - 1, max_height - 1],
+        [0, max_height - 1]
+    ], dtype=np.float32)
+
+    # Compute homography
+    M = cv2.getPerspectiveTransform(src_pts, dst_pts)
+
+    # Apply perspective warp
+    warped = cv2.warpPerspective(image_np, M, (max_width, max_height), flags=cv2.INTER_CUBIC)
+
+    warped_pil = Image.fromarray(warped)
+    return warped_pil, None
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 📸 Perspective Flatten Tool\nUpload an image, click 4 corners (Top-Left → Top-Right → Bottom-Right → Bottom-Left), then flatten!")
+
+    with gr.Row():
+        input_image = gr.Image(label="Upload Image", tool="select", type="pil")
+        output_image = gr.Image(label="Flattened Output")
+
+    coords = gr.State([])
+
+    def collect_points(evt: gr.SelectData, points):
+        if points is None:
+            points = []
+        points.append(evt.index)  # evt.index returns (x, y)
+        if len(points) > 4:
+            points = points[-4:]  # keep only last 4
+        return points, f"Selected {len(points)}/4 points: {points}"
+
+    points_output = gr.Textbox(label="Selected Points", interactive=False)
+
+    input_image.select(fn=collect_points, inputs=coords, outputs=[coords, points_output])
+
+    flatten_btn = gr.Button("🔄 Flatten Image")
+
+    error_box = gr.Textbox(label="Messages", interactive=False)
+
+    flatten_btn.click(fn=flatten_image, inputs=[input_image, coords], outputs=[output_image, error_box])
+
+    gr.Markdown("Tip: Re-upload image to reset point selection.")
+
+demo.launch()