Convert pose encoding to extrinsics (c2w) and move to CPU.
(predictions, images)
| 276 | |
| 277 | |
| 278 | def postprocess(predictions, images): |
| 279 | """Convert pose encoding to extrinsics (c2w) and move to CPU.""" |
| 280 | extrinsic, intrinsic = pose_encoding_to_extri_intri(predictions["pose_enc"], images.shape[-2:]) |
| 281 | |
| 282 | # Convert w2c to c2w |
| 283 | extrinsic_4x4 = torch.zeros((*extrinsic.shape[:-2], 4, 4), device=extrinsic.device, dtype=extrinsic.dtype) |
| 284 | extrinsic_4x4[..., :3, :4] = extrinsic |
| 285 | extrinsic_4x4[..., 3, 3] = 1.0 |
| 286 | extrinsic_4x4 = closed_form_inverse_se3_general(extrinsic_4x4) |
| 287 | extrinsic = extrinsic_4x4[..., :3, :4] |
| 288 | |
| 289 | predictions["extrinsic"] = extrinsic |
| 290 | predictions["intrinsic"] = intrinsic |
| 291 | predictions.pop("pose_enc_list", None) |
| 292 | predictions.pop("images", None) |
| 293 | |
| 294 | print("Moving results to CPU...") |
| 295 | for k in list(predictions.keys()): |
| 296 | if isinstance(predictions[k], torch.Tensor): |
| 297 | predictions[k] = _squeeze_single_batch( |
| 298 | k, predictions[k].to("cpu", non_blocking=True) |
| 299 | ) |
| 300 | images_cpu = images.to("cpu", non_blocking=True) |
| 301 | if torch.cuda.is_available(): |
| 302 | torch.cuda.synchronize() |
| 303 | |
| 304 | return predictions, images_cpu |
| 305 | |
| 306 | |
| 307 | def prepare_for_visualization(predictions, images=None): |
no test coverage detected