finish generate_sequence.py

This commit is contained in:
hofee 2024-08-22 20:27:21 +08:00
parent ff3e89b17c
commit 7cd1954a25
7 changed files with 121 additions and 57 deletions

View File

@ -1,23 +1,24 @@
runners:
runner:
general:
seed: 0
device: cpu
cuda_visible_devices: "0,1,2,3,4,5,6,7"
generate:
voxel_threshold: 0.005
overlap_threshold: 0.3
experiment:
name: debug
root_dir: "experiments"
generate:
- name: OmniObject3d_train
component: OmniObject3d
data_type: train
dataset_list:
- OmniObject3d
datasets:
general:
components:
OmniObject3d:
root_dir: "C:\\Document\\Local Project\\nbv_rec\\output"
root_dir: "C:\\Document\\Local Project\\nbv_rec\\sample_dataset"
output_dir: "C:\\Document\\Local Project\\nbv_rec\\sample_output"

17
core/dataset.py Normal file
View File

@ -0,0 +1,17 @@
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.stereotype as stereotype
@stereotype.dataset("nbv_reconstruction_dataset", comment="unfinished")
class NBVReconstructionDataset(BaseDataset):
def __init__(self, config):
super(NBVReconstructionDataset, self).__init__(config)
self.config = config
def get_datalist(self):
pass
def load_view(path):
pass
def load_data_item(self, idx):
pass

20
core/evaluation.py Normal file
View File

@ -0,0 +1,20 @@
import torch
import PytorchBoot.stereotype as stereotype
@stereotype.evaluation_method("delta_pose_diff", comment="unfinished")
class DeltaPoseDiff:
def __init__(self, config):
pass
def evaluate(self, output_list, data_list):
return
@stereotype.evaluation_method("coverage_rate_increase",comment="unfinished")
class ConverageRateIncrease:
def __init__(self, config):
pass
def evaluate(self, output_list, data_list):
return

View File

@ -5,7 +5,7 @@ import PytorchBoot.stereotype as stereotype
from PytorchBoot.factory.component_factory import ComponentFactory
from PytorchBoot.utils import Log
@stereotype.pipeline("nbv_reconstruction_pipeline")
@stereotype.pipeline("nbv_reconstruction_pipeline", comment="should be tested")
class NBVReconstructionPipeline(nn.Module):
def __init__(self, config):
super(NBVReconstructionPipeline, self).__init__()

View File

@ -1,5 +1,4 @@
from modules.func_lib.samplers import (
cond_pc_sampler,
cond_ode_sampler
)
from modules.func_lib.sde import (

View File

@ -1,16 +1,30 @@
import os
import json
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
import PytorchBoot.stereotype as stereotype
@stereotype.runner("strategy_generator")
from utils.data_load import DataLoadUtil
from utils.reconstruction import ReconstructionUtil
@stereotype.runner("strategy_generator", comment="unfinished")
class StrategyGenerator(Runner):
def __init__(self, config):
super().__init__(config)
self.load_experiment("generate")
def run(self):
self.demo(seq=16,num=100)
dataset_name_list = ConfigManager.get("runner", "dataset_list")
voxel_threshold, overlap_threshold = ConfigManager.get("runner","generate","voxel_threshold"), ConfigManager.get("runner","generate","overlap_threshold")
for dataset_name in dataset_name_list:
root_dir = ConfigManager.get("datasets", dataset_name, "root_dir")
output_dir = ConfigManager.get("datasets", dataset_name, "output_dir")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
scene_idx_list = DataLoadUtil.get_scene_idx_list(root_dir)
for scene_idx in scene_idx_list:
self.generate_sequence(root_dir, output_dir, scene_idx,voxel_threshold, overlap_threshold)
def create_experiment(self, backup_name=None):
super().create_experiment(backup_name)
@ -20,54 +34,34 @@ class StrategyGenerator(Runner):
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
def demo(self, seq, num=100):
import os
from utils.data_load import DataLoadUtil
from utils.reconstruction import ReconstructionUtil
import numpy as np
component = self.config["generate"][0]["component"] #r"C:\Document\Local Project\nbv_rec\output"
data_dir = ConfigManager.get("datasets", "components", component, "root_dir")
model_path = os.path.join(data_dir, f"sequence.{seq}\\world_points.txt")
model_pts = np.loadtxt(model_path)
output_dir = os.path.join(str(self.experiment_path), "output")
def generate_sequence(self,root, output_dir, seq, voxel_threshold, overlap_threshold):
frame_idx_list = DataLoadUtil.get_frame_idx_list(root, seq)
model_pts = DataLoadUtil.load_model_points(root, seq)
pts_list = []
for idx in range(0,num):
path = DataLoadUtil.get_path(data_dir, seq, idx)
for frame_idx in frame_idx_list:
path = DataLoadUtil.get_path(root, seq, frame_idx)
point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, 0.005)
sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, voxel_threshold)
pts_list.append(sampled_point_cloud)
sampled_model_pts = ReconstructionUtil.downsample_point_cloud(model_pts, 0.005)
np.savetxt(os.path.join(output_dir,"sampled_model_points.txt"), sampled_model_pts)
thre = 0.005
useful_view, useless_view = ReconstructionUtil.compute_next_best_view_sequence(model_pts, pts_list, threshold=thre)
print("useful:", useful_view)
print("useless:", useless_view)
selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, useful_view)
downsampled_selected_full_views = ReconstructionUtil.downsample_point_cloud(selected_full_views, thre)
np.savetxt(os.path.join(output_dir,"selected_full_views.txt"), downsampled_selected_full_views)
limited_useful_view, limited_useless_view = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=thre, overlap_threshold=0.3)
print("limited_useful:", limited_useful_view)
print("limited_useless:", limited_useless_view)
limited_selected_full_views = ReconstructionUtil.combine_point_with_view_sequence(pts_list, limited_useful_view)
downsampled_limited_selected_full_views = ReconstructionUtil.downsample_point_cloud(limited_selected_full_views, thre)
np.savetxt(os.path.join(output_dir,"selected_full_views_limited.txt"), downsampled_limited_selected_full_views)
import json
for idx, score in limited_useful_view:
path = DataLoadUtil.get_path(data_dir, seq, idx)
point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
print("saving useful view: ", idx, " | score: ", score)
np.savetxt(os.path.join(output_dir,f"useful_view_{idx}.txt"), point_cloud)
with open(os.path.join(output_dir,f"useful_view.json"), 'w') as f:
json.dump(limited_useful_view, f)
print("seq length: ", len(useful_view), "limited seq length: ", len(limited_useful_view))
limited_useful_view, _ = ReconstructionUtil.compute_next_best_view_sequence_with_overlap(model_pts, pts_list, threshold=voxel_threshold, overlap_threshold=overlap_threshold)
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]
}
output_label_path = DataLoadUtil.get_label_path(output_dir, seq)
with open(output_label_path, 'w') as f:
json.dump(seq_save_data, f)
def generate_data_pairs(self, useful_view):
data_pairs = []
for next_view_idx in range(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

View File

@ -4,6 +4,7 @@ import Imath
import numpy as np
import json
import cv2
import re
class DataLoadUtil:
@ -12,6 +13,38 @@ class DataLoadUtil:
path = os.path.join(root, f"sequence.{scene_idx}", f"step{frame_idx}")
return path
@staticmethod
def get_label_path(root, scene_idx):
path = os.path.join(root, f"sequence.{scene_idx}_label.json")
return path
@staticmethod
def get_scene_idx_list(root):
scene_dir = os.listdir(root)
scene_idx_list = []
for scene in scene_dir:
if "sequence" in scene:
scene_idx = int(re.search(r'\d+', scene).group())
scene_idx_list.append(scene_idx)
return scene_idx_list
@staticmethod
def get_frame_idx_list(root, scene_idx):
scene_path = os.path.join(root, f"sequence.{scene_idx}")
view_dir = os.listdir(scene_path)
seen_frame_idx = set()
for view in view_dir:
if "step" in view:
frame_idx = int(re.search(r'\d+', view).group())
seen_frame_idx.add(frame_idx)
return list(seen_frame_idx)
@staticmethod
def load_model_points(root,scene_idx):
model_path = os.path.join(root, f"sequence.{scene_idx}", "world_points.txt")
model_pts = np.loadtxt(model_path)
return model_pts
@staticmethod
def read_exr_depth(depth_path):
file = OpenEXR.InputFile(depth_path)