UPDATE

app.py

@@ -1,107 +1,117 @@
-# ==============================================================
-# 👗 Saree AI — Content-Aware Image Fitting (Smart Bounding Box)
-# Hugging Face Space — Stable Version (no runtime errors)
-# ==============================================================
-
-import cv2, numpy as np, gradio as gr
-
-# --------------------------------------------------------------
-# Core Function
-# --------------------------------------------------------------
-def content_aware_fit(image, target_size=(512, 512)):
-    """Performs content-aware fitting of saree images preserving ornate regions."""
-    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-    # ----------------------------------------------------------
-    # Step 1: Texture-based "Saliency" detection
-    # ----------------------------------------------------------
-    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
-    lap = cv2.Laplacian(gray, cv2.CV_64F)
-    lap = np.absolute(lap)
-    saliency_map = np.uint8(255 * (lap / np.max(lap)))
-    saliency_map = cv2.GaussianBlur(saliency_map, (7, 7), 0)
-    saliency_map = cv2.normalize(saliency_map, None, 0, 255, cv2.NORM_MINMAX)
-    _, importance_mask = cv2.threshold(saliency_map, 128, 255, cv2.THRESH_BINARY)
-
-    # ----------------------------------------------------------
-    # Step 2: Center of Mass (Weighted by Saliency)
-    # ----------------------------------------------------------
-    M = cv2.moments(importance_mask)
-    if M["m00"] != 0:
-        cx = int(M["m10"] / M["m00"])
-        cy = int(M["m01"] / M["m00"])
-    else:
-        cx, cy = img_rgb.shape[1] // 2, img_rgb.shape[0] // 2
-
-    # ----------------------------------------------------------
-    # Step 3: Smart Bounding Box + Padding
-    # ----------------------------------------------------------
-    coords = np.column_stack(np.where(importance_mask > 0))
-    if coords.shape[0] > 0:
-        x, y, w, h = cv2.boundingRect(coords)
-    else:
-        x, y, w, h = 0, 0, img_rgb.shape[1], img_rgb.shape[0]
-
-    pad_x = int(0.1 * w)
-    pad_y = int(0.1 * h)
-    x1, y1 = max(0, x - pad_x), max(0, y - pad_y)
-    x2, y2 = min(img_rgb.shape[1], x + w + pad_x), min(img_rgb.shape[0], y + h + pad_y)
-    cropped = img_rgb[y1:y2, x1:x2]
-
-    # ----------------------------------------------------------
-    # Step 4: Aspect-ratio-safe Fit + Padding
-    # ----------------------------------------------------------
-    h, w, _ = cropped.shape
-    target_h, target_w = target_size
-    scale = min(target_w / w, target_h / h)
-    new_w, new_h = int(w * scale), int(h * scale)
-    resized = cv2.resize(cropped, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
-
-    pad_x = (target_w - new_w) // 2
-    pad_y = (target_h - new_h) // 2
-    fitted = cv2.copyMakeBorder(
-        resized, pad_y, target_h - new_h - pad_y,
-        pad_x, target_w - new_w - pad_x,
-        cv2.BORDER_CONSTANT, value=[255, 255, 255]
-    )
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+
+# Optional: try to load a Hugging Face dewarping model if installed
+try:
+    from transformers import AutoModel, AutoImageProcessor
+    MODEL_REPO = "richard1231/Document_dewarping_platform"
+    processor = AutoImageProcessor.from_pretrained(MODEL_REPO)
+    model = AutoModel.from_pretrained(MODEL_REPO)
+    model.eval()
+    USE_HF_MODEL = True
+except Exception as e:
+    print("⚠️ Hugging Face model not found, using OpenCV-only version.")
+    USE_HF_MODEL = False
+
+
+# -------------------------------------------------------------
+# 🔹 Perspective correction (OpenCV fallback)
+# -------------------------------------------------------------
+def flatten_perspective(input_image: Image.Image) -> Image.Image:
+    img = np.array(input_image.convert("RGB"))
+    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+    blur = cv2.GaussianBlur(gray, (5, 5), 0)
+    edges = cv2.Canny(blur, 50, 150)
+
+    contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours:
+        return input_image
+    contour = max(contours, key=cv2.contourArea)
+
+    peri = cv2.arcLength(contour, True)
+    approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
+    if len(approx) != 4:
+        return input_image
+
+    pts = np.float32(approx.reshape(4, 2))
+    s = pts.sum(axis=1)
+    rect = np.zeros((4, 2), dtype="float32")
+    rect[0] = pts[np.argmin(s)]
+    rect[2] = pts[np.argmax(s)]
+    diff = np.diff(pts, axis=1)
+    rect[1] = pts[np.argmin(diff)]
+    rect[3] = pts[np.argmax(diff)]
+
+    (tl, tr, br, bl) = rect
+    widthA = np.linalg.norm(br - bl)
+    widthB = np.linalg.norm(tr - tl)
+    heightA = np.linalg.norm(tr - br)
+    heightB = np.linalg.norm(tl - bl)
+    maxWidth, maxHeight = int(max(widthA, widthB)), int(max(heightA, heightB))
+
+    dst = np.array([[0, 0],
+                    [maxWidth - 1, 0],
+                    [maxWidth - 1, maxHeight - 1],
+                    [0, maxHeight - 1]], dtype="float32")
+
+    M = cv2.getPerspectiveTransform(rect, dst)
+    warped = cv2.warpPerspective(img, M, (maxWidth, maxHeight))
+    return Image.fromarray(warped)
 
-    # ----------------------------------------------------------
-    # Step 5: Overlay (for visualization)
-    # ----------------------------------------------------------
-    marked = img_rgb.copy()
-    cv2.circle(marked, (cx, cy), 8, (255, 0, 0), -1)
-    cv2.rectangle(marked, (x1, y1), (x2, y2), (0, 255, 0), 3)
-
-    return (
-        cv2.cvtColor(saliency_map, cv2.COLOR_GRAY2RGB),
-        cv2.cvtColor(importance_mask, cv2.COLOR_GRAY2RGB),
-        marked,
-        cropped,
-        fitted
-    )
 
-# --------------------------------------------------------------
-# Gradio Interface
-# --------------------------------------------------------------
-demo = gr.Interface(
-    fn=content_aware_fit,
-    inputs=gr.Image(type="pil", label="Upload Saree Image"),
-    outputs=[
-        gr.Image(label="Texture / Saliency Map"),
-        gr.Image(label="Importance Mask"),
-        gr.Image(label="Center + Bounding Box"),
-        gr.Image(label="Cropped View"),
-        gr.Image(label="Final Content-Aware Fitted Output")
-    ],
-    title="👗 Saree AI — Content-Aware Image Fitting",
-    description=(
-        "Automatically detects ornate or high-detail regions (borders, pallus) "
-        "and fits the saree image with smart padding to preserve its design. "
-        "Ideal preprocessing step before saree draping or catalog generation."
-    ),
-    allow_flagging="never"
-)
+# -------------------------------------------------------------
+# 🔹 Learned de-warping (Hugging Face model)
+# -------------------------------------------------------------
+@torch.no_grad()
+def flatten_learned(input_image: Image.Image) -> Image.Image:
+    if not USE_HF_MODEL:
+        return flatten_perspective(input_image)
+
+    inputs = processor(images=input_image, return_tensors="pt")
+    outputs = model(**inputs)
+    # Post-process — many HF models return tensors in 0-1 range
+    out_img = outputs.last_hidden_state[0]
+    out_img = (out_img - out_img.min()) / (out_img.max() - out_img.min())
+    out_img = (out_img * 255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
+    return Image.fromarray(out_img)
+
+
+# -------------------------------------------------------------
+# 🔹 Gradio UI
+# -------------------------------------------------------------
+description = """
+## 🧾 Auto Image Flattening (Perspective + Learned Dewarping)
+Upload a **tilted or curved document/fabric photo**.  
+- Default: OpenCV 4-point perspective flattening  
+- Optional: if the **Hugging Face DewarpNet/DocRes model** is available, uses that instead  
+"""
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 📄 Auto Image Flattening (OpenCV / Hugging Face)")
+    gr.Markdown(description)
+
+    with gr.Row():
+        inp = gr.Image(type="pil", label="Upload Image")
+        out = gr.Image(type="pil", label="Flattened Output")
+
+    mode = gr.Radio(["Auto (Use HF if available)", "OpenCV Only"], value="Auto (Use HF if available)", label="Mode")
+
+    def process(img, mode):
+        if mode == "OpenCV Only" or not USE_HF_MODEL:
+            return flatten_perspective(img)
+        return flatten_learned(img)
+
+    btn = gr.Button("Flatten Image")
+    btn.click(process, inputs=[inp, mode], outputs=out)
+
+    gr.Examples(
+        examples=["example1.jpg", "example2.jpg"],
+        inputs=inp,
+        examples_per_page=2,
+    )
 
 if __name__ == "__main__":
     demo.launch()