459 lines
12 KiB
Python
459 lines
12 KiB
Python
# backend/ai_server.py
|
|
|
|
import json
|
|
import os
|
|
import secrets
|
|
from pathlib import Path
|
|
from typing import List, Optional
|
|
|
|
from fastapi import Depends, FastAPI, HTTPException, Request, status
|
|
from pydantic import BaseModel
|
|
from ultralytics import YOLO
|
|
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent
|
|
|
|
|
|
def existing_file(path: Path) -> Optional[Path]:
|
|
try:
|
|
if path.exists() and path.is_file() and path.stat().st_size > 0:
|
|
return path
|
|
except OSError:
|
|
pass
|
|
return None
|
|
|
|
|
|
def resolve_training_root() -> Path:
|
|
env_root = os.environ.get("TRAINING_ROOT", "").strip()
|
|
if env_root:
|
|
root = Path(env_root).expanduser().resolve()
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
return root
|
|
|
|
candidates = [
|
|
# Wenn ai_server.py aus backend/ läuft:
|
|
BASE_DIR / "generated" / "training",
|
|
|
|
# Wenn ai_server.py aus backend/ml/ laufen würde:
|
|
BASE_DIR.parent / "generated" / "training",
|
|
|
|
# Wenn ai_server.py embedded aus Temp läuft, aber backendRoot als cwd gesetzt wurde:
|
|
Path.cwd() / "generated" / "training",
|
|
|
|
# Wenn Working Directory Projektroot ist:
|
|
Path.cwd() / "backend" / "generated" / "training",
|
|
]
|
|
|
|
for root in candidates:
|
|
if (
|
|
existing_file(root / "detection_labels.json")
|
|
or existing_file(root / "detector" / "model" / "best.pt")
|
|
):
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
return root.resolve()
|
|
|
|
# Fallback: Server soll trotzdem starten.
|
|
root = (Path.cwd() / "generated" / "training").resolve()
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
return root
|
|
|
|
|
|
TRAINING_ROOT = resolve_training_root()
|
|
DEFAULT_MODEL_PATH = TRAINING_ROOT / "detector" / "model" / "best.pt"
|
|
|
|
|
|
def resolve_detection_labels_path() -> Path:
|
|
env_path = os.environ.get("DETECTION_LABELS_PATH", "").strip()
|
|
if env_path:
|
|
p = Path(env_path).expanduser().resolve()
|
|
if existing_file(p):
|
|
return p
|
|
raise RuntimeError(f"DETECTION_LABELS_PATH not found: {p}")
|
|
|
|
p = TRAINING_ROOT / "detection_labels.json"
|
|
if existing_file(p):
|
|
return p.resolve()
|
|
|
|
raise RuntimeError(f"detection_labels.json not found: {p}")
|
|
|
|
|
|
def resolve_model_path() -> str:
|
|
env_path = os.environ.get("YOLO_MODEL", "").strip()
|
|
if env_path:
|
|
p = Path(env_path).expanduser().resolve()
|
|
if existing_file(p):
|
|
return str(p)
|
|
raise RuntimeError(f"YOLO_MODEL not found: {p}")
|
|
|
|
if existing_file(DEFAULT_MODEL_PATH):
|
|
return str(DEFAULT_MODEL_PATH)
|
|
|
|
raise RuntimeError(f"YOLO model not found: {DEFAULT_MODEL_PATH}")
|
|
|
|
|
|
# Server darf auch ohne Labels/Model starten.
|
|
DETECTION_LABELS_PATH: Optional[Path] = None
|
|
|
|
LABEL_GROUPS = {
|
|
"people": set(),
|
|
"sexPositions": {"unknown"},
|
|
"bodyParts": set(),
|
|
"objects": set(),
|
|
"clothing": set(),
|
|
}
|
|
|
|
BODY_LABELS = LABEL_GROUPS["bodyParts"]
|
|
OBJECT_LABELS = LABEL_GROUPS["objects"]
|
|
CLOTHING_LABELS = LABEL_GROUPS["clothing"]
|
|
POSITION_LABELS = set()
|
|
|
|
PERSON_LABELS = {
|
|
"person_male",
|
|
"person_female",
|
|
}
|
|
|
|
_MODEL_PATH = ""
|
|
_MODEL_ERROR = ""
|
|
_LABEL_ERROR = ""
|
|
|
|
_DEVICE = os.environ.get("YOLO_DEVICE", "")
|
|
_CONF = float(os.environ.get("YOLO_CONF", "0.25"))
|
|
_BATCH = int(os.environ.get("YOLO_BATCH", "16"))
|
|
_IMGSZ = int(os.environ.get("YOLO_IMGSZ", "640"))
|
|
_HALF = os.environ.get("YOLO_HALF", "0").lower() in {"1", "true", "yes", "on"}
|
|
|
|
model = None
|
|
|
|
app = FastAPI()
|
|
|
|
AI_SERVER_TOKEN = os.environ.get("AI_SERVER_TOKEN", "").strip()
|
|
AI_SERVER_AUTH_REQUIRED = os.environ.get("AI_SERVER_AUTH_REQUIRED", "1").strip().lower() not in {
|
|
"0",
|
|
"false",
|
|
"no",
|
|
"off",
|
|
}
|
|
|
|
|
|
def require_ai_server_auth(request: Request) -> None:
|
|
if not AI_SERVER_AUTH_REQUIRED:
|
|
return
|
|
|
|
if not AI_SERVER_TOKEN:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail="AI server auth token not configured",
|
|
)
|
|
|
|
auth = request.headers.get("authorization", "").strip()
|
|
header_token = request.headers.get("x-ai-server-token", "").strip()
|
|
|
|
provided = ""
|
|
|
|
if auth.lower().startswith("bearer "):
|
|
provided = auth[7:].strip()
|
|
elif header_token:
|
|
provided = header_token
|
|
|
|
if not provided or not secrets.compare_digest(provided, AI_SERVER_TOKEN):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="unauthorized",
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
)
|
|
|
|
|
|
class PredictBatchRequest(BaseModel):
|
|
paths: List[str]
|
|
detectorOnly: bool = False
|
|
imageSize: int = 640
|
|
model: Optional[str] = None
|
|
|
|
|
|
def empty_prediction(source: str = "model_missing") -> dict:
|
|
return {
|
|
"modelAvailable": False,
|
|
"source": source,
|
|
"sexPosition": "unknown",
|
|
"sexPositionScore": 0.0,
|
|
"peoplePresent": [],
|
|
"bodyPartsPresent": [],
|
|
"objectsPresent": [],
|
|
"clothingPresent": [],
|
|
"boxes": [],
|
|
}
|
|
|
|
|
|
def load_label_groups_safe() -> None:
|
|
global DETECTION_LABELS_PATH
|
|
global LABEL_GROUPS
|
|
global BODY_LABELS
|
|
global OBJECT_LABELS
|
|
global CLOTHING_LABELS
|
|
global POSITION_LABELS
|
|
global PERSON_LABELS
|
|
global _LABEL_ERROR
|
|
|
|
try:
|
|
path = resolve_detection_labels_path()
|
|
DETECTION_LABELS_PATH = path
|
|
|
|
with path.open("r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
|
|
LABEL_GROUPS = {
|
|
"people": set(
|
|
str(x).strip().lower()
|
|
for x in data.get("people", [])
|
|
if str(x).strip()
|
|
),
|
|
"sexPositions": set(
|
|
str(x).strip().lower()
|
|
for x in data.get("sexPositions", [])
|
|
if str(x).strip()
|
|
),
|
|
"bodyParts": set(
|
|
str(x).strip().lower()
|
|
for x in data.get("bodyParts", [])
|
|
if str(x).strip()
|
|
),
|
|
"objects": set(
|
|
str(x).strip().lower()
|
|
for x in data.get("objects", [])
|
|
if str(x).strip()
|
|
),
|
|
"clothing": set(
|
|
str(x).strip().lower()
|
|
for x in data.get("clothing", [])
|
|
if str(x).strip()
|
|
),
|
|
}
|
|
|
|
if not LABEL_GROUPS["sexPositions"]:
|
|
LABEL_GROUPS["sexPositions"] = {"unknown"}
|
|
|
|
_LABEL_ERROR = ""
|
|
|
|
except Exception as exc:
|
|
DETECTION_LABELS_PATH = None
|
|
_LABEL_ERROR = str(exc)
|
|
|
|
LABEL_GROUPS = {
|
|
"people": set(),
|
|
"sexPositions": {"unknown"},
|
|
"bodyParts": set(),
|
|
"objects": set(),
|
|
"clothing": set(),
|
|
}
|
|
|
|
BODY_LABELS = LABEL_GROUPS["bodyParts"]
|
|
OBJECT_LABELS = LABEL_GROUPS["objects"]
|
|
CLOTHING_LABELS = LABEL_GROUPS["clothing"]
|
|
|
|
POSITION_LABELS = {
|
|
label for label in LABEL_GROUPS["sexPositions"]
|
|
if label and label != "unknown"
|
|
}
|
|
|
|
PERSON_LABELS = {
|
|
label for label in LABEL_GROUPS["people"]
|
|
if label
|
|
}
|
|
|
|
|
|
def get_model():
|
|
global model
|
|
global _MODEL_PATH
|
|
global _MODEL_ERROR
|
|
|
|
if model is not None:
|
|
return model
|
|
|
|
try:
|
|
path = resolve_model_path()
|
|
loaded = YOLO(path)
|
|
|
|
model = loaded
|
|
_MODEL_PATH = path
|
|
_MODEL_ERROR = ""
|
|
|
|
# Labels erst laden, wenn Inference wirklich gebraucht wird.
|
|
load_label_groups_safe()
|
|
|
|
return model
|
|
|
|
except Exception as exc:
|
|
model = None
|
|
_MODEL_PATH = ""
|
|
_MODEL_ERROR = str(exc)
|
|
return None
|
|
|
|
|
|
def scored(label: str, score: float) -> dict:
|
|
return {
|
|
"label": label,
|
|
"score": float(score),
|
|
}
|
|
|
|
|
|
def best_score(items: list[dict], label: str, score: float) -> None:
|
|
for item in items:
|
|
if item["label"] == label:
|
|
if score > item["score"]:
|
|
item["score"] = float(score)
|
|
return
|
|
|
|
items.append(scored(label, score))
|
|
|
|
|
|
def prediction_from_result(result) -> dict:
|
|
names = result.names or {}
|
|
|
|
boxes_out = []
|
|
people_present = []
|
|
body_parts = []
|
|
objects = []
|
|
clothing = []
|
|
|
|
sex_position = "unknown"
|
|
sex_position_score = 0.0
|
|
|
|
if result.boxes is not None:
|
|
xywhn = result.boxes.xywhn.cpu().tolist()
|
|
cls_values = result.boxes.cls.cpu().tolist()
|
|
conf_values = result.boxes.conf.cpu().tolist()
|
|
|
|
for box_xywhn, cls_id, conf in zip(xywhn, cls_values, conf_values):
|
|
label = str(names.get(int(cls_id), int(cls_id))).strip().lower()
|
|
score = float(conf)
|
|
|
|
if not label:
|
|
continue
|
|
|
|
cx, cy, w, h = [float(v) for v in box_xywhn]
|
|
|
|
x = max(0.0, min(1.0, cx - w / 2.0))
|
|
y = max(0.0, min(1.0, cy - h / 2.0))
|
|
w = max(0.0, min(1.0 - x, w))
|
|
h = max(0.0, min(1.0 - y, h))
|
|
|
|
is_person = label in PERSON_LABELS
|
|
is_body = label in BODY_LABELS
|
|
is_object = label in OBJECT_LABELS
|
|
is_clothing = label in CLOTHING_LABELS
|
|
is_position = label in POSITION_LABELS
|
|
|
|
if is_position:
|
|
if score > sex_position_score:
|
|
sex_position = label
|
|
sex_position_score = score
|
|
continue
|
|
|
|
if not (is_person or is_body or is_object or is_clothing):
|
|
continue
|
|
|
|
boxes_out.append({
|
|
"label": label,
|
|
"score": score,
|
|
"x": x,
|
|
"y": y,
|
|
"w": w,
|
|
"h": h,
|
|
})
|
|
|
|
if is_person:
|
|
best_score(people_present, label, score)
|
|
|
|
if is_body:
|
|
best_score(body_parts, label, score)
|
|
|
|
if is_object:
|
|
best_score(objects, label, score)
|
|
|
|
if is_clothing:
|
|
best_score(clothing, label, score)
|
|
|
|
return {
|
|
"modelAvailable": True,
|
|
"source": f"yolo-server:{Path(_MODEL_PATH).name}",
|
|
"sexPosition": sex_position,
|
|
"sexPositionScore": sex_position_score,
|
|
"peoplePresent": people_present,
|
|
"bodyPartsPresent": body_parts,
|
|
"objectsPresent": objects,
|
|
"clothingPresent": clothing,
|
|
"boxes": boxes_out,
|
|
}
|
|
|
|
|
|
@app.post("/predict-batch", dependencies=[Depends(require_ai_server_auth)])
|
|
def predict_batch(req: PredictBatchRequest):
|
|
paths = [str(path).strip() for path in req.paths if str(path).strip()]
|
|
if not paths:
|
|
return {
|
|
"ok": False,
|
|
"predictions": [],
|
|
"error": "no paths supplied",
|
|
}
|
|
|
|
current_model = get_model()
|
|
if current_model is None:
|
|
return {
|
|
"ok": True,
|
|
"predictions": [empty_prediction("model_missing") for _ in paths],
|
|
"error": _MODEL_ERROR or f"YOLO model not found: {DEFAULT_MODEL_PATH}",
|
|
}
|
|
|
|
if DETECTION_LABELS_PATH is None or _LABEL_ERROR:
|
|
return {
|
|
"ok": False,
|
|
"predictions": [],
|
|
"error": f"detection labels missing: {_LABEL_ERROR}",
|
|
}
|
|
|
|
imgsz = int(req.imageSize or _IMGSZ or 640)
|
|
|
|
try:
|
|
results = current_model.predict(
|
|
source=paths,
|
|
imgsz=imgsz,
|
|
conf=_CONF,
|
|
batch=_BATCH,
|
|
device=_DEVICE or None,
|
|
half=_HALF,
|
|
verbose=False,
|
|
)
|
|
|
|
predictions = [prediction_from_result(result) for result in results]
|
|
|
|
return {
|
|
"ok": True,
|
|
"predictions": predictions,
|
|
}
|
|
|
|
except Exception as exc:
|
|
return {
|
|
"ok": False,
|
|
"predictions": [],
|
|
"error": str(exc),
|
|
}
|
|
|
|
|
|
@app.get("/health", dependencies=[Depends(require_ai_server_auth)])
|
|
def health():
|
|
current_model = get_model()
|
|
names = getattr(current_model, "names", {}) or {} if current_model is not None else {}
|
|
|
|
return {
|
|
"ok": True,
|
|
"ready": current_model is not None,
|
|
"modelAvailable": current_model is not None,
|
|
"model": _MODEL_PATH,
|
|
"modelError": _MODEL_ERROR,
|
|
"expectedModel": str(DEFAULT_MODEL_PATH),
|
|
"trainingRoot": str(TRAINING_ROOT),
|
|
"classCount": len(names),
|
|
"classes": list(names.values())[:80] if isinstance(names, dict) else names,
|
|
"labelConfig": str(DETECTION_LABELS_PATH) if DETECTION_LABELS_PATH else "",
|
|
"labelError": _LABEL_ERROR,
|
|
} |