import os
import sys
import time
import numpy as np
from PIL import Image
import streamlit as st
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
# ── Page config (must be first Streamlit call) ──────────────────────────────
st.set_page_config(
page_title="PatternScan — BOM Detection",
page_icon="⚡",
layout="wide",
initial_sidebar_state="expanded",
)
# ── Global CSS ───────────────────────────────────────────────────────────────
st.markdown("""
""", unsafe_allow_html=True)
# ── Pipeline loader ──────────────────────────────────────────────────────────
@st.cache_resource(show_spinner="Loading DINOv2 ViT-S/14 …")
def _load_pipeline():
from src.pipeline import PatternDetectionPipeline
return PatternDetectionPipeline()
try:
pipeline = _load_pipeline()
_model_ok = True
except Exception as _e:
_model_ok = False
_model_err = str(_e)
# ============================================================
# SIDEBAR
# ============================================================
with st.sidebar:
# Logo
st.markdown("""
⚡ PatternScan
BOM Engineering Drawings
""", unsafe_allow_html=True)
st.divider()
# Navigation
st.markdown("""
📐 Detection
📊 History
⚙ Settings
""", unsafe_allow_html=True)
st.divider()
# Algorithm settings
st.markdown('Algorithm', unsafe_allow_html=True)
manual_mode = st.toggle("Manual mode", value=False,
help="Override automatic threshold selection")
if manual_mode:
ncc_thr = st.slider("NCC Threshold", 0.15, 0.9, 0.28, 0.01,
help="Lower → higher recall, more false positives")
dino_thr = st.slider("DINOv2 Threshold", 0.50, 0.95, 0.84, 0.01,
help="Higher → fewer false positives")
dilate = st.slider("Stroke Dilation", 0, 9, 0, 1,
help="Thicken template strokes for bold drawings")
else:
ncc_thr, dino_thr, dilate = 0.28, 0.84, 0
st.divider()
# Model status
if _model_ok:
st.markdown('● Model Ready', unsafe_allow_html=True)
st.caption("DINOv2 ViT-S/14 · NCC multi-scale")
else:
st.markdown('● Error', unsafe_allow_html=True)
st.caption(_model_err[:80])
st.divider()
st.caption("Zero-shot · No training required")
st.caption("NCC + DINOv2 + Containment NMS")
# ============================================================
# MAIN CONTENT
# ============================================================
# Page header
st.markdown("""
Zero-Shot Pattern Detection
Locate component symbols in BOM engineering drawings —
no training data required. Upload a pattern and drawing, then click Run.
""", unsafe_allow_html=True)
# ── Upload row ───────────────────────────────────────────────────────────────
up1, up2 = st.columns(2, gap="large")
with up1:
st.markdown('Pattern Symbol', unsafe_allow_html=True)
pat_file = st.file_uploader(
"pattern", type=["png", "jpg", "jpeg", "tif", "tiff"],
label_visibility="collapsed", key="pat_up",
)
if pat_file:
pil_p = Image.open(pat_file).convert("RGB")
st.image(pil_p, caption=f"{pil_p.width} × {pil_p.height} px",
use_container_width=True)
else:
st.markdown("""
📂 PNG · JPG · TIFF
""", unsafe_allow_html=True)
with up2:
st.markdown('Engineering Drawing', unsafe_allow_html=True)
drw_file = st.file_uploader(
"drawing", type=["png", "jpg", "jpeg", "tif", "tiff"],
label_visibility="collapsed", key="drw_up",
)
if drw_file:
pil_d = Image.open(drw_file).convert("RGB")
st.image(pil_d, caption=f"{pil_d.width} × {pil_d.height} px",
use_container_width=True)
else:
st.markdown("""
📂 PNG · JPG · TIFF
""", unsafe_allow_html=True)
st.markdown("
", unsafe_allow_html=True)
# ── Run button ───────────────────────────────────────────────────────────────
btn_col, _ = st.columns([1, 2])
with btn_col:
run_clicked = st.button(
"⚡ Run Detection",
type="primary",
use_container_width=True,
disabled=not _model_ok,
)
# ── Detection logic ──────────────────────────────────────────────────────────
if run_clicked:
if pat_file is None or drw_file is None:
st.warning("⚠ Upload both a **pattern image** and an **engineering drawing** first.")
else:
with st.spinner("Running detection pipeline …"):
try:
pat_arr = np.array(Image.open(pat_file).convert("RGB"))
drw_arr = np.array(Image.open(drw_file).convert("RGB"))
t0 = time.time()
if manual_mode:
pipeline.update_thresholds(
ncc_threshold=ncc_thr,
cosine_threshold=dino_thr,
)
pipeline.dilate_pattern = dilate
result = pipeline.detect(pat_arr, drw_arr, return_visualization=True)
else:
result = pipeline.detect_auto(pat_arr, drw_arr, return_visualization=True)
elapsed = time.time() - t0
vis = result.pop("visualization", None)
dets = result.get("detections", [])
n = result["total_detections"]
st.session_state.update({
"result": result, "vis": vis, "elapsed": elapsed,
"n": n, "dets": dets,
})
except Exception as ex:
st.error(f"Detection failed: {ex}")
st.session_state["result"] = None
# ── Results section ───────────────────────────────────────────────────────────
if st.session_state.get("result") is not None:
result = st.session_state["result"]
vis = st.session_state["vis"]
elapsed = st.session_state["elapsed"]
n = st.session_state["n"]
dets = st.session_state["dets"]
st.markdown("---")
# ── Stat cards ──
st.markdown('Detection Summary', unsafe_allow_html=True)
s1, s2, s3, s4 = st.columns(4, gap="medium")
with s1:
status_lbl = "Detected" if n > 0 else "Not found"
st.metric("Instances Found", n)
with s2:
avg_c = round(sum(d["confidence"] for d in dets) / n, 2) if n else 0.0
st.metric("Avg Confidence", f"{avg_c:.2f}")
with s3:
best_d = round(max(d["dino_score"] for d in dets), 4) if n else 0.0
st.metric("Best DINO Score", f"{best_d:.4f}")
with s4:
st.metric("Processing Time", f"{elapsed:.1f} s")
st.markdown("
", unsafe_allow_html=True)
# ── Output image + detection list ──
out1, out2 = st.columns([3, 1], gap="large")
with out1:
st.markdown('Annotated Output', unsafe_allow_html=True)
if vis is not None:
st.image(vis, use_container_width=True,
caption=f"{n} instance(s) detected · {elapsed:.1f} s · "
f"NCC + DINOv2 ViT-S/14")
else:
st.info("No visualization available.")
with out2:
st.markdown('Detections', unsafe_allow_html=True)
if not dets:
st.markdown("""
No patterns detected
""", unsafe_allow_html=True)
else:
for i, det in enumerate(dets):
bb = det["bbox"]
conf = float(det["confidence"])
if conf >= 0.70:
ccolor, cbadge = "#15803D", "badge-green"
elif conf >= 0.55:
ccolor, cbadge = "#B45309", "badge-amber"
else:
ccolor, cbadge = "#DC2626", "badge-red"
st.markdown(f"""
Detection #{i + 1}
({bb['x']}, {bb['y']}) {bb['w']} × {bb['h']} px
{conf:.2f}
NCC {det['ncc_score']:.3f}
DINO {det['dino_score']:.4f}
Scale {det['scale']:.2f}×
""", unsafe_allow_html=True)
# Raw JSON
st.markdown("
", unsafe_allow_html=True)
with st.expander("📄 Raw JSON output"):
st.json(result)