原路径：/SSD2/yangzijing/egoview-v9-foundationhand-s9l10-20260901/v9_fusion/s9_sequence_fit.py
SHA256：bc1661a1c78bfb1a09f49dd3378c96ef58452aa8b62bfc946a922f4c585ba144
1 """Joint RGB-constrained MANO estimation initialized ONLY by S9L10. 2 3 No learned 3D model is loaded here. Use the same licensed MANO decoder as the 4 foundation model, with calibrated K, fixed network wrist Z by default, and shared 5 subject shape. RGB-only inputs may explicitly enable metric palm initialization 6 and joint depth estimation. Every temporal parameter belongs to the full sequence, not independent 7 chunks. Temporal regularization is inside estimation, never an output filter. 8 """ 9 from __future__ import annotations 10 11 import argparse 12 import json 13 import time 14 from dataclasses import asdict, dataclass 15 from pathlib import Path 16 17 import numpy as np 18 19 from .io import sha256_file, write_json, write_npz 20 21 22 @dataclass(frozen=True) 23 class SequenceFitConfig: 24 steps: int = 360 25 pose_learning_rate: float = .012 26 root_learning_rate: float = .002 27 score_threshold: float = .35 28 minimum_joints: int = 4 29 sigma_box_fraction: float = .025 30 sigma_image_fraction: float = .001 31 pose_prior_weight: float = .08 32 temporal_weight: float = .4 33 maximum_pose_delta_rad: float = .7 34 maximum_global_delta_rad: float = float(np.pi) 35 holdout_period: int = 0 36 rigid_palm_initialization: bool = True 37 articulation_temporal_weight: float = .4 38 root_temporal_weight: float = .015 39 shared_subject_shape: bool = True 40 optimize_wrist_depth: bool = False 41 depth_temporal_weight: float = 2.0 42 depth_prior_weight: float = .02 43 kinematic_surface_depth_seed: bool = False 44 respect_observation_track_boundaries: bool = False 45 partial_mask_initialization: bool = False 46 palm_confidence_guard: bool = False 47 48 def __post_init__(self): 49 if (self.steps < 1 or self.temporal_weight < 0 or self.articulation_temporal_weight < 0 50 or self.root_temporal_weight < 0 51 or not 0 < self.maximum_global_delta_rad <= np.pi + 1e-6 52 or self.holdout_period not in (0, 5)): 53 raise ValueError("invalid sequence fitting configuration") 54 55 56 def calibrated_intrinsics(prepared): 57 contract = prepared.get("camera_contract", {}) 58 if prepared.get("hand_rgb_transform") != "rectified_stereo_left": 59 raise ValueError("calibrated S9 replay currently requires rectified stereo left") 60 if not contract.get("real_calibration_used"): 61 raise ValueError("calibrated S9 replay requires declared calibration") 62 k = np.asarray(contract.get("display_intrinsics"), np.float32) 63 if k.shape != (3, 3) or not np.isfinite(k).all() or min(k[0, 0], k[1, 1]) <= 0: 64 raise ValueError("invalid calibrated raster K") 65 return k 66 67 68 def temporal_pairs(frame_ids, present, track_ids=None): 69 """Never regularize distinct observed tracks as one moving physical hand.""" 70 present=np.asarray(present,bool) 71 pairs=present[1:]&present[:-1]&(np.diff(frame_ids)==1)[:,None] 72 if track_ids is not None: 73 tracks=np.asarray(track_ids) 74 if tracks.shape!=present.shape or not np.issubdtype(tracks.dtype,np.integer): 75 raise ValueError('track IDs must be integer [T,S] matching validity') 76 pairs &= (tracks[1:]>=0)&(tracks[:-1]>=0)&(tracks[1:]==tracks[:-1]) 77 return pairs 78 79 80 def evidence_arrays(raw, evidence, k, config): 81 """Use only S9 presence and RGB observations; reject identity mismatch.""" 82 if not np.array_equal(raw["frame_id"], evidence["frame_ids"]): 83 raise ValueError("sequence fit frame identity differs") 84 k = np.asarray(k) 85 if k.shape not in ((3, 3), (len(raw["frame_id"]), 3, 3)) or not np.isfinite(k).all(): 86 raise ValueError("intrinsics must be finite [3,3] or [T,3,3]") 87 keypoints = np.asarray(evidence["v5r_mmpose_keypoints_2d"], np.float32) 88 scores = np.asarray(evidence["v5r_mmpose_keypoint_scores"], np.float32) 89 valid = np.asarray(evidence["valid"], bool) 90 usable = (np.isfinite(keypoints).all(-1) & np.isfinite(scores) 91 & (scores >= config.score_threshold) 92 & np.asarray(evidence["v5r_mmpose_fixed_usable_joints"], bool) 93 & np.asarray(evidence["v5r_mmpose_valid"], bool)[..., None] 94 & valid[..., None]) 95 supported = valid & (usable.sum(-1) >= config.minimum_joints) 96 usable &= supported[..., None] 97 boxes = np.asarray(evidence["boxes_xyxy"], np.float32) 98 diagonal = np.linalg.norm(boxes[..., 2:] - boxes[..., :2], axis=-1) 99 sigma = np.maximum(np.nan_to_num(diagonal), 1) * config.sigma_box_fraction 100 sigma = np.maximum(sigma, np.hypot(*evidence["image_size_hw"]) * config.sigma_image_fraction) 101 roots = np.array(raw["predicted_wrist_root_translation"], copy=True) 102 j = np.asarray(raw["predicted_joints_camera"], np.float32) 103 for f, s in np.argwhere(supported): 104 frame_k = k[f] if k.ndim == 3 else k 105 ok = usable[f, s] 106 desired = (keypoints[f, s, ok] - frame_k[:2, 2]) * j[f, s, ok, 2:] / np.diag(frame_k)[:2] 107 roots[f, s, :2] += np.median(desired - j[f, s, ok, :2], axis=0) 108 return keypoints, scores, usable, supported, sigma, roots 109 110 111 def fit_sequence(raw, evidence, k, mano, config=SequenceFitConfig(), *, device="cuda:0", progress=None, surface_constraints=None, surface_config=None): 112 import torch 113 import torch.nn.functional as F 114 from foundation_hand.geometry.rotations import axis_angle_to_matrix 115 from .s9l10_geometry import stable_matrix_to_axis_angle 116 117 fit_evidence = dict(evidence) 118 if config.holdout_period: 119 # Withhold complete RGB observation frames, including XY initialization. 120 fit_evidence["v5r_mmpose_valid"] = np.array(evidence["v5r_mmpose_valid"], copy=True) 121 fit_evidence["v5r_mmpose_valid"][np.arange(len(raw["frame_id"])) % config.holdout_period == 2] = False 122 keypoints, scores, usable, supported, sigma, roots = evidence_arrays(raw, fit_evidence, k, config) 123 initial_rotations = raw["predicted_mano_rotation_matrices"] 124 palm_report = {"enabled": config.rigid_palm_initialization, "changed_slots": 0} 125 if config.rigid_palm_initialization: 126 from .s9_pose_initialization import initialize_palm_orientation 127 initial_rotations, roots, palm = initialize_palm_orientation( 128 raw["predicted_joints_camera"], initial_rotations, roots, 129 keypoints, usable, scores, k, fixed_wrist_z=not config.optimize_wrist_depth, 130 confidence_guard=config.palm_confidence_guard) 131 changed = palm["changed"] 132 palm_report.update(changed_slots=int(changed.sum()), changed_frame_slots=np.argwhere(changed).tolist(), 133 mean_error_before_px=float(palm["before_px"][changed].mean()) if changed.any() else None, 134 mean_error_after_px=float(palm["after_px"][changed].mean()) if changed.any() else None, 135 method="rigid_palm_SQPNP_candidates_scored_at_fixed_S9_wrist_z", 136 local_finger_rotations_unchanged=True, native_wrist_z_unchanged=not config.optimize_wrist_depth) 137 if config.optimize_wrist_depth:palm_report["method"]="metric_S9_palm_SQPNP_seed_with_calibrated_K" 138 palm_report['confidence_guard'] = config.palm_confidence_guard 139 surface = None 140 surface_seed_changes=0 141 if surface_constraints is not None: 142 if not config.optimize_wrist_depth: 143 raise ValueError('surface-constrained fit requires jointly estimated depth') 144 from .hand_surface import SurfaceLoss 145 surface=SurfaceLoss(surface_constraints,device,**({} if surface_config is None else {'config':surface_config})) 146 for f,s in np.argwhere(supported): 147 ok=surface_constraints['joint_depth_weight'][f,s]>0 148 if ok.sum()>=3: 149 if config.kinematic_surface_depth_seed: 150 from .s9_pose_initialization import surface_wrist_seed 151 local=(raw['predicted_joints_camera'][f,s]-raw['predicted_joints_camera'][f,s,:1])@raw['predicted_mano_rotation_matrices'][f,s,0] 152 relative=local@initial_rotations[f,s,0].T 153 frame_k=k[f] if np.asarray(k).ndim==3 else k 154 candidate,changed=surface_wrist_seed(relative,surface_constraints['joint_depth_m'][f,s], 155 surface_constraints['joint_depth_weight'][f,s],roots[f,s],keypoints[f,s],usable[f,s],frame_k, 156 joint_below_surface_m=surface.config.joint_below_surface_m) 157 if changed:roots[f,s]=candidate;surface_seed_changes+=1 158 else: 159 z=float(np.median(surface_constraints['joint_depth_m'][f,s,ok]))+.01 160 roots[f,s,:2]*=z/max(float(roots[f,s,2]),.02);roots[f,s,2]=z 161 partial_seed_count=0 162 if config.partial_mask_initialization: 163 initial_rotations=np.array(initial_rotations,copy=True) 164 if surface_constraints is None or 'mask_depth_seed_m' not in surface_constraints: 165 raise ValueError('partial-mask initialization requires bound surface observations') 166 # Initialize only observations without enough visible keypoints. No 167 # output interpolation or validity filling; every slot has same-frame SAM. 168 tracks=evidence['observation_track_ids'] 169 for f,s in np.argwhere(np.asarray(evidence['valid'])&~supported&(surface_constraints['mask_weight']>0)): 170 candidates=np.flatnonzero(supported[:,s]&(tracks[:,s]==tracks[f,s])) 171 if len(candidates): 172 near=int(candidates[np.abs(candidates-f).argmin()]) 173 initial_rotations[f,s]=initial_rotations[near,s] 174 relative=(raw['predicted_joints_camera'][f,s]-raw['predicted_joints_camera'][f,s,:1])@raw['predicted_mano_rotation_matrices'][f,s,0] 175 relative=relative@initial_rotations[f,s,0].T 176 z=float(surface_constraints['mask_depth_seed_m'][f,s]) 177 if not .1<=z<=4:z=float(roots[f,s,2]) 178 z=max(.1,z) 179 box=surface_constraints['mask_boxes'][f,s] 180 center=box[:2]+surface_constraints['mask_points'][f,s].mean(0)*(box[2:]-box[:2]) 181 palm=relative[[0,5,9,13,17]].mean(0) 182 frame_k=k[f] if np.asarray(k).ndim==3 else k 183 roots[f,s,:2]=(center-frame_k[:2,2])*(z+palm[2])/np.diag(frame_k)[:2]-palm[:2] 184 roots[f,s,2]=z;partial_seed_count+=1 185 palm_report['partial_mask_initialized_slots']=partial_seed_count 186 tensor = lambda x, dtype=torch.float32: torch.as_tensor(x, dtype=dtype, device=device) 187 target = tensor(np.nan_to_num(keypoints)) 188 weights = tensor(np.where(usable, np.maximum(scores, 0) ** 2, 0)) 189 scale = tensor(sigma)[..., None, None] 190 present = tensor(evidence["valid"], torch.bool) 191 track_ids=None 192 if config.respect_observation_track_boundaries: 193 if 'observation_track_ids' not in evidence:raise ValueError('track-aware fit requires observation_track_ids') 194 track_ids=evidence['observation_track_ids'] 195 pair_array=temporal_pairs(raw['frame_id'],evidence['valid'],track_ids) 196 pair = tensor(pair_array, torch.bool) 197 triple = pair[1:] & pair[:-1] 198 source_rot = tensor(initial_rotations) 199 source_aa = stable_matrix_to_axis_angle(source_rot).detach() 200 source_shape = tensor(raw["predicted_mano_shape"]) 201 # One tracked physical hand has one shape throughout the recording. 202 shape_weight = tensor(supported)[..., None] 203 shared_shape = (source_shape * shape_weight).sum(0) / shape_weight.sum(0).clamp_min(1) 204 missing = ~tensor(supported.any(0), torch.bool) 205 shared_shape[missing] = source_shape.mean(0)[missing] 206 shape = (shared_shape[None].expand(len(source_aa), -1, -1).contiguous() 207 if config.shared_subject_shape else source_shape) 208 fixed_z = tensor(roots if config.optimize_wrist_depth else raw["predicted_wrist_root_translation"])[..., 2:] 209 log_z_initial = fixed_z.clamp_min(.1).log().detach() 210 log_z = torch.nn.Parameter(log_z_initial.clone()) if config.optimize_wrist_depth else None 211 xy_initial = tensor(roots[..., :2]) / fixed_z.clamp_min(.001) 212 xy = torch.nn.Parameter(xy_initial.clone()) 213 delta = torch.nn.Parameter(torch.zeros_like(source_aa)) 214 optimizer = torch.optim.Adam([ 215 {"params": [delta], "lr": config.pose_learning_rate}, 216 {"params": [xy], "lr": config.root_learning_rate}, 217 ]) 218 shape_parameter=None 219 if surface is not None and surface.config.optimize_subject_shape: 220 shape_parameter=torch.nn.Parameter(shared_shape.detach().clone()) 221 optimizer.add_param_group({'params':[shape_parameter],'lr':.004}) 222 if log_z is not None:optimizer.add_param_group({"params":[log_z],"lr":config.root_learning_rate}) 223 matrix = tensor(k) 224 # Entire-sequence temporal terms have no independent window boundaries. 225 def decode(pose, beta, wrist): 226 vertices, joints = [], [] 227 for side, layer in enumerate(mano.layers): 228 v, j = layer(pose[:, side].reshape(-1, 48), beta[:, side]) 229 center = j[:, :1] 230 vertices.append((v-center)/1000 + wrist[:, side, None]) 231 joints.append((j-center)/1000 + wrist[:, side, None]) 232 return torch.stack(vertices, 1), torch.stack(joints, 1) 233 234 def project(j): 235 if matrix.ndim == 3: 236 focal = matrix.diagonal(dim1=-2, dim2=-1)[:, None, None, :2] 237 center = matrix[:, None, None, :2, 2] 238 return j[..., :2] / j[..., 2:].clamp_min(.02) * focal + center 239 return j[..., :2] / j[..., 2:].clamp_min(.02) * matrix.diag()[:2] + matrix[:2, 2] 240 241 history = [] 242 with torch.no_grad(): 243 _, initial_j = decode(source_aa, shape, torch.cat((xy_initial*fixed_z, fixed_z), -1)) 244 # Relative temporal scale comes from subject hand extent, not video tuning. 245 extent = (initial_j - initial_j[:, :, :1]).norm(dim=-1).amax(-1) 246 extent = (extent * tensor(supported)).sum(0) / tensor(supported).sum(0).clamp_min(1) 247 extent = extent.clamp_min(.05) 248 for step in range(config.steps): 249 optimizer.zero_grad(set_to_none=True) 250 if log_z is not None:fixed_z=log_z.exp() 251 if shape_parameter is not None:shape=shape_parameter[None].expand(len(source_aa),-1,-1).contiguous() 252 aa = source_aa + delta 253 wrist = torch.cat((xy*fixed_z, fixed_z), -1) 254 vertices, j = decode(aa, shape, wrist) 255 pix = project(j) 256 residual = (pix-target) / scale 257 reproj = (F.smooth_l1_loss(residual, torch.zeros_like(residual), reduction="none").sum(-1) * weights).sum() / weights.sum().clamp_min(1) 258 rot = axis_angle_to_matrix(aa) 259 prior = ((rot-source_rot).square().sum((-1,-2)) * present[...,None]).sum() / (present.sum()*16).clamp_min(1) 260 local = (j-j[:, :, :1]) / extent[None,:,None,None] 261 acc = local[2:]-2*local[1:-1]+local[:-2] 262 temporal = (acc.square().sum(-1).mean(-1)*triple).sum()/triple.sum().clamp_min(1) 263 # A real palm turn must not be explained by counter-rotating fingers. 264 # Constrain articulation in the palm frame, separately from camera motion. 265 canonical = local @ rot[:, :, 0] 266 finger_acc = canonical[2:]-2*canonical[1:-1]+canonical[:-2] 267 finger_temporal = (finger_acc.square().sum(-1).mean(-1)*triple).sum()/triple.sum().clamp_min(1) 268 # Projected root regularization: weak enough to follow real fast motion. 269 root_px = project(wrist[...,None,:]).squeeze(-2) 270 root_acc = (root_px[2:]-2*root_px[1:-1]+root_px[:-2]) / sigma.mean() 271 root_temporal = (root_acc.square().sum(-1)*triple).sum()/triple.sum().clamp_min(1) 272 loss = (reproj + config.pose_prior_weight*prior + config.temporal_weight*temporal*100 273 + config.articulation_temporal_weight*finger_temporal*100 + config.root_temporal_weight*root_temporal) 274 surface_terms={} 275 if surface is not None: 276 surface_loss,surface_terms=surface(j,vertices,project) 277 loss=loss+surface_loss 278 if shape_parameter is not None:loss=loss+surface.config.shape_prior_weight*(shape_parameter-shared_shape).square().mean() 279 if log_z is not None: 280 depth_acc=log_z[2:]-2*log_z[1:-1]+log_z[:-2] 281 depth_temporal=(depth_acc.square().squeeze(-1)*triple).sum()/triple.sum().clamp_min(1) 282 depth_prior=((log_z-log_z_initial).square().squeeze(-1)*present).sum()/present.sum().clamp_min(1) 283 loss=loss+config.depth_temporal_weight*depth_temporal+config.depth_prior_weight*depth_prior 284 if not torch.isfinite(loss): 285 raise RuntimeError("nonfinite sequence estimation objective") 286 loss.backward() 287 torch.nn.utils.clip_grad_norm_([delta, xy], 10) 288 optimizer.step() 289 with torch.no_grad(): 290 limits = torch.full_like(delta[..., :1], config.maximum_pose_delta_rad) 291 limits[:,:,0] = config.maximum_global_delta_rad 292 delta.mul_(torch.minimum(torch.ones_like(limits), limits/delta.norm(dim=-1,keepdim=True).clamp_min(1e-8))) 293 delta[~present] = 0 294 if shape_parameter is not None:shape_parameter.copy_(torch.maximum(torch.minimum(shape_parameter,shared_shape+2),shared_shape-2)) 295 xy[~present] = xy_initial[~present] 296 if log_z is not None: 297 log_z.clamp_(float(np.log(.1)),float(np.log(4))) 298 log_z[~present]=log_z_initial[~present] 299 if step % 30 == 0 or step+1 == config.steps: 300 record = {"step": step+1, "loss": float(loss.detach()), "reprojection": float(reproj.detach()), 301 "temporal": float(temporal.detach()), "articulation_temporal": float(finger_temporal.detach()), 302 "pose_prior": float(prior.detach())} 303 record.update({key:float(value.detach()) for key,value in surface_terms.items()}) 304 history.append(record) 305 if progress is not None: 306 progress(record) 307 with torch.no_grad(): 308 if log_z is not None:fixed_z=log_z.exp() 309 if shape_parameter is not None:shape=shape_parameter[None].expand(len(source_aa),-1,-1).contiguous() 310 aa = source_aa + delta 311 wrist = torch.cat((xy*fixed_z, fixed_z), -1) 312 v, j = decode(aa, shape, wrist) 313 rotation = axis_angle_to_matrix(aa) 314 output = dict(raw, predicted_joints_camera=j.cpu().numpy(), predicted_vertices_camera=v.cpu().numpy(), 315 predicted_mano_rotation_matrices=rotation.cpu().numpy(), predicted_mano_shape=shape.cpu().numpy(), 316 predicted_wrist_root_translation=wrist.cpu().numpy()) 317 if not config.optimize_wrist_depth: 318 np.testing.assert_array_equal(output["predicted_wrist_root_translation"][...,2], raw["predicted_wrist_root_translation"][...,2]) 319 return output, {"config": asdict(config), "history": history, "supported_slots": int(supported.sum()), 320 "temporal_track_boundaries": {"enabled":config.respect_observation_track_boundaries, 321 "removed_pairs":int((temporal_pairs(raw['frame_id'],evidence['valid'])&~pair_array).sum()), 322 "validity_changed":False}, 323 "surface_wrist_initialization": {"method":"posed_joint_depth_offsets" if config.kinematic_surface_depth_seed else "legacy_median_palm_surface_as_wrist", 324 "changed_slots":surface_seed_changes}, 325 "palm_initialization": palm_report, 326 "shared_shape": (shared_shape if shape_parameter is None else shape_parameter.detach()).cpu().numpy().tolist(), "source": "FoundationHand-S9L10 (1e-6)", 327 "hawor_or_v7_3d_used": False, "native_sensor_K": k.tolist(), "native_wrist_z_unchanged": not config.optimize_wrist_depth, 328 "wrist_depth_policy": ("quality_gated_stereo_surface_and_RGB_joint_estimation" if surface is not None and np.any(surface_constraints['joint_depth_weight']>0) 329 else "calibrated_RGB_and_metric_S9_shape_joint_estimation" if config.optimize_wrist_depth else "fixed_native_S9_wrist_z"), 330 "surface_constraints": None if surface is None else dict(config=asdict(surface.config),mask_slots=int((surface_constraints['mask_weight']>0).sum()), 331 stereo_joint_observations=int((surface_constraints['joint_depth_weight']>0).sum()),source='same-frame SAM3 mask and decoded shared stereo confidence; no hand GT'), 332 "post_output_temporal_filter": False, "method": "joint_full_sequence_RGB_constrained_MANO_estimation", 333 "rgb_2d_source": "existing V9 MMPose observations", "learned_3d_model_count": 1, 334 "test_time_geometry_optimization": True} 335 336 337 def main(): 338 from . import s9l10_worker as worker 339 from .s9l10 import S9L10_CONFIG 340 p = argparse.ArgumentParser(description=__doc__) 341 p.add_argument("--prediction", type=Path, required=True) 342 p.add_argument("--observations", type=Path, required=True) 343 p.add_argument("--prepared", type=Path, required=True) 344 p.add_argument("--output", type=Path, required=True) 345 defaults = SequenceFitConfig() 346 p.add_argument("--steps", type=int, default=defaults.steps) 347 p.add_argument("--temporal-weight", type=float, default=defaults.temporal_weight) 348 p.add_argument("--holdout-period", type=int, default=0) 349 p.add_argument("--no-palm-initialization", action="store_true", help="Controlled comparison with the previous estimator") 350 p.add_argument("--global-rotation-limit", type=float, default=defaults.maximum_global_delta_rad) 351 p.add_argument("--articulation-temporal-weight", type=float, default=defaults.articulation_temporal_weight) 352 p.add_argument("--root-temporal-weight", type=float, default=defaults.root_temporal_weight) 353 p.add_argument("--per-frame-shape", action="store_true", help="Diagnostic ablation; default shares subject shape") 354 args = p.parse_args() 355 args.output.mkdir(parents=True, exist_ok=False) 356 started = time.monotonic() 357 worker._verify_locked_sources() 358 worker._activate_official_runtime() 359 import torch 360 import yaml 361 from foundation_hand.models.mano import ManoCameraLayer 362 torch.set_num_threads(8) 363 cfg = yaml.safe_load(S9L10_CONFIG.read_text()) 364 mano = ManoCameraLayer(cfg["assets"]["mano_root"], fix_left_shapedirs=True).cuda().eval() 365 mano.requires_grad_(False) 366 with np.load(args.prediction, allow_pickle=False) as z: 367 raw = {k:z[k] for k in z.files} 368 names = ("frame_ids", "valid", "boxes_xyxy", "image_size_hw", "v5r_mmpose_keypoints_2d", 369 "v5r_mmpose_keypoint_scores", "v5r_mmpose_fixed_usable_joints", "v5r_mmpose_valid") 370 with np.load(args.observations, allow_pickle=False) as z: 371 evidence = {k:z[k] for k in names} 372 k = calibrated_intrinsics(json.loads(args.prepared.read_text())) 373 output, report = fit_sequence(raw, evidence, k, mano, 374 SequenceFitConfig(steps=args.steps, temporal_weight=args.temporal_weight, holdout_period=args.holdout_period, 375 rigid_palm_initialization=not args.no_palm_initialization, 376 maximum_global_delta_rad=args.global_rotation_limit, 377 articulation_temporal_weight=args.articulation_temporal_weight, 378 root_temporal_weight=args.root_temporal_weight, shared_subject_shape=not args.per_frame_shape), 379 progress=lambda v: print(json.dumps(v), flush=True)) 380 target = args.output / "predictions.npz" 381 write_npz(target, output) 382 report.update(prediction_input=str(args.prediction.resolve()), prediction_input_sha256=sha256_file(args.prediction), 383 observations_input=str(args.observations.resolve()), observations_input_sha256=sha256_file(args.observations), 384 prepared_manifest_sha256=sha256_file(args.prepared), prediction_output_sha256=sha256_file(target), 385 elapsed_seconds=time.monotonic()-started) 386 write_json(args.output / "report.json", report) 387 print(json.dumps({"complete": str(target), "seconds": report["elapsed_seconds"]}), flush=True) 388 389 390 if __name__ == "__main__": 391 main()