619 lines
22 KiB
Python
Executable File
619 lines
22 KiB
Python
Executable File
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.13"
|
|
# dependencies = ["ultralytics", "torch", "opencv-python", "numpy", "pillow"]
|
|
# ///
|
|
import argparse
|
|
from functools import cache
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import List, Tuple, NamedTuple, Optional
|
|
from dataclasses import dataclass
|
|
import math
|
|
from itertools import combinations
|
|
|
|
# --- Heavyweight imports for detection ---
|
|
try:
|
|
from ultralytics import YOLO
|
|
import torch
|
|
import numpy as np
|
|
except Exception: # pragma: no cover - optional
|
|
YOLO = None
|
|
torch = None
|
|
np = None
|
|
|
|
try:
|
|
import cv2
|
|
except Exception: # pragma: no cover - optional
|
|
cv2 = None
|
|
|
|
try:
|
|
from PIL import Image
|
|
except Exception:
|
|
Image = None
|
|
|
|
# --- Constants ---
|
|
# Tolerance for centering/pose checking (Normalized 0-1 space)
|
|
# Minimum confidence/visibility score for a keypoint to be used
|
|
VISIBILITY_THRESH = 0.1
|
|
|
|
# --- Structured Coordinate and Keypoint Types ---
|
|
|
|
|
|
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:
|
|
"""Return midpoint between left and right shoulder if available."""
|
|
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)
|
|
|
|
def hip_midpoint(self) -> Coords:
|
|
"""Return midpoint between left and right hip if available."""
|
|
l = self.left_hip
|
|
r = self.right_hip
|
|
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)
|
|
|
|
def eye_midpoint(self) -> Coords:
|
|
"""Return midpoint between left and right eye if available."""
|
|
l = self.left_eye
|
|
r = self.right_eye
|
|
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)
|
|
|
|
|
|
# --- Dataclasses for Structured Output ---
|
|
|
|
|
|
@dataclass
|
|
class PoseDetectionResult:
|
|
"""Encapsulates results from the pose detection strategy."""
|
|
|
|
boxes: np.ndarray
|
|
keypoints_xyc: np.ndarray
|
|
# List of all detected people's keypoints
|
|
all_pose_kps: List[PoseKeypoints]
|
|
|
|
# --- Centering logic moved into a method (CLEANED) ---
|
|
|
|
@dataclass
|
|
class CenterResult:
|
|
is_centered: bool
|
|
reason: str
|
|
coords: Optional[Coords]
|
|
threshold: float
|
|
|
|
def is_centered(self, center_threshold: float) -> "PoseDetectionResult.CenterResult":
|
|
"""
|
|
Checks if any person's core (nose, shoulder midpoint, or hip midpoint)
|
|
is horizontally centered within the image based on the threshold.
|
|
|
|
Returns: (is_centered, centered_by_point_name, centering_point_coords)
|
|
"""
|
|
cx = 0.5
|
|
band_min = cx - center_threshold
|
|
band_max = cx + center_threshold
|
|
|
|
def _is_in_band(c: Coords) -> bool:
|
|
return c.is_visible and (band_min <= c.x <= band_max)
|
|
|
|
for kps in self.all_pose_kps:
|
|
nose = kps.nose
|
|
shoulder_mid = kps.shoulder_midpoint()
|
|
hip_mid = kps.hip_midpoint()
|
|
eye_mid = kps.eye_midpoint()
|
|
|
|
important_features = [nose, shoulder_mid, eye_mid]
|
|
visible_features = [f for f in important_features if f.is_visible]
|
|
|
|
if not visible_features:
|
|
continue
|
|
|
|
all_centered = all(_is_in_band(f) for f in visible_features)
|
|
if all_centered:
|
|
# Representative point: prefer shoulder, then nose, then hip, then eyes
|
|
pref = None
|
|
for f in (shoulder_mid, nose, hip_mid, eye_mid):
|
|
if f.is_visible:
|
|
pref = f
|
|
break
|
|
if pref is None:
|
|
pref = visible_features[0]
|
|
return PoseDetectionResult.CenterResult(is_centered=True, reason="multiple", coords=pref, threshold=center_threshold)
|
|
|
|
return PoseDetectionResult.CenterResult(is_centered=False, reason="none", coords=None, threshold=center_threshold)
|
|
|
|
def is_torso_centered(self, center_threshold: float) -> Tuple[bool, str, Optional[Coords]]:
|
|
"""
|
|
Determines if the torso (the line passing through the shoulder midpoint
|
|
and the hip midpoint) crosses the central vertical band of the image.
|
|
|
|
Returns: (is_centered, "torso_line" or "none", Coords of intersection/midpoint)
|
|
"""
|
|
band_min = 0.5 - center_threshold
|
|
band_max = 0.5 + center_threshold
|
|
|
|
for kps in self.all_pose_kps:
|
|
l_sh = kps.left_shoulder
|
|
r_sh = kps.right_shoulder
|
|
l_hp = kps.left_hip
|
|
r_hp = kps.right_hip
|
|
|
|
# Need visibility for both shoulders and both hips to form the line
|
|
if not (l_sh.is_visible and r_sh.is_visible and l_hp.is_visible and r_hp.is_visible):
|
|
continue
|
|
|
|
sx = (l_sh.x + r_sh.x) / 2.0
|
|
sy = (l_sh.y + r_sh.y) / 2.0
|
|
hx = (l_hp.x + r_hp.x) / 2.0
|
|
hy = (l_hp.y + r_hp.y) / 2.0
|
|
|
|
seg_min_x = min(sx, hx)
|
|
seg_max_x = max(sx, hx)
|
|
|
|
# Quick reject: if the x-range of the segment doesn't touch the band
|
|
if seg_max_x < band_min or seg_min_x > band_max:
|
|
continue
|
|
|
|
# If either endpoint is already inside the band, return that endpoint/midpoint
|
|
if band_min <= sx <= band_max:
|
|
return True, "torso_line", Coords(x=sx, y=sy, is_visible=True)
|
|
if band_min <= hx <= band_max:
|
|
return True, "torso_line", Coords(x=hx, y=hy, is_visible=True)
|
|
|
|
# Otherwise the segment crosses the band somewhere between the endpoints.
|
|
# Compute intersection with the central vertical line x=0.5 when possible.
|
|
dx = hx - sx
|
|
dy = hy - sy
|
|
if abs(dx) < 1e-6:
|
|
# Vertical segment (x nearly constant) and we already know it intersects band
|
|
mid_x = sx
|
|
mid_y = (sy + hy) / 2.0
|
|
return True, "torso_line", Coords(x=mid_x, y=mid_y, is_visible=True)
|
|
|
|
# param t where x(t) = sx + t*dx == 0.5
|
|
t = (0.5 - sx) / dx
|
|
if 0.0 <= t <= 1.0:
|
|
inter_y = sy + t * dy
|
|
return True, "torso_line", Coords(x=0.5, y=inter_y, is_visible=True)
|
|
|
|
# Fallback: return midpoint of segment if we reach here (shouldn't normally)
|
|
mid_x = (sx + hx) / 2.0
|
|
mid_y = (sy + hy) / 2.0
|
|
return True, "torso_line", Coords(x=mid_x, y=mid_y, is_visible=True)
|
|
|
|
return False, "none", None
|
|
|
|
def is_upright(self, angle_threshold_degrees: float = 20.0) -> bool:
|
|
"""Return True if the torso (or any pair of important features) is approximately vertical.
|
|
|
|
Logic: consider important features (eye_mid, shoulder_mid, hip_mid, nose). If at least two
|
|
visible features form a vector whose angle to vertical is within `angle_threshold_degrees`,
|
|
consider the person upright.
|
|
"""
|
|
|
|
def angle_from_vertical(p1: Coords, p2: Coords) -> float:
|
|
vx = p2.x - p1.x
|
|
vy = p2.y - p1.y
|
|
if abs(vx) < 1e-9 and abs(vy) < 1e-9:
|
|
return 90.0
|
|
# angle between (vx, vy) and vertical (0,1): use atan2(|vx|, |vy|)
|
|
ang_rad = math.atan2(abs(vx), abs(vy))
|
|
return math.degrees(ang_rad)
|
|
|
|
for kps in self.all_pose_kps:
|
|
eye = kps.eye_midpoint()
|
|
shoulder = kps.shoulder_midpoint()
|
|
hip = kps.hip_midpoint()
|
|
nose = kps.nose
|
|
|
|
features = [f for f in (shoulder, hip) if f.is_visible]
|
|
if len(features) < 2:
|
|
continue
|
|
|
|
# Check all pairs; if any pair is near-vertical, return True
|
|
for a_f, b_f in combinations(features, 2):
|
|
a = angle_from_vertical(a_f, b_f)
|
|
if a <= angle_threshold_degrees:
|
|
return True
|
|
|
|
return False
|
|
|
|
def is_laying(self, angle_threshold_degrees: float = 20.0) -> bool:
|
|
"""Return True if the torso (or any pair of important features) is approximately horizontal.
|
|
|
|
Logic: consider important features (eye_mid, shoulder_mid, hip_mid, nose). If at least two
|
|
visible features form a vector whose angle to horizontal is within `angle_threshold_degrees`,
|
|
consider the person laying down.
|
|
"""
|
|
|
|
def angle_from_horizontal(p1: Coords, p2: Coords) -> float:
|
|
vx = p2.x - p1.x
|
|
vy = p2.y - p1.y
|
|
if abs(vx) < 1e-9 and abs(vy) < 1e-9:
|
|
return 90.0
|
|
# angle between (vx, vy) and horizontal (1,0): use atan2(|vy|, |vx|)
|
|
ang_rad = math.atan2(abs(vy), abs(vx))
|
|
return math.degrees(ang_rad)
|
|
|
|
for kps in self.all_pose_kps:
|
|
eye = kps.eye_midpoint()
|
|
shoulder = kps.shoulder_midpoint()
|
|
hip = kps.hip_midpoint()
|
|
nose = kps.nose
|
|
|
|
features = [f for f in (shoulder, hip) if f.is_visible]
|
|
if len(features) < 2:
|
|
continue
|
|
|
|
# Check all pairs; if any pair is near-horizontal, return True
|
|
for a_f, b_f in combinations(features, 2):
|
|
a = angle_from_horizontal(a_f, b_f)
|
|
if a <= angle_threshold_degrees:
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
# --- Utility Functions ---
|
|
|
|
|
|
def read_dims(image_path: Path) -> tuple[int, int]:
|
|
# ... (read_dims implementation remains UNCHANGED) ...
|
|
with image_path.open("rb") as file:
|
|
if file.read(2) != b"\xff\xd8":
|
|
raise ValueError(f"{image_path} is not a valid JPEG file")
|
|
|
|
file.seek(0)
|
|
try:
|
|
img = Image.open(file)
|
|
width, height = img.size
|
|
return width, height
|
|
except Exception:
|
|
file.seek(0)
|
|
while True:
|
|
marker = file.read(1)
|
|
if not marker or marker != b"\xff":
|
|
raise ValueError(f"Invalid JPEG format in {image_path}")
|
|
marker_type = int.from_bytes(file.read(1), byteorder="big")
|
|
length = int.from_bytes(file.read(2), byteorder="big") - 2
|
|
is_sof = 0xC0 <= marker_type <= 0xCF and marker_type not in (0xC4, 0xC8, 0xCC)
|
|
if is_sof:
|
|
file.seek(1, 1)
|
|
height = int.from_bytes(file.read(2), byteorder="big")
|
|
width = int.from_bytes(file.read(2), byteorder="big")
|
|
return width, height
|
|
file.seek(length, 1)
|
|
|
|
|
|
def calc_threshold(width: int, height: int, target_aspect: float = 1080 / 2400) -> float:
|
|
"""Compute a per-image center threshold using its dimensions.
|
|
|
|
For tall images the center area will be wider than the image so we return 0.5 (full width). For
|
|
wide images the square center is narrower and the returned threshold < 0.5.
|
|
"""
|
|
img_aspect = width / height
|
|
|
|
if target_aspect >= img_aspect:
|
|
return 0.5
|
|
|
|
# center_width = target_aspect * height (in pixels)
|
|
center_width = target_aspect * height
|
|
half_width_norm = (center_width / 2.0) / width
|
|
return min(max(half_width_norm, 0.0), 0.5)
|
|
|
|
|
|
@cache
|
|
def make_dir(path: Path) -> None:
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
@cache
|
|
def _yolov8_detector(model_type: str):
|
|
if YOLO is None or torch is None:
|
|
logging.error("YOLO/Torch dependencies are missing.")
|
|
return None
|
|
try:
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
except Exception:
|
|
device = "cpu"
|
|
try:
|
|
if model_type == "pose":
|
|
model = YOLO(Path(__file__).parent / "yolov8n-pose.pt")
|
|
else:
|
|
raise ValueError(f"Unknown model type: {model_type}")
|
|
model.to(device)
|
|
return model
|
|
except Exception:
|
|
logging.exception(f"YOLOv8-{model_type} model initialization failed.")
|
|
return None
|
|
|
|
|
|
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),
|
|
)
|
|
|
|
|
|
# --- Core Detection Function (UNCHANGED) ---
|
|
|
|
|
|
def detect_pose(image_path: Path) -> PoseDetectionResult:
|
|
"""
|
|
Performs pose detection and returns a structured result object.
|
|
"""
|
|
yolo_model = _yolov8_detector("pose")
|
|
all_person_boxes = np.array([])
|
|
all_keypoints_xyc = np.array([])
|
|
all_pose_kps = []
|
|
|
|
if yolo_model is None:
|
|
return PoseDetectionResult(boxes=all_person_boxes, keypoints_xyc=all_keypoints_xyc, all_pose_kps=all_pose_kps)
|
|
|
|
try:
|
|
results = yolo_model(str(image_path), conf=0.65, iou=0.5, verbose=False)
|
|
|
|
if results and results[0].boxes and results[0].keypoints:
|
|
all_person_boxes = results[0].boxes.xyxy.cpu().numpy()
|
|
kps_norm_xy = results[0].keypoints.xyn.cpu().numpy()
|
|
kps_conf = results[0].keypoints.conf.cpu().numpy()
|
|
# Combine into an N_person x 17 x 3 array (x, y, confidence)
|
|
all_keypoints_xyc = np.concatenate([kps_norm_xy, np.expand_dims(kps_conf, axis=2)], axis=2)
|
|
|
|
if all_person_boxes.size > 0:
|
|
for kp_xyc in all_keypoints_xyc:
|
|
# Encapsulate the raw keypoint data for clean access
|
|
all_pose_kps.append(_extract_keypoints(kp_xyc))
|
|
|
|
return PoseDetectionResult(boxes=all_person_boxes, keypoints_xyc=all_keypoints_xyc, all_pose_kps=all_pose_kps)
|
|
|
|
except Exception:
|
|
logging.exception("YOLOv8-Pose detection failed.")
|
|
return PoseDetectionResult(boxes=all_person_boxes, keypoints_xyc=all_keypoints_xyc, all_pose_kps=all_pose_kps)
|
|
|
|
|
|
# --- Debug Drawing Function (UNCHANGED) ---
|
|
|
|
|
|
def draw_debug_image(
|
|
image_path: Path,
|
|
width: int,
|
|
height: int,
|
|
result: PoseDetectionResult,
|
|
center_result: PoseDetectionResult.CenterResult,
|
|
save_path: Path,
|
|
) -> None:
|
|
"""
|
|
Draws the centering zone, bounding boxes, and highlights the successful centering point
|
|
using the coordinates provided by the is_centered method.
|
|
"""
|
|
if cv2 is None or np is None:
|
|
logging.error("OpenCV/Numpy is required for debug but is not available.")
|
|
return
|
|
|
|
# Check for centering first
|
|
is_centered = center_result.is_centered
|
|
centered_by = center_result.reason
|
|
centering_point_coords = center_result.coords
|
|
|
|
img = cv2.imdecode(np.fromfile(str(image_path), dtype=np.uint8), cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
logging.error(f"Could not load image for debugging: {image_path}")
|
|
return
|
|
|
|
# Draw Centering Zone
|
|
cx = width / 2.0
|
|
thresh_px_x = int(width * center_result.threshold)
|
|
x_mid_start = int(cx - thresh_px_x)
|
|
x_mid_end = int(cx + thresh_px_x)
|
|
y_mid_start, y_mid_end = 0, height
|
|
|
|
BOX_THICKNESS = 10
|
|
KP_RADIUS = 10
|
|
KP_THICKNESS = -1
|
|
# VISIBILITY_THRESH is not strictly needed here but kept for clarity
|
|
|
|
overlay = img.copy()
|
|
zone_color = (0, 255, 0) if is_centered else (0, 0, 255) # Green if centered, Red otherwise
|
|
cv2.rectangle(overlay, (x_mid_start, y_mid_start), (x_mid_end, y_mid_end), zone_color, -1)
|
|
alpha = 0.2
|
|
img = cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)
|
|
|
|
# Draw Detections
|
|
for box_idx in range(len(result.boxes)):
|
|
box = result.boxes[box_idx]
|
|
kp_xyc = result.keypoints_xyc[box_idx]
|
|
kp_px = (kp_xyc[:, :2] * np.array([width, height])).astype(int)
|
|
kps_conf = kp_xyc[:, 2]
|
|
|
|
# The bounding box color is based on the global centering status
|
|
person_color = (0, 255, 0) if is_centered else (255, 0, 0)
|
|
|
|
x1, y1, x2, y2 = map(int, box)
|
|
cv2.rectangle(img, (x1, y1), (x2, y2), person_color, BOX_THICKNESS)
|
|
|
|
# Highlight the calculated centering point (Assumes the first person detected is the one that triggered the center check)
|
|
if box_idx == 0 and is_centered and centering_point_coords:
|
|
norm_coords = centering_point_coords
|
|
centering_point_px = (int(norm_coords.x * width), int(norm_coords.y * height))
|
|
|
|
# Highlight the calculated centering point
|
|
cv2.circle(img, centering_point_px, 12, person_color, -1)
|
|
cv2.circle(img, centering_point_px, 6, (255, 255, 255), -1)
|
|
|
|
# Draw all visible keypoints (for context)
|
|
for i in range(len(kp_px)):
|
|
if kps_conf[i] > VISIBILITY_THRESH:
|
|
cv2.circle(img, tuple(kp_px[i]), KP_RADIUS, (255, 255, 0), KP_THICKNESS)
|
|
|
|
# Save the debug image
|
|
|
|
_, buffer = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 60])
|
|
save_path.write_bytes(buffer.tobytes())
|
|
logging.info(f"Saved debug image to {save_path}. Centered by: {centered_by}")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Sort images based on pose centering.")
|
|
parser.add_argument("image_paths", nargs="+", type=Path, help="Path(s) to JPEG image(s)")
|
|
|
|
parser.add_argument(
|
|
"--by-pose",
|
|
action="store_true",
|
|
default=True,
|
|
help="Sort images based on horizontal centering of the person's core (nose/shoulders/hips).",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--debug",
|
|
action="store_true",
|
|
help="Saves a debug image showing the centering zone and detected points, but does NOT move the original file.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--is-upright",
|
|
action="store_true",
|
|
help="Also detect upright posture and move upright images to `_pose_upright`.",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
"""Main function to process images."""
|
|
args = parse_args()
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
|
|
|
for image_path in args.image_paths:
|
|
if not image_path.is_file():
|
|
logging.warning(f"{image_path} is not a file. Skipping.")
|
|
continue
|
|
|
|
if image_path.suffix.lower() not in [".jpg", ".jpeg"]:
|
|
logging.warning(f"{image_path} is not a JPEG file. Skipping.")
|
|
continue
|
|
|
|
try:
|
|
width, height = read_dims(image_path)
|
|
|
|
# 1. Run detection (only data extraction)
|
|
detection_result = detect_pose(image_path)
|
|
|
|
# Compute per-image threshold and run centering logic
|
|
threshold = calc_threshold(width=width, height=height)
|
|
center_res = detection_result.is_centered(center_threshold=threshold)
|
|
is_centered = center_res.is_centered
|
|
centered_by = center_res.reason
|
|
|
|
# Define target path/directory based on centering result
|
|
target_dir = image_path.parent / ("_pose_centered" if is_centered else "_pose_other")
|
|
# If user requested upright detection and the person is upright, override target
|
|
if args.is_upright:
|
|
try:
|
|
if detection_result.is_upright(angle_threshold_degrees=20):
|
|
target_dir = image_path.parent / "_pose_upright"
|
|
elif detection_result.is_laying(angle_threshold_degrees=40):
|
|
target_dir = image_path.parent / "_pose_laying"
|
|
except Exception:
|
|
# If upright detection fails, fall back to normal behavior
|
|
pass
|
|
make_dir(target_dir)
|
|
target_path = (target_dir / image_path.name).with_suffix(".jpg")
|
|
|
|
has_detections = detection_result.boxes.size > 0
|
|
|
|
# --- DEBUG LOGIC (Only output image, no move) ---
|
|
if args.debug:
|
|
if has_detections:
|
|
debug_filename = image_path.stem + "_debug" + image_path.suffix
|
|
debug_path = target_dir / debug_filename
|
|
draw_debug_image(
|
|
image_path=image_path,
|
|
width=width,
|
|
height=height,
|
|
result=detection_result,
|
|
center_result=center_res,
|
|
save_path=debug_path,
|
|
)
|
|
logging.info(f"Processed {image_path} (Debug mode active). File was NOT moved.")
|
|
else:
|
|
logging.info(f"Processed {image_path} (Debug mode active). No person detected, skipping debug output.")
|
|
# --- NON-DEBUG LOGIC (Move file) ---
|
|
else:
|
|
image_path.rename(target_path)
|
|
logging.info(f"Moved {image_path} to {target_path}. Centered by: {centered_by}")
|
|
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except Exception as e:
|
|
logging.exception(f"Error processing {image_path}. Skipping.")
|
|
continue
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|