File size: 2,726 Bytes
d0f5702
9a12c0e
d0f5702
85e2f7c
 
 
 
d0f5702
85e2f7c
d0f5702
9a12c0e
 
 
d0f5702
 
 
 
9a12c0e
 
 
 
 
 
 
 
d0f5702
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9a12c0e
85e2f7c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d0f5702
85e2f7c
 
 
 
 
 
 
 
 
d0f5702
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
"""Gradio-SDK entry point (free accounts / ZeroGPU hardware).

ZeroGPU only accepts apps started through Gradio's launch(), so this starts a small Gradio
page and serves the same FastAPI service (serve/main.py) under /v1. Inference runs on the CPU.

Gradio's own middleware answers GET requests it doesn't recognise with its HTML page, so
requests for /v1 are dispatched to FastAPI at the ASGI-server level, before Gradio sees them.
"""
import logging
import os

os.environ.setdefault("MODEL_PATH", "models/model.onnx")

import gradio as gr  # noqa: E402
from starlette.routing import Mount  # noqa: E402

try:  # ZeroGPU expects at least one @spaces.GPU function; it is never called.
    import spaces

    @spaces.GPU
    def _unused_gpu_placeholder():
        return None
except ImportError:
    pass

from serve.main import create_app  # noqa: E402

api = create_app(load_model=True)

INFO = """
# PCB defect detection API

FastAPI + ONNX Runtime (YOLO11n, static INT8, detection head FP32), CPU inference.
On this Space the API is mounted under **/v1**:

- `POST /v1/predict`: multipart field `file`; query `conf`, and `visualize=true` for an annotated PNG
- `GET /v1/health`, `GET /v1/metrics` (Prometheus), `GET /v1/docs` (OpenAPI)

```bash
curl -F "file=@board.jpg;type=image/jpeg" https://michaelandrew35-pcb-defect-api.hf.space/v1/predict
```

Code, tests, CI and benchmarks: https://gitlab.com/michael.sucahyo/pcb-defect-optimization
"""

with gr.Blocks(title="PCB defect detection API") as demo:
    gr.Markdown(INFO)

PREFIX = "/v1"


class PrefixDispatch:
    """Send /v1 requests to the FastAPI service and everything else to Gradio."""

    def __init__(self, fallback, api):
        self.fallback, self.api = fallback, api

    async def __call__(self, scope, receive, send):
        path = scope.get("path", "")
        if scope["type"] in ("http", "websocket") and (path == PREFIX or path.startswith(PREFIX + "/")):
            scope = dict(scope, root_path=scope.get("root_path", "") + PREFIX)
            return await self.api(scope, receive, send)
        return await self.fallback(scope, receive, send)


demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", "7860")), prevent_thread_lock=True)

demo.app.router.routes.insert(0, Mount(PREFIX, app=api))  # fallback for connections opened before the swap
try:  # uvicorn hands config.loaded_app to every new connection
    cfg = demo.server.config
    cfg.loaded_app = PrefixDispatch(cfg.loaded_app, api)
    logging.getLogger("pcb_serve").info("routing %s to FastAPI ahead of Gradio", PREFIX)
except AttributeError as exc:
    logging.getLogger("pcb_serve").warning("could not install /v1 dispatcher: %s", exc)

demo.block_thread()