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

update basic framework

Browse Source
This commit is contained in:
hofee
2024-08-21 17:11:56 +08:00
parent 73dcd592df
commit f977fd4b8e
29 changed files with 1393 additions and 719 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
runners/runner.py
View File

@@ -1,59 +0,0 @@
import os
import time
from abc import abstractmethod, ABC
import numpy as np
import torch
from configs.config import ConfigManager
class Runner(ABC):
@abstractmethod
def __init__(self, config_path):
ConfigManager.load_config_with(config_path)
ConfigManager.print_config()
seed = ConfigManager.get("settings", "general", "seed")
self.device = ConfigManager.get("settings", "general", "device")
self.cuda_visible_devices = ConfigManager.get("settings","general","cuda_visible_devices")
os.environ["CUDA_VISIBLE_DEVICES"] = self.cuda_visible_devices
self.experiments_config = ConfigManager.get("settings", "experiment")
self.experiment_path = os.path.join(self.experiments_config["root_dir"], self.experiments_config["name"])
np.random.seed(seed)
torch.manual_seed(seed)
lt = time.localtime()
self.file_name = f"{lt.tm_year}_{lt.tm_mon}_{lt.tm_mday}_{lt.tm_hour}h{lt.tm_min}m{lt.tm_sec}s"
@abstractmethod
def run(self):
pass
@abstractmethod
def load_experiment(self, backup_name=None):
if not os.path.exists(self.experiment_path):
print(f"experiments environment {self.experiments_config['name']} does not exists.")
self.create_experiment(backup_name)
else:
print(f"experiments environment {self.experiments_config['name']}")
backup_config_dir = os.path.join(str(self.experiment_path), "configs")
if not os.path.exists(backup_config_dir):
os.makedirs(backup_config_dir)
ConfigManager.backup_config_to(backup_config_dir, self.file_name, backup_name)
@abstractmethod
def create_experiment(self, backup_name=None):
print("creating experiment: " + self.experiments_config["name"])
os.makedirs(self.experiment_path)
backup_config_dir = os.path.join(str(self.experiment_path), "configs")
os.makedirs(backup_config_dir)
ConfigManager.backup_config_to(backup_config_dir, self.file_name, backup_name)
log_dir = os.path.join(str(self.experiment_path), "log")
os.makedirs(log_dir)
cache_dir = os.path.join(str(self.experiment_path), "cache")
os.makedirs(cache_dir)
def print_info(self):
table_size = 80
print("+" + "-" * table_size + "+")
print(f"| Experiment <{self.experiments_config['name']}>")
print("+" + "-" * table_size + "+")

73
runners/strategy_generator.py Normal file
View File

@@ -0,0 +1,73 @@
import os
from PytorchBoot.runners.runner import Runner
from PytorchBoot.config import ConfigManager
import PytorchBoot.stereotype as stereotype
@stereotype.runner("strategy_generator")
class StrategyGenerator(Runner):
def __init__(self, config):
super().__init__(config)
self.load_experiment("generate")
def run(self):
self.demo(seq=16,num=100)
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 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")
pts_list = []
for idx in range(0,num):
path = DataLoadUtil.get_path(data_dir, seq, idx)
point_cloud = DataLoadUtil.get_point_cloud_world_from_path(path)
sampled_point_cloud = ReconstructionUtil.downsample_point_cloud(point_cloud, 0.005)
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))

261
runners/trainer.py
View File

@@ -1,261 +0,0 @@
import os
import sys
import json
from datetime import datetime
import torch
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter
path = os.path.abspath(__file__)
for i in range(2):
path = os.path.dirname(path)
PROJECT_ROOT = path
sys.path.append(PROJECT_ROOT)
from configs.config import ConfigManager
from datasets.dataset_factory import DatasetFactory
from optimizers.optimizer_factory import OptimizerFactory
from evaluations.eval_function_factory import EvalFunctionFactory
from losses.loss_function_factory import LossFunctionFactory
from modules.pipeline import Pipeline
from runners.runner import Runner
from utils.tensorboard_util import TensorboardWriter
from annotations.external_module import EXTERNAL_FREEZE_MODULES
class Trainer(Runner):
CHECKPOINT_DIR_NAME: str = 'checkpoints'
TENSORBOARD_DIR_NAME: str = 'tensorboard'
LOG_DIR_NAME: str = 'log'
def __init__(self, config_path):
super().__init__(config_path)
tensorboard_path = os.path.join(self.experiment_path, Trainer.TENSORBOARD_DIR_NAME)
''' Pipeline '''
self.pipeline_config = ConfigManager.get("settings", "pipeline")
self.parallel = ConfigManager.get("settings","general","parallel")
self.pipeline = Pipeline(self.pipeline_config)
if self.parallel and self.device == "cuda":
self.pipeline = torch.nn.DataParallel(self.pipeline)
self.pipeline = self.pipeline.to(self.device)
''' Experiment '''
self.current_epoch = 0
self.max_epochs = self.experiments_config["max_epochs"]
self.test_first = self.experiments_config["test_first"]
self.load_experiment("train")
''' Train '''
self.train_config = ConfigManager.get("settings", "train")
self.train_dataset_config = self.train_config["dataset"]
self.train_set = DatasetFactory.create(self.train_dataset_config)
self.optimizer = OptimizerFactory.create(self.train_config["optimizer"], self.pipeline.parameters())
self.train_writer = SummaryWriter(
log_dir=os.path.join(tensorboard_path, f"[train]{self.train_dataset_config['name']}"))
''' Test '''
self.test_config = ConfigManager.get("settings", "test")
self.test_dataset_config_list = self.test_config["dataset_list"]
self.test_set_list = []
self.test_writer_list = []
seen_name = set()
for test_dataset_config in self.test_dataset_config_list:
if test_dataset_config["name"] not in seen_name:
seen_name.add(test_dataset_config["name"])
else:
raise ValueError("Duplicate test dataset name: {}".format(test_dataset_config["name"]))
test_set = DatasetFactory.create(test_dataset_config)
test_writer = SummaryWriter(
log_dir=os.path.join(tensorboard_path, f"[test]{test_dataset_config['name']}"))
self.test_set_list.append(test_set)
self.test_writer_list.append(test_writer)
del seen_name
self.print_info()
def run(self):
save_interval = self.experiments_config["save_checkpoint_interval"]
if self.current_epoch != 0:
print("Continue training from epoch {}.".format(self.current_epoch))
else:
print("Start training from initial model.")
if self.test_first:
print("Do test first.")
self.test()
while self.current_epoch < self.max_epochs:
self.current_epoch += 1
self.train()
self.test()
if self.current_epoch % save_interval == 0:
self.save_checkpoint()
self.save_checkpoint(is_last=True)
def train(self):
self.pipeline.train()
train_set_name = self.train_dataset_config["name"]
ratio = self.train_dataset_config["ratio"]
train_loader = self.train_set.get_loader(device="cuda", shuffle=True)
loop = tqdm(enumerate(train_loader), total=len(train_loader))
loader_length = len(train_loader)
for i, data in loop:
self.train_set.process_batch(data, self.device)
loss_dict = self.train_step(data)
loop.set_description(
f'Epoch [{self.current_epoch}/{self.max_epochs}] (Train: {train_set_name}, ratio={ratio})')
loop.set_postfix(loss=loss_dict)
curr_iters = (self.current_epoch - 1) * loader_length + i
TensorboardWriter.write_tensorboard(self.train_writer, "iter", loss_dict, curr_iters)
def train_step(self, data):
self.optimizer.zero_grad()
output = self.pipeline(data, Pipeline.TRAIN_MODE)
total_loss, loss_dict = self.loss_fn(output, data)
total_loss.backward()
self.optimizer.step()
for k, v in loss_dict.items():
loss_dict[k] = round(v, 5)
return loss_dict
def loss_fn(self, output, data):
loss_config = self.train_config["losses"]
loss_dict = {}
total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.device)
for key in loss_config:
weight = loss_config[key]
target_loss_fn = LossFunctionFactory.create(key)
loss = target_loss_fn(output, data)
loss_dict[key] = loss.item()
total_loss += weight * loss
loss_dict['total_loss'] = total_loss.item()
return total_loss, loss_dict
def test(self):
self.pipeline.eval()
with torch.no_grad():
for dataset_idx, test_set in enumerate(self.test_set_list):
eval_list = self.test_dataset_config_list[dataset_idx]["eval_list"]
test_set_name = self.test_dataset_config_list[dataset_idx]["name"]
ratio = self.test_dataset_config_list[dataset_idx]["ratio"]
writer = self.test_writer_list[dataset_idx]
output_list = []
data_list = []
test_loader = test_set.get_loader("cpu")
loop = tqdm(enumerate(test_loader), total=int(len(test_loader)))
for i, data in loop:
test_set.process_batch(data, self.device)
output = self.pipeline(data, Pipeline.TEST_MODE)
output_list.append(output)
data_list.append(data)
loop.set_description(
f'Epoch [{self.current_epoch}/{self.max_epochs}] (Test: {test_set_name}, ratio={ratio})')
result_dict = self.eval_fn(output_list, data_list, eval_list)
TensorboardWriter.write_tensorboard(writer, "epoch", result_dict, self.current_epoch - 1)
@staticmethod
def eval_fn(output_list, data_list, eval_list):
target_eval_fn = EvalFunctionFactory.create(eval_list)
result_dict = target_eval_fn(output_list, data_list)
return result_dict
def get_checkpoint_path(self, is_last=False):
return os.path.join(self.experiment_path, Trainer.CHECKPOINT_DIR_NAME,
"Epoch_{}.pth".format(
self.current_epoch if self.current_epoch != -1 and not is_last else "last"))
def load_checkpoint(self, is_last=False):
self.load(self.get_checkpoint_path(is_last))
print(f"Loaded checkpoint from {self.get_checkpoint_path(is_last)}")
if is_last:
checkpoint_root = os.path.join(self.experiment_path, Trainer.CHECKPOINT_DIR_NAME)
meta_path = os.path.join(checkpoint_root, "meta.json")
if not os.path.exists(meta_path):
raise FileNotFoundError(
"No checkpoint meta.json file in the experiment {}".format(self.experiments_config["name"]))
file_path = os.path.join(checkpoint_root, "meta.json")
with open(file_path, "r") as f:
meta = json.load(f)
self.current_epoch = meta["last_epoch"]
def save_checkpoint(self, is_last=False):
self.save(self.get_checkpoint_path(is_last))
if not is_last:
print(f"Checkpoint at epoch {self.current_epoch} saved to {self.get_checkpoint_path(is_last)}")
else:
meta = {
"last_epoch": self.current_epoch,
"time": str(datetime.now())
}
checkpoint_root = os.path.join(self.experiment_path, Trainer.CHECKPOINT_DIR_NAME)
file_path = os.path.join(checkpoint_root, "meta.json")
with open(file_path, "w") as f:
json.dump(meta, f)
def load_experiment(self, backup_name=None):
super().load_experiment(backup_name)
if self.experiments_config["use_checkpoint"]:
self.current_epoch = self.experiments_config["epoch"]
self.load_checkpoint(is_last=(self.current_epoch == -1))
def create_experiment(self, backup_name=None):
super().create_experiment(backup_name)
ckpt_dir = os.path.join(str(self.experiment_path), Trainer.CHECKPOINT_DIR_NAME)
os.makedirs(ckpt_dir)
tensorboard_dir = os.path.join(str(self.experiment_path), Trainer.TENSORBOARD_DIR_NAME)
os.makedirs(tensorboard_dir)
def load(self, path):
state_dict = torch.load(path)
if self.parallel:
self.pipeline.module.load_state_dict(state_dict)
else:
self.pipeline.load_state_dict(state_dict)
def save(self, path):
if self.parallel:
state_dict = self.pipeline.module.state_dict()
else:
state_dict = self.pipeline.state_dict()
for name, module in self.pipeline.named_modules():
if module.__class__ in EXTERNAL_FREEZE_MODULES:
if name in state_dict:
del state_dict[name]
torch.save(state_dict, path)
def print_info(self):
def print_dataset(config, dataset):
print("\t name: {}".format(config["name"]))
print("\t source: {}".format(config["source"]))
print("\t data_type: {}".format(config["data_type"]))
print("\t total_length: {}".format(len(dataset)))
print("\t ratio: {}".format(config["ratio"]))
print()
super().print_info()
table_size = 70
print(f"{'+' + '-' * (table_size // 2)} Pipeline {'-' * (table_size // 2)}" + '+')
print(self.pipeline)
print(f"{'+' + '-' * (table_size // 2)} Datasets {'-' * (table_size // 2)}" + '+')
print("train dataset: ")
print_dataset(self.train_dataset_config, self.train_set)
for i, test_dataset_config in enumerate(self.test_dataset_config_list):
print(f"test dataset {i}: ")
print_dataset(test_dataset_config, self.test_set_list[i])
print(f"{'+' + '-' * (table_size // 2)}----------{'-' * (table_size // 2)}" + '+')
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, default="configs/train_config.yaml")
args = parser.parse_args()
trainer = Trainer(args.config)
trainer.run()