def predict_density_map(test_root, checkpoint_path, device, index):
model = CANNet().to(device)
model.load_state_dict(torch.load(checkpoint_path))
test_loader = torch.utils.data.DataLoader(ShanghaiTechPartA(
test_root,
transform=transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
downsample=8),
batch_size=args.batch_size)
model.eval()
for i, (img, density_map) in enumerate(test_loader):
if i == index:
img = img.to(device)
density_map = density_map.to(device)
est_density_map = model(img).detach()
est_density_map = est_density_map.squeeze(0).squeeze(
0).cpu().numpy()
plt.imshow(est_density_map, cmap=CM.jet)
break
def main():
global args, best_prec1
best_prec1 = 1e6
args = parser.parse_args()
args.lr = 1e-4
args.batch_size = 12 #26
args.decay = 5 * 1e-4
args.start_epoch = 0
args.epochs = 1000
args.workers = 4
args.seed = int(time.time())
args.print_freq = 4
with open(args.train_json, 'r') as outfile:
train_list = json.load(outfile)
print len(train_list)
with open(args.val_json, 'r') as outfile:
val_list = json.load(outfile)
print len(val_list)
torch.cuda.manual_seed(args.seed)
model = CANNet()
model = model.cuda()
criterion = nn.MSELoss(size_average=False).cuda()
optimizer = torch.optim.Adam(model.parameters(),
args.lr,
weight_decay=args.decay)
###########
if args.best:
print("=> loading best checkpoint '{}'".format(args.best))
checkpoint = torch.load(os.path.join(args.output,
'model_best.pth.tar'))
model.load_state_dict(checkpoint['state_dict'])
best_prec1 = validate(val_list, model, criterion)
print(' * best MAE {mae:.3f} '.format(mae=best_prec1))
if args.pre:
if os.path.isfile(args.pre):
print("=> loading checkpoint '{}'".format(args.pre))
checkpoint = torch.load(args.pre)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.pre))
if args.initial:
args.start_epoch = args.initial
print(args.initial)
print(x)
for epoch in range(args.start_epoch, args.epochs):
train(train_list, model, criterion, optimizer, epoch)
prec1 = validate(val_list, model, criterion)
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
print(' * best MAE {mae:.3f} '.format(mae=best_prec1))
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
},
is_best,
args.output,
filename='checkpoint.pth.tar')
print('Train process finished')
parser.add_argument('test_json', metavar='test', help='path to val json')
parser.add_argument('output', metavar='VAL', help='path output')
args = parser.parse_args()
with open(args.test_json, 'r') as outfile:
img_paths = json.load(outfile)
model = CANNet()
model = model.cuda()
checkpoint = torch.load(os.path.join(args.output, 'model_best.pth.tar'))
model.load_state_dict(checkpoint['state_dict'])
model.eval()
pred = []
gt = []
for i in xrange(len(img_paths)):
img = transform(Image.open(img_paths[i]).convert('RGB')).cuda()
img = img.unsqueeze(0)
h, w = img.shape[2:4]
h_d = h / 2
w_d = w / 2
img_1 = Variable(img[:, :, :h_d, :w_d].cuda())
img_2 = Variable(img[:, :, :h_d, w_d:].cuda())
img_3 = Variable(img[:, :, h_d:, :w_d].cuda())
img_4 = Variable(img[:, :, h_d:, w_d:].cuda())