finish generate_sequence.py
This commit is contained in:
hofee
parent
ff3e89b17c
commit
7cd1954a25
@ -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
17
core/dataset.py
Normal 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
20
core/evaluation.py
Normal 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
|
||||
|
||||
@ -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__()
|
||||
|
||||
@ -1,5 +1,4 @@
|
||||
from modules.func_lib.samplers import (
|
||||
cond_pc_sampler,
|
||||
cond_ode_sampler
|
||||
)
|
||||
from modules.func_lib.sde import (
|
||||
|
||||
@ -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)
|
||||
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
|
||||
|
||||
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))
|
||||
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user