start_time = time.time()

depth_dataset = DepthDataset(root_dir=data,
                             transform=transforms.Compose([ToTensor()]))
train_loader = torch.utils.data.DataLoader(depth_dataset, batchSize)
dataiter = iter(train_loader)
images = dataiter.next()

print("\n Time taken to load Images: %s " % (time.time() - start_time))
print("\n Test Dataset Shape: {shape}".format(shape=np.shape(depth_dataset)))

# ### Importing the Model

from Mobile_model import Model

model = Model().cuda()
model = nn.DataParallel(model)

# Import the Pre-trained Model

model.load_state_dict(torch.load(pretrained_path))
print("\n Loaded MobileNet U-Net Weights successfully\n")

model.eval()

# ### Model Variables (state_dict)

# print("\n\nModel's state_dict:\n\n")
# for param_tensor in model.state_dict():
#     print(param_tensor, "\t", model.state_dict()[param_tensor].size())
示例#2
0
        t2 = torch.from_numpy(t1).float().div(255)
        t3 = model(t2.cuda())
        t4 = t3.detach().cpu().numpy()
        t5 = t4
        # t5 = cv2.resize(t5, (320, 240))
        preds.append(t5[0][0])
    return preds


image_list = glob.glob('rgb_syn_test/*.jpg')
test_images = load_images(image_list)

image_list_gt = glob.glob('gt_syn_test/*.png')
test_images_gt = load_images(image_list_gt)
# model checkpoints
model = Model().cuda()
# model.load_state_dict(torch.load(r'models\01-27-2021_12-51-33-n17183-e20-bs4-lr0.0001\weights.epoch8_model.pth'))
# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# epoch = checkpoint['epoch']
# loss = checkpoint['loss']

optimizer = torch.optim.Adam(model.parameters(), 0.0001)
checkpoint = torch.load(
    r'models\01-27-2021_12-51-33-n17183-e20-bs4-lr0.0001\weights.epoch8_model.pth'
)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()
import datetime

import torch
import torch.nn as nn
import torch.nn.utils as utils
import torchvision.utils as vutils
from tensorboardX import SummaryWriter

writer = SummaryWriter()

# from data import getTrainingTestingData
# from utils import AverageMeter, DepthNorm, colorize

#

model = Model().cuda()
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
    model = nn.DataParallel(model)
#load trained model if needed
#model.load_state_dict(torch.load('/workspace/1.pth'))
print('Model created.')

# [['rr-1.ppm', 'd-1.jpg'], [], []]

depth_dataset = DepthDataset(traincsv=traincsv,
                             transform=transforms.Compose(
                                 [Augmentation(0.5),
                                  ToTensor()]))
train_loader = DataLoader(depth_dataset, batch_size, shuffle=True)
示例#4
0
                        help='Image size of network input')
    parser.add_argument('--data_dir',
                        default='comarision_datasets\input',
                        type=str,
                        help='Data path')
    parser.add_argument(
        '--result_dir',
        default='demo_results',
        type=str,
        help='Directory for saving results, default: demo_results')
    parser.add_argument('--gpu_id',
                        default=0,
                        type=int,
                        help='GPU id, default:0')
    args = parser.parse_args()

    if not os.path.exists(args.result_dir):
        os.makedirs(args.result_dir)

    gpu_id = args.gpu_id
    torch.cuda.device(gpu_id)

    net = Model().cuda()
    net.load_state_dict(torch.load('weights_model.pth'))
    net.eval()

    print('Begin to test ...')
    with torch.no_grad():
        demo(net, args)
    print('Finished!')