Create app.py

app.py

@@ -0,0 +1,267 @@
+"""
+STL Firearm/Knife Screener
+---------------------------
+Upload an STL file. The app rotates a virtual camera fully around the model
+(covering all azimuths, plus dedicated top-down and bottom-up views), renders
+each view as a 2D image, and runs the "Guns-100-11m" YOLOv11 detector
+(wuhp/guns-100-11m) on every rendered frame. Results are aggregated into an
+overall verdict plus an annotated gallery so a detection that only shows up
+from one angle (e.g. the trigger guard only visible from the side) still
+gets caught.
+
+Run locally:
+    pip install -r requirements.txt
+    python app.py
+
+Model weights are pulled automatically from the Hugging Face Hub the first
+time the app runs (needs internet access) and cached locally afterwards.
+"""
+
+import os
+import tempfile
+
+import gradio as gr
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+from matplotlib.colors import LightSource
+import numpy as np
+import pandas as pd
+import trimesh
+from PIL import Image, ImageDraw, ImageFont
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+
+# --------------------------------------------------------------------------
+# Model loading
+# --------------------------------------------------------------------------
+
+MODEL_REPO = "wuhp/guns-100-11m"
+MODEL_FILE = "Guns-100-11m.pt"
+CLASS_NAMES = {0: "Gun", 1: "Knife"}
+
+_model = None
+
+
+def get_model():
+    """Lazily download (once) and cache the YOLO model."""
+    global _model
+    if _model is None:
+        weights_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+        _model = YOLO(weights_path)
+    return _model
+
+
+# --------------------------------------------------------------------------
+# STL -> multi-view rendering
+# --------------------------------------------------------------------------
+
+def build_view_angles(num_azimuths: int, num_elevations: int):
+    """
+    Build a list of (azimuth, elevation) pairs that fully cover the sphere
+    around the object, including explicit top-down and bottom-up shots.
+    """
+    azimuths = np.linspace(0, 360, num_azimuths, endpoint=False)
+
+    # Elevation ring between (but not including) the poles, e.g. for
+    # num_elevations=3 -> [-45, 0, 45]
+    if num_elevations > 0:
+        elevations = np.linspace(-75, 75, num_elevations)
+    else:
+        elevations = np.array([0.0])
+
+    views = []
+    for elev in elevations:
+        for az in azimuths:
+            views.append((float(az), float(elev)))
+
+    # Dedicated poles: straight down (top view) and straight up (bottom view)
+    views.append((0.0, 90.0))   # top
+    views.append((0.0, -90.0))  # bottom
+
+    return views
+
+
+def render_view(tris: np.ndarray, center: np.ndarray, scale: float,
+                 azim: float, elev: float, img_size: int = 640) -> Image.Image:
+    """Render a single (azim, elev) view of the mesh triangles to a PIL image."""
+    fig = plt.figure(figsize=(img_size / 100, img_size / 100), dpi=100)
+    ax = fig.add_subplot(111, projection="3d")
+    ax.set_facecolor("white")
+
+    ls = LightSource(azdeg=315, altdeg=55)
+    facecolors = [(0.62, 0.62, 0.66, 1.0)] * len(tris)
+    poly = Poly3DCollection(
+        tris,
+        facecolors=facecolors,
+        edgecolor=(0.15, 0.15, 0.15, 0.4),
+        linewidths=0.15,
+        shade=True,
+        lightsource=ls,
+    )
+    ax.add_collection3d(poly)
+
+    ax.set_xlim(center[0] - scale, center[0] + scale)
+    ax.set_ylim(center[1] - scale, center[1] + scale)
+    ax.set_zlim(center[2] - scale, center[2] + scale)
+    try:
+        ax.set_box_aspect((1, 1, 1))
+    except Exception:
+        pass
+    ax.set_axis_off()
+    ax.view_init(elev=elev, azim=azim)
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
+
+    fig.canvas.draw()
+    buf = np.asarray(fig.canvas.buffer_rgba())
+    img = Image.fromarray(buf).convert("RGB")
+    plt.close(fig)
+    return img
+
+
+def load_mesh_triangles(stl_path: str):
+    mesh = trimesh.load(stl_path, force="mesh")
+    if not isinstance(mesh, trimesh.Trimesh):
+        # Scene with multiple geometries -> merge them
+        mesh = trimesh.util.concatenate(mesh.dump())
+    tris = mesh.vertices[mesh.faces]
+    center = mesh.centroid
+    scale = float(np.max(mesh.extents)) / 2.0 * 1.05  # small margin
+    return tris, center, scale, mesh
+
+
+# --------------------------------------------------------------------------
+# Detection + aggregation
+# --------------------------------------------------------------------------
+
+def annotate_image(img: Image.Image, boxes, label_prefix=""):
+    """Draw YOLO detection boxes on a copy of the image."""
+    out = img.copy()
+    draw = ImageDraw.Draw(out)
+    try:
+        font = ImageFont.load_default()
+    except Exception:
+        font = None
+
+    for cls_id, conf, xyxy in boxes:
+        x1, y1, x2, y2 = xyxy
+        color = (220, 30, 30) if cls_id == 0 else (30, 120, 220)
+        draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
+        label = f"{CLASS_NAMES.get(cls_id, cls_id)} {conf:.2f}"
+        draw.text((x1 + 3, max(0, y1 - 14)), label, fill=color, font=font)
+    return out
+
+
+def process_stl(stl_file, num_azimuths, num_elevations, conf_threshold, img_size,
+                 progress=gr.Progress()):
+    if stl_file is None:
+        return None, "Please upload an STL file.", None
+
+    progress(0, desc="Loading model...")
+    model = get_model()
+
+    progress(0.05, desc="Loading mesh...")
+    tris, center, scale, mesh = load_mesh_triangles(stl_file)
+
+    views = build_view_angles(int(num_azimuths), int(num_elevations))
+    total = len(views)
+
+    all_detections = []          # rows for the results table
+    annotated_images = []        # (image, caption) for the gallery
+    max_conf_per_class = {0: 0.0, 1: 0.0}
+    any_hit = False
+
+    for i, (az, el) in enumerate(views):
+        progress((0.1 + 0.85 * i / total), desc=f"Rendering + detecting view {i+1}/{total}...")
+        frame = render_view(tris, center, scale, az, el, img_size=int(img_size))
+
+        result = model.predict(source=np.array(frame), conf=float(conf_threshold), verbose=False)[0]
+
+        boxes = []
+        for b in result.boxes:
+            cls_id = int(b.cls.item())
+            conf = float(b.conf.item())
+            xyxy = [float(v) for v in b.xyxy[0].tolist()]
+            boxes.append((cls_id, conf, xyxy))
+            max_conf_per_class[cls_id] = max(max_conf_per_class[cls_id], conf)
+            any_hit = True
+            all_detections.append({
+                "View": f"az={az:.0f} el={el:.0f}",
+                "Class": CLASS_NAMES.get(cls_id, str(cls_id)),
+                "Confidence": round(conf, 3),
+            })
+
+        if boxes:
+            annotated = annotate_image(frame, boxes)
+            caption = f"az={az:.0f}, el={el:.0f} | " + ", ".join(
+                f"{CLASS_NAMES.get(c, c)} {p:.2f}" for c, p, _ in boxes
+            )
+            annotated_images.append((annotated, caption))
+
+    progress(1.0, desc="Done")
+
+    # Build verdict
+    if any_hit:
+        parts = []
+        if max_conf_per_class[0] > 0:
+            parts.append(f"GUN detected (max confidence {max_conf_per_class[0]:.2f})")
+        if max_conf_per_class[1] > 0:
+            parts.append(f"KNIFE detected (max confidence {max_conf_per_class[1]:.2f})")
+        verdict = "\u26a0\ufe0f WEAPON-LIKE OBJECT DETECTED: " + " | ".join(parts)
+        verdict += f"\n\nDetected in {len(annotated_images)} of {total} rendered views."
+    else:
+        verdict = f"\u2705 No firearm or knife detected across {total} rendered views."
+
+    verdict += ("\n\nNote: this is a heuristic screen based on rendered silhouettes "
+                "of the mesh, not a photo of a real object. Treat results as advisory, "
+                "not a certified determination.")
+
+    if not annotated_images:
+        # still show a couple of representative views so the user sees *something*
+        sample_idxs = [0, total // 2] if total > 1 else [0]
+        for idx in sample_idxs:
+            az, el = views[idx]
+            frame = render_view(tris, center, scale, az, el, img_size=int(img_size))
+            annotated_images.append((frame, f"az={az:.0f}, el={el:.0f} (no detection)"))
+
+    df = pd.DataFrame(all_detections) if all_detections else pd.DataFrame(
+        columns=["View", "Class", "Confidence"])
+
+    return annotated_images, verdict, df
+
+
+# --------------------------------------------------------------------------
+# Gradio UI
+# --------------------------------------------------------------------------
+
+with gr.Blocks(title="STL Firearm/Knife Screener") as demo:
+    gr.Markdown(
+        "# 🔫🔪 STL Firearm / Knife Screener\n"
+        "Upload a 3D model (`.stl`). The app spins a virtual camera all the way "
+        "around it — including straight-down and straight-up shots — renders each "
+        "view, and runs the **Guns-100-11m** YOLOv11 detector on every frame."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            stl_input = gr.File(label="Upload STL file", file_types=[".stl"], type="filepath")
+            num_azimuths = gr.Slider(4, 16, value=8, step=1, label="Azimuth steps per elevation ring (more = finer 360° coverage, slower)")
+            num_elevations = gr.Slider(1, 5, value=3, step=1, label="Elevation rings between the poles")
+            conf_threshold = gr.Slider(0.05, 0.95, value=0.25, step=0.05, label="Detection confidence threshold")
+            img_size = gr.Slider(320, 960, value=640, step=32, label="Render resolution (px)")
+            run_btn = gr.Button("Run 360° Scan", variant="primary")
+
+        with gr.Column(scale=2):
+            verdict_out = gr.Textbox(label="Verdict", lines=4)
+            gallery_out = gr.Gallery(label="Views with detections", columns=3, height=500)
+            table_out = gr.Dataframe(label="All detections", wrap=True)
+
+    run_btn.click(
+        fn=process_stl,
+        inputs=[stl_input, num_azimuths, num_elevations, conf_threshold, img_size],
+        outputs=[gallery_out, verdict_out, table_out],
+    )
+
+if __name__ == "__main__":
+    demo.queue().launch()