def estimate_depth(self, new_width, new_height):
original_width = self.images[0].shape[1]
original_height = self.images[0].shape[0]
factor_x = new_width / original_width
factor_y = new_height / original_height
# Resizing the image
self.resize(new_width, new_height)
# convert to pythorch format (data shape changed)
self.pytorch_format()
# model
depthnet = depthNet()
model_data = torch.load('opensource_model.pth.tar')
depthnet.load_state_dict(model_data['state_dict'])
depthnet = depthnet.cuda()
cudnn.benchmark = True
depthnet.eval()
# for warp the image to construct the cost volume
pixel_coordinate = np.indices([new_width, new_height]).astype(
np.float32) # stores row and coloumn matrix
pixel_coordinate = np.concatenate(
(pixel_coordinate, np.ones([1, new_width, new_height])), axis=0) # stores row, coloumn and third D matrix
# Convert individual matrix to a row vector
pixel_coordinate = np.reshape(pixel_coordinate, [3, -1])
KRKiUV_cuda_T, KT_cuda_T = self.camera_matrices(pixel_coordinate, factor_x, factor_y)
depth = depthnet(self.image_cuda, KRKiUV_cuda_T, KT_cuda_T, new_width,
new_height, self.reference_index, self.indices)
return depth
def main():
global tensorboard_path, tum_traindata_file, test_file, model_save_path, tum_validate_file
global epoch_num, iterate_num, best_loss_train, best_loss_validate, sun3d_traindata_root, sun3d_validate_root
trian_writer = SummaryWriter(tensorboard_path)
train_loader = torch.utils.data.DataLoader(torch.utils.data.ConcatDataset([
TumDataset(tum_traindata_file[0]),
TumDataset(tum_traindata_file[1]),
Sun3dDataset(sun3d_traindata_root)
]),
batch_size=8,
shuffle=True,
num_workers=4)
validate_loader = torch.utils.data.DataLoader(
torch.utils.data.ConcatDataset([
TumDataset(tum_validate_file[0], use_augment=False),
TumDataset(tum_validate_file[1], use_augment=False),
Sun3dDataset(sun3d_validate_root, use_augment=False),
]),
batch_size=8,
shuffle=False,
num_workers=4)
print('train data have %d pairs' % len(train_loader))
print('validate data have %d pairs' % len(validate_loader))
# model
depthnet = depthNet()
if resume_train:
pretrained_data = torch.load(resume_train_path)
depthnet.load_state_dict(pretrained_data['state_dict'])
epoch_num = pretrained_data['epoch_num']
iterate_num = pretrained_data['iterate_num']
best_loss_train = pretrained_data['best_loss_train']
best_loss_validate = pretrained_data['best_loss_validate']
print(
'we start training from epoch %d, iteration %d, the best loss in training is %f and in validation is %f'
% (epoch_num, iterate_num, best_loss_train, best_loss_validate))
if (not resume_train) and initialize_train:
pretrained_data = torch.load(initialize_train_path)
depthnet.load_state_dict(pretrained_data['state_dict'])
print('we inittialize training from epoch %d, iteration %d' %
(epoch_num, pretrained_data['iterate_num']))
depthnet = depthnet.cuda()
cudnn.benchmark = True
# optimizer
optimizer = torch.optim.Adam(depthnet.parameters(), lr=1e-3)
optimizer.zero_grad()
#start the epoch
while True:
loss_train = train_one_epoch(depthnet, train_loader, optimizer,
trian_writer)
loss_validate = vaild_one_epoch(depthnet, validate_loader,
trian_writer)
#save the checkpoint
checkpoint_name = model_save_path + 'checkpoint' + '%05d' % epoch_num + '.pth.tar'
torch.save(get_state(depthnet), checkpoint_name)
if best_loss_train < 0 or loss_train < best_loss_train:
shutil.copyfile(checkpoint_name,
model_save_path + 'best_train.pth.tar')
best_loss_train = loss_train
if best_loss_validate < 0 or loss_validate < best_loss_validate:
shutil.copyfile(checkpoint_name,
model_save_path + 'best_validate.pth.tar')
best_loss_validate = loss_validate
# update the index for next update
epoch_num = epoch_num + 1
import numpy as np
from numpy.linalg import inv
import torch
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
from torch import Tensor
from depthNet_model import depthNet
from visualize import *
with open('sample_data.pkl', 'rb') as fp:
sample_datas = pickle.load(fp)
# model
depthnet = depthNet()
model_data = torch.load('opensource_model.pth.tar')
depthnet.load_state_dict(model_data['state_dict'])
depthnet = depthnet.cuda()
cudnn.benchmark = True
depthnet.eval()
# for warp the image to construct the cost volume
pixel_coordinate = np.indices([320, 256]).astype(np.float32)
pixel_coordinate = np.concatenate((pixel_coordinate, np.ones([1, 320, 256])),
axis=0)
pixel_coordinate = np.reshape(pixel_coordinate, [3, -1])
cv2.namedWindow('result')
cv2.moveWindow('result', 200, 200)