import base64 import sys import threading import time from pathlib import Path import cv2 import gradio as gr import numpy as np import torch try: import spaces except ImportError: class spaces: @staticmethod def GPU(*args, **kwargs): def decorator(func): return func return decorator APP_DIR = Path(__file__).resolve().parent CORE_DIR = APP_DIR / "microghost" DEFAULT_MODEL_PATH = APP_DIR / "checkpoints" / "best_microghost_thermal_v3.pth" sys.path.insert(0, str(CORE_DIR)) from inference import ThermalInferenceEngine # noqa: E402 _engine = None _engine_lock = threading.Lock() def get_engine(): global _engine with _engine_lock: if _engine is None: if not DEFAULT_MODEL_PATH.exists(): raise RuntimeError(f"Model checkpoint not found: {DEFAULT_MODEL_PATH}") device = "cuda" if torch.cuda.is_available() else "cpu" _engine = ThermalInferenceEngine(model_path=str(DEFAULT_MODEL_PATH), device=device) return _engine def data_url(image_bgr, ext=".jpg"): ok, encoded = cv2.imencode(ext, image_bgr) if not ok: raise RuntimeError("Failed to encode result image.") payload = base64.b64encode(encoded.tobytes()).decode("ascii") mime = "image/png" if ext == ".png" else "image/jpeg" return f"data:{mime};base64,{payload}" def draw_detections(base_bgr, detections): out = base_bgr.copy() h, w = out.shape[:2] for det in detections: x1, y1, x2, y2 = [int(float(v) * d) for v, d in zip(det["bbox"], [w, h, w, h])] x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(w - 1, x2), min(h - 1, y2) cv2.rectangle(out, (x1, y1), (x2, y2), (31, 214, 128), 2) label = f"{det.get('class', 'target')} {det.get('combined_conf', 0):.2f}" cv2.putText( out, label, (x1, max(18, y1 - 8)), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (31, 214, 128), 2, cv2.LINE_AA, ) return out def normalize_detection(det): return { "class": det.get("class", "unknown"), "confidence": round(float(det.get("combined_conf", det.get("conf", 0.0))), 4), "objectness": round(float(det.get("conf", 0.0)), 4), "temperature_c": float(det.get("temp_c", 0.0)), "laplacian_variance": float(det.get("lap_var", 0.0)), "bbox": [round(float(v), 6) for v in det.get("bbox", [])], "merged_parts": int(det.get("merged_parts", 1)), } def prepare_inputs(rgb_image, thermal_image, conf_thresh): rgb_bgr = None thermal_gray = None if rgb_image is not None: rgb_array = rgb_image.astype(np.uint8) rgb_bgr = cv2.cvtColor(rgb_array, cv2.COLOR_RGB2BGR) if thermal_image is not None: thermal_array = thermal_image.astype(np.uint8) if thermal_array.ndim == 3: thermal_gray = cv2.cvtColor(thermal_array, cv2.COLOR_RGB2GRAY) else: thermal_gray = thermal_array if rgb_bgr is None and thermal_gray is None: raise gr.Error("Upload an RGB image, a thermal image, or both.") if rgb_bgr is not None and thermal_gray is not None: return { "mode": "paired", "model_rgb": cv2.cvtColor(rgb_bgr, cv2.COLOR_BGR2RGB), "model_thermal": thermal_gray, "lap_image": cv2.cvtColor(rgb_bgr, cv2.COLOR_BGR2RGB), "primary_bgr": rgb_bgr, "effective_conf": None if conf_thresh <= 0 else conf_thresh, } if thermal_gray is not None: h, w = thermal_gray.shape[:2] return { "mode": "thermal_only", "model_rgb": np.zeros((h, w, 3), dtype=np.uint8), "model_thermal": thermal_gray, "lap_image": thermal_gray, "primary_bgr": cv2.cvtColor(thermal_gray, cv2.COLOR_GRAY2BGR), "effective_conf": 0.20 if conf_thresh <= 0 else conf_thresh, } h, w = rgb_bgr.shape[:2] return { "mode": "rgb_only", "model_rgb": cv2.cvtColor(rgb_bgr, cv2.COLOR_BGR2RGB), "model_thermal": np.zeros((h, w), dtype=np.uint8), "lap_image": cv2.cvtColor(rgb_bgr, cv2.COLOR_BGR2RGB), "primary_bgr": rgb_bgr, "effective_conf": 0.20 if conf_thresh <= 0 else conf_thresh, } @spaces.GPU(duration=60) def gradio_analyze(rgb_image, thermal_image, conf_thresh, lap_thresh): started = time.perf_counter() inputs = prepare_inputs(rgb_image, thermal_image, float(conf_thresh),) engine = get_engine() detections = engine.detect_confirmed( inputs["model_rgb"], inputs["model_thermal"], lap_image=inputs["lap_image"], lap_thresh=0 if inputs["mode"] == "thermal_only" else float(lap_thresh), conf_threshold=inputs["effective_conf"], ) if inputs["mode"] == "thermal_only": detections = engine.filter_thermal_only_artifacts(detections) thermal_bgr = cv2.applyColorMap(inputs["model_thermal"], cv2.COLORMAP_JET) annotated_primary = draw_detections(inputs["primary_bgr"], detections) annotated_thermal = draw_detections(thermal_bgr, detections) return { "ok": True, "mode": inputs["mode"], "count": len(detections), "elapsed_ms": int((time.perf_counter() - started) * 1000), "thresholds": { "confidence": inputs["effective_conf"], "laplacian": float(lap_thresh), "lap_bypass_confidence": None, }, "detections": [normalize_detection(det) for det in detections], "images": { "annotated_primary": data_url(annotated_primary), "annotated_thermal": data_url(annotated_thermal), }, } with gr.Blocks(title="MicroGhost Thermal Inference") as demo: gr.Markdown("# MicroGhost Thermal Inference") gr.Markdown("Upload RGB, thermal, or both. Use this Space directly, or call it from the Vercel app.") with gr.Row(): rgb_input = gr.Image(label="RGB image", type="numpy", image_mode="RGB") thermal_input = gr.Image(label="Thermal image", type="numpy", image_mode="L") with gr.Accordion("Advanced tuning", open=False): conf_input = gr.Slider(0, 0.9, value=0, step=0.01, label="Confidence override (0 = automatic)") lap_input = gr.Slider(0, 220, value=80, step=5, label="Laplacian threshold") analyze_button = gr.Button("Analyze", variant="primary") json_output = gr.JSON(label="Result") analyze_button.click( gradio_analyze, inputs=[rgb_input, thermal_input, conf_input, lap_input], outputs=[json_output], api_name="gradio_analyze", ) if __name__ == "__main__": demo.queue().launch(server_name="0.0.0.0", server_port=7860)