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
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())
density_1 = model(img_1).data.cpu().numpy()