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))