File size: 5,535 Bytes
43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f 3c9e0f6 43e737f | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | #!/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()
|