def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
data_list = []
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
images, label, size, name = batch
images = Variable(images, volatile=True)
h, w, c = size[0].numpy()
images075 = nn.Upsample(size=(int(h * 0.75), int(w * 0.75)),
mode='bilinear')(images)
images05 = nn.Upsample(size=(int(h * 0.5), int(w * 0.5)),
mode='bilinear')(images)
out100 = model(images.cuda(args.gpu))
out075 = model(images075.cuda(args.gpu))
out05 = model(images05.cuda(args.gpu))
o_h, o_w = out100.size()[2:]
interpo1 = nn.Upsample(size=(o_h, o_w), mode='bilinear')
out_max = torch.max(torch.stack(
[out100, interpo1(out075),
interpo1(out05)]),
dim=0)[0]
output = interp(out_max).cpu().data[0].numpy()
output = output[:, :h, :w]
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
gt = np.asarray(label[0].numpy()[:h, :w], dtype=np.int)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
get_iou(data_list, args.num_classes)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
if torch.cuda.is_available():
args.device = torch.device('cuda')
model = model.cuda()
print()
else:
args.device = torch.device('cpu')
print('the model is operating on {}'.format(args.device))
testloader = data.DataLoader(VOCDataSet(args.data_dir,
args.data_list,
hue_value=args.hue_value,
crop_size=(505, 505),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
data_list = []
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
image, label, size, name = batch
image = image.to(device=args.device)
size = size[0].numpy()
output = model(Variable(image, volatile=True))
output = interp(output).cpu().data[0].numpy()
output = output[:, :size[0], :size[1]]
gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int)
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
get_iou(data_list, args.num_classes, args.hue_value)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
valid_loader = data.DataLoader(BerkeleyDataSet(args.data_dir, args.data_list, mean=IMG_MEAN, scale=False, mirror=False, train=False),
batch_size=1, shuffle=False, pin_memory=True)
test_loader = data.DataLoader(BerkeleyDataTestSet(args.data_dir, args.data_list, mean=IMG_MEAN))
interp = nn.Upsample(size=(321, 321), mode='bilinear', align_corners=True)
data_list = []
# #Evaluation loop for Valid Loader
with torch.no_grad():
for index, batch in enumerate(valid_loader):
if index % 100 == 0:
print('%d processd'%(index))
image, label, name, size = batch
h, w, c = size[0].numpy()
# print(name)
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
# print(output.shape)
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# print(output)
# print(label[0].shape)
ground_truth = np.asarray(label[0].numpy()[:h,:w], dtype=np.int)
show_all(ground_truth, output)
data_list.append([ground_truth.flatten(), output.flatten()])
get_iou(data_list, args.num_classes)
#Evaluation loop for Test Loader
'''
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
data_list = []
with open('result.txt', 'w') as f:
for index, batch in enumerate(testloader, 0):
#print('%d processd'%(index))
image, label, size, name = batch
size = size[0].numpy() #the size of original input image
#print("size:",size)
output = model(Variable(image, volatile=True).cuda(gpu0))
#print("model output size: ",output.size())
output = interp(output).cpu().data[0].numpy()
output = output[:, :size[0], :size[1]]
gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int)
output = output.transpose(1, 2, 0)
prob = softmax(output, axis=2)
entropy = (-prob * np.log(prob)).mean() * 100
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
iou = get_iou(data_list, args.num_classes)
data_list = []
print('{:.4f}, {:.4f}'.format(entropy.item(), iou.item()))
f.write('{:.8f}, {:.8f}\n'.format(entropy.item(), iou.item()))
def iou():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear', align_corners=True)
data_list = []
for index, batch in enumerate(testloader):
if index % 100 == 0:
#print('%d processd'%(index))
sys.stdout.flush()
image, label, size, name = batch
size = size[0].numpy()
with torch.no_grad():
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output[:, :size[0], :size[1]]
gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int)
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
return get_iou(data_list, args.num_classes)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
dataset = GenericDataset(DATA_DIRECTORY, 'val', test_transform)
testloader = data.DataLoader(dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=4,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
data_list = []
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
image, label, size, name = batch
size = size[0].numpy()
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output[:, :size[0], :size[1]]
gt = np.asarray(label[0].numpy()[:size[0], :size[1]], dtype=np.int)
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# show_all(gt, output)
data_list.append([gt.flatten(), output.flatten()])
get_iou(data_list, args.num_classes)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
sigmoid = nn.Sigmoid()
data_list = []
for index, batch in enumerate(testloader):
if index % 1 == 0:
print('%d processd' % (index))
image, size, name = batch
size = size[0].numpy()
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output)
output = output[:, :size[0], :size[1]]
output = torch.from_numpy(output).float()
result = torchvision.transforms.ToPILImage()(output)
path = os.path.join(SR_dir, name[0])
result.save(path)
def main():
"""Create the model and start the evaluation process."""
gpu0 = 0
model = Res_Deeplab(num_classes=NUM_CLASSES)
saved_state_dict = torch.load(RESTORE_FROM)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(VOCDataSet(DATA_DIRECTORY, DATA_LIST_PATH, crop_size=(321, 321), mean=IMG_MEAN, scale=False, mirror=False),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(321, 321), mode='bilinear', align_corners=True) #changed to model 321,321
data_list = []
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd'%(index))
torch.save(batch, SAVE_TO + '/batch' + str(index) + '.pth') # Save the batch
image, label, size, name = batch
size = size[0].numpy()
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output)
torch.save(output, SAVE_TO + '/prediction' + str(index) + '.pth') #Save b11 prediction
output = output.cpu().data[0].numpy()
output = output[:,:size[0],:size[1]]
gt = np.asarray(label[0].numpy()[:size[0],:size[1]], dtype=np.int)
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
data_list.append([gt.flatten(), output.flatten()])
get_iou(data_list, NUM_CLASSES)
def main():
global opt, model, netContent
opt = parser.parse_args()
print(opt)
gpuid = 0
cuda = True
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
if opt.vgg_loss:
print('===> Loading VGG model')
netVGG = models.vgg19()
netVGG.load_state_dict(
model_zoo.load_url(
'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'))
class _content_model(nn.Module):
def __init__(self):
super(_content_model, self).__init__()
self.feature = nn.Sequential(
*list(netVGG.features.children())[:-1])
def forward(self, x):
out = self.feature(x)
return out
netContent = _content_model()
print("===> Building model")
#deeplab
deeplab_res = Res_Deeplab(num_classes=21)
saved_state_dict = torch.load('model/VOC12_scenes_20000.pth')
deeplab_res.load_state_dict(saved_state_dict)
deeplab_res = deeplab_res.eval()
#deeplab_res.eval()
mid = mid_layer()
#SRResNet
print("===> Building model")
model = Net()
finetune = True
if finetune == True:
model_pretrained = torch.load(
'model/model_DIV2K_noBN_96_epoch_36.pth',
map_location=lambda storage, loc: storage)["model"]
index = 0
for (src, dst) in zip(model_pretrained.parameters(),
model.parameters()):
if index != 0:
list(model.parameters())[index].data = src.data
index = index + 1
criterion = nn.MSELoss(size_average=False)
print("===> Setting GPU")
if cuda:
model = model.cuda(gpuid)
mid = mid.cuda(gpuid)
deeplab_res = deeplab_res.cuda(gpuid)
criterion = criterion.cuda(gpuid)
if opt.vgg_loss:
netContent = netContent.cuda(gpuid)
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model"].state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training1")
#root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5'
root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train1_HDF5'
files_num = len(os.listdir(root_dir))
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
#save_checkpoint(model, epoch)
print("===> Loading datasets")
x = random.sample(os.listdir(root_dir), files_num)
for index in range(0, files_num):
train_path = os.path.join(root_dir, x[index])
print("===> Training datasets: '{}'".format(train_path))
train_set = DatasetFromHdf5_OLD(train_path)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
avgloss = train(training_data_loader, optimizer, deeplab_res,
model, mid, criterion, epoch, gpuid)
if epoch % 2 == 0:
save_checkpoint(model, epoch)
SR_dir = join(root_dir, 'SS_Output/SR')
if os.path.isdir(SR_dir):
pass
else:
os.mkdir(SR_dir)
HR_dir = join(root_dir, 'SS_Output/HR')
if os.path.isdir(HR_dir):
pass
else:
os.mkdir(HR_dir)
deeplab_res = Res_Deeplab(num_classes=21)
saved_state_dict = torch.load('model/VOC12_scenes_20000.pth')
deeplab_res.load_state_dict(saved_state_dict)
deeplab_res = deeplab_res.eval()
mid = mid_layer()
criterion = CrossEntropy_Probability()
criterion1 = CrossEntropy2d()
criterion_kl = KL_Loss()
mid = mid.cuda(gpuid)
deeplab_res = deeplab_res.cuda(gpuid)
criterion = criterion.cuda(gpuid)
criterion1 = criterion1.cuda(gpuid)
criterion_kl = criterion_kl.cuda(gpuid)
testloader = data.DataLoader(VOCDataSSSet(root_dir,
DATA_LIST_PATH,
crop_size=(321, 321),
mean=IMG_MEAN,
def main():
"""Create the model and start the training."""
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
torch.manual_seed(args.random_seed)
model = Res_Deeplab(num_classes=args.num_classes)
model = torch.nn.DataParallel(model)
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model.module),
'lr': args.learning_rate
}, {
'params': get_10x_lr_params(model.module),
'lr': 10 * args.learning_rate
}],
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.fine_tune:
# fine tune from coco dataset
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
if i_parts[1] != 'layer5':
new_params[i.replace('Scale', 'module')] = saved_state_dict[i]
model.load_state_dict(new_params)
elif args.restore_from:
if os.path.isfile(args.restore_from):
print("=> loading checkpoint '{}'".format(args.restore_from))
checkpoint = torch.load(args.restore_from)
try:
if args.set_start:
args.start_step = int(
math.ceil(checkpoint['example'] / args.batch_size))
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (step {})".format(
args.restore_from, args.start_step))
except:
model.load_state_dict(checkpoint)
print("=> loaded checkpoint '{}'".format(args.restore_from))
else:
print("=> no checkpoint found at '{}'".format(args.restore_from))
if not args.is_training:
# Frozen BN
# when training, the model will use the running means and the
# running vars of the pretrained model.
# But note that eval() doesn't turn off history tracking.
print(
"Freezing BN layers, that is taking BN as linear transform layer")
model.eval()
else:
model.train()
model.cuda()
cudnn.benchmark = True
trainloader = data.DataLoader(XiangyaTrain(args.data_list,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
color_jitter=args.random_jitter,
rotate=args.random_rotate,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
pixel_losses = AverageMeter()
patch_losses = AverageMeter()
accuracy = AverageMeter()
writer = SummaryWriter(args.snapshot_dir)
cnt = 0
actual_step = args.start_step
while actual_step < args.final_step:
iter_end = timeit.default_timer()
for i_iter, (images, labels, patch_name) in enumerate(trainloader):
actual_step = int(args.start_step + cnt)
data_time.update(timeit.default_timer() - iter_end)
images = Variable(images).cuda()
optimizer.zero_grad()
adjust_learning_rate(optimizer, actual_step)
# pred = interp(model(images))
pred = model(images)
image = images.data.cpu().numpy()[0]
del images
# 0 Normal 1 DG 2 JR
labels = resize_target(labels, pred.size(2))
pixel_loss, patch_loss = loss_calc(pred, labels)
loss = pixel_loss.double() + args.loss_coeff * patch_loss
losses.update(loss.item(), pred.size(0))
pixel_losses.update(pixel_loss.item(), pred.size(0))
patch_losses.update(patch_loss.item(), pred.size(0))
acc = _pixel_accuracy(pred.data.cpu().numpy(),
labels.data.cpu().numpy())
accuracy.update(acc, pred.size(0))
loss.backward()
optimizer.step()
batch_time.update(timeit.default_timer() - iter_end)
iter_end = timeit.default_timer()
if actual_step % args.print_freq == 0:
print(
'iter: [{0}]{1}/{2}\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Pixel Loss {pixel_loss.val:.4f} ({pixel_loss.avg:.4f})\t'
'Patch_Loss {patch_loss.val:.4f} ({patch_loss.avg:.4f})\t'
'Pixel Accuracy {accuracy.val:.3f} ({accuracy.avg:.3f})'.
format(cnt,
actual_step,
args.final_step,
batch_time=batch_time,
data_time=data_time,
loss=losses,
pixel_loss=pixel_losses,
patch_loss=patch_losses,
accuracy=accuracy))
writer.add_scalar("train_loss", losses.avg, actual_step)
writer.add_scalar("pixel_loss", pixel_losses.avg, actual_step)
writer.add_scalar("patch_loss", patch_losses.avg, actual_step)
writer.add_scalar("pixel_accuracy", accuracy.avg, actual_step)
writer.add_scalar("lr", optimizer.param_groups[0]['lr'],
actual_step)
# TODO complete this part using writer
if actual_step % args.save_img_freq == 0:
msk_size = pred.size(2)
image = image.transpose(1, 2, 0)
image = cv2.resize(image, (msk_size, msk_size),
interpolation=cv2.INTER_NEAREST)
image += IMG_MEAN
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
label = labels.data.cpu().numpy()[0]
label = np.repeat(254, msk_size) - label * 127
single_pred = pred.data.cpu().numpy()[0].argmax(axis=0)
single_pred = single_pred * 127
new_im = Image.new('RGB', (msk_size * 3, msk_size))
new_im.paste(Image.fromarray(image.astype('uint8'), 'RGB'),
(0, 0))
new_im.paste(Image.fromarray(single_pred.astype('uint8'), 'L'),
(msk_size, 0))
new_im.paste(Image.fromarray(label.astype('uint8'), 'L'),
(msk_size * 2, 0))
new_im_name = 'B' + str(args.batch_size) + '_S' + str(
actual_step) + '_' + patch_name[0]
new_im_file = os.path.join(args.img_dir, new_im_name)
new_im.save(new_im_file)
if actual_step % args.save_pred_every == 0 and cnt != 0:
print('taking snapshot ...')
torch.save(
{
'example': actual_step * args.batch_size,
'state_dict': model.state_dict()
},
osp.join(
args.snapshot_dir,
'Xiangya_Deeplab_B' + str(args.batch_size) + '_S' +
str(actual_step) + '.pth'))
cnt += 1
print('save the final model ...')
torch.save(
{
'example': actual_step * args.batch_size,
'state_dict': model.state_dict()
},
osp.join(
args.snapshot_dir, 'Xiangya_Deeplab_B' + str(args.batch_size) +
'_S' + str(actual_step) + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
print(
"SRResNet with Semantic_KL_loss training from scratch on VOC 160*160 patches."
) #改参数之前改一下
global opt, model, netContent, deeplab_res, mid, semantic_criterion, semantic_kl_criterion, KL_DivLoss
opt = parser.parse_args()
print(opt)
gpuid = 0
cuda = True
opt.vgg_loss = False
opt.semantic_loss = True
KL_DivLoss = True
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
if opt.vgg_loss:
print('===> Loading VGG model')
netVGG = models.vgg19()
netVGG.load_state_dict(
model_zoo.load_url(
'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth')
) #建立vgg网络 从哪取loss照这个vgg loss写
class _content_model(nn.Module):
def __init__(self):
super(_content_model, self).__init__()
self.feature = nn.Sequential(
*list(netVGG.features.children())[:-1])
def forward(self, x):
out = self.feature(x)
return out
netContent = _content_model()
if opt.semantic_loss:
print('===> Loading DeepLab model')
deeplab_res = Res_Deeplab(num_classes=21)
saved_state_dict = torch.load('model/VOC12_scenes_20000.pth')
deeplab_res.load_state_dict(saved_state_dict)
deeplab_res = deeplab_res.eval()
semantic_criterion = CrossEntropy_Probability()
semantic_kl_criterion = nn.KLDivLoss(size_average=False)
print("===> Building model")
model = Net()
criterion = nn.MSELoss(size_average=False)
print("===> Setting GPU")
if cuda:
model = model.cuda(gpuid)
criterion = criterion.cuda(gpuid)
if opt.vgg_loss:
netContent = netContent.cuda(gpuid)
if opt.semantic_loss:
deeplab_res = deeplab_res.cuda(gpuid)
semantic_criterion = semantic_criterion.cuda(gpuid)
semantic_kl_criterion = semantic_kl_criterion.cuda(gpuid)
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model"].state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training1")
#root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5'
root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train_hrlabel160_HDF5'
files_num = len(os.listdir(root_dir))
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
#save_checkpoint(model, epoch)
print("===> Loading datasets")
x = random.sample(os.listdir(root_dir), files_num)
for index in range(0, files_num):
train_path = os.path.join(root_dir, x[index])
print("===> Training datasets: '{}'".format(train_path))
train_set = DatasetFromHdf5(train_path) #看
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
avgloss = train(training_data_loader, optimizer, model, criterion,
epoch, gpuid)
if epoch % 2 == 0:
save_checkpoint(model, epoch)
gpu0 = 0
os.environ["CUDA_VISIBLE_DEVICES"]=str(gpu0)
cudnn.enabled = True
# Initilize model
model = Res_Deeplab(num_classes=21)
# We use the path to the state dictionary on our own server as a default, can be changed with an argument
pathToTrainedModel = '/root/20000StepsDefaultParametersBatch6/VOC12_scenes_20000.pth'
if args.PathToPretrainedModel:
pathToTrainedModel = PathToPretrainedModel
saved_state_dict = torch.load(pathToTrainedModel)
model.load_state_dict(saved_state_dict)
#prepare for evaluation and use of cuda (GPU)
model.eval()
model.cuda()
image = Variable(image).cuda() #gets and saves a gpu output
interp = nn.Upsample(size=inter_size, mode='bilinear', align_corners=True)
# puts the tensor through the model and returns the prediction (the segmentation)
pred = interp(model(image))
# Converts the prediciton to a image and saves it under the atop defined name (given with the arguments)
def saveImage(pred):
output = pred.cpu().data[0].numpy()
output = output[:,:size[0],:size[1]]
output = output.transpose(1,2,0)
def main():
"""Create the model and start the training."""
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
logger = util.set_logger(args.snapshot_dir, args.log_file, args.debug)
logger.info('start with arguments %s', args)
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
lscale, hscale = map(float, args.train_scale.split(','))
train_scale = (lscale, hscale)
cudnn.enabled = True
# Create network.
model = Res_Deeplab(num_classes=args.num_classes)
#saved_state_dict = torch.load(args.restore_from)
#new_params = model.state_dict().copy()
#for i in saved_state_dict:
# #Scale.layer5.conv2d_list.3.weight
# i_parts = i.split('.')
# # print i_parts
# if not args.num_classes == 21 or not i_parts[1]=='layer5':
# new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
model_urls = {
'resnet18':
'https://download.pytorch.org/models/resnet18-5c106cde.pth',
'resnet34':
'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
'resnet50':
'https://download.pytorch.org/models/resnet50-19c8e357.pth',
'resnet101':
'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
'resnet152':
'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
saved_state_dict = torch.utils.model_zoo.load_url(model_urls['resnet101'])
# coco pretrained parameters:
# saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
#Scale.layer5.conv2d_list.3.weight
i_parts = str(i).split('.')
# print i_parts
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
model.load_state_dict(new_params)
#model.float()
model.eval() # use_global_stats = True
#model.train()
device = torch.device("cuda:" + str(args.gpu))
model.to(device)
cudnn.benchmark = True
trainloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps *
args.batch_size,
crop_size=input_size,
train_scale=train_scale,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN,
std=IMG_STD),
batch_size=args.batch_size,
shuffle=True,
num_workers=5,
pin_memory=args.pin_memory)
optimizer = optim.SGD([{
'params': get_1x_lr_params_NOscale(model),
'lr': args.learning_rate
}, {
'params': get_10x_lr_params(model),
'lr': 10 * args.learning_rate
}],
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
interp = nn.Upsample(size=input_size, mode='bilinear', align_corners=True)
for i_iter, batch in enumerate(trainloader):
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
pred = interp(model(images))
loss = loss_calc(pred, labels)
loss.backward()
optimizer.step()
# print('iter = ', i_iter, 'of', args.num_steps,'completed, loss = ', loss.data.cpu().numpy())
logger.info('iter = {} of {} completed, loss = {:.4f}'.format(
i_iter, args.num_steps,
loss.data.cpu().numpy()))
if i_iter >= args.num_steps - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC12_scenes_' + str(args.num_steps) + '.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'VOC12_scenes_' + str(i_iter) + '.pth'))
end = timeit.default_timer()
print(end - start, 'seconds')
def main():
"""Create the model and start the evaluation process."""
device = torch.device("cuda:" + str(args.gpu))
if not os.path.exists(args.save):
os.makedirs(args.save)
logger = util.set_logger(args.save, args.log_file, args.debug)
logger.info('start with arguments %s', args)
x_num = 0
with open(args.data_list) as f:
for _ in f.readlines():
x_num = x_num + 1
sys.path.insert(0, 'dataset/helpers')
if args.data_src == 'gta' or args.data_src == 'cityscapes':
from labels import id2label, trainId2label
elif args.data_src == 'synthia':
from labels_cityscapes_synthia import id2label, trainId2label
#
label_2_id = 255 * np.ones((256, ))
for l in id2label:
if l in (-1, 255):
continue
label_2_id[l] = id2label[l].trainId
id_2_label = np.array(
[trainId2label[_].id for _ in trainId2label if _ not in (-1, 255)])
valid_labels = sorted(set(id_2_label.ravel()))
scorer = ScoreUpdater(valid_labels, args.num_classes, x_num, logger)
scorer.reset()
if args.model == 'DeeplabRes':
model = Res_Deeplab(num_classes=args.num_classes)
# elif args.model == 'DeeplabVGG':
# model = DeeplabVGG(num_classes=args.num_classes)
# if args.restore_from == RESTORE_FROM:
# args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = str(i).split('.')
# print i_parts
if not i_parts[0] == 'fc':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
else:
loc = "cuda:" + str(args.gpu)
saved_state_dict = torch.load(args.restore_from, map_location=loc)
new_params = saved_state_dict.copy()
model.load_state_dict(new_params)
#model.train()
model.eval()
model.to(device)
testloader = data.DataLoader(GTA5TestDataSet(args.data_dir,
args.data_list,
test_scale=1.0,
test_size=(1024, 512),
mean=IMG_MEAN,
std=IMG_STD,
scale=False,
mirror=False),
batch_size=1,
shuffle=False,
pin_memory=True)
test_scales = [float(_) for _ in str(args.test_scale).split(',')]
h, w = map(int, args.test_image_size.split(','))
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(h, w), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(h, w), mode='bilinear')
test_image_size = (h, w)
mean_rgb = IMG_MEAN[::-1].copy()
std_rgb = IMG_STD[::-1].copy()
with torch.no_grad():
for index, batch in enumerate(testloader):
image, label, _, name = batch
img = image.clone()
num_scales = len(test_scales)
# output_dict = {k: [] for k in range(num_scales)}
for scale_idx in range(num_scales):
if version.parse(torch.__version__) > version.parse('0.4.0'):
image = F.interpolate(image,
scale_factor=test_scales[scale_idx],
mode='bilinear',
align_corners=True)
else:
test_size = (int(h * test_scales[scale_idx]),
int(w * test_scales[scale_idx]))
interp_tmp = nn.Upsample(size=test_size,
mode='bilinear',
align_corners=True)
image = interp_tmp(img)
if args.model == 'DeeplabRes':
output2 = model(image.to(device))
coutput = interp(output2).cpu().data[0].numpy()
if args.test_flipping:
output2 = model(
torch.from_numpy(
image.numpy()[:, :, :, ::-1].copy()).to(device))
coutput = 0.5 * (
coutput +
interp(output2).cpu().data[0].numpy()[:, :, ::-1])
if scale_idx == 0:
output = coutput.copy()
else:
output += coutput
output = output / num_scales
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
pred_label = output.copy()
label = label_2_id[np.asarray(label.numpy(), dtype=np.uint8)]
scorer.update(pred_label.flatten(), label.flatten(), index)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' %
(args.save, name.split('.')[0]))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
interp = nn.Upsample(size=(720, 1280), mode='bilinear', align_corners=True)
data_list = []
'''
Use the following Data Directories and List Path for validation (BDD Dataset)
DATA_DIRECTORY = './BDD_Deepdrive/bdd100k/'
DATA_LIST_PATH = './dataset/list/BDD_val.txt'
'''
# valid_loader = data.DataLoader(BerkeleyDataSet(args.data_dir, args.data_list, mean=IMG_MEAN, scale=False, mirror=False, train=False),
# batch_size=1, shuffle=False, pin_memory=True)
#Evaluation loop for Valid Loader
# with torch.no_grad():
# for index, batch in enumerate(valid_loader):
# if index % 100 == 0:
# print('%d processd'%(index))
# image, label, name, size = batch
# h, w, c = size[0].numpy()
# # print(name)
# output = model(Variable(image).cuda(gpu0))
# output = interp(output).cpu().data[0].numpy()
# # print(output.shape)
# output = output.transpose(1,2,0)
# output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# ground_truth = np.asarray(label[0].numpy()[:h,:w], dtype=np.int)
# show_all(ground_truth, output, name)
# data_list.append([ground_truth.flatten(), output.flatten()])
# get_iou(data_list, args.num_classes)
'''
Use the following codes on the testset.
Use the following Data Directories and List Path for testing (BerkeleyDataTestSet)
DATA_DIRECTORY = './BDD_Deepdrive/bdd100k/'
DATA_LIST_PATH = './dataset/list/BDD_test.txt'
test_loader = data.DataLoader(BerkeleyDataTestSet(args.data_dir, args.data_list, mean=IMG_MEAN))
Use the following for Cityscapes Dataset (CityscapesDataSet)
DATA_DIRECTORY = './CityScapes/stuttgart_01/'
DATA_LIST_PATH = './dataset/list/CityScapesStuttgart_01.txt'
test_loader = data.DataLoader(CityscapesDataSet(args.data_dir, args.data_list, mean=IMG_MEAN))
'''
test_loader = data.DataLoader(
BerkeleyDataTestSet(args.data_dir, args.data_list, mean=IMG_MEAN))
masked_directory = 'D:/PyTorch-DeepLab-Berkeley/BDDMasked/'
# #Evaluation loop for Test Loader
with torch.no_grad():
for index, batch in enumerate(test_loader):
if index % 100 == 0:
print('%d processd' % (index))
image, name, size = batch
h, w, c = size[0].numpy()
# print(name)
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.int)
# show_all(output, name=name)
overlay(
name=name,
original_directory='./BDD_Deepdrive/bdd100k/images/100k/test/',
masked_directory=masked_directory)