# 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" 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, 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 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 try: model = YOLO(str(model_path)) except Exception as e: print(json.dumps({ "available": False, "source": "pose_load_failed", "modelPath": str(model_path), "error": repr(e), "persons": [], }, ensure_ascii=False)) return device = 0 if torch.cuda.is_available() else "cpu" try: results = model.predict( source=str(image_path), conf=float(args.conf), imgsz=int(args.imgsz), verbose=False, device=device, ) except Exception as e: print(json.dumps({ "available": False, "source": "pose_predict_failed", "modelPath": str(model_path), "image": str(image_path), "error": repr(e), "persons": [], }, ensure_ascii=False)) return persons = [] if results: 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 not None: 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, })) print(json.dumps({ "available": True, "source": model_source, "modelPath": str(model_path), "image": str(image_path), "conf": float(args.conf), "imgsz": int(args.imgsz), "device": str(device), "imageWidth": img_w, "imageHeight": img_h, "personCount": len(persons), "persons": persons, }, ensure_ascii=False)) if __name__ == "__main__": main()