Spaces:

zestdapoet
/

Pattern_detection_in_technical_drawings

Running

Pattern_detection_in_technical_drawings / app /web /server.py

Duy

feat: add About page with system spec link to web UI

2bdb3d6 5 months ago

6.83 kB

"""FastAPI backend for BOM Pattern Detection web UI."""
import io
import csv
import base64
import time
import sys
import os
import traceback

# Fix OpenMP conflict on Windows/Anaconda
os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")

from pathlib import Path
from typing import Optional

import numpy as np
from PIL import Image
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# Ensure project root is importable
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT))

from src.pipeline import PatternDetectionPipeline

app = FastAPI(title="BOM Pattern Detection API", version="1.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount static files
STATIC_DIR = Path(__file__).parent / "static"
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

# Cached pipeline instance
_pipeline: Optional[PatternDetectionPipeline] = None


def get_pipeline(config: dict = None) -> PatternDetectionPipeline:
    global _pipeline
    if _pipeline is None:
        print("[Server] Loading pipeline...")
        _pipeline = PatternDetectionPipeline(config=config)
        print("[Server] Pipeline ready.")
    return _pipeline


def upload_to_numpy(file_bytes: bytes) -> np.ndarray:
    img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
    return np.array(img)


def numpy_to_b64(img: np.ndarray) -> str:
    if img.ndim == 2:
        pil_img = Image.fromarray(img, mode="L")
    else:
        pil_img = Image.fromarray(img)
    buf = io.BytesIO()
    pil_img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
    return Response(status_code=204)


@app.get("/", response_class=HTMLResponse)
async def index():
    html_path = Path(__file__).parent / "index.html"
    return html_path.read_text(encoding="utf-8")


@app.get("/spec", response_class=HTMLResponse)
async def system_spec():
    spec_path = ROOT / "design_spec" / "system_spec.html"
    return spec_path.read_text(encoding="utf-8")


def _coerce_bool(val) -> bool:
    """Form values arrive as strings; coerce 'true'/'1'/'on' to bool."""
    if isinstance(val, bool):
        return val
    return str(val).strip().lower() in ("true", "1", "on", "yes")


@app.post("/api/detect")
async def detect(
    pattern: UploadFile = File(...),
    drawing: UploadFile = File(...),
    mode: str = Form("auto"),
    ncc_threshold: float = Form(0.55),
    cosine_threshold: float = Form(0.84),
    final_nms_iou: float = Form(0.4),
    use_vlm: str = Form("false"),
):
    try:
        t_start = time.time()

        pattern_bytes = await pattern.read()
        drawing_bytes = await drawing.read()

        pattern_np = upload_to_numpy(pattern_bytes)
        drawing_np = upload_to_numpy(drawing_bytes)

        pipeline = get_pipeline()
        pipeline.update_thresholds(
            ncc_threshold=ncc_threshold,
            cosine_threshold=cosine_threshold,
            final_nms_iou=final_nms_iou,
        )
        # VLM Stage-3 is toggled at runtime (model lazy-loads on first enable).
        pipeline.use_vlm = _coerce_bool(use_vlm)

        if mode == "auto":
            result = pipeline.detect_auto(pattern_np, drawing_np, return_visualization=True)
        else:
            result = pipeline.detect(pattern_np, drawing_np, return_visualization=True)

        elapsed = round(time.time() - t_start, 2)

        viz_b64 = None
        if "visualization" in result and result["visualization"] is not None:
            viz_b64 = numpy_to_b64(result["visualization"])

        # Also return the original drawing as b64 for display
        drawing_b64 = numpy_to_b64(drawing_np)

        return JSONResponse({
            "success": True,
            "total_detections": result["total_detections"],
            "detections": result["detections"],
            "elapsed": elapsed,
            "visualization": viz_b64,
            "drawing_preview": drawing_b64,
        })

    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")


def _detections_to_csv(detections: list) -> str:
    """Convert detection list to CSV string."""
    buf = io.StringIO()
    writer = csv.writer(buf)
    writer.writerow(["id", "x", "y", "width", "height", "confidence", "ncc_score", "dino_score", "scale", "angle"])
    for i, d in enumerate(detections, 1):
        b = d["bbox"]
        writer.writerow([
            i,
            b["x"], b["y"], b["w"], b["h"],
            round(d.get("confidence", 0), 4),
            round(d.get("ncc_score", 0), 4),
            round(d.get("dino_score", 0), 4),
            round(d.get("scale", 1.0), 4),
            round(d.get("angle", 0), 1),
        ])
    return buf.getvalue()


@app.post("/api/detect/csv")
async def detect_csv(
    pattern: UploadFile = File(...),
    drawing: UploadFile = File(...),
    mode: str = Form("auto"),
    ncc_threshold: float = Form(0.55),
    cosine_threshold: float = Form(0.84),
    final_nms_iou: float = Form(0.4),
    use_vlm: str = Form("false"),
):
    """Run detection and return results as a downloadable CSV file."""
    try:
        pattern_bytes = await pattern.read()
        drawing_bytes = await drawing.read()
        pattern_np = upload_to_numpy(pattern_bytes)
        drawing_np = upload_to_numpy(drawing_bytes)

        pipeline = get_pipeline()
        pipeline.update_thresholds(
            ncc_threshold=ncc_threshold,
            cosine_threshold=cosine_threshold,
            final_nms_iou=final_nms_iou,
        )
        pipeline.use_vlm = _coerce_bool(use_vlm)
        if mode == "auto":
            result = pipeline.detect_auto(pattern_np, drawing_np, return_visualization=False)
        else:
            result = pipeline.detect(pattern_np, drawing_np, return_visualization=False)

        csv_str = _detections_to_csv(result["detections"])
        return StreamingResponse(
            io.BytesIO(csv_str.encode("utf-8")),
            media_type="text/csv",
            headers={"Content-Disposition": "attachment; filename=detections.csv"},
        )
    except Exception as e:
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")


@app.get("/api/health")
async def health():
    return {"status": "ok", "pipeline_loaded": _pipeline is not None}


if __name__ == "__main__":
    # Pre-load pipeline before serving
    get_pipeline()
    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")

	"""FastAPI backend for BOM Pattern Detection web UI."""
	import io
	import csv
	import base64
	import time
	import sys
	import os
	import traceback

	# Fix OpenMP conflict on Windows/Anaconda
	os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")

	from pathlib import Path
	from typing import Optional

	import numpy as np
	from PIL import Image
	from fastapi import FastAPI, File, UploadFile, Form, HTTPException
	from fastapi.staticfiles import StaticFiles
	from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse
	from fastapi.middleware.cors import CORSMiddleware
	import uvicorn

	# Ensure project root is importable
	ROOT = Path(__file__).resolve().parents[2]
	sys.path.insert(0, str(ROOT))

	from src.pipeline import PatternDetectionPipeline

	app = FastAPI(title="BOM Pattern Detection API", version="1.0.0")

	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_methods=["*"],
	allow_headers=["*"],
	)

	# Mount static files
	STATIC_DIR = Path(__file__).parent / "static"
	app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")

	# Cached pipeline instance
	_pipeline: Optional[PatternDetectionPipeline] = None


	def get_pipeline(config: dict = None) -> PatternDetectionPipeline:
	global _pipeline
	if _pipeline is None:
	print("[Server] Loading pipeline...")
	_pipeline = PatternDetectionPipeline(config=config)
	print("[Server] Pipeline ready.")
	return _pipeline


	def upload_to_numpy(file_bytes: bytes) -> np.ndarray:
	img = Image.open(io.BytesIO(file_bytes)).convert("RGB")
	return np.array(img)


	def numpy_to_b64(img: np.ndarray) -> str:
	if img.ndim == 2:
	pil_img = Image.fromarray(img, mode="L")
	else:
	pil_img = Image.fromarray(img)
	buf = io.BytesIO()
	pil_img.save(buf, format="PNG")
	return base64.b64encode(buf.getvalue()).decode()


	@app.get("/favicon.ico", include_in_schema=False)
	async def favicon():
	return Response(status_code=204)


	@app.get("/", response_class=HTMLResponse)
	async def index():
	html_path = Path(__file__).parent / "index.html"
	return html_path.read_text(encoding="utf-8")


	@app.get("/spec", response_class=HTMLResponse)
	async def system_spec():
	spec_path = ROOT / "design_spec" / "system_spec.html"
	return spec_path.read_text(encoding="utf-8")


	def _coerce_bool(val) -> bool:
	"""Form values arrive as strings; coerce 'true'/'1'/'on' to bool."""
	if isinstance(val, bool):
	return val
	return str(val).strip().lower() in ("true", "1", "on", "yes")


	@app.post("/api/detect")
	async def detect(
	pattern: UploadFile = File(...),
	drawing: UploadFile = File(...),
	mode: str = Form("auto"),
	ncc_threshold: float = Form(0.55),
	cosine_threshold: float = Form(0.84),
	final_nms_iou: float = Form(0.4),
	use_vlm: str = Form("false"),
	):
	try:
	t_start = time.time()

	pattern_bytes = await pattern.read()
	drawing_bytes = await drawing.read()

	pattern_np = upload_to_numpy(pattern_bytes)
	drawing_np = upload_to_numpy(drawing_bytes)

	pipeline = get_pipeline()
	pipeline.update_thresholds(
	ncc_threshold=ncc_threshold,
	cosine_threshold=cosine_threshold,
	final_nms_iou=final_nms_iou,
	)
	# VLM Stage-3 is toggled at runtime (model lazy-loads on first enable).
	pipeline.use_vlm = _coerce_bool(use_vlm)

	if mode == "auto":
	result = pipeline.detect_auto(pattern_np, drawing_np, return_visualization=True)
	else:
	result = pipeline.detect(pattern_np, drawing_np, return_visualization=True)

	elapsed = round(time.time() - t_start, 2)

	viz_b64 = None
	if "visualization" in result and result["visualization"] is not None:
	viz_b64 = numpy_to_b64(result["visualization"])

	# Also return the original drawing as b64 for display
	drawing_b64 = numpy_to_b64(drawing_np)

	return JSONResponse({
	"success": True,
	"total_detections": result["total_detections"],
	"detections": result["detections"],
	"elapsed": elapsed,
	"visualization": viz_b64,
	"drawing_preview": drawing_b64,
	})

	except Exception as e:
	traceback.print_exc()
	raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")


	def _detections_to_csv(detections: list) -> str:
	"""Convert detection list to CSV string."""
	buf = io.StringIO()
	writer = csv.writer(buf)
	writer.writerow(["id", "x", "y", "width", "height", "confidence", "ncc_score", "dino_score", "scale", "angle"])
	for i, d in enumerate(detections, 1):
	b = d["bbox"]
	writer.writerow([
	i,
	b["x"], b["y"], b["w"], b["h"],
	round(d.get("confidence", 0), 4),
	round(d.get("ncc_score", 0), 4),
	round(d.get("dino_score", 0), 4),
	round(d.get("scale", 1.0), 4),
	round(d.get("angle", 0), 1),
	])
	return buf.getvalue()


	@app.post("/api/detect/csv")
	async def detect_csv(
	pattern: UploadFile = File(...),
	drawing: UploadFile = File(...),
	mode: str = Form("auto"),
	ncc_threshold: float = Form(0.55),
	cosine_threshold: float = Form(0.84),
	final_nms_iou: float = Form(0.4),
	use_vlm: str = Form("false"),
	):
	"""Run detection and return results as a downloadable CSV file."""
	try:
	pattern_bytes = await pattern.read()
	drawing_bytes = await drawing.read()
	pattern_np = upload_to_numpy(pattern_bytes)
	drawing_np = upload_to_numpy(drawing_bytes)

	pipeline = get_pipeline()
	pipeline.update_thresholds(
	ncc_threshold=ncc_threshold,
	cosine_threshold=cosine_threshold,
	final_nms_iou=final_nms_iou,
	)
	pipeline.use_vlm = _coerce_bool(use_vlm)
	if mode == "auto":
	result = pipeline.detect_auto(pattern_np, drawing_np, return_visualization=False)
	else:
	result = pipeline.detect(pattern_np, drawing_np, return_visualization=False)

	csv_str = _detections_to_csv(result["detections"])
	return StreamingResponse(
	io.BytesIO(csv_str.encode("utf-8")),
	media_type="text/csv",
	headers={"Content-Disposition": "attachment; filename=detections.csv"},
	)
	except Exception as e:
	traceback.print_exc()
	raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")


	@app.get("/api/health")
	async def health():
	return {"status": "ok", "pipeline_loaded": _pipeline is not None}


	if __name__ == "__main__":
	# Pre-load pipeline before serving
	get_pipeline()
	uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")