Spaces:
Sleeping
Sleeping
File size: 4,951 Bytes
5ec223d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """
Geometric body-region segmentation using MediaPipe pose landmarks.
SAM2 segments by pixel similarity, which on a clothed photo yields the tank-top
color region rather than the underlying anatomy. For breast/buttocks regions we
build elliptical masks positioned by pose landmarks (shoulders, hips, knees) so
the result is anatomically located regardless of clothing.
"""
import sys
import numpy as np
from PIL import Image
def log(msg: str):
print(msg, flush=True)
def _ellipse_mask(W: int, H: int, cx: float, cy: float, rx: float, ry: float, angle: float = 0.0) -> np.ndarray:
"""Build a (H, W) bool mask: True inside the ellipse at (cx,cy) with semi-axes (rx,ry), rotated by `angle` rad."""
ys, xs = np.mgrid[0:H, 0:W]
dx = xs - cx
dy = ys - cy
cos_a = np.cos(-angle)
sin_a = np.sin(-angle)
rx_d = (dx * cos_a - dy * sin_a) / max(rx, 1e-3)
ry_d = (dx * sin_a + dy * cos_a) / max(ry, 1e-3)
return (rx_d * rx_d + ry_d * ry_d) <= 1.0
def _bbox_of_mask(mask: np.ndarray) -> list[int]:
rows = np.any(mask, axis=1)
cols = np.any(mask, axis=0)
if not rows.any() or not cols.any():
return [0, 0, 0, 0]
rmin, rmax = np.where(rows)[0][[0, -1]]
cmin, cmax = np.where(cols)[0][[0, -1]]
return [int(cmin), int(rmin), int(cmax - cmin + 1), int(rmax - rmin + 1)]
def segment_anatomy(image: Image.Image, regions: list[str]) -> dict:
"""
Build geometric masks for body regions using MediaPipe pose landmarks.
Supported regions: breast_left, breast_right, buttocks.
Returns {region: {"mask": list[list[bool]], "bbox": [x,y,w,h]}}.
Skips regions for which the required landmarks aren't visible.
"""
from pose import detect_landmarks # reuse the same MediaPipe wrapper
log(f"[Anatomy] Geometric segmentation for {regions} on image {image.size}")
landmarks = detect_landmarks(image)
W, H = image.size
def get_pixel(idx: int) -> tuple[float, float] | None:
lm = landmarks.get(idx)
if not lm or lm.get("visibility", 0) < 0.3:
return None
return float(lm["px"]), float(lm["py"])
sl = get_pixel(11) # left shoulder (image-right side of subject)
sr = get_pixel(12) # right shoulder
hl = get_pixel(23) # left hip
hr = get_pixel(24) # right hip
kl = get_pixel(25) # left knee
kr = get_pixel(26) # right knee
if not (sl and sr and hl and hr):
log("[Anatomy] Missing shoulder/hip landmarks — cannot build geometric masks.")
return {}
# Sternum (midpoint between shoulders)
sx = (sl[0] + sr[0]) / 2
sy = (sl[1] + sr[1]) / 2
# Pelvis (midpoint between hips)
px = (hl[0] + hr[0]) / 2
py = (hl[1] + hr[1]) / 2
shoulder_width = float(np.hypot(sl[0] - sr[0], sl[1] - sr[1]))
torso_height = float(np.hypot(sx - px, sy - py))
torso_angle = float(np.arctan2(py - sy, px - sx) - np.pi / 2) # 0 if perfectly upright
log(f"[Anatomy] shoulders=({sl[0]:.0f},{sl[1]:.0f})↔({sr[0]:.0f},{sr[1]:.0f}) "
f"hips=({hl[0]:.0f},{hl[1]:.0f})↔({hr[0]:.0f},{hr[1]:.0f}) "
f"shoulder_w={shoulder_width:.0f} torso_h={torso_height:.0f}")
results: dict = {}
# ── Breasts: elliptical regions sitting on the upper torso, halfway between
# each shoulder and the sternum, dropped 25% of the torso height down. ──
if "breast_left" in regions or "breast_right" in regions:
# Vertical drop from shoulder line toward pelvis
drop_x = (px - sx) * 0.30
drop_y = (py - sy) * 0.30
for region, shoulder in (("breast_left", sl), ("breast_right", sr)):
if region not in regions:
continue
cx = (shoulder[0] + sx) / 2 + drop_x
cy = (shoulder[1] + sy) / 2 + drop_y
rx = abs(shoulder[0] - sx) * 0.55
ry = torso_height * 0.18
mask = _ellipse_mask(W, H, cx, cy, rx, ry, torso_angle)
log(f"[Anatomy] {region} center=({cx:.0f},{cy:.0f}) rx={rx:.0f} ry={ry:.0f}")
results[region] = {"mask": mask.tolist(), "bbox": _bbox_of_mask(mask)}
# ── Buttocks: ellipse spanning hips, dropped halfway toward knees. ──
if "buttocks" in regions:
if kl and kr:
kx = (kl[0] + kr[0]) / 2
ky = (kl[1] + kr[1]) / 2
cx = (px + kx) / 2 * 0.5 + px * 0.5 # weighted toward hips
cy = py + (ky - py) * 0.25
else:
cx, cy = px, py + torso_height * 0.20
hip_width = float(np.hypot(hl[0] - hr[0], hl[1] - hr[1]))
rx = hip_width * 0.65
ry = torso_height * 0.30
mask = _ellipse_mask(W, H, cx, cy, rx, ry, torso_angle)
log(f"[Anatomy] buttocks center=({cx:.0f},{cy:.0f}) rx={rx:.0f} ry={ry:.0f}")
results["buttocks"] = {"mask": mask.tolist(), "bbox": _bbox_of_mask(mask)}
log(f"[Anatomy] Built {len(results)} geometric masks.")
return results
|