Create app.py

app.py

@@ -0,0 +1,173 @@
+import os
+import io
+import cv2
+import time
+import json
+import math
+import torch
+import numpy as np
+import pandas as pd
+from PIL import Image
+import gradio as gr
+
+# ----------------------------
+# Config
+# ----------------------------
+DEFAULT_MODEL_PATH = os.getenv("MODEL_PATH", "weights/best.pt")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ----------------------------
+# Model loading (edit this)
+# ----------------------------
+_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
+    _model = DummyModel()
+    return _model
+
+# ----------------------------
+# Inference wrapper (edit this)
+# ----------------------------
+def infer(image_bgr: np.ndarray, conf: float = 0.25):
+    """
+    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.
+    """
+    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}
+
+    # Option B (segmentation):
+    # mask = your_segmentation(image_bgr)  # 0/1 uint8
+    # return {"type": "mask", "mask": mask}
+
+    # --------- 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
+    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)
+    out = []
+    for c in cnts:
+        x, y, w, h = cv2.boundingRect(c)
+        if w * h >= min_area:
+            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"])
+    df.to_csv(path, index=False)
+    return path, df
+
+# ----------------------------
+# Gradio handlers
+# ----------------------------
+def process(image: Image.Image, conf: float, draw_x: bool, min_area: int):
+    if image is None:
+        return None, pd.DataFrame(), None
+
+    # PIL -> BGR np
+    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
+    vis = draw_boxes_with_x(img_bgr, boxes) if draw_x else img_bgr.copy()
+    vis_rgb = cv2.cvtColor(vis, cv2.COLOR_BGR2RGB)
+
+    # csv + table
+    csv_path, df = to_csv_file(boxes)
+
+    return Image.fromarray(vis_rgb), df, csv_path
+
+# ----------------------------
+# UI
+# ----------------------------
+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."
+    )
+    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")
+            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_csv = gr.File(label="Download CSV")
+
+    run_btn.click(process, inputs=[inp, conf, draw_x, min_area],
+                  outputs=[out_img, out_table, out_csv])
+
+if __name__ == "__main__":
+    load_model()  # warmup
+    demo.launch()