739 lines
24 KiB
Python
Executable File
739 lines
24 KiB
Python
Executable File
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.13"
|
|
# dependencies = ["ultralytics", "torch", "numpy", "pillow", "bottle", "mediapipe"]
|
|
# ///
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import urllib.request
|
|
import webbrowser
|
|
from dataclasses import dataclass
|
|
from functools import cache
|
|
from pathlib import Path
|
|
from typing import NamedTuple, Optional
|
|
|
|
try:
|
|
import bottle
|
|
except Exception:
|
|
bottle = None
|
|
|
|
try:
|
|
from ultralytics import YOLO
|
|
import torch
|
|
import numpy as np
|
|
except Exception:
|
|
YOLO = None
|
|
torch = None
|
|
np = None
|
|
|
|
try:
|
|
from PIL import Image
|
|
except Exception:
|
|
Image = None
|
|
|
|
try:
|
|
import mediapipe as mp
|
|
except Exception:
|
|
mp = None
|
|
|
|
VISIBILITY_THRESH = 0.1
|
|
FACE_LANDMARK_VISIBILITY_THRESH = 0.5
|
|
FACING_DIRECTION_THRESH = 0.03
|
|
LOOKING_AT_CAMERA_THRESH = 0.015 # max |nose.x - eye_midpoint.x| for frontal face
|
|
HEAD_CROP_Y_THRESH = 0.35 # shoulders must be in top 35% of frame to classify as head-cropped
|
|
NOSE_TO_NOSE_THRESH = 0.15 # normalized image distance
|
|
BODY_INTERSECT_MARGIN = 0.05 # expand each person's bbox by this before overlap test
|
|
EYE_BLINK_THRESH = 0.40 # blendshape score above this → eye closed
|
|
|
|
|
|
class Coords(NamedTuple):
|
|
"""Represents normalized coordinates (0.0 to 1.0) and visibility for a single point."""
|
|
|
|
x: float
|
|
y: float
|
|
is_visible: bool
|
|
|
|
|
|
@dataclass
|
|
class PoseKeypoints:
|
|
"""Holds structured, normalized keypoint data for all 17 COCO points as direct fields."""
|
|
|
|
# 0
|
|
nose: Coords
|
|
# 1-4
|
|
left_eye: Coords
|
|
right_eye: Coords
|
|
left_ear: Coords
|
|
right_ear: Coords
|
|
# 5-6
|
|
left_shoulder: Coords
|
|
right_shoulder: Coords
|
|
# 7-10
|
|
left_elbow: Coords
|
|
right_elbow: Coords
|
|
left_wrist: Coords
|
|
right_wrist: Coords
|
|
# 11-12
|
|
left_hip: Coords
|
|
right_hip: Coords
|
|
# 13-16
|
|
left_knee: Coords
|
|
right_knee: Coords
|
|
left_ankle: Coords
|
|
right_ankle: Coords
|
|
|
|
def shoulder_midpoint(self) -> Coords:
|
|
l = self.left_shoulder
|
|
r = self.right_shoulder
|
|
if l.is_visible and r.is_visible:
|
|
return Coords(x=(l.x + r.x) / 2.0, y=(l.y + r.y) / 2.0, is_visible=True)
|
|
return Coords(x=(l.x + r.x) / 2.0, y=(l.y + r.y) / 2.0, is_visible=False)
|
|
|
|
|
|
@dataclass
|
|
class FaceLandmarks:
|
|
"""5-point face landmarks from the derronqi yolov8-face model."""
|
|
|
|
left_eye: Coords
|
|
right_eye: Coords
|
|
nose: Coords
|
|
left_mouth: Coords
|
|
right_mouth: Coords
|
|
|
|
|
|
def _get_coords(kp_xyc: "np.ndarray", idx: int) -> Coords:
|
|
"""Helper to safely extract Coords from the raw numpy array."""
|
|
x, y, conf = kp_xyc[idx]
|
|
is_visible = conf > VISIBILITY_THRESH
|
|
return Coords(x=x, y=y, is_visible=is_visible)
|
|
|
|
|
|
def _extract_keypoints(kp_xyc: "np.ndarray") -> PoseKeypoints:
|
|
"""Extracts all 17 COCO normalized keypoints and populates the PoseKeypoints dataclass directly."""
|
|
return PoseKeypoints(
|
|
nose=_get_coords(kp_xyc, 0),
|
|
left_eye=_get_coords(kp_xyc, 1),
|
|
right_eye=_get_coords(kp_xyc, 2),
|
|
left_ear=_get_coords(kp_xyc, 3),
|
|
right_ear=_get_coords(kp_xyc, 4),
|
|
left_shoulder=_get_coords(kp_xyc, 5),
|
|
right_shoulder=_get_coords(kp_xyc, 6),
|
|
left_elbow=_get_coords(kp_xyc, 7),
|
|
right_elbow=_get_coords(kp_xyc, 8),
|
|
left_wrist=_get_coords(kp_xyc, 9),
|
|
right_wrist=_get_coords(kp_xyc, 10),
|
|
left_hip=_get_coords(kp_xyc, 11),
|
|
right_hip=_get_coords(kp_xyc, 12),
|
|
left_knee=_get_coords(kp_xyc, 13),
|
|
right_knee=_get_coords(kp_xyc, 14),
|
|
left_ankle=_get_coords(kp_xyc, 15),
|
|
right_ankle=_get_coords(kp_xyc, 16),
|
|
)
|
|
|
|
|
|
_MODEL_URLS = {
|
|
"yolov11n-face.pt": "https://huggingface.co/AdamCodd/YOLOv11n-face-detection/resolve/main/model.pt",
|
|
"yolo11s-pose.pt": "https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11s-pose.pt",
|
|
"yolov8n-face-derronqi.pt": "https://huggingface.co/junjiang/GestureFace/resolve/main/yolov8n-face.pt",
|
|
"face_landmarker.task": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task",
|
|
}
|
|
|
|
|
|
def _ensure_model(filename: str) -> Path:
|
|
dest = Path(__file__).parent / filename
|
|
if not dest.exists():
|
|
url = _MODEL_URLS[filename]
|
|
logging.warning(f"Downloading {filename} from {url} ...")
|
|
urllib.request.urlretrieve(url, dest)
|
|
logging.warning(f"Saved {filename}")
|
|
return dest
|
|
|
|
|
|
def _best_device() -> str:
|
|
try:
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
return "mps"
|
|
except Exception:
|
|
pass
|
|
return "cpu"
|
|
|
|
|
|
@cache
|
|
def _face_detector():
|
|
if YOLO is None or torch is None:
|
|
logging.error("YOLO/Torch dependencies are missing.")
|
|
return None
|
|
try:
|
|
model = YOLO(_ensure_model("yolov11n-face.pt"))
|
|
model.to(_best_device())
|
|
return model
|
|
except Exception:
|
|
logging.exception("Face detector model initialization failed.")
|
|
return None
|
|
|
|
|
|
@cache
|
|
def _pose_detector():
|
|
if YOLO is None or torch is None:
|
|
logging.error("YOLO/Torch dependencies are missing.")
|
|
return None
|
|
try:
|
|
model = YOLO(_ensure_model("yolo11s-pose.pt"))
|
|
model.to(_best_device())
|
|
return model
|
|
except Exception:
|
|
logging.exception("Pose detector model initialization failed.")
|
|
return None
|
|
|
|
|
|
@cache
|
|
def _face_landmark_detector():
|
|
if YOLO is None or torch is None:
|
|
logging.error("YOLO/Torch dependencies are missing.")
|
|
return None
|
|
try:
|
|
model = YOLO(_ensure_model("yolov8n-face-derronqi.pt"))
|
|
model.to(_best_device())
|
|
return model
|
|
except Exception:
|
|
logging.exception("Face landmark detector model initialization failed.")
|
|
return None
|
|
|
|
|
|
@cache
|
|
def _face_landmarker():
|
|
if mp is None:
|
|
logging.error("mediapipe dependency is missing.")
|
|
return None
|
|
try:
|
|
from mediapipe.tasks import python as mp_python
|
|
from mediapipe.tasks.python import vision as mp_vision
|
|
base_options = mp_python.BaseOptions(
|
|
model_asset_path=str(_ensure_model("face_landmarker.task"))
|
|
)
|
|
options = mp_vision.FaceLandmarkerOptions(
|
|
base_options=base_options,
|
|
output_face_blendshapes=True,
|
|
running_mode=mp_vision.RunningMode.IMAGE,
|
|
num_faces=10,
|
|
)
|
|
return mp_vision.FaceLandmarker.create_from_options(options)
|
|
except Exception:
|
|
logging.exception("FaceLandmarker initialization failed.")
|
|
return None
|
|
|
|
|
|
def _run_face_landmarker_model(image_path: Path):
|
|
detector = _face_landmarker()
|
|
if detector is None or Image is None:
|
|
return None
|
|
try:
|
|
img = Image.open(image_path).convert("RGB")
|
|
arr = np.asarray(img)
|
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=arr)
|
|
return detector.detect(mp_image)
|
|
except Exception:
|
|
logging.exception(f"FaceLandmarker failed for {image_path}.")
|
|
return None
|
|
|
|
|
|
def check_eyes_closed(face_landmarker_result) -> bool:
|
|
if face_landmarker_result is None:
|
|
return False
|
|
if not face_landmarker_result.face_blendshapes:
|
|
return False
|
|
for face_blendshapes in face_landmarker_result.face_blendshapes:
|
|
scores = {b.category_name: b.score for b in face_blendshapes}
|
|
if scores.get("eyeBlinkLeft", 0.0) > EYE_BLINK_THRESH or \
|
|
scores.get("eyeBlinkRight", 0.0) > EYE_BLINK_THRESH:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _run_face_model(image_path: Path) -> list:
|
|
model = _face_detector()
|
|
if model is None:
|
|
return []
|
|
try:
|
|
return model(str(image_path), conf=0.5, iou=0.5, verbose=False)
|
|
except Exception:
|
|
logging.exception(f"Face detection failed for {image_path}.")
|
|
return []
|
|
|
|
|
|
def _run_face_landmark_model(image_path: Path) -> list:
|
|
model = _face_landmark_detector()
|
|
if model is None:
|
|
return []
|
|
try:
|
|
return model(str(image_path), conf=0.25, iou=0.5, verbose=False)
|
|
except Exception:
|
|
logging.exception(f"Face landmark detection failed for {image_path}.")
|
|
return []
|
|
|
|
|
|
def _run_pose_model(image_path: Path) -> list:
|
|
model = _pose_detector()
|
|
if model is None:
|
|
return []
|
|
try:
|
|
return model(str(image_path), conf=0.35, iou=0.5, verbose=False)
|
|
except Exception:
|
|
logging.exception(f"Pose detection failed for {image_path}.")
|
|
return []
|
|
|
|
|
|
def count_faces(results: list) -> int:
|
|
try:
|
|
return len(results[0].boxes)
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def detect_poses(results: list) -> list[PoseKeypoints]:
|
|
try:
|
|
if not results or results[0].keypoints is None:
|
|
return []
|
|
kps_norm_xy = results[0].keypoints.xyn.cpu().numpy()
|
|
kps_conf = results[0].keypoints.conf.cpu().numpy()
|
|
all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2)
|
|
return [_extract_keypoints(kp_xyc) for kp_xyc in all_keypoints_xyc]
|
|
except Exception:
|
|
logging.exception("Pose keypoint extraction failed.")
|
|
return []
|
|
|
|
|
|
def detect_face_landmarks(results: list) -> list[FaceLandmarks]:
|
|
try:
|
|
if not results or results[0].keypoints is None:
|
|
return []
|
|
kps_norm_xy = results[0].keypoints.xyn.cpu().numpy()
|
|
kps_conf = results[0].keypoints.conf.cpu().numpy()
|
|
all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2)
|
|
landmarks = []
|
|
for kp_xyc in all_keypoints_xyc:
|
|
def _lm(idx, arr=kp_xyc):
|
|
x, y, conf = arr[idx]
|
|
return Coords(x=x, y=y, is_visible=conf > FACE_LANDMARK_VISIBILITY_THRESH)
|
|
landmarks.append(FaceLandmarks(
|
|
left_eye=_lm(0),
|
|
right_eye=_lm(1),
|
|
nose=_lm(2),
|
|
left_mouth=_lm(3),
|
|
right_mouth=_lm(4),
|
|
))
|
|
return landmarks
|
|
except Exception:
|
|
logging.exception("Face landmark extraction failed.")
|
|
return []
|
|
|
|
|
|
def _face_all_invisible(kps: PoseKeypoints) -> bool:
|
|
return (
|
|
not kps.nose.is_visible
|
|
and not kps.left_eye.is_visible
|
|
and not kps.right_eye.is_visible
|
|
)
|
|
|
|
|
|
def is_turned_back(kps: PoseKeypoints) -> bool:
|
|
shoulder_visible = kps.left_shoulder.is_visible or kps.right_shoulder.is_visible
|
|
return _face_all_invisible(kps) and shoulder_visible and not is_eyes_cropped_out(kps)
|
|
|
|
|
|
def is_eyes_cropped_out(kps: PoseKeypoints) -> bool:
|
|
"""True when face is not visible but shoulders are near the top of the frame,
|
|
indicating the head is above the image boundary."""
|
|
shoulder_visible = kps.left_shoulder.is_visible or kps.right_shoulder.is_visible
|
|
if not _face_all_invisible(kps) or not shoulder_visible:
|
|
return False
|
|
# Use the topmost (lowest y) visible shoulder
|
|
ys = [kp.y for kp in (kps.left_shoulder, kps.right_shoulder) if kp.is_visible]
|
|
return min(ys) < HEAD_CROP_Y_THRESH
|
|
|
|
|
|
def get_facing_x_direction(kps: PoseKeypoints) -> Optional[float]:
|
|
shoulder_mid = kps.shoulder_midpoint()
|
|
if not kps.nose.is_visible or not shoulder_mid.is_visible:
|
|
return None
|
|
return kps.nose.x - shoulder_mid.x
|
|
|
|
|
|
def get_face_yaw(kps: FaceLandmarks) -> Optional[float]:
|
|
"""Nose x offset from eye midpoint. ~0 = frontal, positive = turned right, negative = turned left."""
|
|
if not kps.left_eye.is_visible or not kps.right_eye.is_visible or not kps.nose.is_visible:
|
|
return None
|
|
eye_mid_x = (kps.left_eye.x + kps.right_eye.x) / 2.0
|
|
return kps.nose.x - eye_mid_x
|
|
|
|
|
|
def check_looking_at_camera(all_landmarks: list[FaceLandmarks]) -> bool:
|
|
return any(
|
|
(yaw := get_face_yaw(lm)) is not None and abs(yaw) < LOOKING_AT_CAMERA_THRESH
|
|
for lm in all_landmarks
|
|
)
|
|
|
|
|
|
def check_facing_each_other(all_kps: list[PoseKeypoints]) -> bool:
|
|
classifiable = []
|
|
for kps in all_kps:
|
|
delta = get_facing_x_direction(kps)
|
|
if delta is not None:
|
|
shoulder_mid = kps.shoulder_midpoint()
|
|
classifiable.append((shoulder_mid.x, delta))
|
|
|
|
if len(classifiable) < 2:
|
|
return False
|
|
|
|
classifiable.sort(key=lambda t: t[0])
|
|
|
|
for i in range(len(classifiable)):
|
|
for j in range(i + 1, len(classifiable)):
|
|
left_delta = classifiable[i][1]
|
|
right_delta = classifiable[j][1]
|
|
if left_delta > FACING_DIRECTION_THRESH and right_delta < -FACING_DIRECTION_THRESH:
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def detect_person_boxes(results: list) -> list[tuple[float, float, float, float]]:
|
|
try:
|
|
if not results or results[0].boxes is None:
|
|
return []
|
|
return [tuple(box) for box in results[0].boxes.xyxyn.cpu().tolist()]
|
|
except Exception:
|
|
logging.exception("Person box extraction failed.")
|
|
return []
|
|
|
|
|
|
def _face_boxes_to_body_boxes(
|
|
face_boxes: list[tuple[float, float, float, float]]
|
|
) -> list[tuple[float, float, float, float]]:
|
|
out = []
|
|
for x1, y1, x2, y2 in face_boxes:
|
|
fw = x2 - x1
|
|
fh = y2 - y1
|
|
bx1 = max(0.0, x1 - 0.3 * fw)
|
|
bx2 = min(1.0, x2 + 0.3 * fw)
|
|
by1 = y1
|
|
by2 = min(1.0, y2 + 2.5 * fh) # extend ~2.5 face-heights downward
|
|
out.append((bx1, by1, bx2, by2))
|
|
return out
|
|
|
|
|
|
def check_bodies_intersecting(person_boxes: list[tuple[float, float, float, float]]) -> bool:
|
|
m = BODY_INTERSECT_MARGIN
|
|
for i in range(len(person_boxes)):
|
|
for j in range(i + 1, len(person_boxes)):
|
|
ax1, ay1, ax2, ay2 = person_boxes[i]
|
|
bx1, by1, bx2, by2 = person_boxes[j]
|
|
if ax1 - m < bx2 and ax2 + m > bx1 and ay1 - m < by2 and ay2 + m > by1:
|
|
return True
|
|
return False
|
|
|
|
|
|
def check_nose_to_nose(all_kps: list[PoseKeypoints]) -> bool:
|
|
visible = [(kps.nose.x, kps.nose.y) for kps in all_kps if kps.nose.is_visible]
|
|
for i in range(len(visible)):
|
|
for j in range(i + 1, len(visible)):
|
|
dx = visible[i][0] - visible[j][0]
|
|
dy = visible[i][1] - visible[j][1]
|
|
if (dx*dx + dy*dy) ** 0.5 < NOSE_TO_NOSE_THRESH:
|
|
return True
|
|
return False
|
|
|
|
|
|
def classify_image(image_path: Path, debug_dir: Optional[Path] = None) -> dict:
|
|
face_results = _run_face_model(image_path)
|
|
pose_results = _run_pose_model(image_path)
|
|
face_landmark_results = _run_face_landmark_model(image_path)
|
|
face_mesh_result = _run_face_landmarker_model(image_path)
|
|
|
|
total_faces = count_faces(face_results)
|
|
all_kps = detect_poses(pose_results)
|
|
face_landmarks = detect_face_landmarks(face_landmark_results)
|
|
turned_back = any(is_turned_back(kps) for kps in all_kps)
|
|
facing_each_other = check_facing_each_other(all_kps)
|
|
eyes_cropped_out = any(is_eyes_cropped_out(kps) for kps in all_kps)
|
|
nose_to_nose = check_nose_to_nose(all_kps)
|
|
person_boxes = detect_person_boxes(pose_results)
|
|
if len(person_boxes) < 2:
|
|
raw_face_boxes = []
|
|
try:
|
|
if face_results and face_results[0].boxes is not None:
|
|
raw_face_boxes = [tuple(b) for b in face_results[0].boxes.xyxyn.cpu().tolist()]
|
|
except Exception:
|
|
pass
|
|
person_boxes = _face_boxes_to_body_boxes(raw_face_boxes)
|
|
bodies_intersecting = check_bodies_intersecting(person_boxes)
|
|
looking_at_camera = check_looking_at_camera(face_landmarks)
|
|
|
|
if debug_dir:
|
|
debug_dir.mkdir(parents=True, exist_ok=True)
|
|
for label, results in [("face", face_results), ("pose", pose_results)]:
|
|
if results and results[0].boxes is not None:
|
|
annotated = results[0].plot()
|
|
img = Image.fromarray(annotated[..., ::-1])
|
|
img.save(debug_dir / f"{image_path.stem}_{label}.jpg")
|
|
|
|
return {
|
|
"image_path": str(image_path),
|
|
"total_faces": total_faces,
|
|
"facing_each_other": facing_each_other,
|
|
"turned_back": turned_back,
|
|
"eyes_cropped_out": eyes_cropped_out,
|
|
"nose_to_nose": nose_to_nose,
|
|
"bodies_intersecting": bodies_intersecting,
|
|
"looking_at_camera": looking_at_camera,
|
|
"eyes_closed": check_eyes_closed(face_mesh_result),
|
|
}
|
|
|
|
|
|
HTML_PAGE = """<!DOCTYPE html>
|
|
<html lang="en">
|
|
<head>
|
|
<meta charset="UTF-8">
|
|
<title>classify_image debug</title>
|
|
<style>
|
|
body { background: #111; color: #eee; font-family: monospace; margin: 0; padding: 16px; transition: outline 0.1s; }
|
|
body.over { outline: 3px dashed #4f4; outline-offset: -6px; }
|
|
#drop {
|
|
border: 2px dashed #555; border-radius: 8px; padding: 40px;
|
|
text-align: center; color: #888; margin-bottom: 20px;
|
|
transition: border-color 0.1s, color 0.1s;
|
|
}
|
|
body.over #drop { border-color: #4f4; color: #4f4; }
|
|
#results { display: flex; flex-wrap: wrap; gap: 20px; }
|
|
.card {
|
|
display: flex; flex-direction: row; align-items: flex-start;
|
|
max-width: 960px; background: #1a1a1a; border-radius: 8px; padding: 12px; gap: 16px;
|
|
}
|
|
.img-wrap { position: relative; display: inline-block; flex-shrink: 0; }
|
|
.img-wrap img { display: block; max-width: 560px; max-height: 600px; border-radius: 4px; }
|
|
canvas { position: absolute; top: 0; left: 0; pointer-events: none; }
|
|
pre.data { font-size: 12px; margin: 0; overflow: auto; white-space: pre-wrap; color: #afa; }
|
|
.spinner { color: #888; padding: 20px; }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div id="drop">Drop images here to classify</div>
|
|
<div id="results"></div>
|
|
<script>
|
|
const EDGES = [
|
|
[0,1],[0,2],[1,3],[2,4],
|
|
[5,7],[7,9],[6,8],[8,10],
|
|
[5,6],[5,11],[6,12],[11,12],
|
|
[11,13],[13,15],[12,14],[14,16]
|
|
];
|
|
const COLORS = ["#f55","#5af","#ff5","#5f5","#f5f","#fa5","#5ff"];
|
|
|
|
const drop = document.getElementById("drop");
|
|
const results = document.getElementById("results");
|
|
|
|
let dragCounter = 0;
|
|
document.addEventListener("dragenter", e => { e.preventDefault(); if (++dragCounter === 1) document.body.classList.add("over"); });
|
|
document.addEventListener("dragleave", () => { if (--dragCounter === 0) document.body.classList.remove("over"); });
|
|
document.addEventListener("dragover", e => e.preventDefault());
|
|
document.addEventListener("drop", e => {
|
|
e.preventDefault();
|
|
dragCounter = 0;
|
|
document.body.classList.remove("over");
|
|
for (const file of e.dataTransfer.files) processFile(file);
|
|
});
|
|
|
|
function processFile(file) {
|
|
const card = document.createElement("div");
|
|
card.className = "card";
|
|
card.innerHTML = '<div class="spinner">Processing...</div>';
|
|
results.prepend(card);
|
|
|
|
const fd = new FormData();
|
|
fd.append("image", file, file.name);
|
|
|
|
fetch("/classify", { method: "POST", body: fd })
|
|
.then(r => r.json())
|
|
.then(data => {
|
|
const wrap = document.createElement("div");
|
|
wrap.className = "img-wrap";
|
|
const img = document.createElement("img");
|
|
const canvas = document.createElement("canvas");
|
|
wrap.appendChild(img);
|
|
wrap.appendChild(canvas);
|
|
|
|
const pre = document.createElement("pre");
|
|
pre.className = "data";
|
|
pre.textContent = JSON.stringify(data.classification, null, 2);
|
|
|
|
card.innerHTML = "";
|
|
card.appendChild(wrap);
|
|
card.appendChild(pre);
|
|
|
|
img.onload = () => {
|
|
canvas.width = img.offsetWidth;
|
|
canvas.height = img.offsetHeight;
|
|
const ctx = canvas.getContext("2d");
|
|
drawOverlays(ctx, data.detections, img.offsetWidth, img.offsetHeight);
|
|
};
|
|
img.src = URL.createObjectURL(file);
|
|
})
|
|
.catch(err => { card.innerHTML = '<div class="spinner">Error: ' + err + '</div>'; });
|
|
}
|
|
|
|
function drawOverlays(ctx, detections, W, H) {
|
|
ctx.lineWidth = 2;
|
|
ctx.strokeStyle = "#0f0";
|
|
for (const [x1n, y1n, x2n, y2n] of detections.faces) {
|
|
ctx.strokeRect(x1n * W, y1n * H, (x2n - x1n) * W, (y2n - y1n) * H);
|
|
}
|
|
detections.poses.forEach((kps, pi) => {
|
|
const col = COLORS[pi % COLORS.length];
|
|
ctx.strokeStyle = col;
|
|
ctx.fillStyle = col;
|
|
for (const [a, b] of EDGES) {
|
|
const ka = kps[a], kb = kps[b];
|
|
if (ka.v && kb.v) {
|
|
ctx.beginPath();
|
|
ctx.moveTo(ka.x * W, ka.y * H);
|
|
ctx.lineTo(kb.x * W, kb.y * H);
|
|
ctx.stroke();
|
|
}
|
|
}
|
|
for (const kp of kps) {
|
|
if (kp.v) {
|
|
ctx.beginPath();
|
|
ctx.arc(kp.x * W, kp.y * H, 4, 0, 2 * Math.PI);
|
|
ctx.fill();
|
|
}
|
|
}
|
|
});
|
|
}
|
|
</script>
|
|
</body>
|
|
</html>"""
|
|
|
|
|
|
def _extract_web_data(face_results: list, pose_results: list) -> dict:
|
|
faces = []
|
|
try:
|
|
if face_results and face_results[0].boxes is not None:
|
|
faces = face_results[0].boxes.xyxyn.cpu().tolist()
|
|
except Exception:
|
|
pass
|
|
|
|
poses = []
|
|
try:
|
|
if pose_results and pose_results[0].keypoints is not None:
|
|
kps_norm_xy = pose_results[0].keypoints.xyn.cpu().numpy()
|
|
kps_conf = pose_results[0].keypoints.conf.cpu().numpy()
|
|
all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2)
|
|
for person_kps in all_keypoints_xyc:
|
|
poses.append([
|
|
{"x": float(kp[0]), "y": float(kp[1]), "v": bool(kp[2] > VISIBILITY_THRESH)}
|
|
for kp in person_kps
|
|
])
|
|
except Exception:
|
|
pass
|
|
|
|
return {"faces": faces, "poses": poses}
|
|
|
|
|
|
def _classify_route():
|
|
upload = bottle.request.files.get("image")
|
|
suffix = Path(upload.filename).suffix or ".jpg"
|
|
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
|
|
tmp = Path(f.name)
|
|
upload.save(f)
|
|
try:
|
|
face_results = _run_face_model(tmp)
|
|
pose_results = _run_pose_model(tmp)
|
|
face_landmark_results = _run_face_landmark_model(tmp)
|
|
face_mesh_result = _run_face_landmarker_model(tmp)
|
|
all_kps = detect_poses(pose_results)
|
|
face_landmarks = detect_face_landmarks(face_landmark_results)
|
|
person_boxes = detect_person_boxes(pose_results)
|
|
if len(person_boxes) < 2:
|
|
raw_face_boxes = []
|
|
try:
|
|
if face_results and face_results[0].boxes is not None:
|
|
raw_face_boxes = [tuple(b) for b in face_results[0].boxes.xyxyn.cpu().tolist()]
|
|
except Exception:
|
|
pass
|
|
person_boxes = _face_boxes_to_body_boxes(raw_face_boxes)
|
|
classification = {
|
|
"image": upload.filename,
|
|
"total_faces": count_faces(face_results),
|
|
"facing_each_other": check_facing_each_other(all_kps),
|
|
"turned_back": any(is_turned_back(kps) for kps in all_kps),
|
|
"eyes_cropped_out": any(is_eyes_cropped_out(kps) for kps in all_kps),
|
|
"nose_to_nose": check_nose_to_nose(all_kps),
|
|
"bodies_intersecting": check_bodies_intersecting(person_boxes),
|
|
"looking_at_camera": check_looking_at_camera(face_landmarks),
|
|
"eyes_closed": check_eyes_closed(face_mesh_result),
|
|
}
|
|
detections = _extract_web_data(face_results, pose_results)
|
|
return bottle.HTTPResponse(
|
|
json.dumps({"classification": classification, "detections": detections}),
|
|
content_type="application/json",
|
|
)
|
|
finally:
|
|
tmp.unlink(missing_ok=True)
|
|
|
|
|
|
def web_main():
|
|
if bottle is None:
|
|
print("bottle not installed. Run: pip install bottle", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
app = bottle.Bottle()
|
|
|
|
@app.get("/")
|
|
def index():
|
|
return HTML_PAGE
|
|
|
|
@app.post("/classify")
|
|
def classify_route():
|
|
return _classify_route()
|
|
|
|
port = 7777
|
|
threading.Timer(0.5, lambda: webbrowser.open(f"http://localhost:{port}")).start()
|
|
bottle.run(app, host="localhost", port=port, quiet=True)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Classify images for face-related attributes.")
|
|
parser.add_argument("image_paths", nargs="*", type=Path, help="Path(s) to image(s)")
|
|
parser.add_argument("--debug", action="store_true", help="Save annotated debug images to _debug/ subdirectory")
|
|
parser.add_argument("--web", action="store_true", help="Start debug web server")
|
|
args = parser.parse_args()
|
|
|
|
logging.basicConfig(level=logging.WARNING, stream=sys.stderr)
|
|
|
|
if args.web:
|
|
web_main()
|
|
return
|
|
|
|
if not args.image_paths:
|
|
parser.print_usage(sys.stderr)
|
|
sys.exit(1)
|
|
|
|
for image_path in args.image_paths:
|
|
try:
|
|
debug_dir = (image_path.parent / "_debug") if args.debug else None
|
|
result = classify_image(image_path, debug_dir=debug_dir)
|
|
print(json.dumps(result), flush=True)
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except Exception:
|
|
logging.exception(f"Error processing {image_path}.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|