def save_depth()

in depth_fine_tuning.py [0:0]

• 27 lines of code
• 4 McCabe index (conditional complexity)

    def save_depth(self, dir: str = None, frames=None):
        if dir is None:
            dir = self.out_dir
        if frames is None:
            frames = self.frames

        color_fmt = pjoin(self.base_dir, "color_down", "frame_{:06d}.raw")
        depth_dir = pjoin(dir, "depth")
        depth_fmt = pjoin(depth_dir, "frame_{:06d}")

        dataset = VideoFrameDataset(color_fmt, frames)
        data_loader = DataLoader(
            dataset, batch_size=1, shuffle=False, num_workers=4
        )

        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True

        self.model.eval()

        os.makedirs(depth_dir, exist_ok=True)
        for data in data_loader:
            data = to_device(data)
            stacked_images, metadata = data
            frame_id = metadata["frame_id"][0]

            depth = self.model.forward(stacked_images, metadata)

            depth = depth.detach().cpu().numpy().squeeze()
            inv_depth = 1.0 / depth

            image_io.save_raw_float32_image(
                depth_fmt.format(frame_id) + ".raw", inv_depth)

        with SuppressedStdout():
            visualization.visualize_depth_dir(depth_dir, depth_dir, force=True)