def test_myimagefloder():
show_images = True
scale_factor = 1.0
dataset_folder = '/datasets/hrvs/carla-highres/trainingF'
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
dataset_folder)
loader_eth3 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=2)
inv_t = get_inv_transform()
for left_img, right_img, disp in loader_eth3:
if show_images:
left_img_np = np.array(inv_t(left_img)).astype(np.uint8)
right_img_np = np.array(inv_t(right_img)).astype(np.uint8)
disp_img = cv2.normalize(disp,
None,
alpha=0,
beta=1,
norm_type=cv2.NORM_MINMAX,
dtype=cv2.CV_32F)
cv2.imshow('left_img', left_img_np[:, :, ::-1])
cv2.imshow('right_img', right_img_np[:, :, ::-1])
cv2.imshow('disp_img', disp_img)
cv2.waitKey(0)
assert left_img.shape == (3, 512, 768)
assert right_img.shape == (3, 512, 768)
assert disp.shape == (512, 768)
break
def test_dataloader(dataset_folder):
left, right, disp_l, disp_r = dataloader(dataset_folder)
assert len(left) > 0
assert len(left) == len(right) == len(disp_l) == len(disp_r)
i = 0
for (l, r, dl, dr) in zip(left, right, disp_l, disp_r):
if i == 0:
assert os.path.exists(l)
assert os.path.exists(r)
assert os.path.exists(dl)
i += 1
assert os.path.dirname(l) == os.path.dirname(r) == os.path.dirname(
dl) == os.path.dirname(dr)
assert 'im0' in os.path.basename(l)
assert 'im1' in os.path.basename(r)
assert 'disp0GT' in os.path.basename(dl)
assert 'disp1GT' in os.path.basename(dr)
def init_dataloader(input_args):
batch_size = input_args.batchsize
scale_factor = input_args.maxdisp / 384. # controls training resolution
hrvs_folder = '%s/hrvs/carla-highres/trainingF' % input_args.database
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
hrvs_folder)
loader_carla = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=2)
middlebury_folder = '%s/middlebury/mb-ex-training/trainingF' % input_args.database
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
middlebury_folder)
loader_mb = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=0)
rand_scale = [0.9, 2.4 * scale_factor]
all_left_img, all_right_img, all_left_disp, all_right_disp = lt.dataloader(
'%s/sceneflow/' % input_args.database)
loader_scene = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=rand_scale,
order=2)
# change to trainval when finetuning on KITTI
all_left_img, all_right_img, all_left_disp, _, _, _ = lk15.dataloader(
'%s/kitti15/training/' % input_args.database, split='train')
loader_kitti15 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=rand_scale,
order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader(
'%s/kitti12/training/' % input_args.database)
loader_kitti12 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=rand_scale,
order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader(
'%s/eth3d/' % input_args.database)
loader_eth3d = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=rand_scale,
order=0)
all_left_img, all_right_img, all_left_disp, all_right_disp = lidar_dataloader(
'%s/lidar-hdsm-dataset/' % input_args.database)
loader_lidar = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.5, 1.1 * scale_factor],
rand_bright=[0.8, 1.2],
order=2,
flip_disp_ud=True,
occlusion_size=[10, 25])
data_inuse = torch.utils.data.ConcatDataset(
[loader_carla] * 10 + [loader_mb] * 150 + # 71 pairs
[loader_scene] + # 39K pairs 960x540
[loader_kitti15] + [loader_kitti12] * 24 + [loader_eth3d] * 300 +
[loader_lidar]) # 25K pairs
# airsim ~750
train_dataloader = torch.utils.data.DataLoader(data_inuse,
batch_size=batch_size,
shuffle=True,
num_workers=batch_size,
drop_last=True,
worker_init_fn=_init_fn)
print('%d batches per epoch' % (len(data_inuse) // batch_size))
return train_dataloader
def _init_fn(worker_id):
np.random.seed()
random.seed()
torch.manual_seed(args.seed) # set again
torch.cuda.manual_seed(args.seed)
from dataloader import listfiles as ls
from dataloader import listsceneflow as lt
from dataloader import KITTIloader2015 as lk15
from dataloader import KITTIloader2012 as lk12
from dataloader import MiddleburyLoader as DA
batch_size = args.batchsize
scale_factor = args.maxdisp / 384. # controls training resolution
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader('%s/carla-highres/trainingF'%args.database)
loader_carla = DA.myImageFloder(all_left_img,all_right_img,all_left_disp,right_disparity=all_right_disp, rand_scale=[0.225,0.6*scale_factor], order=2)
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader('%s/mb-ex-training/trainingF'%args.database) # mb-ex
loader_mb = DA.myImageFloder(all_left_img,all_right_img,all_left_disp,right_disparity=all_right_disp, rand_scale=[0.225,0.6*scale_factor], order=0)
all_left_img, all_right_img, all_left_disp, all_right_disp = lt.dataloader('%s/sceneflow/'%args.database)
loader_scene = DA.myImageFloder(all_left_img,all_right_img,all_left_disp,right_disparity=all_right_disp, rand_scale=[0.9,2.4*scale_factor], order=2)
all_left_img, all_right_img, all_left_disp,_,_,_ = lk15.dataloader('%s/kitti_scene/training/'%args.database,typ='train') # trainval
loader_kitti15 = DA.myImageFloder(all_left_img,all_right_img,all_left_disp, rand_scale=[0.9,2.4*scale_factor], order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader('%s/data_stereo_flow/training/'%args.database)
loader_kitti12 = DA.myImageFloder(all_left_img,all_right_img,all_left_disp, rand_scale=[0.9,2.4*scale_factor], order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader('%s/eth3d/'%args.database)
loader_eth3d = DA.myImageFloder(all_left_img,all_right_img,all_left_disp, rand_scale=[0.9,2.4*scale_factor],order=0)
parser.add_argument('--clean', type=float, default=-1,
help='clean up output using entropy estimation')
parser.add_argument('--testres', type=float, default=0.5,
help='test time resolution ratio 0-x')
parser.add_argument('--max_disp', type=float, default=-1,
help='maximum disparity to search for')
parser.add_argument('--level', type=int, default=1,
help='output level of output, default is level 1 (stage 3),\
can also use level 2 (stage 2) or level 3 (stage 1)')
args = parser.parse_args()
# dataloader
from dataloader import listfiles as DA
test_left_img, test_right_img, _, _ = DA.dataloader(args.datapath)
# construct model
model = hsm(128,args.clean,level=args.level)
model = nn.DataParallel(model, device_ids=[0])
model.cuda()
if args.loadmodel is not None:
pretrained_dict = torch.load(args.loadmodel)
pretrained_dict['state_dict'] = {k:v for k,v in pretrained_dict['state_dict'].items() if 'disp' not in k}
model.load_state_dict(pretrained_dict['state_dict'],strict=False)
else:
print('run with random init')
print('Number of model parameters: {}'.format(sum([p.data.nelement() for p in model.parameters()])))
# dry run
def get_training_dataloader(maxdisp, dataset_folder):
scale_factor = maxdisp / 384.
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
'%s/hrvs/carla-highres/trainingF' % dataset_folder)
loader_carla = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=2)
all_left_img, all_right_img, all_left_disp, all_right_disp = ls.dataloader(
'%s/middlebury/mb-ex-training/trainingF' % dataset_folder) # mb-ex
loader_mb = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.225, 0.6 * scale_factor],
rand_bright=[0.8, 1.2],
order=0)
all_left_img, all_right_img, all_left_disp, all_right_disp = lt.dataloader(
'%s/sceneflow/' % dataset_folder)
loader_scene = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
right_disparity=all_right_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=2)
all_left_img, all_right_img, all_left_disp, _, _, _ = lk15.dataloader(
'%s/kitti15/training/' % dataset_folder,
typ='train') # change to trainval when finetuning on KITTI
loader_kitti15 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp = lk12.dataloader(
'%s/kitti12/training/' % dataset_folder)
loader_kitti12 = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp, _ = ls.dataloader(
'%s/eth3d/' % dataset_folder)
loader_eth3d = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.9, 2.4 * scale_factor],
order=0)
all_left_img, all_right_img, all_left_disp = lld.dataloader(
'%s/lidar_dataset/train' % dataset_folder)
loader_lidar = DA.myImageFloder(all_left_img,
all_right_img,
all_left_disp,
rand_scale=[0.5, 1.25 * scale_factor],
rand_bright=[0.8, 1.2],
order=0)
all_dataloaders = [{
'name': 'lidar',
'dl': loader_lidar,
'count': 1
}, {
'name': 'hrvs',
'dl': loader_carla,
'count': 1
}, {
'name': 'middlebury',
'dl': loader_mb,
'count': 1
}, {
'name': 'sceneflow',
'dl': loader_scene,
'count': 1
}, {
'name': 'kitti12',
'dl': loader_kitti12,
'count': 1
}, {
'name': 'kitti15',
'dl': loader_kitti15,
'count': 1
}, {
'name': 'eth3d',
'dl': loader_eth3d,
'count': 1
}]
max_count = 0
for dataloader in all_dataloaders:
max_count = max(max_count, len(dataloader['dl']))
print('=' * 80)
concat_dataloaders = []
for dataloader in all_dataloaders:
dataloader['count'] = max(1, max_count // len(dataloader['dl']))
concat_dataloaders += [dataloader['dl']] * dataloader['count']
print('{name}: {size} (x{count})'.format(name=dataloader['name'],
size=len(dataloader['dl']),
count=dataloader['count']))
data_inuse = torch.utils.data.ConcatDataset(concat_dataloaders)
print('Total dataset size: {}'.format(len(data_inuse)))
print('=' * 80)
return data_inuse