chore: Update everything
This commit is contained in:
Executable
+618
@@ -0,0 +1,618 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user