341 lines
12 KiB
Python
341 lines
12 KiB
Python
# coding=utf-8
|
||
"""总服务器中心 API：收图、静态访问、识别结果落库。
|
||
|
||
本文件不依赖 YOLO / 相机 / 客户端 UI，可单独拷到总服务器运行。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import threading
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
from urllib.parse import parse_qs, unquote, urlparse
|
||
|
||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||
if _HERE not in sys.path:
|
||
sys.path.insert(0, _HERE)
|
||
|
||
try:
|
||
from .media_store import ALLOWED_TYPES, MEDIA_ROOT, MediaStore, local_path_under
|
||
except ImportError:
|
||
from media_store import ALLOWED_TYPES, MEDIA_ROOT, MediaStore, local_path_under
|
||
|
||
|
||
def _cfg():
|
||
host = os.environ.get("DISH_API_HOST", "0.0.0.0")
|
||
port = int(os.environ.get("DISH_API_PORT", "9000"))
|
||
token = os.environ.get("DISH_API_TOKEN", "").strip()
|
||
result_dir = os.environ.get("DISH_RESULT_STORE") or os.path.join(_HERE, "data", "results")
|
||
max_bytes = int(os.environ.get("DISH_MEDIA_MAX_BYTES", str(12 * 1024 * 1024)))
|
||
return host, port, token, result_dir, max_bytes
|
||
|
||
|
||
HOST, PORT, API_TOKEN, RESULT_DIR, MAX_UPLOAD_BYTES = _cfg()
|
||
os.makedirs(RESULT_DIR, exist_ok=True)
|
||
_store = MediaStore()
|
||
_result_lock = threading.Lock()
|
||
|
||
|
||
def reload_runtime():
|
||
"""start.py 加载 env 后再调用，刷新监听配置与存储路径。"""
|
||
global HOST, PORT, API_TOKEN, RESULT_DIR, MAX_UPLOAD_BYTES, _store
|
||
HOST, PORT, API_TOKEN, RESULT_DIR, MAX_UPLOAD_BYTES = _cfg()
|
||
os.makedirs(RESULT_DIR, exist_ok=True)
|
||
public = os.environ.get("DISH_MEDIA_PUBLIC_BASE", "http://127.0.0.1:9000/media").rstrip("/")
|
||
root = os.environ.get("DISH_MEDIA_ROOT") or os.path.join(_HERE, "data", "media")
|
||
_store = MediaStore(root=root, public_base=public)
|
||
|
||
|
||
def _json_bytes(obj, status=200):
|
||
raw = json.dumps(obj, ensure_ascii=False).encode("utf-8")
|
||
return status, raw
|
||
|
||
|
||
def _ok(data=None, message="ok"):
|
||
return _json_bytes({"code": 0, "message": message, "data": data if data is not None else {}})
|
||
|
||
|
||
def _err(http_status, code, message):
|
||
return _json_bytes({"code": code, "message": message, "data": None}, status=http_status)
|
||
|
||
|
||
def _parse_multipart(content_type, body):
|
||
files = {}
|
||
fields = {}
|
||
if not content_type or "multipart/form-data" not in content_type.lower():
|
||
return fields, files
|
||
boundary = None
|
||
for part in content_type.split(";"):
|
||
part = part.strip()
|
||
if part.lower().startswith("boundary="):
|
||
boundary = part.split("=", 1)[1].strip().strip('"')
|
||
break
|
||
if not boundary:
|
||
return fields, files
|
||
marker = b"--" + boundary.encode("utf-8")
|
||
chunks = body.split(marker)
|
||
for chunk in chunks:
|
||
chunk = chunk.lstrip(b"\r\n")
|
||
if not chunk or chunk.startswith(b"--"):
|
||
continue
|
||
if b"\r\n\r\n" not in chunk:
|
||
continue
|
||
header_raw, data = chunk.split(b"\r\n\r\n", 1)
|
||
if data.endswith(b"\r\n"):
|
||
data = data[:-2]
|
||
headers = {}
|
||
for line in header_raw.decode("utf-8", errors="replace").split("\r\n"):
|
||
if ":" in line:
|
||
k, v = line.split(":", 1)
|
||
headers[k.strip().lower()] = v.strip()
|
||
disp = headers.get("content-disposition", "")
|
||
name = None
|
||
filename = None
|
||
for item in disp.split(";"):
|
||
item = item.strip()
|
||
if item.lower().startswith("name="):
|
||
name = item.split("=", 1)[1].strip().strip('"')
|
||
elif item.lower().startswith("filename="):
|
||
filename = item.split("=", 1)[1].strip().strip('"')
|
||
if not name:
|
||
continue
|
||
if filename is not None:
|
||
files[name] = (filename, data)
|
||
else:
|
||
fields[name] = data.decode("utf-8", errors="replace")
|
||
return fields, files
|
||
|
||
|
||
def _auth_ok(handler):
|
||
if not API_TOKEN:
|
||
return True
|
||
header = handler.headers.get("Authorization") or ""
|
||
if header == f"Bearer {API_TOKEN}":
|
||
return True
|
||
if header == API_TOKEN:
|
||
return True
|
||
if handler.headers.get("X-Api-Token") == API_TOKEN:
|
||
return True
|
||
return False
|
||
|
||
|
||
def _guess_mime(path):
|
||
ext = os.path.splitext(path)[1].lower()
|
||
return {
|
||
".jpg": "image/jpeg",
|
||
".jpeg": "image/jpeg",
|
||
".png": "image/png",
|
||
".bmp": "image/bmp",
|
||
".webp": "image/webp",
|
||
".json": "application/json",
|
||
}.get(ext, "application/octet-stream")
|
||
|
||
|
||
def _save_result(payload):
|
||
request_id = str(payload.get("request_id") or "").strip()
|
||
if not request_id:
|
||
return False, "missing request_id", None
|
||
path = os.path.join(RESULT_DIR, f"{request_id}.json")
|
||
duplicate = False
|
||
with _result_lock:
|
||
if os.path.isfile(path):
|
||
duplicate = True
|
||
else:
|
||
with open(path, "w", encoding="utf-8") as fh:
|
||
json.dump(payload, fh, ensure_ascii=False, indent=2)
|
||
return True, None, duplicate
|
||
|
||
|
||
class ApiHandler(BaseHTTPRequestHandler):
|
||
protocol_version = "HTTP/1.1"
|
||
|
||
def log_message(self, fmt, *args):
|
||
sys.stderr.write("[central] " + (fmt % args) + "\n")
|
||
|
||
def _send(self, status, body, content_type="application/json; charset=utf-8"):
|
||
self.send_response(status)
|
||
self.send_header("Content-Type", content_type)
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Access-Control-Allow-Origin", "*")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def _read_body(self, limit=None):
|
||
limit = MAX_UPLOAD_BYTES if limit is None else limit
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
if length > limit:
|
||
return None
|
||
return self.rfile.read(length) if length else b""
|
||
|
||
def do_OPTIONS(self):
|
||
self.send_response(204)
|
||
self.send_header("Access-Control-Allow-Origin", "*")
|
||
self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Base-Code, X-Api-Token")
|
||
self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
|
||
self.end_headers()
|
||
|
||
def do_GET(self):
|
||
parsed = urlparse(self.path)
|
||
path = unquote(parsed.path)
|
||
if path in ("/api/health", "/health"):
|
||
status, body = _ok({
|
||
"status": "up",
|
||
"media_root": _store.root,
|
||
"result_dir": RESULT_DIR,
|
||
})
|
||
self._send(status, body)
|
||
return
|
||
if path.startswith("/media/"):
|
||
rel = path[len("/media/"):]
|
||
try:
|
||
dest = local_path_under(_store.root, rel)
|
||
except ValueError:
|
||
self._send(404, b"not found", "text/plain")
|
||
return
|
||
if not os.path.isfile(dest):
|
||
self._send(404, b"not found", "text/plain")
|
||
return
|
||
with open(dest, "rb") as fh:
|
||
data = fh.read()
|
||
self._send(200, data, _guess_mime(dest))
|
||
return
|
||
if path == "/api/dish/recognize/records":
|
||
qs = parse_qs(parsed.query or "")
|
||
rid = (qs.get("request_id") or [None])[0]
|
||
items = []
|
||
if os.path.isdir(RESULT_DIR):
|
||
for name in sorted(os.listdir(RESULT_DIR), reverse=True):
|
||
if not name.endswith(".json"):
|
||
continue
|
||
fpath = os.path.join(RESULT_DIR, name)
|
||
try:
|
||
with open(fpath, "r", encoding="utf-8") as fh:
|
||
row = json.load(fh)
|
||
except Exception:
|
||
continue
|
||
if rid and row.get("request_id") != rid:
|
||
continue
|
||
items.append(row)
|
||
status, body = _ok({"items": items[:50], "total": len(items)})
|
||
self._send(status, body)
|
||
return
|
||
if path.startswith("/api/dish/recognize/records/"):
|
||
rid = path.rsplit("/", 1)[-1]
|
||
fpath = os.path.join(RESULT_DIR, f"{rid}.json")
|
||
if not os.path.isfile(fpath):
|
||
status, body = _err(404, 40400, "record not found")
|
||
self._send(status, body)
|
||
return
|
||
with open(fpath, "r", encoding="utf-8") as fh:
|
||
row = json.load(fh)
|
||
status, body = _ok(row)
|
||
self._send(status, body)
|
||
return
|
||
status, body = _err(404, 40400, "not found")
|
||
self._send(status, body)
|
||
|
||
def do_POST(self):
|
||
parsed = urlparse(self.path)
|
||
path = parsed.path.rstrip("/") or "/"
|
||
if not _auth_ok(self):
|
||
status, body = _err(401, 40100, "未授权")
|
||
self._send(status, body)
|
||
return
|
||
body = self._read_body()
|
||
if body is None:
|
||
status, body = _err(413, 41300, "文件过大")
|
||
self._send(status, body)
|
||
return
|
||
|
||
if path == "/api/dish/media/upload":
|
||
fields, files = _parse_multipart(self.headers.get("Content-Type", ""), body)
|
||
blob = files.get("file")
|
||
if not blob:
|
||
status, out = _err(400, 40001, "缺少 file")
|
||
self._send(status, out)
|
||
return
|
||
filename, data = blob
|
||
media_type = (fields.get("type") or "source").strip().lower()
|
||
if media_type not in ALLOWED_TYPES:
|
||
status, out = _err(400, 40001, "type 必须是 source / labeled / crop")
|
||
self._send(status, out)
|
||
return
|
||
base_id = fields.get("base_id") or self.headers.get("X-Base-Code") or "BASE"
|
||
request_id = fields.get("request_id") or ""
|
||
try:
|
||
saved = _store.save_bytes(
|
||
data, filename, media_type=media_type, base_id=base_id, request_id=request_id,
|
||
)
|
||
except Exception as exc:
|
||
status, out = _err(500, 50000, str(exc))
|
||
self._send(status, out)
|
||
return
|
||
status, out = _ok({"url": saved["url"], "path": saved["rel_path"], "type": saved["type"]})
|
||
self._send(status, out)
|
||
return
|
||
|
||
if path == "/api/dish/recognize/result":
|
||
try:
|
||
payload = json.loads(body.decode("utf-8"))
|
||
except Exception:
|
||
status, out = _err(400, 40001, "JSON 无法解析")
|
||
self._send(status, out)
|
||
return
|
||
if not isinstance(payload, dict):
|
||
status, out = _err(400, 40001, "body 必须是对象")
|
||
self._send(status, out)
|
||
return
|
||
dishes = payload.get("dishes") or []
|
||
try:
|
||
dish_count = int(payload.get("dish_count"))
|
||
except Exception:
|
||
dish_count = -1
|
||
if dish_count != len(dishes):
|
||
status, out = _err(400, 40001, "dish_count 与 dishes 长度不一致")
|
||
self._send(status, out)
|
||
return
|
||
if not payload.get("request_id") or not payload.get("base_id"):
|
||
status, out = _err(400, 40001, "缺少 request_id 或 base_id")
|
||
self._send(status, out)
|
||
return
|
||
ok, msg, duplicate = _save_result(payload)
|
||
if not ok:
|
||
status, out = _err(400, 40001, msg)
|
||
self._send(status, out)
|
||
return
|
||
status, out = _ok({
|
||
"request_id": payload["request_id"],
|
||
"accepted": True,
|
||
"duplicate": bool(duplicate),
|
||
})
|
||
self._send(status, out)
|
||
return
|
||
|
||
status, out = _err(404, 40400, "not found")
|
||
self._send(status, out)
|
||
|
||
|
||
def serve(host=None, port=None):
|
||
reload_runtime()
|
||
host = HOST if host is None else host
|
||
port = PORT if port is None else port
|
||
httpd = ThreadingHTTPServer((host, port), ApiHandler)
|
||
print(f"[central] listening http://{host}:{port}")
|
||
print("[central] GET /api/health")
|
||
print("[central] POST /api/dish/media/upload")
|
||
print("[central] GET /media/...")
|
||
print("[central] POST /api/dish/recognize/result")
|
||
print(f"[central] media {_store.root}")
|
||
print(f"[central] result {RESULT_DIR}")
|
||
try:
|
||
httpd.serve_forever()
|
||
except KeyboardInterrupt:
|
||
print("\n[central] stopped")
|
||
finally:
|
||
httpd.server_close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
serve()
|