app.py

import os
import cv2
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"

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 stub (kept)
# ----------------------------
_model = None
def load_model(model_path: str = DEFAULT_MODEL_PATH):
    global _model
    if _model is not None:
        return _model
    class Dummy: ...
    _model = Dummy()
    return _model

# ----------------------------
# Helpers
# ----------------------------
def _norm_mae(a: np.ndarray, b: np.ndarray) -> float:
    if a.shape != b.shape:
        b = cv2.resize(b, (a.shape[1], a.shape[0]), interpolation=cv2.INTER_AREA)
    if a.ndim == 3: a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY)
    if b.ndim == 3: b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY)
    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:
    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
    return best_name if best_score < 0.01 else None  # ~99% similar

def draw_boxes_with_x(image_bgr: np.ndarray, boxes, thickness: int = 3):
    img = image_bgr.copy()
    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)
        cv2.line(img, (x1, y1), (x2, y2), color, thickness)
        cv2.line(img, (x1, y2), (x2, y1), color, thickness)
        cv2.putText(img, f"{label}:{score:.2f}", (x1, max(y1-6,0)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2, cv2.LINE_AA)
    return img

def to_csv(df: pd.DataFrame, path="/tmp/defect_report.csv"):
    df.to_csv(path, index=False)
    return path, df

def placeholder_boxes(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 [b for b in boxes if b[4] >= conf]

# ----------------------------
# Gradio handler (filepath input)
# ----------------------------
def process(image_path: str, conf: float, draw_x: bool, min_area: int):
    try:
        if not image_path or not os.path.exists(image_path):
            return None, pd.DataFrame(), None, "Please upload an image."

        # read as BGR
        pil_img = Image.open(image_path).convert("RGB")
        img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

        # Demo match?
        match_name = match_demo_image(img_bgr)
        if match_name is not None:
            meta = DEMO_INPUTS[match_name]
            vis_pil = Image.open(meta["output_path"]).convert("RGB")
            df = pd.DataFrame(meta["rows"], columns=["row", "cell", "defect"])
            csv_path, df = to_csv(df)
            return vis_pil, df, csv_path, meta["summary"]

        # Fallback placeholder
        boxes = placeholder_boxes(img_bgr, conf=conf)
        vis = draw_boxes_with_x(img_bgr, boxes) if draw_x else img_bgr.copy()
        vis_pil = Image.fromarray(cv2.cvtColor(vis, cv2.COLOR_BGR2RGB))
        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)} region(s) by placeholder."
        else:
            df = pd.DataFrame(columns=["x1","y1","x2","y2","score","label"])
            summary = "No defects detected by placeholder."
        csv_path, df = to_csv(df)
        return vis_pil, df, csv_path, summary
    except Exception as e:
        # surface errors to UI instead of crashing
        return None, pd.DataFrame(), None, f"Error: {type(e).__name__}: {e}"

# ----------------------------
# UI
# ----------------------------
with gr.Blocks(title="AI-Driven EL Defect Recognition") as demo:
    gr.Markdown(
        "## Defect Images\n"
        " .\n"
    )
    with gr.Row():
        with gr.Column():
            # Use filepath to avoid large base64 uploads
            inp = gr.Image(type="filepath", 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 (non-demo only)")
            min_area = gr.Slider(10, 5000, value=120, step=10, label="Min defect area (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(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, summary_md])

if __name__ == "__main__":
    load_model()
    # Disable SSR to avoid upload edge-cases
    demo.launch(ssr_mode=False)