albu.MotionBlur(p=0.3),
]),
# Noise augmentation
albu.OneOf([
albu.GaussNoise(var_limit=(10, 30), p=0.5),
albu.IAAAdditiveGaussianNoise(
loc=0, scale=(0.005 * 255, 0.02 * 255), p=0.5)
]),
],
p=0.6)
if args.use_hourglass == True:
model = stackhourglass.PSMNet(int(args.maxdisp), args.use_dilation,
args.use_ssp)
else:
model = basic.PSMNet(int(args.maxdisp), args.use_dilation, args.use_ssp)
if args.cuda:
model = torch.nn.DataParallel(model, device_ids=[0, 1])
model.cuda()
if args.loadmodel is not None:
if Path(args.loadmodel).exists():
print("Loading {:s} ......".format(args.loadmodel))
state_dict = torch.load(args.loadmodel)
step = state_dict['step']
epoch = state_dict['epoch']
model.load_state_dict(state_dict['state_dict'])
print('Restored model, epoch {}, step {}'.format(epoch, step))
else:
print("No trained model detected")
parser.add_argument('--seed',
type=int,
default=666,
metavar='S',
help='random seed (default: 1)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
if args.model == 'stackhourglass':
model = stackhourglass.PSMNet(int(args.maxdisp))
elif args.model == 'basic':
model = basic.PSMNet(int(args.maxdisp))
else:
print('no model')
if args.cuda:
model = torch.nn.DataParallel(model, device_ids=[0, 1])
model.cuda()
if args.loadmodel is not None:
print('load PSMNet')
state_dict = torch.load(args.loadmodel)
step = state_dict['step']
epoch = state_dict['epoch']
model.load_state_dict(state_dict['state_dict'])
print('Restored model, epoch {}, step {}'.format(epoch, step))
from dataloader import KITTI_submission_loader as DA
else:
from dataloader import KITTI_submission_loader2012 as DA
test_left_img, test_right_img = DA.dataloader(args.datapath)
args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
if args.model == 'stackhourglass':
model = psm_net.PSMNet(args.maxdisp)
elif args.model == 'basic':
model = basic_net.PSMNet(args.maxdisp)
else:
print('no model')
model = nn.DataParallel(model, device_ids=[0])
model.cuda()
if args.loadmodel is not None:
print('load PSMNet')
state_dict = torch.load(args.loadmodel)
model.load_state_dict(state_dict['state_dict'])
print('Number of model parameters: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))