chris/nsfwapp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
nsfwapp/backend/ai_server.py
Linrador 01c8850a39 added pose correction
2026-06-29 15:44:05 +02:00

1930 lines
57 KiB
Python
Raw Blame History

# backend\ai_server.py
import json
import math
import os
import secrets
import tempfile
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
KEYPOINT_NAMES = [
"nose",
"left_eye", "right_eye",
"left_ear", "right_ear",
"left_shoulder", "right_shoulder",
"left_elbow", "right_elbow",
"left_wrist", "right_wrist",
"left_hip", "right_hip",
"left_knee", "right_knee",
"left_ankle", "right_ankle",
]
NO_SEX_POSITION_LABEL = "keine"
NO_SEX_POSITION_ALIASES = {
"",
NO_SEX_POSITION_LABEL,
}
def normalize_sex_position_label(value) -> str:
clean = str(value or "").strip().lower()
if clean in NO_SEX_POSITION_ALIASES:
return NO_SEX_POSITION_LABEL
return clean
def is_no_sex_position_label(value) -> bool:
return normalize_sex_position_label(value) == NO_SEX_POSITION_LABEL
def normalize_sex_position_labels(values) -> set[str]:
labels = set()
has_no_position = False
for value in values or []:
clean = normalize_sex_position_label(value)
if is_no_sex_position_label(clean):
has_no_position = True
continue
if clean:
labels.add(clean)
if has_no_position:
labels.add(NO_SEX_POSITION_LABEL)
return labels
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 existing_model_dir(path: Path) -> Optional[Path]:
try:
if path.exists() and path.is_dir() and existing_file(path / "config.json"):
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")
or existing_model_dir(root / "videomae" / "model")
):
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"
DEFAULT_POSE_MODEL_PATH = TRAINING_ROOT / "pose" / "model" / "best.pt"
DEFAULT_VIDEOMAE_MODEL_PATH = TRAINING_ROOT / "videomae" / "model"
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}")
def resolve_pose_model_path() -> Optional[Path]:
env_path = os.environ.get("YOLO_POSE_MODEL", "").strip()
if env_path:
p = Path(env_path).expanduser().resolve()
if existing_file(p):
return p
raise RuntimeError(f"YOLO_POSE_MODEL not found: {p}")
if existing_file(DEFAULT_POSE_MODEL_PATH):
return DEFAULT_POSE_MODEL_PATH.resolve()
return None
def base_pose_model_candidates() -> list[Path]:
return [
BASE_DIR / "yolo26n-pose.pt",
Path.cwd() / "yolo26n-pose.pt",
TRAINING_ROOT / "yolo26n-pose.pt",
TRAINING_ROOT.parent / "yolo26n-pose.pt",
TRAINING_ROOT.parent.parent / "yolo26n-pose.pt",
Path(tempfile.gettempdir()) / "nsfwapp-ml" / "yolo26n-pose.pt",
]
def resolve_base_pose_model_path() -> Optional[Path]:
env_path = os.environ.get("YOLO_BASE_POSE_MODEL", "").strip()
if env_path:
p = Path(env_path).expanduser().resolve()
if existing_file(p):
return p
raise RuntimeError(f"YOLO_BASE_POSE_MODEL not found: {p}")
for candidate in base_pose_model_candidates():
p = existing_file(candidate)
if p:
return p.resolve()
return None
def resolve_videomae_model_path() -> Optional[Path]:
env_path = os.environ.get("VIDEOMAE_MODEL", "").strip()
if env_path:
p = Path(env_path).expanduser().resolve()
if existing_model_dir(p):
return p
raise RuntimeError(f"VIDEOMAE_MODEL not found: {p}")
if existing_model_dir(DEFAULT_VIDEOMAE_MODEL_PATH):
return DEFAULT_VIDEOMAE_MODEL_PATH.resolve()
return None
# Server darf auch ohne Labels/Model starten.
DETECTION_LABELS_PATH: Optional[Path] = None
LABEL_GROUPS = {
"people": set(),
"sexPositions": {NO_SEX_POSITION_LABEL},
"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 = ""
_POSE_MODEL_PATH = ""
_POSE_MODEL_ERROR = ""
_BASE_POSE_MODEL_PATH = ""
_BASE_POSE_MODEL_ERROR = ""
_VIDEOMAE_MODEL_PATH = ""
_VIDEOMAE_MODEL_ERROR = ""
_VIDEOMAE_DEVICE_ACTIVE = ""
_LABEL_ERROR = ""
_DEVICE = os.environ.get("YOLO_DEVICE", "")
_CONF = float(os.environ.get("YOLO_CONF", "0.25"))
_POSE_CONF = float(os.environ.get("YOLO_POSE_CONF", "0.30"))
_BASE_POSE_CONF = float(os.environ.get("YOLO_BASE_POSE_CONF", "0.10"))
_POSE_RELIABLE_MIN_SCORE = float(os.environ.get("YOLO_POSE_RELIABLE_MIN_SCORE", "0.30"))
_POSE_RELIABLE_MIN_KEYPOINTS = int(os.environ.get("YOLO_POSE_RELIABLE_MIN_KEYPOINTS", "6"))
_POSE_RELIABLE_MIN_QUALITY = float(os.environ.get("YOLO_POSE_RELIABLE_MIN_QUALITY", "0.45"))
_POSITION_CONTEXT_MIN_SCORE = float(os.environ.get("YOLO_POSITION_CONTEXT_MIN_SCORE", "0.22"))
_POSITION_CONTEXT_MAX_SCORE = float(os.environ.get("YOLO_POSITION_CONTEXT_MAX_SCORE", "0.44"))
_POSITION_CONTEXT_BOOST_WEIGHT = float(os.environ.get("YOLO_POSITION_CONTEXT_BOOST_WEIGHT", "0.60"))
_POSITION_CONTEXT_OVERRIDE_MARGIN = float(os.environ.get("YOLO_POSITION_CONTEXT_OVERRIDE_MARGIN", "0.16"))
_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"}
_VIDEOMAE_DEVICE = os.environ.get("VIDEOMAE_DEVICE", "auto")
_VIDEOMAE_NUM_FRAMES = int(os.environ.get("VIDEOMAE_NUM_FRAMES", "16"))
model = None
pose_model = None
base_pose_model = None
videomae_model = None
videomae_processor = 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
class PositionClipItem(BaseModel):
time: float = 0.0
start: float = 0.0
end: float = 0.0
paths: List[str]
class PredictPositionClipsRequest(BaseModel):
clips: List[PositionClipItem]
numFrames: int = 16
def empty_prediction(source: str = "model_missing") -> dict:
return {
"modelAvailable": False,
"source": source,
"sexPosition": NO_SEX_POSITION_LABEL,
"sexPositionScore": 0.0,
"peoplePresent": [],
"bodyPartsPresent": [],
"objectsPresent": [],
"clothingPresent": [],
"boxes": [],
"persons": [],
}
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": normalize_sex_position_labels(data.get("sexPositions", [])),
"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"] = {NO_SEX_POSITION_LABEL}
_LABEL_ERROR = ""
except Exception as exc:
DETECTION_LABELS_PATH = None
_LABEL_ERROR = str(exc)
LABEL_GROUPS = {
"people": set(),
"sexPositions": {NO_SEX_POSITION_LABEL},
"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 not is_no_sex_position_label(label)
}
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 get_pose_model():
global pose_model
global _POSE_MODEL_PATH
global _POSE_MODEL_ERROR
if pose_model is not None:
return pose_model
try:
path = resolve_pose_model_path()
if path is None:
pose_model = None
_POSE_MODEL_PATH = ""
_POSE_MODEL_ERROR = ""
return None
loaded = YOLO(str(path))
pose_model = loaded
_POSE_MODEL_PATH = str(path)
_POSE_MODEL_ERROR = ""
return pose_model
except Exception as exc:
pose_model = None
_POSE_MODEL_PATH = ""
_POSE_MODEL_ERROR = str(exc)
return None
def get_base_pose_model():
global base_pose_model
global _BASE_POSE_MODEL_PATH
global _BASE_POSE_MODEL_ERROR
if base_pose_model is not None:
return base_pose_model
try:
path = resolve_base_pose_model_path()
if path is None:
base_pose_model = None
_BASE_POSE_MODEL_PATH = ""
_BASE_POSE_MODEL_ERROR = ""
return None
loaded = YOLO(str(path))
base_pose_model = loaded
_BASE_POSE_MODEL_PATH = str(path)
_BASE_POSE_MODEL_ERROR = ""
return base_pose_model
except Exception as exc:
base_pose_model = None
_BASE_POSE_MODEL_PATH = ""
_BASE_POSE_MODEL_ERROR = str(exc)
return None
def get_videomae_components():
global videomae_model
global videomae_processor
global _VIDEOMAE_MODEL_PATH
global _VIDEOMAE_MODEL_ERROR
global _VIDEOMAE_DEVICE_ACTIVE
if videomae_model is not None and videomae_processor is not None:
return videomae_model, videomae_processor
try:
path = resolve_videomae_model_path()
if path is None:
videomae_model = None
videomae_processor = None
_VIDEOMAE_MODEL_PATH = ""
_VIDEOMAE_MODEL_ERROR = ""
_VIDEOMAE_DEVICE_ACTIVE = ""
return None, None
import torch
from transformers import AutoImageProcessor, VideoMAEForVideoClassification
if str(_VIDEOMAE_DEVICE).lower() == "auto":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
else:
device = torch.device(_VIDEOMAE_DEVICE)
processor = AutoImageProcessor.from_pretrained(path)
loaded = VideoMAEForVideoClassification.from_pretrained(path)
loaded.to(device)
loaded.eval()
videomae_model = loaded
videomae_processor = processor
_VIDEOMAE_MODEL_PATH = str(path)
_VIDEOMAE_MODEL_ERROR = ""
_VIDEOMAE_DEVICE_ACTIVE = str(device)
return videomae_model, videomae_processor
except Exception as exc:
videomae_model = None
videomae_processor = None
_VIDEOMAE_MODEL_PATH = ""
_VIDEOMAE_DEVICE_ACTIVE = ""
_VIDEOMAE_MODEL_ERROR = str(exc)
return None, None
def resample_values(values: list, count: int) -> list:
if not values:
return []
if count <= 1:
return [values[0]]
if len(values) == 1:
return [values[0] for _ in range(count)]
last = len(values) - 1
return [
values[int(round((i * last) / max(1, count - 1)))]
for i in range(count)
]
def load_videomae_clip_frames(paths: list[str], num_frames: int):
from PIL import Image
clean_paths = [
str(path).strip()
for path in paths or []
if str(path).strip()
]
selected = resample_values(clean_paths, max(2, int(num_frames or _VIDEOMAE_NUM_FRAMES or 16)))
frames = []
for path in selected:
with Image.open(path) as img:
frames.append(img.convert("RGB").copy())
return frames
def predict_videomae_clip(clip: PositionClipItem, num_frames: int) -> dict:
current_model, current_processor = get_videomae_components()
if current_model is None or current_processor is None:
return {
"time": float(clip.time or 0.0),
"start": float(clip.start or 0.0),
"end": float(clip.end or 0.0),
"sexPosition": NO_SEX_POSITION_LABEL,
"sexPositionScore": 0.0,
"source": "videomae_missing",
"scores": [],
}
import torch
frames = load_videomae_clip_frames(clip.paths, num_frames)
if not frames:
return {
"time": float(clip.time or 0.0),
"start": float(clip.start or 0.0),
"end": float(clip.end or 0.0),
"sexPosition": NO_SEX_POSITION_LABEL,
"sexPositionScore": 0.0,
"source": "videomae_no_frames",
"scores": [],
}
inputs = current_processor(frames, return_tensors="pt")
device = next(current_model.parameters()).device
pixel_values = inputs["pixel_values"].to(device)
with torch.no_grad():
logits = current_model(pixel_values=pixel_values).logits
probs = torch.softmax(logits, dim=-1)[0].detach().cpu().tolist()
id_to_label = getattr(current_model.config, "id2label", {}) or {}
scores = []
best_label = NO_SEX_POSITION_LABEL
best_score = 0.0
for idx, score in enumerate(probs):
raw_label = id_to_label.get(idx, id_to_label.get(str(idx), idx))
label = normalize_sex_position_label(raw_label)
score = clamp01(score)
scores.append({
"label": label,
"score": score,
})
if score > best_score:
best_label = label
best_score = score
scores.sort(key=lambda item: item["score"], reverse=True)
return {
"time": float(clip.time or 0.0),
"start": float(clip.start or 0.0),
"end": float(clip.end or 0.0),
"sexPosition": best_label,
"sexPositionScore": best_score,
"source": "videomae",
"scores": scores[:10],
}
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 = []
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
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": NO_SEX_POSITION_LABEL,
"sexPositionScore": 0.0,
"peoplePresent": people_present,
"bodyPartsPresent": body_parts,
"objectsPresent": objects,
"clothingPresent": clothing,
"boxes": boxes_out,
"persons": [],
}
def pose_persons_from_result(result) -> list[dict]:
names = result.names or {}
persons = []
if result.boxes is None:
return persons
kpts_xyn = None
kpts_conf = None
if result.keypoints is not None:
try:
kpts_xyn = result.keypoints.xyn.cpu().tolist()
except Exception:
kpts_xyn = None
try:
kpts_conf = result.keypoints.conf.cpu().tolist()
except Exception:
kpts_conf = None
xywhn_values = result.boxes.xywhn.cpu().tolist()
cls_values = result.boxes.cls.cpu().tolist()
conf_values = result.boxes.conf.cpu().tolist()
for i, (box_xywhn, cls_id, conf) in enumerate(
zip(xywhn_values, cls_values, conf_values)
):
label = str(names.get(int(cls_id), int(cls_id))).strip().lower()
score = float(conf)
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))
if w <= 0 or h <= 0:
continue
keypoints = []
if kpts_xyn is not None and i < len(kpts_xyn):
for ki, point in enumerate(kpts_xyn[i]):
if len(point) < 2:
continue
kconf = 0.0
if (
kpts_conf is not None
and i < len(kpts_conf)
and ki < len(kpts_conf[i])
):
kconf = float(kpts_conf[i][ki])
keypoints.append({
"name": KEYPOINT_NAMES[ki] if ki < len(KEYPOINT_NAMES) else str(ki),
"x": float(point[0]),
"y": float(point[1]),
"conf": kconf,
})
persons.append(annotate_pose_person_quality({
"label": label,
"score": score,
"box": {
"x": x,
"y": y,
"w": w,
"h": h,
},
"keypoints": keypoints,
}))
return persons
def clamp01(value) -> float:
try:
n = float(value)
except Exception:
return 0.0
if not math.isfinite(n):
return 0.0
return max(0.0, min(1.0, n))
def is_finite01(value) -> bool:
try:
n = float(value)
except Exception:
return False
return math.isfinite(n) and 0.0 <= n <= 1.0
def normalized_box(box: dict | None) -> dict | None:
if not isinstance(box, dict):
return None
x = clamp01(box.get("x", 0.0))
y = clamp01(box.get("y", 0.0))
w = clamp01(box.get("w", 0.0))
h = clamp01(box.get("h", 0.0))
if w <= 0 or h <= 0:
return None
if x + w > 1:
w = 1 - x
if y + h > 1:
h = 1 - y
if w <= 0 or h <= 0:
return None
out = dict(box)
out.update({"x": x, "y": y, "w": w, "h": h})
return out
def box_area(box: dict) -> float:
clean = normalized_box(box)
if not clean:
return 0.0
return float(clean["w"]) * float(clean["h"])
def box_center(box: dict) -> tuple[float, float] | None:
clean = normalized_box(box)
if not clean:
return None
return float(clean["x"]) + float(clean["w"]) / 2.0, float(clean["y"]) + float(clean["h"]) / 2.0
def box_gap(a: dict, b: dict) -> float:
a = normalized_box(a)
b = normalized_box(b)
if not a or not b:
return 1.0
dx = max(0.0, max(float(a["x"]) - (float(b["x"]) + float(b["w"])), float(b["x"]) - (float(a["x"]) + float(a["w"]))))
dy = max(0.0, max(float(a["y"]) - (float(b["y"]) + float(b["h"])), float(b["y"]) - (float(a["y"]) + float(a["h"]))))
return math.sqrt(dx * dx + dy * dy)
def box_overlap_ratio(a: dict, b: dict) -> float:
a = normalized_box(a)
b = normalized_box(b)
if not a or not b:
return 0.0
left = max(float(a["x"]), float(b["x"]))
top = max(float(a["y"]), float(b["y"]))
right = min(float(a["x"]) + float(a["w"]), float(b["x"]) + float(b["w"]))
bottom = min(float(a["y"]) + float(a["h"]), float(b["y"]) + float(b["h"]))
if right <= left or bottom <= top:
return 0.0
min_area = min(box_area(a), box_area(b))
if min_area <= 0:
return 0.0
return clamp01(((right - left) * (bottom - top)) / min_area)
def box_horizontal_overlap_ratio(a: dict, b: dict) -> float:
a = normalized_box(a)
b = normalized_box(b)
if not a or not b:
return 0.0
left = max(float(a["x"]), float(b["x"]))
right = min(float(a["x"]) + float(a["w"]), float(b["x"]) + float(b["w"]))
if right <= left:
return 0.0
min_width = min(float(a["w"]), float(b["w"]))
if min_width <= 0:
return 0.0
return clamp01((right - left) / min_width)
def is_person_like_label(label: str) -> bool:
clean = str(label or "").strip().lower()
return clean == "person" or clean.startswith("person_")
def boxes_by_label(prediction: dict, *labels: str) -> list[dict]:
wanted = {str(label).strip().lower() for label in labels if str(label).strip()}
out = []
for box in prediction.get("boxes", []) or []:
label = str(box.get("label", "")).strip().lower() if isinstance(box, dict) else ""
if label not in wanted:
continue
clean = normalized_box(box)
if clean:
out.append(clean)
return out
def any_boxes_near(left_boxes: list[dict], right_boxes: list[dict], margin: float) -> bool:
for left in left_boxes:
for right in right_boxes:
if box_gap(left, right) <= margin or box_overlap_ratio(left, right) > 0:
return True
return False
def point_near_box(x: float, y: float, box: dict, margin: float) -> bool:
if not is_finite01(x) or not is_finite01(y):
return False
clean = normalized_box(box)
if not clean:
return False
return (
float(clean["x"]) - margin <= x <= float(clean["x"]) + float(clean["w"]) + margin
and float(clean["y"]) - margin <= y <= float(clean["y"]) + float(clean["h"]) + margin
)
def any_pose_keypoint_near_boxes(
persons: list[dict],
keypoint_names: list[str],
boxes: list[dict],
margin: float,
) -> bool:
if not boxes:
return False
wanted = {str(name).strip().lower() for name in keypoint_names}
for person in reliable_pose_persons(persons):
for point in person.get("keypoints", []) or []:
name = str(point.get("name", "")).strip().lower()
if name not in wanted:
continue
if float(point.get("conf") or 0.0) < 0.20:
continue
x = float(point.get("x") or 0.0)
y = float(point.get("y") or 0.0)
for box in boxes:
if point_near_box(x, y, box, margin):
return True
return False
def pose_keypoint_stats(person: dict) -> tuple[int, float]:
keypoints = person.get("keypoints", []) or []
if not keypoints:
return 0, 0.0
visible = 0
total_conf = 0.0
for point in keypoints:
x = float(point.get("x") or 0.0)
y = float(point.get("y") or 0.0)
conf = float(point.get("conf") or 0.0)
if conf < 0.20 or not is_finite01(x) or not is_finite01(y):
continue
visible += 1
total_conf += clamp01(conf)
if visible == 0:
return 0, 0.0
coverage = clamp01(visible / max(1, len(KEYPOINT_NAMES)))
avg_conf = clamp01(total_conf / visible)
return visible, clamp01(coverage * 0.45 + avg_conf * 0.55)
def pose_keypoint_quality(person: dict) -> float:
_, quality = pose_keypoint_stats(person)
return quality
def annotate_pose_person_quality(person: dict) -> dict:
visible, quality = pose_keypoint_stats(person)
score = clamp01(float(person.get("score") or 0.0))
person["visibleKeypoints"] = visible
person["quality"] = quality
person["reliable"] = (
score >= _POSE_RELIABLE_MIN_SCORE
and visible >= _POSE_RELIABLE_MIN_KEYPOINTS
and quality >= _POSE_RELIABLE_MIN_QUALITY
)
return person
def reliable_pose_persons(persons: list[dict]) -> list[dict]:
out = []
for person in persons:
if not isinstance(person, dict):
continue
if "reliable" not in person:
annotate_pose_person_quality(person)
if bool(person.get("reliable")):
out.append(person)
return out
def scene_person_boxes(prediction: dict, persons: list[dict]) -> list[dict]:
detector_boxes = []
pose_boxes = []
for box in prediction.get("boxes", []) or []:
if not isinstance(box, dict) or not is_person_like_label(box.get("label", "")):
continue
clean = normalized_box(box)
if clean:
detector_boxes.append(clean)
for person in reliable_pose_persons(persons):
clean = normalized_box(person.get("box"))
if clean:
pose_boxes.append(clean)
if detector_boxes and len(detector_boxes) >= len(pose_boxes):
return detector_boxes
return pose_boxes
def scene_person_pair_signals(boxes: list[dict]) -> dict:
signals = {
"close": False,
"overlap": False,
"horizontal": False,
"vertical": False,
"stacked": False,
}
for i, left in enumerate(boxes):
left = normalized_box(left)
if not left:
continue
for raw_right in boxes[i + 1:]:
right = normalized_box(raw_right)
if not right:
continue
gap = box_gap(left, right)
overlap = box_overlap_ratio(left, right)
close = gap <= 0.08 or overlap >= 0.08
if close:
signals["close"] = True
if overlap >= 0.18:
signals["overlap"] = True
if close and float(left["w"]) > float(left["h"]) * 1.15 and float(right["w"]) > float(right["h"]) * 1.15:
signals["horizontal"] = True
if close and float(left["h"]) > float(left["w"]) * 1.25 and float(right["h"]) > float(right["w"]) * 1.25:
signals["vertical"] = True
left_center = box_center(left)
right_center = box_center(right)
if (
left_center
and right_center
and gap <= 0.12
and box_horizontal_overlap_ratio(left, right) >= 0.35
and abs(left_center[1] - right_center[1]) >= 0.12
):
signals["stacked"] = True
return signals
def pose_point(person: dict, name: str, min_conf: float = 0.20) -> tuple[float, float] | None:
wanted = str(name or "").strip().lower()
if not wanted:
return None
for point in person.get("keypoints", []) or []:
point_name = str(point.get("name", "")).strip().lower()
if point_name != wanted:
continue
conf = float(point.get("conf") or 0.0)
x = float(point.get("x") or 0.0)
y = float(point.get("y") or 0.0)
if conf >= min_conf and is_finite01(x) and is_finite01(y):
return x, y
return None
def pose_midpoint(
person: dict,
names: list[str],
min_conf: float = 0.20,
min_points: int = 1,
) -> tuple[float, float] | None:
points = [
point
for name in names
if (point := pose_point(person, name, min_conf)) is not None
]
if len(points) < min_points:
return None
x = sum(point[0] for point in points) / len(points)
y = sum(point[1] for point in points) / len(points)
return x, y
def point_distance(a: tuple[float, float] | None, b: tuple[float, float] | None) -> float:
if not a or not b:
return 0.0
return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)
def pose_person_geometry(person: dict) -> dict:
box = normalized_box(person.get("box"))
box_center_point = box_center(box) if box else None
left_hip = pose_point(person, "left_hip")
right_hip = pose_point(person, "right_hip")
left_knee = pose_point(person, "left_knee")
right_knee = pose_point(person, "right_knee")
hip = pose_midpoint(person, ["left_hip", "right_hip"])
shoulder = pose_midpoint(person, ["left_shoulder", "right_shoulder"])
knee = pose_midpoint(person, ["left_knee", "right_knee"])
head = pose_midpoint(person, ["nose", "left_eye", "right_eye", "left_ear", "right_ear"])
center = hip or box_center_point or (0.5, 0.5)
torso_dx = 0.0
torso_dy = 0.0
torso_len = 0.0
if hip and shoulder:
torso_dx = abs(hip[0] - shoulder[0])
torso_dy = abs(hip[1] - shoulder[1])
torso_len = math.sqrt(torso_dx * torso_dx + torso_dy * torso_dy)
hip_width = point_distance(left_hip, right_hip)
knee_width = point_distance(left_knee, right_knee)
knees_below_hips = bool(knee and hip and knee[1] > hip[1] + 0.045)
knees_wide = bool(
knee_width > 0
and knee_width >= max(0.11, hip_width * 1.12)
)
straddling = knees_below_hips and knees_wide
torso_horizontal = torso_len >= 0.07 and torso_dx >= torso_dy * 1.15
torso_vertical = torso_len >= 0.07 and torso_dy >= torso_dx * 1.15
box_horizontal = bool(box and float(box["w"]) >= float(box["h"]) * 1.05)
box_vertical = bool(box and float(box["h"]) >= float(box["w"]) * 1.25)
lying = torso_horizontal or box_horizontal
upright = torso_vertical or box_vertical
all_fours = torso_horizontal and knees_below_hips
bent_or_kneeling = all_fours or (knees_below_hips and not straddling)
return {
"box": box,
"center": center,
"head": head,
"hip": hip,
"shoulder": shoulder,
"knee": knee,
"lying": lying,
"upright": upright,
"straddling": straddling,
"knees_below_hips": knees_below_hips,
"knees_wide": knees_wide,
"all_fours": all_fours,
"bent_or_kneeling": bent_or_kneeling,
}
def combine_position_score(scores: dict[str, float], label: str, score: float) -> None:
label = normalize_sex_position_label(label)
if is_no_sex_position_label(label) or label not in POSITION_LABELS or score <= 0:
return
score = clamp01(score)
current = clamp01(scores.get(label, 0.0))
scores[label] = clamp01(1 - (1 - current) * (1 - score))
def append_prediction_source(prediction: dict, source: str) -> None:
source = str(source or "").strip()
if not source:
return
current = str(prediction.get("source") or "").strip()
if not current:
prediction["source"] = source
return
if source in {part.strip() for part in current.split("+")}:
return
prediction["source"] = f"{current}+{source}"
def fuse_hybrid_position_scores(
pose_scores: dict[str, float],
context_scores: dict[str, float],
) -> tuple[str, float, bool, bool]:
labels = {
label
for label in set(pose_scores.keys()) | set(context_scores.keys())
if not is_no_sex_position_label(label)
}
best_position = ""
best_score = 0.0
best_has_pose = False
best_has_context = False
best_pose_position = ""
best_pose_score = 0.0
best_pose_has_context = False
for label in labels:
pose_score = clamp01(pose_scores.get(label, 0.0))
context_score = clamp01(context_scores.get(label, 0.0))
has_pose = pose_score > 0
has_context = context_score > 0
score = 0.0
if has_pose:
score = pose_score
if has_context:
boost = clamp01(context_score * _POSITION_CONTEXT_BOOST_WEIGHT)
score = clamp01(1 - (1 - score) * (1 - boost))
elif context_score >= _POSITION_CONTEXT_MIN_SCORE:
score = min(_POSITION_CONTEXT_MAX_SCORE, context_score)
if score > best_score:
best_position = label
best_score = score
best_has_pose = has_pose
best_has_context = has_context
if has_pose and score > best_pose_score:
best_pose_position = label
best_pose_score = score
best_pose_has_context = has_context
if (
best_pose_position
and not best_has_pose
and best_score <= best_pose_score + _POSITION_CONTEXT_OVERRIDE_MARGIN
):
best_position = best_pose_position
best_score = best_pose_score
best_has_pose = True
best_has_context = best_pose_has_context
return best_position, clamp01(best_score), best_has_pose, best_has_context
def add_pose_pair_geometry_scores(scores: dict[str, float], persons: list[dict]) -> None:
def add(label: str, score: float) -> None:
combine_position_score(scores, label, score)
reliable_persons = reliable_pose_persons(persons)
if len(reliable_persons) < 2:
return
geometries = [pose_person_geometry(person) for person in reliable_persons]
for i, left in enumerate(geometries):
left_box = left["box"]
if not left_box:
continue
for right in geometries[i + 1:]:
right_box = right["box"]
if not right_box:
continue
gap = box_gap(left_box, right_box)
overlap = box_overlap_ratio(left_box, right_box)
horizontal_overlap = box_horizontal_overlap_ratio(left_box, right_box)
close = gap <= 0.12 or overlap >= 0.08
if not close:
continue
left_center = left["center"]
right_center = right["center"]
top = left if left_center[1] <= right_center[1] else right
bottom = right if top is left else left
top_above = top["center"][1] <= bottom["center"][1] - 0.055
strong_stack = horizontal_overlap >= 0.35 and (top_above or overlap >= 0.22)
top_has_rider_shape = (
top["straddling"]
or top["knees_wide"]
or (top["upright"] and top["knees_below_hips"])
)
if strong_stack and top_has_rider_shape:
add("cowgirl", 0.20)
add("reverse_cowgirl", 0.17)
if bottom["lying"]:
add("cowgirl", 0.12)
add("reverse_cowgirl", 0.10)
if top["straddling"]:
add("cowgirl", 0.08)
add("reverse_cowgirl", 0.06)
if strong_stack and bottom["lying"]:
add("missionary", 0.14)
if not top["straddling"]:
add("missionary", 0.08)
both_lying = left["lying"] and right["lying"]
same_level = abs(left_center[1] - right_center[1]) <= 0.15
side_by_side = abs(left_center[0] - right_center[0]) >= 0.10
if both_lying and (same_level or side_by_side):
add("spooning", 0.18)
if overlap >= 0.10:
add("prone_bone", 0.07)
left_bent = left["all_fours"] or left["bent_or_kneeling"]
right_bent = right["all_fours"] or right["bent_or_kneeling"]
if left_bent != right_bent:
other = right if left_bent else left
add("doggy", 0.16)
if other["upright"]:
add("doggy", 0.06)
add("standing_doggy", 0.07)
if left["lying"] or right["lying"]:
add("prone_bone", 0.06)
elif left["all_fours"] or right["all_fours"]:
add("doggy", 0.13)
if left["lying"] or right["lying"]:
add("prone_bone", 0.07)
def add_pose_scene_context_scores(scores: dict[str, float], prediction: dict, persons: list[dict]) -> None:
def add(label: str, score: float) -> None:
combine_position_score(scores, label, score)
reliable_persons = reliable_pose_persons(persons)
person_boxes = scene_person_boxes(prediction, reliable_persons)
person_count = max(len(person_boxes), len(reliable_persons))
pair = scene_person_pair_signals(person_boxes)
add_pose_pair_geometry_scores(scores, reliable_persons)
penis_boxes = boxes_by_label(prediction, "penis")
pussy_boxes = boxes_by_label(prediction, "pussy", "vagina", "vulva", "labia")
ass_boxes = boxes_by_label(prediction, "ass", "anus")
breast_boxes = boxes_by_label(prediction, "breasts")
tongue_boxes = boxes_by_label(prediction, "tongue")
toy_boxes = boxes_by_label(prediction, "dildo", "vibrator", "strapon", "buttplug")
has_penis = bool(penis_boxes)
has_pussy = bool(pussy_boxes)
has_ass = bool(ass_boxes)
has_breasts = bool(breast_boxes)
has_tongue = bool(tongue_boxes)
has_toy = bool(toy_boxes)
head_names = ["nose", "left_eye", "right_eye", "left_ear", "right_ear"]
hand_names = ["left_wrist", "right_wrist"]
hip_names = ["left_hip", "right_hip"]
head_near_penis = any_pose_keypoint_near_boxes(reliable_persons, head_names, penis_boxes, 0.09)
head_near_pussy = any_pose_keypoint_near_boxes(reliable_persons, head_names, pussy_boxes, 0.09)
head_near_ass = any_pose_keypoint_near_boxes(reliable_persons, head_names, ass_boxes, 0.09)
hand_near_penis = any_pose_keypoint_near_boxes(reliable_persons, hand_names, penis_boxes, 0.08)
hand_near_pussy = any_pose_keypoint_near_boxes(reliable_persons, hand_names, pussy_boxes, 0.08)
hand_near_toy = any_pose_keypoint_near_boxes(reliable_persons, hand_names, toy_boxes, 0.08)
hips_near_genitals = any_pose_keypoint_near_boxes(reliable_persons, hip_names, penis_boxes + pussy_boxes, 0.08)
if person_count >= 2:
for label, score in [
("missionary", 0.04),
("doggy", 0.04),
("cowgirl", 0.04),
("reverse_cowgirl", 0.04),
("standing_doggy", 0.04),
("spooning", 0.04),
("69", 0.03),
]:
add(label, score)
if pair["close"]:
for label in ["missionary", "doggy", "cowgirl", "spooning"]:
add(label, 0.04)
if pair["overlap"]:
add("missionary", 0.05)
add("cowgirl", 0.05)
add("prone_bone", 0.04)
if pair["horizontal"]:
add("spooning", 0.12)
add("prone_bone", 0.07)
if pair["vertical"]:
add("standing", 0.08)
add("standing_doggy", 0.10)
if pair["stacked"]:
add("missionary", 0.07)
add("cowgirl", 0.07)
add("reverse_cowgirl", 0.06)
add("facesitting", 0.05)
if has_penis and has_pussy:
for label, score in [
("missionary", 0.08),
("doggy", 0.08),
("cowgirl", 0.08),
("reverse_cowgirl", 0.07),
("prone_bone", 0.06),
("standing_doggy", 0.06),
("spooning", 0.05),
]:
add(label, score)
if any_boxes_near(penis_boxes, pussy_boxes, 0.09) or hips_near_genitals:
add("missionary", 0.08)
add("doggy", 0.08)
add("cowgirl", 0.08)
add("reverse_cowgirl", 0.07)
if has_penis and (head_near_penis or has_tongue):
add("blowjob", 0.16)
if head_near_penis:
add("blowjob", 0.10)
if has_pussy and (head_near_pussy or has_tongue or any_boxes_near(tongue_boxes, pussy_boxes, 0.08)):
add("cunnilingus", 0.16)
if head_near_pussy or any_boxes_near(tongue_boxes, pussy_boxes, 0.08):
add("cunnilingus", 0.10)
if has_penis and hand_near_penis:
add("handjob", 0.18)
if has_pussy and hand_near_pussy:
add("fingering", 0.18)
if has_penis and has_breasts and any_boxes_near(penis_boxes, breast_boxes, 0.10):
add("boobjob", 0.20)
if has_toy:
add("toy_play", 0.12)
if (
hand_near_toy
or any_boxes_near(toy_boxes, pussy_boxes, 0.10)
or any_boxes_near(toy_boxes, penis_boxes, 0.10)
or any_boxes_near(toy_boxes, ass_boxes, 0.10)
):
add("toy_play", 0.12)
if person_count >= 2 and (head_near_ass or head_near_pussy):
if has_ass:
add("facesitting", 0.14)
if has_pussy and has_penis:
add("69", 0.10)
if has_ass and has_pussy and pair["horizontal"]:
add("doggy", 0.07)
add("prone_bone", 0.07)
def best_hybrid_pose_scene_position(
prediction: dict,
position_persons: list[dict],
context_persons: list[dict],
) -> tuple[str, float, bool, bool]:
pose_scores: dict[str, float] = {}
context_scores: dict[str, float] = {}
reliable_position_persons = reliable_pose_persons(position_persons)
reliable_context_persons = reliable_pose_persons(context_persons)
if not reliable_context_persons:
reliable_context_persons = reliable_position_persons
for person in reliable_position_persons:
label = normalize_sex_position_label(person.get("label"))
score = float(person.get("score") or 0.0)
if is_no_sex_position_label(label) or label not in POSITION_LABELS:
continue
combine_position_score(pose_scores, label, score)
quality = pose_keypoint_quality(person)
if quality > 0:
combine_position_score(pose_scores, label, 0.04 * quality)
add_pose_scene_context_scores(context_scores, prediction, reliable_context_persons)
best_position, best_score, has_pose_signal, has_context_signal = fuse_hybrid_position_scores(
pose_scores,
context_scores,
)
if best_position and (has_pose_signal or best_score >= _POSITION_CONTEXT_MIN_SCORE):
return best_position, clamp01(best_score), has_pose_signal, has_context_signal
return "", 0.0, False, has_context_signal
def best_pose_scene_position(prediction: dict, persons: list[dict]) -> tuple[str, float, bool, bool]:
return best_hybrid_pose_scene_position(prediction, persons, persons)
def apply_pose_result_to_prediction(prediction: dict, result) -> dict:
persons = pose_persons_from_result(result)
if persons:
prediction["persons"] = persons
best_position, best_score_value, _has_pose_signal, has_context_signal = best_pose_scene_position(
prediction,
persons,
)
if best_position:
prediction["sexPosition"] = best_position
prediction["sexPositionScore"] = best_score_value
append_prediction_source(prediction, "yolo_pose")
if has_context_signal:
append_prediction_source(prediction, "box_context")
return prediction
def predict_pose_results(model, paths: list[str], imgsz: int, conf: float):
return model.predict(
source=paths,
imgsz=imgsz,
conf=conf,
batch=_BATCH,
device=_DEVICE or None,
half=_HALF,
verbose=False,
)
def has_reliable_pose_persons(persons: list[dict]) -> bool:
return bool(reliable_pose_persons(persons))
def apply_hybrid_pose_persons_to_prediction(
prediction: dict,
position_persons: list[dict],
context_persons: list[dict],
) -> None:
display_persons = context_persons or position_persons
if display_persons:
prediction["persons"] = display_persons
best_position, best_score_value, has_pose_signal, has_context_signal = best_hybrid_pose_scene_position(
prediction,
position_persons,
context_persons,
)
if best_position:
prediction["sexPosition"] = best_position
prediction["sexPositionScore"] = best_score_value
if position_persons or context_persons:
append_prediction_source(prediction, "yolo_pose")
if context_persons:
append_prediction_source(prediction, "base_pose")
if has_pose_signal:
append_prediction_source(prediction, "pose_position")
if has_context_signal:
append_prediction_source(prediction, "box_context")
def apply_pose_batch_to_predictions(paths: list[str], predictions: list[dict], imgsz: int) -> None:
global _POSE_MODEL_ERROR
global _BASE_POSE_MODEL_ERROR
position_persons_by_index: dict[int, list[dict]] = {}
context_persons_by_index: dict[int, list[dict]] = {}
current_pose_model = get_pose_model()
if current_pose_model is not None:
try:
pose_results = predict_pose_results(current_pose_model, paths, imgsz, _POSE_CONF)
for index, pose_result in enumerate(pose_results):
position_persons_by_index[index] = pose_persons_from_result(pose_result)
except Exception as exc:
_POSE_MODEL_ERROR = str(exc)
needs_base_pose = [
index
for index in range(len(paths))
if not has_reliable_pose_persons(position_persons_by_index.get(index, []))
]
current_base_pose_model = get_base_pose_model()
if current_base_pose_model is not None and needs_base_pose:
try:
base_paths = [paths[index] for index in needs_base_pose]
base_results = predict_pose_results(current_base_pose_model, base_paths, imgsz, _BASE_POSE_CONF)
for original_index, base_result in zip(needs_base_pose, base_results):
context_persons_by_index[original_index] = pose_persons_from_result(base_result)
except Exception as exc:
_BASE_POSE_MODEL_ERROR = str(exc)
for index, prediction in enumerate(predictions):
position_persons = position_persons_by_index.get(index, [])
context_persons = context_persons_by_index.get(index, [])
apply_hybrid_pose_persons_to_prediction(
prediction,
position_persons,
context_persons,
)
def pose_model_status() -> dict:
try:
expected = resolve_pose_model_path()
expected_text = str(expected) if expected else str(DEFAULT_POSE_MODEL_PATH)
exists = expected is not None
error = _POSE_MODEL_ERROR
except Exception as exc:
expected_text = str(DEFAULT_POSE_MODEL_PATH)
exists = False
error = str(exc)
return {
"poseModelAvailable": exists,
"poseModelLoaded": pose_model is not None,
"poseModel": _POSE_MODEL_PATH or (expected_text if exists else ""),
"poseModelError": error,
"expectedPoseModel": expected_text,
}
def base_pose_model_status() -> dict:
try:
expected = resolve_base_pose_model_path()
expected_text = str(expected) if expected else "yolo26n-pose.pt"
exists = expected is not None
error = _BASE_POSE_MODEL_ERROR
except Exception as exc:
expected_text = "yolo26n-pose.pt"
exists = False
error = str(exc)
return {
"basePoseModelAvailable": exists,
"basePoseModelLoaded": base_pose_model is not None,
"basePoseModel": _BASE_POSE_MODEL_PATH or (expected_text if exists else ""),
"basePoseModelError": error,
"expectedBasePoseModel": expected_text,
}
def videomae_model_status() -> dict:
try:
expected = resolve_videomae_model_path()
expected_text = str(expected) if expected else str(DEFAULT_VIDEOMAE_MODEL_PATH)
exists = expected is not None
error = _VIDEOMAE_MODEL_ERROR
except Exception as exc:
expected_text = str(DEFAULT_VIDEOMAE_MODEL_PATH)
exists = False
error = str(exc)
return {
"videoMAEModelAvailable": exists,
"videoMAEModelLoaded": videomae_model is not None,
"videoMAEModel": _VIDEOMAE_MODEL_PATH or (expected_text if exists else ""),
"videoMAEModelError": error,
"expectedVideoMAEModel": expected_text,
"videoMAEDevice": _VIDEOMAE_DEVICE_ACTIVE,
}
def detector_model_status(load: bool = False) -> dict:
current_model = get_model() if load else model
names = getattr(current_model, "names", {}) or {} if current_model is not None else {}
expected_text = str(DEFAULT_MODEL_PATH)
exists = False
error = _MODEL_ERROR
try:
expected = resolve_model_path()
expected_text = str(expected)
exists = True
except Exception as exc:
if not error:
error = str(exc)
return {
"modelAvailable": exists,
"modelLoaded": current_model is not None,
"model": _MODEL_PATH or (expected_text if exists else ""),
"modelError": error,
"expectedModel": expected_text,
"classCount": len(names),
"classes": list(names.values())[:80] if isinstance(names, dict) else names,
}
@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",
}
imgsz = int(req.imageSize or _IMGSZ or 640)
current_model = get_model()
if current_model is None:
predictions = [empty_prediction("detector_model_missing") for _ in paths]
if not req.detectorOnly:
apply_pose_batch_to_predictions(paths, predictions, imgsz)
return {
"ok": True,
"available": False,
"modelAvailable": False,
"predictions": predictions,
"modelError": _MODEL_ERROR,
"expectedModel": str(DEFAULT_MODEL_PATH),
}
if DETECTION_LABELS_PATH is None or _LABEL_ERROR:
return {
"ok": False,
"predictions": [],
"error": f"detection labels missing: {_LABEL_ERROR}",
}
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]
if not req.detectorOnly:
apply_pose_batch_to_predictions(paths, predictions, imgsz)
return {
"ok": True,
"available": True,
"modelAvailable": True,
"predictions": predictions,
}
except Exception as exc:
return {
"ok": False,
"predictions": [],
"error": str(exc),
}
@app.post("/predict-position-clips", dependencies=[Depends(require_ai_server_auth)])
def predict_position_clips(req: PredictPositionClipsRequest):
clips = [clip for clip in req.clips if clip.paths]
if not clips:
return {
"ok": True,
"available": False,
"predictions": [],
"error": "no clips supplied",
}
current_model, current_processor = get_videomae_components()
if current_model is None or current_processor is None:
return {
"ok": True,
"available": False,
"predictions": [],
"error": _VIDEOMAE_MODEL_ERROR or f"VideoMAE model not found: {DEFAULT_VIDEOMAE_MODEL_PATH}",
}
predictions = []
for clip in clips:
try:
predictions.append(predict_videomae_clip(clip, int(req.numFrames or _VIDEOMAE_NUM_FRAMES or 16)))
except Exception as exc:
predictions.append({
"time": float(clip.time or 0.0),
"start": float(clip.start or 0.0),
"end": float(clip.end or 0.0),
"sexPosition": NO_SEX_POSITION_LABEL,
"sexPositionScore": 0.0,
"source": "videomae_predict_failed",
"error": repr(exc),
"scores": [],
})
return {
"ok": True,
"available": True,
"predictions": predictions,
}
@app.get("/health", dependencies=[Depends(require_ai_server_auth)])
def health():
status_payload = {
"ok": True,
"ready": True,
"trainingRoot": str(TRAINING_ROOT),
"labelConfig": str(DETECTION_LABELS_PATH) if DETECTION_LABELS_PATH else "",
"labelError": _LABEL_ERROR,
}
status_payload.update(detector_model_status(load=False))
status_payload.update(pose_model_status())
status_payload.update(base_pose_model_status())
status_payload.update(videomae_model_status())
return status_payload
@app.post("/reload", dependencies=[Depends(require_ai_server_auth)])
def reload_model():
global model
global pose_model
global base_pose_model
global videomae_model
global videomae_processor
global _MODEL_PATH
global _MODEL_ERROR
global _POSE_MODEL_PATH
global _POSE_MODEL_ERROR
global _BASE_POSE_MODEL_PATH
global _BASE_POSE_MODEL_ERROR
global _VIDEOMAE_MODEL_PATH
global _VIDEOMAE_MODEL_ERROR
global _VIDEOMAE_DEVICE_ACTIVE
global DETECTION_LABELS_PATH
model = None
pose_model = None
base_pose_model = None
videomae_model = None
videomae_processor = None
_MODEL_PATH = ""
_MODEL_ERROR = ""
_POSE_MODEL_PATH = ""
_POSE_MODEL_ERROR = ""
_BASE_POSE_MODEL_PATH = ""
_BASE_POSE_MODEL_ERROR = ""
_VIDEOMAE_MODEL_PATH = ""
_VIDEOMAE_MODEL_ERROR = ""
_VIDEOMAE_DEVICE_ACTIVE = ""
DETECTION_LABELS_PATH = None
status_payload = {
"ok": True,
"ready": True,
"trainingRoot": str(TRAINING_ROOT),
"labelConfig": str(DETECTION_LABELS_PATH) if DETECTION_LABELS_PATH else "",
"labelError": _LABEL_ERROR,
}
status_payload.update(detector_model_status(load=True))
status_payload.update(pose_model_status())
status_payload.update(base_pose_model_status())
status_payload.update(videomae_model_status())
return status_payload
Reference in New Issue View Git Blame Copy Permalink