import os, json, cv2, numpy as np, gradio as gr from dataclasses import dataclass from typing import List, Dict, Tuple # --- Quiet Ultralytics settings warning and ensure a writeable config dir --- os.environ.setdefault("YOLO_CONFIG_DIR", "/tmp/Ultralytics") os.makedirs("/tmp/Ultralytics", exist_ok=True) # ---------- Config ---------- @dataclass class InspectConfig: min_box_area: int = 2000 aspect_min: float = 0.35 aspect_max: float = 3.0 size_tol_pct: float = 20.0 appearance_thresh: float = 0.30 solidity_thresh: float = 0.90 seal_center_band: float = 0.25 seal_min_length_frac: float = 0.6 canny1: int = 50 canny2: int = 150 morph_kernel: int = 5 debug: bool = False CFG = InspectConfig() # ---------- Helpers ---------- def roi_histogram(img_bgr: np.ndarray) -> np.ndarray: hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) hist = cv2.calcHist([hsv], [0, 1], None, [32, 32], [0, 180, 0, 256]) cv2.normalize(hist, hist) return hist def bhattacharyya(h1: np.ndarray, h2: np.ndarray) -> float: return cv2.compareHist(h1, h2, cv2.HISTCMP_BHATTACHARYYA) # ---------- YOLO (lazy load) ---------- from ultralytics import YOLO _YOLO = None # Prefer a local weights file; or pull from a model repo if provided YOLO_MODEL_PATH = os.getenv("YOLO_MODEL_PATH", "weights/best.pt") YOLO_MODEL_REPO = os.getenv("YOLO_MODEL_REPO", "") # e.g. "vaibhavi18092002/myralens-yolo" YOLO_MODEL_FILE = os.getenv("YOLO_MODEL_FILE", "best.pt") def _ensure_weights(): """Ensure YOLO_MODEL_PATH exists; if not, try to download from a HF model repo.""" global YOLO_MODEL_PATH if os.path.exists(YOLO_MODEL_PATH): return if YOLO_MODEL_REPO: try: from huggingface_hub import hf_hub_download YOLO_MODEL_PATH = hf_hub_download( repo_id=YOLO_MODEL_REPO, filename=YOLO_MODEL_FILE, repo_type="model" ) except Exception as e: print("[WARN] Could not download weights:", e) def _get_yolo(): global _YOLO if _YOLO is None: _ensure_weights() # Fallback to yolov8n.pt if custom weights not found weights = YOLO_MODEL_PATH if os.path.exists(YOLO_MODEL_PATH) else "yolov8n.pt" _YOLO = YOLO(weights) print("[INFO] Loaded YOLO weights from:", weights) return _YOLO USE_YOLO = True # ---------- Detection + rules ---------- def detect_boxes(frame: np.ndarray, cfg: InspectConfig = CFG) -> List[Dict]: if USE_YOLO: res = _get_yolo().predict(source=frame, verbose=False, conf=0.25)[0] boxes = [] H, W = frame.shape[:2] if res.boxes is None or len(res.boxes) == 0: return boxes for x1, y1, x2, y2 in res.boxes.xyxy.cpu().numpy(): x1, y1, x2, y2 = map(int, [x1, y1, x2, y2]) x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(W - 1, x2), min(H - 1, y2) w, h = max(1, x2 - x1), max(1, y2 - y1) roi = frame[y1:y2, x1:x2].copy() area = float(w * h) aspect = max(w, h) / (min(w, h) + 1e-6) boxes.append({ "bbox_xywh": [x1, y1, w, h], "box_pts": [[x1, y1], [x2, y1], [x2, y2], [x1, y2]], "area": area, "aspect": aspect, "solidity": 1.0, "roi": roi }) return boxes # (Edge-based fallback if you ever disable YOLO) gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) blur = cv2.GaussianBlur(gray, (5, 5), 0) edges = cv2.Canny(blur, cfg.canny1, cfg.canny2) kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (cfg.morph_kernel, cfg.morph_kernel)) closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=2) cnts, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) proposals: List[Dict] = [] for c in cnts: area = cv2.contourArea(c) if area < cfg.min_box_area: continue rect = cv2.minAreaRect(c); (w, h) = rect[1] if w == 0 or h == 0: continue aspect = max(w, h) / (min(w, h) + 1e-6) if not (cfg.aspect_min <= aspect <= cfg.aspect_max): continue box = cv2.boxPoints(rect).astype(int) hull = cv2.convexHull(c) solidity = float(area) / (cv2.contourArea(hull) + 1e-6) x, y, ww, hh = cv2.boundingRect(c) roi = frame[y:y + hh, x:x + ww].copy() proposals.append({ "bbox_xywh": [int(x), int(y), int(ww), int(hh)], "box_pts": box.tolist(), "area": float(area), "aspect": float(aspect), "solidity": float(solidity), "roi": roi }) return proposals def check_size_uniformity(boxes: List[Dict], cfg: InspectConfig = CFG) -> None: if not boxes: return areas = np.array([b["area"] for b in boxes]) med = np.median(areas); tol = (cfg.size_tol_pct / 100.0) * med for b in boxes: b["size_ok"] = (abs(b["area"] - med) <= tol) b["size_note"] = {"median_area": float(med), "tol": float(tol)} def check_appearance_uniformity(boxes: List[Dict], cfg: InspectConfig = CFG) -> None: if not boxes: return hists = [roi_histogram(b["roi"]) for b in boxes] areas = np.array([b["area"] for b in boxes]) ref_idx = int(np.argsort(areas)[len(areas)//2]); ref_hist = hists[ref_idx] for b, h in zip(boxes, hists): dist = bhattacharyya(ref_hist, h) b["appearance_ok"] = (dist <= CFG.appearance_thresh) b["appearance_dist"] = float(dist) def check_damage(boxes: List[Dict], cfg: InspectConfig = CFG) -> None: for b in boxes: b["damage_ok"] = (b["solidity"] >= cfg.solidity_thresh) def check_seal(boxes: List[Dict], cfg: InspectConfig = CFG) -> None: for b in boxes: roi = b["roi"] if roi.size == 0: b["seal_ok"] = False; continue h, w = roi.shape[:2] band_h = int(h * cfg.seal_center_band) y1 = max(0, h//2 - band_h//2); y2 = min(h, h//2 + band_h//2) band = roi[y1:y2, :] g = cv2.cvtColor(band, cv2.COLOR_BGR2GRAY) g = cv2.GaussianBlur(g, (3,3), 0) edges = cv2.Canny(g, 50, 150) lines = cv2.HoughLinesP(edges, 1, np.pi/180, 80, minLineLength=int(cfg.seal_min_length_frac*w), maxLineGap=10) b["seal_ok"] = lines is not None b["seal_lines_found"] = 0 if lines is None else int(len(lines)) def score_and_annotate(frame: np.ndarray, boxes: List[Dict]) -> Tuple[np.ndarray, List[Dict]]: out = frame.copy(); results = [] for idx, b in enumerate(boxes): x, y, w, h = b["bbox_xywh"]; issues = [] if not b.get("size_ok", True): issues.append("size_outlier") if not b.get("appearance_ok", True): issues.append("appearance_diff") if not b.get("damage_ok", True): issues.append("possible_damage") if not b.get("seal_ok", True): issues.append("seal_missing") status = "PASS" if len(issues) == 0 else "FAIL" color = (0, 200, 0) if status == "PASS" else (0, 0, 255) cv2.rectangle(out, (x, y), (x+w, y+h), color, 2) cv2.putText(out, f"{status}:{idx}", (x, max(0, y-8)), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) if issues: cv2.putText(out, ",".join(issues), (x, y+h+18), cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2) results.append({ "id": idx, "bbox_xywh": b["bbox_xywh"], "status": status, "issues": issues, "metrics": { "area": b["area"], "aspect": b["aspect"], "solidity": b["solidity"], "appearance_dist": float(b.get("appearance_dist", 0.0)), "seal_lines_found": int(b.get("seal_lines_found", 0)) } }) overall = "PASS" if all(r["status"] == "PASS" for r in results) else "FAIL" cv2.putText(out, f"Batch: {overall}", (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0,200,0) if overall=="PASS" else (0,0,255), 2) return out, results def inspect_frame(frame: np.ndarray, cfg: InspectConfig = CFG): boxes = detect_boxes(frame, cfg) check_size_uniformity(boxes, cfg) check_appearance_uniformity(boxes, cfg) check_damage(boxes, cfg) check_seal(boxes, cfg) annotated, per_box = score_and_annotate(frame, boxes) report = { "boxes": per_box, "counts": { "total": len(per_box), "pass": sum(1 for r in per_box if r["status"] == "PASS"), "fail": sum(1 for r in per_box if r["status"] == "FAIL"), } } return annotated, report def run_on_image(image_path: str, out_img_path: str, out_json_path: str, cfg: InspectConfig = CFG): os.makedirs(os.path.dirname(out_img_path), exist_ok=True) os.makedirs(os.path.dirname(out_json_path), exist_ok=True) frame = cv2.imread(image_path) if frame is None: raise FileNotFoundError(f"Could not read image: {image_path}") annotated, report = inspect_frame(frame, cfg) cv2.imwrite(out_img_path, annotated) with open(out_json_path, "w") as f: json.dump(report, f, indent=2) def run_on_video(video_path: str, out_video_path: str, out_jsonl_path: str, cfg: InspectConfig = CFG): cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise FileNotFoundError(f"Could not open video: {video_path}") os.makedirs(os.path.dirname(out_video_path), exist_ok=True) os.makedirs(os.path.dirname(out_jsonl_path), exist_ok=True) fps = cap.get(cv2.CAP_PROP_FPS) or 25.0 w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)); h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) writer = cv2.VideoWriter(out_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) frame_idx = 0 with open(out_jsonl_path, "w") as jf: while True: ret, frame = cap.read() if not ret: break annotated, report = inspect_frame(frame, cfg) writer.write(annotated) jf.write(json.dumps({"frame": frame_idx, **report}) + "\n") frame_idx += 1 cap.release(); writer.release() # ---------- Gradio UI ---------- os.makedirs("uploads", exist_ok=True) os.makedirs("outputs", exist_ok=True) def ui_image(img, size_tol=20.0, appearance_thr=0.30, solidity_thr=0.90, seal_frac=0.60): CFG.size_tol_pct = float(size_tol) CFG.appearance_thresh = float(appearance_thr) CFG.solidity_thresh = float(solidity_thr) CFG.seal_min_length_frac = float(seal_frac) inp = "uploads/img.jpg"; out_img = "outputs/img_annot.jpg"; out_json = "outputs/img_report.json" cv2.imwrite(inp, cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) run_on_image(inp, out_img, out_json, CFG) with open(out_json) as f: report = f.read() return out_img, report, out_img, out_json def ui_video(video_path, size_tol=20.0, appearance_thr=0.30, solidity_thr=0.90, seal_frac=0.60): CFG.size_tol_pct = float(size_tol) CFG.appearance_thresh = float(appearance_thr) CFG.solidity_thresh = float(solidity_thr) CFG.seal_min_length_frac = float(seal_frac) out_vid = "outputs/vid_annot.mp4"; out_jsonl = "outputs/vid_report.jsonl" run_on_video(video_path, out_vid, out_jsonl, CFG) head = [] try: with open(out_jsonl) as f: for _ in range(8): line = f.readline() if not line: break head.append(line.strip()) except: pass return out_vid, "\n".join(head), out_vid, out_jsonl theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="slate") custom_css = """ .gradio-container { max-width: 100% !important; padding: 0 24px !important; } #header { background: linear-gradient(135deg, #4f46e5 0%, #7c3aed 100%); color: #fff; padding: 22px 24px; border-radius: 16px; margin: 18px 0; } #header h1 { font-size: 26px; margin: 0 0 6px 0; line-height: 1.2; } #header p { margin: 0; opacity: .9; } .section { background: #fff; border-radius: 14px; box-shadow: 0 6px 20px rgba(0,0,0,.06); padding: 16px; } """ with gr.Blocks(theme=theme, css=custom_css, fill_height=True) as demo: gr.HTML(""" """) with gr.Row(): with gr.Column(scale=3, elem_classes=["section"]): size_tol = gr.Slider(5, 50, value=20, step=1, label="Size tolerance (±%)") appearance_thr = gr.Slider(0.05, 0.80, value=0.30, step=0.01, label="Appearance threshold (lower = stricter)") solidity_thr = gr.Slider(0.70, 0.99, value=0.90, step=0.01, label="Damage sensitivity (solidity ≥)") seal_frac = gr.Slider(0.30, 0.90, value=0.60, step=0.01, label="Seal min length (fraction of width)") with gr.Tab("Image"): with gr.Row(): with gr.Column(scale=1, elem_classes=["section"]): img_in = gr.Image(type="numpy", label="Upload image", height=420) img_btn = gr.Button("Run inspection", variant="primary") with gr.Column(scale=1, elem_classes=["section"]): img_out = gr.Image(label="Annotated image", height=420) with gr.Row(): with gr.Column(scale=1, elem_classes=["section"]): json_out = gr.Textbox(label="JSON report", lines=18) with gr.Column(scale=1, elem_classes=["section"]): dl_img = gr.File(label="Download annotated image") dl_json = gr.File(label="Download JSON report") img_btn.click( ui_image, inputs=[img_in, size_tol, appearance_thr, solidity_thr, seal_frac], outputs=[img_out, json_out, dl_img, dl_json], concurrency_limit=2 # Gradio 4: per-event limit ) with gr.Tab("Video"): with gr.Row(): with gr.Column(scale=1, elem_classes=["section"]): vid_in = gr.Video(label="Upload MP4", height=420) vid_btn = gr.Button("Run inspection", variant="primary") with gr.Column(scale=1, elem_classes=["section"]): vid_out = gr.Video(label="Annotated video", height=420) with gr.Row(): with gr.Column(scale=1, elem_classes=["section"]): vjson_out = gr.Textbox(label="Report (first lines)", lines=18) with gr.Column(scale=1, elem_classes=["section"]): dl_vid = gr.File(label="Download annotated video") dl_jsonl = gr.File(label="Download JSONL report") vid_btn.click( ui_video, inputs=[vid_in, size_tol, appearance_thr, solidity_thr, seal_frac], outputs=[vid_out, vjson_out, dl_vid, dl_jsonl], concurrency_limit=1 # keep video runs serialized ) # In Spaces: no share=True, no host/port tweaks needed. demo.queue(max_size=32).launch()