from __future__ import annotations import math import pickle from typing import Optional, Sequence import numpy as np import torch from tqdm import tqdm from camera import Camera, PerspectiveCamera from camera.camera import focal2fov, fov2focal from flame_solver import FlameSolverOps, SolverContext from common.config import INVISIBLE_LOG_VAR, GNFlameOptimizerConfig from common.utils import load_expression_blendshapes, load_mesh_from_obj from common.types import FittingOutput, InferenceOutput EYE_REGION_VERTICES = [ # eyelid 807, 814, 822, 824, 827, 991, 994, 995, 1093, 1096, 1113, 1170, 1175, 1195, 1200, 1201, 1202, 1216, 1218, 1329, 1331, 1336, 1340, 1343, 1344, 1354, 1355, 1357, 3833, 3855, 2264, 2267, 2271, 2273, 2276, 2355, 2358, 2359, 2402, 2403, 2406, 2436, 2437, 2447, 2448, 2449, 2450, 2451, 2453, 2485, 2486, 2487, 2491, 2494, 2495, 2505, 2506, 2508, 3632, 3689, # iris 4478, 4482, 4486, 4490, 4494, 4498, 4502, 4506, 4510, 4514, 4518, 4522, 4526, 4530, 4534, 4538, 4542, 4546, 4550, 4554, 4558, 4562, 4566, 4570, 4574, 4578, 4582, 4586, 4590, 4594, 4599, 4603, 3932, 3936, 3940, 3944, 3948, 3952, 3956, 3960, 3964, 3968, 3972, 3976, 3980, 3984, 3988, 3992, 3996, 4000, 4004, 4008, 4012, 4016, 4020, 4024, 4028, 4032, 4036, 4040, 4044, 4048, 4053, 4057 ] class GNFlameOptimizer: HISTORY_COMPONENTS = ("align_loss", "depth_loss", "reg_expr", "reg_pose", "reg_ident", "total_loss") def __init__(self, config: GNFlameOptimizerConfig): self.config = config self.device = torch.device(config.device) self._load_flame_assets(config.flame) self.solver_ops = FlameSolverOps(num_expressions=self.num_expressions) self.solver_context = SolverContext.create(self.solver_ops.extension) self.key_vertex_ids = torch.tensor([], dtype=torch.long, device=self.device) # key_vertex_ids = np.unique(load_mesh_from_obj("weights/flame/flame_key_region.obj")[1].flatten()) # self.key_vertex_ids = torch.from_numpy(key_vertex_ids).to(dtype=torch.int64, device=self.device) self._empty_uncertainty = torch.empty((0,), dtype=torch.float32, device=self.device) self._eye_region_ids = torch.tensor(EYE_REGION_VERTICES, dtype=torch.long, device=self.device) self._eye_boost_log_var_delta = -2.0 * math.log(max(float(config.eye_region_boost_factor), 1e-8)) def _load_flame_assets(self, flame_config) -> None: self.num_expressions = flame_config.resolved_num_expressions() self.dim_x = self.num_expressions + 18 self.pose_start = self.num_expressions self.trans_start = self.num_expressions + 15 with open(flame_config.flame_model_path, "rb") as f: data = pickle.load(f, encoding="latin1") self.v_template = torch.from_numpy(np.asarray(data["v_template"])).to(dtype=torch.float32, device=self.device) self.faces = torch.from_numpy(np.asarray(data["f"]).astype(np.int64)).to(dtype=torch.int64, device=self.device) self.pose_dirs = torch.from_numpy(np.asarray(data["posedirs"])).to(dtype=torch.float32, device=self.device) self.lbs_weights = torch.from_numpy(np.asarray(data["weights"])).to(dtype=torch.float32, device=self.device) j_regressor_coo = data["J_regressor"].tocoo() self.J_regressor_row = torch.from_numpy(np.asarray(j_regressor_coo.row)).to(dtype=torch.int32, device=self.device) self.J_regressor_col = torch.from_numpy(np.asarray(j_regressor_coo.col)).to(dtype=torch.int32, device=self.device) self.J_regressor_values = torch.from_numpy(np.asarray(j_regressor_coo.data)).to(dtype=torch.float32, device=self.device) num_pose_basis = self.pose_dirs.shape[-1] shapedirs_np = np.asarray(data["shapedirs"]).astype(np.float32, copy=False) self.shape_dirs = torch.from_numpy(np.ascontiguousarray(shapedirs_np[..., :300])).to( dtype=torch.float32, device=self.device ).view(-1, 300).T.contiguous() expr_dirs_np = load_expression_blendshapes(data, flame_config) # (V, 3, num_expressions) self.expr_dirs = torch.from_numpy(expr_dirs_np).to(dtype=torch.float32, device=self.device) \ .view(-1, self.num_expressions).T.contiguous() self.pose_dirs = self.pose_dirs.reshape([-1, num_pose_basis]).T.contiguous() j_regressor_dense = torch.from_numpy(np.asarray(data["J_regressor"].todense())).to(dtype=torch.float32, device=self.device) self.joints_dirs = ( j_regressor_dense @ self.shape_dirs.view(300, 5023, 3).permute(1, 2, 0).reshape(5023, 300 * 3) ).view(5, 3, 300).permute(2, 0, 1).contiguous() def _apply_eye_region_boost(self, vertex_coord_log_var: torch.Tensor | None) -> torch.Tensor | None: if vertex_coord_log_var is None or not self.config.boost_eye_region: return vertex_coord_log_var vertex_coord_log_var = vertex_coord_log_var.clone() vertex_coord_log_var[..., self._eye_region_ids, :] += self._eye_boost_log_var_delta return vertex_coord_log_var def _calc_canonical_flame(self, identity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: return self.solver_ops.calc_canonical( identity, self.v_template, self.shape_dirs, self.J_regressor_row, self.J_regressor_col, self.J_regressor_values, ) def _calc_posed_flame( self, v_canonical: torch.Tensor, j_canonical: torch.Tensor, x: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: return self.solver_ops.forward_vertices( x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) @staticmethod def _check_finite(tensor: torch.Tensor | None, label: str) -> bool: """Return True if *tensor* is fully finite. Otherwise print a NaN/inf diagnostic line including stats over the finite slice, and return False. ``None`` is treated as finite (nothing to check).""" if tensor is None or not torch.is_tensor(tensor): return True finite_mask = torch.isfinite(tensor) if bool(finite_mask.all().item()): return True bad = (~finite_mask) n_bad = int(bad.sum().item()) n_total = int(tensor.numel()) n_nan = int(torch.isnan(tensor).sum().item()) n_pinf = int(torch.isposinf(tensor).sum().item()) n_ninf = int(torch.isneginf(tensor).sum().item()) finite_vals = tensor[finite_mask].detach().to(torch.float32) if finite_vals.numel() > 0: amin = float(finite_vals.min().item()) amax = float(finite_vals.max().item()) amean = float(finite_vals.mean().item()) stats = f"finite[min={amin:.4g} max={amax:.4g} mean={amean:.4g}]" else: stats = "no finite values" print( f"[GN][NaN!!] {label}: bad={n_bad}/{n_total} " f"(nan={n_nan}, +inf={n_pinf}, -inf={n_ninf}) " f"shape={tuple(tensor.shape)} dtype={tensor.dtype}; {stats}", flush=True, ) return False def _stack_camera_matrices(self, cameras: Sequence[Camera]) -> tuple[torch.Tensor, torch.Tensor]: if len(cameras) == 0: raise ValueError("At least one camera is required.") view_mats = [camera.get_w2v for camera in cameras] viewproj_mats = [camera.get_full_proj for camera in cameras] return ( torch.stack(view_mats, dim=0).to(self.device, torch.float32), torch.stack(viewproj_mats, dim=0).to(self.device, torch.float32) ) def _single_camera_matrices(self, camera: Camera) -> tuple[torch.Tensor, torch.Tensor]: view_mat = camera.get_w2v.to(self.device, torch.float32).unsqueeze(0) viewproj_mat = camera.get_full_proj.to(self.device, torch.float32).unsqueeze(0) return view_mat, viewproj_mat def _build_config_camera(self, image_width: int, image_height: int, fov_y: float) -> PerspectiveCamera: rot = self.config.camera_rotation.reshape(3, 3) pos = self.config.camera_position.reshape(3) return PerspectiveCamera( fov_y=math.radians(float(fov_y)), rot=rot, pos=pos, width=int(image_width), height=int(image_height), ) def _fov_y_to_focal(self, image_height: int, fov_y: float) -> float: return fov2focal(math.radians(float(fov_y)), float(image_height)) def _focal_to_fov_y(self, image_height: int, focal_y: float) -> float: return math.degrees(focal2fov(float(focal_y), float(image_height))) def _evaluate_energy( self, camera_view_mat: torch.Tensor, camera_viewproj_mat: torch.Tensor, vertex_coord: torch.Tensor, vertex_depth: torch.Tensor, vertex_coord_log_var: torch.Tensor | None, vertex_depth_log_var: torch.Tensor | None, identity: torch.Tensor, x: torch.Tensor, ) -> torch.Tensor: num_views = int(camera_view_mat.shape[0]) identity = identity.to(device=self.device, dtype=torch.float32).reshape(300) x = x.to(device=self.device, dtype=torch.float32) vertex_indices = self.key_vertex_ids num_items = int(vertex_indices.numel()) if vertex_indices.numel() > 0 else int(vertex_coord.shape[-2]) # CUDA assemble_jacobian writes obs_dim columns per item: 3 (uv + depth) # when depth_weight > 0, else 2 (uv only). See assemble_jacobian.cu:39, 205. obs_dim = 3 if float(self.config.relative_depth_weight) > 0.0 else 2 num_data_rows = num_views * num_items * obs_dim chunk_size = int(self.config.energy_chunk_size) def energy_from_residual(residual: torch.Tensor) -> torch.Tensor: if residual.ndim == 1: data_residual = residual[:num_data_rows].reshape(-1, obs_dim) return torch.square(data_residual).sum() if residual.ndim == 2: data_residual = residual[:, :num_data_rows].reshape(residual.shape[0], -1, obs_dim) return torch.square(data_residual).sum(dim=(1, 2)) raise ValueError("x residual must be [R] or [T,R].") v_canonical, j_canonical = self._calc_canonical_flame(identity) if x.ndim == 1: flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian( x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) residual, _ = self.solver_ops.assemble_x_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, x, flame_v_jacobian, flame_j_jacobian, self.config.correspondence_weight, self.config.relative_depth_weight, self.config.expression_regularization, self.config.pose_regularization, vertex_indices, ) return energy_from_residual(residual) if chunk_size <= 0: raise ValueError("sequence_energy_chunk_size must be positive.") batch_size = int(x.shape[0]) chunk_energies: list[torch.Tensor] = [] for start in range(0, batch_size, chunk_size): end = min(start + chunk_size, batch_size) x_chunk = x[start:end] vertex_coord_chunk = vertex_coord[start:end] vertex_depth_chunk = vertex_depth[start:end] if vertex_coord_log_var is None or vertex_depth_log_var is None: vertex_coord_log_var_chunk = vertex_coord_log_var vertex_depth_log_var_chunk = vertex_depth_log_var else: vertex_coord_log_var_chunk = vertex_coord_log_var[start:end] vertex_depth_log_var_chunk = vertex_depth_log_var[start:end] flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian( x_chunk, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) residual, _ = self.solver_ops.assemble_x_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord_chunk, vertex_depth_chunk, vertex_coord_log_var_chunk, vertex_depth_log_var_chunk, x_chunk, flame_v_jacobian, flame_j_jacobian, self.config.correspondence_weight, self.config.relative_depth_weight, self.config.expression_regularization, self.config.pose_regularization, vertex_indices, ) chunk_energies.append(energy_from_residual(residual)) return torch.cat(chunk_energies, dim=0) def _register_wi_fov_search( self, image_width: int, image_height: int, vertex_coord: torch.Tensor, vertex_depth: torch.Tensor, vertex_coord_log_var: torch.Tensor | None, vertex_depth_log_var: torch.Tensor | None, init_identity: Optional[torch.Tensor] = None, init_x: Optional[torch.Tensor] = None, pose_init: bool = True, ) -> tuple[float, torch.Tensor, torch.Tensor, torch.Tensor]: pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(init_identity) def evaluate_fov(fov_y: float) -> float: camera = self._build_config_camera(image_width, image_height, fov_y) camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera) # Keep the search objective comparable across candidate cameras. if pose_init: pose_x = self._register_pose( camera_view_mat, camera_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, pose_v_canonical, pose_j_canonical, init_x=init_x, ) else: pose_x = init_x identity, x = self._register( camera_view_mat, camera_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, init_identity=init_identity, init_x=pose_x, num_iters=self.config.fov_registration_iterations, ) return float( self._evaluate_energy( camera_view_mat, camera_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, identity, x, ).mean().item() ) l = float(self.config.fov_y_min) r = float(self.config.fov_y_max) phi = 0.5 * (math.sqrt(5.0) - 1.0) x1 = r - phi * (r - l) x2 = l + phi * (r - l) e1 = evaluate_fov(x1) e2 = evaluate_fov(x2) for iter_id in range(self.config.fov_search_iterations): if r - l < self.config.fov_search_epsilon: break if e1 < e2: r = x2 x2 = x1 e2 = e1 x1 = r - phi * (r - l) e1 = evaluate_fov(x1) else: l = x1 x1 = x2 e1 = e2 x2 = l + phi * (r - l) e2 = evaluate_fov(x2) best_fov_y = 0.5 * (l + r) best_focal = self._fov_y_to_focal(image_height, best_fov_y) final_camera = self._build_config_camera(image_width, image_height, best_fov_y) final_view_mat, final_viewproj_mat = self._single_camera_matrices(final_camera) if pose_init: pose_x = self._register_pose( final_view_mat, final_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, pose_v_canonical, pose_j_canonical, init_x=init_x, ) else: pose_x = init_x final_identity, final_x = self._register( final_view_mat, final_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, init_identity=init_identity, init_x=pose_x, ) return best_fov_y, final_identity, final_x, pose_x def _register( self, camera_view_mat: torch.Tensor, camera_viewproj_mat: torch.Tensor, vertex_coord: torch.Tensor, vertex_depth: torch.Tensor, vertex_coord_log_var: torch.Tensor | None = None, vertex_depth_log_var: torch.Tensor | None = None, init_identity: Optional[torch.Tensor] = None, init_x: Optional[torch.Tensor] = None, num_iters: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: num_views = int(camera_view_mat.shape[-3]) if vertex_coord.ndim == 2 and vertex_depth.ndim == 2: shape = (self.dim_x,) elif vertex_coord.ndim == 3 and vertex_depth.ndim == 3: if vertex_coord.shape[0] == num_views: shape = (self.dim_x,) else: shape = (vertex_coord.shape[0], self.dim_x) elif vertex_coord.ndim == 4 and vertex_depth.ndim == 4: shape = (vertex_coord.shape[0], self.dim_x) else: raise ValueError("Invalid shape of vertex_coord or vertex_depth.") if init_identity is None: identity = torch.zeros((300,), dtype=torch.float32, device=self.device) else: identity = init_identity.to(device=self.device, dtype=torch.float32).reshape(300).clone() if init_x is None: x = torch.zeros(shape, dtype=torch.float32, device=self.device) else: x = init_x.to(device=self.device, dtype=torch.float32).reshape(shape).clone() effective_num_iters = self.config.registration_iterations if num_iters is None else int(num_iters) num_batch_frames = x.shape[0] if x.ndim == 2 else 1 for it in range(effective_num_iters): v_canonical, j_canonical = self._calc_canonical_flame(identity) damping_lambda = self.config.damping flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian( x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights ) residual, jacobian = self.solver_ops.assemble_x_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, x, flame_v_jacobian, flame_j_jacobian, self.config.correspondence_weight, self.config.relative_depth_weight, self.config.expression_regularization, self.config.pose_regularization, self.key_vertex_ids, ) x = x + self.solver_ops.solve_normal_equation( self.solver_context, jacobian, residual, damping_lambda, ) flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.identity_jacobian( x, v_canonical, j_canonical, self.joints_dirs, self.shape_dirs, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) residual, jacobian = self.solver_ops.assemble_identity_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, identity if init_identity is None else identity - init_identity, flame_v_jacobian, flame_j_jacobian, self.config.correspondence_weight, self.config.relative_depth_weight, self.config.identity_regularization, self.key_vertex_ids, ) identity = identity + self.solver_ops.solve_normal_equation( self.solver_context, jacobian, residual, damping_lambda * num_batch_frames, ) return identity, x def _register_pose( self, camera_view_mat: torch.Tensor, camera_viewproj_mat: torch.Tensor, vertex_coord: torch.Tensor, vertex_depth: torch.Tensor, vertex_coord_log_var: torch.Tensor | None, vertex_depth_log_var: torch.Tensor | None, v_canonical: torch.Tensor, j_canonical: torch.Tensor, init_x: Optional[torch.Tensor] = None, num_iters: int | None = None, ) -> torch.Tensor: num_views = int(camera_view_mat.shape[0]) if vertex_coord.ndim == 2 and vertex_depth.ndim == 2: shape = (self.dim_x,) elif vertex_coord.ndim == 3 and vertex_depth.ndim == 3: if vertex_coord.shape[0] == num_views: shape = (self.dim_x,) else: shape = (vertex_coord.shape[0], self.dim_x) elif vertex_coord.ndim == 4 and vertex_depth.ndim == 4: shape = (vertex_coord.shape[0], self.dim_x) else: raise ValueError("Invalid shape of vertex_coord or vertex_depth.") if init_x is None: x = torch.zeros(shape, dtype=torch.float32, device=self.device) else: x = init_x.to(device=self.device, dtype=torch.float32).reshape(shape).clone() opt_dims = torch.tensor( [ self.pose_start, self.pose_start + 1, self.pose_start + 2, self.trans_start, self.trans_start + 1, self.trans_start + 2, ], dtype=torch.int64, device=self.device, ) effective_num_iters = self.config.pose_registration_iterations if num_iters is None else int(num_iters) for _ in range(effective_num_iters): flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian( x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) residual, jacobian = self.solver_ops.assemble_x_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, x, flame_v_jacobian, flame_j_jacobian, alignment_weight=1.0, rel_depth_weight=0.0, exp_reg_weight=0.0, pose_reg_weight=0.0, vertex_indices=self.key_vertex_ids, ) delta = self.solver_ops.solve_normal_equation( self.solver_context, jacobian[..., opt_dims], residual, self.config.pose_damping, ) x[..., opt_dims] = x[..., opt_dims] + delta return x def _track( self, camera_view_mat: torch.Tensor, camera_viewproj_mat: torch.Tensor, vertex_coord: torch.Tensor, vertex_depth: torch.Tensor, vertex_coord_log_var: torch.Tensor | None, vertex_depth_log_var: torch.Tensor | None, v_canonical: torch.Tensor, j_canonical: torch.Tensor, x: torch.Tensor, ) -> tuple[torch.Tensor, int]: # Root-pose pre-step: push bulk rotation into root joint first root_pose_dims = torch.tensor( [self.pose_start, self.pose_start + 1, self.pose_start + 2, self.trans_start, self.trans_start + 1, self.trans_start + 2], dtype=torch.int64, device=self.device, ) for _ in range(self.config.root_pose_tracking_iterations): flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian( x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) residual, jacobian = self.solver_ops.assemble_x_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, x, flame_v_jacobian, flame_j_jacobian, self.config.correspondence_weight, self.config.relative_depth_weight, 0.0, 0.0, self.key_vertex_ids, ) delta = self.solver_ops.solve_normal_equation( self.solver_context, jacobian[..., root_pose_dims], residual, self.config.damping, ) x[..., root_pose_dims] = x[..., root_pose_dims] + delta # Full optimization with early stopping prev_energy = float('inf') steps = 0 for _ in range(self.config.tracking_iterations): flame_v, flame_j, flame_v_jacobian, flame_j_jacobian = self.solver_ops.expression_jacobian( x, v_canonical, j_canonical, self.expr_dirs, self.pose_dirs, self.lbs_weights, ) residual, jacobian = self.solver_ops.assemble_x_system( camera_view_mat, camera_viewproj_mat, flame_v, flame_j, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, x, flame_v_jacobian, flame_j_jacobian, self.config.correspondence_weight, self.config.relative_depth_weight, self.config.expression_regularization, self.config.pose_regularization, self.key_vertex_ids, ) energy = float(torch.sum(residual ** 2).item()) if prev_energy - energy < self.config.tracking_early_stop_ratio * prev_energy: break prev_energy = energy steps += 1 x = x + self.solver_ops.solve_normal_equation( self.solver_context, jacobian, residual, self.config.damping, ) return x, steps def _track_pass( self, camera_view_mat: torch.Tensor, camera_viewproj_mat: torch.Tensor, identity: torch.Tensor, init_x: torch.Tensor, vertex_coords: torch.Tensor, vertex_depths: torch.Tensor, vertex_coord_log_vars: torch.Tensor | None, vertex_depth_log_vars: torch.Tensor | None, desc: str = "track_pass" ) -> tuple[torch.Tensor, torch.Tensor]: num_frames = vertex_coords.shape[0] xs = torch.zeros((num_frames, self.dim_x), dtype=torch.float32, device=self.device) per_frame_init = init_x.ndim == 2 if not per_frame_init: x = init_x.clone() v_canonical, j_canonical = self._calc_canonical_flame(identity) step_counts = [] for frame_id in tqdm(range(num_frames), desc=desc, leave=False): if per_frame_init: x = init_x[frame_id].clone() if vertex_coord_log_vars is None or vertex_depth_log_vars is None: frame_projected_log_var = vertex_coord_log_vars frame_rel_depth_log_var = vertex_depth_log_vars else: frame_projected_log_var = vertex_coord_log_vars[frame_id] frame_rel_depth_log_var = vertex_depth_log_vars[frame_id] x, steps = self._track( camera_view_mat, camera_viewproj_mat, vertex_coords[frame_id], vertex_depths[frame_id], frame_projected_log_var, frame_rel_depth_log_var, v_canonical, j_canonical, x, ) step_counts.append(steps) xs[frame_id] = x print(f" [{desc}] steps: min={min(step_counts)} max={max(step_counts)} mean={np.mean(step_counts):.1f}") return xs def _select_keyframes(self, identity: torch.Tensor, xs: torch.Tensor, num_keyframes: int) -> torch.Tensor: num_frames = xs.shape[0] xs_eval = xs.clone() xs_eval[:, self.trans_start:self.dim_x] = 0 v_canonical, j_canonical = self._calc_canonical_flame(identity) vertices, _ = self._calc_posed_flame(v_canonical, j_canonical, xs_eval) flat_vertices = vertices.flatten(1) key_frames = [] key_frame = int(np.random.randint(num_frames)) dists = torch.sum((flat_vertices - flat_vertices[key_frame].unsqueeze(0)) ** 2, dim=-1) dists[key_frame] = -1.0 key_frames.append(key_frame) for _ in range(num_keyframes - 1): key_frame = int(torch.argmax(dists).item()) dists_to_this = torch.sum((flat_vertices - flat_vertices[key_frame].unsqueeze(0)) ** 2, dim=-1) dists = torch.minimum(dists, dists_to_this) dists[key_frame] = -1.0 key_frames.append(key_frame) return torch.as_tensor(key_frames, dtype=torch.long, device=self.device) def fit_frame(self, inference: InferenceOutput, camera: Camera | None = None, init_identity: np.ndarray | None = None, ) -> FittingOutput: if inference.vertex_coord.ndim != 2 or inference.vertex_coord.shape[-1] != 2: raise ValueError("fit_frame expects inference.vertex_coord with shape [5023, 2].") if inference.vertex_depth.ndim != 2 or inference.vertex_depth.shape[-1] != 1: raise ValueError("fit_frame expects inference.vertex_depth with shape [5023, 1].") if self.config.use_uncertainty_weights: if inference.vertex_coord_log_var.ndim != 2 or inference.vertex_coord_log_var.shape[-1] != 1: raise ValueError("fit_frame expects inference.vertex_coord_log_var with shape [5023, 1].") if inference.vertex_depth_log_var.ndim != 2 or inference.vertex_depth_log_var.shape[-1] != 1: raise ValueError("fit_frame expects inference.vertex_depth_log_var with shape [5023, 1].") vertex_coord = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32) vertex_depth = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32) vertex_coord[..., 0] /= float(inference.image_width) vertex_coord[..., 1] /= float(inference.image_height) if self.config.use_uncertainty_weights: vertex_coord_log_var: torch.Tensor | None = torch.from_numpy(inference.vertex_coord_log_var).to(device=self.device, dtype=torch.float32) vertex_depth_log_var: torch.Tensor | None = torch.from_numpy(inference.vertex_depth_log_var).to(device=self.device, dtype=torch.float32) else: vertex_coord_log_var = None vertex_depth_log_var = None vertex_coord_log_var = self._apply_eye_region_boost(vertex_coord_log_var) if init_identity is None: pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device) else: pose_identity = torch.from_numpy(init_identity).to(dtype=torch.float32, device=self.device) camera_fov_y = None if camera is None: camera_fov_y, identity, x, pose_x = self._register_wi_fov_search( inference.image_width, inference.image_height, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, init_identity=pose_identity, ) else: camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera) pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(pose_identity) pose_x = self._register_pose( camera_view_mat, camera_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, pose_v_canonical, pose_j_canonical, ) identity, x = self._register( camera_view_mat, camera_viewproj_mat, vertex_coord, vertex_depth, vertex_coord_log_var, vertex_depth_log_var, init_identity=pose_identity, init_x=pose_x, ) return FittingOutput(identity=identity, x=x, pose_x=pose_x, camera_fov_y=camera_fov_y) def fit_frame_mv(self, inferences: list[InferenceOutput], cameras: list[Camera]) -> FittingOutput: if len(inferences) == 0: raise ValueError("fit_frame_mv expects at least one InferenceOutput.") if len(inferences) != len(cameras): raise ValueError("fit_frame_mv expects the same number of inferences and cameras.") if any(camera is None for camera in cameras): raise ValueError("fit_frame_mv requires an explicit camera for each view.") vertex_coords = [] vertex_depths = [] vertex_coord_log_vars = [] vertex_depth_log_vars = [] for inference in inferences: if inference.vertex_coord.ndim != 2 or inference.vertex_coord.shape[-1] != 2: raise ValueError("fit_frame_mv expects each inference.vertex_coord with shape [5023, 2].") if inference.vertex_depth.ndim != 2 or inference.vertex_depth.shape[-1] != 1: raise ValueError("fit_frame_mv expects each inference.vertex_depth with shape [5023, 1].") if self.config.use_uncertainty_weights: if inference.vertex_coord_log_var.ndim != 2 or inference.vertex_coord_log_var.shape[-1] != 1: raise ValueError("fit_frame_mv expects each inference.vertex_coord_log_var with shape [5023, 1].") if inference.vertex_depth_log_var.ndim != 2 or inference.vertex_depth_log_var.shape[-1] != 1: raise ValueError("fit_frame_mv expects each inference.vertex_depth_log_var with shape [5023, 1].") vertex_coord = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32) vertex_depth = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32) vertex_coord[..., 0] /= float(inference.image_width) vertex_coord[..., 1] /= float(inference.image_height) vertex_coords.append(vertex_coord) vertex_depths.append(vertex_depth) if self.config.use_uncertainty_weights: vertex_coord_log_vars.append( torch.from_numpy(inference.vertex_coord_log_var).to(device=self.device, dtype=torch.float32) ) vertex_depth_log_vars.append( torch.from_numpy(inference.vertex_depth_log_var).to(device=self.device, dtype=torch.float32) ) vertex_coords_tensor = torch.stack(vertex_coords, dim=0) # [C, V, 2] vertex_depths_tensor = torch.stack(vertex_depths, dim=0) # [C, V, 1] if self.config.use_uncertainty_weights: vertex_coord_log_vars_tensor: torch.Tensor | None = torch.stack(vertex_coord_log_vars, dim=0) # [C, V, 1] vertex_depth_log_vars_tensor: torch.Tensor | None = torch.stack(vertex_depth_log_vars, dim=0) # [C, V, 1] else: vertex_coord_log_vars_tensor = None vertex_depth_log_vars_tensor = None vertex_coord_log_vars_tensor = self._apply_eye_region_boost(vertex_coord_log_vars_tensor) camera_view_mat, camera_viewproj_mat = self._stack_camera_matrices(cameras) pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device) pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(pose_identity) pose_x = self._register_pose( camera_view_mat, camera_viewproj_mat, vertex_coords_tensor, vertex_depths_tensor, vertex_coord_log_vars_tensor, vertex_depth_log_vars_tensor, pose_v_canonical, pose_j_canonical, ) identity, x = self._register( camera_view_mat, camera_viewproj_mat, vertex_coords_tensor, vertex_depths_tensor, vertex_coord_log_vars_tensor, vertex_depth_log_vars_tensor, init_x=pose_x, ) return FittingOutput(identity=identity, x=x) def fit_sequence(self, inference: InferenceOutput, camera: Camera | None = None, init_identity: np.ndarray | None = None, ) -> FittingOutput: if inference.vertex_coord.ndim != 3 or inference.vertex_coord.shape[-1] != 2: raise ValueError("fit_sequence expects inference.vertex_coord with shape [T, 5023, 2].") if inference.vertex_depth.ndim != 3 or inference.vertex_depth.shape[-1] != 1: raise ValueError("fit_sequence expects inference.vertex_depth with shape [T, 5023, 1].") if self.config.use_uncertainty_weights: if inference.vertex_coord_log_var.ndim != 3 or inference.vertex_coord_log_var.shape[-1] != 1: raise ValueError("fit_sequence expects inference.vertex_coord_log_var with shape [T, 5023, 1].") if inference.vertex_depth_log_var.ndim != 3 or inference.vertex_depth_log_var.shape[-1] != 1: raise ValueError("fit_sequence expects inference.vertex_depth_log_var with shape [T, 5023, 1].") vertex_coords = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32) vertex_depths = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32) vertex_coords[..., 0] /= float(inference.image_width) vertex_coords[..., 1] /= float(inference.image_height) if self.config.use_uncertainty_weights: vertex_coord_log_vars: torch.Tensor | None = torch.from_numpy(inference.vertex_coord_log_var).to(device=self.device, dtype=torch.float32) vertex_depth_log_vars: torch.Tensor | None = torch.from_numpy(inference.vertex_depth_log_var).to(device=self.device, dtype=torch.float32) else: vertex_coord_log_vars = None vertex_depth_log_vars = None vertex_coord_log_vars = self._apply_eye_region_boost(vertex_coord_log_vars) if vertex_coord_log_vars is not None: frame0_coord_log_var: torch.Tensor | None = vertex_coord_log_vars[0] frame0_depth_log_var: torch.Tensor | None = vertex_depth_log_vars[0] else: frame0_coord_log_var = None frame0_depth_log_var = None fit_camera = camera is None if init_identity is None: init_identity = torch.zeros((300,), dtype=torch.float32, device=self.device) else: init_identity = torch.from_numpy(init_identity).to(dtype=torch.float32, device=self.device) init_v_canonical, init_j_canonical = self._calc_canonical_flame(init_identity) print("[GN] register first frame") camera_fov_y = None if fit_camera: camera_fov_y, identity, x0, pose_x0 = self._register_wi_fov_search( inference.image_width, inference.image_height, vertex_coords[0], vertex_depths[0], frame0_coord_log_var, frame0_depth_log_var, init_identity=init_identity, ) print(f"[GN] optimized camera fov_y after first frame register: {camera_fov_y:.6f} deg") camera = self._build_config_camera(inference.image_width, inference.image_height, camera_fov_y) camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera) else: camera_view_mat, camera_viewproj_mat = self._single_camera_matrices(camera) pose_x0 = self._register_pose( camera_view_mat, camera_viewproj_mat, vertex_coords[0], vertex_depths[0], frame0_coord_log_var, frame0_depth_log_var, init_v_canonical, init_j_canonical, ) identity, x0 = self._register( camera_view_mat, camera_viewproj_mat, vertex_coords[0], vertex_depths[0], frame0_coord_log_var, frame0_depth_log_var, init_identity=init_identity, init_x=pose_x0, ) xs = self._track_pass( camera_view_mat, camera_viewproj_mat, identity, x0, vertex_coords, vertex_depths, vertex_coord_log_vars, vertex_depth_log_vars, desc="[GN] track pass 0" ) num_frames = xs.shape[0] for round_id in range(self.config.identity_refinement_rounds): print(f"[GN] refine round {round_id + 1}/{self.config.identity_refinement_rounds}") # select keyframes if num_frames > self.config.num_keyframes: key_frames = self._select_keyframes(identity, xs, self.config.num_keyframes) xs_key = xs[key_frames] vertex_coords_key = vertex_coords[key_frames] vertex_depths_key = vertex_depths[key_frames] if vertex_coord_log_vars is None or vertex_depth_log_vars is None: vertex_coord_log_vars_key = None vertex_depth_log_vars_key = None else: vertex_coord_log_vars_key = vertex_coord_log_vars[key_frames] vertex_depth_log_vars_key = vertex_depth_log_vars[key_frames] else: xs_key = xs vertex_coords_key = vertex_coords vertex_depths_key = vertex_depths vertex_coord_log_vars_key = vertex_coord_log_vars vertex_depth_log_vars_key = vertex_depth_log_vars identity, _ = self._register( camera_view_mat, camera_viewproj_mat, vertex_coords_key, vertex_depths_key, vertex_coord_log_vars_key, vertex_depth_log_vars_key, init_identity=identity, init_x=xs_key ) # tracking pass xs = self._track_pass( camera_view_mat, camera_viewproj_mat, identity, xs, vertex_coords, vertex_depths, vertex_coord_log_vars, vertex_depth_log_vars, desc=f"[GN] track pass {round_id + 1}" ) return FittingOutput(identity=identity, x=xs, camera_fov_y=camera_fov_y) def fit_sequence_mv(self, inferences: list[InferenceOutput], cameras: list[Camera]) -> FittingOutput: if len(inferences) == 0: raise ValueError("fit_sequence_mv expects at least one InferenceOutput.") if len(inferences) != len(cameras): raise ValueError("fit_sequence_mv expects the same number of inferences and cameras.") if any(camera is None for camera in cameras): raise ValueError("fit_sequence_mv requires an explicit camera for each view.") first_inference = inferences[0] if first_inference.vertex_coord.ndim != 3 or first_inference.vertex_coord.shape[-1] != 2: raise ValueError("fit_sequence_mv expects each inference.vertex_coord with shape [T, 5023, 2].") if first_inference.vertex_depth.ndim != 3 or first_inference.vertex_depth.shape[-1] != 1: raise ValueError("fit_sequence_mv expects each inference.vertex_depth with shape [T, 5023, 1].") if self.config.use_uncertainty_weights: if first_inference.vertex_coord_log_var.ndim != 3 or first_inference.vertex_coord_log_var.shape[-1] != 1: raise ValueError("fit_sequence_mv expects each inference.vertex_coord_log_var with shape [T, 5023, 1].") if first_inference.vertex_depth_log_var.ndim != 3 or first_inference.vertex_depth_log_var.shape[-1] != 1: raise ValueError("fit_sequence_mv expects each inference.vertex_depth_log_var with shape [T, 5023, 1].") num_views = len(inferences) num_frames = first_inference.vertex_coord.shape[0] num_vertices = first_inference.vertex_coord.shape[1] vertex_coords_tensor = torch.empty((num_frames, num_views, num_vertices, 2), dtype=torch.float32, device=self.device) vertex_depths_tensor = torch.empty((num_frames, num_views, num_vertices, 1), dtype=torch.float32, device=self.device) if self.config.use_uncertainty_weights: vertex_coord_log_vars_tensor: torch.Tensor | None = torch.empty( (num_frames, num_views, num_vertices, 1), dtype=torch.float32, device=self.device ) vertex_depth_log_vars_tensor: torch.Tensor | None = torch.empty( (num_frames, num_views, num_vertices, 1), dtype=torch.float32, device=self.device ) else: vertex_coord_log_vars_tensor = None vertex_depth_log_vars_tensor = None for view_id, inference in enumerate(inferences): if inference.vertex_coord.ndim != 3 or inference.vertex_coord.shape[-1] != 2: raise ValueError("fit_sequence_mv expects each inference.vertex_coord with shape [T, 5023, 2].") if inference.vertex_depth.ndim != 3 or inference.vertex_depth.shape[-1] != 1: raise ValueError("fit_sequence_mv expects each inference.vertex_depth with shape [T, 5023, 1].") if self.config.use_uncertainty_weights: if inference.vertex_coord_log_var.ndim != 3 or inference.vertex_coord_log_var.shape[-1] != 1: raise ValueError("fit_sequence_mv expects each inference.vertex_coord_log_var with shape [T, 5023, 1].") if inference.vertex_depth_log_var.ndim != 3 or inference.vertex_depth_log_var.shape[-1] != 1: raise ValueError("fit_sequence_mv expects each inference.vertex_depth_log_var with shape [T, 5023, 1].") if inference.vertex_coord.shape[0] != num_frames: raise ValueError("fit_sequence_mv expects all views to share the same frame count.") vertex_coords = torch.from_numpy(inference.vertex_coord).to(device=self.device, dtype=torch.float32) vertex_depths = torch.from_numpy(inference.vertex_depth).to(device=self.device, dtype=torch.float32) vertex_coords[..., 0] /= float(inference.image_width) vertex_coords[..., 1] /= float(inference.image_height) vertex_coords_tensor[:, view_id] = vertex_coords vertex_depths_tensor[:, view_id] = vertex_depths if self.config.use_uncertainty_weights: vertex_coord_log_vars_tensor[:, view_id] = torch.from_numpy(inference.vertex_coord_log_var).to( device=self.device, dtype=torch.float32 ) vertex_depth_log_vars_tensor[:, view_id] = torch.from_numpy(inference.vertex_depth_log_var).to( device=self.device, dtype=torch.float32 ) vertex_coord_log_vars_tensor = self._apply_eye_region_boost(vertex_coord_log_vars_tensor) if self.config.use_uncertainty_weights: frame0_coord_log_var = vertex_coord_log_vars_tensor[0] frame0_depth_log_var = vertex_depth_log_vars_tensor[0] else: frame0_coord_log_var = None frame0_depth_log_var = None camera_view_mat, camera_viewproj_mat = self._stack_camera_matrices(cameras) print("[GN] register first frame (multi-view)") pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device) pose_v_canonical, pose_j_canonical = self._calc_canonical_flame(pose_identity) pose_x0 = self._register_pose( camera_view_mat, camera_viewproj_mat, vertex_coords_tensor[0], vertex_depths_tensor[0], frame0_coord_log_var, frame0_depth_log_var, pose_v_canonical, pose_j_canonical, ) self._check_finite(pose_x0, "after _register_pose: pose_x0") identity, x0 = self._register( camera_view_mat, camera_viewproj_mat, vertex_coords_tensor[0], vertex_depths_tensor[0], frame0_coord_log_var, frame0_depth_log_var, init_x=pose_x0, ) self._check_finite(identity, "after first _register: identity") self._check_finite(x0, "after first _register: x0") xs = self._track_pass( camera_view_mat, camera_viewproj_mat, identity, x0, vertex_coords_tensor, vertex_depths_tensor, vertex_coord_log_vars_tensor, vertex_depth_log_vars_tensor, desc="[GN] track pass 0 (multi-view)" ) self._check_finite(xs, "after track pass 0: xs") num_frames = xs.shape[0] for round_id in range(self.config.identity_refinement_rounds): print(f"[GN] refine round {round_id + 1}/{self.config.identity_refinement_rounds} (multi-view)") # select key frames if num_frames > self.config.num_keyframes: key_frames = self._select_keyframes(identity, xs, self.config.num_keyframes) xs_key = xs[key_frames] vertex_coords_key = vertex_coords_tensor[key_frames] vertex_depths_key = vertex_depths_tensor[key_frames] if vertex_coord_log_vars_tensor is None or vertex_depth_log_vars_tensor is None: vertex_coord_log_vars_key = None vertex_depth_log_vars_key = None else: vertex_coord_log_vars_key = vertex_coord_log_vars_tensor[key_frames] vertex_depth_log_vars_key = vertex_depth_log_vars_tensor[key_frames] else: xs_key = xs vertex_coords_key = vertex_coords_tensor vertex_depths_key = vertex_depths_tensor vertex_coord_log_vars_key = vertex_coord_log_vars_tensor vertex_depth_log_vars_key = vertex_depth_log_vars_tensor self._check_finite(xs_key, f"refine {round_id + 1}: xs_key (after keyframe select)") # register pass identity, _ = self._register( camera_view_mat, camera_viewproj_mat, vertex_coords_key, vertex_depths_key, vertex_coord_log_vars_key, vertex_depth_log_vars_key, init_identity=identity, init_x=xs_key, ) self._check_finite(identity, f"refine {round_id + 1}: identity (after _register)") # tracking pass xs = self._track_pass( camera_view_mat, camera_viewproj_mat, identity, xs, vertex_coords_tensor, vertex_depths_tensor, vertex_coord_log_vars_tensor, vertex_depth_log_vars_tensor, desc=f"[GN] track pass {round_id + 1} (multi-view)" ) self._check_finite(xs, f"refine {round_id + 1}: xs (after track pass)") return FittingOutput(identity=identity, x=xs) # ---- Robust multi-view tracking (adaptive per-frame camera selection) ---- def _track_pass_robust( self, all_cameras: list[Camera], identity: torch.Tensor, init_x: torch.Tensor, per_view_coords: list[torch.Tensor], per_view_depths: list[torch.Tensor], per_view_coord_log_vars: list[torch.Tensor] | None, per_view_depth_log_vars: list[torch.Tensor] | None, per_frame_views: list[list[int]], frame_reliable: torch.Tensor, desc: str = "track_pass_robust", ) -> torch.Tensor: num_frames = len(per_frame_views) xs = torch.zeros((num_frames, self.dim_x), dtype=torch.float32, device=self.device) per_frame_init = init_x.ndim == 2 if not per_frame_init: x = init_x.clone() v_canonical, j_canonical = self._calc_canonical_flame(identity) step_counts = [] for t in tqdm(range(num_frames), desc=desc, leave=False): if per_frame_init: x = init_x[t].clone() if not frame_reliable[t]: xs[t] = x continue views = per_frame_views[t] cam_view, cam_proj = self._stack_camera_matrices([all_cameras[v] for v in views]) frame_coords = torch.stack([per_view_coords[v][t] for v in views]) frame_depths = torch.stack([per_view_depths[v][t] for v in views]) if per_view_coord_log_vars is not None: frame_coord_lv = torch.stack([per_view_coord_log_vars[v][t] for v in views]) frame_depth_lv = torch.stack([per_view_depth_log_vars[v][t] for v in views]) else: frame_coord_lv = None frame_depth_lv = None x, steps = self._track( cam_view, cam_proj, frame_coords, frame_depths, frame_coord_lv, frame_depth_lv, v_canonical, j_canonical, x, ) step_counts.append(steps) xs[t] = x print(f" [{desc}] steps: min={min(step_counts)} max={max(step_counts)} mean={np.mean(step_counts):.1f}") return xs def _gather_register_data( self, keyframe_indices: torch.Tensor, per_frame_views: list[list[int]], all_cameras: list[Camera], per_view_coords: list[torch.Tensor], per_view_depths: list[torch.Tensor], per_view_coord_log_vars: list[torch.Tensor] | None, per_view_depth_log_vars: list[torch.Tensor] | None, num_vertices: int, ): key_ts = keyframe_indices.cpu().tolist() n_keys = len(key_ts) per_key_views = [per_frame_views[t] for t in key_ts] c_max = max(len(v) for v in per_key_views) cam_views_list = [] cam_projs_list = [] for ki in range(n_keys): views = per_key_views[ki] cv, cp = self._stack_camera_matrices([all_cameras[v] for v in views]) if len(views) < c_max: pad = c_max - len(views) cv = torch.cat([cv, cv[:1].expand(pad, -1, -1)], dim=0) cp = torch.cat([cp, cp[:1].expand(pad, -1, -1)], dim=0) cam_views_list.append(cv) cam_projs_list.append(cp) cam_view = torch.stack(cam_views_list) cam_proj = torch.stack(cam_projs_list) coords = torch.full((n_keys, c_max, num_vertices, 2), 0.5, dtype=torch.float32, device=self.device) depths = torch.zeros((n_keys, c_max, num_vertices, 1), dtype=torch.float32, device=self.device) coord_lvs = torch.full((n_keys, c_max, num_vertices, 1), INVISIBLE_LOG_VAR, dtype=torch.float32, device=self.device) depth_lvs = torch.full((n_keys, c_max, num_vertices, 1), INVISIBLE_LOG_VAR, dtype=torch.float32, device=self.device) for ki, t in enumerate(key_ts): for ci, v in enumerate(per_key_views[ki]): coords[ki, ci] = per_view_coords[v][t] depths[ki, ci] = per_view_depths[v][t] if per_view_coord_log_vars is not None: coord_lvs[ki, ci] = per_view_coord_log_vars[v][t] depth_lvs[ki, ci] = per_view_depth_log_vars[v][t] if per_view_coord_log_vars is None: return cam_view, cam_proj, coords, depths, None, None return cam_view, cam_proj, coords, depths, coord_lvs, depth_lvs def fit_sequence_mv_robust( self, inferences: list[InferenceOutput], cameras: list[Camera], init_x: Optional[torch.Tensor] = None, n_max_views: int = 16, n_min_views: int = 3, ) -> FittingOutput: num_all_views = len(inferences) num_frames = inferences[0].vertex_coord.shape[0] num_vertices = inferences[0].vertex_coord.shape[1] # --- Per-view data to GPU tensors --- per_view_coords: list[torch.Tensor] = [] per_view_depths: list[torch.Tensor] = [] per_view_coord_log_vars: list[torch.Tensor] | None = [] per_view_depth_log_vars: list[torch.Tensor] | None = [] for inference in inferences: vc = torch.from_numpy(np.asarray(inference.vertex_coord)).to( device=self.device, dtype=torch.float32) vd = torch.from_numpy(np.asarray(inference.vertex_depth)).to( device=self.device, dtype=torch.float32) vc[..., 0] /= float(inference.image_width) vc[..., 1] /= float(inference.image_height) per_view_coords.append(vc) per_view_depths.append(vd) if self.config.use_uncertainty_weights: per_view_coord_log_vars.append( torch.from_numpy(np.asarray(inference.vertex_coord_log_var)).to( device=self.device, dtype=torch.float32)) per_view_depth_log_vars.append( torch.from_numpy(np.asarray(inference.vertex_depth_log_var)).to( device=self.device, dtype=torch.float32)) if self.config.use_uncertainty_weights: per_view_coord_log_vars = [self._apply_eye_region_boost(lv) for lv in per_view_coord_log_vars] else: per_view_coord_log_vars = None per_view_depth_log_vars = None # --- Per-frame view selection --- frame_valid_matrix = np.ones((num_frames, num_all_views), dtype=bool) for v, inf in enumerate(inferences): if inf.frame_valid is not None: frame_valid_matrix[:, v] = np.asarray(inf.frame_valid) per_frame_views: list[list[int]] = [] frame_reliable = torch.ones(num_frames, dtype=torch.bool, device=self.device) for t in range(num_frames): valid = [v for v in range(num_all_views) if frame_valid_matrix[t, v]] if len(valid) < n_min_views: print(f"[GN][warn] frame {t}: only {len(valid)} valid views (< {n_min_views}), marked unreliable") frame_reliable[t] = False per_frame_views.append([]) elif len(valid) > n_max_views: per_frame_views.append(valid[:n_max_views]) else: per_frame_views.append(valid) n_reliable = int(frame_reliable.sum().item()) view_counts = [len(v) for v in per_frame_views] print(f"[GN] {n_reliable}/{num_frames} reliable frames, " f"views per frame: mean={np.mean(view_counts):.1f} min={np.min(view_counts)} max={np.max(view_counts)} " f"(n_max_views={n_max_views}, n_min_views={n_min_views})") # --- Register first reliable frame --- first_t = int(torch.argmax(frame_reliable.int()).item()) first_views = per_frame_views[first_t] first_cam_view, first_cam_proj = self._stack_camera_matrices( [cameras[v] for v in first_views]) first_coords = torch.stack([per_view_coords[v][first_t] for v in first_views]) first_depths = torch.stack([per_view_depths[v][first_t] for v in first_views]) if per_view_coord_log_vars is not None: first_coord_lv = torch.stack([per_view_coord_log_vars[v][first_t] for v in first_views]) first_depth_lv = torch.stack([per_view_depth_log_vars[v][first_t] for v in first_views]) else: first_coord_lv = None first_depth_lv = None print(f"[GN] register frame {first_t} with {len(first_views)} views (robust)") pose_identity = torch.zeros((300,), dtype=torch.float32, device=self.device) pose_v, pose_j = self._calc_canonical_flame(pose_identity) pose_x0 = self._register_pose( first_cam_view, first_cam_proj, first_coords, first_depths, first_coord_lv, first_depth_lv, pose_v, pose_j, init_x=init_x, ) identity, x0 = self._register( first_cam_view, first_cam_proj, first_coords, first_depths, first_coord_lv, first_depth_lv, init_x=pose_x0, ) # --- Track pass 0 --- xs = self._track_pass_robust( cameras, identity, x0, per_view_coords, per_view_depths, per_view_coord_log_vars, per_view_depth_log_vars, per_frame_views, frame_reliable, desc="[GN] track pass 0 (robust)", ) # --- Refine rounds --- reliable_indices = torch.where(frame_reliable)[0] register_keyframes: list[list[int]] = [] for round_id in range(self.config.identity_refinement_rounds): print(f"[GN] refine round {round_id + 1}/{self.config.identity_refinement_rounds} (robust)") xs_reliable = xs[reliable_indices] n_keys = min(self.config.num_keyframes, len(reliable_indices)) if len(reliable_indices) > n_keys: key_in_reliable = self._select_keyframes(identity, xs_reliable, n_keys) key_frames = reliable_indices[key_in_reliable] else: key_frames = reliable_indices register_keyframes.append(key_frames.cpu().tolist()) cam_view, cam_proj, reg_coords, reg_depths, reg_coord_lv, reg_depth_lv = \ self._gather_register_data( key_frames, per_frame_views, cameras, per_view_coords, per_view_depths, per_view_coord_log_vars, per_view_depth_log_vars, num_vertices, ) xs_key = xs[key_frames] print(f"[GN] {len(key_frames)} keyframes, {cam_view.shape[1]} cameras for register pass") identity, _ = self._register( cam_view, cam_proj, reg_coords, reg_depths, reg_coord_lv, reg_depth_lv, init_identity=identity, init_x=xs_key, ) xs = self._track_pass_robust( cameras, identity, xs, per_view_coords, per_view_depths, per_view_coord_log_vars, per_view_depth_log_vars, per_frame_views, frame_reliable, desc=f"[GN] track pass {round_id + 1} (robust)", ) return FittingOutput( identity=identity, x=xs, per_frame_views=per_frame_views, frame_reliable=frame_reliable.cpu(), register_keyframes=register_keyframes, ) def decode(self, output: FittingOutput) -> tuple[torch.Tensor, torch.Tensor]: identity = output.identity.to(device=self.device, dtype=torch.float32).reshape(300) x = output.x.to(device=self.device, dtype=torch.float32) squeeze_batch = False if x.ndim == 1: x = x.unsqueeze(0) squeeze_batch = True if x.ndim != 2 or x.shape[-1] != self.dim_x: raise ValueError( f"decode expects output.x with shape [{self.dim_x}] or [T, {self.dim_x}]." ) v_canonical, j_canonical = self._calc_canonical_flame(identity) vertices, joints = self._calc_posed_flame(v_canonical, j_canonical, x) if squeeze_batch: return vertices[0], joints[0] return vertices, joints