def build(input_reader_config,
model_config,
training,
voxel_generator,
target_assigner=None):
"""Builds a tensor dictionary based on the InputReader config.
Args:
input_reader_config: A input_reader_pb2.InputReader object.
Returns:
A tensor dict based on the input_reader_config.
Raises:
ValueError: On invalid input reader proto.
ValueError: If no input paths are specified.
"""
if not isinstance(input_reader_config, input_reader_pb2.InputReader):
raise ValueError('input_reader_config not of type '
'input_reader_pb2.InputReader.')
generate_bev = model_config.use_bev
without_reflectivity = model_config.without_reflectivity
num_point_features = model_config.num_point_features
out_size_factor = model_config.rpn.layer_strides[
0] // model_config.rpn.upsample_strides[0]
cfg = input_reader_config
db_sampler_cfg = input_reader_config.database_sampler
db_sampler = None
if len(db_sampler_cfg.sample_groups) > 0: # enable sample
db_sampler = dbsampler_builder.build(db_sampler_cfg)
u_db_sampler_cfg = input_reader_config.unlabeled_database_sampler
u_db_sampler = None
if len(u_db_sampler_cfg.sample_groups) > 0: # enable sample
u_db_sampler = dbsampler_builder.build(u_db_sampler_cfg)
grid_size = voxel_generator.grid_size
# [352, 400]
feature_map_size = grid_size[:2] // out_size_factor
feature_map_size = [*feature_map_size, 1][::-1]
prep_func = partial(
prep_pointcloud,
root_path=cfg.kitti_root_path,
class_names=list(cfg.class_names),
voxel_generator=voxel_generator,
target_assigner=target_assigner,
training=training,
max_voxels=cfg.max_number_of_voxels,
remove_outside_points=False,
remove_unknown=cfg.remove_unknown_examples,
create_targets=training,
shuffle_points=cfg.shuffle_points,
gt_rotation_noise=list(cfg.groundtruth_rotation_uniform_noise),
gt_loc_noise_std=list(cfg.groundtruth_localization_noise_std),
global_rotation_noise=list(cfg.global_rotation_uniform_noise),
global_scaling_noise=list(cfg.global_scaling_uniform_noise),
global_loc_noise_std=(0.2, 0.2, 0.2),
global_random_rot_range=list(
cfg.global_random_rotation_range_per_object),
db_sampler=db_sampler,
unlabeled_db_sampler=u_db_sampler,
generate_bev=generate_bev,
without_reflectivity=without_reflectivity,
num_point_features=num_point_features,
anchor_area_threshold=cfg.anchor_area_threshold,
gt_points_drop=cfg.groundtruth_points_drop_percentage,
gt_drop_max_keep=cfg.groundtruth_drop_max_keep_points,
remove_points_after_sample=cfg.remove_points_after_sample,
remove_environment=cfg.remove_environment,
use_group_id=cfg.use_group_id,
out_size_factor=out_size_factor)
dataset = KittiDataset(info_path=cfg.kitti_info_path,
root_path=cfg.kitti_root_path,
num_point_features=num_point_features,
target_assigner=target_assigner,
feature_map_size=feature_map_size,
prep_func=prep_func)
return dataset
for epoch in range(1, args.epochs + 1):
lr_scheduler.step(epoch)
total_it = train_one_epoch(model, train_loader, optimizer, epoch, lr_scheduler, total_it, tb_log, log_f)
if epoch % args.ckpt_save_interval == 0:
with torch.no_grad():
avg_iou = eval_one_epoch(model, eval_loader, epoch, tb_log, log_f)
ckpt_name = os.path.join(ckpt_dir, 'checkpoint_epoch_%d' % epoch)
save_checkpoint(model, epoch, ckpt_name)
if __name__ == '__main__':
MODEL = importlib.import_module(args.net) # import network module
model = MODEL.get_model(input_channels=0)
eval_set = KittiDataset(root_dir='data', mode='EVAL', split='val')
eval_loader = DataLoader(eval_set, batch_size=args.batch_size, shuffle=False, pin_memory=True,
num_workers=args.workers, collate_fn=eval_set.collate_batch)
if args.mode == 'train':
train_set = KittiDataset(root_dir='data', mode='TRAIN', split='train')
train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, pin_memory=True,
num_workers=args.workers, collate_fn=train_set.collate_batch)
# output dir config
output_dir = os.path.join(args.output_dir, args.extra_tag)
os.makedirs(output_dir, exist_ok=True)
tb_log.configure(os.path.join(output_dir, 'tensorboard'))
ckpt_dir = os.path.join(output_dir, 'ckpt')
os.makedirs(ckpt_dir, exist_ok=True)
log_file = os.path.join(output_dir, 'log.txt')
import numpy as np
from pathlib import Path
from torchvision.transforms import Compose
from data.dataset import KittiDataset
from data.paths import kitti_paths
from data.transforms import WhiteningImage
if __name__ == '__main__':
trans = Compose([WhiteningImage()])
dataset = KittiDataset(kitti_paths(
(Path('/home/morsic/datasets/kitti_flow/training'))),
transforms=trans)
mean_l = np.zeros((3, ))
std_l = np.zeros((3, ))
mean_r = np.zeros((3, ))
std_r = np.zeros((3, ))
for triplet in dataset:
L, R, _ = triplet
mean_l += np.mean(L, axis=(0, 1))
std_l += np.std(L, axis=(0, 1))
mean_r += np.mean(R, axis=(0, 1))
std_r += np.std(R, axis=(0, 1))
print(f'Mean left: {mean_l / len(dataset)}')
print(f'Std left: {std_l / len(dataset)}')
print(f'Mean right: {mean_r / len(dataset)}')
parser = argparse.ArgumentParser(description='Create HDF5 dataset storage')
parser.add_argument('root',
type=str,
help='Dataset root',
default='/home/morsic/datasets/kitti_flow/training')
parser.add_argument('--dest',
type=str,
help='Path to HDF5',
default='kitti.hdf5')
if __name__ == '__main__':
args = parser.parse_args()
root = Path(args.root)
paths = kitti_paths(root)
dataset = KittiDataset(paths)
offset, created = 0, False
f = h5py.File(args.dest, "w")
for i, example in tqdm(enumerate(dataset), total=len(dataset)):
patches = generate_examples(example)
N, P, C, H, W = patches.shape
if not created:
dset = f.create_dataset('triplets',
data=patches,
maxshape=(None, P, C, H, W))
else:
new_dset = f.get('triplets')
new_dset.resize(len(new_dset) + N, axis=0)
new_dset[offset:offset + N, :, :, :, :] = patches
created = True
def build_dataset(input_reader_config,
model_config,
training,
voxel_generator,
target_assigner=None):
"""Builds a tensor dictionary based on the InputReader config.
Returns:
A tensor dict based on the input_reader_config.
"""
generate_bev = model_config.POST_PROCESSING.use_bev
without_reflectivity = model_config.WITHOUT_REFLECTIVITY
num_point_features = model_config.NUM_POINT_FEATURES
out_size_factor = model_config.BACKBONE.layer_strides[
0] // model_config.BACKBONE.upsample_strides[0]
cfg = input_reader_config
db_sampler_cfg = input_reader_config.DATABASE_SAMPLER
db_sampler = None
if len(db_sampler_cfg.sample_groups) > 0:
db_sampler = build_dbsampler(db_sampler_cfg)
try:
u_db_sampler_cfg = input_reader_config.UNLABELED_DATABASE_SAMPLER
u_db_sampler = None
if len(u_db_sampler_cfg.sample_groups) > 0:
u_db_sampler = build_dbsampler(u_db_sampler_cfg)
except:
u_db_sampler = None
grid_size = voxel_generator.grid_size #[352,400]
feature_map_size = grid_size[:2] // out_size_factor
feature_map_size = [*feature_map_size, 1][::-1]
prep_func = partial(
prep_pointcloud,
root_path=cfg.KITTI_ROOT_PATH,
class_names=cfg.CLASS_NAMES,
voxel_generator=voxel_generator,
target_assigner=target_assigner,
training=training,
max_voxels=cfg.MAX_NUMBER_OF_VOXELS,
remove_outside_points=False,
create_targets=training,
shuffle_points=cfg.SHUFFLE_POINTS,
gt_rotation_noise=cfg.GROUNDTRUTH_ROTATION_UNIFORM_NOISE,
gt_loc_noise_std=cfg.GROUNDTRUTH_LOCALIZATION_NOISE_STD,
global_rotation_noise=cfg.GLOBAL_ROTATION_UNIFORM_NOISE,
global_scaling_noise=cfg.GLOBAL_SCALING_UNIFORM_NOISE,
global_loc_noise_std=(0.2, 0.2, 0.2),
global_random_rot_range=cfg.GLOBAL_RANDOM_ROTATION_RANGE_PER_OBJECT,
db_sampler=db_sampler,
unlabeled_db_sampler=u_db_sampler,
generate_bev=generate_bev,
without_reflectivity=without_reflectivity,
num_point_features=num_point_features,
anchor_area_threshold=cfg.ANCHOR_AREA_THRESHOLD,
gt_points_drop=cfg.GROUNDTRUTH_POINTS_DROP_PERCENTAGE,
gt_drop_max_keep=cfg.GROUNDTRUTH_DROP_MAX_KEEP_POINTS,
remove_points_after_sample=cfg.REMOVE_POINTS_AFTER_SAMPLE,
remove_environment=cfg.REMOVE_ENVIRONMENT,
use_group_id=False,
out_size_factor=out_size_factor)
dataset = KittiDataset(info_path=cfg.KITTI_INFO_PATH,
root_path=cfg.KITTI_ROOT_PATH,
num_point_features=num_point_features,
target_assigner=target_assigner,
feature_map_size=feature_map_size,
prep_func=prep_func)
return dataset
from models.regression import ResnetRegression
from models.classification import SoftArgmin
from data.dataset_paths import get_all_kitti_paths
from data.dataset import KittiDataset
from data.transform import Padding, RandomCrop, ToTensor
model = StereoRegression(
ResnetUnary(),
CostVolumeConcat(),
ResnetRegression(),
SoftArgmin(),
)
size = (1, 3, 512, 256)
model.cuda()
# disparity = out.data.cpu().numpy().squeeze()
transform = Compose((RandomCrop((256, 512)), ToTensor()))
dataset = KittiDataset(get_all_kitti_paths(), transform)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4)
example = dataset[0]
l, r, d = example['left'], example['right'], example['disparity']
disp = model(Variable(l.unsqueeze_(0).cuda()),
Variable(r.unsqueeze_(0).cuda()))