PATTERN

app.py

@@ -1,66 +1,95 @@
-import gradio as gr
-import cv2, numpy as np
+# ==========================================================
+# 🪔 Saree Pattern Extractor (DINOv2 / ViT Patch Clustering)
+# Optimized for HF Free CPU Tier
+# ==========================================================
+import os, zipfile, io, cv2, numpy as np, torch
 from PIL import Image
+from sklearn.cluster import KMeans
+import gradio as gr
+from transformers import AutoImageProcessor, AutoModel
+import matplotlib.pyplot as plt
 
-# -------------------------------------------------
-# 🌈  Color-Preserving Fold Remover (LAB domain)
-# -------------------------------------------------
-def remove_folds_color_safe(image, intensity=0.5, enhance=False):
-    if image is None:
-        return None
+# -----------------------------
+# 1️⃣  Load DINOv2-small model
+# -----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_id = "facebook/dinov2-small"
+processor = AutoImageProcessor.from_pretrained(model_id)
+model = AutoModel.from_pretrained(model_id).to(device).eval()
 
-    # Convert to numpy / LAB
-    img_rgb = np.array(image.convert("RGB"))
-    img_lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)
-    L, A, B = cv2.split(img_lab)
+# -----------------------------
+# 2️⃣  Helper functions
+# -----------------------------
+def upscale_label_map(lbl, target_size):
+    h, w = lbl.shape
+    return cv2.resize(lbl.astype(np.uint8), target_size, interpolation=cv2.INTER_NEAREST)
 
-    # --- (1) Smooth illumination field ---
-    ksize = int(61 + intensity * 140)           # adaptive blur
-    ksize = ksize + 1 if ksize % 2 == 0 else ksize
-    L_blur = cv2.GaussianBlur(L, (ksize, ksize), 0)
+def extract_patterns(image: Image.Image, K: int = 8):
+    # Resize to keep under memory limits
+    image = image.convert("RGB")
+    img_small = image.copy()
+    img_small.thumbnail((480,480))
 
-    # --- (2) Lighting flattening (division in luminance only) ---
-    L_float = L.astype(np.float32)
-    L_norm = (L_float / (L_blur + 1e-6)) * np.mean(L_blur)
-    L_norm = np.clip(L_norm, 0, 255).astype(np.uint8)
+    inputs = processor(images=img_small, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs, output_hidden_states=True)
+        feats = outputs.last_hidden_state.squeeze(0)[1:].cpu().numpy()
 
-    # --- (3) Local contrast restore (CLAHE) ---
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-    L_final = clahe.apply(L_norm)
+    grid = int(np.sqrt(feats.shape[0]))
+    km = KMeans(n_clusters=K, random_state=0, n_init="auto").fit(feats)
+    labels = km.labels_.reshape(grid, grid)
 
-    # --- (4) Optional weave-enhancement ---
-    if enhance:
-        sharp = cv2.GaussianBlur(L_final, (0, 0), 3)
-        L_final = cv2.addWeighted(L_final, 1.3, sharp, -0.3, 0)
+    # Rebuild overlay
+    lbl_map = upscale_label_map(labels, img_small.size)
+    colors = plt.cm.tab10(np.linspace(0,1,K))[:,:3]
+    overlay = np.zeros((*lbl_map.shape,3))
+    for k in range(K):
+        overlay[lbl_map==k] = colors[k]
+    overlay = (overlay*255).astype(np.uint8)
+    blend = cv2.addWeighted(np.array(img_small), 0.6, overlay, 0.4, 0)
 
-    # --- (5) Merge color back and convert to RGB ---
-    img_lab_eq = cv2.merge([L_final, A, B])
-    result = cv2.cvtColor(img_lab_eq, cv2.COLOR_LAB2RGB)
+    # Patch export
+    outdir = "patterns"
+    os.makedirs(outdir, exist_ok=True)
+    ph, pw = np.array(img_small.size)//grid
+    for k in range(K):
+        mask = (labels==k)
+        coords = np.argwhere(mask)
+        for i,(y,x) in enumerate(coords):
+            y0,y1 = int(y*ph), int((y+1)*ph)
+            x0,x1 = int(x*pw), int((x+1)*pw)
+            patch = img_small.crop((x0,y0,x1,y1))
+            patch.save(f"{outdir}/cluster{k}_patch{i}.png")
 
-    return Image.fromarray(result)
+    # Zip patches
+    zip_path = "patterns.zip"
+    with zipfile.ZipFile(zip_path,"w",zipfile.ZIP_DEFLATED) as zf:
+        for fn in os.listdir(outdir):
+            zf.write(os.path.join(outdir,fn), fn)
 
-# -------------------------------------------------
-# 🎛️  Gradio UI
-# -------------------------------------------------
-with gr.Blocks(title="🪄 Saree Fold Remover — Color-Safe LAB Edition") as demo:
-    gr.Markdown("""
-    ## 🪄 Saree Fold Remover — *Color-Safe Edition*
-    Removes folds & lighting shadows **without changing original colors**.<br>
-    Works beautifully on silk, cotton and zari fabrics.<br>
-    Adjust *Fold Intensity* for stronger corrections; enable *Enhance Texture* for sharper weave.
-    """)
+    return Image.fromarray(blend), zip_path
 
-    with gr.Row():
-        inp = gr.Image(label="Upload Saree Image", type="pil")
-        intensity = gr.Slider(
-            0.0, 1.0, value=0.5, step=0.05,
-            label="Fold Intensity",
-            info="Higher = stronger flattening"
-        )
-        enhance = gr.Checkbox(label="✨ Enhance Texture", value=False)
+# -----------------------------
+# 3️⃣  Gradio Interface
+# -----------------------------
+title = "🪔 Saree Pattern Extractor (DINOv2)"
+desc = "Upload a saree image → cluster similar motifs & textures → preview overlay → download all patch crops as ZIP."
 
-    out = gr.Image(label="Flat, Fold-Free Output")
-    run = gr.Button("🚀 Remove Folds")
-    run.click(fn=remove_folds_color_safe, inputs=[inp, intensity, enhance], outputs=out)
+demo = gr.Interface(
+    fn=extract_patterns,
+    inputs=[
+        gr.Image(label="Upload Saree Image", type="pil"),
+        gr.Slider(3,12,value=8,step=1,label="Number of Pattern Clusters")
+    ],
+    outputs=[
+        gr.Image(label="Cluster Overlay"),
+        gr.File(label="Download ZIP of Crops")
+    ],
+    title=title,
+    description=desc,
+    allow_flagging="never",
+    cache_examples=False
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()