nbv_reconstruction/runners/strategy_generator.py

import os
import json

import numpy as np
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
from PytorchBoot.utils import Log
import PytorchBoot.stereotype as stereotype
from PytorchBoot.status import status_manager

from utils.data_load import DataLoadUtil
from utils.reconstruction import ReconstructionUtil
from utils.pts import PtsUtil

@stereotype.runner("strategy_generator")
class StrategyGenerator(Runner):
    def __init__(self, config):
        super().__init__(config)
        self.load_experiment("generate")
        self.status_info = {
            "status_manager": status_manager,
            "app_name": "generate",
            "runner_name": "strategy_generator"
        }
        
    def run(self):
        dataset_name_list =  ConfigManager.get("runner", "generate", "dataset_list")
        voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold")
        self.save_pts = ConfigManager.get("runner","generate","save_points")
        for dataset_idx in range(len(dataset_name_list)):
            dataset_name = dataset_name_list[dataset_idx]
            status_manager.set_progress("generate", "strategy_generator", "dataset", dataset_idx, len(dataset_name_list))
            root_dir = ConfigManager.get("datasets", dataset_name, "root_dir")
            model_dir = ConfigManager.get("datasets", dataset_name, "model_dir")
            scene_name_list = os.listdir(root_dir)
            cnt = 0
            total = len(scene_name_list)
            for scene_name in scene_name_list:
                Log.info(f"({dataset_name})Processing [{cnt}/{total}]: {scene_name}")
                status_manager.set_progress("generate", "strategy_generator", "scene", cnt, total)
                self.generate_sequence(root_dir, model_dir, scene_name,voxel_threshold, overlap_threshold)
                
                cnt += 1
            status_manager.set_progress("generate", "strategy_generator", "scene", total, total)
        status_manager.set_progress("generate", "strategy_generator", "dataset", len(dataset_name_list), len(dataset_name_list))
    
    def create_experiment(self, backup_name=None):
        super().create_experiment(backup_name)
        output_dir = os.path.join(str(self.experiment_path), "output")
        os.makedirs(output_dir)
    
    def load_experiment(self, backup_name=None):
        super().load_experiment(backup_name)
    
    def generate_sequence(self, root, model_dir, scene_name, voxel_threshold, overlap_threshold):
        status_manager.set_status("generate", "strategy_generator", "scene", scene_name)
        frame_num = DataLoadUtil.get_scene_seq_length(root, scene_name)
        model_pts = DataLoadUtil.load_original_model_points(model_dir, scene_name)
        down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
        obj_pose = DataLoadUtil.load_target_object_pose(root, scene_name)
        down_sampled_transformed_model_pts = PtsUtil.transform_point_cloud(down_sampled_model_pts, obj_pose)
        pts_list = []
        
        for frame_idx in range(frame_num):
            path = DataLoadUtil.get_path(root, scene_name, frame_idx)
            status_manager.set_progress("generate", "strategy_generator", "loading frame", frame_idx, frame_num)
            point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
            sampled_point_cloud = PtsUtil.voxel_downsample_point_cloud(point_cloud, voxel_threshold)
            if self.save_pts:
                pts_dir = os.path.join(root,scene_name, "pts")
                if not os.path.exists(pts_dir):
                    os.makedirs(pts_dir)
                np.savetxt(os.path.join(pts_dir, f"{frame_idx}.txt"), sampled_point_cloud)
            pts_list.append(sampled_point_cloud)
        status_manager.set_progress("generate", "strategy_generator", "loading frame", frame_num, frame_num)
        
        limited_useful_view, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(down_sampled_transformed_model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold, status_info=self.status_info)
        data_pairs = self.generate_data_pairs(limited_useful_view)
        seq_save_data = {
            "data_pairs": data_pairs,
            "best_sequence": limited_useful_view,
            "max_coverage_rate": limited_useful_view[-1][1]
        }
    
        status_manager.set_status("generate", "strategy_generator", "max_coverage_rate", limited_useful_view[-1][1])
        Log.success(f"Scene <{scene_name}> Finished, Max Coverage Rate: {limited_useful_view[-1][1]}, Best Sequence length: {len(limited_useful_view)}")

        output_label_path = DataLoadUtil.get_label_path(root, scene_name)
        with open(output_label_path, 'w') as f:
            json.dump(seq_save_data, f)
        
        DataLoadUtil.save_downsampled_world_model_points(root, scene_name, down_sampled_transformed_model_pts)

    def generate_data_pairs(self, useful_view):
        data_pairs = []
        for next_view_idx in range(1, len(useful_view)):
            scanned_views = useful_view[:next_view_idx]
            next_view = useful_view[next_view_idx]
            data_pairs.append((scanned_views, next_view))
        return data_pairs