(opt: AllConfigs, args: argparse.Namespace)
| 70 | print(f"Decoder checkpoint path not found or not provided: {decoder_path}") |
| 71 | |
| 72 | def main(opt: AllConfigs, args: argparse.Namespace): |
| 73 | device = 'cuda' |
| 74 | print(f"Using device: {device}") |
| 75 | GAP = args.frame_gap |
| 76 | |
| 77 | torch.set_float32_matmul_precision('high') |
| 78 | |
| 79 | model_opt = deepcopy(opt) |
| 80 | model_opt.input_frames = 1 |
| 81 | model_opt.output_frames = GAP + 1 |
| 82 | model_opt.epoch = 0 |
| 83 | model = SplatModel(model_opt).to(device) |
| 84 | load_model_weights(model, opt.resume, device, opt.compile) |
| 85 | model.eval() |
| 86 | |
| 87 | frames_dir = args.input_frames_path |
| 88 | depths_dir = args.input_depths_path |
| 89 | |
| 90 | if not os.path.isdir(frames_dir): |
| 91 | print(f"Error: Input frames directory not found: {frames_dir}") |
| 92 | return |
| 93 | |
| 94 | if depths_dir and not os.path.isdir(depths_dir): |
| 95 | print(f"Warning: Input depths directory not found: {depths_dir}") |
| 96 | depths_dir = None |
| 97 | |
| 98 | output_dir = args.output_dir if args.output_dir else os.path.join(opt.workspace, "inference_output", os.path.basename(frames_dir)) |
| 99 | os.makedirs(output_dir, exist_ok=True) |
| 100 | print(f"Output directory: {output_dir}") |
| 101 | |
| 102 | image_extensions = ('.png', '.jpg', '.jpeg') |
| 103 | all_frame_files = sorted([ |
| 104 | os.path.join(frames_dir, f) for f in os.listdir(frames_dir) |
| 105 | if f.lower().endswith(image_extensions) |
| 106 | ]) |
| 107 | |
| 108 | depth_extensions = ('.png', '.exr', '.npy') |
| 109 | all_depth_files = [] |
| 110 | if depths_dir: |
| 111 | all_depth_files = sorted([ |
| 112 | os.path.join(depths_dir, f) for f in os.listdir(depths_dir) |
| 113 | if f.lower().endswith(depth_extensions) |
| 114 | ]) |
| 115 | |
| 116 | if not all_frame_files: |
| 117 | print(f"Error: No frame files found in {frames_dir}") |
| 118 | return |
| 119 | |
| 120 | if opt.enable_depth and not all_depth_files: |
| 121 | print(f"Warning: No depth files found, using zero depths.") |
| 122 | |
| 123 | total_available_frames = len(all_frame_files) |
| 124 | frame_gap = args.frame_gap |
| 125 | num_output_frames = frame_gap + 1 |
| 126 | |
| 127 | selected_indices = list(range(0, total_available_frames, frame_gap)) |
| 128 | |
| 129 | if len(selected_indices) < 2: |
no test coverage detected