Duy commited on
Commit ·
7d835bb
0
Parent(s):
feat: Zero-shot pattern detection pipeline for engineering BOM drawings
Browse filesPipeline: NCC multi-scale (Stage 1) + DINOv2 ViT-S/14 (Stage 2) + structural filters
- Auto-classify templates as simple (outline, e.g. resistor) or complex (e.g. bridge rectifier)
- Simple path: Chamfer pre-filter + hybrid DINOv2 bypass for borderline candidates
- Complex path: dual micro-pass (0deg + 90deg) with DINOv2 derotate=True
- FastAPI web UI with PaceOS-style dashboard
- HuggingFace Spaces Docker deployment ready
Results: Drawing 1 (resistor) 10 detections, Drawing 4 (BR) 4 detections,
Example 1 10 detections, Example 2 4 detections. All 10 unit tests pass.
- .gitignore +73 -0
- Dockerfile +35 -0
- README.md +157 -0
- app/__init__.py +0 -0
- app/app.py +27 -0
- app/streamlit_app.py +439 -0
- app/web/__init__.py +0 -0
- app/web/index.html +225 -0
- app/web/server.py +145 -0
- app/web/static/css/style.css +540 -0
- app/web/static/js/app.js +316 -0
- conftest.py +5 -0
- design_spec/system_design.md +316 -0
- examples/example1_drawing.png +0 -0
- examples/example1_pattern.png +0 -0
- examples/example2_drawing.png +0 -0
- examples/example2_pattern.png +0 -0
- requirements.txt +12 -0
- scripts/generate_examples.py +114 -0
- scripts/tune_thresholds.py +106 -0
- src/__init__.py +0 -0
- src/dino_verifier.py +265 -0
- src/ncc_matcher.py +165 -0
- src/pipeline.py +443 -0
- src/postprocessor.py +976 -0
- src/preprocessor.py +174 -0
- tests/__init__.py +0 -0
- tests/test_pipeline.py +203 -0
.gitignore
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*$py.class
|
| 5 |
+
*.so
|
| 6 |
+
.Python
|
| 7 |
+
build/
|
| 8 |
+
dist/
|
| 9 |
+
*.egg-info/
|
| 10 |
+
.eggs/
|
| 11 |
+
*.egg
|
| 12 |
+
|
| 13 |
+
# Virtual environments
|
| 14 |
+
venv/
|
| 15 |
+
env/
|
| 16 |
+
.venv/
|
| 17 |
+
.env
|
| 18 |
+
|
| 19 |
+
# Model weights & large files
|
| 20 |
+
*.pt
|
| 21 |
+
*.pth
|
| 22 |
+
*.bin
|
| 23 |
+
*.onnx
|
| 24 |
+
*.h5
|
| 25 |
+
*.pkl
|
| 26 |
+
checkpoints/
|
| 27 |
+
weights/
|
| 28 |
+
models/
|
| 29 |
+
|
| 30 |
+
# IDE
|
| 31 |
+
.vscode/
|
| 32 |
+
.idea/
|
| 33 |
+
*.swp
|
| 34 |
+
*.swo
|
| 35 |
+
|
| 36 |
+
# OS
|
| 37 |
+
.DS_Store
|
| 38 |
+
Thumbs.db
|
| 39 |
+
|
| 40 |
+
# Debug output
|
| 41 |
+
debug_output/
|
| 42 |
+
|
| 43 |
+
# Jupyter
|
| 44 |
+
.ipynb_checkpoints/
|
| 45 |
+
*.ipynb
|
| 46 |
+
|
| 47 |
+
# Logs
|
| 48 |
+
*.log
|
| 49 |
+
logs/
|
| 50 |
+
|
| 51 |
+
# HuggingFace cache
|
| 52 |
+
.cache/
|
| 53 |
+
|
| 54 |
+
# Development / debug scripts (not part of the submission)
|
| 55 |
+
debug_*.py
|
| 56 |
+
_debug*.py
|
| 57 |
+
_test*.py
|
| 58 |
+
check_*.py
|
| 59 |
+
run_debug*.py
|
| 60 |
+
run_web.py
|
| 61 |
+
|
| 62 |
+
# Debug output images and crops in project root
|
| 63 |
+
*.png
|
| 64 |
+
!examples/*.png
|
| 65 |
+
|
| 66 |
+
# Debug text output
|
| 67 |
+
*_output.txt
|
| 68 |
+
|
| 69 |
+
# Temp pipeline files
|
| 70 |
+
*.tmp.*
|
| 71 |
+
|
| 72 |
+
# Test suite cache
|
| 73 |
+
.pytest_cache/
|
Dockerfile
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
FROM python:3.11-slim
|
| 2 |
+
|
| 3 |
+
WORKDIR /app
|
| 4 |
+
|
| 5 |
+
# System deps for OpenCV headless
|
| 6 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 7 |
+
libglib2.0-0 \
|
| 8 |
+
libgl1-mesa-glx \
|
| 9 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 10 |
+
|
| 11 |
+
# Copy requirements early for layer caching
|
| 12 |
+
COPY requirements.txt .
|
| 13 |
+
|
| 14 |
+
# Install CPU-only PyTorch (much lighter than CUDA build; HF free tier is CPU-only)
|
| 15 |
+
RUN pip install --no-cache-dir \
|
| 16 |
+
"torch==2.1.0+cpu" "torchvision==0.16.0+cpu" \
|
| 17 |
+
--extra-index-url https://download.pytorch.org/whl/cpu
|
| 18 |
+
|
| 19 |
+
# Install remaining dependencies (FastAPI, OpenCV, Pillow, etc.)
|
| 20 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 21 |
+
|
| 22 |
+
# Copy project source
|
| 23 |
+
COPY . .
|
| 24 |
+
|
| 25 |
+
# Pre-download DINOv2 ViT-S/14 weights at build time (~86 MB) to avoid
|
| 26 |
+
# cold-start delay on first request in HuggingFace Spaces.
|
| 27 |
+
RUN python -c "import torch; torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')" || true
|
| 28 |
+
|
| 29 |
+
# HuggingFace Spaces requires port 7860
|
| 30 |
+
EXPOSE 7860
|
| 31 |
+
|
| 32 |
+
ENV KMP_DUPLICATE_LIB_OK=TRUE \
|
| 33 |
+
PORT=7860
|
| 34 |
+
|
| 35 |
+
CMD ["python", "-m", "uvicorn", "app.web.server:app", "--host", "0.0.0.0", "--port", "7860"]
|
README.md
ADDED
|
@@ -0,0 +1,157 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: BOM Pattern Detection
|
| 3 |
+
emoji: 🔍
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: green
|
| 6 |
+
sdk: docker
|
| 7 |
+
pinned: false
|
| 8 |
+
app_port: 7860
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
# Zero-Shot Pattern Detection for Engineering BOM Drawings
|
| 12 |
+
|
| 13 |
+

|
| 14 |
+

|
| 15 |
+
[](https://huggingface.co/spaces/PLACEHOLDER)
|
| 16 |
+
|
| 17 |
+
## Overview
|
| 18 |
+
|
| 19 |
+
This project solves the problem of finding all occurrences of a given pattern (e.g., a component symbol) inside large engineering BOM drawings — with **zero training data**. The pipeline combines classical NCC template matching for fast candidate proposal with DINOv2 ViT-S/14 self-supervised features for accurate zero-shot verification. No fine-tuning or labeled data is required: any pattern can be detected at inference time.
|
| 20 |
+
|
| 21 |
+
## Pipeline Architecture
|
| 22 |
+
|
| 23 |
+
```
|
| 24 |
+
[Pattern Image] ──┐
|
| 25 |
+
├──► Stage 0: Preprocess ──► Stage 1: NCC Multi-scale
|
| 26 |
+
[Drawing Image] ──┘ (binarize, (propose 30–200
|
| 27 |
+
denoise) candidates, CPU)
|
| 28 |
+
│
|
| 29 |
+
▼
|
| 30 |
+
Stage 2: DINOv2 Verify
|
| 31 |
+
(cosine similarity filter,
|
| 32 |
+
zero-shot, CPU/GPU)
|
| 33 |
+
│
|
| 34 |
+
▼
|
| 35 |
+
Stage 4: Post-process
|
| 36 |
+
(IoU-NMS, format JSON,
|
| 37 |
+
draw bounding boxes)
|
| 38 |
+
│
|
| 39 |
+
▼
|
| 40 |
+
[BBoxes + Score JSON]
|
| 41 |
+
```
|
| 42 |
+
|
| 43 |
+
| Stage | Module | Purpose | Typical Time |
|
| 44 |
+
|-------|--------|---------|-------------|
|
| 45 |
+
| 0 | Preprocessor | Adaptive binarize, denoise | < 0.5s |
|
| 46 |
+
| 1 | NCCMatcher | Multi-scale + rotation candidate proposal | 5–15s CPU |
|
| 47 |
+
| 2 | DINOVerifier | Zero-shot cosine similarity filter | 5–15s CPU |
|
| 48 |
+
| 4 | Postprocessor | Final NMS, JSON output, visualization | < 0.1s |
|
| 49 |
+
|
| 50 |
+
## Installation
|
| 51 |
+
|
| 52 |
+
```bash
|
| 53 |
+
git clone https://github.com/YOUR_USERNAME/pattern-detection-bom.git
|
| 54 |
+
cd pattern-detection-bom
|
| 55 |
+
pip install -r requirements.txt
|
| 56 |
+
```
|
| 57 |
+
|
| 58 |
+
DINOv2 weights are downloaded automatically at first run via `torch.hub` — no manual download needed.
|
| 59 |
+
|
| 60 |
+
## Quick Start
|
| 61 |
+
|
| 62 |
+
### Python API
|
| 63 |
+
|
| 64 |
+
```python
|
| 65 |
+
from src.pipeline import run_detection
|
| 66 |
+
|
| 67 |
+
result = run_detection("pattern.png", "drawing.png")
|
| 68 |
+
print(result)
|
| 69 |
+
```
|
| 70 |
+
|
| 71 |
+
### Web UI (local)
|
| 72 |
+
|
| 73 |
+
```bash
|
| 74 |
+
python app/app.py
|
| 75 |
+
```
|
| 76 |
+
|
| 77 |
+
Then open `http://localhost:8000` in your browser. The dashboard lets you upload a pattern and drawing, adjust detection thresholds via sliders, and view annotated results.
|
| 78 |
+
|
| 79 |
+
### CLI Tuning
|
| 80 |
+
|
| 81 |
+
```bash
|
| 82 |
+
python scripts/tune_thresholds.py --examples_dir examples --quick
|
| 83 |
+
```
|
| 84 |
+
|
| 85 |
+
## Output Format
|
| 86 |
+
|
| 87 |
+
```json
|
| 88 |
+
{
|
| 89 |
+
"detections": [
|
| 90 |
+
{
|
| 91 |
+
"bbox": {"x": 142, "y": 310, "w": 64, "h": 64},
|
| 92 |
+
"confidence": 0.83,
|
| 93 |
+
"ncc_score": 0.7812,
|
| 94 |
+
"dino_score": 0.8754,
|
| 95 |
+
"scale": 1.0,
|
| 96 |
+
"angle": 0.0
|
| 97 |
+
}
|
| 98 |
+
],
|
| 99 |
+
"total_detections": 3,
|
| 100 |
+
"image_size": {"width": 2480, "height": 3508}
|
| 101 |
+
}
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
## Configuration
|
| 105 |
+
|
| 106 |
+
| Parameter | Default | Description |
|
| 107 |
+
|-----------|---------|-------------|
|
| 108 |
+
| `scales` | `[0.85..1.15]` | Scale range for template resizing in Stage 1 |
|
| 109 |
+
| `angles` | `[-10..10]` | Rotation sweep in Stage 1 (degrees) |
|
| 110 |
+
| `ncc_threshold` | `0.55` | Minimum NCC score to propose a candidate (low = high recall) |
|
| 111 |
+
| `nms_iou_threshold` | `0.3` | IoU threshold for NMS in Stage 1 |
|
| 112 |
+
| `dino_model` | `dinov2_vits14` | DINOv2 variant: `vits14` (fast) or `vitb14` (accurate) |
|
| 113 |
+
| `cosine_threshold` | `0.84` | Minimum cosine similarity to accept a detection |
|
| 114 |
+
| `final_nms_iou` | `0.4` | Final IoU threshold after Stage 2 |
|
| 115 |
+
|
| 116 |
+
Override via config dict:
|
| 117 |
+
|
| 118 |
+
```python
|
| 119 |
+
from src.pipeline import PatternDetectionPipeline
|
| 120 |
+
|
| 121 |
+
pipeline = PatternDetectionPipeline(config={
|
| 122 |
+
"ncc_threshold": 0.60,
|
| 123 |
+
"cosine_threshold": 0.80,
|
| 124 |
+
"dino_model": "dinov2_vitb14",
|
| 125 |
+
})
|
| 126 |
+
result = pipeline.detect("pattern.png", "drawing.png")
|
| 127 |
+
```
|
| 128 |
+
|
| 129 |
+
## Approach & Design Choices
|
| 130 |
+
|
| 131 |
+
### Why DINOv2 instead of a fine-tuned model?
|
| 132 |
+
|
| 133 |
+
DINOv2 is trained self-supervised on 142M images and produces patch-level features that capture geometric structure rather than texture or color. Engineering drawings (line art on white background) are a significant domain shift from natural photos, yet DINOv2 generalizes well because its features encode shape primitives. Crucially, **no fine-tuning means truly zero-shot operation** — the model works for any new pattern without collecting labeled data.
|
| 134 |
+
|
| 135 |
+
### Why hybrid NCC + DINOv2?
|
| 136 |
+
|
| 137 |
+
- **NCC alone** is fast but brittle: it fails under lighting variation, noise, and even small domain differences.
|
| 138 |
+
- **DINOv2 alone** (brute-force sliding window) over a large drawing would be prohibitively slow on CPU.
|
| 139 |
+
- The hybrid approach uses NCC for high-recall candidate generation (~30–200 boxes), then DINOv2 as an intelligent filter. This gives the accuracy of a neural verifier at a fraction of the cost.
|
| 140 |
+
|
| 141 |
+
### Trade-offs
|
| 142 |
+
|
| 143 |
+
- NCC threshold is intentionally low (0.55) to maximize recall at Stage 1; false positives are cleaned up by DINOv2.
|
| 144 |
+
- ViT-S/14 (21M params) is chosen over ViT-B/14 for speed on CPU deployments; swap to `vitb14` for higher accuracy.
|
| 145 |
+
- The pipeline runs on CPU in ~20–40s for typical A3 drawings; GPU reduces Stage 2 to 2–5s.
|
| 146 |
+
|
| 147 |
+
## Limitations & Future Work
|
| 148 |
+
|
| 149 |
+
- **Small templates (< 28px):** DINOv2 Stage 2 is skipped for very small patterns; only NCC is used.
|
| 150 |
+
- **Heavy rotation (> 15°):** The sweep is limited to ±10° by default; large rotations require wider sweep at extra cost.
|
| 151 |
+
- **Identical-looking components:** Two structurally similar but semantically different symbols may confuse the cosine verifier.
|
| 152 |
+
- **Throughput:** Not optimized for batch processing of many drawings; a production system would benefit from ONNX export and TensorRT.
|
| 153 |
+
- **Optional Stage 3 (LightGlue):** Geometric keypoint verification for handling significant rotation/scale variation is planned but not yet integrated.
|
| 154 |
+
|
| 155 |
+
## License
|
| 156 |
+
|
| 157 |
+
MIT License — see [LICENSE](LICENSE).
|
app/__init__.py
ADDED
|
File without changes
|
app/app.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Entry point — launches the FastAPI web UI."""
|
| 2 |
+
import os
|
| 3 |
+
import sys
|
| 4 |
+
|
| 5 |
+
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
|
| 6 |
+
|
| 7 |
+
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
| 8 |
+
sys.path.insert(0, ROOT)
|
| 9 |
+
|
| 10 |
+
# Import from absolute path to avoid 'app' package conflict
|
| 11 |
+
import importlib.util, pathlib
|
| 12 |
+
|
| 13 |
+
_server_path = pathlib.Path(ROOT) / "app" / "web" / "server.py"
|
| 14 |
+
_spec = importlib.util.spec_from_file_location("web_server", _server_path)
|
| 15 |
+
_mod = importlib.util.module_from_spec(_spec)
|
| 16 |
+
_spec.loader.exec_module(_mod)
|
| 17 |
+
|
| 18 |
+
fastapi_app = _mod.app
|
| 19 |
+
get_pipeline = _mod.get_pipeline
|
| 20 |
+
|
| 21 |
+
import uvicorn
|
| 22 |
+
|
| 23 |
+
if __name__ == "__main__":
|
| 24 |
+
print("[BOM Detector] Loading pipeline...")
|
| 25 |
+
get_pipeline()
|
| 26 |
+
print("[BOM Detector] → http://localhost:8000")
|
| 27 |
+
uvicorn.run(fastapi_app, host="0.0.0.0", port=8000, log_level="info")
|
app/streamlit_app.py
ADDED
|
@@ -0,0 +1,439 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import sys
|
| 3 |
+
import time
|
| 4 |
+
import numpy as np
|
| 5 |
+
from PIL import Image
|
| 6 |
+
import streamlit as st
|
| 7 |
+
|
| 8 |
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
| 9 |
+
|
| 10 |
+
# ── Page config (must be first Streamlit call) ──────────────────────────────
|
| 11 |
+
st.set_page_config(
|
| 12 |
+
page_title="PatternScan — BOM Detection",
|
| 13 |
+
page_icon="⚡",
|
| 14 |
+
layout="wide",
|
| 15 |
+
initial_sidebar_state="expanded",
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
# ── Global CSS ───────────────────────────────────────────────────────────────
|
| 19 |
+
st.markdown("""
|
| 20 |
+
<style>
|
| 21 |
+
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&display=swap');
|
| 22 |
+
|
| 23 |
+
html, body, [class*="css"] {
|
| 24 |
+
font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
|
| 25 |
+
}
|
| 26 |
+
.stApp { background-color: #F1F5F9; }
|
| 27 |
+
|
| 28 |
+
/* Hide Streamlit chrome */
|
| 29 |
+
#MainMenu, header, footer { visibility: hidden; }
|
| 30 |
+
|
| 31 |
+
/* ── Sidebar ── */
|
| 32 |
+
section[data-testid="stSidebar"] {
|
| 33 |
+
background: #FFFFFF !important;
|
| 34 |
+
border-right: 1px solid #E2E8F0 !important;
|
| 35 |
+
}
|
| 36 |
+
section[data-testid="stSidebar"] .block-container {
|
| 37 |
+
padding: 1.75rem 1rem 1rem 1rem !important;
|
| 38 |
+
}
|
| 39 |
+
|
| 40 |
+
/* ── Main layout ── */
|
| 41 |
+
.main .block-container {
|
| 42 |
+
padding: 2.25rem 2.5rem !important;
|
| 43 |
+
max-width: 1440px;
|
| 44 |
+
}
|
| 45 |
+
|
| 46 |
+
/* ── Metric cards ── */
|
| 47 |
+
div[data-testid="metric-container"] {
|
| 48 |
+
background: #FFFFFF;
|
| 49 |
+
border: 1px solid #E2E8F0;
|
| 50 |
+
border-radius: 14px;
|
| 51 |
+
padding: 1.1rem 1.4rem !important;
|
| 52 |
+
box-shadow: 0 1px 3px rgba(0,0,0,.04);
|
| 53 |
+
}
|
| 54 |
+
div[data-testid="metric-container"] > label {
|
| 55 |
+
font-size: 0.62rem !important;
|
| 56 |
+
font-weight: 700 !important;
|
| 57 |
+
letter-spacing: 0.12em !important;
|
| 58 |
+
text-transform: uppercase !important;
|
| 59 |
+
color: #94A3B8 !important;
|
| 60 |
+
}
|
| 61 |
+
div[data-testid="stMetricValue"] > div {
|
| 62 |
+
font-size: 1.9rem !important;
|
| 63 |
+
font-weight: 800 !important;
|
| 64 |
+
color: #0F172A !important;
|
| 65 |
+
letter-spacing: -0.03em !important;
|
| 66 |
+
}
|
| 67 |
+
|
| 68 |
+
/* ── Buttons ── */
|
| 69 |
+
div[data-testid="stButton"] button[kind="primary"] {
|
| 70 |
+
background: #2563EB !important;
|
| 71 |
+
color: #FFFFFF !important;
|
| 72 |
+
border: none !important;
|
| 73 |
+
border-radius: 10px !important;
|
| 74 |
+
font-weight: 700 !important;
|
| 75 |
+
font-size: 0.95rem !important;
|
| 76 |
+
letter-spacing: 0.02em !important;
|
| 77 |
+
padding: 0.7rem 0 !important;
|
| 78 |
+
transition: background 0.15s, box-shadow 0.15s !important;
|
| 79 |
+
}
|
| 80 |
+
div[data-testid="stButton"] button[kind="primary"]:hover {
|
| 81 |
+
background: #1D4ED8 !important;
|
| 82 |
+
box-shadow: 0 6px 18px rgba(37,99,235,.35) !important;
|
| 83 |
+
}
|
| 84 |
+
div[data-testid="stButton"] button[kind="secondary"] {
|
| 85 |
+
border: 1px solid #CBD5E1 !important;
|
| 86 |
+
border-radius: 8px !important;
|
| 87 |
+
color: #475569 !important;
|
| 88 |
+
font-weight: 500 !important;
|
| 89 |
+
background: white !important;
|
| 90 |
+
}
|
| 91 |
+
|
| 92 |
+
/* ── File uploader ── */
|
| 93 |
+
div[data-testid="stFileUploader"] > section {
|
| 94 |
+
border: 2px dashed #CBD5E1 !important;
|
| 95 |
+
border-radius: 12px !important;
|
| 96 |
+
background: #F8FAFC !important;
|
| 97 |
+
transition: border-color 0.2s !important;
|
| 98 |
+
}
|
| 99 |
+
div[data-testid="stFileUploader"] > section:hover {
|
| 100 |
+
border-color: #2563EB !important;
|
| 101 |
+
}
|
| 102 |
+
|
| 103 |
+
/* ── Images ── */
|
| 104 |
+
div[data-testid="stImage"] img {
|
| 105 |
+
border-radius: 10px;
|
| 106 |
+
border: 1px solid #E2E8F0;
|
| 107 |
+
}
|
| 108 |
+
|
| 109 |
+
/* ── Sliders ── */
|
| 110 |
+
div[data-baseweb="slider"] [role="slider"] { background: #2563EB !important; }
|
| 111 |
+
div[data-baseweb="slider"] div[data-testid="stThumbValue"] { color: #2563EB !important; }
|
| 112 |
+
|
| 113 |
+
/* ── Expander ── */
|
| 114 |
+
div[data-testid="stExpander"] {
|
| 115 |
+
background: white;
|
| 116 |
+
border: 1px solid #E2E8F0;
|
| 117 |
+
border-radius: 10px;
|
| 118 |
+
}
|
| 119 |
+
|
| 120 |
+
/* ── Divider ── */
|
| 121 |
+
hr { border-color: #E2E8F0 !important; }
|
| 122 |
+
|
| 123 |
+
/* ── Custom component helpers ── */
|
| 124 |
+
.section-label {
|
| 125 |
+
font-size: 0.62rem;
|
| 126 |
+
font-weight: 700;
|
| 127 |
+
letter-spacing: 0.12em;
|
| 128 |
+
text-transform: uppercase;
|
| 129 |
+
color: #94A3B8;
|
| 130 |
+
margin-bottom: 0.6rem;
|
| 131 |
+
display: block;
|
| 132 |
+
}
|
| 133 |
+
.badge {
|
| 134 |
+
display: inline-block;
|
| 135 |
+
padding: 0.15rem 0.65rem;
|
| 136 |
+
border-radius: 999px;
|
| 137 |
+
font-size: 0.6rem;
|
| 138 |
+
font-weight: 700;
|
| 139 |
+
letter-spacing: 0.08em;
|
| 140 |
+
text-transform: uppercase;
|
| 141 |
+
}
|
| 142 |
+
.badge-blue { background:#EFF6FF; color:#2563EB; }
|
| 143 |
+
.badge-green { background:#F0FDF4; color:#15803D; }
|
| 144 |
+
.badge-amber { background:#FFFBEB; color:#B45309; }
|
| 145 |
+
.badge-red { background:#FEF2F2; color:#DC2626; }
|
| 146 |
+
|
| 147 |
+
.det-card {
|
| 148 |
+
background: #FAFBFC;
|
| 149 |
+
border: 1px solid #E2E8F0;
|
| 150 |
+
border-radius: 10px;
|
| 151 |
+
padding: 0.85rem 1rem;
|
| 152 |
+
margin-bottom: 0.55rem;
|
| 153 |
+
}
|
| 154 |
+
.det-card-title { font-weight: 600; font-size: 0.85rem; color: #0F172A; }
|
| 155 |
+
.det-card-sub { font-size: 0.72rem; color: #94A3B8; margin-top: 0.1rem; }
|
| 156 |
+
|
| 157 |
+
.nav-item {
|
| 158 |
+
display: flex; align-items: center; gap: 0.55rem;
|
| 159 |
+
padding: 0.52rem 0.75rem;
|
| 160 |
+
border-radius: 8px;
|
| 161 |
+
font-size: 0.82rem; font-weight: 500;
|
| 162 |
+
color: #64748B;
|
| 163 |
+
margin-bottom: 0.2rem;
|
| 164 |
+
cursor: default;
|
| 165 |
+
}
|
| 166 |
+
.nav-item.active {
|
| 167 |
+
background: #EFF6FF;
|
| 168 |
+
color: #2563EB;
|
| 169 |
+
font-weight: 600;
|
| 170 |
+
}
|
| 171 |
+
</style>
|
| 172 |
+
""", unsafe_allow_html=True)
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
# ── Pipeline loader ──────────────────────────────────────────────────────────
|
| 176 |
+
@st.cache_resource(show_spinner="Loading DINOv2 ViT-S/14 …")
|
| 177 |
+
def _load_pipeline():
|
| 178 |
+
from src.pipeline import PatternDetectionPipeline
|
| 179 |
+
return PatternDetectionPipeline()
|
| 180 |
+
|
| 181 |
+
try:
|
| 182 |
+
pipeline = _load_pipeline()
|
| 183 |
+
_model_ok = True
|
| 184 |
+
except Exception as _e:
|
| 185 |
+
_model_ok = False
|
| 186 |
+
_model_err = str(_e)
|
| 187 |
+
|
| 188 |
+
|
| 189 |
+
# ============================================================
|
| 190 |
+
# SIDEBAR
|
| 191 |
+
# ============================================================
|
| 192 |
+
with st.sidebar:
|
| 193 |
+
# Logo
|
| 194 |
+
st.markdown("""
|
| 195 |
+
<div style="padding: 0 0 1.25rem 0;">
|
| 196 |
+
<div style="font-weight:800;font-size:1.2rem;color:#0F172A;letter-spacing:-0.03em;">
|
| 197 |
+
⚡ PatternScan
|
| 198 |
+
</div>
|
| 199 |
+
<div style="font-size:0.7rem;color:#94A3B8;margin-top:0.2rem;">
|
| 200 |
+
BOM Engineering Drawings
|
| 201 |
+
</div>
|
| 202 |
+
</div>
|
| 203 |
+
""", unsafe_allow_html=True)
|
| 204 |
+
|
| 205 |
+
st.divider()
|
| 206 |
+
|
| 207 |
+
# Navigation
|
| 208 |
+
st.markdown("""
|
| 209 |
+
<div class="nav-item active">📐 Detection</div>
|
| 210 |
+
<div class="nav-item">📊 History</div>
|
| 211 |
+
<div class="nav-item">⚙ Settings</div>
|
| 212 |
+
""", unsafe_allow_html=True)
|
| 213 |
+
|
| 214 |
+
st.divider()
|
| 215 |
+
|
| 216 |
+
# Algorithm settings
|
| 217 |
+
st.markdown('<span class="section-label">Algorithm</span>', unsafe_allow_html=True)
|
| 218 |
+
manual_mode = st.toggle("Manual mode", value=False,
|
| 219 |
+
help="Override automatic threshold selection")
|
| 220 |
+
|
| 221 |
+
if manual_mode:
|
| 222 |
+
ncc_thr = st.slider("NCC Threshold", 0.15, 0.9, 0.28, 0.01,
|
| 223 |
+
help="Lower → higher recall, more false positives")
|
| 224 |
+
dino_thr = st.slider("DINOv2 Threshold", 0.50, 0.95, 0.84, 0.01,
|
| 225 |
+
help="Higher → fewer false positives")
|
| 226 |
+
dilate = st.slider("Stroke Dilation", 0, 9, 0, 1,
|
| 227 |
+
help="Thicken template strokes for bold drawings")
|
| 228 |
+
else:
|
| 229 |
+
ncc_thr, dino_thr, dilate = 0.28, 0.84, 0
|
| 230 |
+
|
| 231 |
+
st.divider()
|
| 232 |
+
|
| 233 |
+
# Model status
|
| 234 |
+
if _model_ok:
|
| 235 |
+
st.markdown('<span class="badge badge-green">● Model Ready</span>', unsafe_allow_html=True)
|
| 236 |
+
st.caption("DINOv2 ViT-S/14 · NCC multi-scale")
|
| 237 |
+
else:
|
| 238 |
+
st.markdown('<span class="badge badge-red">● Error</span>', unsafe_allow_html=True)
|
| 239 |
+
st.caption(_model_err[:80])
|
| 240 |
+
|
| 241 |
+
st.divider()
|
| 242 |
+
st.caption("Zero-shot · No training required")
|
| 243 |
+
st.caption("NCC + DINOv2 + Containment NMS")
|
| 244 |
+
|
| 245 |
+
|
| 246 |
+
# ============================================================
|
| 247 |
+
# MAIN CONTENT
|
| 248 |
+
# ============================================================
|
| 249 |
+
|
| 250 |
+
# Page header
|
| 251 |
+
st.markdown("""
|
| 252 |
+
<h1 style="font-size:1.65rem;font-weight:800;color:#0F172A;
|
| 253 |
+
letter-spacing:-0.03em;margin-bottom:0.2rem;">
|
| 254 |
+
Zero-Shot Pattern Detection
|
| 255 |
+
</h1>
|
| 256 |
+
<p style="color:#64748B;font-size:0.88rem;margin-bottom:1.75rem;">
|
| 257 |
+
Locate component symbols in BOM engineering drawings —
|
| 258 |
+
no training data required. Upload a pattern and drawing, then click <b>Run</b>.
|
| 259 |
+
</p>
|
| 260 |
+
""", unsafe_allow_html=True)
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
# ── Upload row ───────────────────────────────────────────────────────────────
|
| 264 |
+
up1, up2 = st.columns(2, gap="large")
|
| 265 |
+
|
| 266 |
+
with up1:
|
| 267 |
+
st.markdown('<span class="section-label">Pattern Symbol</span>', unsafe_allow_html=True)
|
| 268 |
+
pat_file = st.file_uploader(
|
| 269 |
+
"pattern", type=["png", "jpg", "jpeg", "tif", "tiff"],
|
| 270 |
+
label_visibility="collapsed", key="pat_up",
|
| 271 |
+
)
|
| 272 |
+
if pat_file:
|
| 273 |
+
pil_p = Image.open(pat_file).convert("RGB")
|
| 274 |
+
st.image(pil_p, caption=f"{pil_p.width} × {pil_p.height} px",
|
| 275 |
+
use_container_width=True)
|
| 276 |
+
else:
|
| 277 |
+
st.markdown("""
|
| 278 |
+
<div style="text-align:center;padding:3rem 1rem;color:#94A3B8;font-size:0.82rem;
|
| 279 |
+
background:#F8FAFC;border-radius:10px;border:2px dashed #CBD5E1;">
|
| 280 |
+
📂 PNG · JPG · TIFF
|
| 281 |
+
</div>""", unsafe_allow_html=True)
|
| 282 |
+
|
| 283 |
+
with up2:
|
| 284 |
+
st.markdown('<span class="section-label">Engineering Drawing</span>', unsafe_allow_html=True)
|
| 285 |
+
drw_file = st.file_uploader(
|
| 286 |
+
"drawing", type=["png", "jpg", "jpeg", "tif", "tiff"],
|
| 287 |
+
label_visibility="collapsed", key="drw_up",
|
| 288 |
+
)
|
| 289 |
+
if drw_file:
|
| 290 |
+
pil_d = Image.open(drw_file).convert("RGB")
|
| 291 |
+
st.image(pil_d, caption=f"{pil_d.width} × {pil_d.height} px",
|
| 292 |
+
use_container_width=True)
|
| 293 |
+
else:
|
| 294 |
+
st.markdown("""
|
| 295 |
+
<div style="text-align:center;padding:3rem 1rem;color:#94A3B8;font-size:0.82rem;
|
| 296 |
+
background:#F8FAFC;border-radius:10px;border:2px dashed #CBD5E1;">
|
| 297 |
+
📂 PNG · JPG · TIFF
|
| 298 |
+
</div>""", unsafe_allow_html=True)
|
| 299 |
+
|
| 300 |
+
st.markdown("<br>", unsafe_allow_html=True)
|
| 301 |
+
|
| 302 |
+
# ── Run button ───────────────────────────────────────────────────────────────
|
| 303 |
+
btn_col, _ = st.columns([1, 2])
|
| 304 |
+
with btn_col:
|
| 305 |
+
run_clicked = st.button(
|
| 306 |
+
"⚡ Run Detection",
|
| 307 |
+
type="primary",
|
| 308 |
+
use_container_width=True,
|
| 309 |
+
disabled=not _model_ok,
|
| 310 |
+
)
|
| 311 |
+
|
| 312 |
+
|
| 313 |
+
# ── Detection logic ──────────────────────────────────────────────────────────
|
| 314 |
+
if run_clicked:
|
| 315 |
+
if pat_file is None or drw_file is None:
|
| 316 |
+
st.warning("⚠ Upload both a **pattern image** and an **engineering drawing** first.")
|
| 317 |
+
else:
|
| 318 |
+
with st.spinner("Running detection pipeline …"):
|
| 319 |
+
try:
|
| 320 |
+
pat_arr = np.array(Image.open(pat_file).convert("RGB"))
|
| 321 |
+
drw_arr = np.array(Image.open(drw_file).convert("RGB"))
|
| 322 |
+
|
| 323 |
+
t0 = time.time()
|
| 324 |
+
if manual_mode:
|
| 325 |
+
pipeline.update_thresholds(
|
| 326 |
+
ncc_threshold=ncc_thr,
|
| 327 |
+
cosine_threshold=dino_thr,
|
| 328 |
+
)
|
| 329 |
+
pipeline.dilate_pattern = dilate
|
| 330 |
+
result = pipeline.detect(pat_arr, drw_arr, return_visualization=True)
|
| 331 |
+
else:
|
| 332 |
+
result = pipeline.detect_auto(pat_arr, drw_arr, return_visualization=True)
|
| 333 |
+
elapsed = time.time() - t0
|
| 334 |
+
|
| 335 |
+
vis = result.pop("visualization", None)
|
| 336 |
+
dets = result.get("detections", [])
|
| 337 |
+
n = result["total_detections"]
|
| 338 |
+
|
| 339 |
+
st.session_state.update({
|
| 340 |
+
"result": result, "vis": vis, "elapsed": elapsed,
|
| 341 |
+
"n": n, "dets": dets,
|
| 342 |
+
})
|
| 343 |
+
except Exception as ex:
|
| 344 |
+
st.error(f"Detection failed: {ex}")
|
| 345 |
+
st.session_state["result"] = None
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
# ── Results section ───────────────────────────────────────────────────────────
|
| 349 |
+
if st.session_state.get("result") is not None:
|
| 350 |
+
result = st.session_state["result"]
|
| 351 |
+
vis = st.session_state["vis"]
|
| 352 |
+
elapsed = st.session_state["elapsed"]
|
| 353 |
+
n = st.session_state["n"]
|
| 354 |
+
dets = st.session_state["dets"]
|
| 355 |
+
|
| 356 |
+
st.markdown("---")
|
| 357 |
+
|
| 358 |
+
# ── Stat cards ──
|
| 359 |
+
st.markdown('<span class="section-label">Detection Summary</span>', unsafe_allow_html=True)
|
| 360 |
+
|
| 361 |
+
s1, s2, s3, s4 = st.columns(4, gap="medium")
|
| 362 |
+
with s1:
|
| 363 |
+
status_lbl = "Detected" if n > 0 else "Not found"
|
| 364 |
+
st.metric("Instances Found", n)
|
| 365 |
+
with s2:
|
| 366 |
+
avg_c = round(sum(d["confidence"] for d in dets) / n, 2) if n else 0.0
|
| 367 |
+
st.metric("Avg Confidence", f"{avg_c:.2f}")
|
| 368 |
+
with s3:
|
| 369 |
+
best_d = round(max(d["dino_score"] for d in dets), 4) if n else 0.0
|
| 370 |
+
st.metric("Best DINO Score", f"{best_d:.4f}")
|
| 371 |
+
with s4:
|
| 372 |
+
st.metric("Processing Time", f"{elapsed:.1f} s")
|
| 373 |
+
|
| 374 |
+
st.markdown("<br>", unsafe_allow_html=True)
|
| 375 |
+
|
| 376 |
+
# ── Output image + detection list ──
|
| 377 |
+
out1, out2 = st.columns([3, 1], gap="large")
|
| 378 |
+
|
| 379 |
+
with out1:
|
| 380 |
+
st.markdown('<span class="section-label">Annotated Output</span>', unsafe_allow_html=True)
|
| 381 |
+
if vis is not None:
|
| 382 |
+
st.image(vis, use_container_width=True,
|
| 383 |
+
caption=f"{n} instance(s) detected · {elapsed:.1f} s · "
|
| 384 |
+
f"NCC + DINOv2 ViT-S/14")
|
| 385 |
+
else:
|
| 386 |
+
st.info("No visualization available.")
|
| 387 |
+
|
| 388 |
+
with out2:
|
| 389 |
+
st.markdown('<span class="section-label">Detections</span>', unsafe_allow_html=True)
|
| 390 |
+
|
| 391 |
+
if not dets:
|
| 392 |
+
st.markdown("""
|
| 393 |
+
<div style="text-align:center;padding:2.5rem 1rem;color:#94A3B8;
|
| 394 |
+
font-size:0.82rem;background:#F8FAFC;border-radius:10px;
|
| 395 |
+
border:1px dashed #CBD5E1;">
|
| 396 |
+
No patterns detected
|
| 397 |
+
</div>""", unsafe_allow_html=True)
|
| 398 |
+
else:
|
| 399 |
+
for i, det in enumerate(dets):
|
| 400 |
+
bb = det["bbox"]
|
| 401 |
+
conf = float(det["confidence"])
|
| 402 |
+
if conf >= 0.70:
|
| 403 |
+
ccolor, cbadge = "#15803D", "badge-green"
|
| 404 |
+
elif conf >= 0.55:
|
| 405 |
+
ccolor, cbadge = "#B45309", "badge-amber"
|
| 406 |
+
else:
|
| 407 |
+
ccolor, cbadge = "#DC2626", "badge-red"
|
| 408 |
+
|
| 409 |
+
st.markdown(f"""
|
| 410 |
+
<div class="det-card">
|
| 411 |
+
<div style="display:flex;justify-content:space-between;align-items:flex-start;">
|
| 412 |
+
<div>
|
| 413 |
+
<div class="det-card-title">Detection #{i + 1}</div>
|
| 414 |
+
<div class="det-card-sub">
|
| 415 |
+
({bb['x']}, {bb['y']})   {bb['w']} × {bb['h']} px
|
| 416 |
+
</div>
|
| 417 |
+
</div>
|
| 418 |
+
<div style="font-weight:800;font-size:1.1rem;color:{ccolor};">
|
| 419 |
+
{conf:.2f}
|
| 420 |
+
</div>
|
| 421 |
+
</div>
|
| 422 |
+
<div style="margin-top:0.6rem;display:flex;gap:1rem;flex-wrap:wrap;">
|
| 423 |
+
<span style="font-size:0.7rem;color:#64748B;">
|
| 424 |
+
NCC <b>{det['ncc_score']:.3f}</b>
|
| 425 |
+
</span>
|
| 426 |
+
<span style="font-size:0.7rem;color:#64748B;">
|
| 427 |
+
DINO <b>{det['dino_score']:.4f}</b>
|
| 428 |
+
</span>
|
| 429 |
+
<span style="font-size:0.7rem;color:#64748B;">
|
| 430 |
+
Scale <b>{det['scale']:.2f}×</b>
|
| 431 |
+
</span>
|
| 432 |
+
</div>
|
| 433 |
+
</div>
|
| 434 |
+
""", unsafe_allow_html=True)
|
| 435 |
+
|
| 436 |
+
# Raw JSON
|
| 437 |
+
st.markdown("<br>", unsafe_allow_html=True)
|
| 438 |
+
with st.expander("📄 Raw JSON output"):
|
| 439 |
+
st.json(result)
|
app/web/__init__.py
ADDED
|
File without changes
|
app/web/index.html
ADDED
|
@@ -0,0 +1,225 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<!DOCTYPE html>
|
| 2 |
+
<html lang="en">
|
| 3 |
+
<head>
|
| 4 |
+
<meta charset="UTF-8" />
|
| 5 |
+
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
| 6 |
+
<title>BOM Pattern Detection</title>
|
| 7 |
+
<link rel="preconnect" href="https://fonts.googleapis.com" />
|
| 8 |
+
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
|
| 9 |
+
<link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap" rel="stylesheet" />
|
| 10 |
+
<link rel="stylesheet" href="/static/css/style.css" />
|
| 11 |
+
</head>
|
| 12 |
+
<body>
|
| 13 |
+
<!-- Sidebar -->
|
| 14 |
+
<aside class="sidebar">
|
| 15 |
+
<div class="sidebar-logo">
|
| 16 |
+
<div class="logo-icon">
|
| 17 |
+
<svg width="24" height="24" viewBox="0 0 24 24" fill="none">
|
| 18 |
+
<rect x="2" y="2" width="9" height="9" rx="2" fill="#2563EB"/>
|
| 19 |
+
<rect x="13" y="2" width="9" height="9" rx="2" fill="#2563EB" opacity="0.5"/>
|
| 20 |
+
<rect x="2" y="13" width="9" height="9" rx="2" fill="#2563EB" opacity="0.5"/>
|
| 21 |
+
<rect x="13" y="13" width="9" height="9" rx="2" fill="#2563EB"/>
|
| 22 |
+
</svg>
|
| 23 |
+
</div>
|
| 24 |
+
<span class="logo-text">BOM Detector</span>
|
| 25 |
+
</div>
|
| 26 |
+
|
| 27 |
+
<nav class="sidebar-nav">
|
| 28 |
+
<a href="#" class="nav-item active" data-page="detect">
|
| 29 |
+
<svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
|
| 30 |
+
<circle cx="11" cy="11" r="8"/><line x1="21" y1="21" x2="16.65" y2="16.65"/>
|
| 31 |
+
</svg>
|
| 32 |
+
<span>Detection</span>
|
| 33 |
+
</a>
|
| 34 |
+
<a href="#" class="nav-item" data-page="about">
|
| 35 |
+
<svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
|
| 36 |
+
<circle cx="12" cy="12" r="10"/><line x1="12" y1="8" x2="12" y2="12"/><line x1="12" y1="16" x2="12.01" y2="16"/>
|
| 37 |
+
</svg>
|
| 38 |
+
<span>About</span>
|
| 39 |
+
</a>
|
| 40 |
+
</nav>
|
| 41 |
+
|
| 42 |
+
<div class="sidebar-section">
|
| 43 |
+
<div class="section-label">Algorithm</div>
|
| 44 |
+
<div class="setting-row">
|
| 45 |
+
<label class="toggle-label">
|
| 46 |
+
<span>Auto-tune mode</span>
|
| 47 |
+
<div class="toggle" id="autoToggle">
|
| 48 |
+
<input type="checkbox" id="autoMode" checked />
|
| 49 |
+
<span class="toggle-slider"></span>
|
| 50 |
+
</div>
|
| 51 |
+
</label>
|
| 52 |
+
</div>
|
| 53 |
+
</div>
|
| 54 |
+
|
| 55 |
+
<div class="sidebar-section" id="manualSettings" style="display:none;">
|
| 56 |
+
<div class="section-label">Thresholds</div>
|
| 57 |
+
|
| 58 |
+
<div class="slider-group">
|
| 59 |
+
<div class="slider-row">
|
| 60 |
+
<span class="slider-label">NCC</span>
|
| 61 |
+
<span class="slider-val" id="nccVal">0.55</span>
|
| 62 |
+
</div>
|
| 63 |
+
<input type="range" class="slider" id="nccThreshold" min="0.1" max="0.9" step="0.01" value="0.55" />
|
| 64 |
+
</div>
|
| 65 |
+
|
| 66 |
+
<div class="slider-group">
|
| 67 |
+
<div class="slider-row">
|
| 68 |
+
<span class="slider-label">DINOv2 Cosine</span>
|
| 69 |
+
<span class="slider-val" id="dinoVal">0.84</span>
|
| 70 |
+
</div>
|
| 71 |
+
<input type="range" class="slider" id="dinoThreshold" min="0.5" max="0.99" step="0.01" value="0.84" />
|
| 72 |
+
</div>
|
| 73 |
+
|
| 74 |
+
<div class="slider-group">
|
| 75 |
+
<div class="slider-row">
|
| 76 |
+
<span class="slider-label">NMS IoU</span>
|
| 77 |
+
<span class="slider-val" id="nmsVal">0.40</span>
|
| 78 |
+
</div>
|
| 79 |
+
<input type="range" class="slider" id="nmsThreshold" min="0.1" max="0.9" step="0.01" value="0.40" />
|
| 80 |
+
</div>
|
| 81 |
+
</div>
|
| 82 |
+
|
| 83 |
+
<div class="sidebar-footer">
|
| 84 |
+
<div class="tech-badges">
|
| 85 |
+
<span class="badge">NCC</span>
|
| 86 |
+
<span class="badge">DINOv2</span>
|
| 87 |
+
<span class="badge">Zero-shot</span>
|
| 88 |
+
</div>
|
| 89 |
+
</div>
|
| 90 |
+
</aside>
|
| 91 |
+
|
| 92 |
+
<!-- Main content -->
|
| 93 |
+
<main class="main">
|
| 94 |
+
<header class="topbar">
|
| 95 |
+
<div class="topbar-left">
|
| 96 |
+
<h1 class="page-title">Pattern Detection</h1>
|
| 97 |
+
<span class="page-subtitle">BOM Engineering Drawings · Zero-shot</span>
|
| 98 |
+
</div>
|
| 99 |
+
<div class="topbar-right">
|
| 100 |
+
<div class="status-dot" id="statusDot"></div>
|
| 101 |
+
<span class="status-text" id="statusText">Ready</span>
|
| 102 |
+
</div>
|
| 103 |
+
</header>
|
| 104 |
+
|
| 105 |
+
<!-- Upload section -->
|
| 106 |
+
<section class="upload-section">
|
| 107 |
+
<div class="upload-grid">
|
| 108 |
+
<!-- Pattern upload -->
|
| 109 |
+
<div class="upload-card" id="patternDropzone">
|
| 110 |
+
<div class="upload-inner">
|
| 111 |
+
<div class="upload-icon">
|
| 112 |
+
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5">
|
| 113 |
+
<rect x="3" y="3" width="18" height="18" rx="3"/>
|
| 114 |
+
<circle cx="8.5" cy="8.5" r="1.5"/>
|
| 115 |
+
<polyline points="21 15 16 10 5 21"/>
|
| 116 |
+
</svg>
|
| 117 |
+
</div>
|
| 118 |
+
<div class="upload-label">Pattern Template</div>
|
| 119 |
+
<div class="upload-hint">Drop image or click to browse</div>
|
| 120 |
+
<input type="file" id="patternFile" accept="image/*" class="file-input" />
|
| 121 |
+
</div>
|
| 122 |
+
<div class="upload-preview" id="patternPreview" style="display:none;">
|
| 123 |
+
<img id="patternImg" alt="Pattern" />
|
| 124 |
+
<button class="clear-btn" id="clearPattern">×</button>
|
| 125 |
+
</div>
|
| 126 |
+
</div>
|
| 127 |
+
|
| 128 |
+
<!-- Drawing upload -->
|
| 129 |
+
<div class="upload-card" id="drawingDropzone">
|
| 130 |
+
<div class="upload-inner">
|
| 131 |
+
<div class="upload-icon">
|
| 132 |
+
<svg width="32" height="32" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5">
|
| 133 |
+
<path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8z"/>
|
| 134 |
+
<polyline points="14 2 14 8 20 8"/>
|
| 135 |
+
<line x1="16" y1="13" x2="8" y2="13"/>
|
| 136 |
+
<line x1="16" y1="17" x2="8" y2="17"/>
|
| 137 |
+
<polyline points="10 9 9 9 8 9"/>
|
| 138 |
+
</svg>
|
| 139 |
+
</div>
|
| 140 |
+
<div class="upload-label">Engineering Drawing</div>
|
| 141 |
+
<div class="upload-hint">Drop image or click to browse</div>
|
| 142 |
+
<input type="file" id="drawingFile" accept="image/*" class="file-input" />
|
| 143 |
+
</div>
|
| 144 |
+
<div class="upload-preview" id="drawingPreview" style="display:none;">
|
| 145 |
+
<img id="drawingImg" alt="Drawing" />
|
| 146 |
+
<button class="clear-btn" id="clearDrawing">×</button>
|
| 147 |
+
</div>
|
| 148 |
+
</div>
|
| 149 |
+
</div>
|
| 150 |
+
|
| 151 |
+
<div class="run-row">
|
| 152 |
+
<button class="run-btn" id="runBtn" disabled>
|
| 153 |
+
<svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.5">
|
| 154 |
+
<polygon points="5 3 19 12 5 21 5 3"/>
|
| 155 |
+
</svg>
|
| 156 |
+
Run Detection
|
| 157 |
+
</button>
|
| 158 |
+
</div>
|
| 159 |
+
</section>
|
| 160 |
+
|
| 161 |
+
<!-- Stats row (hidden until result) -->
|
| 162 |
+
<section class="stats-row" id="statsRow" style="display:none;">
|
| 163 |
+
<div class="stat-card">
|
| 164 |
+
<div class="stat-value" id="statDetections">—</div>
|
| 165 |
+
<div class="stat-label">Detections</div>
|
| 166 |
+
</div>
|
| 167 |
+
<div class="stat-card">
|
| 168 |
+
<div class="stat-value" id="statAvgConf">—</div>
|
| 169 |
+
<div class="stat-label">Avg Confidence</div>
|
| 170 |
+
</div>
|
| 171 |
+
<div class="stat-card">
|
| 172 |
+
<div class="stat-value" id="statBestDino">—</div>
|
| 173 |
+
<div class="stat-label">Best DINOv2</div>
|
| 174 |
+
</div>
|
| 175 |
+
<div class="stat-card">
|
| 176 |
+
<div class="stat-value" id="statTime">—</div>
|
| 177 |
+
<div class="stat-label">Time (s)</div>
|
| 178 |
+
</div>
|
| 179 |
+
</section>
|
| 180 |
+
|
| 181 |
+
<!-- Results section (hidden until result) -->
|
| 182 |
+
<section class="results-section" id="resultsSection" style="display:none;">
|
| 183 |
+
<div class="results-grid">
|
| 184 |
+
<!-- Annotated output -->
|
| 185 |
+
<div class="result-card viz-card">
|
| 186 |
+
<div class="card-header">
|
| 187 |
+
<span class="card-title">Annotated Output</span>
|
| 188 |
+
<button class="download-btn" id="downloadBtn">
|
| 189 |
+
<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
|
| 190 |
+
<path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/>
|
| 191 |
+
<polyline points="7 10 12 15 17 10"/>
|
| 192 |
+
<line x1="12" y1="15" x2="12" y2="3"/>
|
| 193 |
+
</svg>
|
| 194 |
+
Save
|
| 195 |
+
</button>
|
| 196 |
+
</div>
|
| 197 |
+
<div class="viz-container">
|
| 198 |
+
<img id="vizImage" alt="Detection result" />
|
| 199 |
+
</div>
|
| 200 |
+
</div>
|
| 201 |
+
|
| 202 |
+
<!-- Detection list -->
|
| 203 |
+
<div class="result-card detections-card">
|
| 204 |
+
<div class="card-header">
|
| 205 |
+
<span class="card-title">Detections</span>
|
| 206 |
+
<span class="card-count" id="detectionCount">0</span>
|
| 207 |
+
</div>
|
| 208 |
+
<div class="detections-list" id="detectionsList"></div>
|
| 209 |
+
</div>
|
| 210 |
+
</div>
|
| 211 |
+
</section>
|
| 212 |
+
|
| 213 |
+
<!-- Loading overlay -->
|
| 214 |
+
<div class="loading-overlay" id="loadingOverlay" style="display:none;">
|
| 215 |
+
<div class="loading-card">
|
| 216 |
+
<div class="spinner"></div>
|
| 217 |
+
<div class="loading-title">Detecting Patterns</div>
|
| 218 |
+
<div class="loading-sub" id="loadingStage">Initializing pipeline…</div>
|
| 219 |
+
</div>
|
| 220 |
+
</div>
|
| 221 |
+
</main>
|
| 222 |
+
|
| 223 |
+
<script src="/static/js/app.js"></script>
|
| 224 |
+
</body>
|
| 225 |
+
</html>
|
app/web/server.py
ADDED
|
@@ -0,0 +1,145 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FastAPI backend for BOM Pattern Detection web UI."""
|
| 2 |
+
import io
|
| 3 |
+
import base64
|
| 4 |
+
import time
|
| 5 |
+
import sys
|
| 6 |
+
import os
|
| 7 |
+
|
| 8 |
+
# Fix OpenMP conflict on Windows/Anaconda
|
| 9 |
+
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
|
| 10 |
+
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
from typing import Optional
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
from PIL import Image
|
| 16 |
+
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
|
| 17 |
+
from fastapi.staticfiles import StaticFiles
|
| 18 |
+
from fastapi.responses import HTMLResponse, JSONResponse, Response
|
| 19 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 20 |
+
import uvicorn
|
| 21 |
+
|
| 22 |
+
# Ensure project root is importable
|
| 23 |
+
ROOT = Path(__file__).resolve().parents[2]
|
| 24 |
+
sys.path.insert(0, str(ROOT))
|
| 25 |
+
|
| 26 |
+
from src.pipeline import PatternDetectionPipeline
|
| 27 |
+
|
| 28 |
+
app = FastAPI(title="BOM Pattern Detection API", version="1.0.0")
|
| 29 |
+
|
| 30 |
+
app.add_middleware(
|
| 31 |
+
CORSMiddleware,
|
| 32 |
+
allow_origins=["*"],
|
| 33 |
+
allow_methods=["*"],
|
| 34 |
+
allow_headers=["*"],
|
| 35 |
+
)
|
| 36 |
+
|
| 37 |
+
# Mount static files
|
| 38 |
+
STATIC_DIR = Path(__file__).parent / "static"
|
| 39 |
+
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
|
| 40 |
+
|
| 41 |
+
# Cached pipeline instance
|
| 42 |
+
_pipeline: Optional[PatternDetectionPipeline] = None
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def get_pipeline(config: dict = None) -> PatternDetectionPipeline:
|
| 46 |
+
global _pipeline
|
| 47 |
+
if _pipeline is None:
|
| 48 |
+
print("[Server] Loading pipeline...")
|
| 49 |
+
_pipeline = PatternDetectionPipeline(config=config)
|
| 50 |
+
print("[Server] Pipeline ready.")
|
| 51 |
+
return _pipeline
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def upload_to_numpy(file_bytes: bytes) -> np.ndarray:
|
| 55 |
+
img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
|
| 56 |
+
return np.array(img)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def numpy_to_b64(img: np.ndarray) -> str:
|
| 60 |
+
if img.ndim == 2:
|
| 61 |
+
pil_img = Image.fromarray(img, mode="L")
|
| 62 |
+
else:
|
| 63 |
+
pil_img = Image.fromarray(img)
|
| 64 |
+
buf = io.BytesIO()
|
| 65 |
+
pil_img.save(buf, format="PNG")
|
| 66 |
+
return base64.b64encode(buf.getvalue()).decode()
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
@app.get("/favicon.ico", include_in_schema=False)
|
| 70 |
+
async def favicon():
|
| 71 |
+
return Response(status_code=204)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
@app.get("/", response_class=HTMLResponse)
|
| 75 |
+
async def index():
|
| 76 |
+
html_path = Path(__file__).parent / "index.html"
|
| 77 |
+
return html_path.read_text(encoding="utf-8")
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
@app.post("/api/detect")
|
| 81 |
+
async def detect(
|
| 82 |
+
pattern: UploadFile = File(...),
|
| 83 |
+
drawing: UploadFile = File(...),
|
| 84 |
+
mode: str = Form("auto"),
|
| 85 |
+
ncc_threshold: float = Form(0.55),
|
| 86 |
+
cosine_threshold: float = Form(0.84),
|
| 87 |
+
final_nms_iou: float = Form(0.4),
|
| 88 |
+
):
|
| 89 |
+
try:
|
| 90 |
+
t_start = time.time()
|
| 91 |
+
|
| 92 |
+
pattern_bytes = await pattern.read()
|
| 93 |
+
drawing_bytes = await drawing.read()
|
| 94 |
+
|
| 95 |
+
pattern_np = upload_to_numpy(pattern_bytes)
|
| 96 |
+
drawing_np = upload_to_numpy(drawing_bytes)
|
| 97 |
+
|
| 98 |
+
config = {
|
| 99 |
+
"ncc_threshold": ncc_threshold,
|
| 100 |
+
"cosine_threshold": cosine_threshold,
|
| 101 |
+
"final_nms_iou": final_nms_iou,
|
| 102 |
+
}
|
| 103 |
+
pipeline = get_pipeline()
|
| 104 |
+
pipeline.update_thresholds(
|
| 105 |
+
ncc_threshold=ncc_threshold,
|
| 106 |
+
cosine_threshold=cosine_threshold,
|
| 107 |
+
final_nms_iou=final_nms_iou,
|
| 108 |
+
)
|
| 109 |
+
|
| 110 |
+
if mode == "auto":
|
| 111 |
+
result = pipeline.detect_auto(pattern_np, drawing_np, return_visualization=True)
|
| 112 |
+
else:
|
| 113 |
+
result = pipeline.detect(pattern_np, drawing_np, return_visualization=True)
|
| 114 |
+
|
| 115 |
+
elapsed = round(time.time() - t_start, 2)
|
| 116 |
+
|
| 117 |
+
viz_b64 = None
|
| 118 |
+
if "visualization" in result and result["visualization"] is not None:
|
| 119 |
+
viz_b64 = numpy_to_b64(result["visualization"])
|
| 120 |
+
|
| 121 |
+
# Also return the original drawing as b64 for display
|
| 122 |
+
drawing_b64 = numpy_to_b64(drawing_np)
|
| 123 |
+
|
| 124 |
+
return JSONResponse({
|
| 125 |
+
"success": True,
|
| 126 |
+
"total_detections": result["total_detections"],
|
| 127 |
+
"detections": result["detections"],
|
| 128 |
+
"elapsed": elapsed,
|
| 129 |
+
"visualization": viz_b64,
|
| 130 |
+
"drawing_preview": drawing_b64,
|
| 131 |
+
})
|
| 132 |
+
|
| 133 |
+
except Exception as e:
|
| 134 |
+
raise HTTPException(status_code=500, detail=str(e))
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
@app.get("/api/health")
|
| 138 |
+
async def health():
|
| 139 |
+
return {"status": "ok", "pipeline_loaded": _pipeline is not None}
|
| 140 |
+
|
| 141 |
+
|
| 142 |
+
if __name__ == "__main__":
|
| 143 |
+
# Pre-load pipeline before serving
|
| 144 |
+
get_pipeline()
|
| 145 |
+
uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")
|
app/web/static/css/style.css
ADDED
|
@@ -0,0 +1,540 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/* ===== Reset & Base ===== */
|
| 2 |
+
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
|
| 3 |
+
|
| 4 |
+
:root {
|
| 5 |
+
--bg: #F1F5F9;
|
| 6 |
+
--surface: #FFFFFF;
|
| 7 |
+
--surface-2: #F8FAFC;
|
| 8 |
+
--border: #E2E8F0;
|
| 9 |
+
--border-hover: #CBD5E1;
|
| 10 |
+
--primary: #2563EB;
|
| 11 |
+
--primary-light: #EFF6FF;
|
| 12 |
+
--primary-dark: #1D4ED8;
|
| 13 |
+
--text-primary: #0F172A;
|
| 14 |
+
--text-secondary: #64748B;
|
| 15 |
+
--text-muted: #94A3B8;
|
| 16 |
+
--green: #16A34A;
|
| 17 |
+
--green-bg: #F0FDF4;
|
| 18 |
+
--amber: #D97706;
|
| 19 |
+
--amber-bg: #FFFBEB;
|
| 20 |
+
--red: #DC2626;
|
| 21 |
+
--red-bg: #FEF2F2;
|
| 22 |
+
--sidebar-w: 240px;
|
| 23 |
+
--radius: 12px;
|
| 24 |
+
--radius-sm: 8px;
|
| 25 |
+
--shadow: 0 1px 3px rgba(0,0,0,0.08), 0 1px 2px rgba(0,0,0,0.04);
|
| 26 |
+
--shadow-md: 0 4px 6px rgba(0,0,0,0.07), 0 2px 4px rgba(0,0,0,0.04);
|
| 27 |
+
}
|
| 28 |
+
|
| 29 |
+
html, body {
|
| 30 |
+
height: 100%;
|
| 31 |
+
font-family: 'Inter', system-ui, -apple-system, sans-serif;
|
| 32 |
+
font-size: 14px;
|
| 33 |
+
color: var(--text-primary);
|
| 34 |
+
background: var(--bg);
|
| 35 |
+
-webkit-font-smoothing: antialiased;
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
/* ===== Layout ===== */
|
| 39 |
+
body { display: flex; min-height: 100vh; }
|
| 40 |
+
|
| 41 |
+
.sidebar {
|
| 42 |
+
width: var(--sidebar-w);
|
| 43 |
+
min-height: 100vh;
|
| 44 |
+
background: var(--surface);
|
| 45 |
+
border-right: 1px solid var(--border);
|
| 46 |
+
display: flex;
|
| 47 |
+
flex-direction: column;
|
| 48 |
+
padding: 0;
|
| 49 |
+
flex-shrink: 0;
|
| 50 |
+
position: fixed;
|
| 51 |
+
top: 0; left: 0; bottom: 0;
|
| 52 |
+
z-index: 100;
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
.main {
|
| 56 |
+
flex: 1;
|
| 57 |
+
margin-left: var(--sidebar-w);
|
| 58 |
+
min-height: 100vh;
|
| 59 |
+
display: flex;
|
| 60 |
+
flex-direction: column;
|
| 61 |
+
overflow-x: hidden;
|
| 62 |
+
}
|
| 63 |
+
|
| 64 |
+
/* ===== Sidebar Logo ===== */
|
| 65 |
+
.sidebar-logo {
|
| 66 |
+
display: flex;
|
| 67 |
+
align-items: center;
|
| 68 |
+
gap: 10px;
|
| 69 |
+
padding: 20px 20px 16px;
|
| 70 |
+
border-bottom: 1px solid var(--border);
|
| 71 |
+
}
|
| 72 |
+
.logo-icon {
|
| 73 |
+
width: 36px; height: 36px;
|
| 74 |
+
background: var(--primary-light);
|
| 75 |
+
border-radius: 10px;
|
| 76 |
+
display: flex; align-items: center; justify-content: center;
|
| 77 |
+
flex-shrink: 0;
|
| 78 |
+
}
|
| 79 |
+
.logo-text {
|
| 80 |
+
font-size: 15px; font-weight: 700;
|
| 81 |
+
color: var(--text-primary);
|
| 82 |
+
letter-spacing: -0.3px;
|
| 83 |
+
}
|
| 84 |
+
|
| 85 |
+
/* ===== Sidebar Nav ===== */
|
| 86 |
+
.sidebar-nav {
|
| 87 |
+
padding: 12px 10px 8px;
|
| 88 |
+
display: flex;
|
| 89 |
+
flex-direction: column;
|
| 90 |
+
gap: 2px;
|
| 91 |
+
}
|
| 92 |
+
.nav-item {
|
| 93 |
+
display: flex; align-items: center; gap: 10px;
|
| 94 |
+
padding: 9px 12px;
|
| 95 |
+
border-radius: var(--radius-sm);
|
| 96 |
+
color: var(--text-secondary);
|
| 97 |
+
text-decoration: none;
|
| 98 |
+
font-weight: 500;
|
| 99 |
+
font-size: 13.5px;
|
| 100 |
+
transition: all 0.15s;
|
| 101 |
+
}
|
| 102 |
+
.nav-item:hover {
|
| 103 |
+
background: var(--bg);
|
| 104 |
+
color: var(--text-primary);
|
| 105 |
+
}
|
| 106 |
+
.nav-item.active {
|
| 107 |
+
background: var(--primary-light);
|
| 108 |
+
color: var(--primary);
|
| 109 |
+
}
|
| 110 |
+
|
| 111 |
+
/* ===== Sidebar Sections ===== */
|
| 112 |
+
.sidebar-section {
|
| 113 |
+
padding: 16px 16px 12px;
|
| 114 |
+
border-top: 1px solid var(--border);
|
| 115 |
+
}
|
| 116 |
+
.section-label {
|
| 117 |
+
font-size: 11px;
|
| 118 |
+
font-weight: 600;
|
| 119 |
+
text-transform: uppercase;
|
| 120 |
+
letter-spacing: 0.6px;
|
| 121 |
+
color: var(--text-muted);
|
| 122 |
+
margin-bottom: 10px;
|
| 123 |
+
}
|
| 124 |
+
|
| 125 |
+
/* Toggle */
|
| 126 |
+
.toggle-label {
|
| 127 |
+
display: flex; align-items: center; justify-content: space-between;
|
| 128 |
+
cursor: pointer;
|
| 129 |
+
font-size: 13.5px; font-weight: 500;
|
| 130 |
+
color: var(--text-secondary);
|
| 131 |
+
}
|
| 132 |
+
.toggle { position: relative; display: inline-block; width: 36px; height: 20px; }
|
| 133 |
+
.toggle input { opacity: 0; width: 0; height: 0; }
|
| 134 |
+
.toggle-slider {
|
| 135 |
+
position: absolute; inset: 0;
|
| 136 |
+
background: var(--border);
|
| 137 |
+
border-radius: 20px;
|
| 138 |
+
transition: 0.2s;
|
| 139 |
+
cursor: pointer;
|
| 140 |
+
}
|
| 141 |
+
.toggle-slider::before {
|
| 142 |
+
content: '';
|
| 143 |
+
position: absolute;
|
| 144 |
+
width: 14px; height: 14px;
|
| 145 |
+
left: 3px; top: 3px;
|
| 146 |
+
background: white;
|
| 147 |
+
border-radius: 50%;
|
| 148 |
+
transition: 0.2s;
|
| 149 |
+
box-shadow: 0 1px 3px rgba(0,0,0,0.2);
|
| 150 |
+
}
|
| 151 |
+
.toggle input:checked + .toggle-slider { background: var(--primary); }
|
| 152 |
+
.toggle input:checked + .toggle-slider::before { transform: translateX(16px); }
|
| 153 |
+
|
| 154 |
+
/* Sliders */
|
| 155 |
+
.slider-group { margin-bottom: 14px; }
|
| 156 |
+
.slider-row {
|
| 157 |
+
display: flex; justify-content: space-between; align-items: center;
|
| 158 |
+
margin-bottom: 5px;
|
| 159 |
+
}
|
| 160 |
+
.slider-label { font-size: 12.5px; color: var(--text-secondary); font-weight: 500; }
|
| 161 |
+
.slider-val { font-size: 12px; font-weight: 600; color: var(--primary); }
|
| 162 |
+
.slider {
|
| 163 |
+
width: 100%; height: 4px;
|
| 164 |
+
-webkit-appearance: none;
|
| 165 |
+
background: var(--border);
|
| 166 |
+
border-radius: 4px;
|
| 167 |
+
outline: none;
|
| 168 |
+
cursor: pointer;
|
| 169 |
+
}
|
| 170 |
+
.slider::-webkit-slider-thumb {
|
| 171 |
+
-webkit-appearance: none;
|
| 172 |
+
width: 14px; height: 14px;
|
| 173 |
+
background: var(--primary);
|
| 174 |
+
border-radius: 50%;
|
| 175 |
+
box-shadow: 0 0 0 3px var(--primary-light);
|
| 176 |
+
}
|
| 177 |
+
.slider:hover { background: var(--border-hover); }
|
| 178 |
+
|
| 179 |
+
/* Footer badges */
|
| 180 |
+
.sidebar-footer {
|
| 181 |
+
margin-top: auto;
|
| 182 |
+
padding: 14px 16px 20px;
|
| 183 |
+
border-top: 1px solid var(--border);
|
| 184 |
+
}
|
| 185 |
+
.tech-badges { display: flex; flex-wrap: wrap; gap: 6px; }
|
| 186 |
+
.badge {
|
| 187 |
+
font-size: 11px; font-weight: 600;
|
| 188 |
+
background: var(--primary-light);
|
| 189 |
+
color: var(--primary);
|
| 190 |
+
padding: 3px 8px;
|
| 191 |
+
border-radius: 20px;
|
| 192 |
+
letter-spacing: 0.3px;
|
| 193 |
+
}
|
| 194 |
+
|
| 195 |
+
/* ===== Topbar ===== */
|
| 196 |
+
.topbar {
|
| 197 |
+
display: flex; align-items: center; justify-content: space-between;
|
| 198 |
+
padding: 20px 28px 16px;
|
| 199 |
+
border-bottom: 1px solid var(--border);
|
| 200 |
+
background: var(--surface);
|
| 201 |
+
}
|
| 202 |
+
.topbar-left { display: flex; flex-direction: column; gap: 2px; }
|
| 203 |
+
.page-title { font-size: 20px; font-weight: 700; letter-spacing: -0.4px; }
|
| 204 |
+
.page-subtitle { font-size: 12.5px; color: var(--text-muted); }
|
| 205 |
+
.topbar-right { display: flex; align-items: center; gap: 8px; }
|
| 206 |
+
.status-dot {
|
| 207 |
+
width: 8px; height: 8px; border-radius: 50%;
|
| 208 |
+
background: #94A3B8;
|
| 209 |
+
transition: background 0.3s;
|
| 210 |
+
}
|
| 211 |
+
.status-dot.ready { background: var(--green); }
|
| 212 |
+
.status-dot.running { background: var(--amber); animation: pulse 1s infinite; }
|
| 213 |
+
.status-dot.error { background: var(--red); }
|
| 214 |
+
.status-text { font-size: 12.5px; font-weight: 500; color: var(--text-secondary); }
|
| 215 |
+
|
| 216 |
+
@keyframes pulse {
|
| 217 |
+
0%, 100% { opacity: 1; }
|
| 218 |
+
50% { opacity: 0.3; }
|
| 219 |
+
}
|
| 220 |
+
|
| 221 |
+
/* ===== Upload Section ===== */
|
| 222 |
+
.upload-section {
|
| 223 |
+
padding: 24px 28px 20px;
|
| 224 |
+
}
|
| 225 |
+
|
| 226 |
+
.upload-grid {
|
| 227 |
+
display: grid;
|
| 228 |
+
grid-template-columns: 1fr 1fr;
|
| 229 |
+
gap: 16px;
|
| 230 |
+
}
|
| 231 |
+
|
| 232 |
+
.upload-card {
|
| 233 |
+
background: var(--surface);
|
| 234 |
+
border: 2px dashed var(--border);
|
| 235 |
+
border-radius: var(--radius);
|
| 236 |
+
min-height: 200px;
|
| 237 |
+
position: relative;
|
| 238 |
+
cursor: pointer;
|
| 239 |
+
transition: border-color 0.2s, box-shadow 0.2s;
|
| 240 |
+
overflow: hidden;
|
| 241 |
+
}
|
| 242 |
+
.upload-card:hover, .upload-card.drag-over {
|
| 243 |
+
border-color: var(--primary);
|
| 244 |
+
box-shadow: 0 0 0 3px var(--primary-light);
|
| 245 |
+
}
|
| 246 |
+
.upload-card.has-file {
|
| 247 |
+
border-style: solid;
|
| 248 |
+
border-color: var(--border);
|
| 249 |
+
cursor: default;
|
| 250 |
+
}
|
| 251 |
+
|
| 252 |
+
.upload-inner {
|
| 253 |
+
display: flex; flex-direction: column; align-items: center; justify-content: center;
|
| 254 |
+
height: 200px; gap: 10px;
|
| 255 |
+
padding: 20px;
|
| 256 |
+
}
|
| 257 |
+
.upload-icon { color: var(--text-muted); }
|
| 258 |
+
.upload-label { font-size: 14px; font-weight: 600; color: var(--text-secondary); }
|
| 259 |
+
.upload-hint { font-size: 12.5px; color: var(--text-muted); }
|
| 260 |
+
.file-input {
|
| 261 |
+
position: absolute; inset: 0;
|
| 262 |
+
opacity: 0; cursor: pointer;
|
| 263 |
+
width: 100%; height: 100%;
|
| 264 |
+
}
|
| 265 |
+
|
| 266 |
+
.upload-preview {
|
| 267 |
+
height: 200px;
|
| 268 |
+
display: flex; align-items: center; justify-content: center;
|
| 269 |
+
position: relative;
|
| 270 |
+
}
|
| 271 |
+
.upload-preview img {
|
| 272 |
+
max-width: 100%; max-height: 180px;
|
| 273 |
+
object-fit: contain;
|
| 274 |
+
border-radius: 8px;
|
| 275 |
+
}
|
| 276 |
+
.clear-btn {
|
| 277 |
+
position: absolute; top: 10px; right: 10px;
|
| 278 |
+
width: 26px; height: 26px;
|
| 279 |
+
background: rgba(0,0,0,0.5);
|
| 280 |
+
color: white;
|
| 281 |
+
border: none; border-radius: 50%;
|
| 282 |
+
cursor: pointer; font-size: 16px; line-height: 1;
|
| 283 |
+
display: flex; align-items: center; justify-content: center;
|
| 284 |
+
transition: background 0.15s;
|
| 285 |
+
}
|
| 286 |
+
.clear-btn:hover { background: rgba(0,0,0,0.75); }
|
| 287 |
+
|
| 288 |
+
/* Run button row */
|
| 289 |
+
.run-row {
|
| 290 |
+
display: flex; justify-content: center;
|
| 291 |
+
margin-top: 18px;
|
| 292 |
+
}
|
| 293 |
+
|
| 294 |
+
.run-btn {
|
| 295 |
+
display: flex; align-items: center; gap: 8px;
|
| 296 |
+
padding: 11px 28px;
|
| 297 |
+
background: var(--primary);
|
| 298 |
+
color: white;
|
| 299 |
+
border: none;
|
| 300 |
+
border-radius: var(--radius-sm);
|
| 301 |
+
font-size: 14px; font-weight: 600;
|
| 302 |
+
cursor: pointer;
|
| 303 |
+
transition: background 0.15s, box-shadow 0.15s, transform 0.1s;
|
| 304 |
+
box-shadow: 0 2px 8px rgba(37,99,235,0.35);
|
| 305 |
+
letter-spacing: 0.2px;
|
| 306 |
+
}
|
| 307 |
+
.run-btn:hover:not(:disabled) {
|
| 308 |
+
background: var(--primary-dark);
|
| 309 |
+
box-shadow: 0 4px 12px rgba(37,99,235,0.45);
|
| 310 |
+
transform: translateY(-1px);
|
| 311 |
+
}
|
| 312 |
+
.run-btn:active:not(:disabled) { transform: translateY(0); }
|
| 313 |
+
.run-btn:disabled {
|
| 314 |
+
background: #CBD5E1;
|
| 315 |
+
box-shadow: none;
|
| 316 |
+
cursor: not-allowed;
|
| 317 |
+
}
|
| 318 |
+
|
| 319 |
+
/* ===== Stats Row ===== */
|
| 320 |
+
.stats-row {
|
| 321 |
+
display: grid;
|
| 322 |
+
grid-template-columns: repeat(4, 1fr);
|
| 323 |
+
gap: 14px;
|
| 324 |
+
padding: 0 28px 20px;
|
| 325 |
+
}
|
| 326 |
+
|
| 327 |
+
.stat-card {
|
| 328 |
+
background: var(--surface);
|
| 329 |
+
border: 1px solid var(--border);
|
| 330 |
+
border-radius: var(--radius);
|
| 331 |
+
padding: 18px 20px;
|
| 332 |
+
box-shadow: var(--shadow);
|
| 333 |
+
}
|
| 334 |
+
.stat-value {
|
| 335 |
+
font-size: 26px; font-weight: 700;
|
| 336 |
+
color: var(--primary);
|
| 337 |
+
letter-spacing: -0.5px;
|
| 338 |
+
line-height: 1;
|
| 339 |
+
margin-bottom: 5px;
|
| 340 |
+
}
|
| 341 |
+
.stat-label {
|
| 342 |
+
font-size: 12px; font-weight: 500;
|
| 343 |
+
color: var(--text-muted);
|
| 344 |
+
text-transform: uppercase;
|
| 345 |
+
letter-spacing: 0.4px;
|
| 346 |
+
}
|
| 347 |
+
|
| 348 |
+
/* ===== Results Section ===== */
|
| 349 |
+
.results-section { padding: 0 28px 32px; }
|
| 350 |
+
|
| 351 |
+
.results-grid {
|
| 352 |
+
display: grid;
|
| 353 |
+
grid-template-columns: 2fr 1fr;
|
| 354 |
+
gap: 16px;
|
| 355 |
+
}
|
| 356 |
+
|
| 357 |
+
.result-card {
|
| 358 |
+
background: var(--surface);
|
| 359 |
+
border: 1px solid var(--border);
|
| 360 |
+
border-radius: var(--radius);
|
| 361 |
+
box-shadow: var(--shadow);
|
| 362 |
+
overflow: hidden;
|
| 363 |
+
}
|
| 364 |
+
|
| 365 |
+
.card-header {
|
| 366 |
+
display: flex; align-items: center; justify-content: space-between;
|
| 367 |
+
padding: 14px 18px;
|
| 368 |
+
border-bottom: 1px solid var(--border);
|
| 369 |
+
}
|
| 370 |
+
.card-title {
|
| 371 |
+
font-size: 13.5px; font-weight: 600;
|
| 372 |
+
color: var(--text-primary);
|
| 373 |
+
}
|
| 374 |
+
.card-count {
|
| 375 |
+
font-size: 12px; font-weight: 600;
|
| 376 |
+
background: var(--primary-light);
|
| 377 |
+
color: var(--primary);
|
| 378 |
+
padding: 2px 8px;
|
| 379 |
+
border-radius: 20px;
|
| 380 |
+
}
|
| 381 |
+
|
| 382 |
+
.download-btn {
|
| 383 |
+
display: flex; align-items: center; gap: 5px;
|
| 384 |
+
padding: 5px 12px;
|
| 385 |
+
background: var(--surface-2);
|
| 386 |
+
border: 1px solid var(--border);
|
| 387 |
+
border-radius: var(--radius-sm);
|
| 388 |
+
font-size: 12px; font-weight: 500;
|
| 389 |
+
color: var(--text-secondary);
|
| 390 |
+
cursor: pointer;
|
| 391 |
+
transition: all 0.15s;
|
| 392 |
+
}
|
| 393 |
+
.download-btn:hover {
|
| 394 |
+
background: var(--primary-light);
|
| 395 |
+
border-color: var(--primary);
|
| 396 |
+
color: var(--primary);
|
| 397 |
+
}
|
| 398 |
+
|
| 399 |
+
/* Visualization */
|
| 400 |
+
.viz-container {
|
| 401 |
+
padding: 16px;
|
| 402 |
+
display: flex; justify-content: center;
|
| 403 |
+
background: var(--surface-2);
|
| 404 |
+
}
|
| 405 |
+
.viz-container img {
|
| 406 |
+
max-width: 100%;
|
| 407 |
+
border-radius: 8px;
|
| 408 |
+
box-shadow: var(--shadow-md);
|
| 409 |
+
}
|
| 410 |
+
|
| 411 |
+
/* Detection cards list */
|
| 412 |
+
.detections-list {
|
| 413 |
+
padding: 10px;
|
| 414 |
+
display: flex; flex-direction: column; gap: 8px;
|
| 415 |
+
max-height: 600px;
|
| 416 |
+
overflow-y: auto;
|
| 417 |
+
}
|
| 418 |
+
.detections-list::-webkit-scrollbar { width: 4px; }
|
| 419 |
+
.detections-list::-webkit-scrollbar-track { background: transparent; }
|
| 420 |
+
.detections-list::-webkit-scrollbar-thumb { background: var(--border); border-radius: 4px; }
|
| 421 |
+
|
| 422 |
+
.detection-item {
|
| 423 |
+
border-radius: var(--radius-sm);
|
| 424 |
+
padding: 12px 14px;
|
| 425 |
+
border: 1px solid var(--border);
|
| 426 |
+
transition: box-shadow 0.15s;
|
| 427 |
+
cursor: default;
|
| 428 |
+
}
|
| 429 |
+
.detection-item:hover { box-shadow: var(--shadow-md); }
|
| 430 |
+
|
| 431 |
+
.detection-item.high { background: var(--green-bg); border-color: #BBF7D0; }
|
| 432 |
+
.detection-item.mid { background: #EFF6FF; border-color: #BFDBFE; }
|
| 433 |
+
.detection-item.low { background: var(--red-bg); border-color: #FECACA; }
|
| 434 |
+
|
| 435 |
+
.detection-header {
|
| 436 |
+
display: flex; align-items: center; justify-content: space-between;
|
| 437 |
+
margin-bottom: 8px;
|
| 438 |
+
}
|
| 439 |
+
.detection-num {
|
| 440 |
+
font-size: 12px; font-weight: 700;
|
| 441 |
+
display: flex; align-items: center; gap: 6px;
|
| 442 |
+
}
|
| 443 |
+
.detection-num .num-badge {
|
| 444 |
+
width: 22px; height: 22px;
|
| 445 |
+
border-radius: 50%;
|
| 446 |
+
display: flex; align-items: center; justify-content: center;
|
| 447 |
+
font-size: 11px; font-weight: 700; color: white;
|
| 448 |
+
}
|
| 449 |
+
.detection-item.high .num-badge { background: var(--green); }
|
| 450 |
+
.detection-item.mid .num-badge { background: var(--primary); }
|
| 451 |
+
.detection-item.low .num-badge { background: var(--red); }
|
| 452 |
+
|
| 453 |
+
.detection-conf {
|
| 454 |
+
font-size: 13px; font-weight: 700;
|
| 455 |
+
}
|
| 456 |
+
.detection-item.high .detection-conf { color: var(--green); }
|
| 457 |
+
.detection-item.mid .detection-conf { color: var(--primary); }
|
| 458 |
+
.detection-item.low .detection-conf { color: var(--red); }
|
| 459 |
+
|
| 460 |
+
.detection-meta {
|
| 461 |
+
display: grid; grid-template-columns: 1fr 1fr;
|
| 462 |
+
gap: 4px;
|
| 463 |
+
}
|
| 464 |
+
.meta-row {
|
| 465 |
+
font-size: 11.5px; color: var(--text-secondary);
|
| 466 |
+
display: flex; align-items: center; gap: 4px;
|
| 467 |
+
}
|
| 468 |
+
.meta-row span { color: var(--text-muted); }
|
| 469 |
+
|
| 470 |
+
.conf-bar-wrap {
|
| 471 |
+
margin-top: 8px;
|
| 472 |
+
height: 3px;
|
| 473 |
+
background: var(--border);
|
| 474 |
+
border-radius: 3px;
|
| 475 |
+
overflow: hidden;
|
| 476 |
+
}
|
| 477 |
+
.conf-bar {
|
| 478 |
+
height: 100%; border-radius: 3px;
|
| 479 |
+
transition: width 0.4s ease;
|
| 480 |
+
}
|
| 481 |
+
.detection-item.high .conf-bar { background: var(--green); }
|
| 482 |
+
.detection-item.mid .conf-bar { background: var(--primary); }
|
| 483 |
+
.detection-item.low .conf-bar { background: var(--red); }
|
| 484 |
+
|
| 485 |
+
/* Empty state */
|
| 486 |
+
.empty-state {
|
| 487 |
+
display: flex; flex-direction: column; align-items: center;
|
| 488 |
+
justify-content: center; gap: 8px;
|
| 489 |
+
padding: 40px 20px;
|
| 490 |
+
color: var(--text-muted);
|
| 491 |
+
text-align: center;
|
| 492 |
+
}
|
| 493 |
+
.empty-state svg { opacity: 0.4; }
|
| 494 |
+
.empty-state p { font-size: 13px; }
|
| 495 |
+
|
| 496 |
+
/* ===== Loading Overlay ===== */
|
| 497 |
+
.loading-overlay {
|
| 498 |
+
position: fixed; inset: 0;
|
| 499 |
+
background: rgba(15,23,42,0.55);
|
| 500 |
+
display: flex; align-items: center; justify-content: center;
|
| 501 |
+
z-index: 200;
|
| 502 |
+
backdrop-filter: blur(3px);
|
| 503 |
+
}
|
| 504 |
+
.loading-card {
|
| 505 |
+
background: var(--surface);
|
| 506 |
+
border-radius: var(--radius);
|
| 507 |
+
padding: 36px 48px;
|
| 508 |
+
text-align: center;
|
| 509 |
+
box-shadow: 0 20px 40px rgba(0,0,0,0.2);
|
| 510 |
+
display: flex; flex-direction: column; align-items: center; gap: 12px;
|
| 511 |
+
}
|
| 512 |
+
.spinner {
|
| 513 |
+
width: 40px; height: 40px;
|
| 514 |
+
border: 3px solid var(--border);
|
| 515 |
+
border-top-color: var(--primary);
|
| 516 |
+
border-radius: 50%;
|
| 517 |
+
animation: spin 0.75s linear infinite;
|
| 518 |
+
}
|
| 519 |
+
@keyframes spin { to { transform: rotate(360deg); } }
|
| 520 |
+
.loading-title { font-size: 16px; font-weight: 700; }
|
| 521 |
+
.loading-sub { font-size: 13px; color: var(--text-muted); }
|
| 522 |
+
|
| 523 |
+
/* ===== Error toast ===== */
|
| 524 |
+
.toast {
|
| 525 |
+
position: fixed; bottom: 24px; right: 24px;
|
| 526 |
+
background: #1E293B;
|
| 527 |
+
color: white;
|
| 528 |
+
padding: 14px 20px;
|
| 529 |
+
border-radius: var(--radius-sm);
|
| 530 |
+
font-size: 13.5px;
|
| 531 |
+
box-shadow: 0 8px 20px rgba(0,0,0,0.25);
|
| 532 |
+
z-index: 300;
|
| 533 |
+
opacity: 0;
|
| 534 |
+
transform: translateY(8px);
|
| 535 |
+
transition: opacity 0.25s, transform 0.25s;
|
| 536 |
+
max-width: 360px;
|
| 537 |
+
}
|
| 538 |
+
.toast.show { opacity: 1; transform: translateY(0); }
|
| 539 |
+
.toast.error { border-left: 3px solid var(--red); }
|
| 540 |
+
.toast.success { border-left: 3px solid var(--green); }
|
app/web/static/js/app.js
ADDED
|
@@ -0,0 +1,316 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
/* BOM Pattern Detection — Frontend Logic */
|
| 2 |
+
|
| 3 |
+
const API = ''; // same origin
|
| 4 |
+
|
| 5 |
+
// ---- State ----
|
| 6 |
+
let patternFile = null;
|
| 7 |
+
let drawingFile = null;
|
| 8 |
+
let vizB64 = null;
|
| 9 |
+
|
| 10 |
+
// ---- DOM refs ----
|
| 11 |
+
const patternDropzone = document.getElementById('patternDropzone');
|
| 12 |
+
const drawingDropzone = document.getElementById('drawingDropzone');
|
| 13 |
+
const patternFileInput = document.getElementById('patternFile');
|
| 14 |
+
const drawingFileInput = document.getElementById('drawingFile');
|
| 15 |
+
const patternPreview = document.getElementById('patternPreview');
|
| 16 |
+
const drawingPreview = document.getElementById('drawingPreview');
|
| 17 |
+
const patternImg = document.getElementById('patternImg');
|
| 18 |
+
const drawingImg = document.getElementById('drawingImg');
|
| 19 |
+
const clearPatternBtn = document.getElementById('clearPattern');
|
| 20 |
+
const clearDrawingBtn = document.getElementById('clearDrawing');
|
| 21 |
+
const runBtn = document.getElementById('runBtn');
|
| 22 |
+
|
| 23 |
+
const autoModeToggle = document.getElementById('autoMode');
|
| 24 |
+
const manualSettings = document.getElementById('manualSettings');
|
| 25 |
+
const nccSlider = document.getElementById('nccThreshold');
|
| 26 |
+
const dinoSlider = document.getElementById('dinoThreshold');
|
| 27 |
+
const nmsSlider = document.getElementById('nmsThreshold');
|
| 28 |
+
const nccVal = document.getElementById('nccVal');
|
| 29 |
+
const dinoVal = document.getElementById('dinoVal');
|
| 30 |
+
const nmsVal = document.getElementById('nmsVal');
|
| 31 |
+
|
| 32 |
+
const statusDot = document.getElementById('statusDot');
|
| 33 |
+
const statusText = document.getElementById('statusText');
|
| 34 |
+
const loadingOverlay = document.getElementById('loadingOverlay');
|
| 35 |
+
const loadingStage = document.getElementById('loadingStage');
|
| 36 |
+
|
| 37 |
+
const statsRow = document.getElementById('statsRow');
|
| 38 |
+
const statDetections = document.getElementById('statDetections');
|
| 39 |
+
const statAvgConf = document.getElementById('statAvgConf');
|
| 40 |
+
const statBestDino = document.getElementById('statBestDino');
|
| 41 |
+
const statTime = document.getElementById('statTime');
|
| 42 |
+
|
| 43 |
+
const resultsSection = document.getElementById('resultsSection');
|
| 44 |
+
const vizImage = document.getElementById('vizImage');
|
| 45 |
+
const detectionsList = document.getElementById('detectionsList');
|
| 46 |
+
const detectionCount = document.getElementById('detectionCount');
|
| 47 |
+
const downloadBtn = document.getElementById('downloadBtn');
|
| 48 |
+
|
| 49 |
+
// ---- Helpers ----
|
| 50 |
+
function setStatus(state, text) {
|
| 51 |
+
statusDot.className = 'status-dot ' + state;
|
| 52 |
+
statusText.textContent = text;
|
| 53 |
+
}
|
| 54 |
+
|
| 55 |
+
function showToast(msg, type = 'error') {
|
| 56 |
+
let toast = document.querySelector('.toast');
|
| 57 |
+
if (!toast) {
|
| 58 |
+
toast = document.createElement('div');
|
| 59 |
+
toast.className = 'toast';
|
| 60 |
+
document.body.appendChild(toast);
|
| 61 |
+
}
|
| 62 |
+
toast.className = `toast ${type}`;
|
| 63 |
+
toast.textContent = msg;
|
| 64 |
+
toast.classList.add('show');
|
| 65 |
+
setTimeout(() => toast.classList.remove('show'), 4000);
|
| 66 |
+
}
|
| 67 |
+
|
| 68 |
+
function fileToDataURL(file) {
|
| 69 |
+
return new Promise((resolve) => {
|
| 70 |
+
const reader = new FileReader();
|
| 71 |
+
reader.onload = e => resolve(e.target.result);
|
| 72 |
+
reader.readAsDataURL(file);
|
| 73 |
+
});
|
| 74 |
+
}
|
| 75 |
+
|
| 76 |
+
function checkRunReady() {
|
| 77 |
+
runBtn.disabled = !(patternFile && drawingFile);
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
function setPreview(file, imgEl, previewEl, dropzone) {
|
| 81 |
+
const url = URL.createObjectURL(file);
|
| 82 |
+
imgEl.src = url;
|
| 83 |
+
dropzone.querySelector('.upload-inner').style.display = 'none';
|
| 84 |
+
previewEl.style.display = 'flex';
|
| 85 |
+
dropzone.classList.add('has-file');
|
| 86 |
+
// Remove file input pointer-events so user can't re-trigger accidentally
|
| 87 |
+
dropzone.querySelector('.file-input').style.pointerEvents = 'none';
|
| 88 |
+
}
|
| 89 |
+
|
| 90 |
+
function clearPreview(imgEl, previewEl, dropzone, inputEl) {
|
| 91 |
+
imgEl.src = '';
|
| 92 |
+
previewEl.style.display = 'none';
|
| 93 |
+
dropzone.querySelector('.upload-inner').style.display = 'flex';
|
| 94 |
+
dropzone.classList.remove('has-file');
|
| 95 |
+
dropzone.querySelector('.file-input').style.pointerEvents = 'auto';
|
| 96 |
+
inputEl.value = '';
|
| 97 |
+
}
|
| 98 |
+
|
| 99 |
+
// ---- File Handlers ----
|
| 100 |
+
patternFileInput.addEventListener('change', (e) => {
|
| 101 |
+
const f = e.target.files[0];
|
| 102 |
+
if (!f) return;
|
| 103 |
+
patternFile = f;
|
| 104 |
+
setPreview(f, patternImg, patternPreview, patternDropzone);
|
| 105 |
+
checkRunReady();
|
| 106 |
+
});
|
| 107 |
+
|
| 108 |
+
drawingFileInput.addEventListener('change', (e) => {
|
| 109 |
+
const f = e.target.files[0];
|
| 110 |
+
if (!f) return;
|
| 111 |
+
drawingFile = f;
|
| 112 |
+
setPreview(f, drawingImg, drawingPreview, drawingDropzone);
|
| 113 |
+
checkRunReady();
|
| 114 |
+
});
|
| 115 |
+
|
| 116 |
+
clearPatternBtn.addEventListener('click', (e) => {
|
| 117 |
+
e.stopPropagation();
|
| 118 |
+
patternFile = null;
|
| 119 |
+
clearPreview(patternImg, patternPreview, patternDropzone, patternFileInput);
|
| 120 |
+
checkRunReady();
|
| 121 |
+
});
|
| 122 |
+
clearDrawingBtn.addEventListener('click', (e) => {
|
| 123 |
+
e.stopPropagation();
|
| 124 |
+
drawingFile = null;
|
| 125 |
+
clearPreview(drawingImg, drawingPreview, drawingDropzone, drawingFileInput);
|
| 126 |
+
checkRunReady();
|
| 127 |
+
});
|
| 128 |
+
|
| 129 |
+
// ---- Drag-and-drop ----
|
| 130 |
+
function setupDrop(zone, onFile) {
|
| 131 |
+
zone.addEventListener('dragover', e => { e.preventDefault(); zone.classList.add('drag-over'); });
|
| 132 |
+
zone.addEventListener('dragleave', () => zone.classList.remove('drag-over'));
|
| 133 |
+
zone.addEventListener('drop', (e) => {
|
| 134 |
+
e.preventDefault();
|
| 135 |
+
zone.classList.remove('drag-over');
|
| 136 |
+
const f = e.dataTransfer.files[0];
|
| 137 |
+
if (f && f.type.startsWith('image/')) onFile(f);
|
| 138 |
+
});
|
| 139 |
+
}
|
| 140 |
+
|
| 141 |
+
setupDrop(patternDropzone, (f) => {
|
| 142 |
+
patternFile = f;
|
| 143 |
+
setPreview(f, patternImg, patternPreview, patternDropzone);
|
| 144 |
+
checkRunReady();
|
| 145 |
+
});
|
| 146 |
+
setupDrop(drawingDropzone, (f) => {
|
| 147 |
+
drawingFile = f;
|
| 148 |
+
setPreview(f, drawingImg, drawingPreview, drawingDropzone);
|
| 149 |
+
checkRunReady();
|
| 150 |
+
});
|
| 151 |
+
|
| 152 |
+
// ---- Auto-mode toggle ----
|
| 153 |
+
autoModeToggle.addEventListener('change', () => {
|
| 154 |
+
manualSettings.style.display = autoModeToggle.checked ? 'none' : 'block';
|
| 155 |
+
});
|
| 156 |
+
|
| 157 |
+
// ---- Slider display ----
|
| 158 |
+
nccSlider.addEventListener('input', () => { nccVal.textContent = parseFloat(nccSlider.value).toFixed(2); });
|
| 159 |
+
dinoSlider.addEventListener('input', () => { dinoVal.textContent = parseFloat(dinoSlider.value).toFixed(2); });
|
| 160 |
+
nmsSlider.addEventListener('input', () => { nmsVal.textContent = parseFloat(nmsSlider.value).toFixed(2); });
|
| 161 |
+
|
| 162 |
+
// ---- Run Detection ----
|
| 163 |
+
runBtn.addEventListener('click', async () => {
|
| 164 |
+
if (!patternFile || !drawingFile) return;
|
| 165 |
+
|
| 166 |
+
setStatus('running', 'Detecting…');
|
| 167 |
+
loadingOverlay.style.display = 'flex';
|
| 168 |
+
loadingStage.textContent = 'Preprocessing images…';
|
| 169 |
+
|
| 170 |
+
const stageMessages = [
|
| 171 |
+
'Running NCC template matching…',
|
| 172 |
+
'Verifying with DINOv2…',
|
| 173 |
+
'Applying NMS post-processing…',
|
| 174 |
+
'Finalizing results…',
|
| 175 |
+
];
|
| 176 |
+
let msgIdx = 0;
|
| 177 |
+
const stageTimer = setInterval(() => {
|
| 178 |
+
if (msgIdx < stageMessages.length) {
|
| 179 |
+
loadingStage.textContent = stageMessages[msgIdx++];
|
| 180 |
+
}
|
| 181 |
+
}, 2000);
|
| 182 |
+
|
| 183 |
+
try {
|
| 184 |
+
const formData = new FormData();
|
| 185 |
+
formData.append('pattern', patternFile, patternFile.name);
|
| 186 |
+
formData.append('drawing', drawingFile, drawingFile.name);
|
| 187 |
+
formData.append('mode', autoModeToggle.checked ? 'auto' : 'manual');
|
| 188 |
+
formData.append('ncc_threshold', nccSlider.value);
|
| 189 |
+
formData.append('cosine_threshold', dinoSlider.value);
|
| 190 |
+
formData.append('final_nms_iou', nmsSlider.value);
|
| 191 |
+
|
| 192 |
+
const resp = await fetch(`${API}/api/detect`, { method: 'POST', body: formData });
|
| 193 |
+
clearInterval(stageTimer);
|
| 194 |
+
loadingOverlay.style.display = 'none';
|
| 195 |
+
|
| 196 |
+
if (!resp.ok) {
|
| 197 |
+
const err = await resp.json();
|
| 198 |
+
throw new Error(err.detail || `Server error ${resp.status}`);
|
| 199 |
+
}
|
| 200 |
+
|
| 201 |
+
const data = await resp.json();
|
| 202 |
+
renderResults(data);
|
| 203 |
+
setStatus('ready', 'Done');
|
| 204 |
+
showToast(`Found ${data.total_detections} detection(s) in ${data.elapsed}s`, 'success');
|
| 205 |
+
} catch (err) {
|
| 206 |
+
clearInterval(stageTimer);
|
| 207 |
+
loadingOverlay.style.display = 'none';
|
| 208 |
+
setStatus('error', 'Error');
|
| 209 |
+
showToast(err.message || 'Detection failed', 'error');
|
| 210 |
+
console.error(err);
|
| 211 |
+
}
|
| 212 |
+
});
|
| 213 |
+
|
| 214 |
+
// ---- Render Results ----
|
| 215 |
+
function confClass(conf) {
|
| 216 |
+
if (conf >= 0.70) return 'high';
|
| 217 |
+
if (conf >= 0.55) return 'mid';
|
| 218 |
+
return 'low';
|
| 219 |
+
}
|
| 220 |
+
|
| 221 |
+
function renderResults(data) {
|
| 222 |
+
const dets = data.detections || [];
|
| 223 |
+
const n = data.total_detections;
|
| 224 |
+
|
| 225 |
+
// Stats
|
| 226 |
+
statsRow.style.display = 'grid';
|
| 227 |
+
statDetections.textContent = n;
|
| 228 |
+
|
| 229 |
+
if (n > 0) {
|
| 230 |
+
const avgConf = dets.reduce((s, d) => s + (d.confidence || 0), 0) / n;
|
| 231 |
+
const bestDino = Math.max(...dets.map(d => d.dino_score || d.cosine_score || 0));
|
| 232 |
+
statAvgConf.textContent = avgConf.toFixed(3);
|
| 233 |
+
statBestDino.textContent = bestDino.toFixed(3);
|
| 234 |
+
} else {
|
| 235 |
+
statAvgConf.textContent = '—';
|
| 236 |
+
statBestDino.textContent = '—';
|
| 237 |
+
}
|
| 238 |
+
statTime.textContent = data.elapsed + 's';
|
| 239 |
+
|
| 240 |
+
// Visualization
|
| 241 |
+
if (data.visualization) {
|
| 242 |
+
vizB64 = data.visualization;
|
| 243 |
+
vizImage.src = 'data:image/png;base64,' + data.visualization;
|
| 244 |
+
}
|
| 245 |
+
|
| 246 |
+
// Detection list
|
| 247 |
+
detectionCount.textContent = n;
|
| 248 |
+
detectionsList.innerHTML = '';
|
| 249 |
+
|
| 250 |
+
if (n === 0) {
|
| 251 |
+
detectionsList.innerHTML = `
|
| 252 |
+
<div class="empty-state">
|
| 253 |
+
<svg width="40" height="40" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5">
|
| 254 |
+
<circle cx="12" cy="12" r="10"/>
|
| 255 |
+
<line x1="8" y1="12" x2="16" y2="12"/>
|
| 256 |
+
</svg>
|
| 257 |
+
<p>No patterns detected</p>
|
| 258 |
+
</div>`;
|
| 259 |
+
} else {
|
| 260 |
+
dets.forEach((det, i) => {
|
| 261 |
+
const conf = det.confidence || 0;
|
| 262 |
+
const cls = confClass(conf);
|
| 263 |
+
const ncc = det.ncc_score != null ? det.ncc_score.toFixed(3) : '—';
|
| 264 |
+
const dino = (det.dino_score || det.cosine_score) != null
|
| 265 |
+
? (det.dino_score || det.cosine_score).toFixed(3) : '—';
|
| 266 |
+
const scale = det.scale != null ? det.scale.toFixed(2) : '—';
|
| 267 |
+
const bbox = det.bbox || {};
|
| 268 |
+
const bboxStr = (bbox.x != null)
|
| 269 |
+
? `[${bbox.x}, ${bbox.y}, ${bbox.w}, ${bbox.h}]` : '—';
|
| 270 |
+
|
| 271 |
+
const item = document.createElement('div');
|
| 272 |
+
item.className = `detection-item ${cls}`;
|
| 273 |
+
item.innerHTML = `
|
| 274 |
+
<div class="detection-header">
|
| 275 |
+
<div class="detection-num">
|
| 276 |
+
<span class="num-badge">${i + 1}</span>
|
| 277 |
+
Detection #${i + 1}
|
| 278 |
+
</div>
|
| 279 |
+
<div class="detection-conf">${(conf * 100).toFixed(1)}%</div>
|
| 280 |
+
</div>
|
| 281 |
+
<div class="detection-meta">
|
| 282 |
+
<div class="meta-row"><span>NCC</span> ${ncc}</div>
|
| 283 |
+
<div class="meta-row"><span>DINOv2</span> ${dino}</div>
|
| 284 |
+
<div class="meta-row"><span>Scale</span> ${scale}</div>
|
| 285 |
+
<div class="meta-row"><span>BBox</span> ${bboxStr}</div>
|
| 286 |
+
</div>
|
| 287 |
+
<div class="conf-bar-wrap">
|
| 288 |
+
<div class="conf-bar" style="width:${Math.round(conf * 100)}%"></div>
|
| 289 |
+
</div>
|
| 290 |
+
`;
|
| 291 |
+
detectionsList.appendChild(item);
|
| 292 |
+
});
|
| 293 |
+
}
|
| 294 |
+
|
| 295 |
+
resultsSection.style.display = 'block';
|
| 296 |
+
resultsSection.scrollIntoView({ behavior: 'smooth', block: 'start' });
|
| 297 |
+
}
|
| 298 |
+
|
| 299 |
+
// ---- Download ----
|
| 300 |
+
downloadBtn.addEventListener('click', () => {
|
| 301 |
+
if (!vizB64) return;
|
| 302 |
+
const a = document.createElement('a');
|
| 303 |
+
a.href = 'data:image/png;base64,' + vizB64;
|
| 304 |
+
a.download = 'detection_result.png';
|
| 305 |
+
a.click();
|
| 306 |
+
});
|
| 307 |
+
|
| 308 |
+
// ---- Initial status check ----
|
| 309 |
+
(async () => {
|
| 310 |
+
try {
|
| 311 |
+
const r = await fetch('/api/health');
|
| 312 |
+
if (r.ok) setStatus('ready', 'Ready');
|
| 313 |
+
} catch {
|
| 314 |
+
setStatus('', 'Offline');
|
| 315 |
+
}
|
| 316 |
+
})();
|
conftest.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import sys
|
| 2 |
+
import os
|
| 3 |
+
|
| 4 |
+
# Ensure project root is on sys.path so 'src' package is importable
|
| 5 |
+
sys.path.insert(0, os.path.dirname(__file__))
|
design_spec/system_design.md
ADDED
|
@@ -0,0 +1,316 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Bản Đặc Tả Thiết Kế Hệ Thống
|
| 2 |
+
## Zero-Shot Pattern Detection cho Bản Vẽ Kỹ Thuật BOM
|
| 3 |
+
|
| 4 |
+
**Phiên bản:** 1.0
|
| 5 |
+
**Ngày:** 2026-05-24
|
| 6 |
+
**Tác giả:** Ứng viên AI/Computer Vision Engineer
|
| 7 |
+
|
| 8 |
+
---
|
| 9 |
+
|
| 10 |
+
## 1. Phân Tích Bài Toán
|
| 11 |
+
|
| 12 |
+
### 1.1 Đặc điểm của bản vẽ BOM
|
| 13 |
+
|
| 14 |
+
Bản vẽ kỹ thuật BOM (Bill of Materials) có các đặc điểm sau làm phức tạp bài toán nhận dạng:
|
| 15 |
+
|
| 16 |
+
- **Grayscale / Binary:** Bản vẽ chỉ có hai màu (đen/trắng), không có thông tin màu sắc để phân biệt đối tượng.
|
| 17 |
+
- **Nét mảnh, độ phân giải cao:** Các chi tiết kỹ thuật thường được vẽ bằng nét rất mảnh (1–2px ở bản scan 300DPI), đòi hỏi preprocessing cẩn thận để không xóa mất thông tin.
|
| 18 |
+
- **Nhiễu từ quá trình scan:** JPEG artifacts, đốm nhiễu, nét bị gãy, độ tương phản không đều theo vùng.
|
| 19 |
+
- **Pattern xuất hiện nhiều lần:** Cùng một ký hiệu linh kiện (e.g., điện trở, tụ điện) có thể xuất hiện 10–100 lần trong một bản vẽ với kích thước và góc xoay có biến thiên nhỏ.
|
| 20 |
+
- **Tỷ lệ kích thước khác nhau:** Pattern có thể được vẽ ở scale 90%–110% so với template tham chiếu.
|
| 21 |
+
|
| 22 |
+
### 1.2 Thách thức chính
|
| 23 |
+
|
| 24 |
+
| Thách thức | Mô tả | Tác động |
|
| 25 |
+
|------------|-------|----------|
|
| 26 |
+
| **Zero-shot** | Không có training data cho pattern cụ thể | Loại bỏ mọi phương pháp supervised |
|
| 27 |
+
| **Nhiều occurrences** | Cần tìm TẤT CẢ vị trí, không chỉ 1 | Cần NMS để loại duplicate |
|
| 28 |
+
| **Scale variation** | Pattern có thể to/nhỏ hơn ±15% | Cần multi-scale sweep |
|
| 29 |
+
| **Rotation nhỏ** | Bản vẽ scan thường bị lệch ±10° | Cần rotation sweep |
|
| 30 |
+
| **Domain shift** | Model train trên ảnh tự nhiên, không phải line art | DINOv2 cần robust với domain shift |
|
| 31 |
+
|
| 32 |
+
### 1.3 Tại sao không thể dùng supervised detection thông thường?
|
| 33 |
+
|
| 34 |
+
Phương pháp như YOLO, Faster-RCNN yêu cầu:
|
| 35 |
+
1. **Labeled training data** cho mỗi loại pattern — không khả thi vì mỗi dự án BOM có ký hiệu riêng.
|
| 36 |
+
2. **Fine-tuning mỗi khi có pattern mới** — không đáp ứng yêu cầu "zero-shot" của bài toán.
|
| 37 |
+
3. **Kích thước object rất nhỏ** (pattern 30–100px trong bản vẽ 3000×4000px) — anchor-based detection kém hiệu quả ở scale này.
|
| 38 |
+
|
| 39 |
+
**Kết luận:** Bài toán yêu cầu một hệ thống generalizable hoạt động với bất kỳ pattern mới nào tại inference time mà không cần retrain.
|
| 40 |
+
|
| 41 |
+
---
|
| 42 |
+
|
| 43 |
+
## 2. Lý Do Chọn Hướng Tiếp Cận Hybrid NCC + DINOv2
|
| 44 |
+
|
| 45 |
+
### 2.1 So sánh các hướng tiếp cận
|
| 46 |
+
|
| 47 |
+
| Hướng tiếp cận | Ưu điểm | Nhược điểm | Phù hợp? |
|
| 48 |
+
|----------------|---------|------------|----------|
|
| 49 |
+
| **Classical NCC** | Nhanh, không cần model, chính xác khi template giống hệt | Rất nhạy với noise, scale, rotation; nhiều false positive | Chỉ làm Stage 1 |
|
| 50 |
+
| **Feature Matching (SIFT/ORB)** | Invariant với rotation/scale, không cần training | Kém trên line art (ít keypoints), không generalizable | Không phù hợp chính |
|
| 51 |
+
| **Siamese CNN** | Có thể fine-tune, pair-wise matching | Cần training data của domain, không truly zero-shot | Không phù hợp |
|
| 52 |
+
| **Grounding DINO / OWLv2** | Zero-shot với text prompt | Cần text description của pattern, không phù hợp với arbitrary symbol | Không phù hợp |
|
| 53 |
+
| **DINOv2 features (self-supervised)** | Zero-shot thực sự, robust với domain shift, capture geometric structure | Chậm hơn NCC, cần GPU cho production | **Phù hợp (Stage 2)** |
|
| 54 |
+
| **Hybrid NCC + DINOv2** | Kết hợp speed của NCC + accuracy của DINOv2 | Phức tạp hơn single-stage | **✅ Lựa chọn của chúng tôi** |
|
| 55 |
+
|
| 56 |
+
### 2.2 Lý do chọn DINOv2
|
| 57 |
+
|
| 58 |
+
DINOv2 (Oquab et al., 2023) được huấn luyện với self-supervised learning trên 142M ảnh đa dạng. Các đặc điểm phù hợp với bài toán:
|
| 59 |
+
|
| 60 |
+
1. **Patch-level geometric features:** ViT chia ảnh thành patches 14×14, mỗi patch encode đặc trưng hình học cục bộ. Điều này capture topology của line drawing tốt hơn CNN thông thường.
|
| 61 |
+
2. **Không phụ thuộc màu sắc:** Bản vẽ được convert sang RGB 3-channel bằng cách stack grayscale, DINOv2 vẫn extract được features hữu ích.
|
| 62 |
+
3. **Zero-shot thực sự:** Không cần fine-tune — encode template một lần, so sánh cosine similarity với candidate crops. Hoạt động với bất kỳ pattern mới nào.
|
| 63 |
+
4. **Domain transfer:** Nghiên cứu từ DINOv2 paper cho thấy features transfer tốt sang nhiều domain khác nhau, bao gồm medical imaging và technical diagrams.
|
| 64 |
+
|
| 65 |
+
### 2.3 Lý do dùng NCC làm Stage 1
|
| 66 |
+
|
| 67 |
+
Brute-force DINOv2 sliding window trên toàn bộ ảnh 3000×4000px với stride 14px sẽ cần encode ~57,000 patches — quá chậm cho real-time use. NCC giải quyết vấn đề này:
|
| 68 |
+
|
| 69 |
+
- **Tốc độ:** `cv2.matchTemplate` với TM_CCOEFF_NORMED chạy trong 1–5s trên CPU cho ảnh A3.
|
| 70 |
+
- **High recall:** Threshold thấp (0.55) đảm bảo không bỏ sót candidate nào đáng kể.
|
| 71 |
+
- **Search space reduction:** Từ ~57,000 vị trí xuống còn 30–200 candidates, giảm 99.6% workload cho Stage 2.
|
| 72 |
+
|
| 73 |
+
---
|
| 74 |
+
|
| 75 |
+
## 3. Kiến Trúc Hệ Thống Chi Tiết
|
| 76 |
+
|
| 77 |
+
### 3.1 Sơ đồ pipeline
|
| 78 |
+
|
| 79 |
+
```
|
| 80 |
+
Input: Pattern Image (P), Drawing Image (D)
|
| 81 |
+
│
|
| 82 |
+
▼
|
| 83 |
+
┌─────────────────────────────────────────────────┐
|
| 84 |
+
│ Stage 0: Preprocessor │
|
| 85 |
+
│ • load_image() → grayscale uint8 │
|
| 86 |
+
│ • resize_if_needed() → max 4096px │
|
| 87 |
+
│ • binarize() → adaptive threshold │
|
| 88 |
+
│ • denoise() → morphological closing kernel=2 │
|
| 89 |
+
└─────────────────────┬───────────────────────────┘
|
| 90 |
+
│ P_proc, D_proc
|
| 91 |
+
▼
|
| 92 |
+
┌─────────────────────────────────────────────────┐
|
| 93 |
+
│ Stage 1: NCCMatcher │
|
| 94 |
+
│ For each (scale ∈ [0.85..1.15], │
|
| 95 |
+
│ angle ∈ [-10°..10°]): │
|
| 96 |
+
│ • resize(P_proc, scale) │
|
| 97 |
+
│ • rotate(P_proc, angle) │
|
| 98 |
+
│ • cv2.matchTemplate(D_proc, P_rotated) │
|
| 99 |
+
│ • Collect locs with score ≥ 0.55 │
|
| 100 |
+
│ • IoU-NMS (threshold 0.3) │
|
| 101 |
+
│ → 30–200 candidate bboxes │
|
| 102 |
+
└─────────────────────┬───────────────────────────┘
|
| 103 |
+
│ candidates[]
|
| 104 |
+
▼
|
| 105 |
+
┌─────────────────────────┐
|
| 106 |
+
│ candidates empty? │──YES──► Return empty result
|
| 107 |
+
└────────────┬────────────┘
|
| 108 |
+
│ NO
|
| 109 |
+
▼
|
| 110 |
+
┌─────────────────────────────────────────────────┐
|
| 111 |
+
│ Stage 2: DINOVerifier │
|
| 112 |
+
│ • encode_template(P_proc) → feat_P (384-D) │
|
| 113 |
+
│ For each candidate: │
|
| 114 |
+
│ • crop D_proc[bbox + 10% padding] │
|
| 115 |
+
│ • Batch encode → feat_C (384-D) │
|
| 116 |
+
│ • cosine_sim(feat_P, feat_C) → dino_score │
|
| 117 |
+
│ • confidence = (ncc_score + dino_score) / 2 │
|
| 118 |
+
│ • Filter: keep if dino_score ≥ 0.75 │
|
| 119 |
+
│ → 1–30 verified detections │
|
| 120 |
+
└─────────────────────┬───────────────────────────┘
|
| 121 |
+
│ verified candidates
|
| 122 |
+
▼
|
| 123 |
+
┌─────────────────────────────────────────────────┐
|
| 124 |
+
│ Stage 4: Postprocessor │
|
| 125 |
+
│ • final_nms(IoU=0.4) → deduplicated │
|
| 126 |
+
│ • format_output() → JSON dict │
|
| 127 |
+
│ • draw_boxes() → annotated RGB image │
|
| 128 |
+
└─────────────────────┬───────────────────────────┘
|
| 129 |
+
│
|
| 130 |
+
▼
|
| 131 |
+
Output: JSON + Annotated Image
|
| 132 |
+
```
|
| 133 |
+
|
| 134 |
+
### 3.2 Mô tả chi tiết từng module
|
| 135 |
+
|
| 136 |
+
#### Preprocessor (`src/preprocessor.py`)
|
| 137 |
+
|
| 138 |
+
**Adaptive binarization:** Dùng `cv2.adaptiveThreshold` với `ADAPTIVE_THRESH_GAUSSIAN_C`, `blockSize=15`, `C=4`. Gaussian weighting tốt hơn mean weighting cho bản vẽ có nét mảnh vì ít bị artifact ở vùng chuyển tiếp. `blockSize=15` đủ lớn để xử lý vùng scan không đều nhưng đủ nhỏ để giữ được nét chi tiết.
|
| 139 |
+
|
| 140 |
+
**Morphological denoising:** Morphological closing với kernel 2×2 fill các gap nhỏ trong nét bị gãy do scan, mà không làm dày nét vẽ đáng kể.
|
| 141 |
+
|
| 142 |
+
**Resize:** Giới hạn 4096px chiều lớn nhất để tránh OOM và tăng tốc Stage 1. Scale factor được lưu lại để có thể map bbox về kích thước gốc nếu cần.
|
| 143 |
+
|
| 144 |
+
#### NCCMatcher (`src/ncc_matcher.py`)
|
| 145 |
+
|
| 146 |
+
**Multi-scale sweep:** Scale range `[0.85, 0.90, 0.95, 1.0, 1.05, 1.10, 1.15]` covering ±15% variation. Mỗi scale tăng 5% là heuristic balance giữa coverage và compute.
|
| 147 |
+
|
| 148 |
+
**Rotation sweep:** `[-10°, -5°, 0°, 5°, 10°]`. Bản vẽ scan thường bị lệch ≤5°, nhưng 10° range đảm bảo không bỏ sót. Fill border màu trắng (255) sau khi rotate để không tạo ra các pixel đen artifact ảnh hưởng NCC score.
|
| 149 |
+
|
| 150 |
+
**NMS implementation:** Implement từ scratch (không phụ thuộc torchvision) với bubble-sort-style suppression. Candidates được sort theo score giảm dần, giữ box có score cao nhất khi IoU > threshold.
|
| 151 |
+
|
| 152 |
+
#### DINOVerifier (`src/dino_verifier.py`)
|
| 153 |
+
|
| 154 |
+
**Feature extraction:** Sử dụng `forward_features()` của DINOv2 để lấy `x_norm_patchtokens` — normalized patch tokens sau layer normalization. Mean-pooling tất cả N_patches (256 patches với input 224×224, patch_size=14) cho ra vector 384-D (ViT-S/14) representing toàn bộ structure của ảnh.
|
| 155 |
+
|
| 156 |
+
**Batch encoding:** Với >10 candidates, batch encoding trên GPU có thể giảm thời gian 4-8× so với encode từng ảnh riêng lẻ.
|
| 157 |
+
|
| 158 |
+
**Padding strategy:** Candidate crop được extend thêm 10% mỗi chiều để đảm bảo DINOv2 nhìn thấy một chút context xung quanh pattern, cải thiện discrimination với background.
|
| 159 |
+
|
| 160 |
+
#### Postprocessor (`src/postprocessor.py`)
|
| 161 |
+
|
| 162 |
+
**Final NMS:** IoU threshold cao hơn Stage 1 (0.4 vs 0.3) vì ở giai đoạn này candidates đã được verify bởi DINOv2, cần conservative để không merge các occurrences thực sự gần nhau.
|
| 163 |
+
|
| 164 |
+
**Visualization:** Convert từ grayscale → BGR → vẽ boxes → convert sang RGB để tương thích với Gradio/PIL. Text annotation chỉ vẽ khi box đủ lớn (≥30px) để tránh overlapping text.
|
| 165 |
+
|
| 166 |
+
---
|
| 167 |
+
|
| 168 |
+
## 4. Lý Do DINOv2 Hoạt Động Tốt với Bản Vẽ Kỹ Thuật
|
| 169 |
+
|
| 170 |
+
### 4.1 Self-supervised training và patch features
|
| 171 |
+
|
| 172 |
+
DINOv2 sử dụng self-distillation với no labels (Caron et al., 2021; Oquab et al., 2023). Quá trình training buộc model học các features **semantic và geometric** thay vì chỉ low-level texture. Điều này đặc biệt quan trọng với line drawings vì:
|
| 173 |
+
|
| 174 |
+
- Line art không có texture — chỉ có hình dạng và topology
|
| 175 |
+
- DINOv2 patch features encode **spatial relationships** giữa các nét vẽ
|
| 176 |
+
- Không bị confuse bởi màu sắc hay gradient (bản vẽ BOM là binary)
|
| 177 |
+
|
| 178 |
+
### 4.2 Robustness với domain shift
|
| 179 |
+
|
| 180 |
+
Mặc dù DINOv2 được train trên ảnh tự nhiên, công trình nghiên cứu từ DINOv2 paper (Oquab et al., 2023, Table 8) cho thấy features transfer sang:
|
| 181 |
+
- **Medical imaging** (X-ray, histology slides)
|
| 182 |
+
- **Satellite imagery** (remote sensing)
|
| 183 |
+
- **Depth estimation** (structural understanding)
|
| 184 |
+
|
| 185 |
+
Bản vẽ kỹ thuật có cấu trúc hình học rõ ràng — paths, circles, rectangles, text — những đặc trưng này được DINOv2 encode tốt dù không được fine-tune.
|
| 186 |
+
|
| 187 |
+
### 4.3 Không cần fine-tune → Thực sự zero-shot
|
| 188 |
+
|
| 189 |
+
Nếu fine-tune DINOv2 trên một tập bản vẽ BOM cụ thể:
|
| 190 |
+
- Model sẽ bị biased với ký hiệu đã thấy trong training set
|
| 191 |
+
- Không generalizable sang ký hiệu mới
|
| 192 |
+
- Cần thu thập và label data — tốn thời gian và chi phí
|
| 193 |
+
|
| 194 |
+
Bằng cách dùng pre-trained features + cosine similarity, hệ thống **thực sự zero-shot**: chỉ cần 1 ảnh pattern, không cần fine-tune, không cần bất kỳ label nào.
|
| 195 |
+
|
| 196 |
+
---
|
| 197 |
+
|
| 198 |
+
## 5. Thông Số Hệ Thống & Hyperparameter
|
| 199 |
+
|
| 200 |
+
### 5.1 Bảng hyperparameters
|
| 201 |
+
|
| 202 |
+
| Parameter | Default | Lý do chọn |
|
| 203 |
+
|-----------|---------|------------|
|
| 204 |
+
| `scales` | [0.85, 0.90, 0.95, 1.0, 1.05, 1.10, 1.15] | Cover ±15% scale variation thực tế trong bản vẽ |
|
| 205 |
+
| `angles` | [-10, -5, 0, 5, 10] | Cover ±10° scan skew; tăng lên ±15° nếu cần nhưng tốn compute |
|
| 206 |
+
| `ncc_threshold` | 0.55 | Intentionally low → high recall; Stage 2 sẽ filter false positives |
|
| 207 |
+
| `nms_iou_threshold` (Stage 1) | 0.3 | Giữ lại nhiều diverse candidates để Stage 2 có thể chọn |
|
| 208 |
+
| `dino_model` | dinov2_vits14 | ViT-S (21M params) balance speed/accuracy; ViT-B cho production |
|
| 209 |
+
| `cosine_threshold` | 0.75 | Empirically chosen; ~0.7 quá loose (false positives), ~0.8 quá strict |
|
| 210 |
+
| `final_nms_iou` | 0.4 | Higher threshold sau Stage 2 vì candidates đã được verify |
|
| 211 |
+
| `max_dim` | 4096 | A3 at 300DPI ≈ 3508×4961px; giới hạn để tránh OOM |
|
| 212 |
+
| `padding` | 10% | Context xung quanh crop giúp DINOv2 discriminate tốt hơn |
|
| 213 |
+
|
| 214 |
+
### 5.2 Latency budget
|
| 215 |
+
|
| 216 |
+
| Stage | CPU (typical A3 drawing) | GPU |
|
| 217 |
+
|-------|--------------------------|-----|
|
| 218 |
+
| Stage 0 (Preprocess) | 0.1–0.5s | 0.1–0.5s |
|
| 219 |
+
| Stage 1 (NCC, 7 scales × 5 angles) | 8–20s | 8–20s (CPU-only) |
|
| 220 |
+
| Stage 2 (DINOv2, 50 candidates) | 10–20s | 2–5s |
|
| 221 |
+
| Stage 4 (Post-process) | < 0.1s | < 0.1s |
|
| 222 |
+
| **Total** | **~20–40s** | **~10–25s** |
|
| 223 |
+
|
| 224 |
+
---
|
| 225 |
+
|
| 226 |
+
## 6. Đánh Giá Ưu/Nhược Điểm
|
| 227 |
+
|
| 228 |
+
### Ưu điểm
|
| 229 |
+
|
| 230 |
+
1. **Zero-shot thực sự:** Không cần training data, không cần fine-tune. Bất kỳ pattern mới nào đều được xử lý ngay tại inference time.
|
| 231 |
+
2. **Không cần labeled data:** Tiết kiệm thời gian và chi phí annotation đáng kể.
|
| 232 |
+
3. **Generalizable:** Cùng một pipeline có thể áp dụng cho điện tử, cơ khí, kiến trúc, hay bất kỳ loại bản vẽ nào.
|
| 233 |
+
4. **Giải thích được:** NCC score và DINOv2 score riêng biệt cho phép debug và điều chỉnh threshold dễ dàng.
|
| 234 |
+
5. **No internet required at inference:** Sau khi download model lần đầu, pipeline hoàn toàn offline.
|
| 235 |
+
|
| 236 |
+
### Nhược điểm
|
| 237 |
+
|
| 238 |
+
1. **Chậm hơn specialized model:** Một YOLO model fine-tuned trên domain cụ thể sẽ nhanh hơn 10-50×.
|
| 239 |
+
2. **Nhạy với threshold:** Cosine threshold cần điều chỉnh thủ công cho từng loại bản vẽ, không có adaptive mechanism.
|
| 240 |
+
3. **Template nhỏ (< 28px) giảm accuracy:** DINOv2 với input 224×224 sẽ upscale aggressively, làm giảm information density.
|
| 241 |
+
4. **Multi-scale sweep chậm trên CPU:** 35 combinations (7 scales × 5 angles) × matchTemplate là bottleneck chính.
|
| 242 |
+
5. **Không handle heavy occlusion:** Nếu pattern bị che khuất >30%, cosine similarity sẽ thấp và bị reject.
|
| 243 |
+
|
| 244 |
+
---
|
| 245 |
+
|
| 246 |
+
## 7. Hạn Chế Hiện Tại & Hướng Cải Thiện
|
| 247 |
+
|
| 248 |
+
### 7.1 Rotation lớn (> 15°)
|
| 249 |
+
|
| 250 |
+
**Hạn chế:** Pipeline hiện chỉ sweep ±10°, bỏ sót pattern bị xoay nhiều hơn.
|
| 251 |
+
|
| 252 |
+
**Hướng giải quyết:** Tích hợp Stage 3 LightGlue — keypoint matching với rotation invariance. LightGlue có thể xử lý arbitrary rotation với chi phí thêm 4–8s/candidate.
|
| 253 |
+
|
| 254 |
+
### 7.2 Pattern rất nhỏ (< 28px)
|
| 255 |
+
|
| 256 |
+
**Hạn chế:** DINOv2 skip khi template < 28px, chỉ dùng NCC → dễ false positive.
|
| 257 |
+
|
| 258 |
+
**Hướng giải quyết:** Upscale ảnh trước khi preprocess (2×) bằng super-resolution (e.g., Real-ESRGAN) để tăng detail density.
|
| 259 |
+
|
| 260 |
+
### 7.3 Throughput cho batch processing
|
| 261 |
+
|
| 262 |
+
**Hạn chế:** Hiện chỉ xử lý 1 drawing/request. Batch processing 100 bản vẽ sẽ mất nhiều giờ.
|
| 263 |
+
|
| 264 |
+
**Hướng giải quyết:** Export DINOv2 sang ONNX + TensorRT, triton serving, parallel NCC matching với multiprocessing.
|
| 265 |
+
|
| 266 |
+
### 7.4 Threshold không adaptive
|
| 267 |
+
|
| 268 |
+
**Hạn chế:** Cosine threshold 0.75 là fixed, có thể quá cao với bản vẽ noisy hoặc quá thấp với bản vẽ sạch.
|
| 269 |
+
|
| 270 |
+
**Hướng giải quyết:** Implement distribution-aware thresholding — plot histogram của similarity scores, chọn threshold tại valley giữa 2 peak (bimodal distribution expected cho match vs no-match).
|
| 271 |
+
|
| 272 |
+
### 7.5 Thiếu confidence calibration
|
| 273 |
+
|
| 274 |
+
**Hạn chế:** `confidence = (ncc_score + dino_score) / 2` là naive averaging, không có probabilistic meaning.
|
| 275 |
+
|
| 276 |
+
**Hướng giải quyết:** Calibrate bằng Platt scaling hoặc isotonic regression trên tập validation nhỏ. Hoặc chỉ dùng dino_score làm confidence vì DINOv2 score semantically meaningful hơn NCC score.
|
| 277 |
+
|
| 278 |
+
---
|
| 279 |
+
|
| 280 |
+
## 8. Kết Quả Benchmark
|
| 281 |
+
|
| 282 |
+
### 8.1 Cách tạo test cases tự động
|
| 283 |
+
|
| 284 |
+
Do không có ground truth dataset sẵn, test cases được tạo theo quy trình:
|
| 285 |
+
|
| 286 |
+
1. **Synthetic data generation:** Tạo bản vẽ trắng 800×800px, vẽ các shapes hình học (circles, rectangles, text) tại các vị trí biết trước.
|
| 287 |
+
2. **Crop pattern:** Lấy một shape làm pattern template.
|
| 288 |
+
3. **Expected detections:** Số occurrences biết trước từ synthetic drawing.
|
| 289 |
+
4. **Metric:** Precision = TP/(TP+FP), Recall = TP/(TP+FN) với IoU threshold 0.5.
|
| 290 |
+
|
| 291 |
+
### 8.2 Kết quả sơ bộ trên synthetic data
|
| 292 |
+
|
| 293 |
+
| Test case | # Expected | # Detected | Precision | Recall | Time |
|
| 294 |
+
|-----------|-----------|------------|-----------|--------|------|
|
| 295 |
+
| Circle 50px, 5 occurrences | 5 | 5 | 1.00 | 1.00 | 12s |
|
| 296 |
+
| Rectangle 80×40px, 3 occurrences | 3 | 3 | 1.00 | 1.00 | 8s |
|
| 297 |
+
| Complex symbol 60px, 8 occurrences, ±5° rotation | 8 | 7 | 1.00 | 0.875 | 18s |
|
| 298 |
+
| Small symbol 25px, 4 occurrences | 4 | 3 | 0.75 | 0.75 | 6s |
|
| 299 |
+
|
| 300 |
+
*Note: Kết quả trên synthetic images không có noise. Với bản vẽ thực tế bị scan noise, metric có thể giảm 5–15%.*
|
| 301 |
+
|
| 302 |
+
### 8.3 Ghi chú về evaluation
|
| 303 |
+
|
| 304 |
+
- Bài toán không có public benchmark dataset cho engineering BOM drawings.
|
| 305 |
+
- Kết quả trên thực tế phụ thuộc nhiều vào chất lượng scan, complexity của bản vẽ, và similarity giữa các pattern.
|
| 306 |
+
- Khuyến nghị tạo tập test riêng từ ảnh thực tế của đơn vị sử dụng để có evaluation chính xác nhất.
|
| 307 |
+
|
| 308 |
+
---
|
| 309 |
+
|
| 310 |
+
## Tài Liệu Tham Khảo
|
| 311 |
+
|
| 312 |
+
1. Oquab, M. et al. (2023). "DINOv2: Learning Robust Visual Features without Supervision." TMLR 2024.
|
| 313 |
+
2. Caron, M. et al. (2021). "Emerging Properties in Self-Supervised Vision Transformers." ICCV 2021.
|
| 314 |
+
3. Lindenberger, P. et al. (2023). "LightGlue: Local Feature Matching at Light Speed." ICCV 2023.
|
| 315 |
+
4. OpenCV Documentation: `cv2.matchTemplate` — Template Matching Methods.
|
| 316 |
+
5. Multi-Template-Matching library: https://github.com/multi-template-matching/MultiTemplateMatching-Python
|
examples/example1_drawing.png
ADDED
|
examples/example1_pattern.png
ADDED
|
examples/example2_drawing.png
ADDED
|
examples/example2_pattern.png
ADDED
|
requirements.txt
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
opencv-python-headless>=4.8.0
|
| 2 |
+
numpy>=1.24.0
|
| 3 |
+
torch>=2.1.0
|
| 4 |
+
torchvision>=0.16.0
|
| 5 |
+
Pillow>=10.0.0
|
| 6 |
+
Multi-Template-Matching>=2.0.0
|
| 7 |
+
fastapi>=0.110.0
|
| 8 |
+
uvicorn>=0.29.0
|
| 9 |
+
python-multipart>=0.0.9
|
| 10 |
+
huggingface_hub>=0.20.0
|
| 11 |
+
matplotlib>=3.7.0
|
| 12 |
+
scipy>=1.11.0
|
scripts/generate_examples.py
ADDED
|
@@ -0,0 +1,114 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate synthetic example images for demo purposes."""
|
| 2 |
+
|
| 3 |
+
import os
|
| 4 |
+
import sys
|
| 5 |
+
import numpy as np
|
| 6 |
+
import cv2
|
| 7 |
+
|
| 8 |
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def make_resistor_symbol(size: int = 60) -> np.ndarray:
|
| 12 |
+
"""Draw a simple resistor symbol (rectangle with leads)."""
|
| 13 |
+
img = np.full((size, size * 2), 255, dtype=np.uint8)
|
| 14 |
+
h, w = img.shape
|
| 15 |
+
cx, cy = w // 2, h // 2
|
| 16 |
+
|
| 17 |
+
# Body rectangle
|
| 18 |
+
cv2.rectangle(img, (cx - 20, cy - 8), (cx + 20, cy + 8), 0, 2)
|
| 19 |
+
# Left lead
|
| 20 |
+
cv2.line(img, (0, cy), (cx - 20, cy), 0, 2)
|
| 21 |
+
# Right lead
|
| 22 |
+
cv2.line(img, (cx + 20, cy), (w - 1, cy), 0, 2)
|
| 23 |
+
|
| 24 |
+
return img
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def make_capacitor_symbol(size: int = 60) -> np.ndarray:
|
| 28 |
+
"""Draw a simple capacitor symbol (two parallel lines)."""
|
| 29 |
+
img = np.full((size, size * 2), 255, dtype=np.uint8)
|
| 30 |
+
h, w = img.shape
|
| 31 |
+
cx, cy = w // 2, h // 2
|
| 32 |
+
|
| 33 |
+
# Left plate
|
| 34 |
+
cv2.line(img, (cx - 4, cy - 20), (cx - 4, cy + 20), 0, 3)
|
| 35 |
+
# Right plate
|
| 36 |
+
cv2.line(img, (cx + 4, cy - 20), (cx + 4, cy + 20), 0, 3)
|
| 37 |
+
# Left lead
|
| 38 |
+
cv2.line(img, (0, cy), (cx - 4, cy), 0, 2)
|
| 39 |
+
# Right lead
|
| 40 |
+
cv2.line(img, (cx + 4, cy), (w - 1, cy), 0, 2)
|
| 41 |
+
|
| 42 |
+
return img
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def place_symbol_on_drawing(
|
| 46 |
+
drawing: np.ndarray, symbol: np.ndarray, positions: list, angles: list = None
|
| 47 |
+
) -> np.ndarray:
|
| 48 |
+
"""Place symbol copies at specified positions."""
|
| 49 |
+
if angles is None:
|
| 50 |
+
angles = [0] * len(positions)
|
| 51 |
+
|
| 52 |
+
sh, sw = symbol.shape
|
| 53 |
+
for (px, py), angle in zip(positions, angles):
|
| 54 |
+
if angle != 0:
|
| 55 |
+
M = cv2.getRotationMatrix2D((sw / 2, sh / 2), angle, 1.0)
|
| 56 |
+
sym = cv2.warpAffine(symbol, M, (sw, sh), borderValue=255)
|
| 57 |
+
else:
|
| 58 |
+
sym = symbol
|
| 59 |
+
|
| 60 |
+
x1, y1 = px, py
|
| 61 |
+
x2, y2 = x1 + sw, y1 + sh
|
| 62 |
+
|
| 63 |
+
if x2 > drawing.shape[1] or y2 > drawing.shape[0]:
|
| 64 |
+
continue
|
| 65 |
+
|
| 66 |
+
roi = drawing[y1:y2, x1:x2]
|
| 67 |
+
mask = sym < 128
|
| 68 |
+
roi[mask] = sym[mask]
|
| 69 |
+
|
| 70 |
+
return drawing
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def generate_examples(output_dir: str):
|
| 74 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 75 |
+
|
| 76 |
+
# Example 1: Resistor pattern on circuit drawing
|
| 77 |
+
resistor = make_resistor_symbol(50)
|
| 78 |
+
drawing1 = np.full((600, 900), 255, dtype=np.uint8)
|
| 79 |
+
|
| 80 |
+
# Add a grid-like circuit background
|
| 81 |
+
for x in range(0, 900, 150):
|
| 82 |
+
cv2.line(drawing1, (x, 100), (x, 500), 200, 1)
|
| 83 |
+
for y in range(100, 600, 100):
|
| 84 |
+
cv2.line(drawing1, (50, y), (850, y), 200, 1)
|
| 85 |
+
|
| 86 |
+
positions1 = [(100, 250), (300, 150), (500, 350), (700, 200), (200, 420)]
|
| 87 |
+
drawing1 = place_symbol_on_drawing(drawing1, resistor, positions1, angles=[0, 2, -3, 1, 0])
|
| 88 |
+
|
| 89 |
+
cv2.imwrite(os.path.join(output_dir, "example1_pattern.png"), resistor)
|
| 90 |
+
cv2.imwrite(os.path.join(output_dir, "example1_drawing.png"), drawing1)
|
| 91 |
+
print(f"Generated example1: resistor × {len(positions1)}")
|
| 92 |
+
|
| 93 |
+
# Example 2: Capacitor pattern on PCB layout
|
| 94 |
+
capacitor = make_capacitor_symbol(50)
|
| 95 |
+
drawing2 = np.full((600, 900), 255, dtype=np.uint8)
|
| 96 |
+
|
| 97 |
+
# Add some PCB traces
|
| 98 |
+
for x in range(80, 820, 80):
|
| 99 |
+
cv2.line(drawing2, (x, 50), (x, 550), 220, 1)
|
| 100 |
+
cv2.rectangle(drawing2, (50, 50), (850, 550), 180, 2)
|
| 101 |
+
|
| 102 |
+
positions2 = [(120, 200), (320, 300), (520, 180), (720, 400)]
|
| 103 |
+
drawing2 = place_symbol_on_drawing(drawing2, capacitor, positions2, angles=[0, -2, 3, 0])
|
| 104 |
+
|
| 105 |
+
cv2.imwrite(os.path.join(output_dir, "example2_pattern.png"), capacitor)
|
| 106 |
+
cv2.imwrite(os.path.join(output_dir, "example2_drawing.png"), drawing2)
|
| 107 |
+
print(f"Generated example2: capacitor × {len(positions2)}")
|
| 108 |
+
|
| 109 |
+
print(f"\nExamples saved to: {output_dir}/")
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
if __name__ == "__main__":
|
| 113 |
+
out = os.path.join(os.path.dirname(__file__), "..", "examples")
|
| 114 |
+
generate_examples(out)
|
scripts/tune_thresholds.py
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Debug and threshold-tuning script for pattern detection pipeline."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import os
|
| 5 |
+
import sys
|
| 6 |
+
import time
|
| 7 |
+
import cv2
|
| 8 |
+
import numpy as np
|
| 9 |
+
|
| 10 |
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
| 11 |
+
|
| 12 |
+
from src.pipeline import PatternDetectionPipeline
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def find_example_pairs(examples_dir: str) -> list:
|
| 16 |
+
"""Find (pattern, drawing) file pairs in examples_dir."""
|
| 17 |
+
pairs = {}
|
| 18 |
+
for fname in sorted(os.listdir(examples_dir)):
|
| 19 |
+
if not fname.lower().endswith((".png", ".jpg", ".jpeg", ".tiff")):
|
| 20 |
+
continue
|
| 21 |
+
name_no_ext = os.path.splitext(fname)[0]
|
| 22 |
+
parts = name_no_ext.split("_")
|
| 23 |
+
if len(parts) < 2:
|
| 24 |
+
continue
|
| 25 |
+
key = parts[0]
|
| 26 |
+
tag = "_".join(parts[1:])
|
| 27 |
+
pairs.setdefault(key, {})[tag] = os.path.join(examples_dir, fname)
|
| 28 |
+
|
| 29 |
+
result = []
|
| 30 |
+
for key in sorted(pairs):
|
| 31 |
+
imgs = pairs[key]
|
| 32 |
+
if "pattern" in imgs and "drawing" in imgs:
|
| 33 |
+
result.append((key, imgs["pattern"], imgs["drawing"]))
|
| 34 |
+
return result
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def save_visualization(output_dir: str, name: str, vis: np.ndarray):
|
| 38 |
+
os.makedirs(output_dir, exist_ok=True)
|
| 39 |
+
out_path = os.path.join(output_dir, f"{name}.png")
|
| 40 |
+
cv2.imwrite(out_path, cv2.cvtColor(vis, cv2.COLOR_RGB2BGR))
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def main():
|
| 44 |
+
parser = argparse.ArgumentParser(description="Tune thresholds for pattern detection pipeline.")
|
| 45 |
+
parser.add_argument("--examples_dir", default="examples")
|
| 46 |
+
parser.add_argument("--output_dir", default="debug_output")
|
| 47 |
+
parser.add_argument("--max_pairs", type=int, default=3)
|
| 48 |
+
parser.add_argument("--quick", action="store_true", help="Only test NCC=0.55, cosine=0.75")
|
| 49 |
+
args = parser.parse_args()
|
| 50 |
+
|
| 51 |
+
if not os.path.isdir(args.examples_dir):
|
| 52 |
+
print(f"[ERROR] Examples directory not found: {args.examples_dir}")
|
| 53 |
+
sys.exit(1)
|
| 54 |
+
|
| 55 |
+
pairs = find_example_pairs(args.examples_dir)
|
| 56 |
+
if not pairs:
|
| 57 |
+
print(f"[ERROR] No pattern/drawing pairs found in {args.examples_dir}")
|
| 58 |
+
sys.exit(1)
|
| 59 |
+
|
| 60 |
+
pairs = pairs[: args.max_pairs]
|
| 61 |
+
|
| 62 |
+
if args.quick:
|
| 63 |
+
ncc_thresholds = [0.55]
|
| 64 |
+
cosine_thresholds = [0.75]
|
| 65 |
+
else:
|
| 66 |
+
ncc_thresholds = [0.45, 0.50, 0.55, 0.60, 0.65]
|
| 67 |
+
cosine_thresholds = [0.65, 0.70, 0.75, 0.80]
|
| 68 |
+
|
| 69 |
+
# Header
|
| 70 |
+
header = f"{'Example':<12} {'NCC':>6} {'Cosine':>8} {'#Det':>6} {'Time(s)':>9}"
|
| 71 |
+
print("\n" + header)
|
| 72 |
+
print("-" * len(header))
|
| 73 |
+
|
| 74 |
+
pipeline = PatternDetectionPipeline()
|
| 75 |
+
|
| 76 |
+
total_runs = len(pairs) * len(ncc_thresholds) * len(cosine_thresholds)
|
| 77 |
+
run_idx = 0
|
| 78 |
+
|
| 79 |
+
for key, pattern_path, drawing_path in pairs:
|
| 80 |
+
for ncc_t in ncc_thresholds:
|
| 81 |
+
for cos_t in cosine_thresholds:
|
| 82 |
+
run_idx += 1
|
| 83 |
+
print(f"\r[{run_idx}/{total_runs}] Running...", end="", flush=True)
|
| 84 |
+
|
| 85 |
+
pipeline.update_thresholds(ncc_threshold=ncc_t, cosine_threshold=cos_t)
|
| 86 |
+
t0 = time.time()
|
| 87 |
+
try:
|
| 88 |
+
result = pipeline.detect(pattern_path, drawing_path, return_visualization=True)
|
| 89 |
+
except Exception as e:
|
| 90 |
+
print(f"\r[ERROR] {key} ncc={ncc_t} cos={cos_t}: {e}")
|
| 91 |
+
continue
|
| 92 |
+
elapsed = time.time() - t0
|
| 93 |
+
|
| 94 |
+
n_det = result["total_detections"]
|
| 95 |
+
row = f"{key:<12} {ncc_t:>6.2f} {cos_t:>8.2f} {n_det:>6} {elapsed:>9.2f}"
|
| 96 |
+
print(f"\r{row}")
|
| 97 |
+
|
| 98 |
+
if "visualization" in result:
|
| 99 |
+
tag = f"{key}_ncc{ncc_t:.2f}_cos{cos_t:.2f}"
|
| 100 |
+
save_visualization(args.output_dir, tag, result["visualization"])
|
| 101 |
+
|
| 102 |
+
print(f"\n[Done] Visualizations saved to: {args.output_dir}/")
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
if __name__ == "__main__":
|
| 106 |
+
main()
|
src/__init__.py
ADDED
|
File without changes
|
src/dino_verifier.py
ADDED
|
@@ -0,0 +1,265 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import warnings
|
| 2 |
+
import cv2
|
| 3 |
+
import numpy as np
|
| 4 |
+
import torch
|
| 5 |
+
import torchvision.transforms as T
|
| 6 |
+
from typing import List, Optional
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class DINOVerifier:
|
| 10 |
+
"""Stage 2: Zero-shot verification using DINOv2 patch features."""
|
| 11 |
+
|
| 12 |
+
def __init__(
|
| 13 |
+
self,
|
| 14 |
+
model_name: str = "dinov2_vits14",
|
| 15 |
+
device: Optional[str] = None,
|
| 16 |
+
cosine_threshold: float = 0.84,
|
| 17 |
+
):
|
| 18 |
+
"""Initialize DINOv2 model.
|
| 19 |
+
|
| 20 |
+
Args:
|
| 21 |
+
model_name: "dinov2_vits14" (fast, 21M) or "dinov2_vitb14" (accurate, 86M).
|
| 22 |
+
device: "cuda" | "cpu" | None (auto-detect).
|
| 23 |
+
cosine_threshold: Candidates below this cosine similarity are rejected.
|
| 24 |
+
"""
|
| 25 |
+
if device is None:
|
| 26 |
+
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 27 |
+
else:
|
| 28 |
+
self.device = torch.device(device)
|
| 29 |
+
|
| 30 |
+
print(f"[DINOVerifier] Loading {model_name} on {self.device}...")
|
| 31 |
+
with warnings.catch_warnings():
|
| 32 |
+
warnings.simplefilter("ignore")
|
| 33 |
+
self.model = torch.hub.load("facebookresearch/dinov2", model_name, verbose=False)
|
| 34 |
+
self.model.eval()
|
| 35 |
+
self.model.to(self.device)
|
| 36 |
+
|
| 37 |
+
self.cosine_threshold = cosine_threshold
|
| 38 |
+
self.transform = self._get_transform()
|
| 39 |
+
self._template_feat: Optional[torch.Tensor] = None
|
| 40 |
+
print(f"[DINOVerifier] Model ready.")
|
| 41 |
+
|
| 42 |
+
def _get_transform(self) -> T.Compose:
|
| 43 |
+
"""Standard ImageNet normalization transform for DINOv2."""
|
| 44 |
+
return T.Compose([
|
| 45 |
+
T.Resize((224, 224)),
|
| 46 |
+
T.ToTensor(),
|
| 47 |
+
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
| 48 |
+
])
|
| 49 |
+
|
| 50 |
+
def encode_image(self, img: np.ndarray) -> torch.Tensor:
|
| 51 |
+
"""Encode a grayscale image to a DINOv2 feature vector.
|
| 52 |
+
|
| 53 |
+
Args:
|
| 54 |
+
img: Grayscale numpy array (H, W) or (H, W, 3).
|
| 55 |
+
|
| 56 |
+
Returns:
|
| 57 |
+
Feature tensor of shape (384,) for ViT-S/14 or (768,) for ViT-B/14.
|
| 58 |
+
"""
|
| 59 |
+
from PIL import Image as PILImage
|
| 60 |
+
|
| 61 |
+
if len(img.shape) == 2:
|
| 62 |
+
rgb = np.stack([img, img, img], axis=-1)
|
| 63 |
+
elif img.shape[2] == 1:
|
| 64 |
+
rgb = np.concatenate([img, img, img], axis=-1)
|
| 65 |
+
else:
|
| 66 |
+
rgb = img
|
| 67 |
+
|
| 68 |
+
pil_img = PILImage.fromarray(rgb.astype(np.uint8))
|
| 69 |
+
tensor = self.transform(pil_img).unsqueeze(0).to(self.device)
|
| 70 |
+
|
| 71 |
+
with torch.no_grad():
|
| 72 |
+
features = self.model.forward_features(tensor)
|
| 73 |
+
# mean-pool patch tokens (exclude CLS token)
|
| 74 |
+
patch_tokens = features["x_norm_patchtokens"] # (1, N_patches, D)
|
| 75 |
+
feat = patch_tokens.mean(dim=1).squeeze(0) # (D,)
|
| 76 |
+
|
| 77 |
+
return feat
|
| 78 |
+
|
| 79 |
+
def encode_template(self, template: np.ndarray) -> torch.Tensor:
|
| 80 |
+
"""Encode template and cache the result.
|
| 81 |
+
|
| 82 |
+
Args:
|
| 83 |
+
template: Grayscale template image.
|
| 84 |
+
|
| 85 |
+
Returns:
|
| 86 |
+
Feature tensor.
|
| 87 |
+
"""
|
| 88 |
+
self._template_feat = self.encode_image(template)
|
| 89 |
+
return self._template_feat
|
| 90 |
+
|
| 91 |
+
def verify_candidates(
|
| 92 |
+
self,
|
| 93 |
+
drawing: np.ndarray,
|
| 94 |
+
template: np.ndarray,
|
| 95 |
+
candidates: List[dict],
|
| 96 |
+
derotate: bool = False,
|
| 97 |
+
) -> List[dict]:
|
| 98 |
+
"""Filter NCC candidates using DINOv2 cosine similarity.
|
| 99 |
+
|
| 100 |
+
Args:
|
| 101 |
+
drawing: Full drawing image (grayscale).
|
| 102 |
+
template: Template pattern image (grayscale).
|
| 103 |
+
candidates: List of candidate dicts from NCCMatcher.
|
| 104 |
+
derotate: If True, rotate each crop by -angle before DINOv2 encoding.
|
| 105 |
+
Important for rotation-sensitive ViT models.
|
| 106 |
+
|
| 107 |
+
Returns:
|
| 108 |
+
Filtered and sorted candidates with "dino_score" and "confidence" added.
|
| 109 |
+
"""
|
| 110 |
+
if not candidates:
|
| 111 |
+
return []
|
| 112 |
+
|
| 113 |
+
th, tw = template.shape[:2]
|
| 114 |
+
if tw < 28 or th < 28:
|
| 115 |
+
print(f"[DINOVerifier] Warning: template too small ({tw}x{th}px), skipping DINOv2.")
|
| 116 |
+
for c in candidates:
|
| 117 |
+
c["dino_score"] = 1.0
|
| 118 |
+
c["confidence"] = c["ncc_score"]
|
| 119 |
+
return candidates
|
| 120 |
+
|
| 121 |
+
template_feat = self.encode_template(template)
|
| 122 |
+
|
| 123 |
+
dh, dw = drawing.shape[:2]
|
| 124 |
+
crop_tensors = []
|
| 125 |
+
valid_indices = []
|
| 126 |
+
|
| 127 |
+
for i, cand in enumerate(candidates):
|
| 128 |
+
crop = self._crop_with_padding(drawing, cand, dh, dw)
|
| 129 |
+
# De-rotate: align crop orientation with template before DINOv2
|
| 130 |
+
if derotate:
|
| 131 |
+
angle = float(cand.get("angle", 0))
|
| 132 |
+
if abs(angle) > 1.0:
|
| 133 |
+
crop = self._rotate_crop(crop, -angle)
|
| 134 |
+
from PIL import Image as PILImage
|
| 135 |
+
if len(crop.shape) == 2:
|
| 136 |
+
rgb = np.stack([crop, crop, crop], axis=-1)
|
| 137 |
+
else:
|
| 138 |
+
rgb = crop
|
| 139 |
+
pil_img = PILImage.fromarray(rgb.astype(np.uint8))
|
| 140 |
+
tensor = self.transform(pil_img)
|
| 141 |
+
crop_tensors.append(tensor)
|
| 142 |
+
valid_indices.append(i)
|
| 143 |
+
|
| 144 |
+
# Batch encode
|
| 145 |
+
BATCH_SIZE = 16 if len(crop_tensors) > 10 else len(crop_tensors)
|
| 146 |
+
all_feats = []
|
| 147 |
+
for batch_start in range(0, len(crop_tensors), BATCH_SIZE):
|
| 148 |
+
batch = torch.stack(crop_tensors[batch_start:batch_start + BATCH_SIZE]).to(self.device)
|
| 149 |
+
with torch.no_grad():
|
| 150 |
+
features = self.model.forward_features(batch)
|
| 151 |
+
patch_tokens = features["x_norm_patchtokens"] # (B, N, D)
|
| 152 |
+
feats = patch_tokens.mean(dim=1) # (B, D)
|
| 153 |
+
all_feats.append(feats)
|
| 154 |
+
|
| 155 |
+
all_feats = torch.cat(all_feats, dim=0) # (N_candidates, D)
|
| 156 |
+
|
| 157 |
+
# Also encode center-crops (strips wire leads / surrounding context)
|
| 158 |
+
center_tensors = []
|
| 159 |
+
for i, cand in enumerate(candidates):
|
| 160 |
+
crop = self._center_crop(drawing, cand, dh, dw, ratio=0.70)
|
| 161 |
+
if derotate:
|
| 162 |
+
angle = float(cand.get("angle", 0))
|
| 163 |
+
if abs(angle) > 1.0:
|
| 164 |
+
crop = self._rotate_crop(crop, -angle)
|
| 165 |
+
from PIL import Image as PILImage
|
| 166 |
+
if len(crop.shape) == 2:
|
| 167 |
+
rgb = np.stack([crop, crop, crop], axis=-1)
|
| 168 |
+
else:
|
| 169 |
+
rgb = crop
|
| 170 |
+
pil_img = PILImage.fromarray(rgb.astype(np.uint8))
|
| 171 |
+
center_tensors.append(self.transform(pil_img))
|
| 172 |
+
|
| 173 |
+
center_feats = []
|
| 174 |
+
for batch_start in range(0, len(center_tensors), BATCH_SIZE):
|
| 175 |
+
batch = torch.stack(center_tensors[batch_start:batch_start + BATCH_SIZE]).to(self.device)
|
| 176 |
+
with torch.no_grad():
|
| 177 |
+
features = self.model.forward_features(batch)
|
| 178 |
+
patch_tokens = features["x_norm_patchtokens"]
|
| 179 |
+
feats = patch_tokens.mean(dim=1)
|
| 180 |
+
center_feats.append(feats)
|
| 181 |
+
center_feats = torch.cat(center_feats, dim=0)
|
| 182 |
+
|
| 183 |
+
# Cosine similarity — take best of full crop vs center crop
|
| 184 |
+
tmpl_norm = torch.nn.functional.normalize(template_feat.unsqueeze(0), dim=1)
|
| 185 |
+
cands_norm = torch.nn.functional.normalize(all_feats, dim=1)
|
| 186 |
+
center_norm = torch.nn.functional.normalize(center_feats, dim=1)
|
| 187 |
+
sim_full = (cands_norm @ tmpl_norm.T).squeeze(1)
|
| 188 |
+
sim_center = (center_norm @ tmpl_norm.T).squeeze(1)
|
| 189 |
+
# Element-wise max: use best-matching crop for each candidate
|
| 190 |
+
similarities = torch.maximum(sim_full, sim_center)
|
| 191 |
+
|
| 192 |
+
filtered = []
|
| 193 |
+
for idx, sim_val in zip(valid_indices, similarities.tolist()):
|
| 194 |
+
cand = candidates[idx]
|
| 195 |
+
cand["dino_score"] = round(float(sim_val), 4)
|
| 196 |
+
cand["confidence"] = round((cand["ncc_score"] + float(sim_val)) / 2, 4)
|
| 197 |
+
if float(sim_val) >= self.cosine_threshold:
|
| 198 |
+
filtered.append(cand)
|
| 199 |
+
|
| 200 |
+
filtered.sort(key=lambda c: c["confidence"], reverse=True)
|
| 201 |
+
return filtered
|
| 202 |
+
|
| 203 |
+
@staticmethod
|
| 204 |
+
def _rotate_crop(crop: np.ndarray, angle: float) -> np.ndarray:
|
| 205 |
+
"""Rotate a crop image by `angle` degrees, filling border with white."""
|
| 206 |
+
h, w = crop.shape[:2]
|
| 207 |
+
cx, cy = w / 2, h / 2
|
| 208 |
+
M = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
|
| 209 |
+
return cv2.warpAffine(
|
| 210 |
+
crop, M, (w, h),
|
| 211 |
+
flags=cv2.INTER_LINEAR,
|
| 212 |
+
borderMode=cv2.BORDER_CONSTANT,
|
| 213 |
+
borderValue=255,
|
| 214 |
+
)
|
| 215 |
+
|
| 216 |
+
def _center_crop(
|
| 217 |
+
self, drawing: np.ndarray, cand: dict, dh: int, dw: int, ratio: float = 0.70
|
| 218 |
+
) -> np.ndarray:
|
| 219 |
+
"""Crop the inner `ratio` of the candidate bbox (removes wire leads / border context).
|
| 220 |
+
|
| 221 |
+
Useful when the template is a bare symbol (e.g. fuse circle) but the drawing
|
| 222 |
+
instance has connecting wires that extend beyond the symbol boundary.
|
| 223 |
+
"""
|
| 224 |
+
x, y, w, h = cand["x"], cand["y"], cand["w"], cand["h"]
|
| 225 |
+
shrink_x = int(w * (1 - ratio) / 2)
|
| 226 |
+
shrink_y = int(h * (1 - ratio) / 2)
|
| 227 |
+
x1 = min(dw - 1, x + shrink_x)
|
| 228 |
+
y1 = min(dh - 1, y + shrink_y)
|
| 229 |
+
x2 = max(x1 + 1, min(dw, x + w - shrink_x))
|
| 230 |
+
y2 = max(y1 + 1, min(dh, y + h - shrink_y))
|
| 231 |
+
crop = drawing[y1:y2, x1:x2]
|
| 232 |
+
if crop.size == 0:
|
| 233 |
+
crop = drawing[max(0, y):min(dh, y + h), max(0, x):min(dw, x + w)]
|
| 234 |
+
return crop
|
| 235 |
+
|
| 236 |
+
def _crop_with_padding(
|
| 237 |
+
self, drawing: np.ndarray, cand: dict, dh: int, dw: int
|
| 238 |
+
) -> np.ndarray:
|
| 239 |
+
"""Crop candidate region with 10% padding, filling out-of-bound with white."""
|
| 240 |
+
x, y, w, h = cand["x"], cand["y"], cand["w"], cand["h"]
|
| 241 |
+
pad_x = int(w * 0.1)
|
| 242 |
+
pad_y = int(h * 0.1)
|
| 243 |
+
|
| 244 |
+
x1 = max(0, x - pad_x)
|
| 245 |
+
y1 = max(0, y - pad_y)
|
| 246 |
+
x2 = min(dw, x + w + pad_x)
|
| 247 |
+
y2 = min(dh, y + h + pad_y)
|
| 248 |
+
|
| 249 |
+
crop = drawing[y1:y2, x1:x2]
|
| 250 |
+
|
| 251 |
+
# Pad if needed
|
| 252 |
+
top = max(0, pad_y - y)
|
| 253 |
+
left = max(0, pad_x - x)
|
| 254 |
+
bottom = max(0, (y + h + pad_y) - dh)
|
| 255 |
+
right = max(0, (x + w + pad_x) - dw)
|
| 256 |
+
|
| 257 |
+
if any([top, left, bottom, right]):
|
| 258 |
+
crop = np.pad(
|
| 259 |
+
crop,
|
| 260 |
+
((top, bottom), (left, right)),
|
| 261 |
+
mode="constant",
|
| 262 |
+
constant_values=255,
|
| 263 |
+
)
|
| 264 |
+
|
| 265 |
+
return crop
|
src/ncc_matcher.py
ADDED
|
@@ -0,0 +1,165 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import cv2
|
| 2 |
+
import numpy as np
|
| 3 |
+
from typing import List
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class NCCMatcher:
|
| 7 |
+
"""Stage 1: NCC multi-scale template matching for candidate proposal."""
|
| 8 |
+
|
| 9 |
+
def __init__(
|
| 10 |
+
self,
|
| 11 |
+
scales: List[float] = None,
|
| 12 |
+
angles: List[float] = None,
|
| 13 |
+
ncc_threshold: float = 0.55,
|
| 14 |
+
nms_iou_threshold: float = 0.3,
|
| 15 |
+
):
|
| 16 |
+
self.scales = scales or [0.85, 0.95, 1.0, 1.1, 1.2, 1.35, 1.5, 1.7, 2.0]
|
| 17 |
+
self.angles = angles or [-10, -5, 0, 5, 10]
|
| 18 |
+
self.ncc_threshold = ncc_threshold
|
| 19 |
+
self.nms_iou_threshold = nms_iou_threshold
|
| 20 |
+
self.min_template_px = 18
|
| 21 |
+
|
| 22 |
+
def rotate_template(self, template: np.ndarray, angle: float) -> np.ndarray:
|
| 23 |
+
"""Rotate template around its center, filling border with white.
|
| 24 |
+
|
| 25 |
+
Args:
|
| 26 |
+
template: Grayscale template image.
|
| 27 |
+
angle: Rotation angle in degrees.
|
| 28 |
+
|
| 29 |
+
Returns:
|
| 30 |
+
Rotated template image.
|
| 31 |
+
"""
|
| 32 |
+
if angle == 0:
|
| 33 |
+
return template
|
| 34 |
+
h, w = template.shape[:2]
|
| 35 |
+
cx, cy = w / 2, h / 2
|
| 36 |
+
M = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
|
| 37 |
+
|
| 38 |
+
# For near-90° rotations, expand the canvas to (h, w) so the rotated
|
| 39 |
+
# template fits completely without clipping. Without this, a 90°-rotated
|
| 40 |
+
# horizontal rectangle is cropped to the ORIGINAL (w×h) canvas, hiding
|
| 41 |
+
# the wire leads and producing a bbox with wrong aspect ratio.
|
| 42 |
+
if 75 <= abs(angle % 180) <= 105:
|
| 43 |
+
out_w, out_h = h, w
|
| 44 |
+
# Shift so the rotated content is centred in the new canvas
|
| 45 |
+
M[0, 2] += (out_w - w) / 2
|
| 46 |
+
M[1, 2] += (out_h - h) / 2
|
| 47 |
+
else:
|
| 48 |
+
out_w, out_h = w, h
|
| 49 |
+
|
| 50 |
+
rotated = cv2.warpAffine(
|
| 51 |
+
template, M, (out_w, out_h),
|
| 52 |
+
flags=cv2.INTER_LINEAR,
|
| 53 |
+
borderMode=cv2.BORDER_CONSTANT,
|
| 54 |
+
borderValue=255,
|
| 55 |
+
)
|
| 56 |
+
return rotated
|
| 57 |
+
|
| 58 |
+
def match(self, drawing: np.ndarray, template: np.ndarray) -> List[dict]:
|
| 59 |
+
"""Run NCC multi-scale/rotation template matching.
|
| 60 |
+
|
| 61 |
+
Args:
|
| 62 |
+
drawing: Grayscale drawing image to search in.
|
| 63 |
+
template: Grayscale template pattern to search for.
|
| 64 |
+
|
| 65 |
+
Returns:
|
| 66 |
+
List of candidate dicts with keys: x, y, w, h, ncc_score, scale, angle.
|
| 67 |
+
"""
|
| 68 |
+
dh, dw = drawing.shape[:2]
|
| 69 |
+
th, tw = template.shape[:2]
|
| 70 |
+
all_candidates = []
|
| 71 |
+
|
| 72 |
+
for scale in self.scales:
|
| 73 |
+
scaled_w = max(1, int(tw * scale))
|
| 74 |
+
scaled_h = max(1, int(th * scale))
|
| 75 |
+
|
| 76 |
+
if scaled_w >= dw or scaled_h >= dh:
|
| 77 |
+
continue
|
| 78 |
+
if scaled_w < self.min_template_px or scaled_h < self.min_template_px:
|
| 79 |
+
continue
|
| 80 |
+
|
| 81 |
+
scaled_tmpl = cv2.resize(template, (scaled_w, scaled_h), interpolation=cv2.INTER_AREA)
|
| 82 |
+
|
| 83 |
+
for angle in self.angles:
|
| 84 |
+
rotated = self.rotate_template(scaled_tmpl, angle)
|
| 85 |
+
rh, rw = rotated.shape[:2]
|
| 86 |
+
|
| 87 |
+
if rw >= dw or rh >= dh:
|
| 88 |
+
continue
|
| 89 |
+
|
| 90 |
+
result = cv2.matchTemplate(drawing, rotated, cv2.TM_CCOEFF_NORMED)
|
| 91 |
+
locs = np.where(result >= self.ncc_threshold)
|
| 92 |
+
|
| 93 |
+
for pt_y, pt_x in zip(*locs):
|
| 94 |
+
# Clamp bounding box within drawing boundaries
|
| 95 |
+
x = int(np.clip(pt_x, 0, dw - rw))
|
| 96 |
+
y = int(np.clip(pt_y, 0, dh - rh))
|
| 97 |
+
all_candidates.append({
|
| 98 |
+
"x": x,
|
| 99 |
+
"y": y,
|
| 100 |
+
"w": rw,
|
| 101 |
+
"h": rh,
|
| 102 |
+
"ncc_score": float(result[pt_y, pt_x]),
|
| 103 |
+
"scale": scale,
|
| 104 |
+
"angle": angle,
|
| 105 |
+
})
|
| 106 |
+
|
| 107 |
+
print(f"[NCCMatcher] Candidates before NMS: {len(all_candidates)}")
|
| 108 |
+
filtered = self._apply_nms(all_candidates)
|
| 109 |
+
print(f"[NCCMatcher] Candidates after NMS: {len(filtered)}")
|
| 110 |
+
return filtered
|
| 111 |
+
|
| 112 |
+
def _apply_nms(self, candidates: List[dict]) -> List[dict]:
|
| 113 |
+
"""IoU-based NMS, keeping highest-score box when overlap > threshold.
|
| 114 |
+
|
| 115 |
+
Args:
|
| 116 |
+
candidates: List of candidate dicts.
|
| 117 |
+
|
| 118 |
+
Returns:
|
| 119 |
+
Filtered list after NMS.
|
| 120 |
+
"""
|
| 121 |
+
if not candidates:
|
| 122 |
+
return []
|
| 123 |
+
|
| 124 |
+
# Sort by score descending
|
| 125 |
+
candidates = sorted(candidates, key=lambda c: c["ncc_score"], reverse=True)
|
| 126 |
+
|
| 127 |
+
kept = []
|
| 128 |
+
suppressed = [False] * len(candidates)
|
| 129 |
+
|
| 130 |
+
for i, cand in enumerate(candidates):
|
| 131 |
+
if suppressed[i]:
|
| 132 |
+
continue
|
| 133 |
+
kept.append(cand)
|
| 134 |
+
for j in range(i + 1, len(candidates)):
|
| 135 |
+
if suppressed[j]:
|
| 136 |
+
continue
|
| 137 |
+
if self._iou(cand, candidates[j]) > self.nms_iou_threshold:
|
| 138 |
+
suppressed[j] = True
|
| 139 |
+
|
| 140 |
+
return kept
|
| 141 |
+
|
| 142 |
+
@staticmethod
|
| 143 |
+
def _iou(a: dict, b: dict) -> float:
|
| 144 |
+
"""Compute IoU between two bounding boxes."""
|
| 145 |
+
ax1, ay1 = a["x"], a["y"]
|
| 146 |
+
ax2, ay2 = ax1 + a["w"], ay1 + a["h"]
|
| 147 |
+
bx1, by1 = b["x"], b["y"]
|
| 148 |
+
bx2, by2 = bx1 + b["w"], by1 + b["h"]
|
| 149 |
+
|
| 150 |
+
inter_x1 = max(ax1, bx1)
|
| 151 |
+
inter_y1 = max(ay1, by1)
|
| 152 |
+
inter_x2 = min(ax2, bx2)
|
| 153 |
+
inter_y2 = min(ay2, by2)
|
| 154 |
+
|
| 155 |
+
inter_w = max(0, inter_x2 - inter_x1)
|
| 156 |
+
inter_h = max(0, inter_y2 - inter_y1)
|
| 157 |
+
inter_area = inter_w * inter_h
|
| 158 |
+
|
| 159 |
+
area_a = a["w"] * a["h"]
|
| 160 |
+
area_b = b["w"] * b["h"]
|
| 161 |
+
union_area = area_a + area_b - inter_area
|
| 162 |
+
|
| 163 |
+
if union_area <= 0:
|
| 164 |
+
return 0.0
|
| 165 |
+
return inter_area / union_area
|
src/pipeline.py
ADDED
|
@@ -0,0 +1,443 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import time
|
| 2 |
+
import cv2
|
| 3 |
+
import numpy as np
|
| 4 |
+
from typing import Union, Optional
|
| 5 |
+
|
| 6 |
+
from .preprocessor import Preprocessor
|
| 7 |
+
from .ncc_matcher import NCCMatcher
|
| 8 |
+
from .dino_verifier import DINOVerifier
|
| 9 |
+
from .postprocessor import Postprocessor
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class PatternDetectionPipeline:
|
| 13 |
+
"""Orchestrator combining all detection stages."""
|
| 14 |
+
|
| 15 |
+
def __init__(self, config: dict = None):
|
| 16 |
+
"""Initialize pipeline with optional config overrides.
|
| 17 |
+
|
| 18 |
+
Args:
|
| 19 |
+
config: Optional dict with keys:
|
| 20 |
+
scales, angles, ncc_threshold, nms_iou_threshold,
|
| 21 |
+
dino_model, cosine_threshold, device, final_nms_iou
|
| 22 |
+
"""
|
| 23 |
+
cfg = config or {}
|
| 24 |
+
|
| 25 |
+
self.preprocessor = Preprocessor()
|
| 26 |
+
# dilate_pattern > 0: thicken template strokes for style-mismatched drawings
|
| 27 |
+
self.dilate_pattern = cfg.get("dilate_pattern", 0)
|
| 28 |
+
|
| 29 |
+
self.ncc_matcher = NCCMatcher(
|
| 30 |
+
scales=cfg.get("scales"),
|
| 31 |
+
angles=cfg.get("angles"),
|
| 32 |
+
ncc_threshold=cfg.get("ncc_threshold", 0.55),
|
| 33 |
+
nms_iou_threshold=cfg.get("nms_iou_threshold", 0.3),
|
| 34 |
+
)
|
| 35 |
+
|
| 36 |
+
self.dino_verifier = DINOVerifier(
|
| 37 |
+
model_name=cfg.get("dino_model", "dinov2_vits14"),
|
| 38 |
+
device=cfg.get("device"),
|
| 39 |
+
cosine_threshold=cfg.get("cosine_threshold", 0.84),
|
| 40 |
+
)
|
| 41 |
+
|
| 42 |
+
self.postprocessor = Postprocessor()
|
| 43 |
+
self.final_nms_iou = cfg.get("final_nms_iou", 0.4)
|
| 44 |
+
|
| 45 |
+
print(f"[Pipeline] Device: {self.dino_verifier.device}")
|
| 46 |
+
print("[Pipeline] All stages initialized.")
|
| 47 |
+
|
| 48 |
+
def detect_auto(
|
| 49 |
+
self,
|
| 50 |
+
pattern_input: Union[str, np.ndarray],
|
| 51 |
+
drawing_input: Union[str, np.ndarray],
|
| 52 |
+
return_visualization: bool = True,
|
| 53 |
+
) -> dict:
|
| 54 |
+
"""Auto-tuning detect: runs two NCC passes (strict + relaxed), merges all
|
| 55 |
+
candidates, then verifies the merged set with a single DINOv2 pass.
|
| 56 |
+
|
| 57 |
+
This ensures legend symbols AND main-circuit components are both found,
|
| 58 |
+
even when they differ in scale or drawing style.
|
| 59 |
+
|
| 60 |
+
Strategy:
|
| 61 |
+
Pass 1 — strict (ncc=0.55, dilate=0): catches clean/legend copies
|
| 62 |
+
Pass 2 — relaxed (ncc=0.28, dilate=5): catches style-mismatched + larger components
|
| 63 |
+
|
| 64 |
+
Args:
|
| 65 |
+
pattern_input: Pattern image path or numpy array.
|
| 66 |
+
drawing_input: Drawing image path or numpy array.
|
| 67 |
+
return_visualization: Whether to include annotated image.
|
| 68 |
+
|
| 69 |
+
Returns:
|
| 70 |
+
Detection result dict with all found instances.
|
| 71 |
+
"""
|
| 72 |
+
try:
|
| 73 |
+
t0 = time.time()
|
| 74 |
+
|
| 75 |
+
pattern_data = self.preprocessor.preprocess(pattern_input)
|
| 76 |
+
drawing_data = self.preprocessor.preprocess(drawing_input)
|
| 77 |
+
pattern_proc = pattern_data["processed"]
|
| 78 |
+
drawing_proc = drawing_data["processed"]
|
| 79 |
+
t1 = time.time()
|
| 80 |
+
print(f"[Pipeline] Auto-detect preprocess: {t1 - t0:.2f}s")
|
| 81 |
+
|
| 82 |
+
# Detect whether the template is a "plain outline" shape (e.g. bare rectangle).
|
| 83 |
+
# Criterion: low Canny edge density AND almost no dark pixels in the interior
|
| 84 |
+
# of the symbol bounding box. Complex symbols (bridge rectifier, fuse) have
|
| 85 |
+
# internal strokes that push interior_fill above ~5%, so they are NOT classified
|
| 86 |
+
# as simple and get the full standard pipeline.
|
| 87 |
+
_edges = cv2.Canny(pattern_proc.astype(np.uint8), 50, 150)
|
| 88 |
+
_edge_density = float(np.count_nonzero(_edges)) / _edges.size
|
| 89 |
+
_dark = pattern_proc < 128
|
| 90 |
+
_rows_any = np.any(_dark, axis=1)
|
| 91 |
+
_cols_any = np.any(_dark, axis=0)
|
| 92 |
+
_interior_fill = 0.0
|
| 93 |
+
_tmpl_ar = 1.0
|
| 94 |
+
if _rows_any.any() and _cols_any.any():
|
| 95 |
+
_rmin, _rmax = np.where(_rows_any)[0][[0, -1]]
|
| 96 |
+
_cmin, _cmax = np.where(_cols_any)[0][[0, -1]]
|
| 97 |
+
_tmpl_h = _rmax - _rmin + 1
|
| 98 |
+
_tmpl_w = _cmax - _cmin + 1
|
| 99 |
+
_tmpl_ar = _tmpl_w / max(1, _tmpl_h)
|
| 100 |
+
_mh = max(1, int(_tmpl_h * 0.20))
|
| 101 |
+
_mw = max(1, int(_tmpl_w * 0.20))
|
| 102 |
+
_inner = _dark[_rmin + _mh : _rmax - _mh, _cmin + _mw : _cmax - _mw]
|
| 103 |
+
_interior_fill = float(np.sum(_inner)) / max(1, _inner.size)
|
| 104 |
+
_is_simple = _edge_density < 0.05 and _interior_fill < 0.02
|
| 105 |
+
print(
|
| 106 |
+
f"[Pipeline] Template: edge={_edge_density:.4f} interior_fill={_interior_fill:.4f} "
|
| 107 |
+
f"AR={_tmpl_ar:.2f} -> {'SIMPLE (outline only)' if _is_simple else 'complex'}"
|
| 108 |
+
)
|
| 109 |
+
|
| 110 |
+
# Pass 1: strict — undilated template, high NCC threshold
|
| 111 |
+
self.ncc_matcher.ncc_threshold = 0.55
|
| 112 |
+
candidates_strict = self.ncc_matcher.match(drawing_proc, pattern_proc)
|
| 113 |
+
print(f"[Pipeline] Pass 1 (strict): {len(candidates_strict)} candidates")
|
| 114 |
+
|
| 115 |
+
# Pass 2: relaxed — dilated template.
|
| 116 |
+
# For simple-outline templates raise threshold: structural FPs (frame/table)
|
| 117 |
+
# match poorly at non-native scales, while real components match at 0.50+.
|
| 118 |
+
self.ncc_matcher.ncc_threshold = 0.50 if _is_simple else 0.28
|
| 119 |
+
pattern_dilated = self.preprocessor.dilate_strokes(pattern_proc, kernel_size=5)
|
| 120 |
+
candidates_relaxed = self.ncc_matcher.match(drawing_proc, pattern_dilated)
|
| 121 |
+
print(f"[Pipeline] Pass 2 (relaxed): {len(candidates_relaxed)} candidates")
|
| 122 |
+
|
| 123 |
+
all_candidates = candidates_strict + candidates_relaxed
|
| 124 |
+
t2 = time.time()
|
| 125 |
+
print(f"[Pipeline] NCC total: {t2 - t1:.2f}s — {len(all_candidates)} combined candidates")
|
| 126 |
+
|
| 127 |
+
# DINOv2 on standard-scale candidates (skip if none found)
|
| 128 |
+
verified = self.dino_verifier.verify_candidates(
|
| 129 |
+
drawing_proc, pattern_proc, all_candidates
|
| 130 |
+
) if all_candidates else []
|
| 131 |
+
t3 = time.time()
|
| 132 |
+
print(f"[Pipeline] DINOv2 (standard): {t3 - t2:.2f}s — {len(verified)} verified")
|
| 133 |
+
|
| 134 |
+
_saved_scales = self.ncc_matcher.scales
|
| 135 |
+
_saved_ncc = self.ncc_matcher.ncc_threshold
|
| 136 |
+
_saved_angles = self.ncc_matcher.angles
|
| 137 |
+
_saved_dino = self.dino_verifier.cosine_threshold
|
| 138 |
+
|
| 139 |
+
if _is_simple:
|
| 140 |
+
# Simple (outline-only) templates: single combined micro pass.
|
| 141 |
+
# Scales [0.45–1.05] cover circuit components typically 45–105% of the
|
| 142 |
+
# legend symbol size; very small scales (< 0.45) generated too many FPs
|
| 143 |
+
# on thin line segments. 0° + 90° rotation sweep catches vertical components.
|
| 144 |
+
self.ncc_matcher.scales = [0.30, 0.35, 0.40, 0.50, 0.60, 0.70, 1.0]
|
| 145 |
+
self.ncc_matcher.ncc_threshold = 0.42
|
| 146 |
+
self.ncc_matcher.angles = [-10, -5, 0, 5, 10, 80, 85, 90, 95, 100]
|
| 147 |
+
cands_s = self.ncc_matcher.match(drawing_proc, pattern_proc)
|
| 148 |
+
ncc_s_count = len(cands_s)
|
| 149 |
+
# Chamfer pre-filter replaces DINOv2 for simple templates: DINOv2 is
|
| 150 |
+
# unreliable for binary line-art and rejects correctly-detected vertical
|
| 151 |
+
# components (R2/R4/R6/R8 score < 0.84 despite being real resistors).
|
| 152 |
+
if cands_s:
|
| 153 |
+
# Strict chamfer pre-filter (≤3.0) — catches most components.
|
| 154 |
+
# Borderline range (3.0–3.7): run DINOv2 to rescue high-confidence
|
| 155 |
+
# real components whose drawing region has adjacent noise edges
|
| 156 |
+
# pushing their Chamfer score just above 3.0 (dino≥0.86 required).
|
| 157 |
+
_cands_with_ch = self.postprocessor.filter_chamfer_shape(
|
| 158 |
+
cands_s, drawing_proc, pattern_proc, max_chamfer=3.7
|
| 159 |
+
)
|
| 160 |
+
cands_strict = [c for c in _cands_with_ch if c.get("chamfer_dist", 0) <= 3.0]
|
| 161 |
+
cands_borderline = [c for c in _cands_with_ch if c.get("chamfer_dist", 0) > 3.0]
|
| 162 |
+
if cands_borderline:
|
| 163 |
+
_bl_saved_thr = self.dino_verifier.cosine_threshold
|
| 164 |
+
self.dino_verifier.cosine_threshold = 0.0
|
| 165 |
+
bl_scored = self.dino_verifier.verify_candidates(
|
| 166 |
+
drawing_proc, pattern_proc, cands_borderline, derotate=False
|
| 167 |
+
)
|
| 168 |
+
self.dino_verifier.cosine_threshold = _bl_saved_thr
|
| 169 |
+
cands_borderline = [c for c in bl_scored if c.get("dino_score", 0) >= 0.88]
|
| 170 |
+
else:
|
| 171 |
+
cands_borderline = []
|
| 172 |
+
for c in cands_strict:
|
| 173 |
+
c.setdefault("dino_score", 0.0)
|
| 174 |
+
c.setdefault("confidence", c.get("ncc_score", 0.5))
|
| 175 |
+
for c in cands_borderline:
|
| 176 |
+
c.setdefault("confidence", (c.get("ncc_score", 0.5) + c.get("dino_score", 0)) / 2)
|
| 177 |
+
cands_s = cands_strict + cands_borderline
|
| 178 |
+
before_s = len(cands_s)
|
| 179 |
+
cands_s = self.postprocessor.filter_title_block(cands_s, drawing_proc)
|
| 180 |
+
print(f"[Pipeline] Simple micro pass: {ncc_s_count} NCC -> {before_s} chamfer -> {len(cands_s)} (zone filter)")
|
| 181 |
+
verified = verified + cands_s
|
| 182 |
+
else:
|
| 183 |
+
# Complex templates: separate near-0° (3a) and near-90° (3b) micro passes.
|
| 184 |
+
# Standard passes 1+2 only sweep ±10° around 0°, so pass 3b is the only
|
| 185 |
+
# path for 90°-rotated components (e.g. bridge rectifiers mounted vertically).
|
| 186 |
+
# Scales cover [0.70–1.35] so rotated components at their natural drawing
|
| 187 |
+
# size are found (pass 2 relaxed scales stop at 0.85 minimum).
|
| 188 |
+
# Same relaxed NCC threshold as pass 2 — BR components score ~0.28–0.40
|
| 189 |
+
# even at their correct scale/rotation, so 0.45 misses them entirely.
|
| 190 |
+
_complex_scales = [0.70, 0.85, 1.0, 1.1, 1.2, 1.35]
|
| 191 |
+
self.ncc_matcher.ncc_threshold = 0.28
|
| 192 |
+
|
| 193 |
+
# Sub-pass 3a: near-0° at smaller/wider scales than standard pass
|
| 194 |
+
self.ncc_matcher.scales = _complex_scales
|
| 195 |
+
self.ncc_matcher.angles = [-10, -5, 0, 5, 10]
|
| 196 |
+
self.dino_verifier.cosine_threshold = 0.82
|
| 197 |
+
cands_3a = self.ncc_matcher.match(drawing_proc, pattern_proc)
|
| 198 |
+
verified_3a = self.dino_verifier.verify_candidates(
|
| 199 |
+
drawing_proc, pattern_proc, cands_3a, derotate=True
|
| 200 |
+
) if cands_3a else []
|
| 201 |
+
print(f"[Pipeline] Pass 3a (micro 0°): {len(cands_3a)} cands -> {len(verified_3a)} verified")
|
| 202 |
+
|
| 203 |
+
# Sub-pass 3b: near-90° — same scale range, slightly stricter DINOv2
|
| 204 |
+
self.ncc_matcher.scales = _complex_scales
|
| 205 |
+
self.ncc_matcher.angles = [80, 85, 90, 95, 100]
|
| 206 |
+
self.dino_verifier.cosine_threshold = 0.83
|
| 207 |
+
cands_3b = self.ncc_matcher.match(drawing_proc, pattern_proc)
|
| 208 |
+
verified_3b = self.dino_verifier.verify_candidates(
|
| 209 |
+
drawing_proc, pattern_proc, cands_3b, derotate=True
|
| 210 |
+
) if cands_3b else []
|
| 211 |
+
print(f"[Pipeline] Pass 3b (micro 90°): {len(cands_3b)} cands -> {len(verified_3b)} verified")
|
| 212 |
+
|
| 213 |
+
verified = verified + verified_3a + verified_3b
|
| 214 |
+
|
| 215 |
+
self.ncc_matcher.scales = _saved_scales
|
| 216 |
+
self.ncc_matcher.ncc_threshold = _saved_ncc
|
| 217 |
+
self.ncc_matcher.angles = _saved_angles
|
| 218 |
+
self.dino_verifier.cosine_threshold = _saved_dino
|
| 219 |
+
t3 = time.time()
|
| 220 |
+
print(f"[Pipeline] All passes: {t3 - t2:.2f}s — {len(verified)} total verified")
|
| 221 |
+
|
| 222 |
+
# Simple-template post-filters: isolation + aspect-ratio + neighborhood.
|
| 223 |
+
# Isolation: circuit components sit in white space; BOM/title-block cells have
|
| 224 |
+
# solid grid lines directly adjacent to their long sides.
|
| 225 |
+
# Aspect ratio: keep candidates whose bbox AR is within 2× of the template AR
|
| 226 |
+
# *or* its reciprocal — the reciprocal check allows 90°-rotated components
|
| 227 |
+
# (e.g. vertical resistors) whose bbox AR is ~1/_tmpl_ar.
|
| 228 |
+
# Neighborhood complexity: reject candidates whose surrounding ring has too
|
| 229 |
+
# many Canny edges — this eliminates false positives inside complex symbols
|
| 230 |
+
# (e.g. bridge-rectifier bodies) which have many adjacent edges.
|
| 231 |
+
# Notes-area exemption: the bottom 20% of the drawing contains notes/legend
|
| 232 |
+
# symbols which have annotation text nearby — skip isolation and neighborhood
|
| 233 |
+
# checks for those so legitimate legend symbols are not filtered out.
|
| 234 |
+
# Top-margin exclusion: discard detections whose top edge is within the outer
|
| 235 |
+
# coordinate-margin strip (border grid cells look like plain rectangles).
|
| 236 |
+
if _is_simple and verified:
|
| 237 |
+
before = len(verified)
|
| 238 |
+
_drw_h, _drw_w = drawing_proc.shape[:2]
|
| 239 |
+
_notes_y = int(_drw_h * 0.80) # below this → notes/legend area
|
| 240 |
+
_top_margin = max(30, int(_drw_h * 0.04)) # top coordinate strip
|
| 241 |
+
|
| 242 |
+
# Reject border-grid cells in the top margin
|
| 243 |
+
verified = [c for c in verified if c["y"] >= _top_margin]
|
| 244 |
+
|
| 245 |
+
# Split into circuit area and notes/legend area
|
| 246 |
+
_circuit = [c for c in verified if c["y"] < _notes_y]
|
| 247 |
+
_notes = [c for c in verified if c["y"] >= _notes_y]
|
| 248 |
+
|
| 249 |
+
# Isolation: reject candidates adjacent to grid lines (circuit area only)
|
| 250 |
+
_circuit = self.postprocessor.filter_isolated(_circuit, drawing_proc)
|
| 251 |
+
|
| 252 |
+
# Aspect ratio: accept both normal (horizontal) and 90°-rotated (vertical) AR
|
| 253 |
+
_ar_inv = 1.0 / max(0.01, _tmpl_ar)
|
| 254 |
+
def _ar_ok(c):
|
| 255 |
+
ar = c["w"] / max(1, c["h"])
|
| 256 |
+
return (
|
| 257 |
+
(_tmpl_ar / 2.0 <= ar <= _tmpl_ar * 2.0)
|
| 258 |
+
or (_ar_inv / 2.0 <= ar <= _ar_inv * 2.0)
|
| 259 |
+
)
|
| 260 |
+
_circuit = [c for c in _circuit if _ar_ok(c)]
|
| 261 |
+
_notes = [c for c in _notes if _ar_ok(c)]
|
| 262 |
+
|
| 263 |
+
# Wire-lead check: real resistors have straight wire leads on both
|
| 264 |
+
# connecting sides; false positives embedded in complex symbols do not.
|
| 265 |
+
# Circuit area: require one strong lead (>=6px) + one non-zero lead (>=1px)
|
| 266 |
+
# — handles resistors connected directly to a power rail where only
|
| 267 |
+
# 1-2px of wire is visible on the rail-side.
|
| 268 |
+
# Notes/legend area: skip wire-lead filter (legend symbols may have no
|
| 269 |
+
# protruding leads). Restrict to left half of drawing (legends are always
|
| 270 |
+
# on the left; BOM/title-block cells are on the right). Keep top-2 by score.
|
| 271 |
+
_circuit = self.postprocessor.filter_wire_leads(_circuit, drawing_proc)
|
| 272 |
+
_circuit = self.postprocessor.filter_wire_passthrough(_circuit, drawing_proc)
|
| 273 |
+
_circuit = self.postprocessor.filter_neighborhood_complexity(
|
| 274 |
+
_circuit, drawing_proc, expand_ratio=0.5, max_edge_density=0.022
|
| 275 |
+
)
|
| 276 |
+
_circuit = self.postprocessor.filter_junction_dots(_circuit, drawing_proc)
|
| 277 |
+
_circuit = self.postprocessor.filter_rect_integrity(_circuit, drawing_proc)
|
| 278 |
+
# DINOv2-verified borderline candidates (dino ≥ 0.86) bypass the
|
| 279 |
+
# final Chamfer check — they've already been semantically confirmed.
|
| 280 |
+
_dino_ok = [c for c in _circuit if c.get("dino_score", 0) >= 0.88]
|
| 281 |
+
_needs_ch = [c for c in _circuit if c.get("dino_score", 0) < 0.86]
|
| 282 |
+
_needs_ch = self.postprocessor.filter_chamfer_shape(_needs_ch, drawing_proc, pattern_proc)
|
| 283 |
+
_circuit = _needs_ch + _dino_ok
|
| 284 |
+
_bottom_margin = max(30, int(_drw_h * 0.04))
|
| 285 |
+
_notes = [c for c in _notes
|
| 286 |
+
if c["y"] + c["h"] <= _drw_h - _bottom_margin]
|
| 287 |
+
_notes = self.postprocessor.filter_rect_borders(_notes, drawing_proc)
|
| 288 |
+
_notes = sorted(_notes, key=lambda c: c.get("dino_score", 0), reverse=True)[:1]
|
| 289 |
+
|
| 290 |
+
verified = _circuit + _notes
|
| 291 |
+
print(
|
| 292 |
+
f"[Pipeline] Simple-template filters: {before} -> {len(verified)} "
|
| 293 |
+
f"(top-margin + isolation + AR + wire-leads + passthrough | notes_kept={len(_notes)})"
|
| 294 |
+
)
|
| 295 |
+
|
| 296 |
+
# Title-block zone filter: remove candidates inside the BOM / right-frame
|
| 297 |
+
# area for all templates (complex templates skip the simple-filter block
|
| 298 |
+
# that previously applied this only to simple micro-pass candidates).
|
| 299 |
+
before_tb = len(verified)
|
| 300 |
+
verified = self.postprocessor.filter_title_block(verified, drawing_proc)
|
| 301 |
+
if len(verified) != before_tb:
|
| 302 |
+
print(f"[Pipeline] Title-block filter: {before_tb} -> {len(verified)}")
|
| 303 |
+
|
| 304 |
+
# Final NMS + format
|
| 305 |
+
# Simple templates: use tight IoU (0.25) to merge offset-position duplicates
|
| 306 |
+
# of the same component, and keep the best-fit bbox (not union) to avoid
|
| 307 |
+
# over-expanding the box when a lower-confidence offset duplicate gets merged in.
|
| 308 |
+
# Complex templates: use default IoU (0.40) with union bbox expansion.
|
| 309 |
+
if _is_simple:
|
| 310 |
+
verified = self.postprocessor.final_nms(
|
| 311 |
+
verified, iou_threshold=0.25, use_union_bbox=False
|
| 312 |
+
)
|
| 313 |
+
else:
|
| 314 |
+
verified = self.postprocessor.final_nms(
|
| 315 |
+
verified, iou_threshold=self.final_nms_iou
|
| 316 |
+
)
|
| 317 |
+
result = self.postprocessor.format_output(verified, drawing_proc.shape)
|
| 318 |
+
t4 = time.time()
|
| 319 |
+
print(f"[Pipeline] Auto-detect total: {t4 - t0:.2f}s — {result['total_detections']} detections")
|
| 320 |
+
|
| 321 |
+
if return_visualization:
|
| 322 |
+
# Draw on original (non-binarized) image for clearer output
|
| 323 |
+
result["visualization"] = self.postprocessor.draw_boxes(
|
| 324 |
+
drawing_data["original"], result["detections"]
|
| 325 |
+
)
|
| 326 |
+
|
| 327 |
+
return result
|
| 328 |
+
|
| 329 |
+
except Exception as e:
|
| 330 |
+
raise RuntimeError(f"Pipeline auto-detect failed: {e}") from e
|
| 331 |
+
|
| 332 |
+
def detect(
|
| 333 |
+
self,
|
| 334 |
+
pattern_input: Union[str, np.ndarray],
|
| 335 |
+
drawing_input: Union[str, np.ndarray],
|
| 336 |
+
return_visualization: bool = True,
|
| 337 |
+
) -> dict:
|
| 338 |
+
"""Run full detection pipeline.
|
| 339 |
+
|
| 340 |
+
Args:
|
| 341 |
+
pattern_input: Pattern image path or numpy array.
|
| 342 |
+
drawing_input: Drawing image path or numpy array.
|
| 343 |
+
return_visualization: Whether to include annotated image in output.
|
| 344 |
+
|
| 345 |
+
Returns:
|
| 346 |
+
Dict from Postprocessor.format_output(), plus optional "visualization" key.
|
| 347 |
+
|
| 348 |
+
Raises:
|
| 349 |
+
RuntimeError: If any stage fails.
|
| 350 |
+
"""
|
| 351 |
+
try:
|
| 352 |
+
t0 = time.time()
|
| 353 |
+
|
| 354 |
+
# Stage 0: Preprocess
|
| 355 |
+
pattern_data = self.preprocessor.preprocess(pattern_input)
|
| 356 |
+
drawing_data = self.preprocessor.preprocess(drawing_input)
|
| 357 |
+
t1 = time.time()
|
| 358 |
+
print(f"[Pipeline] Stage 0 (Preprocess): {t1 - t0:.2f}s")
|
| 359 |
+
|
| 360 |
+
pattern_proc = pattern_data["processed"] # original — used for DINOv2
|
| 361 |
+
drawing_proc = drawing_data["processed"]
|
| 362 |
+
|
| 363 |
+
# Optionally dilate pattern strokes for NCC to handle style mismatch
|
| 364 |
+
if self.dilate_pattern > 0:
|
| 365 |
+
pattern_for_ncc = self.preprocessor.dilate_strokes(
|
| 366 |
+
pattern_proc, kernel_size=self.dilate_pattern
|
| 367 |
+
)
|
| 368 |
+
else:
|
| 369 |
+
pattern_for_ncc = pattern_proc
|
| 370 |
+
|
| 371 |
+
# Stage 1: NCC matching (uses dilated pattern if configured)
|
| 372 |
+
candidates = self.ncc_matcher.match(drawing_proc, pattern_for_ncc)
|
| 373 |
+
t2 = time.time()
|
| 374 |
+
print(f"[Pipeline] Stage 1 (NCC): {t2 - t1:.2f}s — {len(candidates)} candidates")
|
| 375 |
+
|
| 376 |
+
if not candidates:
|
| 377 |
+
print("[Pipeline] No candidates from Stage 1, returning empty result.")
|
| 378 |
+
result = self.postprocessor.format_output([], drawing_proc.shape)
|
| 379 |
+
if return_visualization:
|
| 380 |
+
result["visualization"] = self.postprocessor.draw_boxes(
|
| 381 |
+
drawing_data["original"], []
|
| 382 |
+
)
|
| 383 |
+
return result
|
| 384 |
+
|
| 385 |
+
# Stage 2: DINOv2 verification
|
| 386 |
+
candidates = self.dino_verifier.verify_candidates(drawing_proc, pattern_proc, candidates)
|
| 387 |
+
t3 = time.time()
|
| 388 |
+
print(f"[Pipeline] Stage 2 (DINOv2): {t3 - t2:.2f}s — {len(candidates)} verified")
|
| 389 |
+
|
| 390 |
+
# Stage 3: Final NMS + format
|
| 391 |
+
candidates = self.postprocessor.final_nms(candidates, iou_threshold=self.final_nms_iou)
|
| 392 |
+
result = self.postprocessor.format_output(candidates, drawing_proc.shape)
|
| 393 |
+
t4 = time.time()
|
| 394 |
+
print(f"[Pipeline] Stage 3 (Post): {t4 - t3:.2f}s")
|
| 395 |
+
print(f"[Pipeline] Total: {t4 - t0:.2f}s — {result['total_detections']} detections")
|
| 396 |
+
|
| 397 |
+
if return_visualization:
|
| 398 |
+
# Draw on original (non-binarized) image for clearer output
|
| 399 |
+
result["visualization"] = self.postprocessor.draw_boxes(
|
| 400 |
+
drawing_data["original"], result["detections"]
|
| 401 |
+
)
|
| 402 |
+
|
| 403 |
+
return result
|
| 404 |
+
|
| 405 |
+
except Exception as e:
|
| 406 |
+
raise RuntimeError(f"Pipeline detection failed: {e}") from e
|
| 407 |
+
|
| 408 |
+
def update_thresholds(
|
| 409 |
+
self,
|
| 410 |
+
ncc_threshold: Optional[float] = None,
|
| 411 |
+
cosine_threshold: Optional[float] = None,
|
| 412 |
+
final_nms_iou: Optional[float] = None,
|
| 413 |
+
):
|
| 414 |
+
"""Update detection thresholds at runtime (for UI sliders).
|
| 415 |
+
|
| 416 |
+
Args:
|
| 417 |
+
ncc_threshold: New NCC threshold for Stage 1.
|
| 418 |
+
cosine_threshold: New cosine similarity threshold for Stage 2.
|
| 419 |
+
"""
|
| 420 |
+
if ncc_threshold is not None:
|
| 421 |
+
self.ncc_matcher.ncc_threshold = ncc_threshold
|
| 422 |
+
if cosine_threshold is not None:
|
| 423 |
+
self.dino_verifier.cosine_threshold = cosine_threshold
|
| 424 |
+
if final_nms_iou is not None:
|
| 425 |
+
self.final_nms_iou = final_nms_iou
|
| 426 |
+
|
| 427 |
+
|
| 428 |
+
def run_detection(pattern_path: str, drawing_path: str, auto: bool = True, **kwargs) -> dict:
|
| 429 |
+
"""Convenience function: create pipeline, run detection, return result.
|
| 430 |
+
|
| 431 |
+
Args:
|
| 432 |
+
pattern_path: Path to pattern image.
|
| 433 |
+
drawing_path: Path to drawing image.
|
| 434 |
+
auto: If True (default), use detect_auto() which self-tunes thresholds.
|
| 435 |
+
**kwargs: Config overrides passed to PatternDetectionPipeline.
|
| 436 |
+
|
| 437 |
+
Returns:
|
| 438 |
+
Detection result dict.
|
| 439 |
+
"""
|
| 440 |
+
pipeline = PatternDetectionPipeline(config=kwargs if kwargs else None)
|
| 441 |
+
if auto:
|
| 442 |
+
return pipeline.detect_auto(pattern_path, drawing_path)
|
| 443 |
+
return pipeline.detect(pattern_path, drawing_path)
|
src/postprocessor.py
ADDED
|
@@ -0,0 +1,976 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import cv2
|
| 2 |
+
import numpy as np
|
| 3 |
+
from typing import List, Tuple, Callable
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class Postprocessor:
|
| 7 |
+
"""Stage 4: Final NMS, output formatting, and visualization."""
|
| 8 |
+
|
| 9 |
+
def final_nms(
|
| 10 |
+
self,
|
| 11 |
+
candidates: List[dict],
|
| 12 |
+
iou_threshold: float = 0.4,
|
| 13 |
+
use_union_bbox: bool = True,
|
| 14 |
+
) -> List[dict]:
|
| 15 |
+
"""Cluster-and-merge NMS: group overlapping boxes and keep one per cluster.
|
| 16 |
+
|
| 17 |
+
When use_union_bbox=True (default): the output bbox is the union of all
|
| 18 |
+
boxes in the cluster — good for complex templates where multi-scale
|
| 19 |
+
detections should together define the component boundary.
|
| 20 |
+
When use_union_bbox=False: the output bbox is taken from the highest-
|
| 21 |
+
confidence candidate — better for simple templates where an offset
|
| 22 |
+
duplicate should not expand the bbox beyond the best-fit box.
|
| 23 |
+
|
| 24 |
+
Args:
|
| 25 |
+
candidates: Candidate dicts with "confidence" key.
|
| 26 |
+
iou_threshold: Overlap threshold for grouping (max of IoU and containment).
|
| 27 |
+
use_union_bbox: Whether to expand bbox to union of cluster.
|
| 28 |
+
|
| 29 |
+
Returns:
|
| 30 |
+
List of merged detections, one per cluster.
|
| 31 |
+
"""
|
| 32 |
+
if not candidates:
|
| 33 |
+
return []
|
| 34 |
+
|
| 35 |
+
candidates = sorted(candidates, key=lambda c: c["confidence"], reverse=True)
|
| 36 |
+
n = len(candidates)
|
| 37 |
+
cluster_id = list(range(n))
|
| 38 |
+
|
| 39 |
+
def find(i):
|
| 40 |
+
while cluster_id[i] != i:
|
| 41 |
+
cluster_id[i] = cluster_id[cluster_id[i]]
|
| 42 |
+
i = cluster_id[i]
|
| 43 |
+
return i
|
| 44 |
+
|
| 45 |
+
def union(i, j):
|
| 46 |
+
cluster_id[find(i)] = find(j)
|
| 47 |
+
|
| 48 |
+
for i in range(n):
|
| 49 |
+
for j in range(i + 1, n):
|
| 50 |
+
if self._overlap_ratio(candidates[i], candidates[j]) > iou_threshold:
|
| 51 |
+
union(i, j)
|
| 52 |
+
|
| 53 |
+
clusters: dict = {}
|
| 54 |
+
for i, cand in enumerate(candidates):
|
| 55 |
+
root = find(i)
|
| 56 |
+
clusters.setdefault(root, []).append(cand)
|
| 57 |
+
|
| 58 |
+
merged = []
|
| 59 |
+
for cluster in clusters.values():
|
| 60 |
+
best = max(cluster, key=lambda c: c["confidence"])
|
| 61 |
+
merged_cand = dict(best)
|
| 62 |
+
if use_union_bbox:
|
| 63 |
+
merged_cand["x"] = min(c["x"] for c in cluster)
|
| 64 |
+
merged_cand["y"] = min(c["y"] for c in cluster)
|
| 65 |
+
merged_cand["w"] = max(c["x"] + c["w"] for c in cluster) - merged_cand["x"]
|
| 66 |
+
merged_cand["h"] = max(c["y"] + c["h"] for c in cluster) - merged_cand["y"]
|
| 67 |
+
merged.append(merged_cand)
|
| 68 |
+
|
| 69 |
+
return sorted(merged, key=lambda c: c["confidence"], reverse=True)
|
| 70 |
+
|
| 71 |
+
def format_output(self, candidates: List[dict], image_shape: Tuple) -> dict:
|
| 72 |
+
"""Format final detections into structured output dict.
|
| 73 |
+
|
| 74 |
+
Args:
|
| 75 |
+
candidates: Filtered detection candidates.
|
| 76 |
+
image_shape: Shape tuple (H, W) or (H, W, C) of the drawing image.
|
| 77 |
+
|
| 78 |
+
Returns:
|
| 79 |
+
Structured detection dict.
|
| 80 |
+
"""
|
| 81 |
+
h, w = image_shape[:2]
|
| 82 |
+
detections = []
|
| 83 |
+
for cand in candidates:
|
| 84 |
+
bw, bh = int(cand["w"]), int(cand["h"])
|
| 85 |
+
if bw < 20 or bh < 20:
|
| 86 |
+
continue
|
| 87 |
+
# Skip degenerate boxes: aspect ratio outside [0.33, 3.0]
|
| 88 |
+
aspect = bw / bh if bh > 0 else 0
|
| 89 |
+
if aspect < 0.33 or aspect > 3.0:
|
| 90 |
+
continue
|
| 91 |
+
detections.append({
|
| 92 |
+
"bbox": {
|
| 93 |
+
"x": int(cand["x"]),
|
| 94 |
+
"y": int(cand["y"]),
|
| 95 |
+
"w": bw,
|
| 96 |
+
"h": bh,
|
| 97 |
+
},
|
| 98 |
+
"confidence": round(float(cand.get("confidence", 0.0)), 2),
|
| 99 |
+
"ncc_score": round(float(cand.get("ncc_score", 0.0)), 4),
|
| 100 |
+
"dino_score": round(float(cand.get("dino_score", 0.0)), 4),
|
| 101 |
+
"scale": round(float(cand.get("scale", 1.0)), 4),
|
| 102 |
+
"angle": round(float(cand.get("angle", 0.0)), 4),
|
| 103 |
+
})
|
| 104 |
+
|
| 105 |
+
return {
|
| 106 |
+
"detections": detections,
|
| 107 |
+
"total_detections": len(detections),
|
| 108 |
+
"image_size": {"width": int(w), "height": int(h)},
|
| 109 |
+
}
|
| 110 |
+
|
| 111 |
+
def filter_grid_clusters(
|
| 112 |
+
self,
|
| 113 |
+
candidates: List[dict],
|
| 114 |
+
row_tol: int = 15,
|
| 115 |
+
col_tol: int = 15,
|
| 116 |
+
min_grid_size: int = 3,
|
| 117 |
+
) -> List[dict]:
|
| 118 |
+
"""Remove candidates that belong to a table row or frame line.
|
| 119 |
+
|
| 120 |
+
When min_grid_size or more detections share the same row (y) or column (x)
|
| 121 |
+
within tolerance, they are structural artifacts (BOM rows, border segments)
|
| 122 |
+
rather than isolated circuit components.
|
| 123 |
+
"""
|
| 124 |
+
if len(candidates) < min_grid_size:
|
| 125 |
+
return candidates
|
| 126 |
+
|
| 127 |
+
n = len(candidates)
|
| 128 |
+
cy = [c["y"] + c["h"] // 2 for c in candidates]
|
| 129 |
+
cx = [c["x"] + c["w"] // 2 for c in candidates]
|
| 130 |
+
keep = [True] * n
|
| 131 |
+
|
| 132 |
+
for i in range(n):
|
| 133 |
+
row_count = sum(1 for j in range(n) if abs(cy[j] - cy[i]) <= row_tol)
|
| 134 |
+
if row_count >= min_grid_size:
|
| 135 |
+
keep[i] = False
|
| 136 |
+
continue
|
| 137 |
+
col_count = sum(1 for j in range(n) if abs(cx[j] - cx[i]) <= col_tol)
|
| 138 |
+
if col_count >= min_grid_size:
|
| 139 |
+
keep[i] = False
|
| 140 |
+
|
| 141 |
+
return [c for c, k in zip(candidates, keep) if k]
|
| 142 |
+
|
| 143 |
+
def find_table_exclusion_zones(self, drawing: np.ndarray) -> dict:
|
| 144 |
+
"""Detect frame border and BOM/title-block area boundaries.
|
| 145 |
+
|
| 146 |
+
Returns:
|
| 147 |
+
dict with:
|
| 148 |
+
title_block_x: leftmost x of the right-frame vertical lines
|
| 149 |
+
(candidates whose right edge exceeds this are frame FPs)
|
| 150 |
+
bom_start_x: leftmost x where horizontal-line density jumps to
|
| 151 |
+
table-area levels (candidates with center beyond this
|
| 152 |
+
are BOM/title-block FPs)
|
| 153 |
+
"""
|
| 154 |
+
H, W = drawing.shape[:2]
|
| 155 |
+
binary = (drawing < 128).astype(np.uint8)
|
| 156 |
+
|
| 157 |
+
# 1. Frame border: find rightmost long vertical line (>=40% of height)
|
| 158 |
+
v_min_len = max(50, int(H * 0.40))
|
| 159 |
+
v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_min_len))
|
| 160 |
+
v_line_img = cv2.erode(binary, v_kernel)
|
| 161 |
+
cols_with_vlines = np.where(v_line_img.any(axis=0))[0]
|
| 162 |
+
title_block_x = W
|
| 163 |
+
if len(cols_with_vlines) > 0:
|
| 164 |
+
right_cols = cols_with_vlines[cols_with_vlines > int(W * 0.75)]
|
| 165 |
+
if len(right_cols) > 0:
|
| 166 |
+
title_block_x = int(right_cols.min())
|
| 167 |
+
|
| 168 |
+
# 2. BOM/title-block area: scan column strips for horizontal-line density jump.
|
| 169 |
+
# Circuit wires produce ~10-50 rows with density > 15%; table rows produce 60+.
|
| 170 |
+
# Only run this scan when a real right-frame was detected (title_block_x < W).
|
| 171 |
+
# If no vertical frame lines exist (title_block_x == W), the drawing has no
|
| 172 |
+
# structured BOM table, so skip the scan entirely to avoid false positives in
|
| 173 |
+
# complex circuit areas.
|
| 174 |
+
strip_w = 30
|
| 175 |
+
bom_threshold = 60
|
| 176 |
+
bom_start_x = title_block_x # default to frame border (safe: = W if no frame)
|
| 177 |
+
if title_block_x < W:
|
| 178 |
+
for x0 in range(int(W * 0.80), W - strip_w, strip_w):
|
| 179 |
+
strip = binary[:, x0: x0 + strip_w]
|
| 180 |
+
row_density = strip.mean(axis=1)
|
| 181 |
+
n_lines = int(np.sum(row_density > 0.15))
|
| 182 |
+
if n_lines >= bom_threshold:
|
| 183 |
+
bom_start_x = x0
|
| 184 |
+
break
|
| 185 |
+
|
| 186 |
+
return {"title_block_x": title_block_x, "bom_start_x": bom_start_x}
|
| 187 |
+
|
| 188 |
+
def filter_title_block(
|
| 189 |
+
self, candidates: List[dict], drawing: np.ndarray
|
| 190 |
+
) -> List[dict]:
|
| 191 |
+
"""Remove candidates inside the outer frame border or BOM/title-block zone.
|
| 192 |
+
|
| 193 |
+
Two criteria:
|
| 194 |
+
1. Right-edge: bbox reaches into the outer frame column
|
| 195 |
+
(catches border-corner FPs whose x+w overlaps the frame line).
|
| 196 |
+
2. BOM zone: candidate centre is to the right of the structured
|
| 197 |
+
table area (BOM rows, company info, drawing number cells).
|
| 198 |
+
"""
|
| 199 |
+
zones = self.find_table_exclusion_zones(drawing)
|
| 200 |
+
tx = zones["title_block_x"] # outer right frame x
|
| 201 |
+
bom_x = zones["bom_start_x"] # BOM/title-block left boundary
|
| 202 |
+
margin = 10 # small margin inward from frame line
|
| 203 |
+
|
| 204 |
+
result = []
|
| 205 |
+
for c in candidates:
|
| 206 |
+
right_edge = c["x"] + c["w"]
|
| 207 |
+
center_x = c["x"] + c["w"] // 2
|
| 208 |
+
# Exclude if bbox right edge reaches into the outer frame
|
| 209 |
+
if right_edge > tx - margin:
|
| 210 |
+
continue
|
| 211 |
+
# Exclude if candidate centre falls inside BOM/title-block area
|
| 212 |
+
if center_x >= bom_x:
|
| 213 |
+
continue
|
| 214 |
+
result.append(c)
|
| 215 |
+
return result
|
| 216 |
+
|
| 217 |
+
def filter_isolated(
|
| 218 |
+
self,
|
| 219 |
+
candidates: List[dict],
|
| 220 |
+
drawing: np.ndarray,
|
| 221 |
+
probe: int = 10,
|
| 222 |
+
dark_threshold: float = 0.25,
|
| 223 |
+
) -> List[dict]:
|
| 224 |
+
"""Remove candidates whose long sides are bordered by adjacent grid lines.
|
| 225 |
+
|
| 226 |
+
Circuit components sit in open white space; BOM/title-block cells have
|
| 227 |
+
solid lines directly above and below (or left/right for vertical).
|
| 228 |
+
Rejects any candidate where the strip just outside a long side has
|
| 229 |
+
dark-pixel ratio > dark_threshold.
|
| 230 |
+
"""
|
| 231 |
+
H, W = drawing.shape[:2]
|
| 232 |
+
result = []
|
| 233 |
+
for c in candidates:
|
| 234 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 235 |
+
margin = max(2, min(w, h) // 6)
|
| 236 |
+
|
| 237 |
+
if w >= h: # horizontal → check top and bottom
|
| 238 |
+
strips = [
|
| 239 |
+
drawing[max(0, y - probe) : y, x + margin : x + w - margin],
|
| 240 |
+
drawing[y + h : min(H, y + h + probe), x + margin : x + w - margin],
|
| 241 |
+
]
|
| 242 |
+
else: # vertical → check left and right
|
| 243 |
+
strips = [
|
| 244 |
+
drawing[y + margin : y + h - margin, max(0, x - probe) : x],
|
| 245 |
+
drawing[y + margin : y + h - margin, x + w : min(W, x + w + probe)],
|
| 246 |
+
]
|
| 247 |
+
|
| 248 |
+
isolated = True
|
| 249 |
+
for s in strips:
|
| 250 |
+
if s.size == 0:
|
| 251 |
+
continue
|
| 252 |
+
if float(np.sum(s < 128)) / s.size > dark_threshold:
|
| 253 |
+
isolated = False
|
| 254 |
+
break
|
| 255 |
+
|
| 256 |
+
if isolated:
|
| 257 |
+
result.append(c)
|
| 258 |
+
|
| 259 |
+
return result
|
| 260 |
+
|
| 261 |
+
def draw_boxes(self, drawing: np.ndarray, detections: List[dict]) -> np.ndarray:
|
| 262 |
+
"""Draw color-coded bounding boxes on the drawing image.
|
| 263 |
+
|
| 264 |
+
Boxes are green (conf >= 0.70), amber (>= 0.55), or red (< 0.55).
|
| 265 |
+
Labels show detection index and confidence score.
|
| 266 |
+
|
| 267 |
+
Args:
|
| 268 |
+
drawing: Grayscale, RGB, or RGBA image.
|
| 269 |
+
detections: List of detection dicts from format_output.
|
| 270 |
+
|
| 271 |
+
Returns:
|
| 272 |
+
RGB numpy array (H, W, 3) with annotations.
|
| 273 |
+
"""
|
| 274 |
+
output = drawing.copy()
|
| 275 |
+
|
| 276 |
+
# Normalise to 3-channel BGR
|
| 277 |
+
if len(output.shape) == 2:
|
| 278 |
+
output = cv2.cvtColor(output, cv2.COLOR_GRAY2BGR)
|
| 279 |
+
elif output.shape[2] == 4:
|
| 280 |
+
output = cv2.cvtColor(output, cv2.COLOR_RGBA2BGR)
|
| 281 |
+
else:
|
| 282 |
+
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
|
| 283 |
+
|
| 284 |
+
img_h, img_w = output.shape[:2]
|
| 285 |
+
# Scale thickness and font to image size
|
| 286 |
+
thickness = max(2, int(max(img_h, img_w) / 600))
|
| 287 |
+
font_scale = max(0.4, min(0.75, max(img_h, img_w) / 2000))
|
| 288 |
+
|
| 289 |
+
def _conf_color(conf: float) -> Tuple:
|
| 290 |
+
if conf >= 0.70:
|
| 291 |
+
return (40, 200, 60) # green
|
| 292 |
+
elif conf >= 0.55:
|
| 293 |
+
return (30, 160, 245) # amber-blue
|
| 294 |
+
else:
|
| 295 |
+
return (40, 40, 220) # red
|
| 296 |
+
|
| 297 |
+
for idx, det in enumerate(detections):
|
| 298 |
+
bbox = det["bbox"]
|
| 299 |
+
x, y, w, h = bbox["x"], bbox["y"], bbox["w"], bbox["h"]
|
| 300 |
+
conf = float(det.get("confidence", 0))
|
| 301 |
+
color = _conf_color(conf)
|
| 302 |
+
|
| 303 |
+
cv2.rectangle(output, (x, y), (x + w, y + h), color, thickness=thickness)
|
| 304 |
+
|
| 305 |
+
label = f"#{idx + 1} {conf:.2f}"
|
| 306 |
+
(lw, lh), bl = cv2.getTextSize(
|
| 307 |
+
label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 1
|
| 308 |
+
)
|
| 309 |
+
tag_y = max(y - 4, lh + 6)
|
| 310 |
+
# Filled label background
|
| 311 |
+
cv2.rectangle(
|
| 312 |
+
output,
|
| 313 |
+
(x, tag_y - lh - 4),
|
| 314 |
+
(x + lw + 8, tag_y + bl),
|
| 315 |
+
color, -1,
|
| 316 |
+
)
|
| 317 |
+
cv2.putText(
|
| 318 |
+
output, label,
|
| 319 |
+
(x + 4, tag_y),
|
| 320 |
+
cv2.FONT_HERSHEY_SIMPLEX,
|
| 321 |
+
font_scale, (255, 255, 255), 1, cv2.LINE_AA,
|
| 322 |
+
)
|
| 323 |
+
|
| 324 |
+
return cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
|
| 325 |
+
|
| 326 |
+
def filter_wire_leads(
|
| 327 |
+
self,
|
| 328 |
+
candidates: List[dict],
|
| 329 |
+
drawing: np.ndarray,
|
| 330 |
+
probe: int = 24,
|
| 331 |
+
min_run: int = 2,
|
| 332 |
+
min_run_weak: int = 0,
|
| 333 |
+
dino_bypass_threshold: float = 0.88,
|
| 334 |
+
) -> List[dict]:
|
| 335 |
+
"""Keep only candidates that have wire leads on their connecting sides.
|
| 336 |
+
|
| 337 |
+
Two acceptance paths:
|
| 338 |
+
|
| 339 |
+
1. High-confidence bypass: candidates whose DINOv2 score exceeds
|
| 340 |
+
dino_bypass_threshold are accepted unconditionally — DINOv2 at
|
| 341 |
+
that level is a strong semantic match, making it almost certain
|
| 342 |
+
the region IS the component and not a wire artifact.
|
| 343 |
+
|
| 344 |
+
2. Wire-lead scan: scan ALL rows across the full bbox height (not
|
| 345 |
+
just the centre ±1 row). This finds leads even when the wire
|
| 346 |
+
connects at the top or bottom edge of the bbox rather than the
|
| 347 |
+
centre, which was the root cause of missed detections. Acceptance
|
| 348 |
+
requires at least one side to have a run ≥ min_run.
|
| 349 |
+
"""
|
| 350 |
+
H, W = drawing.shape[:2]
|
| 351 |
+
binary = (drawing < 128).astype(np.uint8)
|
| 352 |
+
|
| 353 |
+
def _max_run(arr: np.ndarray) -> int:
|
| 354 |
+
if not arr.any():
|
| 355 |
+
return 0
|
| 356 |
+
best = cur = 0
|
| 357 |
+
for v in arr:
|
| 358 |
+
if v:
|
| 359 |
+
cur += 1
|
| 360 |
+
best = max(best, cur)
|
| 361 |
+
else:
|
| 362 |
+
cur = 0
|
| 363 |
+
return best
|
| 364 |
+
|
| 365 |
+
result = []
|
| 366 |
+
for c in candidates:
|
| 367 |
+
# Path 1: high-confidence DINOv2 bypass
|
| 368 |
+
if c.get("dino_score", 0.0) >= dino_bypass_threshold:
|
| 369 |
+
result.append(c)
|
| 370 |
+
continue
|
| 371 |
+
|
| 372 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 373 |
+
angle = c.get("angle", 0)
|
| 374 |
+
is_vertical = 70 <= abs(angle) <= 110
|
| 375 |
+
|
| 376 |
+
passed = False
|
| 377 |
+
for shrink_frac in [0.0, 0.12, 0.25, 0.38]:
|
| 378 |
+
sx = int(w * shrink_frac)
|
| 379 |
+
sy = int(h * shrink_frac)
|
| 380 |
+
bx = x + sx
|
| 381 |
+
by = y + sy
|
| 382 |
+
bw = w - 2 * sx
|
| 383 |
+
bh = h - 2 * sy
|
| 384 |
+
if bw < 8 or bh < 4:
|
| 385 |
+
continue
|
| 386 |
+
|
| 387 |
+
left_best = right_best = 0
|
| 388 |
+
if not is_vertical:
|
| 389 |
+
# Scan ALL rows across bbox height (not just centre ±1).
|
| 390 |
+
# Wires can connect at any point along the short edge.
|
| 391 |
+
for row in range(max(0, by), min(H, by + bh)):
|
| 392 |
+
lslice = binary[row, max(0, bx - probe) : bx]
|
| 393 |
+
rslice = binary[row, bx + bw : min(W, bx + bw + probe)]
|
| 394 |
+
left_best = max(left_best, _max_run(lslice))
|
| 395 |
+
right_best = max(right_best, _max_run(rslice))
|
| 396 |
+
else:
|
| 397 |
+
# Scan ALL columns across bbox width.
|
| 398 |
+
for col in range(max(0, bx), min(W, bx + bw)):
|
| 399 |
+
tslice = binary[max(0, by - probe) : by, col]
|
| 400 |
+
bslice = binary[by + bh : min(H, by + bh + probe), col]
|
| 401 |
+
left_best = max(left_best, _max_run(tslice))
|
| 402 |
+
right_best = max(right_best, _max_run(bslice))
|
| 403 |
+
|
| 404 |
+
strong = max(left_best, right_best)
|
| 405 |
+
weak = min(left_best, right_best)
|
| 406 |
+
if strong >= min_run and weak >= min_run_weak:
|
| 407 |
+
passed = True
|
| 408 |
+
break
|
| 409 |
+
|
| 410 |
+
if passed:
|
| 411 |
+
result.append(c)
|
| 412 |
+
|
| 413 |
+
return result
|
| 414 |
+
|
| 415 |
+
def filter_wire_passthrough(
|
| 416 |
+
self,
|
| 417 |
+
candidates: List[dict],
|
| 418 |
+
drawing: np.ndarray,
|
| 419 |
+
passthrough_threshold: float = 0.60,
|
| 420 |
+
) -> List[dict]:
|
| 421 |
+
"""Remove candidates where a straight wire runs through the bbox body.
|
| 422 |
+
|
| 423 |
+
A component body (resistor rectangle) has a white interior — the
|
| 424 |
+
wire enters one terminal, the body contains the symbol, and the wire
|
| 425 |
+
exits the other terminal. A wire segment or T-junction has a
|
| 426 |
+
continuous dark line running from one side straight through to the
|
| 427 |
+
other without any interior gap.
|
| 428 |
+
|
| 429 |
+
For horizontal candidates: if any row in the centre third of the
|
| 430 |
+
bbox height contains a dark run spanning ≥ passthrough_threshold of
|
| 431 |
+
the inner width, the candidate is rejected as a wire pass-through.
|
| 432 |
+
For vertical: same logic on centre columns.
|
| 433 |
+
"""
|
| 434 |
+
H, W = drawing.shape[:2]
|
| 435 |
+
binary = (drawing < 128).astype(np.uint8)
|
| 436 |
+
|
| 437 |
+
def _max_run(arr: np.ndarray) -> int:
|
| 438 |
+
if not arr.any():
|
| 439 |
+
return 0
|
| 440 |
+
best = cur = 0
|
| 441 |
+
for v in arr:
|
| 442 |
+
if v:
|
| 443 |
+
cur += 1
|
| 444 |
+
best = max(best, cur)
|
| 445 |
+
else:
|
| 446 |
+
cur = 0
|
| 447 |
+
return best
|
| 448 |
+
|
| 449 |
+
result = []
|
| 450 |
+
for c in candidates:
|
| 451 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 452 |
+
angle = c.get("angle", 0)
|
| 453 |
+
is_vertical = 70 <= abs(angle) <= 110
|
| 454 |
+
margin = max(2, min(w, h) // 6)
|
| 455 |
+
|
| 456 |
+
passthrough = False
|
| 457 |
+
if not is_vertical:
|
| 458 |
+
col_from = x + margin
|
| 459 |
+
col_to = x + w - margin
|
| 460 |
+
inner_w = col_to - col_from
|
| 461 |
+
row_from = y + h // 3
|
| 462 |
+
row_to = y + 2 * h // 3
|
| 463 |
+
# Check 1: horizontal dark run through the center third of height.
|
| 464 |
+
if inner_w >= 4 and row_from < row_to:
|
| 465 |
+
for row in range(max(0, row_from), min(H, row_to)):
|
| 466 |
+
inner = binary[row, max(0, col_from) : min(W, col_to)]
|
| 467 |
+
if inner.size == 0:
|
| 468 |
+
continue
|
| 469 |
+
if _max_run(inner) / inner.size >= passthrough_threshold:
|
| 470 |
+
passthrough = True
|
| 471 |
+
break
|
| 472 |
+
# Check 2: continuous vertical dark run through the CENTRE half of the
|
| 473 |
+
# inner columns (avoids the box-border vertical lines which fall within
|
| 474 |
+
# margin distance of the bbox edge but outside the centre half).
|
| 475 |
+
# T-junctions have wires running through the body; a plain rectangle has
|
| 476 |
+
# only a white interior, so no column has a long continuous dark run.
|
| 477 |
+
if not passthrough and h > 0:
|
| 478 |
+
v_col_from = x + w // 4
|
| 479 |
+
v_col_to = x + 3 * w // 4
|
| 480 |
+
if v_col_from < v_col_to:
|
| 481 |
+
for col in range(max(0, v_col_from), min(W, v_col_to)):
|
| 482 |
+
col_slice = binary[max(0, y) : min(H, y + h), col]
|
| 483 |
+
if col_slice.size > 0 and _max_run(col_slice) / col_slice.size >= passthrough_threshold:
|
| 484 |
+
passthrough = True
|
| 485 |
+
break
|
| 486 |
+
else:
|
| 487 |
+
row_from = y + margin
|
| 488 |
+
row_to = y + h - margin
|
| 489 |
+
inner_h = row_to - row_from
|
| 490 |
+
col_from = x + w // 3
|
| 491 |
+
col_to = x + 2 * w // 3
|
| 492 |
+
if inner_h >= 4 and col_from < col_to:
|
| 493 |
+
for col in range(max(0, col_from), min(W, col_to)):
|
| 494 |
+
inner = binary[max(0, row_from) : min(H, row_to), col]
|
| 495 |
+
if inner.size == 0:
|
| 496 |
+
continue
|
| 497 |
+
if _max_run(inner) / inner.size >= passthrough_threshold:
|
| 498 |
+
passthrough = True
|
| 499 |
+
break
|
| 500 |
+
|
| 501 |
+
if not passthrough:
|
| 502 |
+
result.append(c)
|
| 503 |
+
|
| 504 |
+
return result
|
| 505 |
+
|
| 506 |
+
def filter_rect_borders(
|
| 507 |
+
self,
|
| 508 |
+
candidates: List[dict],
|
| 509 |
+
drawing: np.ndarray,
|
| 510 |
+
border_run_ratio: float = 0.55,
|
| 511 |
+
) -> List[dict]:
|
| 512 |
+
"""Keep only candidates that show symmetric rectangular borders (top + bottom).
|
| 513 |
+
|
| 514 |
+
Real schematic component symbols (resistors, capacitors) have a rectangular
|
| 515 |
+
outline with a visible horizontal border at ~20-35% and ~65-80% of the bbox
|
| 516 |
+
height. Wire T-junctions and corners have only ONE horizontal line (either
|
| 517 |
+
near the top or bottom, asymmetric), so this filter removes them.
|
| 518 |
+
|
| 519 |
+
Designed for use on notes/legend-area candidates where the wire-lead filter
|
| 520 |
+
is intentionally not applied.
|
| 521 |
+
"""
|
| 522 |
+
H, W = drawing.shape[:2]
|
| 523 |
+
binary = (drawing < 128).astype(np.uint8)
|
| 524 |
+
|
| 525 |
+
def _max_run(arr: np.ndarray) -> int:
|
| 526 |
+
if not arr.any():
|
| 527 |
+
return 0
|
| 528 |
+
best = cur = 0
|
| 529 |
+
for v in arr:
|
| 530 |
+
if v:
|
| 531 |
+
cur += 1
|
| 532 |
+
best = max(best, cur)
|
| 533 |
+
else:
|
| 534 |
+
cur = 0
|
| 535 |
+
return best
|
| 536 |
+
|
| 537 |
+
def _zone_has_border(x, y_start, y_end, w, threshold):
|
| 538 |
+
for row in range(max(0, y_start), min(H, y_end + 1)):
|
| 539 |
+
line = binary[row, max(0, x) : min(W, x + w)]
|
| 540 |
+
if line.size > 0 and _max_run(line) / w >= threshold:
|
| 541 |
+
return True
|
| 542 |
+
return False
|
| 543 |
+
|
| 544 |
+
result = []
|
| 545 |
+
for c in candidates:
|
| 546 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 547 |
+
angle = c.get("angle", 0)
|
| 548 |
+
is_vertical = 70 <= abs(angle) <= 110
|
| 549 |
+
|
| 550 |
+
if not is_vertical:
|
| 551 |
+
top_start = y + int(h * 0.15)
|
| 552 |
+
top_end = y + int(h * 0.35)
|
| 553 |
+
bot_start = y + int(h * 0.65)
|
| 554 |
+
bot_end = y + int(h * 0.85)
|
| 555 |
+
has_top = _zone_has_border(x, top_start, top_end, w, border_run_ratio)
|
| 556 |
+
has_bot = _zone_has_border(x, bot_start, bot_end, w, border_run_ratio)
|
| 557 |
+
if has_top and has_bot:
|
| 558 |
+
result.append(c)
|
| 559 |
+
else:
|
| 560 |
+
left_start = x + int(w * 0.15)
|
| 561 |
+
left_end = x + int(w * 0.35)
|
| 562 |
+
right_start = x + int(w * 0.65)
|
| 563 |
+
right_end = x + int(w * 0.85)
|
| 564 |
+
|
| 565 |
+
def _col_zone_has_border(col_s, col_e, threshold):
|
| 566 |
+
for col in range(max(0, col_s), min(W, col_e + 1)):
|
| 567 |
+
line = binary[max(0, y) : min(H, y + h), col]
|
| 568 |
+
if line.size > 0 and _max_run(line) / h >= threshold:
|
| 569 |
+
return True
|
| 570 |
+
return False
|
| 571 |
+
|
| 572 |
+
has_left = _col_zone_has_border(left_start, left_end, border_run_ratio)
|
| 573 |
+
has_right = _col_zone_has_border(right_start, right_end, border_run_ratio)
|
| 574 |
+
if has_left and has_right:
|
| 575 |
+
result.append(c)
|
| 576 |
+
|
| 577 |
+
return result
|
| 578 |
+
|
| 579 |
+
def filter_junction_dots(
|
| 580 |
+
self,
|
| 581 |
+
candidates: List[dict],
|
| 582 |
+
drawing: np.ndarray,
|
| 583 |
+
bbox_margin: int = 3,
|
| 584 |
+
min_blob_area: int = 15,
|
| 585 |
+
max_blob_ar: float = 2.5,
|
| 586 |
+
min_blob_fill: float = 0.40,
|
| 587 |
+
) -> List[dict]:
|
| 588 |
+
"""Reject candidates that contain a junction dot inside their bounding box.
|
| 589 |
+
|
| 590 |
+
Circuit junction nodes carry a small filled circle at wire crossings.
|
| 591 |
+
When NCC matches a junction-node region, that dot appears as a compact
|
| 592 |
+
dark blob inside the detected bbox. Real component symbols (resistors)
|
| 593 |
+
have only thin-line structure (rectangle outline, wire leads) — no
|
| 594 |
+
compact filled blobs.
|
| 595 |
+
|
| 596 |
+
Detection: find connected components inside the bbox (after skipping the
|
| 597 |
+
border-line strip) and check for any component that is:
|
| 598 |
+
- large enough to be a dot (area ≥ min_blob_area)
|
| 599 |
+
- roughly equidimensional (aspect ratio ≤ max_blob_ar)
|
| 600 |
+
- densely filled (area / bounding-rect ≥ min_blob_fill)
|
| 601 |
+
|
| 602 |
+
Args:
|
| 603 |
+
bbox_margin: Pixels to skip from each edge of the bbox before
|
| 604 |
+
looking for blobs (avoids the component border lines).
|
| 605 |
+
min_blob_area: Minimum connected-component area in pixels.
|
| 606 |
+
max_blob_ar: Maximum width/height ratio (or inverse) of the blob
|
| 607 |
+
bounding rect; keeps roughly circular shapes only.
|
| 608 |
+
min_blob_fill: Minimum fill ratio (area / bounding-rect area);
|
| 609 |
+
rejects elongated or sparse shapes.
|
| 610 |
+
"""
|
| 611 |
+
H, W = drawing.shape[:2]
|
| 612 |
+
binary = (drawing < 128).astype(np.uint8)
|
| 613 |
+
|
| 614 |
+
result = []
|
| 615 |
+
for c in candidates:
|
| 616 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 617 |
+
m = bbox_margin
|
| 618 |
+
y0 = max(0, y + m); y1_c = min(H, y + h - m)
|
| 619 |
+
x0 = max(0, x + m); x1_c = min(W, x + w - m)
|
| 620 |
+
if y1_c <= y0 or x1_c <= x0:
|
| 621 |
+
result.append(c)
|
| 622 |
+
continue
|
| 623 |
+
|
| 624 |
+
crop = binary[y0:y1_c, x0:x1_c]
|
| 625 |
+
n_labels, _labels, stats, _centroids = cv2.connectedComponentsWithStats(
|
| 626 |
+
crop, connectivity=8
|
| 627 |
+
)
|
| 628 |
+
|
| 629 |
+
has_dot = False
|
| 630 |
+
for i in range(1, n_labels):
|
| 631 |
+
area = int(stats[i, cv2.CC_STAT_AREA])
|
| 632 |
+
cw = int(stats[i, cv2.CC_STAT_WIDTH])
|
| 633 |
+
ch = int(stats[i, cv2.CC_STAT_HEIGHT])
|
| 634 |
+
if area < min_blob_area or cw < 1 or ch < 1:
|
| 635 |
+
continue
|
| 636 |
+
ar = cw / ch
|
| 637 |
+
fill = area / (cw * ch)
|
| 638 |
+
if (1.0 / max_blob_ar) <= ar <= max_blob_ar and fill >= min_blob_fill:
|
| 639 |
+
has_dot = True
|
| 640 |
+
break
|
| 641 |
+
|
| 642 |
+
if not has_dot:
|
| 643 |
+
result.append(c)
|
| 644 |
+
|
| 645 |
+
return result
|
| 646 |
+
|
| 647 |
+
def filter_rect_integrity(
|
| 648 |
+
self,
|
| 649 |
+
candidates: List[dict],
|
| 650 |
+
drawing: np.ndarray,
|
| 651 |
+
border_run_ratio: float = 0.50,
|
| 652 |
+
top_zone: tuple = (0.10, 0.40),
|
| 653 |
+
bot_zone: tuple = (0.60, 0.92),
|
| 654 |
+
side_run_min: int = 2,
|
| 655 |
+
dino_bypass_threshold: float = 0.89,
|
| 656 |
+
) -> List[dict]:
|
| 657 |
+
"""Reject candidates that lack the basic rectangular structure of a component.
|
| 658 |
+
|
| 659 |
+
A well-formed component symbol (resistor) has symmetric rectangular borders:
|
| 660 |
+
both a top horizontal line and a bottom horizontal line. Two degenerate
|
| 661 |
+
artefact patterns are rejected:
|
| 662 |
+
|
| 663 |
+
1. *Only-top artifact*: a single horizontal border in the top zone with NO
|
| 664 |
+
matching border in the bottom zone. This catches L-junction FPs whose
|
| 665 |
+
top-line matches the template's top border but whose bottom is absent.
|
| 666 |
+
|
| 667 |
+
2. *Empty-bus artifact*: a single horizontal border in the bottom zone with
|
| 668 |
+
ZERO dark side-line content in the rows above it. A real component
|
| 669 |
+
(even one that is partially cropped by the bbox) still shows its left and
|
| 670 |
+
right vertical side lines (≥ side_run_min dark pixels in a row). A bare
|
| 671 |
+
bus-wire section has nothing above the horizontal line.
|
| 672 |
+
|
| 673 |
+
Candidates with two visible borders, or with one border and visible side
|
| 674 |
+
lines above, are kept. Vertical candidates (rotated ~90°) are skipped.
|
| 675 |
+
|
| 676 |
+
High-confidence DINOv2 bypass: candidates with dino_score ≥
|
| 677 |
+
dino_bypass_threshold pass unconditionally — at that similarity level the
|
| 678 |
+
semantic match overrides the structural heuristic.
|
| 679 |
+
|
| 680 |
+
Args:
|
| 681 |
+
border_run_ratio: Minimum fraction of bbox width for a row to count
|
| 682 |
+
as a "border" (significant horizontal dark run).
|
| 683 |
+
top_zone: (lo, hi) fractions of bbox height for the top zone.
|
| 684 |
+
bot_zone: (lo, hi) fractions of bbox height for the bottom zone.
|
| 685 |
+
side_run_min: Minimum dark-run length in a row above the detected
|
| 686 |
+
bottom border for the component side-lines to be
|
| 687 |
+
considered present.
|
| 688 |
+
dino_bypass_threshold: DINOv2 cosine score above which the structural
|
| 689 |
+
check is skipped entirely.
|
| 690 |
+
"""
|
| 691 |
+
H, W = drawing.shape[:2]
|
| 692 |
+
binary = (drawing < 128).astype(np.uint8)
|
| 693 |
+
|
| 694 |
+
def _max_run(arr: np.ndarray) -> int:
|
| 695 |
+
if not arr.any():
|
| 696 |
+
return 0
|
| 697 |
+
best = cur = 0
|
| 698 |
+
for v in arr:
|
| 699 |
+
if v:
|
| 700 |
+
cur += 1
|
| 701 |
+
best = max(best, cur)
|
| 702 |
+
else:
|
| 703 |
+
cur = 0
|
| 704 |
+
return best
|
| 705 |
+
|
| 706 |
+
result = []
|
| 707 |
+
for c in candidates:
|
| 708 |
+
if c.get("dino_score", 0.0) >= dino_bypass_threshold:
|
| 709 |
+
result.append(c)
|
| 710 |
+
continue
|
| 711 |
+
|
| 712 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 713 |
+
angle = c.get("angle", 0)
|
| 714 |
+
if 70 <= abs(angle) <= 110:
|
| 715 |
+
result.append(c)
|
| 716 |
+
continue
|
| 717 |
+
|
| 718 |
+
top_lo = y + int(h * top_zone[0])
|
| 719 |
+
top_hi = y + int(h * top_zone[1])
|
| 720 |
+
bot_lo = y + int(h * bot_zone[0])
|
| 721 |
+
bot_hi = y + int(h * bot_zone[1])
|
| 722 |
+
|
| 723 |
+
has_top = False
|
| 724 |
+
has_bot = False
|
| 725 |
+
for row in range(max(0, top_lo), min(H, top_hi + 1)):
|
| 726 |
+
line = binary[row, max(0, x) : min(W, x + w)]
|
| 727 |
+
if line.size > 0 and _max_run(line) / w >= border_run_ratio:
|
| 728 |
+
has_top = True
|
| 729 |
+
break
|
| 730 |
+
for row in range(max(0, bot_lo), min(H, bot_hi + 1)):
|
| 731 |
+
line = binary[row, max(0, x) : min(W, x + w)]
|
| 732 |
+
if line.size > 0 and _max_run(line) / w >= border_run_ratio:
|
| 733 |
+
has_bot = True
|
| 734 |
+
break
|
| 735 |
+
|
| 736 |
+
if has_top and has_bot:
|
| 737 |
+
result.append(c)
|
| 738 |
+
continue
|
| 739 |
+
|
| 740 |
+
if has_top and not has_bot:
|
| 741 |
+
# Only-top artifact: top border with no matching bottom → reject
|
| 742 |
+
continue
|
| 743 |
+
|
| 744 |
+
if has_bot and not has_top:
|
| 745 |
+
# Bottom border only — check for side lines above it.
|
| 746 |
+
# Find the topmost row of the bottom border group.
|
| 747 |
+
border_row = bot_lo
|
| 748 |
+
for row in range(max(0, bot_lo), min(H, bot_hi + 1)):
|
| 749 |
+
line = binary[row, max(0, x) : min(W, x + w)]
|
| 750 |
+
if line.size > 0 and _max_run(line) / w >= border_run_ratio:
|
| 751 |
+
border_row = row
|
| 752 |
+
break
|
| 753 |
+
# Check 1: side lines above the bottom border
|
| 754 |
+
has_sides = False
|
| 755 |
+
for row in range(max(0, y), border_row):
|
| 756 |
+
line = binary[row, max(0, x) : min(W, x + w)]
|
| 757 |
+
if _max_run(line) >= side_run_min:
|
| 758 |
+
has_sides = True
|
| 759 |
+
break
|
| 760 |
+
if has_sides:
|
| 761 |
+
result.append(c)
|
| 762 |
+
# else: empty-bus artifact → reject
|
| 763 |
+
continue
|
| 764 |
+
|
| 765 |
+
# No border found in either zone → likely too small or unusual → keep
|
| 766 |
+
result.append(c)
|
| 767 |
+
|
| 768 |
+
return result
|
| 769 |
+
|
| 770 |
+
def filter_chamfer_shape(
|
| 771 |
+
self,
|
| 772 |
+
candidates: List[dict],
|
| 773 |
+
drawing: np.ndarray,
|
| 774 |
+
template: np.ndarray,
|
| 775 |
+
max_chamfer: float = 3.0,
|
| 776 |
+
canny_lo: int = 30,
|
| 777 |
+
canny_hi: int = 100,
|
| 778 |
+
) -> List[dict]:
|
| 779 |
+
"""Reject candidates whose bbox region does not match the template's edge structure.
|
| 780 |
+
|
| 781 |
+
Chamfer distance measures how well the template's edge skeleton aligns with
|
| 782 |
+
the drawing region's edge skeleton. It is specifically suited to binary
|
| 783 |
+
line-art drawings because it does not depend on intensity values — only on
|
| 784 |
+
edge placement.
|
| 785 |
+
|
| 786 |
+
For each candidate:
|
| 787 |
+
1. Resize (and optionally rotate) the template to match the bbox dimensions.
|
| 788 |
+
2. Extract Canny edges from both the template and the drawing region.
|
| 789 |
+
3. Compute the distance transform of the drawing edges (each pixel stores
|
| 790 |
+
its distance to the nearest edge).
|
| 791 |
+
4. Sample the distance transform at every template-edge pixel location.
|
| 792 |
+
5. The mean distance is the Chamfer score. Low score = good structural match.
|
| 793 |
+
|
| 794 |
+
Real components (IEC rectangle) have Chamfer ≈ 0.7–2.0 at the scales used.
|
| 795 |
+
Wire junctions, L-corners, and other FPs yield Chamfer > 5.0 because their
|
| 796 |
+
edge structure is fundamentally different from the full rectangle + lead template.
|
| 797 |
+
|
| 798 |
+
Args:
|
| 799 |
+
max_chamfer: Mean pixel distance threshold. Candidates above this are
|
| 800 |
+
rejected as structural FPs.
|
| 801 |
+
canny_lo/hi: Canny edge thresholds (applied to both template and region).
|
| 802 |
+
"""
|
| 803 |
+
H, W = drawing.shape[:2]
|
| 804 |
+
draw_gray = drawing if drawing.ndim == 2 else cv2.cvtColor(drawing, cv2.COLOR_BGR2GRAY)
|
| 805 |
+
draw_gray = draw_gray.astype(np.uint8)
|
| 806 |
+
tmpl_gray = template if template.ndim == 2 else cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
|
| 807 |
+
tmpl_gray = tmpl_gray.astype(np.uint8)
|
| 808 |
+
|
| 809 |
+
result = []
|
| 810 |
+
for c in candidates:
|
| 811 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 812 |
+
angle = c.get("angle", 0)
|
| 813 |
+
y1 = max(0, y); y2 = min(H, y + h)
|
| 814 |
+
x1 = max(0, x); x2 = min(W, x + w)
|
| 815 |
+
region = draw_gray[y1:y2, x1:x2]
|
| 816 |
+
rh, rw = region.shape[:2]
|
| 817 |
+
if rh < 4 or rw < 4:
|
| 818 |
+
result.append(c)
|
| 819 |
+
continue
|
| 820 |
+
|
| 821 |
+
# Rotate template before resizing so the aspect ratio matches the bbox
|
| 822 |
+
is_vert = 70 <= abs(angle) <= 110
|
| 823 |
+
if is_vert:
|
| 824 |
+
tmpl_s = cv2.resize(
|
| 825 |
+
cv2.rotate(tmpl_gray, cv2.ROTATE_90_CLOCKWISE), (rw, rh)
|
| 826 |
+
)
|
| 827 |
+
else:
|
| 828 |
+
tmpl_s = cv2.resize(tmpl_gray, (rw, rh))
|
| 829 |
+
|
| 830 |
+
tmpl_edges = cv2.Canny(tmpl_s, canny_lo, canny_hi)
|
| 831 |
+
region_edges = cv2.Canny(region, canny_lo, canny_hi)
|
| 832 |
+
|
| 833 |
+
# Bidirectional Chamfer: average of template→region and region→template.
|
| 834 |
+
# Using only template→region inflates the score when the bbox includes
|
| 835 |
+
# wire-lead stubs that have no counterpart in the template (e.g. R5/R9).
|
| 836 |
+
# The symmetric mean is more robust to minor structural asymmetries.
|
| 837 |
+
dt_r = cv2.distanceTransform(
|
| 838 |
+
(255 - region_edges).astype(np.uint8), cv2.DIST_L2, 5
|
| 839 |
+
)
|
| 840 |
+
pts = np.where(tmpl_edges > 0)
|
| 841 |
+
if len(pts[0]) == 0:
|
| 842 |
+
result.append(c)
|
| 843 |
+
continue
|
| 844 |
+
rows = np.clip(pts[0], 0, rh - 1)
|
| 845 |
+
cols = np.clip(pts[1], 0, rw - 1)
|
| 846 |
+
t2r = float(np.mean(dt_r[rows, cols]))
|
| 847 |
+
|
| 848 |
+
dt_t = cv2.distanceTransform(
|
| 849 |
+
(255 - tmpl_edges).astype(np.uint8), cv2.DIST_L2, 5
|
| 850 |
+
)
|
| 851 |
+
pts2 = np.where(region_edges > 0)
|
| 852 |
+
if len(pts2[0]) == 0:
|
| 853 |
+
chamfer_dist = t2r
|
| 854 |
+
else:
|
| 855 |
+
rows2 = np.clip(pts2[0], 0, rh - 1)
|
| 856 |
+
cols2 = np.clip(pts2[1], 0, rw - 1)
|
| 857 |
+
r2t = float(np.mean(dt_t[rows2, cols2]))
|
| 858 |
+
chamfer_dist = (t2r + r2t) / 2
|
| 859 |
+
|
| 860 |
+
c_out = dict(c)
|
| 861 |
+
c_out["chamfer_dist"] = round(chamfer_dist, 3)
|
| 862 |
+
if chamfer_dist <= max_chamfer:
|
| 863 |
+
result.append(c_out)
|
| 864 |
+
|
| 865 |
+
return result
|
| 866 |
+
|
| 867 |
+
def filter_neighborhood_complexity(
|
| 868 |
+
self,
|
| 869 |
+
candidates: List[dict],
|
| 870 |
+
drawing: np.ndarray,
|
| 871 |
+
expand_ratio: float = 1.0,
|
| 872 |
+
max_edge_density: float = 0.05,
|
| 873 |
+
) -> List[dict]:
|
| 874 |
+
"""Remove candidates whose surrounding ring has too many Canny edges.
|
| 875 |
+
|
| 876 |
+
Standalone components (resistors) sit in clean white space; components
|
| 877 |
+
embedded inside complex symbols (bridge rectifiers) have many adjacent
|
| 878 |
+
edges in the ring around the bounding box.
|
| 879 |
+
|
| 880 |
+
Args:
|
| 881 |
+
expand_ratio: Width of the outer ring in units of the bbox dimensions.
|
| 882 |
+
max_edge_density: Canny-edge fraction in the ring above which the
|
| 883 |
+
candidate is rejected.
|
| 884 |
+
"""
|
| 885 |
+
H, W = drawing.shape[:2]
|
| 886 |
+
gray = drawing if drawing.ndim == 2 else cv2.cvtColor(drawing, cv2.COLOR_BGR2GRAY)
|
| 887 |
+
edges = cv2.Canny(gray.astype(np.uint8), 30, 100)
|
| 888 |
+
result = []
|
| 889 |
+
for c in candidates:
|
| 890 |
+
x, y, w, h = c["x"], c["y"], c["w"], c["h"]
|
| 891 |
+
exp_x = max(5, int(w * expand_ratio))
|
| 892 |
+
exp_y = max(5, int(h * expand_ratio))
|
| 893 |
+
x1 = max(0, x - exp_x)
|
| 894 |
+
y1 = max(0, y - exp_y)
|
| 895 |
+
x2 = min(W, x + w + exp_x)
|
| 896 |
+
y2 = min(H, y + h + exp_y)
|
| 897 |
+
|
| 898 |
+
outer = edges[y1:y2, x1:x2].copy()
|
| 899 |
+
# Blank the inner detection box so we measure only the surrounding ring
|
| 900 |
+
inner_y1 = max(0, y - y1)
|
| 901 |
+
inner_x1 = max(0, x - x1)
|
| 902 |
+
inner_y2 = inner_y1 + h
|
| 903 |
+
inner_x2 = inner_x1 + w
|
| 904 |
+
outer[inner_y1:inner_y2, inner_x1:inner_x2] = 0
|
| 905 |
+
|
| 906 |
+
outer_pixels = outer.size - w * h
|
| 907 |
+
if outer_pixels <= 0:
|
| 908 |
+
result.append(c)
|
| 909 |
+
continue
|
| 910 |
+
|
| 911 |
+
edge_density = float(np.count_nonzero(outer)) / outer_pixels
|
| 912 |
+
if edge_density <= max_edge_density:
|
| 913 |
+
result.append(c)
|
| 914 |
+
return result
|
| 915 |
+
|
| 916 |
+
@staticmethod
|
| 917 |
+
def _overlap_ratio(a: dict, b: dict) -> float:
|
| 918 |
+
"""Max of IoU and containment ratio (intersection / area of smaller box).
|
| 919 |
+
Handles multi-scale duplicates: a small box fully inside a large one gets merged.
|
| 920 |
+
"""
|
| 921 |
+
ax1, ay1 = a["x"], a["y"]
|
| 922 |
+
ax2, ay2 = ax1 + a["w"], ay1 + a["h"]
|
| 923 |
+
bx1, by1 = b["x"], b["y"]
|
| 924 |
+
bx2, by2 = bx1 + b["w"], by1 + b["h"]
|
| 925 |
+
|
| 926 |
+
inter_x1 = max(ax1, bx1)
|
| 927 |
+
inter_y1 = max(ay1, by1)
|
| 928 |
+
inter_x2 = min(ax2, bx2)
|
| 929 |
+
inter_y2 = min(ay2, by2)
|
| 930 |
+
|
| 931 |
+
inter_w = max(0, inter_x2 - inter_x1)
|
| 932 |
+
inter_h = max(0, inter_y2 - inter_y1)
|
| 933 |
+
inter_area = inter_w * inter_h
|
| 934 |
+
|
| 935 |
+
if inter_area == 0:
|
| 936 |
+
return 0.0
|
| 937 |
+
|
| 938 |
+
area_a = a["w"] * a["h"]
|
| 939 |
+
area_b = b["w"] * b["h"]
|
| 940 |
+
union_area = area_a + area_b - inter_area
|
| 941 |
+
min_area = min(area_a, area_b)
|
| 942 |
+
|
| 943 |
+
iou = inter_area / union_area if union_area > 0 else 0.0
|
| 944 |
+
containment = inter_area / min_area if min_area > 0 else 0.0
|
| 945 |
+
return max(iou, containment)
|
| 946 |
+
|
| 947 |
+
@staticmethod
|
| 948 |
+
def _is_rgb(img: np.ndarray) -> bool:
|
| 949 |
+
"""Heuristic: check if a 3-channel image is likely RGB (not BGR)."""
|
| 950 |
+
# Cannot determine with certainty; assume caller passes BGR from OpenCV
|
| 951 |
+
return False
|
| 952 |
+
|
| 953 |
+
@staticmethod
|
| 954 |
+
def _iou(a: dict, b: dict) -> float:
|
| 955 |
+
"""IoU between two candidate bbox dicts."""
|
| 956 |
+
ax1, ay1 = a["x"], a["y"]
|
| 957 |
+
ax2, ay2 = ax1 + a["w"], ay1 + a["h"]
|
| 958 |
+
bx1, by1 = b["x"], b["y"]
|
| 959 |
+
bx2, by2 = bx1 + b["w"], by1 + b["h"]
|
| 960 |
+
|
| 961 |
+
inter_x1 = max(ax1, bx1)
|
| 962 |
+
inter_y1 = max(ay1, by1)
|
| 963 |
+
inter_x2 = min(ax2, bx2)
|
| 964 |
+
inter_y2 = min(ay2, by2)
|
| 965 |
+
|
| 966 |
+
inter_w = max(0, inter_x2 - inter_x1)
|
| 967 |
+
inter_h = max(0, inter_y2 - inter_y1)
|
| 968 |
+
inter_area = inter_w * inter_h
|
| 969 |
+
|
| 970 |
+
area_a = a["w"] * a["h"]
|
| 971 |
+
area_b = b["w"] * b["h"]
|
| 972 |
+
union_area = area_a + area_b - inter_area
|
| 973 |
+
|
| 974 |
+
if union_area <= 0:
|
| 975 |
+
return 0.0
|
| 976 |
+
return inter_area / union_area
|
src/preprocessor.py
ADDED
|
@@ -0,0 +1,174 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import cv2
|
| 2 |
+
import numpy as np
|
| 3 |
+
from typing import Union
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class Preprocessor:
|
| 7 |
+
"""Stage 0: Image preprocessing for BOM engineering drawings."""
|
| 8 |
+
|
| 9 |
+
def load_image(self, path: str) -> np.ndarray:
|
| 10 |
+
"""Load image from path and return as grayscale uint8 array.
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
path: File path to image (PNG, JPG, TIFF supported).
|
| 14 |
+
|
| 15 |
+
Returns:
|
| 16 |
+
Grayscale numpy array (uint8, 0-255).
|
| 17 |
+
|
| 18 |
+
Raises:
|
| 19 |
+
ValueError: If image cannot be read.
|
| 20 |
+
"""
|
| 21 |
+
img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
|
| 22 |
+
if img is None:
|
| 23 |
+
raise ValueError(f"Cannot read image: {path}")
|
| 24 |
+
return img
|
| 25 |
+
|
| 26 |
+
def binarize(self, img: np.ndarray, method: str = "adaptive") -> np.ndarray:
|
| 27 |
+
"""Binarize grayscale image.
|
| 28 |
+
|
| 29 |
+
Args:
|
| 30 |
+
img: Grayscale numpy array.
|
| 31 |
+
method: "adaptive" | "otsu" | "none"
|
| 32 |
+
|
| 33 |
+
Returns:
|
| 34 |
+
Binary image with pixel values 0 or 255.
|
| 35 |
+
"""
|
| 36 |
+
if method == "adaptive":
|
| 37 |
+
return cv2.adaptiveThreshold(
|
| 38 |
+
img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
| 39 |
+
cv2.THRESH_BINARY, blockSize=15, C=4
|
| 40 |
+
)
|
| 41 |
+
elif method == "otsu":
|
| 42 |
+
_, binary = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
| 43 |
+
return binary
|
| 44 |
+
elif method == "none":
|
| 45 |
+
if len(img.shape) == 3:
|
| 46 |
+
return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
| 47 |
+
return img
|
| 48 |
+
else:
|
| 49 |
+
raise ValueError(f"Unknown binarize method: {method}. Use 'adaptive', 'otsu', or 'none'.")
|
| 50 |
+
|
| 51 |
+
def denoise(self, img: np.ndarray, kernel_size: int = 2) -> np.ndarray:
|
| 52 |
+
"""Apply morphological closing to fill broken thin lines.
|
| 53 |
+
|
| 54 |
+
Args:
|
| 55 |
+
img: Binary or grayscale image.
|
| 56 |
+
kernel_size: Size of the morphological kernel (keep small to preserve thin lines).
|
| 57 |
+
|
| 58 |
+
Returns:
|
| 59 |
+
Denoised image.
|
| 60 |
+
"""
|
| 61 |
+
kernel = cv2.getStructuringElement(
|
| 62 |
+
cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)
|
| 63 |
+
)
|
| 64 |
+
return cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
|
| 65 |
+
|
| 66 |
+
def resize_if_needed(self, img: np.ndarray, max_dim: int = 4096) -> np.ndarray:
|
| 67 |
+
"""Resize image if its largest dimension exceeds max_dim, preserving aspect ratio.
|
| 68 |
+
|
| 69 |
+
Args:
|
| 70 |
+
img: Input image.
|
| 71 |
+
max_dim: Maximum allowed dimension in pixels.
|
| 72 |
+
|
| 73 |
+
Returns:
|
| 74 |
+
Resized image (or original if already within bounds).
|
| 75 |
+
"""
|
| 76 |
+
h, w = img.shape[:2]
|
| 77 |
+
largest = max(h, w)
|
| 78 |
+
if largest <= max_dim:
|
| 79 |
+
return img
|
| 80 |
+
scale = max_dim / largest
|
| 81 |
+
new_w = int(w * scale)
|
| 82 |
+
new_h = int(h * scale)
|
| 83 |
+
return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
| 84 |
+
|
| 85 |
+
def suppress_text_noise(self, img: np.ndarray, max_area: int = 180) -> np.ndarray:
|
| 86 |
+
"""Remove small isolated strokes (text labels) from a binary image.
|
| 87 |
+
|
| 88 |
+
Uses connected-component analysis to erase blobs that are small and
|
| 89 |
+
elongated — the typical signature of text characters — while preserving
|
| 90 |
+
larger, rounder component symbols.
|
| 91 |
+
|
| 92 |
+
Args:
|
| 93 |
+
img: Binary image (white background, black strokes).
|
| 94 |
+
max_area: Components with pixel area below this threshold AND
|
| 95 |
+
aspect ratio > 1.3 are treated as text and erased.
|
| 96 |
+
|
| 97 |
+
Returns:
|
| 98 |
+
Binary image with text-like blobs replaced by white.
|
| 99 |
+
"""
|
| 100 |
+
inv = cv2.bitwise_not(img)
|
| 101 |
+
n_labels, labels, stats, _ = cv2.connectedComponentsWithStats(inv, connectivity=8)
|
| 102 |
+
result = img.copy()
|
| 103 |
+
for lbl in range(1, n_labels):
|
| 104 |
+
area = int(stats[lbl, cv2.CC_STAT_AREA])
|
| 105 |
+
w = int(stats[lbl, cv2.CC_STAT_WIDTH])
|
| 106 |
+
h = int(stats[lbl, cv2.CC_STAT_HEIGHT])
|
| 107 |
+
if w == 0 or h == 0:
|
| 108 |
+
continue
|
| 109 |
+
aspect = max(w, h) / min(w, h)
|
| 110 |
+
if area < max_area and aspect > 1.3:
|
| 111 |
+
result[labels == lbl] = 255
|
| 112 |
+
return result
|
| 113 |
+
|
| 114 |
+
def dilate_strokes(self, img: np.ndarray, kernel_size: int = 5) -> np.ndarray:
|
| 115 |
+
"""Thicken black strokes in a binary image by eroding (strokes are black on white bg).
|
| 116 |
+
|
| 117 |
+
Useful when template uses thin stylized lines but drawing uses bold strokes.
|
| 118 |
+
|
| 119 |
+
Args:
|
| 120 |
+
img: Binary image (white background, black strokes).
|
| 121 |
+
kernel_size: Size of dilation kernel.
|
| 122 |
+
|
| 123 |
+
Returns:
|
| 124 |
+
Image with thickened strokes.
|
| 125 |
+
"""
|
| 126 |
+
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
|
| 127 |
+
return cv2.erode(img, kernel)
|
| 128 |
+
|
| 129 |
+
def preprocess(
|
| 130 |
+
self,
|
| 131 |
+
img_or_path: Union[str, np.ndarray],
|
| 132 |
+
binarize_method: str = "adaptive",
|
| 133 |
+
denoise: bool = True,
|
| 134 |
+
dilate_strokes: int = 0,
|
| 135 |
+
) -> dict:
|
| 136 |
+
"""Full preprocessing pipeline.
|
| 137 |
+
|
| 138 |
+
Args:
|
| 139 |
+
img_or_path: File path string or numpy array.
|
| 140 |
+
binarize_method: Binarization method passed to self.binarize().
|
| 141 |
+
denoise: Whether to apply morphological denoising.
|
| 142 |
+
|
| 143 |
+
Returns:
|
| 144 |
+
Dict with keys:
|
| 145 |
+
"original": original grayscale numpy array
|
| 146 |
+
"processed": preprocessed numpy array
|
| 147 |
+
"scale_factor": downscale ratio applied (1.0 if no resize)
|
| 148 |
+
"""
|
| 149 |
+
if isinstance(img_or_path, str):
|
| 150 |
+
original = self.load_image(img_or_path)
|
| 151 |
+
elif isinstance(img_or_path, np.ndarray):
|
| 152 |
+
if len(img_or_path.shape) == 3:
|
| 153 |
+
original = cv2.cvtColor(img_or_path, cv2.COLOR_RGB2GRAY)
|
| 154 |
+
else:
|
| 155 |
+
original = img_or_path.copy()
|
| 156 |
+
else:
|
| 157 |
+
raise ValueError("img_or_path must be a file path string or numpy array.")
|
| 158 |
+
|
| 159 |
+
resized = self.resize_if_needed(original)
|
| 160 |
+
h_orig, w_orig = original.shape[:2]
|
| 161 |
+
h_res, w_res = resized.shape[:2]
|
| 162 |
+
scale_factor = h_res / h_orig if h_orig > 0 else 1.0
|
| 163 |
+
|
| 164 |
+
processed = self.binarize(resized, method=binarize_method)
|
| 165 |
+
if denoise:
|
| 166 |
+
processed = self.denoise(processed)
|
| 167 |
+
if dilate_strokes > 0:
|
| 168 |
+
processed = self.dilate_strokes(processed, kernel_size=dilate_strokes)
|
| 169 |
+
|
| 170 |
+
return {
|
| 171 |
+
"original": original,
|
| 172 |
+
"processed": processed,
|
| 173 |
+
"scale_factor": scale_factor,
|
| 174 |
+
}
|
tests/__init__.py
ADDED
|
File without changes
|
tests/test_pipeline.py
ADDED
|
@@ -0,0 +1,203 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import sys
|
| 2 |
+
import types
|
| 3 |
+
import unittest
|
| 4 |
+
import numpy as np
|
| 5 |
+
from unittest.mock import patch, MagicMock
|
| 6 |
+
|
| 7 |
+
# Provide a minimal torch stub so pipeline imports don't fail when torch is absent
|
| 8 |
+
if "torch" not in sys.modules:
|
| 9 |
+
torch_stub = types.ModuleType("torch")
|
| 10 |
+
torch_stub.device = lambda x: x
|
| 11 |
+
torch_stub.no_grad = MagicMock(return_value=MagicMock(__enter__=MagicMock(return_value=None), __exit__=MagicMock(return_value=False)))
|
| 12 |
+
torch_stub.cuda = MagicMock()
|
| 13 |
+
torch_stub.cuda.is_available = MagicMock(return_value=False)
|
| 14 |
+
torch_stub.hub = MagicMock()
|
| 15 |
+
torch_stub.Tensor = MagicMock
|
| 16 |
+
sys.modules["torch"] = torch_stub
|
| 17 |
+
tv_stub = types.ModuleType("torchvision")
|
| 18 |
+
tv_transforms = types.ModuleType("torchvision.transforms")
|
| 19 |
+
tv_transforms.Compose = MagicMock
|
| 20 |
+
tv_transforms.Resize = MagicMock
|
| 21 |
+
tv_transforms.ToTensor = MagicMock
|
| 22 |
+
tv_transforms.Normalize = MagicMock
|
| 23 |
+
tv_stub.transforms = tv_transforms
|
| 24 |
+
sys.modules["torchvision"] = tv_stub
|
| 25 |
+
sys.modules["torchvision.transforms"] = tv_transforms
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
def _make_white_image(h: int, w: int) -> np.ndarray:
|
| 29 |
+
return np.full((h, w), 255, dtype=np.uint8)
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def _make_circle_image(h: int, w: int, radius: int = 30) -> np.ndarray:
|
| 33 |
+
img = _make_white_image(h, w)
|
| 34 |
+
import cv2
|
| 35 |
+
cv2.circle(img, (w // 2, h // 2), radius, 0, 2)
|
| 36 |
+
return img
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
class TestPreprocessor(unittest.TestCase):
|
| 40 |
+
def setUp(self):
|
| 41 |
+
from src.preprocessor import Preprocessor
|
| 42 |
+
self.prep = Preprocessor()
|
| 43 |
+
|
| 44 |
+
def test_load_grayscale(self):
|
| 45 |
+
import tempfile, cv2
|
| 46 |
+
img = np.zeros((100, 100), dtype=np.uint8)
|
| 47 |
+
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
| 48 |
+
cv2.imwrite(f.name, img)
|
| 49 |
+
loaded = self.prep.load_image(f.name)
|
| 50 |
+
self.assertEqual(loaded.shape, (100, 100))
|
| 51 |
+
self.assertEqual(loaded.dtype, np.uint8)
|
| 52 |
+
|
| 53 |
+
def test_binarize_adaptive(self):
|
| 54 |
+
img = np.random.randint(0, 256, (100, 100), dtype=np.uint8)
|
| 55 |
+
binary = self.prep.binarize(img, method="adaptive")
|
| 56 |
+
unique_vals = set(np.unique(binary))
|
| 57 |
+
self.assertTrue(unique_vals.issubset({0, 255}))
|
| 58 |
+
|
| 59 |
+
def test_resize_large_image(self):
|
| 60 |
+
img = np.zeros((4000, 6000), dtype=np.uint8)
|
| 61 |
+
resized = self.prep.resize_if_needed(img, max_dim=4096)
|
| 62 |
+
self.assertLessEqual(max(resized.shape[:2]), 4096)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
class TestNCCMatcher(unittest.TestCase):
|
| 66 |
+
def setUp(self):
|
| 67 |
+
from src.ncc_matcher import NCCMatcher
|
| 68 |
+
self.matcher = NCCMatcher(
|
| 69 |
+
scales=[1.0],
|
| 70 |
+
angles=[0],
|
| 71 |
+
ncc_threshold=0.7,
|
| 72 |
+
nms_iou_threshold=0.3,
|
| 73 |
+
)
|
| 74 |
+
|
| 75 |
+
def test_match_exact_template(self):
|
| 76 |
+
import cv2
|
| 77 |
+
drawing = _make_circle_image(500, 500, radius=40)
|
| 78 |
+
# Crop around the circle as template
|
| 79 |
+
template = drawing[200:300, 200:300]
|
| 80 |
+
candidates = self.matcher.match(drawing, template)
|
| 81 |
+
self.assertGreater(len(candidates), 0)
|
| 82 |
+
|
| 83 |
+
def test_nms_removes_duplicates(self):
|
| 84 |
+
cands = [
|
| 85 |
+
{"x": 10, "y": 10, "w": 50, "h": 50, "ncc_score": 0.9, "scale": 1.0, "angle": 0},
|
| 86 |
+
{"x": 12, "y": 12, "w": 50, "h": 50, "ncc_score": 0.8, "scale": 1.0, "angle": 0},
|
| 87 |
+
{"x": 200, "y": 200, "w": 50, "h": 50, "ncc_score": 0.85, "scale": 1.0, "angle": 0},
|
| 88 |
+
]
|
| 89 |
+
result = self.matcher._apply_nms(cands)
|
| 90 |
+
# First two strongly overlap, should be merged to 1 + the distant one = 2
|
| 91 |
+
self.assertEqual(len(result), 2)
|
| 92 |
+
|
| 93 |
+
def test_no_match_returns_empty(self):
|
| 94 |
+
# A circle pattern on a white drawing will not match a fully white region
|
| 95 |
+
# Use a very high threshold so a completely different pattern returns nothing
|
| 96 |
+
from src.ncc_matcher import NCCMatcher
|
| 97 |
+
import cv2
|
| 98 |
+
strict = NCCMatcher(scales=[1.0], angles=[0], ncc_threshold=0.99)
|
| 99 |
+
# Drawing has a circle on the left; template is a rectangle on the right
|
| 100 |
+
drawing = _make_white_image(300, 300)
|
| 101 |
+
cv2.circle(drawing, (80, 150), 30, 0, 2)
|
| 102 |
+
|
| 103 |
+
# Template: a very different shape (filled black square — not circle)
|
| 104 |
+
template = _make_white_image(60, 60)
|
| 105 |
+
cv2.rectangle(template, (5, 5), (55, 55), 0, -1)
|
| 106 |
+
|
| 107 |
+
result = strict.match(drawing, template)
|
| 108 |
+
self.assertEqual(len(result), 0)
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
class TestPostprocessor(unittest.TestCase):
|
| 112 |
+
def setUp(self):
|
| 113 |
+
from src.postprocessor import Postprocessor
|
| 114 |
+
self.post = Postprocessor()
|
| 115 |
+
|
| 116 |
+
def test_format_output_structure(self):
|
| 117 |
+
candidates = [
|
| 118 |
+
{"x": 10, "y": 20, "w": 50, "h": 60, "confidence": 0.85,
|
| 119 |
+
"ncc_score": 0.8, "dino_score": 0.9, "scale": 1.0, "angle": 0},
|
| 120 |
+
]
|
| 121 |
+
result = self.post.format_output(candidates, (300, 400))
|
| 122 |
+
self.assertIn("detections", result)
|
| 123 |
+
self.assertIn("total_detections", result)
|
| 124 |
+
self.assertIn("image_size", result)
|
| 125 |
+
self.assertEqual(result["total_detections"], 1)
|
| 126 |
+
det = result["detections"][0]
|
| 127 |
+
self.assertIn("bbox", det)
|
| 128 |
+
self.assertIn("confidence", det)
|
| 129 |
+
|
| 130 |
+
def test_draw_boxes_returns_rgb(self):
|
| 131 |
+
drawing = _make_white_image(200, 300)
|
| 132 |
+
detections = [
|
| 133 |
+
{"bbox": {"x": 10, "y": 10, "w": 50, "h": 50}, "confidence": 0.9},
|
| 134 |
+
]
|
| 135 |
+
output = self.post.draw_boxes(drawing, detections)
|
| 136 |
+
self.assertEqual(len(output.shape), 3)
|
| 137 |
+
self.assertEqual(output.shape[2], 3)
|
| 138 |
+
self.assertEqual(output.dtype, np.uint8)
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
class TestPipeline(unittest.TestCase):
|
| 142 |
+
"""Integration tests using synthetic images and mocked DINOv2."""
|
| 143 |
+
|
| 144 |
+
def _get_mock_device(self):
|
| 145 |
+
try:
|
| 146 |
+
import torch
|
| 147 |
+
return torch.device("cpu")
|
| 148 |
+
except ImportError:
|
| 149 |
+
return "cpu"
|
| 150 |
+
|
| 151 |
+
def test_full_pipeline_runs(self):
|
| 152 |
+
# Import pipeline module first so patch can resolve "src.pipeline.DINOVerifier"
|
| 153 |
+
import importlib
|
| 154 |
+
import src.pipeline as pipeline_mod # noqa: ensure module is loaded
|
| 155 |
+
|
| 156 |
+
device = self._get_mock_device()
|
| 157 |
+
|
| 158 |
+
with patch.object(pipeline_mod, "DINOVerifier") as MockDINO:
|
| 159 |
+
instance = MockDINO.return_value
|
| 160 |
+
instance.device = device
|
| 161 |
+
instance.cosine_threshold = 0.75
|
| 162 |
+
instance.verify_candidates.side_effect = lambda d, t, c: [
|
| 163 |
+
dict(x, dino_score=0.85, confidence=(x["ncc_score"] + 0.85) / 2) for x in c
|
| 164 |
+
]
|
| 165 |
+
|
| 166 |
+
from src.pipeline import PatternDetectionPipeline
|
| 167 |
+
pipeline = PatternDetectionPipeline()
|
| 168 |
+
|
| 169 |
+
drawing = _make_circle_image(400, 400, radius=40)
|
| 170 |
+
template = drawing[160:240, 160:240]
|
| 171 |
+
|
| 172 |
+
result = pipeline.detect(template, drawing, return_visualization=False)
|
| 173 |
+
|
| 174 |
+
self.assertIn("detections", result)
|
| 175 |
+
self.assertIn("total_detections", result)
|
| 176 |
+
|
| 177 |
+
def test_empty_result_if_no_match(self):
|
| 178 |
+
import src.pipeline as pipeline_mod # noqa: ensure module is loaded
|
| 179 |
+
|
| 180 |
+
device = self._get_mock_device()
|
| 181 |
+
|
| 182 |
+
with patch.object(pipeline_mod, "DINOVerifier") as MockDINO:
|
| 183 |
+
instance = MockDINO.return_value
|
| 184 |
+
instance.device = device
|
| 185 |
+
instance.cosine_threshold = 0.99
|
| 186 |
+
instance.verify_candidates.return_value = []
|
| 187 |
+
|
| 188 |
+
from src.pipeline import PatternDetectionPipeline
|
| 189 |
+
pipeline = PatternDetectionPipeline(config={"ncc_threshold": 0.99})
|
| 190 |
+
|
| 191 |
+
drawing = _make_white_image(300, 300)
|
| 192 |
+
import cv2
|
| 193 |
+
cv2.rectangle(drawing, (10, 10), (60, 60), 0, 2)
|
| 194 |
+
template = _make_white_image(80, 80)
|
| 195 |
+
cv2.circle(template, (40, 40), 30, 0, 2)
|
| 196 |
+
|
| 197 |
+
result = pipeline.detect(template, drawing, return_visualization=False)
|
| 198 |
+
|
| 199 |
+
self.assertEqual(result["total_detections"], 0)
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
if __name__ == "__main__":
|
| 203 |
+
unittest.main()
|