hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

upd

Browse Source
This commit is contained in:
hofee 2024-11-07 19:33:18 +08:00
parent 5bcd0fc6e3
commit 2cd811c1b7
8 changed files with 77 additions and 105 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
configs/local/inference_config.yaml
View File

@ -14,8 +14,8 @@ runner:
dataset_list:
- OmniObject3d_test
blender_script_path: "C:\\Document\\Local Project\\nbv_rec\\blender\\data_renderer.py"
output_dir: "C:\\Document\\Datasets\\inference_scan_pts_overlap_global_full_on_testset"
blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py"
output_dir: "/media/hofee/data/data/new_inference_test_output"
pipeline: nbv_reconstruction_pipeline
voxel_size: 0.003
@ -34,10 +34,10 @@ dataset:
# load_from_preprocess: True
OmniObject3d_test:
root_dir: "C:\\Document\\Datasets\\inference_test"
model_dir: "C:\\Document\\Datasets\\scaled_object_meshes"
root_dir: "/media/hofee/data/data/new_testset_output"
model_dir: "/media/hofee/data/data/scaled_object_meshes"
source: seq_reconstruction_dataset_preprocessed
split_file: "C:\\Document\\Datasets\\data_list\\OmniObject3d_test.txt"
# split_file: "C:\\Document\\Datasets\\data_list\\OmniObject3d_test.txt"
type: test
filter_degree: 75
eval_list:

12
configs/local/view_generate_config.yaml
View File

@ -7,17 +7,19 @@ runner:
name: debug
root_dir: experiments
generate:
port: 5002
from: 600
port: 5000
from: 0
to: -1 # -1 means all
object_dir: /media/hofee/data/data/object_meshes_part1
object_dir: /media/hofee/data/data/scaled_object_meshes
table_model_path: "/media/hofee/data/data/others/table.obj"
output_dir: /media/hofee/repository/data_part_1
output_dir: /media/hofee/data/data/new_testset
object_list_path: /media/hofee/data/data/OmniObject3d_test.txt
use_list: True
binocular_vision: true
plane_size: 10
max_views: 512
min_views: 128
random_view_ratio: 0.02
random_view_ratio: 0.01
min_cam_table_included_degree: 20
max_diag: 0.7
min_diag: 0.01

86
core/seq_dataset.py
View File

@ -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,7 +47,8 @@ class SeqReconstructionDataset(BaseDataset):
with open(self.split_file_path, "r") as f:
for line in f:
scene_name = line.strip()
scene_name_list.append(scene_name)
if os.path.exists(os.path.join(self.root_dir, scene_name)):
scene_name_list.append(scene_name)
return scene_name_list
def get_scene_name_list(self):
@ -58,29 +59,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 +122,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 +134,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 +151,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
@ -179,40 +168,37 @@ class SeqReconstructionDataset(BaseDataset):
# -------------- Debug ---------------- #
if __name__ == "__main__":
#import ipdb; ipdb.set_trace()
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)

3
core/seq_dataset_preprocessed.py
View File

@ -22,14 +22,13 @@ class SeqReconstructionDatasetPreprocessed(BaseDataset):
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

4
preprocess/pack_upload_data.py
View File

@ -29,8 +29,8 @@ def pack_all_scenes(root, scene_list, output_dir):
pack_scene_data(root, scene, output_dir)
if __name__ == "__main__":
root = r"H:\AI\Datasets\nbv_rec_part2"
output_dir = r"H:\AI\Datasets\upload_part2"
root = r"/media/hofee/repository/data_part_1"
output_dir = r"/media/hofee/repository/upload_part1"
scene_list = os.listdir(root)
from_idx = 0
to_idx = len(scene_list)

10
preprocess/preprocessor.py
View File

@ -164,10 +164,10 @@ def save_scene_data(root, scene, scene_idx=0, scene_total=1,file_type="txt"):
if __name__ == "__main__":
#root = "/media/hofee/repository/new_data_with_normal"
root = r"H:\AI\Datasets\nbv_rec_part2"
root = "/media/hofee/data/data/new_testset"
scene_list = os.listdir(root)
from_idx = 0 # 1000
to_idx = 600 # 1500
to_idx = len(scene_list) # 1500
cnt = 0
@ -179,7 +179,11 @@ if __name__ == "__main__":
print(f"Scene {scene} has been processed")
cnt+=1
continue
save_scene_data(root, scene, cnt, total, file_type="npy")
try:
save_scene_data(root, scene, cnt, total, file_type="npy")
except Exception as e:
print(f"Error occurred when processing scene {scene}")
print(e)
cnt+=1
end = time.time()
print(f"Time cost: {end-start}")

54
runners/inferencer.py
View File

@ -76,6 +76,8 @@ class Inferencer(Runner):
for i in tqdm(range(total), desc=f"Processing {test_set_name}", ncols=100):
data = test_set.__getitem__(i)
scene_name = data["scene_name"]
if scene_name != "omniobject3d-book_004":
continue
inference_result_path = os.path.join(self.output_dir, test_set_name, f"{scene_name}.pkl")
if os.path.exists(inference_result_path):
Log.info(f"Inference result already exists for scene: {scene_name}")
@ -87,7 +89,7 @@ class Inferencer(Runner):
status_manager.set_progress("inference", "inferencer", f"dataset", len(self.test_set_list), len(self.test_set_list))
def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 7):
def predict_sequence(self, data, cr_increase_threshold=0, overlap_area_threshold=25, scan_points_threshold=10, max_iter=50, max_retry = 5):
scene_name = data["scene_name"]
Log.info(f"Processing scene: {scene_name}")
status_manager.set_status("inference", "inferencer", "scene", scene_name)
@ -110,10 +112,13 @@ class Inferencer(Runner):
input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(data["first_scanned_n_to_world_pose_9d"], dtype=torch.float32).to(self.device)]
input_data["mode"] = namespace.Mode.TEST
input_pts_N = input_data["combined_scanned_pts"].shape[1]
root = os.path.dirname(scene_path)
display_table_info = DataLoadUtil.get_display_table_info(root, scene_name)
radius = display_table_info["radius"]
scan_points = np.asarray(ReconstructionUtil.generate_scan_points(display_table_top=0,display_table_radius=radius))
first_frame_target_pts, first_frame_target_normals, first_frame_scan_points_indices = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, scan_points, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
scanned_view_pts = [first_frame_target_pts]
history_indices = [first_frame_scan_points_indices]
@ -124,6 +129,7 @@ class Inferencer(Runner):
retry = 0
pred_cr_seq = [last_pred_cr]
success = 0
last_pts_num = PtsUtil.voxel_downsample_point_cloud(data["first_scanned_pts"][0], 0.002).shape[0]
import time
while len(pred_cr_seq) < max_iter and retry < max_retry:
start_time = time.time()
@ -146,7 +152,7 @@ class Inferencer(Runner):
curr_overlap_area_threshold = overlap_area_threshold * 0.5
downsampled_new_target_pts = PtsUtil.voxel_downsample_point_cloud(new_target_pts, voxel_threshold)
overlap, new_added_pts_num = ReconstructionUtil.check_overlap(downsampled_new_target_pts, down_sampled_model_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
overlap, _ = ReconstructionUtil.check_overlap(downsampled_new_target_pts, down_sampled_model_pts, overlap_area_threshold = curr_overlap_area_threshold, voxel_size=voxel_threshold, require_new_added_pts_num = True)
if not overlap:
retry += 1
retry_overlap_pose.append(pred_pose.cpu().numpy().tolist())
@ -170,31 +176,22 @@ class Inferencer(Runner):
continue
start_time = time.time()
pred_cr, covered_pts_num = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
pred_cr, _ = self.compute_coverage_rate(scanned_view_pts, new_target_pts, down_sampled_model_pts, threshold=voxel_threshold)
end_time = time.time()
print(f"Time taken for coverage rate computation: {end_time - start_time} seconds")
print(pred_cr, last_pred_cr, " max: ", data["seq_max_coverage_rate"])
print("new added pts num: ", new_added_pts_num)
if pred_cr >= data["seq_max_coverage_rate"] - 1e-3:
print("max coverage rate reached!: ", pred_cr)
success += 1
elif new_added_pts_num < 5:
#success += 1
print("min added pts num reached!: ", new_added_pts_num)
if pred_cr <= last_pred_cr + cr_increase_threshold:
retry += 1
retry_duplication_pose.append(pred_pose.cpu().numpy().tolist())
continue
retry = 0
pred_cr_seq.append(pred_cr)
scanned_view_pts.append(new_target_pts)
down_sampled_new_pts_world = PtsUtil.random_downsample_point_cloud(new_target_pts, input_pts_N)
new_pts = down_sampled_new_pts_world
input_data["scanned_n_to_world_pose_9d"] = [torch.cat([input_data["scanned_n_to_world_pose_9d"][0], pred_pose_9d], dim=0)]
combined_scanned_pts = np.concatenate([input_data["combined_scanned_pts"][0].cpu().numpy(), new_pts], axis=0)
combined_scanned_pts = np.vstack(scanned_view_pts)
voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_pts, 0.002)
random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, input_pts_N)
input_data["combined_scanned_pts"] = torch.tensor(random_downsampled_combined_scanned_pts_np, dtype=torch.float32).unsqueeze(0).to(self.device)
@ -202,29 +199,12 @@ class Inferencer(Runner):
if success > 3:
break
last_pred_cr = pred_cr
input_data["scanned_n_to_world_pose_9d"] = input_data["scanned_n_to_world_pose_9d"][0].cpu().numpy().tolist()
result = {
"pred_pose_9d_seq": input_data["scanned_n_to_world_pose_9d"],
"combined_scanned_pts": input_data["combined_scanned_pts"],
"target_pts_seq": scanned_view_pts,
"coverage_rate_seq": pred_cr_seq,
"max_coverage_rate": data["seq_max_coverage_rate"],
"pred_max_coverage_rate": max(pred_cr_seq),
"scene_name": scene_name,
"retry_no_pts_pose": retry_no_pts_pose,
"retry_duplication_pose": retry_duplication_pose,
"retry_overlap_pose": retry_overlap_pose,
"best_seq_len": data["best_seq_len"],
}
self.stat_result[scene_name] = {
"coverage_rate_seq": pred_cr_seq,
"pred_max_coverage_rate": max(pred_cr_seq),
"pred_seq_len": len(pred_cr_seq),
}
print('success rate: ', max(pred_cr_seq))
return result
pts_num = voxel_downsampled_combined_scanned_pts_np.shape[0]
if pts_num - last_pts_num < 10 and pred_cr < data["seq_max_coverage_rate"] - 1e-3:
retry += 1
retry_duplication_pose.append(pred_pose.cpu().numpy().tolist())
print("delta pts num < 10:", pts_num, last_pts_num)
last_pts_num = pts_num
def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
if new_pts is not None:

3
utils/render.py
View File

@ -86,9 +86,10 @@ class RenderUtil:
json.dump(params, f)
start_time = time.time()
result = subprocess.run([
'blender', '-b', '-P', script_path, '--', temp_dir
'/home/hofee/blender-4.0.2-linux-x64/blender', '-b', '-P', script_path, '--', temp_dir
], capture_output=True, text=True)
end_time = time.time()
print(f"-- Time taken for blender: {end_time - start_time} seconds")
path = os.path.join(temp_dir, "tmp")
cam_info = DataLoadUtil.load_cam_info(path, binocular=True)