def setup(opt, checkpoint):
if checkpoint:
model = checkpoint['model'].cuda(0)
return model
elif opt.retrain:
assert os.path.exists(opt.retrain), f'Model not found: {opt.retrain}'
print(f'\n\n --> [Retrain] Loading model from: models/{opt.retrain}')
model = torch.load(opt.retrain)['model'].cuda(0)
return model
elif opt.refine:
if not opt.val_only:
print(f'\n\n --> Creating model from: models/RefineNet.py')
model = RefineNet.RefineNet(opt)
else:
print(f'\n\n --> Loading model from: {opt.refine_dir}')
model = torch.load(opt.refine_dir)['model'].cuda(0)
return model
else:
if not opt.val_only:
print(f'\n\n --> Creating model from: models/CoarseNet.py')
model = CoarseNet.CoarseNet(opt)
else:
print(f'\n\n --> Loading model from: {opt.pred_dir}')
model = torch.load(opt.pred_dir)['model'].cuda(0)
return model
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model).cuda(0)
return model
vnet = VisibNet(vinp, bn='test')
vpred = tf.logical_and(tf.greater(vnet.pred, .5), valid)
vpredf = tf.to_float(vpred) * 0. + 1.
# set up coarsenet
if prm.input_attr == 'depth':
cinp = pdepth * vpredf
elif prm.input_attr == 'depth_rgb':
cinp = tf.concat((pdepth * vpredf, prgb * vpredf / 127.5 - 1.), axis=3)
elif prm.input_attr == 'depth_sift':
cinp = tf.concat((pdepth * vpredf, psift * vpredf / 127.5 - 1.), axis=3)
elif prm.input_attr == 'depth_sift_rgb':
cinp = tf.concat((pdepth * vpredf, psift * vpredf / 127.5 - 1.,
prgb * vpredf / 127.5 - 1.),
axis=3)
cnet = CoarseNet(cinp, bn='test')
cpred = cnet.pred
# set up refinenet
rinp = tf.concat((cpred, cinp), axis=3)
rnet = RefineNet(rinp, bn='train')
rpred = rnet.pred
# scale outputs
cpred = (cpred + 1.) * 127.5
rpred = (rpred + 1.) * 127.5
################################################################################
# Run Graph
sess = tf.Session()
axis=3)
cinp_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 132]
cinp_b0 = tf.Variable(tf.zeros(cinp_sz, dtype=tf.float32))
cinp_b1 = tf.Variable(tf.zeros(cinp_sz, dtype=tf.float32))
cgt = gt_rgb
cgt_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 3]
cgt_b0 = tf.Variable(tf.zeros(cgt_sz, dtype=tf.float32))
cgt_b1 = tf.Variable(tf.zeros(cgt_sz, dtype=tf.float32))
tldr_fetchOp = [cinp_b0.assign(cinp).op, cgt_b0.assign(cgt).op]
vldr_fetchOp = [cinp_b1.assign(cinp).op, cgt_b1.assign(cgt).op]
tldr_swapOp = [cinp_b1.assign(cinp_b0).op, cgt_b1.assign(cgt_b0).op]
# Init coarsenet
C = CoarseNet(cinp_b1, bn='train', outp_act=False)
cpred = (C.pred + 1.) * 127.5
# Init perceptual network
pinp = tf.concat((cgt_b1, cpred), axis=0)
P = VGG16(pinp, stop_layer='conv3_3')
ppred = P.pred
#########################################################################
# Set optimizer
cvars = C.trainable_variables()
optC = tf.train.AdamOptimizer(prm.adam_lr, prm.adam_mom, epsilon=prm.adam_eps)
# Set C loss
cpixloss = tf.reduce_mean(tf.abs(cgt_b1 - cpred))
cinp = proj_depth * vpred
rinp_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 4]
elif prm.input_attr == 'depth_sift':
cinp = tf.concat((proj_depth * vpred, proj_sift * vpred / 127.5 - 1.),
axis=3)
rinp_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 132]
elif prm.input_attr == 'depth_rgb':
cinp = tf.concat((proj_depth * vpred, proj_rgb * vpred / 127.5 - 1.),
axis=3)
rinp_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 7]
elif prm.input_attr == 'depth_sift_rgb':
cinp = tf.concat((proj_depth * vpred, proj_sift * vpred / 127.5 - 1.,
proj_rgb * vpred / 127.5 - 1.),
axis=3)
rinp_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 135]
C = CoarseNet(cinp, bn='test', outp_act=True)
cpred = (C.pred + 1.) * 127.5
# Set up pre-fetching for RefineNet
rinp = tf.concat((cpred, cinp), axis=3)
rinp_b0 = tf.Variable(tf.zeros(rinp_sz, dtype=tf.float32))
rinp_b1 = tf.Variable(tf.zeros(rinp_sz, dtype=tf.float32))
rgt = gt_rgb
rgt_sz = [prm.batch_size, prm.crop_size, prm.crop_size, 3]
rgt_b0 = tf.Variable(tf.zeros(rgt_sz, dtype=tf.float32))
rgt_b1 = tf.Variable(tf.zeros(rgt_sz, dtype=tf.float32))
tldr_fetchOp = [rinp_b0.assign(rinp).op, rgt_b0.assign(rgt).op]
vldr_fetchOp = [rinp_b1.assign(rinp).op, rgt_b1.assign(rgt).op]
tldr_swapOp = [rinp_b1.assign(rinp_b0).op, rgt_b1.assign(rgt_b0).op]