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)