batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.workers,
drop_last=True)
# ========================================================== #
# ================CREATE NETWORK AND OPTIMIZER============== #
net = UNet(use_occ=opt.use_occ)
net.apply(kaiming_init)
weights_normal_init(net.output_layer, 0.001)
net.cuda()
optimizer = optim.Adam(net.parameters(), lr=opt.lr)
if opt.resume:
start_epoch = load_checkpoint(net, optimizer, opt.checkpoint) + 1
else:
start_epoch = 1
gamma = create_gamma_matrix(480, 640, 600, 600)
gamma = torch.from_numpy(gamma).float().cuda()
# ========================================================== #
# =================== DEFINE TRAIN ========================= #
def train(data_loader, net, optimizer):
net.train()
end = time.time()
for i, data in enumerate(data_loader):
# load data and label
'rb'))
ours = np.array(ours) * 10
ours = ours[:, eigen_crop[0]:eigen_crop[1], eigen_crop[2]:eigen_crop[3]]
return ours
# ========================================================== #
# ================CREATE NETWORK AND OPTIMIZER============== #
net = UNet(use_occ=opt.use_occ,
no_contour=opt.no_contour,
only_contour=opt.only_contour,
use_aux=(opt.use_normal or opt.use_img))
optimizer = optim.Adam(net.parameters(), lr=opt.lr)
load_checkpoint(net, optimizer, opt.checkpoint)
net.cuda()
net.eval()
# ========================================================== #
# load in occlusion list
occ_list = sorted(
[name for name in os.listdir(opt.occ_dir) if name.endswith(".npy")])
# load in normal list
normal_list = sorted([
name for name in os.listdir(opt.data_dir) if name.endswith("-normal.png")
])
# load in rgb list
img_list = sorted(