nbv_reconstruction/core/nbv_dataset.py

import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
from PytorchBoot.config import ConfigManager
from PytorchBoot.utils.log_util import Log
import torch
import os
import sys

sys.path.append(r"/data/hofee/project/nbv_rec/nbv_reconstruction")

from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil


@stereotype.dataset("nbv_reconstruction_dataset")
class NBVReconstructionDataset(BaseDataset):
    def __init__(self, config):
        super(NBVReconstructionDataset, self).__init__(config)
        self.config = config
        self.root_dir = config["root_dir"]
        self.split_file_path = config["split_file"]
        self.scene_name_list = self.load_scene_name_list()
        self.datalist = self.get_datalist()

        self.pts_num = config["pts_num"]
        self.type = config["type"]
        self.cache = config.get("cache")
        self.load_from_preprocess = config.get("load_from_preprocess", False)

        if self.type == namespace.Mode.TEST:
            #self.model_dir = config["model_dir"]
            self.filter_degree = config["filter_degree"]
        if self.type == namespace.Mode.TRAIN:
            scale_ratio = 100
            self.datalist = self.datalist*scale_ratio
        if self.cache:
            expr_root = ConfigManager.get("runner", "experiment", "root_dir")
            expr_name = ConfigManager.get("runner", "experiment", "name")
            self.cache_dir = os.path.join(expr_root, expr_name, "cache")
            # self.preprocess_cache()

    def load_scene_name_list(self):
        scene_name_list = []
        with open(self.split_file_path, "r") as f:
            for line in f:
                scene_name = line.strip()
                scene_name_list.append(scene_name)
        return scene_name_list

    def get_datalist(self):
        datalist = []
        for scene_name in self.scene_name_list:
            seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
            scene_max_coverage_rate = 0
            max_coverage_rate_list = []

            for seq_idx in range(seq_num):
                label_path = DataLoadUtil.get_label_path(
                    self.root_dir, scene_name, seq_idx
                )
                label_data = DataLoadUtil.load_label(label_path)
                max_coverage_rate = label_data["max_coverage_rate"]
                if max_coverage_rate > scene_max_coverage_rate:
                    scene_max_coverage_rate = max_coverage_rate
                max_coverage_rate_list.append(max_coverage_rate)
                
            if max_coverage_rate_list:
                mean_coverage_rate = np.mean(max_coverage_rate_list)

            for seq_idx in range(seq_num):
                label_path = DataLoadUtil.get_label_path(
                    self.root_dir, scene_name, seq_idx
                )
                label_data = DataLoadUtil.load_label(label_path)
                if max_coverage_rate_list[seq_idx] > mean_coverage_rate - 0.1:
                    for data_pair in label_data["data_pairs"]:
                        scanned_views = data_pair[0]
                        next_best_view = data_pair[1]
                        datalist.append(
                            {
                                "scanned_views": scanned_views,
                                "next_best_view": next_best_view,
                                "seq_max_coverage_rate": max_coverage_rate,
                                "scene_name": scene_name,
                                "label_idx": seq_idx,
                                "scene_max_coverage_rate": scene_max_coverage_rate,
                            }
                        )
        return datalist

    def preprocess_cache(self):
        Log.info("preprocessing cache...")
        for item_idx in range(len(self.datalist)):
            self.__getitem__(item_idx)
        Log.success("finish preprocessing cache.")

    def load_from_cache(self, scene_name, curr_frame_idx):
        cache_name = f"{scene_name}_{curr_frame_idx}.txt"
        cache_path = os.path.join(self.cache_dir, cache_name)
        if os.path.exists(cache_path):
            data = np.loadtxt(cache_path)
            return data
        else:
            return None

    def save_to_cache(self, scene_name, curr_frame_idx, data):
        cache_name = f"{scene_name}_{curr_frame_idx}.txt"
        cache_path = os.path.join(self.cache_dir, cache_name)
        try:
            np.savetxt(cache_path, data)
        except Exception as e:
            Log.error(f"Save cache failed: {e}")

    def __getitem__(self, index):
        data_item_info = self.datalist[index]
        scanned_views = data_item_info["scanned_views"]
        nbv = data_item_info["next_best_view"]
        max_coverage_rate = data_item_info["seq_max_coverage_rate"]
        scene_name = data_item_info["scene_name"]
        (
            scanned_views_pts,
            scanned_coverages_rate,
            scanned_n_to_world_pose,
        ) = ([], [], [])
        for view in scanned_views:
            frame_idx = view[0]
            coverage_rate = view[1]
            view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
            cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
            
            n_to_world_pose = cam_info["cam_to_world"]  
            target_point_cloud = (
                DataLoadUtil.load_from_preprocessed_pts(view_path)
            )
            downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
                target_point_cloud, self.pts_num
            )
            scanned_views_pts.append(downsampled_target_point_cloud)
            scanned_coverages_rate.append(coverage_rate)
            n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
                np.asarray(n_to_world_pose[:3, :3])
            )
            n_to_world_trans = n_to_world_pose[:3, 3]
            n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
            scanned_n_to_world_pose.append(n_to_world_9d)


        nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
        nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
        cam_info = DataLoadUtil.load_cam_info(nbv_path)
        best_frame_to_world = cam_info["cam_to_world"]

        best_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
            np.asarray(best_frame_to_world[:3, :3])
        )
        best_to_world_trans = best_frame_to_world[:3, 3]
        best_to_world_9d = np.concatenate(
            [best_to_world_6d, best_to_world_trans], axis=0
        )

        combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
        fps_downsampled_combined_scanned_pts, fps_idx = PtsUtil.fps_downsample_point_cloud(
            combined_scanned_views_pts, self.pts_num, require_idx=True
        )
        
        combined_scanned_views_pts_mask = np.zeros(len(combined_scanned_views_pts), dtype=np.uint8)
        
        start_idx = 0
        for i in range(len(scanned_views_pts)):
            end_idx = start_idx + len(scanned_views_pts[i])
            combined_scanned_views_pts_mask[start_idx:end_idx] = i
            start_idx = end_idx
        fps_downsampled_combined_scanned_pts_mask = combined_scanned_views_pts_mask[fps_idx]
        
        data_item = {
            "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
            "scanned_pts_mask": np.asarray(fps_downsampled_combined_scanned_pts_mask,dtype=np.uint8), # Ndarray(N), range(0, S)
            "combined_scanned_pts": np.asarray(fps_downsampled_combined_scanned_pts, dtype=np.float32), # Ndarray(N x 3)
            "scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
            "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
            "best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1)
            "best_to_world_pose_9d": np.asarray(best_to_world_9d, dtype=np.float32), # Ndarray(9)
            "seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
            "scene_name": scene_name, # String
        }

        return data_item

    def __len__(self):
        return len(self.datalist)

    def get_collate_fn(self):
        def collate_fn(batch):
            collate_data = {}
            
            ''' ------ Varialbe Length ------ '''
            
            collate_data["scanned_pts"] = [
                torch.tensor(item["scanned_pts"]) for item in batch
            ]
            collate_data["scanned_n_to_world_pose_9d"] = [
                torch.tensor(item["scanned_n_to_world_pose_9d"]) for item in batch
            ]
            
            ''' ------ Fixed Length ------ '''
            
            collate_data["best_to_world_pose_9d"] = torch.stack(
                [torch.tensor(item["best_to_world_pose_9d"]) for item in batch]
            )
            collate_data["combined_scanned_pts"] = torch.stack(
                [torch.tensor(item["combined_scanned_pts"]) for item in batch]
            )
            collate_data["scanned_pts_mask"] = torch.stack(
                [torch.tensor(item["scanned_pts_mask"]) for item in batch]
            )
                
            for key in batch[0].keys():
                if key not in [
                    "scanned_pts",
                    "scanned_pts_mask",
                    "scanned_n_to_world_pose_9d",
                    "best_to_world_pose_9d",
                    "combined_scanned_pts",
                ]:
                    collate_data[key] = [item[key] for item in batch]
            return collate_data

        return collate_fn


# -------------- Debug ---------------- #
if __name__ == "__main__":
    import torch

    seed = 0
    torch.manual_seed(seed)
    np.random.seed(seed)
    config = {
        "root_dir": "/data/hofee/data/packed_preprocessed_data",
        "source": "nbv_reconstruction_dataset",
        "split_file": "/data/hofee/data/OmniObject3d_train.txt",
        "load_from_preprocess": True,
        "ratio": 0.5,
        "batch_size": 2,
        "filter_degree": 75,
        "num_workers": 0,
        "pts_num": 4096,
        "type": namespace.Mode.TRAIN,
    }
    ds = NBVReconstructionDataset(config)
    print(len(ds))
    # ds.__getitem__(10)
    dl = ds.get_loader(shuffle=True)
    for idx, data in enumerate(dl):
        data = ds.process_batch(data, "cuda:0")
        print(data)
        # ------  Debug Start ------
        import ipdb

        ipdb.set_trace()
        # ------  Debug End ------