322 lines
8.5 KiB
Python
322 lines
8.5 KiB
Python
# backend/ml/predict_pose_model.py
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from ultralytics import YOLO
|
|
import torch
|
|
|
|
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",
|
|
]
|
|
|
|
BASE_MODEL_NAME = "yolo26n-pose.pt"
|
|
BASE_POSE_CONF = 0.10
|
|
POSE_KEYPOINT_MIN_CONFIDENCE = 0.20
|
|
POSE_RELIABLE_MIN_SCORE = 0.30
|
|
POSE_RELIABLE_MIN_KEYPOINTS = 6
|
|
POSE_RELIABLE_MIN_QUALITY = 0.45
|
|
|
|
|
|
def clamp01(value):
|
|
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):
|
|
try:
|
|
n = float(value)
|
|
except Exception:
|
|
return False
|
|
|
|
return math.isfinite(n) and 0.0 <= n <= 1.0
|
|
|
|
|
|
def pose_keypoint_stats(person):
|
|
keypoints = person.get("keypoints", []) or []
|
|
if not keypoints:
|
|
return 0, 0.0
|
|
|
|
visible = 0
|
|
total_conf = 0.0
|
|
|
|
for point in keypoints:
|
|
x = point.get("x", 0.0)
|
|
y = point.get("y", 0.0)
|
|
conf = float(point.get("conf") or 0.0)
|
|
|
|
if conf < POSE_KEYPOINT_MIN_CONFIDENCE 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 annotate_pose_person_quality(person):
|
|
visible, quality = pose_keypoint_stats(person)
|
|
score = clamp01(person.get("score", 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 existing_file(path):
|
|
try:
|
|
p = Path(path).expanduser().resolve()
|
|
if p.exists() and p.is_file() and p.stat().st_size > 0:
|
|
return p
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def base_model_candidates(root):
|
|
script_dir = Path(__file__).resolve().parent
|
|
|
|
return [
|
|
script_dir / BASE_MODEL_NAME,
|
|
script_dir.parent / BASE_MODEL_NAME,
|
|
Path.cwd() / BASE_MODEL_NAME,
|
|
root / BASE_MODEL_NAME,
|
|
root.parent / BASE_MODEL_NAME,
|
|
root.parent.parent / BASE_MODEL_NAME,
|
|
]
|
|
|
|
|
|
def resolve_model_path(root, requested):
|
|
if requested:
|
|
p = existing_file(requested)
|
|
if p:
|
|
return p, "yolo_pose_model"
|
|
return Path(requested).expanduser(), "pose_missing"
|
|
|
|
trained = root / "pose" / "model" / "best.pt"
|
|
p = existing_file(trained)
|
|
if p:
|
|
return p, "yolo_pose"
|
|
|
|
for candidate in base_model_candidates(root):
|
|
p = existing_file(candidate)
|
|
if p:
|
|
return p, "yolo26_pose_base"
|
|
|
|
return trained, "pose_missing"
|
|
|
|
|
|
def has_reliable_persons(persons):
|
|
return any(bool(person.get("reliable")) for person in persons or [])
|
|
|
|
|
|
def persons_from_results(results, img_w, img_h):
|
|
persons = []
|
|
|
|
if not results:
|
|
return persons
|
|
|
|
r = results[0]
|
|
names = r.names or {}
|
|
|
|
kpts_xyn = None
|
|
kpts_conf = None
|
|
|
|
if r.keypoints is not None:
|
|
try:
|
|
kpts_xyn = r.keypoints.xyn.cpu().numpy()
|
|
except Exception:
|
|
kpts_xyn = None
|
|
try:
|
|
kpts_conf = r.keypoints.conf.cpu().numpy()
|
|
except Exception:
|
|
kpts_conf = None
|
|
|
|
if r.boxes is None:
|
|
return persons
|
|
|
|
for i, b in enumerate(r.boxes):
|
|
score = float(b.conf[0].item())
|
|
label = ""
|
|
try:
|
|
cls_id = int(b.cls[0].item())
|
|
label = str(names.get(cls_id, cls_id)).strip().lower()
|
|
except Exception:
|
|
label = ""
|
|
|
|
x1, y1, x2, y2 = [float(v) for v in b.xyxy[0].tolist()]
|
|
x1 = max(0.0, min(float(img_w), x1))
|
|
y1 = max(0.0, min(float(img_h), y1))
|
|
x2 = max(0.0, min(float(img_w), x2))
|
|
y2 = max(0.0, min(float(img_h), y2))
|
|
|
|
if x2 <= x1 or y2 <= y1:
|
|
continue
|
|
|
|
box = {
|
|
"x": x1 / img_w,
|
|
"y": y1 / img_h,
|
|
"w": (x2 - x1) / img_w,
|
|
"h": (y2 - y1) / img_h,
|
|
}
|
|
|
|
keypoints = []
|
|
if kpts_xyn is not None and i < len(kpts_xyn):
|
|
person_kpts = kpts_xyn[i]
|
|
for ki, (kx, ky) in enumerate(person_kpts):
|
|
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])
|
|
|
|
name = KEYPOINT_NAMES[ki] if ki < len(KEYPOINT_NAMES) else str(ki)
|
|
keypoints.append({
|
|
"name": name,
|
|
"x": float(kx),
|
|
"y": float(ky),
|
|
"conf": kconf,
|
|
})
|
|
|
|
persons.append(annotate_pose_person_quality({
|
|
"label": label,
|
|
"score": score,
|
|
"box": box,
|
|
"keypoints": keypoints,
|
|
}))
|
|
|
|
return persons
|
|
|
|
|
|
def predict_persons(model_path, image_path, img_w, img_h, conf, imgsz, device):
|
|
model = YOLO(str(model_path))
|
|
results = model.predict(
|
|
source=str(image_path),
|
|
conf=float(conf),
|
|
imgsz=int(imgsz),
|
|
verbose=False,
|
|
device=device,
|
|
)
|
|
return persons_from_results(results, img_w, img_h)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--root", required=True)
|
|
parser.add_argument("--image", required=True)
|
|
parser.add_argument("--model", default="")
|
|
parser.add_argument("--conf", type=float, default=0.30)
|
|
parser.add_argument("--imgsz", type=int, default=640)
|
|
args = parser.parse_args()
|
|
|
|
root = Path(args.root)
|
|
image_path = Path(args.image)
|
|
|
|
model_path, model_source = resolve_model_path(root, args.model)
|
|
if not model_path.exists():
|
|
print(json.dumps({
|
|
"available": False,
|
|
"source": model_source,
|
|
"modelPath": str(model_path),
|
|
"persons": [],
|
|
}, ensure_ascii=False))
|
|
return
|
|
|
|
img = Image.open(image_path).convert("RGB")
|
|
img_w, img_h = img.size
|
|
|
|
device = 0 if torch.cuda.is_available() else "cpu"
|
|
effective_model_path = model_path
|
|
effective_model_source = model_source
|
|
effective_conf = float(args.conf)
|
|
fallback_error = ""
|
|
|
|
try:
|
|
persons = predict_persons(model_path, image_path, img_w, img_h, effective_conf, args.imgsz, device)
|
|
except Exception as e:
|
|
print(json.dumps({
|
|
"available": False,
|
|
"source": "pose_failed",
|
|
"modelPath": str(model_path),
|
|
"image": str(image_path),
|
|
"error": repr(e),
|
|
"persons": [],
|
|
}, ensure_ascii=False))
|
|
return
|
|
|
|
if not has_reliable_persons(persons) and model_source != "yolo26_pose_base":
|
|
for candidate in base_model_candidates(root):
|
|
base_model_path = existing_file(candidate)
|
|
if not base_model_path or base_model_path == model_path:
|
|
continue
|
|
|
|
try:
|
|
base_conf = min(float(args.conf), BASE_POSE_CONF)
|
|
base_persons = predict_persons(
|
|
base_model_path,
|
|
image_path,
|
|
img_w,
|
|
img_h,
|
|
base_conf,
|
|
args.imgsz,
|
|
device,
|
|
)
|
|
if base_persons:
|
|
persons = base_persons
|
|
effective_model_path = base_model_path
|
|
effective_model_source = "yolo26_pose_base_fallback"
|
|
effective_conf = base_conf
|
|
break
|
|
except Exception as e:
|
|
fallback_error = repr(e)
|
|
|
|
payload = {
|
|
"available": True,
|
|
"source": effective_model_source,
|
|
"modelPath": str(effective_model_path),
|
|
"image": str(image_path),
|
|
"conf": effective_conf,
|
|
"imgsz": int(args.imgsz),
|
|
"device": str(device),
|
|
"imageWidth": img_w,
|
|
"imageHeight": img_h,
|
|
"personCount": len(persons),
|
|
"persons": persons,
|
|
}
|
|
|
|
if fallback_error:
|
|
payload["fallbackError"] = fallback_error
|
|
|
|
print(json.dumps(payload, ensure_ascii=False))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|