Spaces:
Build error
Build error
| 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 ---------- | |
| 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(""" | |
| <div id="header"> | |
| <h1>Myraa Lens — Package Inspection</h1> | |
| <p>Upload an image or MP4 to receive annotated output and a structured report.</p> | |
| </div> | |
| """) | |
| 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() | |