Update app.py

app.py

@@ -1,9 +1,5 @@
 import os
-import io
 import cv2
-import time
-import json
-import math
 import torch
 import numpy as np
 import pandas as pd
@@ -16,92 +12,105 @@ import gradio as gr
 DEFAULT_MODEL_PATH = os.getenv("MODEL_PATH", "weights/best.pt")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+DEMO_INPUTS = {
+    "input_1.png": {
+        "output_path": "output_1.png",
+        "summary": (
+            "**Detected defects (input_1.png):**\n"
+            "- Row 1: Cell 1, Cell 3, Cell 5\n"
+            "- Row 2: None\n"
+            "- Row 3: Cell 1"
+        ),
+        "rows": [
+            {"row": 1, "cell": 1, "defect": "defect"},
+            {"row": 1, "cell": 3, "defect": "defect"},
+            {"row": 1, "cell": 5, "defect": "defect"},
+            {"row": 3, "cell": 1, "defect": "defect"},
+        ],
+    },
+    "input_2.png": {
+        "output_path": "output_2.png",
+        "summary": (
+            "**Detected defects (input_2.png):**\n"
+            "- Row 1: Cell 4 (crack), Cell 5 (large crack)\n"
+            "- Row 2: Cell 1 (crack), Cell 5 (multiple cracks)\n"
+            "- Row 3: Cell 1 (dark defect), Cell 2 (fine crack)"
+        ),
+        "rows": [
+            {"row": 1, "cell": 4, "defect": "crack"},
+            {"row": 1, "cell": 5, "defect": "large crack"},
+            {"row": 2, "cell": 1, "defect": "crack"},
+            {"row": 2, "cell": 5, "defect": "multiple cracks"},
+            {"row": 3, "cell": 1, "defect": "dark defect"},
+            {"row": 3, "cell": 2, "defect": "fine crack"},
+        ],
+    },
+}
+
 # ----------------------------
-# Model loading (edit this)
+# Model loading (kept placeholder)
 # ----------------------------
 _model = None
-
 def load_model(model_path: str = DEFAULT_MODEL_PATH):
-    """
-    Load your trained model once when the Space boots.
-    Replace the placeholder with your code.
-    """
     global _model
     if _model is not None:
         return _model
 
-    # >>> YOUR MODEL HERE <<<
-    # Example (PyTorch scripted/ckpt):
-    # ckpt = torch.load(model_path, map_location=DEVICE)
-    # model = MyNet(...)
-    # model.load_state_dict(ckpt["state_dict"] if "state_dict" in ckpt else ckpt)
-    # model.to(DEVICE).eval()
-    #
-    # For YOLO-like:
-    # from ultralytics import YOLO
-    # model = YOLO(model_path)
-
-    # Placeholder “no-model” to keep UI running:
     class DummyModel:
-        def __init__(self):
-            pass
+        pass
+
     _model = DummyModel()
     return _model
 
+
 # ----------------------------
-# Inference wrapper (edit this)
+# Helpers
 # ----------------------------
-def infer(image_bgr: np.ndarray, conf: float = 0.25):
+def _norm_mae(a: np.ndarray, b: np.ndarray) -> float:
+    """Normalized mean absolute error in [0,1] for matching demo images."""
+    if a.shape != b.shape:
+        b = cv2.resize(b, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_AREA)
+    a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY) if a.ndim == 3 else a
+    b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY) if b.ndim == 3 else b
+    a = a.astype(np.float32) / 255.0
+    b = b.astype(np.float32) / 255.0
+    return float(np.mean(np.abs(a - b)))
+
+
+def match_demo_image(upload_bgr: np.ndarray) -> str | None:
     """
-    Return defects as a list of boxes: [x1,y1,x2,y2,score,label]
-    OR return a binary mask (H,W) where 1=defect.
-    Edit this to call your model.
+    Return 'input_1.png' or 'input_2.png' if the uploaded image matches,
+    else None. Uses simple normalized MAE with resizing.
     """
-    model = load_model()
-
-    # >>> YOUR MODEL HERE <<<
-    # Option A (detection):
-    # results = model(image_bgr[..., ::-1])  # example if model expects RGB
-    # boxes = [[x1,y1,x2,y2,score,"defect"], ...]
-    # return {"type": "boxes", "boxes": boxes}
+    best_name, best_score = None, 1.0
+    for fname in DEMO_INPUTS.keys():
+        if not os.path.exists(fname):
+            continue
+        ref = cv2.imread(fname, cv2.IMREAD_COLOR)
+        if ref is None:
+            continue
+        score = _norm_mae(upload_bgr, ref)
+        if score < best_score:
+            best_score, best_name = score, fname
 
-    # Option B (segmentation):
-    # mask = your_segmentation(image_bgr)  # 0/1 uint8
-    # return {"type": "mask", "mask": mask}
+    # Loose threshold to cope with minor metadata/encoding differences
+    return best_name if best_score < 0.01 else None  # ~99% similar
 
-    # --------- PLACEHOLDER (edge blobs as fake defects) ---------
-    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
-    e = cv2.Canny(gray, 50, 150)
-    cnts, _ = cv2.findContours(e, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    boxes = []
-    h, w = gray.shape[:2]
-    for c in cnts:
-        x, y, bw, bh = cv2.boundingRect(c)
-        if bw * bh < max(0.0005 * w * h, 150):  # skip tiny
-            continue
-        boxes.append([x, y, x + bw, y + bh, 0.5, "defect"])
-        if len(boxes) >= 20:
-            break
-    return {"type": "boxes", "boxes": boxes}
 
-# ----------------------------
-# Utilities
-# ----------------------------
 def draw_boxes_with_x(image_bgr: np.ndarray, boxes, thickness: int = 3):
     img = image_bgr.copy()
-    color = (0, 0, 255)  # red in BGR
+    color = (0, 0, 255)
     for (x1, y1, x2, y2, score, label) in boxes:
         x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
         cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness)
-        # draw X
         cv2.line(img, (x1, y1), (x2, y2), color, thickness)
         cv2.line(img, (x1, y2), (x2, y1), color, thickness)
-        # label
         txt = f"{label}:{score:.2f}"
         cv2.putText(img, txt, (x1, max(y1 - 6, 0)),
                     cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
     return img
 
+
 def boxes_from_mask(mask: np.ndarray, min_area: int = 50):
     mask = (mask > 0).astype(np.uint8)
     cnts, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -112,37 +121,72 @@ def boxes_from_mask(mask: np.ndarray, min_area: int = 50):
             out.append([x, y, x + w, y + h, 1.0, "defect"])
     return out
 
-def to_csv_file(rows, path="/tmp/defect_report.csv"):
-    df = pd.DataFrame(rows, columns=["x1", "y1", "x2", "y2", "score", "label"])
+
+def to_csv_file_from_df(df: pd.DataFrame, path="/tmp/defect_report.csv"):
     df.to_csv(path, index=False)
     return path, df
 
+
+# ----------------------------
+# Placeholder inference for non-demo images
+# ----------------------------
+def infer_placeholder(image_bgr: np.ndarray, conf: float = 0.25):
+    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
+    e = cv2.Canny(gray, 50, 150)
+    cnts, _ = cv2.findContours(e, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    boxes = []
+    h, w = gray.shape[:2]
+    for c in cnts:
+        x, y, bw, bh = cv2.boundingRect(c)
+        if bw * bh < max(0.0005 * w * h, 150):
+            continue
+        boxes.append([x, y, x + bw, y + bh, 0.5, "defect"])
+        if len(boxes) >= 20:
+            break
+    return boxes
+
+
 # ----------------------------
-# Gradio handlers
+# Gradio handler
 # ----------------------------
 def process(image: Image.Image, conf: float, draw_x: bool, min_area: int):
     if image is None:
-        return None, pd.DataFrame(), None
+        return None, pd.DataFrame(), None, "Please upload an image."
 
-    # PIL -> BGR np
+    # PIL -> BGR
     img_rgb = np.array(image.convert("RGB"))
     img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
 
-    res = infer(img_bgr, conf=conf)
-
-    if res["type"] == "mask":
-        boxes = boxes_from_mask(res["mask"], min_area=min_area)
-    else:
-        boxes = [b for b in res["boxes"] if b[4] >= conf]
-
-    # draw
+    # 1) Check if this is one of our demo images
+    match_name = match_demo_image(img_bgr)
+    if match_name is not None:
+        meta = DEMO_INPUTS[match_name]
+        # Use the pre-rendered annotated image
+        vis_pil = Image.open(meta["output_path"]).convert("RGB")
+        df = pd.DataFrame(meta["rows"], columns=["row", "cell", "defect"])
+        csv_path, df = to_csv_file_from_df(df)
+        return vis_pil, df, csv_path, meta["summary"]
+
+    # 2) Otherwise fall back to placeholder detection
+    boxes = infer_placeholder(img_bgr, conf=conf)
+    boxes = [b for b in boxes if b[4] >= conf]
     vis = draw_boxes_with_x(img_bgr, boxes) if draw_x else img_bgr.copy()
     vis_rgb = cv2.cvtColor(vis, cv2.COLOR_BGR2RGB)
+    vis_pil = Image.fromarray(vis_rgb)
+
+    # Build a simple table for non-demo case
+    if boxes:
+        df = pd.DataFrame(
+            [{"x1": b[0], "y1": b[1], "x2": b[2], "y2": b[3], "score": b[4], "label": b[5]} for b in boxes]
+        )
+        summary = f"Detected {len(boxes)} defect region(s)."
+    else:
+        df = pd.DataFrame(columns=["x1", "y1", "x2", "y2", "score", "label"])
+        summary = "No defects detected by placeholder."
 
-    # csv + table
-    csv_path, df = to_csv_file(boxes)
+    csv_path, df = to_csv_file_from_df(df)
+    return vis_pil, df, csv_path, summary
 
-    return Image.fromarray(vis_rgb), df, csv_path
 
 # ----------------------------
 # UI
@@ -150,24 +194,26 @@ def process(image: Image.Image, conf: float, draw_x: bool, min_area: int):
 with gr.Blocks(title="AI-Driven EL Defect Recognition") as demo:
     gr.Markdown(
         "## AI-Driven Defect Recognition in EL Images\n"
-        "Upload an electroluminescence (EL) image. The app detects defective cells, "
-        "draws a red square with an X, and provides a CSV report."
+        "Upload an electroluminescence (EL) image. For the demo, uploading **input_1.png** "
+        "or **input_2.png** will return your predefined annotated results and a CSV.\n"
+        "For other images, a lightweight placeholder detector runs."
     )
     with gr.Row():
         with gr.Column():
             inp = gr.Image(type="pil", label="Upload EL image")
             conf = gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="Confidence threshold")
-            draw_x = gr.Checkbox(True, label="Draw red box + X")
+            draw_x = gr.Checkbox(True, label="Draw red box + X (non-demo only)")
             min_area = gr.Slider(10, 5000, value=120, step=10, label="Min defect area (pixels, for masks)")
             run_btn = gr.Button("Run inference", variant="primary")
         with gr.Column():
             out_img = gr.Image(type="pil", label="Annotated output")
-            out_table = gr.Dataframe(headers=["x1","y1","x2","y2","score","label"], label="Defect report (preview)")
+            out_table = gr.Dataframe(label="Defect report (preview)")
             out_csv = gr.File(label="Download CSV")
+            summary_md = gr.Markdown()
 
     run_btn.click(process, inputs=[inp, conf, draw_x, min_area],
-                  outputs=[out_img, out_table, out_csv])
+                  outputs=[out_img, out_table, out_csv, summary_md])
 
 if __name__ == "__main__":
-    load_model()  # warmup
+    load_model()
     demo.launch()