#!/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()