545 lines
24 KiB
Python
545 lines
24 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from scipy import ndimage
|
|
|
|
|
|
SHAPE_COMPONENTS = 45
|
|
SCALE_COMPONENTS = 28
|
|
|
|
|
|
def _selected_people(pose_data: Dict[str, Any], person_index: int) -> List[Tuple[int, int, Dict[str, Any]]]:
|
|
selected = []
|
|
for frame_index, frame in enumerate(pose_data.get("frames", [])):
|
|
if person_index < 0:
|
|
indices = range(len(frame))
|
|
elif person_index < len(frame):
|
|
indices = (person_index,)
|
|
else:
|
|
indices = ()
|
|
for index in indices:
|
|
selected.append((frame_index, index, frame[index]))
|
|
return selected
|
|
|
|
|
|
def profile_from_pose_data(pose_data: Dict[str, Any], person_index: int = 0) -> Dict[str, Any]:
|
|
people = _selected_people(pose_data, person_index)
|
|
shapes = [np.asarray(p["shape_params"], dtype=np.float32) for _, _, p in people if p.get("shape_params") is not None]
|
|
scales = [np.asarray(p["scale_params"], dtype=np.float32) for _, _, p in people if p.get("scale_params") is not None]
|
|
if not shapes or not scales:
|
|
raise ValueError(f"SAM3D BodyMod: person {person_index} has no usable shape and scale parameters")
|
|
if any(v.shape != (SHAPE_COMPONENTS,) for v in shapes):
|
|
raise ValueError("SAM3D BodyMod: expected 45 MHR shape parameters per person")
|
|
if any(v.shape != (SCALE_COMPONENTS,) for v in scales):
|
|
raise ValueError("SAM3D BodyMod: expected 28 MHR scale parameters per person")
|
|
return {
|
|
"shape_params": np.median(np.stack(shapes), axis=0).astype(np.float32),
|
|
"scale_params": np.median(np.stack(scales), axis=0).astype(np.float32),
|
|
"reference_count": min(len(shapes), len(scales)),
|
|
"preset": "real_human",
|
|
"fitted": False,
|
|
}
|
|
|
|
|
|
def blend_coefficients(source: np.ndarray, target: np.ndarray, strength: float) -> np.ndarray:
|
|
source = np.asarray(source, dtype=np.float32)
|
|
target = np.asarray(target, dtype=np.float32)
|
|
return (source + float(strength) * (target - source)).astype(np.float32)
|
|
|
|
def exaggerate_profile(profile: Dict[str, Any], strength: float) -> Dict[str, Any]:
|
|
out = dict(profile)
|
|
out["shape_params"] = (np.asarray(profile["shape_params"], dtype=np.float32) * float(strength)).astype(np.float32)
|
|
out["scale_params"] = (np.asarray(profile["scale_params"], dtype=np.float32) * float(strength)).astype(np.float32)
|
|
out["recognition_exaggeration"] = float(strength)
|
|
return out
|
|
|
|
|
|
def _body_type_target(base_shape: torch.Tensor, body_type: str) -> torch.Tensor:
|
|
low = base_shape.amin(dim=0)
|
|
high = base_shape.amax(dim=0)
|
|
height = (high[1] - low[1]).clamp(min=1.0)
|
|
y = (base_shape[:, 1] - low[1]) / height
|
|
x = base_shape[:, 0]
|
|
z = base_shape[:, 2] - base_shape[:, 2].median()
|
|
torso = (x.abs() < height * 0.22).to(base_shape.dtype)
|
|
body = (y < 0.84).to(base_shape.dtype)
|
|
|
|
if body_type == "ectomorph":
|
|
radial = -0.11 * body * (0.45 + 0.55 * torso)
|
|
vertical = 0.025 * body
|
|
elif body_type == "mesomorph":
|
|
chest = torch.exp(-0.5 * ((y - 0.68) / 0.12).square())
|
|
waist = torch.exp(-0.5 * ((y - 0.48) / 0.09).square())
|
|
thighs = torch.exp(-0.5 * ((y - 0.30) / 0.11).square())
|
|
radial = body * (0.16 * chest - 0.045 * waist + 0.06 * thighs)
|
|
vertical = torch.zeros_like(y)
|
|
elif body_type == "endomorph":
|
|
abdomen = torch.exp(-0.5 * ((y - 0.49) / 0.17).square())
|
|
hips = torch.exp(-0.5 * ((y - 0.34) / 0.11).square())
|
|
radial = body * torso * (0.17 * abdomen + 0.09 * hips)
|
|
vertical = -0.015 * body
|
|
else:
|
|
raise ValueError(f"SAM3D BodyMod: unknown body type '{body_type}'")
|
|
|
|
target = torch.zeros_like(base_shape)
|
|
target[:, 0] = x * radial
|
|
target[:, 2] = z * radial
|
|
target[:, 1] = (base_shape[:, 1] - (low[1] + high[1]) * 0.5) * vertical
|
|
return target
|
|
|
|
|
|
def body_type_profile(inner: Any, profile: Dict[str, Any], body_type: str, strength: float) -> Dict[str, Any]:
|
|
out = dict(profile)
|
|
shape = np.asarray(profile["shape_params"], dtype=np.float32)
|
|
scale = np.asarray(profile["scale_params"], dtype=np.float32)
|
|
if shape.shape != (SHAPE_COMPONENTS,) or scale.shape != (SCALE_COMPONENTS,):
|
|
raise ValueError("SAM3D BodyMod: invalid shape profile")
|
|
|
|
amount = float(strength)
|
|
if body_type == "detected" or amount == 0.0:
|
|
out["body_type"] = body_type
|
|
out["body_type_strength"] = amount
|
|
return out
|
|
if body_type == "maximize_detected":
|
|
out["shape_params"] = (shape * (1.0 + amount)).astype(np.float32)
|
|
out["scale_params"] = (scale * (1.0 + amount)).astype(np.float32)
|
|
else:
|
|
rig = inner.head_pose.mhr
|
|
basis = rig.identity_basis.detach().float().reshape(SHAPE_COMPONENTS, -1)
|
|
target = _body_type_target(rig.base_shape.detach().float(), body_type).reshape(-1)
|
|
gram = basis @ basis.T
|
|
ridge = gram.diagonal().mean() * 1e-4
|
|
direction = torch.linalg.solve(
|
|
gram + torch.eye(SHAPE_COMPONENTS, device=gram.device, dtype=gram.dtype) * ridge,
|
|
basis @ target,
|
|
)
|
|
direction = direction * (3.0 / direction.abs().max().clamp(min=3.0))
|
|
out["shape_params"] = (shape + amount * direction.cpu().numpy()).astype(np.float32)
|
|
out["scale_params"] = scale.copy()
|
|
out["body_type"] = body_type
|
|
out["body_type_strength"] = amount
|
|
return out
|
|
|
|
|
|
def _mask_polarity(mask: np.ndarray, polarity: str) -> np.ndarray:
|
|
direct = mask > 0.5
|
|
if polarity == "white_is_subject":
|
|
return direct
|
|
if polarity == "comfy_alpha":
|
|
return ~direct
|
|
|
|
border = np.concatenate([direct[0], direct[-1], direct[:, 0], direct[:, -1]])
|
|
direct_area = float(direct.mean())
|
|
inverse_area = 1.0 - direct_area
|
|
|
|
def score(area: float, border_fraction: float) -> float:
|
|
area_penalty = max(0.0, 0.01 - area) * 20.0 + max(0.0, area - 0.85) * 10.0
|
|
return border_fraction * 3.0 + area_penalty
|
|
|
|
direct_score = score(direct_area, float(border.mean()))
|
|
inverse_score = score(inverse_area, float((~border).mean()))
|
|
return direct if direct_score <= inverse_score else ~direct
|
|
|
|
|
|
def _largest_filled_component(mask: np.ndarray) -> np.ndarray:
|
|
mask = ndimage.binary_closing(mask, iterations=2)
|
|
mask = ndimage.binary_fill_holes(mask)
|
|
labels, count = ndimage.label(mask)
|
|
if count == 0:
|
|
return np.zeros_like(mask, dtype=bool)
|
|
sizes = np.bincount(labels.reshape(-1))
|
|
sizes[0] = 0
|
|
return labels == int(sizes.argmax())
|
|
|
|
|
|
def build_foreground_masks(
|
|
image: torch.Tensor,
|
|
pose_data: Dict[str, Any],
|
|
person_index: int,
|
|
foreground_mask: Optional[torch.Tensor] = None,
|
|
mask_polarity: str = "auto",
|
|
auto_mask_threshold: float = 0.08,
|
|
) -> List[np.ndarray]:
|
|
images = image.detach().float().cpu().numpy()
|
|
B, H, W, _ = images.shape
|
|
masks_np = None if foreground_mask is None else foreground_mask.detach().float().cpu().numpy()
|
|
if masks_np is not None and masks_np.ndim == 2:
|
|
masks_np = masks_np[None]
|
|
|
|
out = []
|
|
for frame_index in range(B):
|
|
if masks_np is not None:
|
|
raw = masks_np[min(frame_index, masks_np.shape[0] - 1)]
|
|
if raw.shape != (H, W):
|
|
raw_t = torch.from_numpy(raw).view(1, 1, *raw.shape)
|
|
raw = F.interpolate(raw_t, size=(H, W), mode="bilinear", align_corners=False)[0, 0].numpy()
|
|
mask = _mask_polarity(raw, mask_polarity)
|
|
else:
|
|
frame = images[frame_index, :, :, :3]
|
|
border = np.concatenate([frame[0], frame[-1], frame[:, 0], frame[:, -1]], axis=0)
|
|
background = np.median(border, axis=0)
|
|
distance = np.linalg.norm(frame - background[None, None, :], axis=-1)
|
|
mask = distance > float(auto_mask_threshold)
|
|
|
|
frame_people = pose_data.get("frames", [])[frame_index] if frame_index < len(pose_data.get("frames", [])) else []
|
|
if person_index < len(frame_people):
|
|
bbox = np.asarray(frame_people[person_index].get("bbox", []), dtype=np.float32).reshape(-1)
|
|
if bbox.size >= 4:
|
|
x1, y1, x2, y2 = bbox[:4]
|
|
pad_x = max(4.0, (x2 - x1) * 0.08)
|
|
pad_y = max(4.0, (y2 - y1) * 0.08)
|
|
ix1 = max(0, int(np.floor(x1 - pad_x)))
|
|
iy1 = max(0, int(np.floor(y1 - pad_y)))
|
|
ix2 = min(W, int(np.ceil(x2 + pad_x)))
|
|
iy2 = min(H, int(np.ceil(y2 + pad_y)))
|
|
clipped = np.zeros_like(mask, dtype=bool)
|
|
clipped[iy1:iy2, ix1:ix2] = mask[iy1:iy2, ix1:ix2]
|
|
mask = clipped
|
|
|
|
mask = _largest_filled_component(mask)
|
|
if float(mask.mean()) < 0.005:
|
|
mask = np.zeros((H, W), dtype=bool)
|
|
out.append(mask)
|
|
return out
|
|
|
|
|
|
def _principal_point(person: Dict[str, Any]) -> Tuple[float, float]:
|
|
kp2d = np.asarray(person["pred_keypoints_2d"], dtype=np.float32)
|
|
kp3d = np.asarray(person["pred_keypoints_3d"], dtype=np.float32)
|
|
cam_t = np.asarray(person["pred_cam_t"], dtype=np.float32).reshape(3)
|
|
focal = float(np.asarray(person["focal_length"], dtype=np.float32).reshape(-1)[0])
|
|
x, y, z = kp3d[0] + cam_t
|
|
z = max(float(z), 1e-6)
|
|
return float(kp2d[0, 0] - focal * x / z), float(kp2d[0, 1] - focal * y / z)
|
|
|
|
|
|
def _mask_targets(mask: np.ndarray, bins: int = 18) -> Optional[Dict[str, np.ndarray]]:
|
|
ys, xs = np.nonzero(mask)
|
|
if xs.size == 0:
|
|
return None
|
|
H, W = mask.shape
|
|
top, bottom = int(ys.min()), int(ys.max())
|
|
centers_px = np.linspace(top, bottom, bins, dtype=np.float32)
|
|
half_band = max(1, int(round((bottom - top + 1) / (bins * 1.5))))
|
|
left, right = [], []
|
|
for center in centers_px:
|
|
y1 = max(0, int(round(center)) - half_band)
|
|
y2 = min(H, int(round(center)) + half_band + 1)
|
|
_, band_x = np.nonzero(mask[y1:y2])
|
|
if band_x.size == 0:
|
|
row = mask[min(H - 1, max(0, int(round(center))))]
|
|
band_x = np.nonzero(row)[0]
|
|
left.append(float(band_x.min()) / max(W - 1, 1))
|
|
right.append(float(band_x.max()) / max(W - 1, 1))
|
|
outside = ndimage.distance_transform_edt(~mask).astype(np.float32)
|
|
outside /= max(float(np.hypot(H, W)), 1.0)
|
|
return {
|
|
"centers": centers_px / max(H - 1, 1),
|
|
"left": np.asarray(left, dtype=np.float32),
|
|
"right": np.asarray(right, dtype=np.float32),
|
|
"top": np.asarray(float(top) / max(H - 1, 1), dtype=np.float32),
|
|
"bottom": np.asarray(float(bottom) / max(H - 1, 1), dtype=np.float32),
|
|
"outside": outside,
|
|
}
|
|
|
|
|
|
def fit_profile_to_masks(
|
|
inner: Any,
|
|
pose_data: Dict[str, Any],
|
|
profile: Dict[str, Any],
|
|
masks: List[np.ndarray],
|
|
person_index: int,
|
|
preset: str,
|
|
fit_strength: float,
|
|
iterations: int,
|
|
progress: Optional[Callable[[], None]] = None,
|
|
) -> Dict[str, Any]:
|
|
if iterations <= 0 or fit_strength <= 0.0:
|
|
out = dict(profile)
|
|
out["preset"] = preset
|
|
return out
|
|
|
|
people = _selected_people(pose_data, person_index)
|
|
usable = []
|
|
for frame_index, _, person in people:
|
|
if frame_index >= len(masks):
|
|
continue
|
|
target = _mask_targets(masks[frame_index])
|
|
if target is not None:
|
|
usable.append((frame_index, person, target))
|
|
if not usable:
|
|
out = dict(profile)
|
|
out["preset"] = preset
|
|
return out
|
|
if len(usable) > 4:
|
|
keep = np.linspace(0, len(usable) - 1, 4).round().astype(np.int64)
|
|
usable = [usable[int(i)] for i in keep]
|
|
|
|
head = inner.head_pose
|
|
rig = head.mhr
|
|
device = head.scale_mean.device
|
|
dtype = head.scale_mean.dtype
|
|
H, W = pose_data["image_size"]
|
|
N = len(usable)
|
|
|
|
base_mp = torch.from_numpy(np.stack([
|
|
np.asarray(person["mhr_model_params"], dtype=np.float32) for _, person, _ in usable
|
|
])).to(device=device, dtype=dtype)
|
|
if base_mp.shape[1] != 204:
|
|
raise ValueError(f"SAM3D BodyMod: expected 204 packed MHR model parameters, got {base_mp.shape[1]}")
|
|
base_pose = base_mp[:, :136]
|
|
expr = torch.from_numpy(np.stack([
|
|
np.asarray(person["expr_params"], dtype=np.float32) for _, person, _ in usable
|
|
])).to(device=device, dtype=dtype)
|
|
camera = torch.from_numpy(np.stack([
|
|
np.asarray(person["pred_cam_t"], dtype=np.float32) for _, person, _ in usable
|
|
])).to(device=device, dtype=dtype)
|
|
focal = torch.tensor([
|
|
float(np.asarray(person["focal_length"]).reshape(-1)[0]) for _, person, _ in usable
|
|
], device=device, dtype=dtype)
|
|
principal = torch.tensor([
|
|
_principal_point(person) for _, person, _ in usable
|
|
], device=device, dtype=dtype)
|
|
|
|
centers = torch.from_numpy(np.stack([target["centers"] for _, _, target in usable])).to(device=device, dtype=dtype)
|
|
target_left = torch.from_numpy(np.stack([target["left"] for _, _, target in usable])).to(device=device, dtype=dtype)
|
|
target_right = torch.from_numpy(np.stack([target["right"] for _, _, target in usable])).to(device=device, dtype=dtype)
|
|
target_top = torch.tensor([float(target["top"]) for _, _, target in usable], device=device, dtype=dtype)
|
|
target_bottom = torch.tensor([float(target["bottom"]) for _, _, target in usable], device=device, dtype=dtype)
|
|
outside = torch.from_numpy(np.stack([target["outside"] for _, _, target in usable])).unsqueeze(1).to(device=device, dtype=dtype)
|
|
|
|
initial_shape = torch.from_numpy(np.asarray(profile["shape_params"], dtype=np.float32)).to(device=device, dtype=dtype)
|
|
initial_scale = torch.from_numpy(np.asarray(profile["scale_params"], dtype=np.float32)).to(device=device, dtype=dtype)
|
|
cartoon = preset == "cartoon"
|
|
shape_limit = 5.0 if cartoon else 1.5
|
|
scale_limit = 4.0 if cartoon else 1.0
|
|
regularization = 0.015 if cartoon else 0.12
|
|
tau = 0.012 if cartoon else 0.009
|
|
sigma = 0.035
|
|
|
|
def loss_for(offset: torch.Tensor) -> torch.Tensor:
|
|
shape_offset, scale_offset = torch.split(offset, [SHAPE_COMPONENTS, SCALE_COMPONENTS])
|
|
shape = initial_shape + shape_limit * shape_offset
|
|
scale = initial_scale + scale_limit * scale_offset
|
|
scales = head.scale_mean + scale @ head.scale_comps
|
|
model_params = torch.cat([base_pose, scales.unsqueeze(0).expand(N, -1)], dim=1)
|
|
verts_cm, _ = rig(
|
|
shape.unsqueeze(0).expand(N, -1), model_params, expr,
|
|
apply_correctives=False,
|
|
)
|
|
verts = verts_cm / 100.0
|
|
verts = torch.stack([verts[..., 0], -verts[..., 1], -verts[..., 2]], dim=-1)
|
|
camera_verts = verts + camera[:, None, :]
|
|
z = camera_verts[..., 2].clamp(min=0.05)
|
|
u = (principal[:, None, 0] + focal[:, None] * camera_verts[..., 0] / z) / max(W - 1, 1)
|
|
v = (principal[:, None, 1] + focal[:, None] * camera_verts[..., 1] / z) / max(H - 1, 1)
|
|
|
|
band = -0.5 * ((v[:, None, :] - centers[:, :, None]) / sigma).square()
|
|
band = band + torch.where(camera_verts[:, None, :, 2] > 0.05, 0.0, -1e4)
|
|
left_weights = torch.softmax(band - u[:, None, :] / tau, dim=-1)
|
|
right_weights = torch.softmax(band + u[:, None, :] / tau, dim=-1)
|
|
model_left = (left_weights * u[:, None, :]).sum(dim=-1)
|
|
model_right = (right_weights * u[:, None, :]).sum(dim=-1)
|
|
|
|
top_weights = torch.softmax(-v / tau, dim=-1)
|
|
bottom_weights = torch.softmax(v / tau, dim=-1)
|
|
model_top = (top_weights * v).sum(dim=-1)
|
|
model_bottom = (bottom_weights * v).sum(dim=-1)
|
|
|
|
grid = torch.stack([u * 2.0 - 1.0, v * 2.0 - 1.0], dim=-1).unsqueeze(1)
|
|
outside_sample = F.grid_sample(
|
|
outside, grid, mode="bilinear", padding_mode="border", align_corners=True,
|
|
)[:, 0, 0]
|
|
|
|
edge_loss = ((model_left - target_left).square() + (model_right - target_right).square()).mean()
|
|
height_loss = ((model_top - target_top).square() + (model_bottom - target_bottom).square()).mean()
|
|
outside_loss = outside_sample.square().mean()
|
|
delta_loss = shape_offset.square().mean() + scale_offset.square().mean()
|
|
return 8.0 * edge_loss + 2.0 * height_loss + 0.5 * outside_loss + regularization * delta_loss
|
|
|
|
offset = torch.zeros(SHAPE_COMPONENTS + SCALE_COMPONENTS, device=device, dtype=dtype)
|
|
first_loss = float(loss_for(offset).float().cpu())
|
|
best_loss = first_loss
|
|
best_offset = offset.clone()
|
|
rng = np.random.default_rng(0)
|
|
base_rate = 0.10 if cartoon else 0.06
|
|
for step in range(int(iterations)):
|
|
perturbation = 0.12 / ((step + 1) ** 0.101)
|
|
direction_np = rng.integers(0, 2, size=offset.numel(), dtype=np.int8).astype(np.float32) * 2.0 - 1.0
|
|
direction = torch.from_numpy(direction_np).to(device=device, dtype=dtype)
|
|
plus = (offset + perturbation * direction).clamp(-1.0, 1.0)
|
|
minus = (offset - perturbation * direction).clamp(-1.0, 1.0)
|
|
loss_plus = loss_for(plus)
|
|
loss_minus = loss_for(minus)
|
|
loss_plus_value = float(loss_plus.float().cpu())
|
|
loss_minus_value = float(loss_minus.float().cpu())
|
|
if loss_plus_value < best_loss:
|
|
best_loss = loss_plus_value
|
|
best_offset = plus.clone()
|
|
if loss_minus_value < best_loss:
|
|
best_loss = loss_minus_value
|
|
best_offset = minus.clone()
|
|
gradient = ((loss_plus - loss_minus) / (2.0 * perturbation)) * direction
|
|
gradient_scale = gradient.square().mean().sqrt().clamp(min=1e-8)
|
|
rate = base_rate / (1.0 + 0.04 * step)
|
|
offset = (offset - rate * gradient / gradient_scale).clamp(-1.0, 1.0)
|
|
if progress is not None:
|
|
progress()
|
|
|
|
last_loss = float(loss_for(offset).float().cpu())
|
|
if last_loss < best_loss:
|
|
best_loss = last_loss
|
|
best_offset = offset.clone()
|
|
offset = best_offset
|
|
last_loss = best_loss
|
|
shape_offset, scale_offset = torch.split(offset, [SHAPE_COMPONENTS, SCALE_COMPONENTS])
|
|
fitted_shape = initial_shape + shape_limit * shape_offset
|
|
fitted_scale = initial_scale + scale_limit * scale_offset
|
|
strength = float(fit_strength)
|
|
shape_out = initial_shape + strength * (fitted_shape - initial_shape)
|
|
scale_out = initial_scale + strength * (fitted_scale - initial_scale)
|
|
out = dict(profile)
|
|
out.update({
|
|
"shape_params": shape_out.float().cpu().numpy(),
|
|
"scale_params": scale_out.float().cpu().numpy(),
|
|
"preset": preset,
|
|
"fitted": True,
|
|
"fit_loss_before": first_loss,
|
|
"fit_loss_after": last_loss,
|
|
})
|
|
return out
|
|
|
|
|
|
def _anchor_point(keypoints: np.ndarray, anchor: str) -> np.ndarray:
|
|
if anchor == "pelvis":
|
|
return keypoints[[9, 10]].mean(axis=0)
|
|
if anchor == "feet":
|
|
return keypoints[[13, 14]].mean(axis=0)
|
|
return np.zeros(3, dtype=np.float32)
|
|
|
|
|
|
def regenerate_pose_data(
|
|
inner: Any,
|
|
pose_data: Dict[str, Any],
|
|
profile: Dict[str, Any],
|
|
shape_strength: float,
|
|
proportion_strength: float,
|
|
person_index: int,
|
|
anchor: str,
|
|
batch_size: int,
|
|
progress: Optional[Callable[[], None]] = None,
|
|
) -> Dict[str, Any]:
|
|
target_shape = np.asarray(profile["shape_params"], dtype=np.float32)
|
|
target_scale = np.asarray(profile["scale_params"], dtype=np.float32)
|
|
if target_shape.shape != (SHAPE_COMPONENTS,) or target_scale.shape != (SCALE_COMPONENTS,):
|
|
raise ValueError("SAM3D BodyMod: invalid shape profile")
|
|
|
|
entries = _selected_people(pose_data, person_index)
|
|
needed = (
|
|
"global_rot", "body_pose_params", "hand_pose_params", "shape_params",
|
|
"scale_params", "expr_params", "pred_cam_t", "focal_length",
|
|
"pred_keypoints_2d", "pred_keypoints_3d",
|
|
)
|
|
entries = [(fi, pi, p) for fi, pi, p in entries if all(p.get(key) is not None for key in needed)]
|
|
if not entries:
|
|
raise ValueError(f"SAM3D BodyMod: no complete person records found for person_index={person_index}")
|
|
|
|
frames_out = [[dict(person) for person in frame] for frame in pose_data["frames"]]
|
|
head = inner.head_pose
|
|
device = head.scale_mean.device
|
|
batch_size = max(1, int(batch_size))
|
|
|
|
for start in range(0, len(entries), batch_size):
|
|
chunk = entries[start:start + batch_size]
|
|
persons = [p for _, _, p in chunk]
|
|
shapes = np.stack([
|
|
blend_coefficients(p["shape_params"], target_shape, shape_strength) for p in persons
|
|
])
|
|
scales = np.stack([
|
|
blend_coefficients(p["scale_params"], target_scale, proportion_strength) for p in persons
|
|
])
|
|
global_rot = torch.from_numpy(np.stack([np.asarray(p["global_rot"], dtype=np.float32) for p in persons])).to(device)
|
|
body_pose = torch.from_numpy(np.stack([np.asarray(p["body_pose_params"], dtype=np.float32) for p in persons])).to(device)
|
|
hand_pose = torch.from_numpy(np.stack([np.asarray(p["hand_pose_params"], dtype=np.float32) for p in persons])).to(device)
|
|
shape_t = torch.from_numpy(shapes).to(device)
|
|
scale_t = torch.from_numpy(scales).to(device)
|
|
expr = torch.from_numpy(np.stack([np.asarray(p["expr_params"], dtype=np.float32) for p in persons])).to(device)
|
|
|
|
verts, kp3d_full, joint_coords, model_params, joint_rotmats = head.mhr_forward(
|
|
global_trans=torch.zeros_like(global_rot),
|
|
global_rot=global_rot,
|
|
body_pose_params=body_pose,
|
|
hand_pose_params=hand_pose,
|
|
scale_params=scale_t,
|
|
shape_params=shape_t,
|
|
expr_params=expr,
|
|
return_keypoints=True,
|
|
return_joint_coords=True,
|
|
return_model_params=True,
|
|
return_joint_rotations=True,
|
|
)
|
|
|
|
verts = torch.stack([verts[..., 0], -verts[..., 1], -verts[..., 2]], dim=-1)
|
|
kp3d = kp3d_full[:, :70]
|
|
kp3d = torch.stack([kp3d[..., 0], -kp3d[..., 1], -kp3d[..., 2]], dim=-1)
|
|
face3d = kp3d_full[:, 70:]
|
|
face3d = torch.stack([face3d[..., 0], -face3d[..., 1], -face3d[..., 2]], dim=-1)
|
|
joint_coords = torch.stack([
|
|
joint_coords[..., 0], -joint_coords[..., 1], -joint_coords[..., 2]
|
|
], dim=-1)
|
|
|
|
cameras_np = np.stack([np.asarray(p["pred_cam_t"], dtype=np.float32) for p in persons])
|
|
if anchor in ("pelvis", "feet"):
|
|
old_anchor = np.stack([
|
|
_anchor_point(np.asarray(p["pred_keypoints_3d"], dtype=np.float32), anchor) for p in persons
|
|
])
|
|
new_anchor = np.stack([
|
|
_anchor_point(kp3d[i].detach().float().cpu().numpy(), anchor) for i in range(len(persons))
|
|
])
|
|
cameras_np = cameras_np + old_anchor - new_anchor
|
|
cameras = torch.from_numpy(cameras_np).to(device=device, dtype=kp3d.dtype)
|
|
focals = torch.tensor([
|
|
float(np.asarray(p["focal_length"]).reshape(-1)[0]) for p in persons
|
|
], device=device, dtype=kp3d.dtype)
|
|
principals = torch.tensor([
|
|
_principal_point(p) for p in persons
|
|
], device=device, dtype=kp3d.dtype)
|
|
|
|
def project(points: torch.Tensor) -> torch.Tensor:
|
|
camera_points = points + cameras[:, None, :]
|
|
z = camera_points[..., 2].clamp(min=1e-6)
|
|
u = principals[:, None, 0] + focals[:, None] * camera_points[..., 0] / z
|
|
v = principals[:, None, 1] + focals[:, None] * camera_points[..., 1] / z
|
|
return torch.stack([u, v], dim=-1)
|
|
|
|
kp2d = project(kp3d)
|
|
face2d = project(face3d)
|
|
arrays = {
|
|
"pred_vertices": verts.detach().float().cpu().numpy(),
|
|
"pred_keypoints_3d": kp3d.detach().float().cpu().numpy(),
|
|
"pred_keypoints_2d": kp2d.detach().float().cpu().numpy(),
|
|
"pred_face_keypoints_3d": face3d.detach().float().cpu().numpy(),
|
|
"pred_face_keypoints_2d": face2d.detach().float().cpu().numpy(),
|
|
"pred_joint_coords": joint_coords.detach().float().cpu().numpy(),
|
|
"pred_global_rots": joint_rotmats.detach().float().cpu().numpy(),
|
|
"mhr_model_params": model_params.detach().float().cpu().numpy(),
|
|
}
|
|
|
|
for index, (frame_index, index_in_frame, _) in enumerate(chunk):
|
|
person_out = frames_out[frame_index][index_in_frame]
|
|
person_out["shape_params"] = shapes[index]
|
|
person_out["scale_params"] = scales[index]
|
|
person_out["pred_cam_t"] = cameras_np[index].astype(np.float32)
|
|
for key, value in arrays.items():
|
|
person_out[key] = value[index]
|
|
if progress is not None:
|
|
progress()
|
|
|
|
out = dict(pose_data)
|
|
out["frames"] = frames_out
|
|
return out
|