captchakraken 3.2.0 → 4.0.0

package/python/src/captchakraken/pointer/server.py

@@ -0,0 +1,146 @@
+"""The /pointer HTTP endpoint: one board per request, generation 3's answer back.
+
+    python -m captchakraken.pointer.server --model /models/pointer-0.8b \\
+        --detector /models/detector_r18.pt --port 8001
+
+POST /v1/pointer
+    {"image": "<base64 PNG/JPEG>"}                    a still, as captured
+    {"frames": ["<base64 JPEG>", ...], "fps": 5.0}    a video, `clipvideo`-cut
+    {"marked": false}                                 optional: skip the detector
+    {"avoid": [answer, ...]}                          optional: answers already refused on this
+                                                      board; the likeliest other ones come back
+-> {"actions": [{...}, {...}, {...}],                 the three likeliest distinct answers, best first
+    "answer": {...},                                 the first action, kept for older callers
+    "elements": [{"id", "class", "box", "score"}], "size": [W, H]}
+
+`elements` are in the pixels of the board as the model saw it (`size`), numbered
+as the answer's ids are. The client sends no prompt: the server owns it.
+"""
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+import json
+import os
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Dict
+
+import numpy as np
+from PIL import Image
+
+from . import layout as L
+
+PATH = "/v1/pointer"
+#: Refused answers one request may name: more than a solve's loops ever perform on one board.
+MAX_AVOID = 32
+
+
+def _image(b64: str) -> Image.Image:
+    return Image.open(io.BytesIO(base64.b64decode(b64))).convert("RGB")
+
+
+def parse(body: Dict) -> Dict:
+    """A request body -> PointerService.solve's arguments. Raises ValueError on a bad one."""
+    if ("image" in body) == ("frames" in body):
+        raise ValueError("send exactly one of image or frames")
+    marked = bool(body.get("marked", True))
+    avoid = body.get("avoid") or []
+    if not isinstance(avoid, list) or not all(isinstance(a, dict) for a in avoid):
+        raise ValueError("avoid is a list of earlier answers")
+    if len(avoid) > MAX_AVOID:
+        raise ValueError(f"{len(avoid)} answers to avoid; at most {MAX_AVOID}")
+    history = body.get("history") or []
+    if "image" in body:
+        if len(history) > L.MAX_HISTORY:
+            raise ValueError(f"{len(history)} earlier screens; at most {L.MAX_HISTORY}")
+        return {"image": _image(body["image"]), "marked": marked, "avoid": avoid,
+                "history": [_image(h) for h in history]}
+    if history:
+        raise ValueError("a video board carries no earlier screens")
+    from .. import clipvideo as CV
+    if abs(float(body.get("fps", CV.FPS)) - CV.FPS) > 1e-6:
+        raise ValueError(f"frames must be cut at {CV.FPS} fps by clipvideo, got {body.get('fps')}")
+    frames = [np.asarray(Image.open(io.BytesIO(base64.b64decode(f))).convert("RGB")) for f in body["frames"]]
+    if len(frames) < 2:
+        raise ValueError("a video needs at least two frames")
+    return {"frames": frames, "marked": False, "avoid": avoid}
+
+
+def reply(result: Dict) -> Dict:
+    return {"actions": result.get("answers") or [result["answer"]],
+            "answer": result["answer"], "size": list(result["size"]),
+            "elements": [{"id": n, "class": e["class"], "box": [round(v, 1) for v in e["box"]],
+                          "score": e["score"]} for n, e in enumerate(result["elements"], 1)]}
+
+
+def make_handler(service):
+    lock = threading.Lock()   # one engine; vLLM batches inside a call, not across threads
+
+    class Handler(BaseHTTPRequestHandler):
+        protocol_version = "HTTP/1.1"
+
+        def _send(self, code: int, obj: Dict) -> None:
+            raw = json.dumps(obj).encode()
+            self.send_response(code)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(raw)))
+            self.end_headers()
+            self.wfile.write(raw)
+
+        def _authorized(self) -> bool:
+            key = os.getenv("VLLM_API_KEY", "")
+            return not key or self.headers.get("Authorization") == f"Bearer {key}"
+
+        def do_GET(self):  # noqa: N802 — http.server's name
+            path = self.path.rstrip("/")
+            if path == "/health":
+                return self._send(200, {})
+            if not self._authorized():
+                return self._send(401, {"error": "Unauthorized"})
+            if path == "/v1/models":
+                model = os.getenv("CAPTCHA_LORA_NAME", "captcha-pointer")
+                return self._send(200, {"object": "list", "data": [
+                    {"id": model, "object": "model", "created": 0, "owned_by": "captchakraken"},
+                ]})
+            return self._send(404, {"error": {"message": f"no route {self.path}"}})
+
+        def do_POST(self):  # noqa: N802 — http.server's name
+            if self.path.rstrip("/") != PATH:
+                return self._send(404, {"error": {"message": f"no route {self.path}"}})
+            if not self._authorized():
+                return self._send(401, {"error": "Unauthorized"})
+            try:
+                body = json.loads(self.rfile.read(int(self.headers.get("Content-Length") or 0)))
+                args = parse(body)
+            except (ValueError, KeyError, OSError) as e:
+                return self._send(400, {"error": {"message": str(e), "type": "invalid_request"}})
+            with lock:
+                result = service.solve(**args)
+            return self._send(200, reply(result))
+
+        def log_message(self, fmt, *args):  # quiet: one line per request is the caller's to log
+            pass
+
+    return Handler
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True, help="an exported pointer checkpoint")
+    ap.add_argument("--detector", default=None, help="the element detector; without it every board goes plain")
+    ap.add_argument("--letters", default=None,
+                    help="the letter detector: a typed code is the one it and the model read together")
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--port", type=int, default=8001)
+    ap.add_argument("--gpu-memory-utilization", type=float, default=0.45)
+    a = ap.parse_args()
+    from .service import PointerService
+    svc = PointerService(a.model, detector=a.detector, letters=a.letters,
+                         gpu_memory_utilization=a.gpu_memory_utilization)
+    ThreadingHTTPServer((a.host, a.port), make_handler(svc)).serve_forever()
+
+
+if __name__ == "__main__":
+    main()