backend.py

#!/usr/bin/env python3
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
import os
import tempfile
import threading
from pathlib import Path
from urllib.parse import urlparse


MODEL_ID = "Roboflow/rf-detr-medium"
HOST = os.environ.get("HOST", "0.0.0.0")
PORT = int(os.environ.get("PORT", "7860"))
ROOT = Path(__file__).resolve().parent

detector = None
detector_lock = threading.Lock()


def get_detector():
    global detector

    if detector is None:
        with detector_lock:
            if detector is None:
                from transformers import pipeline

                detector = pipeline(
                    "object-detection",
                    model=MODEL_ID,
                    device_map="auto",
                )

    return detector


def box_area(result):
    box = result.get("box") or {}
    width = max(0, box.get("xmax", 0) - box.get("xmin", 0))
    height = max(0, box.get("ymax", 0) - box.get("ymin", 0))
    return width * height


class Handler(BaseHTTPRequestHandler):
    def end_headers(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header(
            "Access-Control-Allow-Headers",
            "Content-Type, X-Confidence-Threshold, X-Focus-Mode",
        )
        self.send_header("Cross-Origin-Embedder-Policy", "require-corp")
        self.send_header("Cross-Origin-Opener-Policy", "same-origin")
        self.send_header("Cross-Origin-Resource-Policy", "cross-origin")
        super().end_headers()

    def do_OPTIONS(self):
        self.send_response(HTTPStatus.NO_CONTENT)
        self.end_headers()

    def do_HEAD(self):
        if not self.send_static(include_body=False):
            self.send_response(HTTPStatus.NOT_FOUND)
            self.end_headers()

    def do_GET(self):
        path = urlparse(self.path).path

        if path == "/health":
            self.send_json({
                "ok": True,
                "model": MODEL_ID,
                "loaded": detector is not None,
            })
            return

        if self.send_static():
            return

        self.send_json({"error": "not found"}, status=HTTPStatus.NOT_FOUND)

    def do_POST(self):
        if urlparse(self.path).path != "/detect":
            self.send_json({"error": "not found"}, status=404)
            return

        length = int(self.headers.get("Content-Length", "0"))
        image_bytes = self.rfile.read(length)
        threshold = float(self.headers.get("X-Confidence-Threshold", "0.75"))
        focus_mode = self.headers.get("X-Focus-Mode", "true") == "true"
        image_path = None

        try:
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as image_file:
                image_file.write(image_bytes)
                image_path = Path(image_file.name)

            raw_results = get_detector()(str(image_path), threshold=threshold)
            results = []

            for item in raw_results:
                result = {
                    "label": item["label"],
                    "score": float(item["score"]),
                    "box": item["box"],
                }
                result["area"] = box_area(result)
                results.append(result)

            if focus_mode:
                results = [item for item in results if item["label"] != "person"]

            results = [
                item for item in results
                if item["area"] >= 2500
            ]
            results.sort(key=lambda item: item["score"] * item["area"], reverse=True)
            results = results[:1 if focus_mode else 3]

            self.send_json({"results": results})
        except Exception as exc:
            self.send_json({"error": str(exc)}, status=HTTPStatus.INTERNAL_SERVER_ERROR)
        finally:
            if image_path is not None:
                try:
                    image_path.unlink(missing_ok=True)
                except Exception:
                    pass

    def send_json(self, data, status=200):
        body = json.dumps(data).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def send_static(self, include_body=True):
        path = urlparse(self.path).path
        files = {
            "/": ("index.html", "text/html; charset=utf-8"),
            "/index.html": ("index.html", "text/html; charset=utf-8"),
            "/worker.js": ("worker.js", "application/javascript; charset=utf-8"),
        }

        if path not in files:
            return False

        filename, content_type = files[path]
        file_path = ROOT / filename
        if not file_path.is_file():
            return False

        body = file_path.read_bytes()
        self.send_response(HTTPStatus.OK)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        if include_body:
            self.wfile.write(body)
        return True

    def log_message(self, format, *args):
        print("%s - %s" % (self.address_string(), format % args))


if __name__ == "__main__":
    print(f"Starting RF-DETR backend on http://{HOST}:{PORT}")
    print(f"Model will load on first /detect request: {MODEL_ID}")
    ThreadingHTTPServer((HOST, PORT), Handler).serve_forever()