def LogProgress(model, writer, test_loader, global_step):
model.eval()
sequential = test_loader
sample_batched = next(iter(sequential))
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# print()
# if epoch == 0:
# writer.add_image('Train.1.Image', vutils.make_grid(image.data, nrow=6, normalize=True), epoch)
# if epoch == 0:
# writer.add_image('Train.2.Depth', colorize(vutils.make_grid(depth.data, nrow=6, normalize=False)), global_step)
output = model(image)
plt.imshow(output[0].detach().cpu().numpy().squeeze(), cmap='plasma')
plt.show()
# Predict
# output = model(image)
# edge = sobel_(output)
# pred_edge = sobel_(depth_n)
# Compute the loss
# l_sobel = nn.L1Loss()(edge, pred_edge)
# l_depth = nn.L1Loss()(output, depth_n)
# l_ssim = torch.clamp((1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0, 1)
# if torch.isnan(l_sobel):
# print(1)
# print(torch.isnan(l_sobel), torch.isnan(l_depth), torch.isnan(l_ssim))
# writer.add_scalar('Test/L1', l_depth.item(), global_step)
# writer.add_scalar('Test/SSIM', l_ssim.item(), global_step)
# writer.add_scalar('Test/EDGE', l_sobel.item(), global_step)
# normal = Norm(output)
# normal = (normal + 1.0) / 2.
# normal = normal[3, :, :, :]
# normal = normal.detach().cpu().numpy().astype(np.uint8)
# normal = np.transpose(normal, (1, 2, 0))
# print()
# plt.imshow(normal)
# plt.show()
output = DepthNorm(output)
# writer.add_image('Train.3.Normal', vutils.make_grid(normal.data, nrow=6, normalize=False), epoch)
# writer.add_image('Test.2.Ours', colorize(vutils.make_grid(output.data, nrow=6, normalize=False)), global_step)
# writer.add_image('Test.3.Diff', colorize(vutils.make_grid(torch.abs(output - depth).data, nrow=6, normalize=False)),
# global_step)
del image
del depth
del output
def get_sample_image(self, idx):
batch_x, batch_y = np.zeros(self.shape_rgbd), np.zeros(
self.shape_depth)
samples = [ ('void_1500/data/classroom6/image/1552696011.5523.png', 'void_1500/data/classroom6/ground_truth_orig/1552696011.5523.png'),\
('void_1500/data/mechanical_lab3/image/1552096515.0227.png', 'void_1500/data/mechanical_lab3/ground_truth_orig/1552096515.0227.png'),\
('void_1500/data/stairs4/image/1552695287.0660.png', 'void_1500/data/stairs4/ground_truth_orig/1552695287.0660.png'),\
('void_1500/data/office3/image/1552625426.4194.png', 'void_1500/data/office3/ground_truth_orig/1552625426.4194.png'),\
('void_1500/data/desktop2/image/1552625303.1627.png', 'void_1500/data/desktop2/ground_truth_orig/1552625303.1627.png'),\
('void_1500/data/plants3/image/1552695221.0711.png', 'void_1500/data/plants3/ground_truth_orig/1552695221.0711.png') ]
i = 0
index = min(idx, 5)
sample = samples[index]
x = np.clip(
np.asarray(Image.open(self.data_root + "/" + sample[0])).reshape(
480, 640, 3) / 255, 0, 1)
#iz = DepthNorm(np.clip(np.asarray(Image.open( os.path.join(self.data_root, sample[0]).replace('image', 'interp_depth') ))/256.0*100,10.0,1000.0), maxDepth=self.maxDepth)
if self.dont_interpolate:
iz = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'sparse_depth'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth, self.maxDepth),
maxDepth=self.maxDepth)
else:
iz = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'interp_depth'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth, self.maxDepth),
maxDepth=self.maxDepth)
vm = None
y = np.asarray(
np.asarray(Image.open(self.data_root + "/" + sample[1])) / 256.0)
y = np.clip(
y.reshape(480, 640, 1) * settings.DEPTH_SCALE, self.minDepth,
self.maxDepth) # fill missing pixels and convert to cm
y = DepthNorm(y, maxDepth=self.maxDepth)
if self.channels == 5:
vm = np.array(Image.open(
os.path.join(self.data_root,
sample[0]).replace('image', 'validity_map')),
dtype=np.float32)
assert (np.all(np.unique(vm) == [0, 256]))
vm[vm > 0] = 1
batch_x[i] = np.stack([x[:, :, 0], x[:, :, 1], x[:, :, 2], iz, vm],
axis=-1)
else:
batch_x[i] = np.stack([x[:, :, 0], x[:, :, 1], x[:, :, 2], iz],
axis=-1)
batch_y[i] = nyu_resize(y, 240)
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx):
batch_x, batch_y = np.zeros(self.shape_rgbd), np.zeros(
self.shape_depth)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x = np.clip(
np.asarray(Image.open(self.data_root + "/" +
sample[0])).reshape(480, 640, 3) / 255,
0, 1)
#iz = DepthNorm(np.clip(np.asarray(Image.open( os.path.join(self.data_root, sample[0]).replace('image', 'interp_depth') ))/256.0*100,10.0,1000.0), maxDepth=self.maxDepth)
if self.dont_interpolate:
iz = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'sparse_depth'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth, self.maxDepth),
maxDepth=self.maxDepth)
else:
iz = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'interp_depth'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth, self.maxDepth),
maxDepth=self.maxDepth)
vm = None
y = np.asarray(
np.asarray(Image.open(self.data_root + "/" + sample[1])) /
256.0)
y = np.clip(
y.reshape(480, 640, 1) * settings.DEPTH_SCALE, self.minDepth,
self.maxDepth) # fill missing pixels and convert to cm
y = DepthNorm(y, maxDepth=self.maxDepth)
if self.channels == 5:
vm = np.array(Image.open(
os.path.join(self.data_root,
sample[0]).replace('image', 'validity_map')),
dtype=np.float32)
assert (np.all(np.unique(vm) == [0, 256]))
vm[vm > 0] = 1
batch_x[i] = np.stack(
[x[:, :, 0], x[:, :, 1], x[:, :, 2], iz, vm], axis=-1)
else:
batch_x[i] = np.stack([x[:, :, 0], x[:, :, 1], x[:, :, 2], iz],
axis=-1)
batch_y[i] = nyu_resize(y, 240)
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def on_epoch_end(self, epoch, logs=None):
if not test_set == None:
# Samples using current model
import matplotlib.pyplot as plt
from skimage.transform import resize
plasma = plt.get_cmap('plasma')
minDepth, maxDepth = 10, 1000
train_samples = []
test_samples = []
for i in range(self.num_samples):
x_train, y_train = train_generator.__getitem__(self.train_idx[i], False)
x_test, y_test = test_generator[self.test_idx[i]]
x_train, y_train = x_train[0], np.clip(DepthNorm(y_train[0], maxDepth=1000), minDepth, maxDepth) / maxDepth
x_test, y_test = x_test[0], np.clip(DepthNorm(y_test[0], maxDepth=1000), minDepth, maxDepth) / maxDepth
h, w = y_train.shape[0], y_train.shape[1]
rgb_train = resize(x_train, (h,w), preserve_range=True, mode='reflect', anti_aliasing=True)
rgb_test = resize(x_test, (h,w), preserve_range=True, mode='reflect', anti_aliasing=True)
gt_train = plasma(y_train[:,:,0])[:,:,:3]
gt_test = plasma(y_test[:,:,0])[:,:,:3]
predict_train = plasma(predict(model, x_train, minDepth=minDepth, maxDepth=maxDepth)[0,:,:,0])[:,:,:3]
predict_test = plasma(predict(model, x_test, minDepth=minDepth, maxDepth=maxDepth)[0,:,:,0])[:,:,:3]
train_samples.append(np.vstack([rgb_train, gt_train, predict_train]))
test_samples.append(np.vstack([rgb_test, gt_test, predict_test]))
self.writer.add_summary(tf.Summary(value=[tf.Summary.Value(tag='Train', image=make_image(255 * np.hstack(train_samples)))]), epoch)
self.writer.add_summary(tf.Summary(value=[tf.Summary.Value(tag='Test', image=make_image(255 * np.hstack(test_samples)))]), epoch)
# Metrics
e = evaluate(model, test_set['rgb'], test_set['depth'], test_set['crop'], batch_size=6, verbose=True)
logs.update({'rel': e[3]})
logs.update({'rms': e[4]})
logs.update({'log10': e[5]})
if (e[4] < self.bestErr):
modelSavePath = runPath + '/model_epoch_%02d_rms_%.2f.h5' % (epoch, e[4])
print("Epoch %02d: rms improved from %f to %f, saving model to %s" % (epoch, self.bestErr, e[4], modelSavePath))
self.bestErr = e[4]
basemodel.save(modelSavePath)
super().on_epoch_end(epoch, logs)
def __getitem__(self, idx):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x = np.clip(
np.asarray(Image.open(self.data_root + "/" +
sample[0])).reshape(480, 640, 3) / 255,
0, 1)
y = np.asarray(
np.asarray(Image.open(self.data_root + "/" + sample[1])) /
256.0)
y = np.clip(
y.reshape(480, 640, 1) * settings.DEPTH_SCALE, self.minDepth,
self.maxDepth) # fill missing pixels and convert to cm
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(x, 480)
batch_y[i] = nyu_resize(y, 240)
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x = np.clip(
np.asarray(
Image.open(BytesIO(self.data[sample[0].replace(
"\r", "")]))).reshape(480, 640, 3) / 255, 0, 1)
y = np.asarray(Image.open(
BytesIO(self.data[sample[1].replace("\r", "")])),
dtype=np.float32).reshape(
480, 640, 1).copy().astype(float) / 10.0
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(x, 480)
batch_y[i] = nyu_resize(y, 240)
# DEBUG:
# self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
# exit()
return batch_x, batch_y
def LogProgress(model, writer, test_loader, epoch):
model.eval()
sequential = test_loader
sample_batched = next(iter(sequential))
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
if epoch == 0:
writer.add_image('Train.1.Image',
vutils.make_grid(image.data, nrow=6, normalize=True),
epoch)
if epoch == 0:
writer.add_image(
'Train.2.Depth',
colorize(vutils.make_grid(depth.data, nrow=6, normalize=False)),
epoch)
output = DepthNorm(model(image))
writer.add_image(
'Train.3.Ours',
colorize(vutils.make_grid(output.data, nrow=6, normalize=False)),
epoch)
writer.add_image(
'Train.3.Diff',
colorize(
vutils.make_grid(torch.abs(output - depth).data,
nrow=6,
normalize=False)), epoch)
del image
del depth
del output
def __getitem__(self, idx, is_apply_policy=True):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
# Useful for validation
if self.is_skip_policy: is_apply_policy = False
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
rgb_sample = cv2.imdecode(
np.asarray(self.data['x/{}'.format(sample)]), 1)
depth_sample = self.data['y/{}'.format(sample)]
depth_sample = resize(depth_sample,
(self.shape_depth[1], self.shape_depth[2]),
preserve_range=True,
mode='reflect',
anti_aliasing=True)
x = np.clip(rgb_sample / 255, 0, 1)
y = np.clip(depth_sample, 10, self.maxDepth)
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = x
batch_y[i] = y
if is_apply_policy:
batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i],self.maxDepth)/self.maxDepth,0,1), index, i)
return batch_x, batch_y
def __getitem__(self, idx, is_apply_policy=True):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x = np.clip(
np.asarray(Image.open(BytesIO(self.data[sample[0]]))).reshape(
480, 640, 3) / 255, 0, 1)
y = np.clip(
np.asarray(Image.open(BytesIO(self.data[sample[1]]))).reshape(
480, 640, 1) / 255 * self.maxDepth, 0, self.maxDepth)
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(x, 480)
batch_y[i] = nyu_resize(y, 240)
if is_apply_policy:
batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx):
shape_rgb = self.shape_rgb
shape_depth = self.shape_depth
batch_x = np.zeros(shape_rgb[:3] + (5 if self.locations else 3, ),
np.float32)
if self.locations:
batch_y = [] # size can vary
else:
batch_y = np.zeros(shape_depth[:3] + (1, ), np.float32)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x = np.clip(
np.asarray(Image.open(BytesIO(self.data[sample[0]]))).reshape(
480, 640, 3) / 255, 0, 1)
y = np.asarray(Image.open(BytesIO(self.data[sample[1]])),
dtype=np.float32).reshape(
480, 640, 1).copy().astype(float) / 10.0
y = DepthNorm(y, maxDepth=self.maxDepth)
# adjust sizes
if self.nus is not None:
u, v = sampling(y[:, :, 0], self.shape_rgb[1:3], self.nus)
locations = [v, u]
x = x[v.astype(int), u.astype(int), :]
if not self.locations:
u, v = sampling(y[:, :, 0], self.shape_depth[1:3],
self.nus)
y = y[v.astype(int), u.astype(int), :]
else:
shape = x.shape
locations = np.mgrid[:shape[0] - 1:shape_rgb[1] *
1j, :shape[1] - 1:shape_rgb[2] * 1j]
sample = np.round(locations).astype(int)
x = x[sample[0], sample[1], :]
if not self.locations:
sample = np.round(
np.mgrid[:shape[0] - 1:shape_depth[1] * 1j, :shape[1] -
1:shape_depth[2] * 1j]).astype(int)
y = y[sample[0], sample[1]]
batch_x[i, ..., :3] = x
if self.locations:
batch_x[i, ..., 3:] = np.stack(locations, -1)
batch_y.append(y)
else:
batch_y[i, ..., 0] = y
if self.locations:
max_x, max_y = np.max([x.shape[:2] for x in batch_y], axis=0)
batch_y = np.stack([
np.pad(x, [(0, max_x - x.shape[0]), (0, max_y - x.shape[1]),
(0, 0)], 'constant') for x in batch_y
])
return batch_x, batch_y
def __getitem__(self, idx):
batch_x, batch_sz, batch_y = np.zeros(self.shape_rgb), np.zeros(
self.shape_sz), np.zeros(self.shape_depth)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
im = np.clip(
np.asarray(Image.open(self.data_root + "/" + sample[0])) / 255,
0, 1).reshape(480, 640, 3)
#iz = np.clip(np.asarray(Image.open( os.path.join(self.data_root, sample[0]).replace('image', 'interp_depth') ))/256.0/10.0,0,1).reshape(480,640)
iz = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'interp_depth'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth,
self.maxDepth).reshape(480, 640, 1),
maxDepth=self.maxDepth)
vm = np.array(Image.open(
os.path.join(self.data_root,
sample[0]).replace('image', 'validity_map')),
dtype=np.float32).reshape(480, 640, 1)
assert (np.all(np.unique(vm) == [0, 256]))
vm[vm > 0] = 1
gt = np.asarray(
np.asarray(Image.open(self.data_root + "/" + sample[1])) /
256.0)
#y[y <= 0] = 0.0
#v = y.astype(np.float32)
#v[y > 0] = 1.0
#v[y > 10] = 0.0
#y = np.clip(interpolate_depth(y, v).reshape(480,640,1)*100, 10.0, 1000.0) # fill missing pixels and convert to cm
gt = np.clip(
gt.reshape(480, 640, 1) * settings.DEPTH_SCALE, self.minDepth,
self.maxDepth) # fill missing pixels and convert to cm
gt = DepthNorm(gt, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(im, 480)
batch_sz[i] = np.stack([iz, vm], axis=-1).reshape(480, 640, 2)
#batch_x[i] = np.stack([im[:,:,0], im[:,:,1], im[:,:,2], iz, vm], axis=-1).reshape(480,640,5)
batch_y[i] = nyu_resize(gt, 480)
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return [batch_x, batch_sz], batch_y
def __getitem__(self, index):
"""Generate one batch of data
:param index: index of the batch
:return: X and y when fitting. X only when predicting
"""
# Generate indexes of the batch
current_indexes = list(
range(index * self.batch_size, (index + 1) * self.batch_size))
img_paths_temp = self.img_paths[current_indexes]
# Generate data
X = []
y = []
for path in img_paths_temp:
_X = cv2.cvtColor(cv2.imread(self.base_path + f"/images/{path}"),
cv2.COLOR_BGR2RGB)
_y = rgb_to_depth(cv2.imread(self.base_path + f"/depth/{path}"))
_y = 1000.0 * _y
if (np.random.random() < self.augmentation_rate):
_X = augment(_X)
if (np.random.random() < 0.5) and self.augmentation_rate:
_X, _y = flip(_X, _y)
_y = np.clip(_y, self.min_depth, self.max_depth)
_y = DepthNorm(_y, maxDepth=self.max_depth)
_y = resize(_y, (_X.shape[0] // 2, _X.shape[1] // 2),
preserve_range=True,
mode='reflect',
anti_aliasing=True)
_y = _y.reshape(_y.shape[0], _y.shape[1], 1)
#_y = np.log(_y)
X.append(_X)
y.append(_y)
if self.to_fit:
return (np.array(X) /
255).astype('float32'), np.array(y).astype('float32')
else:
return np.array(X).astype('float32')
def on_epoch_end(self, epoch, logs=None):
if not test_set == None:
# Samples using current model
import matplotlib.pyplot as plt
from skimage.transform import resize
plasma = plt.get_cmap('plasma')
minDepth, maxDepth = settings.MIN_DEPTH*settings.DEPTH_SCALE, settings.MAX_DEPTH*settings.DEPTH_SCALE
train_samples = []
test_samples = []
for i in range(self.num_samples):
x_train, y_train = train_generator.__getitem__(self.train_idx[i], False)
x_test, y_test = test_generator[self.test_idx[i]]
x_train, y_train = x_train[0], np.clip(DepthNorm(y_train[0], maxDepth=maxDepth), minDepth, maxDepth) / maxDepth
x_test, y_test = x_test[0], np.clip(DepthNorm(y_test[0], maxDepth=maxDepth), minDepth, maxDepth) / maxDepth
h, w = y_train.shape[0], y_train.shape[1]
rgb_train = resize(x_train, (h,w), preserve_range=True, mode='reflect', anti_aliasing=True)
rgb_test = resize(x_test, (h,w), preserve_range=True, mode='reflect', anti_aliasing=True)
gt_train = plasma(y_train[:,:,0])[:,:,:3]
gt_test = plasma(y_test[:,:,0])[:,:,:3]
predict_train = plasma(predict(model, x_train, minDepth=minDepth, maxDepth=maxDepth)[0,:,:,0])[:,:,:3]
predict_test = plasma(predict(model, x_test, minDepth=minDepth, maxDepth=maxDepth)[0,:,:,0])[:,:,:3]
train_samples.append(np.vstack([rgb_train, gt_train, predict_train]))
test_samples.append(np.vstack([rgb_test, gt_test, predict_test]))
self.writer.add_summary(tf.Summary(value=[tf.Summary.Value(tag='Train', image=make_image(255 * np.hstack(train_samples)))]), epoch)
self.writer.add_summary(tf.Summary(value=[tf.Summary.Value(tag='Test', image=make_image(255 * np.hstack(test_samples)))]), epoch)
# Metrics
e = evaluate(model, test_set['rgb'], test_set['depth'], test_set['crop'], batch_size=6, verbose=True)
logs.update({'rel': e[3]})
logs.update({'rms': e[4]})
logs.update({'log10': e[5]})
super().on_epoch_end(epoch, logs)
def __getitem__(self, idx, is_apply_policy=True, showImage=False):
batch_x, batch_y = np.zeros( self.shape_rgb ), np.zeros( self.shape_depth )
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N-1)
# get path of the rgb and ground truth image
sample = self.dataset[index]
rgb_path = sample[0].strip()
gt_path = sample[1].strip()
x = np.array(Image.open(rgb_path)).reshape(480,640,3)
y = np.array(Image.open(gt_path)).reshape(480,640,1)
if (self.sigmaRGB > 0):
sz = int(np.ceil(6*np.ceil(self.sigmaRGB)) + 1)
x = cv2.GaussianBlur(x, (sz,sz), self.sigmaRGB)
if (self.sigmaD > 0):
sz = int(np.ceil(6*np.ceil(self.sigmaD)) + 1)
y = np.expand_dims(cv2.GaussianBlur(y, (sz,sz), self.sigmaD), 2)
if (showImage and i==0):
plt.subplot(1,2,1)
plt.imshow(x)
plt.title(self.sigmaRGB)
plt.show(block=False)
plt.subplot(1,2,2)
plt.imshow(np.squeeze(y))
plt.title(self.sigmaD)
plt.show(block=False)
plt.waitforbuttonpress()
x = createBorder(x, 8, 255)
y = createBorder(y, 8, 64)
x = np.clip(x/255,0,1)
y = np.clip(y/255*self.maxDepth,1,self.maxDepth)
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(x, 480)
batch_y[i] = nyu_resize(y, 240)
if is_apply_policy: batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def main():
parser = arg.ArgumentParser(
description="Test the model that has been trained")
parser.add_argument("--checkpoint",
"-c",
type=str,
help="path to checkpoint")
parser.add_argument("--device", "-d", type=str, default="cuda")
parser.add_argument("--data",
type=str,
default="examples/",
help="Path to dataset zip file")
args = parser.parse_args()
if len(args.checkpoint) and not os.path.isfile(args.checkpoint):
raise FileNotFoundError("{} no such file".format(args.checkpoint))
device = torch.device("cuda" if args.device == "cuda" else "cpu")
print("Using device: {}".format(device))
# Initializing the model and loading the pretrained model
model = DenseDepth(encoder_pretrained=False)
ckpt = torch.load(args.checkpoint)
model.load_state_dict(ckpt["model_state_dict"])
model = model.to(device)
print("model load from checkpoint complete ...")
# Get Test Images
img_list = glob(args.data + "*.png")
# Set model to eval mode
model.eval()
# Begin testing loop
print("Begin Test Loop ...")
for idx, img_name in enumerate(img_list):
img = load_images([img_name])
img = torch.Tensor(img).float().to(device)
print("Processing {}, Tensor Shape: {}".format(img_name, img.shape))
with torch.no_grad():
preds = DepthNorm(model(img).squeeze(0))
output = colorize(preds.data)
output = output.transpose((1, 2, 0))
cv2.imwrite(img_name.split(".")[0] + "_result.png", output)
print("Processing {} done.".format(img_name))
def __getitem__(self, idx):
batch_x, batch_y = np.zeros( self.shape_rgb ), np.zeros( self.shape_depth )
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N-1)
sample = self.dataset[index]
x = np.clip(np.asarray(Image.open( BytesIO(self.data[sample[0]]))).reshape(480,640,3)/255,0,1)
y = np.asarray(Image.open(BytesIO(self.data[sample[1]])), dtype=np.float32).reshape(480,640,1).copy().astype(float) / 10.0
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(x, 480)
batch_y[i] = nyu_resize(y, 240)
def __getitem__(self, idx):
batch_x1, batch_x2, batch_y = np.zeros(self.shape_depth), np.zeros(
self.shape_depth), np.zeros(self.shape_depth)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x1 = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'prediction'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth,
self.maxDepth).reshape(480, 640, 1),
maxDepth=self.maxDepth)
x2 = DepthNorm(np.clip(
np.asarray(
Image.open(
os.path.join(self.data_root, sample[0]).replace(
'image', 'interp_depth'))) / 256.0 *
settings.DEPTH_SCALE, self.minDepth,
self.maxDepth).reshape(480, 640, 1),
maxDepth=self.maxDepth)
y = DepthNorm(np.clip(
np.asarray(Image.open(os.path.join(self.data_root, sample[1])))
/ 256.0 * settings.DEPTH_SCALE, self.minDepth,
self.maxDepth).reshape(480, 640, 1),
maxDepth=self.maxDepth)
batch_x1[i] = nyu_resize(x1, 240)
batch_x1[i] = nyu_resize(x2, 240)
batch_y[i] = nyu_resize(y, 240)
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return [batch_x1, batch_x2], batch_y
def estimate(self, image_data):
image = Image.open(BytesIO(image_data))
if image.size != (640, 480):
image = image.resize((640, 480), Image.BICUBIC)
image = image.convert('RGB')
input_data = np.clip(np.asarray(image, dtype=float) / 255, 0, 1)
x = np.expand_dims(input_data, axis=0)
inputs = x
payload = Payload(self._estimate, (inputs,))
predictions = self.manager.send_task(payload)
outputs = np.clip(DepthNorm(predictions, maxDepth=self.max_depth), self.min_depth, self.max_depth) / self.max_depth
return np.squeeze(outputs[0])
def __getitem__(self, idx, is_apply_policy=True):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(
self.shape_depth_reduced)
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
#x = np.clip(np.asarray(Image.open( BytesIO(self.data[sample[0]]) )).reshape(self.shape_rgb[1:])/255,0,1)
x = np.clip(
cv2.resize(
np.asarray(Image.open(BytesIO(self.data[sample[0]]))),
self.shape_rgb[1:3]).reshape(self.shape_rgb[1:]) / 255, 0,
1)
#x = np.clip(np.asarray(Image.open( BytesIO(self.data[sample[0]]) ))/255,0,1)
#y = np.clip(np.asarray(Image.open( BytesIO(self.data[sample[1]]) )).reshape(self.orig_shape_depth[1:])/255*self.maxDepth,0,self.maxDepth)
#y = np.clip(np.asarray(cv2.imdecode( np.fromstring( BytesIO(self.data[sample[1]]).read() , np.uint8), 1 )).reshape(self.orig_shape_depth[1:])/255*self.maxDepth,0,self.maxDepth)
#y = np.clip(cv2.resize(np.asarray(cv2.imdecode( np.fromstring( BytesIO(self.data[sample[1]]).read() , np.uint8), 1 )), self.orig_shape_depth[1:3]).reshape(self.orig_shape_depth[1:])/255*self.maxDepth,0,self.maxDepth)
#y=np.asarray(cv2.imdecode( np.fromstring( BytesIO(self.data[sample[1]]).read() , np.uint8), 1 ))
y = cv2.resize(
np.asarray(
cv2.imdecode(
np.fromstring(
BytesIO(self.data[sample[1]]).read(), np.uint8),
-1)),
(self.orig_shape_depth[2], self.orig_shape_depth[1]))
y = np.clip(y[:, :, np.newaxis], 0, self.maxDepth)
#y = np.clip(np.asarray(cv2.imdecode( np.fromstring( BytesIO(self.data[sample[1]]).read() , np.uint8), 1 ))/255*self.maxDepth,0,self.maxDepth)
y[y == 0] = self.maxDepth
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = own_resize(x, self.shape_rgb[1], self.shape_rgb[2])
batch_y[i] = own_resize(y, self.shape_depth_reduced[1],
self.shape_depth_reduced[2])
if is_apply_policy:
batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def LogProgress(model, writer, test_loader, epoch, device):
"""To record intermediate results of training"""
model.eval()
sequential = test_loader
sample_batched = next(iter(sequential))
image = torch.Tensor(sample_batched["image"]).to(device)
depth = torch.Tensor(sample_batched["depth"]).to(device)
if epoch == 0:
writer.add_image(
"Train.1.Image",
vision_utils.make_grid(image.data, nrow=6, normalize=True),
epoch,
)
if epoch == 0:
writer.add_image(
"Train.2.Image",
colorize(
vision_utils.make_grid(depth.data, nrow=6, normalize=False)),
epoch,
)
output = DepthNorm(model(image))
writer.add_image(
"Train.3.Ours",
colorize(vision_utils.make_grid(output.data, nrow=6, normalize=False)),
epoch,
)
writer.add_image(
"Train.4.Diff",
colorize(
vision_utils.make_grid(torch.abs(output - depth).data,
nrow=6,
normalize=False)),
epoch,
)
del image
del depth
del output
def __getitem__(self, idx):
batch_x, batch_y = np.zeros( self.shape_rgb ), np.zeros( self.shape_depth )
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N-1)
sample = self.dataset[index]
x = np.clip(np.asarray(Image.open( '../datasets/' + sample[0] ).resize([384,384])).reshape(384,384,3)/255 , 0, 1)
y = np.clip(np.asarray(Image.open( '../datasets/' + sample[1]).convert(mode='L').resize([192,192])).reshape(192,192,1)*1.0 , 0, self.maxDepth)
#print(np.min(y), np.max(y))
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = iris_resize(x, 384)
batch_y[i] = iris_resize(y, 192)
#print(np.min(batch_y[i]), np.max(batch_y[i]))
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx, is_apply_policy=True):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
rgb_sample = np.asarray(Image.open(BytesIO(self.data[sample[0]])))
rgb_sample = resize(rgb_sample,
(self.shape_rgb[1], self.shape_rgb[2]),
preserve_range=True,
mode='reflect',
anti_aliasing=True)
depth_sample = np.asarray(Image.open(BytesIO(
self.data[sample[1]])))
depth_sample = resize(depth_sample,
(self.shape_depth[1], self.shape_depth[2]),
preserve_range=True,
mode='reflect',
anti_aliasing=True)
#normalise input for efficientlite models
x = np.clip((rgb_sample / 127.5) - 1.0, 0, 1)
y = np.clip(depth_sample, 10, self.maxDepth)
y = y[..., np.newaxis]
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = x
batch_y[i] = y
if is_apply_policy:
batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
for i in range(self.batch_size):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
x = np.clip(
np.asarray(Image.open(BytesIO(self.data[sample[0]]))).reshape(
320, 320, 3) / 255, 0, 1)
y = np.clip(
np.asarray(Image.open(BytesIO(self.data[sample[1]]))).reshape(
320, 320, 1) / 255 * self.maxDepth, 0, self.maxDepth)
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = eyemodel_resize(x, 320)
batch_y[i] = eyemodel_resize(y, 160)
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx, is_apply_policy=True):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
# print(sample)
# print(self.data.keys())
# import os
# print(os.getcwd())
# print(os.path.exists(prefix + sample[0]))
prefix = 'dense_depth_inputs/'
x = np.clip(
np.asarray(Image.open(BytesIO(
self.data[prefix + sample[0]]))).reshape(480, 640, 3) /
255, 0, 1)
y = np.clip(
np.asarray(
Image.open(BytesIO(
self.data[prefix + sample[1][:-1]]))).reshape(
480, 640, 1) / 255 * self.maxDepth, 0,
self.maxDepth) #
y = cv2.bitwise_not(y)
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize(x, 480)
batch_y[i] = nyu_resize(y, 240)
if is_apply_policy:
batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def __getitem__(self, idx, is_apply_policy=True):
batch_x, batch_y = np.zeros(self.shape_rgb), np.zeros(self.shape_depth)
# Augmentation of RGB images
for i in range(batch_x.shape[0]):
index = min((idx * self.batch_size) + i, self.N - 1)
sample = self.dataset[index]
basewidth = 1920
baseheight = 1056
x_img = Image.open(BytesIO(self.data[sample[0]])).resize(
(basewidth, baseheight), Image.ANTIALIAS)
x = np.clip(
np.asarray(x_img).reshape(basewidth, baseheight, 3) / 255, 0,
1)
y_img = Image.open(BytesIO(self.data[sample[1]])).resize(
(basewidth, baseheight), Image.ANTIALIAS)
y = np.clip(
np.asarray(y_img).reshape(basewidth, baseheight, 1) / 255 *
self.maxDepth, 0.1, self.maxDepth)
y = DepthNorm(y, maxDepth=self.maxDepth)
batch_x[i] = nyu_resize_x(x, baseheight)
batch_y[i] = nyu_resize_y(y, baseheight / 2)
if is_apply_policy:
batch_x[i], batch_y[i] = self.policy(batch_x[i], batch_y[i])
# DEBUG:
#self.policy.debug_img(batch_x[i], np.clip(DepthNorm(batch_y[i])/maxDepth,0,1), idx, i)
#exit()
return batch_x, batch_y
def test_model(save_dir, save_img=False, evaluate=True):
if not os.path.exists('%s/testimg' % save_dir):
os.makedirs('%s/testimg' % save_dir)
# load saved model
model = Model_rgbd().cuda()
model.load_state_dict(
torch.load(os.path.join(save_dir, 'model_1up_TXrefixed/epoch-5.pth')))
# model.load_state_dict(torch.load(os.path.join(save_dir, 'epoch-19.pth')))
model.eval()
print('model loaded for evaluation.')
# Load data
test_loader = getTestingDataOnly(batch_size=1)
train_loader_l, test_loader_l = getTranslucentData(batch_size=1)
with torch.cuda.device(0):
model.eval()
tot_len = len(
test_loader_l) # min(len(test_loader), len(test_loader_l))
testiter = iter(test_loader)
testiter_l = iter(test_loader_l)
for i in range(tot_len):
# print("Iteration "+str(i)+". loop start:")
try:
sample_batched = next(testiter)
sample_batched_l = next(testiter_l)
except StopIteration:
print(' (almost) end of iteration: %d.' % i)
break
print('/=/=/=/=/=/ iter %02d /=/=/=/=/' % i)
# (1) Pretext task : test and save
image_nyu = torch.autograd.Variable(sample_batched['image'].cuda())
depth_nyu = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
mask_raw = torch.autograd.Variable(sample_batched_l['mask'].cuda())
depth_nyu_n = DepthNorm(depth_nyu)
# # Apply random mask to it
ordered_index = list(
range(depth_nyu.shape[0])
) # NOTE: NYU test batch size shouldn't be bigger than lucent's.
mask_new = mask_raw[ordered_index, :, :, :]
depth_nyu_masked = resize2d(depth_nyu_n, (480, 640)) * mask_new
# if i <= 1:
# print('====/ %02d /====' % i)
# print(image_nyu.shape)
# print(" " + str(torch.max(image_nyu)) + " " + str(torch.min(image_nyu)))
# print(depth_nyu.shape)
# print(" " + str(torch.max(depth_nyu)) + " " + str(torch.min(depth_nyu)))
# print(mask_new.shape)
# print(" " + str(torch.max(mask_new)) + " " + str(torch.min(mask_new)))
# Predict
(htped_out_t1, _) = model(image_nyu, depth_nyu_masked)
depth_out_t1 = DepthNorm(htped_out_t1)
dn_resized = resize2d(depth_nyu, (240, 320))
if save_img:
# Save image
vutils.save_image(depth_out_t1,
'%s/testimg/1out_%02d.png' % (save_dir, i),
normalize=True,
range=(0, 1000))
if not os.path.exists('%s/testimg/1in_000000_%02d.png' %
(save_dir, i)):
vutils.save_image(depth_nyu_masked,
'%s/testimg/1in_%02d.png' %
(save_dir, i),
normalize=True,
range=(0, 1000))
save_error_image(depth_out_t1 - dn_resized,
'%s/testimg/1diff_%02d.png' % (save_dir, i),
normalize=True,
range=(-300, 300))
del image_nyu, depth_nyu, htped_out_t1, depth_out_t1, dn_resized
# (2) Main task : test and save
image = torch.autograd.Variable(sample_batched_l['image'].cuda())
depth_in = torch.autograd.Variable(
sample_batched_l['depth_raw'].cuda())
htped_in = DepthNorm(depth_in)
depth_gt = torch.autograd.Variable(
sample_batched_l['depth_truth'].cuda(non_blocking=True))
(_, htped_out_t2) = model(image, htped_in)
depth_out_t2 = DepthNorm(htped_out_t2)
mask_small = resize2dmask(mask_raw, (240, 320))
obj_mask = thresh_mask(depth_gt, resize2d(depth_in, (240, 320)))
# print(" " + str(torch.max(depth_out_t2)) + " " + str(torch.min(depth_out_t2)))
# print(" " + str(torch.max(depth_gt)) + " " + str(torch.min(depth_gt)))
# print(" " + str(torch.max(depth_in)) + " " + str(torch.min(depth_in)))
if i == 0:
(s0, s1, s2, s3) = depth_out_t2.size()
# https://stackoverflow.com/questions/22392497/how-to-add-a-new-row-to-an-empty-numpy-array
true_y = np.empty((0, s1, s2, s3), float)
raw_y = np.empty((0, s1, s2, s3), float)
pred_y = np.empty((0, s1, s2, s3), float)
mask_y = np.empty((0, s1, s2, s3), float)
objmask_y = np.empty((0, s1, s2, s3), float)
if evaluate:
true_y = np.append(true_y, depth_gt.cpu().numpy(), axis=0)
raw_y = np.append(raw_y,
resize2d(depth_in, (240, 320)).cpu().numpy(),
axis=0)
pred_y = np.append(pred_y,
depth_out_t2.detach().cpu().numpy(),
axis=0)
mask_y = np.append(mask_y, mask_small.cpu().numpy(), axis=0)
objmask_y = np.append(objmask_y,
obj_mask.cpu().numpy(),
axis=0)
# dl = depth_in.cpu().numpy()
# hl = htped_in.cpu().numpy()
# dr = resize2d(depth_in, (240, 320)).cpu().numpy()
# hr = resize2d(htped_in, (240, 320)).cpu().numpy()
# do = depth_out_t2.cpu().detach().numpy()
# gr = depth_gt.cpu().numpy()
#
# print(" Depth input (original size):" + str(np.min(dl)) + "~" + str(np.max(dl)) + " (" + str(np.mean(dl)) + ")")
# print(" Depth Normed (original size):" + str(np.min(hl)) + "~" + str(np.max(hl)) + " (" + str(np.mean(hl)) + ")")
#
# print(" Depth input (resized):" + str(np.min(dr)) + "~" + str(np.max(dr)) + " (" + str(np.mean(dr)) + ")")
# print(" Depth Normed (resized):" + str(np.min(hr)) + "~" + str(np.max(hr)) + " (" + str(np.mean(hr)) + ")")
#
# print(" Output converted to depth:" + str(np.min(do)) + "~" + str(np.max(do)) + " (" + str(np.mean(do)) + ")")
# print(" GT depth (original size):" + str(np.min(gr)) + "~" + str(np.max(gr)) + " (" + str(np.mean(gr)) + ")")
if save_img:
if not os.path.exists('%s/testimg/2truth_000000_%02d.png' %
(save_dir, i)):
vutils.save_image(depth_in,
'%s/testimg/2in_%02d.png' %
(save_dir, i),
normalize=True,
range=(0, 500))
vutils.save_image(resize2d(depth_in, (240, 320)),
'%s/testimg/2in_s_%02d.png' %
(save_dir, i),
normalize=True,
range=(0, 500))
vutils.save_image(depth_gt,
'%s/testimg/2truth_%02d.png' %
(save_dir, i),
normalize=True,
range=(0, 500))
vutils.save_image(depth_out_t2,
'%s/testimg/2out_%02d.png' % (save_dir, i),
normalize=True,
range=(0, 500))
save_error_image(resize2d(depth_out_t2, (480, 640)) - depth_in,
'%s/testimg/2corr_%02d.png' % (save_dir, i),
normalize=True,
range=(-50, 50),
mask=mask_raw)
save_error_image(depth_out_t2 - depth_gt,
'%s/testimg/2diff_%02d.png' % (save_dir, i),
normalize=True,
range=(-50, 50),
mask=mask_small)
vutils.save_image(mask_small,
'%s/testimg/2_mask_%02d.png' % (save_dir, i),
normalize=True,
range=(-0.5, 1.5))
vutils.save_image(obj_mask,
'%s/testimg/2_objmask_%02d.png' %
(save_dir, i),
normalize=True,
range=(-0.5, 1.5))
del image, htped_in, depth_in, depth_gt, depth_out_t2, mask_raw, mask_small
if evaluate:
eo = eo_r = 0
print(
'# \ta1 \ta2 \ta3 \tabsrel\trmse \tlog10 | \timprovements--> '
)
for j in range(len(true_y)):
# errors = compute_errors(true_y[j], pred_y[j], mask_y[j])
errors_object = compute_errors(true_y[j], pred_y[j],
mask_y[j] * objmask_y[j])
# errors_r = compute_errors(true_y[j], raw_y[j], mask_y[j])
errors_object_r = compute_errors(true_y[j], raw_y[j],
mask_y[j] * objmask_y[j])
eo = eo + errors_object
eo_r = eo_r + errors_object_r
print('{j:2d} | \t'
'{e[1]:.4f}\t'
'{e[2]:.4f}\t'
'{e[3]:.4f}\t'
'{e[4]:.4f}\t'
'{e[5]:.3f}\t'
'{e[6]:.4f} | \t'
'{f1[1]:+.3f}\t'
'{f1[2]:+.3f}\t'
'{f1[3]:+.3f}\t'
'{f2[4]:+.3f}\t'
'{f2[5]:+.3f}\t'
'{f2[6]:+.3f}'.format(
j=j,
e=errors_object,
f1=(1 - errors_object_r) / (1 - errors_object) - 1,
f2=errors_object_r / errors_object - 1))
eo = eo / len(true_y)
eo_r = eo_r / len(true_y)
print('\ntotal \t'
'{e[1]:.4f}\t'
'{e[2]:.4f}\t'
'{e[3]:.4f}\t'
'{e[4]:.4f}\t'
'{e[5]:.3f}\t'
'{e[6]:.4f} | \t'
'{f1[1]:+.3f}\t'
'{f1[2]:+.3f}\t'
'{f1[3]:+.3f}\t'
'{f2[4]:+.3f}\t'
'{f2[5]:+.3f}\t'
'{f2[6]:+.3f}'.format(e=eo,
f1=(1 - eo_r) / (1 - eo) - 1,
f2=eo_r / eo - 1))
def main():
# Arguments
parser = argparse.ArgumentParser(description='High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('-c', '--configFile', required=True, help='Path to config yaml file', metavar='path/to/config')
args = parser.parse_args()
CONFIG_FILE_PATH = args.configFile
with open(CONFIG_FILE_PATH) as fd:
config_yaml = oyaml.load(fd) # Returns an ordered dict. Used for printing
config = AttrDict(config_yaml)
print(colored('Config being used for training:\n{}\n\n'.format(oyaml.dump(config_yaml)), 'green'))
# Create a new directory to save logs
runs = sorted(glob.glob(os.path.join(config.train.logsDir, 'exp-*')))
prev_run_id = int(runs[-1].split('-')[-1]) if runs else 0
MODEL_LOG_DIR = os.path.join(config.train.logsDir, 'exp-{:03d}'.format(prev_run_id + 1))
CHECKPOINT_DIR = os.path.join(MODEL_LOG_DIR, 'checkpoints')
os.makedirs(CHECKPOINT_DIR)
print('Saving logs to folder: ' + colored('"{}"'.format(MODEL_LOG_DIR), 'blue'))
# Save a copy of config file in the logs
shutil.copy(CONFIG_FILE_PATH, os.path.join(MODEL_LOG_DIR, 'config.yaml'))
# Create a tensorboard object and Write config to tensorboard
writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph')
string_out = io.StringIO()
oyaml.dump(config_yaml, string_out, default_flow_style=False)
config_str = string_out.getvalue().split('\n')
string = ''
for line in config_str:
string = string + ' ' + line + '\n\r'
writer.add_text('Config', string, global_step=None)
# Create model
model = Model()
print('Model created.')
# to continue training from a checkpoint
if config.train.continueTraining:
print('Transfer Learning enabled. Model State to be loaded from a prev checkpoint...')
if not os.path.isfile(config.train.pathPrevCheckpoint):
raise ValueError('Invalid path to the given weights file for transfer learning.\
The file {} does not exist'.format(config.train.pathPrevCheckpoint))
CHECKPOINT = torch.load(config.train.pathPrevCheckpoint, map_location='cpu')
if 'model_state_dict' in CHECKPOINT:
# Newer weights file with various dicts
print(colored('Continuing training from checkpoint...Loaded data from checkpoint:', 'green'))
print('Config Used to train Checkpoint:\n', oyaml.dump(CHECKPOINT['config']), '\n')
print('From Checkpoint: Last Epoch Loss:', CHECKPOINT['epoch_loss'], '\n\n')
model.load_state_dict(CHECKPOINT['model_state_dict'])
elif 'state_dict' in CHECKPOINT:
# reading original authors checkpoints
if config.train.model != 'rednet':
# original author deeplab checkpoint
CHECKPOINT['state_dict'].pop('decoder.last_conv.8.weight')
CHECKPOINT['state_dict'].pop('decoder.last_conv.8.bias')
else:
# rednet checkpoint
# print(CHECKPOINT['state_dict'].keys())
CHECKPOINT['state_dict'].pop('final_deconv.weight')
CHECKPOINT['state_dict'].pop('final_deconv.bias')
CHECKPOINT['state_dict'].pop('out5_conv.weight')
CHECKPOINT['state_dict'].pop('out5_conv.bias')
CHECKPOINT['state_dict'].pop('out4_conv.weight')
CHECKPOINT['state_dict'].pop('out4_conv.bias')
CHECKPOINT['state_dict'].pop('out3_conv.weight')
CHECKPOINT['state_dict'].pop('out3_conv.bias')
CHECKPOINT['state_dict'].pop('out2_conv.weight')
CHECKPOINT['state_dict'].pop('out2_conv.bias')
model.load_state_dict(CHECKPOINT['state_dict'], strict=False)
else:
# Old checkpoint containing only model's state_dict()
model.load_state_dict(CHECKPOINT)
# Enable Multi-GPU training
print("Let's use", torch.cuda.device_count(), "GPUs!")
if torch.cuda.device_count() > 1:
print('Multiple GPUs being used, can\'t save model graph to Tensorboard')
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# Training parameters
optimizer = torch.optim.Adam( model.parameters(), config.train.optimAdam.learningRate )
batch_size = config.train.batchSize
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader_list = []
test_loader_list = []
for dataset in config.train.datasetsTrain:
train_data = getTrainingTestingData('rgb', 'train', dataset.images, dataset.labels)
train_loader_list.append(train_data)
for dataset in config.train.datasetsVal:
print(dataset.images)
test_data = getTrainingTestingData('rgb', 'eval', dataset.images, dataset.labels)
test_loader_list.append(test_data)
train_loader = DataLoader(torch.utils.data.ConcatDataset(train_loader_list), batch_size, num_workers=config.train.numWorkers, shuffle=True, drop_last=True, pin_memory=True)
test_loader = DataLoader(torch.utils.data.ConcatDataset(test_loader_list), batch_size, num_workers=config.train.numWorkers, shuffle=False, drop_last=True, pin_memory=True)
print(len(torch.utils.data.ConcatDataset(train_loader_list)))
print(len(train_loader))
print(len(test_loader))
# Create a tensorboard object and Write config to tensorboard
writer = SummaryWriter(MODEL_LOG_DIR, comment='create-graph')
# Loss
l1_criterion = nn.L1Loss()
total_iter_num = 0
# Start training...
for epoch in range(config.train.numEpochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Log the current Epoch Number
writer.add_scalar('data/Epoch Number', epoch, total_iter_num)
# Switch to train mode
model.train()
end = time.time()
running_loss = 0.0
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
total_iter_num += 1
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm( depth )
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp((1 - ssim(output, depth_n, val_range = 1000.0 / 10.0)) * 0.5, 0, 1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val*(N - i))))
# Log progress
niter = epoch*N+i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'
.format(epoch, i, N, batch_time=batch_time, loss=losses, eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 50 == 0:
LogProgress(model, writer, test_loader, niter)
# Log Epoch Loss
epoch_loss = running_loss / (len(train_loader))
writer.add_scalar('data/Train Epoch Loss', epoch_loss, total_iter_num)
print('\nTrain Epoch Loss: {:.4f}'.format(epoch_loss))
metrics = compute_errors(depth_n, output)
print(metrics)
for keys, values in metrics.items():
print(str(keys) + ':' + str(values))
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
# Save the model checkpoint every N epochs
if (epoch % config.train.saveModelInterval) == 0:
filename = os.path.join(CHECKPOINT_DIR, 'checkpoint-epoch-{:04d}.pth'.format(epoch))
if torch.cuda.device_count() > 1:
model_params = model.module.state_dict() # Saving nn.DataParallel model
else:
model_params = model.state_dict()
torch.save(
{
'model_state_dict': model_params,
'optimizer_state_dict': optimizer.state_dict(),
'epoch': epoch,
'total_iter_num': total_iter_num,
'epoch_loss': epoch_loss,
'config': config_yaml
}, filename)
def main():
# Arguments
parser = argparse.ArgumentParser(
description=
'High Quality Monocular Depth Estimation via Transfer Learning')
parser.add_argument('--epochs',
default=20,
type=int,
help='number of total epochs to run')
parser.add_argument('--lr',
'--learning-rate',
default=0.0001,
type=float,
help='initial learning rate')
parser.add_argument('--bs', default=4, type=int, help='batch size')
args = parser.parse_args()
# Create model
model = Model().cuda()
print('Model created.')
# Training parameters
optimizer = torch.optim.Adam(model.parameters(), args.lr)
batch_size = args.bs
prefix = 'densenet_' + str(batch_size)
# Load data
train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)
# Logging
writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
prefix, args.lr, args.epochs, args.bs),
flush_secs=30)
# Loss
l1_criterion = nn.L1Loss()
# Start training...
for epoch in range(args.epochs):
batch_time = AverageMeter()
losses = AverageMeter()
N = len(train_loader)
# Switch to train mode
model.train()
end = time.time()
for i, sample_batched in enumerate(train_loader):
optimizer.zero_grad()
# Prepare sample and target
image = torch.autograd.Variable(sample_batched['image'].cuda())
depth = torch.autograd.Variable(
sample_batched['depth'].cuda(non_blocking=True))
# Normalize depth
depth_n = DepthNorm(depth)
# Predict
output = model(image)
# Compute the loss
l_depth = l1_criterion(output, depth_n)
l_ssim = torch.clamp(
(1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
1)
loss = (1.0 * l_ssim) + (0.1 * l_depth)
# Update step
losses.update(loss.data.item(), image.size(0))
loss.backward()
optimizer.step()
# Measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
eta = str(datetime.timedelta(seconds=int(batch_time.val *
(N - i))))
# Log progress
niter = epoch * N + i
if i % 5 == 0:
# Print to console
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
'ETA {eta}\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
epoch,
i,
N,
batch_time=batch_time,
loss=losses,
eta=eta))
# Log to tensorboard
writer.add_scalar('Train/Loss', losses.val, niter)
if i % 300 == 0:
LogProgress(model, writer, test_loader, niter)
# Record epoch's intermediate results
LogProgress(model, writer, test_loader, niter)
writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
def main():
# CLI arguments
parser = arg.ArgumentParser(
description="Training method for\t"
"High Quality Monocular Depth Estimation via Transfer Learning")
parser.add_argument(
"--epochs",
"-e",
default=20,
type=int,
help="total number of epochs to run training for",
)
parser.add_argument("--lr",
"-l",
default=0.0001,
type=float,
help="initial learning rate")
parser.add_argument("--batch",
"-b",
default=8,
type=int,
help="Batch size")
parser.add_argument(
"--checkpoint",
"-c",
default="",
type=str,
help="path to last saved checkpoint to resume training from",
)
parser.add_argument(
"--resume_epoch",
"-r",
default=-1,
type=int,
help="epoch to resume training from",
)
parser.add_argument(
"--device",
"-d",
default="cuda",
type=str,
help="device to run training on. Use CUDA",
)
parser.add_argument("--enc_pretrain",
"-p",
default=True,
type=bool,
help="Use pretrained encoder")
parser.add_argument("--data",
default="data/nyu_data.zip",
type=str,
help="path to dataset")
parser.add_argument("--theta",
"-t",
default=0.1,
type=float,
help="coeff for L1 (depth) Loss")
parser.add_argument("--save",
"-s",
default="",
type=str,
help="location to save checkpoints in")
args = parser.parse_args()
# Some sanity checks
if len(args.save) > 0 and not args.save.endswith("/"):
raise ValueError(
"save location should be path to directory or empty. (Must end with /"
)
if len(args.save) > 0 and not os.path.isdir(args.save):
raise NotADirectoryError("{} not a dir path".format(args.save))
# Load data
print("Loading Data ...")
trainloader, testloader = getTrainingTestingData(args.data,
batch_size=args.batch)
print("Dataloaders ready ...")
num_trainloader = len(trainloader)
num_testloader = len(testloader)
# Training utils
batch_size = args.batch
model_prefix = "densedepth_"
device = torch.device("cuda:0" if args.device == "cuda" else "cpu")
theta = args.theta
save_count = 0
epoch_loss = []
batch_loss = []
sum_loss = 0
# loading from checkpoint if provided
if len(args.checkpoint) > 0:
print("Loading from checkpoint ...")
model, optimizer, start_epoch = init_or_load_model(
depthmodel=DenseDepth,
enc_pretrain=args.enc_pretrain,
epochs=args.epochs,
lr=args.lr,
ckpt=args.checkpoint,
device=device,
)
print("Resuming from: epoch #{}".format(start_epoch))
else:
print("Initializing fresh model ...")
model, optimizer, start_epoch = init_or_load_model(
depthmodel=DenseDepth,
enc_pretrain=args.enc_pretrain,
epochs=args.epochs,
lr=args.lr,
ckpt=None,
device=device,
)
# Logging
writer = SummaryWriter(
comment="{}-learning_rate:{}-epoch:{}-batch_size:{}".format(
model_prefix, args.lr, args.epochs, args.batch))
# Loss functions
l1_criterion = nn.L1Loss()
# Starting training
print("Device: ", device)
print("Starting training ... ")
for epoch in range(start_epoch, args.epochs):
model.train()
model = model.to(device)
batch_time = AverageMeter()
loss_meter = AverageMeter()
epoch_start = time.time()
end = time.time()
for idx, batch in enumerate(trainloader):
optimizer.zero_grad()
image_x = torch.Tensor(batch["image"]).to(device)
depth_y = torch.Tensor(batch["depth"]).to(device=device)
normalized_depth_y = DepthNorm(depth_y)
preds = model(image_x)
# calculating the losses
l1_loss = l1_criterion(preds, normalized_depth_y)
ssim_loss = torch.clamp(
(1 - ssim_criterion(preds, normalized_depth_y, 1000.0 / 10.0))
* 0.5,
min=0,
max=1,
)
gradient_loss = gradient_criterion(normalized_depth_y,
preds,
device=device)
net_loss = ((1.0 * ssim_loss) + (1.0 * torch.mean(gradient_loss)) +
(theta * torch.mean(l1_loss)))
loss_meter.update(net_loss.data.item(), image_x.size(0))
net_loss.backward()
optimizer.step()
# Time metrics
batch_time.update(time.time() - end)
end = time.time()
eta = str(
datetime.timedelta(seconds=int(batch_time.val *
(num_trainloader - idx))))
# Logging
num_iters = epoch * num_trainloader + idx
if idx % 5 == 0:
print(
"Epoch: #{0} Batch: {1}/{2}\t"
"Time (current/total) {batch_time.val:.3f}/{batch_time.sum:.3f}\t"
"eta {eta}\t"
"LOSS (current/average) {loss.val:.4f}/{loss.avg:.4f}\t".
format(
epoch,
idx,
num_trainloader,
batch_time=batch_time,
eta=eta,
loss=loss_meter,
))
writer.add_scalar("Train/Loss", loss_meter.val, num_iters)
# if idx % 300 == 0:
# LogProgress(model, writer, testloader, num_iters, device)
# print(torch.cuda.memory_allocated()/1e+9)
del image_x
del depth_y
del preds
# print(torch.cuda.memory_allocated()/1e+9)
if epoch % 1 == 0:
print("----------------------------------\n"
"Epoch: #{0}, Avg. Net Loss: {avg_loss:.4f}\n"
"----------------------------------".format(
epoch, avg_loss=loss_meter.avg))
torch.save(
{
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optim_state_dict": optimizer.state_dict(),
"loss": loss_meter.avg,
},
args.save +
"ckpt_{}_{}.pth".format(epoch, int(loss_meter.avg * 100)),
)
# model = model.to(device)
LogProgress(model, writer, testloader, num_iters, device)
if epoch % 5 == 0:
torch.save(
{
"epoch": epoch,
"model_state_dict": model.cpu().state_dict(),
"optim_state_dict": optimizer.state_dict(),
"loss": loss_meter.avg,
},
args.save +
"ckpt_{}_{}.pth".format(epoch, int(loss_meter.avg * 100)),
)
def LogProgress(model, writer, test_loader, test_loader_l, epoch, n, save_dir, HEIGHT, WIDTH, print_task1_rmse=True):
with torch.no_grad():
with torch.cuda.device(0):
torch.cuda.empty_cache()
model.eval()
tot_len = len(test_loader_l) # min(len(test_loader), len(test_loader_l))
testiter = iter(test_loader)
testiter_l = iter(test_loader_l)
task1_rmse = 0
mse_criterion = nn.MSELoss()
for i in range(tot_len):
# print(">>>i:" + str(i))
# print("Iteration "+str(i)+". loop start:")
try:
sample_batched = next(testiter)
sample_batched_l = next(testiter_l)
except StopIteration:
print(' (almost) end of iteration.')
continue
# (1) Pretext task : test and save
image_nyu = torch.autograd.Variable(sample_batched['image'].cuda())
depth_nyu = torch.autograd.Variable(sample_batched['depth'].cuda(non_blocking=True))
# print(" " + str(torch.max(depth_nyu)) + " " + str(torch.min(depth_nyu)))
mask_raw = torch.autograd.Variable(sample_batched_l['mask'].cuda())
depth_nyu_n = DepthNorm(depth_nyu)
# Apply random mask to it
ordered_index = list(range(depth_nyu.shape[0])) # NOTE: NYU test batch size shouldn't be bigger than lucent's.
mask_new = mask_raw[ordered_index, :, :, :]
depth_nyu_masked = resize2d(depth_nyu_n, (HEIGHT, WIDTH)) * mask_new
# Predict
(depth_out_t1n, _) = model(image_nyu, depth_nyu_masked)
depth_out_t1 = DepthNorm(depth_out_t1n)
dn_resized = resize2d(depth_nyu, (int(HEIGHT/2), int(WIDTH/2)))
# Save image
vutils.save_image(depth_out_t1, '%s/img/1out_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(0, 1000))
if not os.path.exists('%s/img/1in_000000_%02d.png' % (save_dir, i)):
vutils.save_image(depth_nyu_masked, '%s/img/1in_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(0, 1000))
save_error_image(depth_out_t1 - dn_resized, '%s/img/1diff_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(-300, 300))
if print_task1_rmse:
task1_rmse = task1_rmse + torch.sqrt(mse_criterion(dn_resized, depth_out_t1))
del image_nyu, depth_nyu, depth_out_t1n, depth_out_t1, dn_resized
torch.cuda.empty_cache()
# (2) Main task : test and save
image = torch.autograd.Variable(sample_batched_l['image'].cuda())
depth_in = torch.autograd.Variable(sample_batched_l['depth_raw'].cuda())
htped_in = DepthNorm(depth_in)
depth_gt = torch.autograd.Variable(sample_batched_l['depth_truth'].cuda(non_blocking=True))
# print('====//====')
# print(image.shape)
# print(" " + str(torch.max(image)) + " " + str(torch.min(image)))
# print(depth_in.shape)
# print(" " + str(torch.max(depth_in)) + " " + str(torch.min(depth_in)))
# print(depth.shape)
# print(" " + str(torch.max(depth)) + " " + str(torch.min(depth)))
if epoch == 0: writer.add_image('Train.1.Image', vutils.make_grid(image.data, nrow=6, normalize=True), epoch)
if epoch == 0: writer.add_image('Train.2.Depth', colorize(vutils.make_grid(depth_gt.data, nrow=6, normalize=False)), epoch)
(_, depth_out_t2n) = model(image, htped_in)
depth_out_t2 = DepthNorm(depth_out_t2n)
writer.add_image('Train.3.Ours', colorize(vutils.make_grid(depth_out_t2.data, nrow=6, normalize=False)), epoch)
writer.add_image('Train.3.Diff', colorize(vutils.make_grid(torch.abs(depth_out_t2-depth_gt).data, nrow=6, normalize=False)), epoch)
# dl = depth_in.cpu().numpy()
# hl = htped_in.cpu().numpy()
# dr = resize2d(depth_in, (int(HEIGHT/2), int(WIDTH/2))).cpu().numpy()
# hr = resize2d(htped_in, (int(HEIGHT/2), int(WIDTH/2))).cpu().numpy()
# do = depth_out_t2.cpu().detach().numpy()
# gr = depth_gt.cpu().numpy()
# print('/=/=/=/=/=/')
# print(" Depth input (original size):" + str(np.min(dl)) + "~" + str(np.max(dl)) + " (" + str(np.mean(dl)) + ")")
# print(" Depth Normed (original size):" + str(np.min(hl)) + "~" + str(np.max(hl)) + " (" + str(np.mean(hl)) + ")")
#
# print(" Depth input (resized):" + str(np.min(dr)) + "~" + str(np.max(dr)) + " (" + str(np.mean(dr)) + ")")
# print(" Depth Normed (resized):" + str(np.min(hr)) + "~" + str(np.max(hr)) + " (" + str(np.mean(hr)) + ")")
#
# print(" Output converted to depth:" + str(np.min(do)) + "~" + str(np.max(do)) + " (" + str(np.mean(do)) + ")")
# print(" GT depth (original size):" + str(np.min(gr)) + "~" + str(np.max(gr)) + " (" + str(np.mean(gr)) + ")")
if not os.path.exists('%s/img/2truth_000000_%02d.png' % (save_dir, i)):
vutils.save_image(depth_in, '%s/img/2inDS_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(0, 500))
vutils.save_image(DepthNorm(htped_in), '%s/img/2inFF_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(0, 500))
vutils.save_image(depth_gt, '%s/img/2truth_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(0, 500))
vutils.save_image(depth_out_t2, '%s/img/2out_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(0, 500))
mask_small = resize2d(mask_raw, (int(HEIGHT/2), int(WIDTH/2)))
save_error_image(resize2d(depth_out_t2, (HEIGHT, WIDTH)) - depth_in, '%s/img/2corr_%06d_%02d.png'
% (save_dir, n, i), normalize=True, range=(-50, 50), mask=mask_raw)
save_error_image(depth_out_t2 - depth_gt, '%s/img/2diff_%06d_%02d.png' % (save_dir, n, i), normalize=True, range=(-50, 50), mask=mask_small)
del image, htped_in, depth_in, depth_gt, depth_out_t2n, depth_out_t2, mask_raw, mask_small
torch.cuda.empty_cache()
if print_task1_rmse:
mse_criterion = nn.MSELoss()
print('> Test RMSE for Task1 = {rmse:.2f}'
.format(rmse=task1_rmse/tot_len))