upd ab_global_only
This commit is contained in:
@@ -8,7 +8,7 @@ import torch
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
|
||||
sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction")
|
||||
|
||||
from utils.data_load import DataLoadUtil
|
||||
from utils.pose import PoseUtil
|
||||
@@ -47,6 +47,8 @@ class SeqReconstructionDataset(BaseDataset):
|
||||
with open(self.split_file_path, "r") as f:
|
||||
for line in f:
|
||||
scene_name = line.strip()
|
||||
if not os.path.exists(os.path.join(self.root_dir, scene_name)):
|
||||
continue
|
||||
scene_name_list.append(scene_name)
|
||||
return scene_name_list
|
||||
|
||||
@@ -58,29 +60,19 @@ class SeqReconstructionDataset(BaseDataset):
|
||||
total = len(self.scene_name_list)
|
||||
for idx, scene_name in enumerate(self.scene_name_list):
|
||||
print(f"processing {scene_name} ({idx}/{total})")
|
||||
seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
|
||||
scene_max_coverage_rate = 0
|
||||
max_coverage_rate_list = []
|
||||
scene_max_cr_idx = 0
|
||||
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
|
||||
scene_max_cr_idx = seq_idx
|
||||
max_coverage_rate_list.append(max_coverage_rate)
|
||||
best_label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx)
|
||||
best_label_data = DataLoadUtil.load_label(best_label_path)
|
||||
first_frame = best_label_data["best_sequence"][0]
|
||||
best_seq_len = len(best_label_data["best_sequence"])
|
||||
frame_len = DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)
|
||||
|
||||
for i in range(frame_len):
|
||||
path = DataLoadUtil.get_path(self.root_dir, scene_name, i)
|
||||
pts = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
|
||||
if pts.shape[0] == 0:
|
||||
continue
|
||||
datalist.append({
|
||||
"scene_name": scene_name,
|
||||
"first_frame": first_frame,
|
||||
"best_seq_len": best_seq_len,
|
||||
"max_coverage_rate": scene_max_coverage_rate,
|
||||
"first_frame": i,
|
||||
"best_seq_len": -1,
|
||||
"max_coverage_rate": 1.0,
|
||||
"label_idx": scene_max_cr_idx,
|
||||
})
|
||||
return datalist
|
||||
@@ -131,8 +123,7 @@ class SeqReconstructionDataset(BaseDataset):
|
||||
scanned_n_to_world_pose,
|
||||
) = ([], [], [])
|
||||
view = data_item_info["first_frame"]
|
||||
frame_idx = view[0]
|
||||
coverage_rate = view[1]
|
||||
frame_idx = view
|
||||
view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
|
||||
cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
|
||||
|
||||
@@ -144,7 +135,7 @@ class SeqReconstructionDataset(BaseDataset):
|
||||
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])
|
||||
)
|
||||
@@ -161,7 +152,6 @@ class SeqReconstructionDataset(BaseDataset):
|
||||
gt_pts = self.seq_combined_pts(scene_name, frame_list)
|
||||
data_item = {
|
||||
"first_scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
|
||||
"first_scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
|
||||
"first_scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
|
||||
"seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
|
||||
"best_seq_len": best_seq_len, # Int
|
||||
@@ -180,39 +170,35 @@ class SeqReconstructionDataset(BaseDataset):
|
||||
# -------------- Debug ---------------- #
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import pickle
|
||||
import os
|
||||
|
||||
seed = 0
|
||||
torch.manual_seed(seed)
|
||||
np.random.seed(seed)
|
||||
'''
|
||||
OmniObject3d_test:
|
||||
root_dir: "H:\\AI\\Datasets\\packed_test_data"
|
||||
model_dir: "H:\\AI\\Datasets\\scaled_object_meshes"
|
||||
source: seq_reconstruction_dataset
|
||||
split_file: "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt"
|
||||
type: test
|
||||
filter_degree: 75
|
||||
eval_list:
|
||||
- pose_diff
|
||||
- coverage_rate_increase
|
||||
ratio: 0.1
|
||||
batch_size: 1
|
||||
num_workers: 12
|
||||
pts_num: 8192
|
||||
load_from_preprocess: True
|
||||
'''
|
||||
|
||||
config = {
|
||||
"root_dir": "H:\\AI\\Datasets\\packed_test_data",
|
||||
"root_dir": "/media/hofee/data/data/new_testset",
|
||||
"source": "seq_reconstruction_dataset",
|
||||
"split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
|
||||
"split_file": "/media/hofee/data/data/OmniObject3d_test.txt",
|
||||
"load_from_preprocess": True,
|
||||
"ratio": 1,
|
||||
"filter_degree": 75,
|
||||
"num_workers": 0,
|
||||
"pts_num": 8192,
|
||||
"type": "test",
|
||||
"type": namespace.Mode.TEST,
|
||||
}
|
||||
ds = SeqReconstructionDataset(config)
|
||||
print(len(ds))
|
||||
print(ds.__getitem__(10))
|
||||
|
||||
output_dir = "/media/hofee/data/data/new_testset_output"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
ds = SeqReconstructionDataset(config)
|
||||
for i in tqdm(range(len(ds)), desc="processing dataset"):
|
||||
output_path = os.path.join(output_dir, f"item_{i}.pkl")
|
||||
item = ds.__getitem__(i)
|
||||
for key, value in item.items():
|
||||
if isinstance(value, np.ndarray):
|
||||
item[key] = value.tolist()
|
||||
#import ipdb; ipdb.set_trace()
|
||||
with open(output_path, "wb") as f:
|
||||
pickle.dump(item, f)
|
||||
|
||||
@@ -15,21 +15,19 @@ from utils.data_load import DataLoadUtil
|
||||
from utils.pose import PoseUtil
|
||||
from utils.pts import PtsUtil
|
||||
|
||||
|
||||
@stereotype.dataset("seq_reconstruction_dataset_preprocessed")
|
||||
class SeqReconstructionDatasetPreprocessed(BaseDataset):
|
||||
def __init__(self, config):
|
||||
super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
|
||||
self.config = config
|
||||
self.root_dir = config["root_dir"]
|
||||
self.real_root_dir = r"H:\AI\Datasets\packed_test_data"
|
||||
self.real_root_dir = r"/media/hofee/data/data/new_testset"
|
||||
self.item_list = os.listdir(self.root_dir)
|
||||
|
||||
def __getitem__(self, index):
|
||||
data = pickle.load(open(os.path.join(self.root_dir, self.item_list[index]), "rb"))
|
||||
data_item = {
|
||||
"first_scanned_pts": np.asarray(data["first_scanned_pts"], dtype=np.float32), # Ndarray(S x Nv x 3)
|
||||
"first_scanned_coverage_rate": data["first_scanned_coverage_rate"], # List(S): Float, range(0, 1)
|
||||
"first_scanned_n_to_world_pose_9d": np.asarray(data["first_scanned_n_to_world_pose_9d"], dtype=np.float32), # Ndarray(S x 9)
|
||||
"seq_max_coverage_rate": data["seq_max_coverage_rate"], # Float, range(0, 1)
|
||||
"best_seq_len": data["best_seq_len"], # Int
|
||||
@@ -43,7 +41,6 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset):
|
||||
def __len__(self):
|
||||
return len(self.item_list)
|
||||
|
||||
|
||||
# -------------- Debug ---------------- #
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
|
||||
Reference in New Issue
Block a user