image = Image.open(img)
w, h = image.size
scale = np.array([w, h, w, h])
x, _, _ = transform(image)
x = x[:, :, (2, 1, 0)] if use_cv2 else x
x = torch.from_numpy(x).permute(2, 0, 1).unsqueeze(0)
if cfg.test_cuda: x = x.cuda()
with torch.no_grad():
y = net(x)
for i in range(y.size(1)):
idx = (y[0, i, :, 0] > 0.5)
dets = y[0, i][idx].view(-1, 5)
if dets.numel() == 0:
continue
print('Find {} {} for {}.'.format(dets.size(0), labelmap[i],
img.split('/')[-1]))
score, loc = dets[:, 0], dets[:, 1:].cpu().numpy() * scale
for k in range(len(score)):
label = '{} {:.2f}'.format(labelmap[i], score[k])
draw_box(image, label, loc[k], i)
if use_cv2:
cv2.imwrite(os.path.join(save_path, img.split('/')[-1]), image)
else:
image.save(os.path.join(save_path, img.split('/')[-1]), quality=90)
if __name__ == '__main__':
img_list = ['../dataset/000004.jpg']
file = get_output_dir(cfg.output_folder, 'demo')
demo(img_list, file)
cls_dets[j, 0], cls_dets[j, 1], cls_dets[j, 2], cls_dets[j, 3]))
print('Testing image {:d}/{:d} {:.3f}s....'.format(i + 1, num_images, detect_time))
# evaluate map
def test_map(files):
annopath = os.path.join(cfg.voc_root, 'VOC2007', 'Annotations', '{:s}.xml')
imagesetfile = os.path.join(cfg.voc_root, 'VOC2007', 'ImageSets', 'Main', 'test.txt')
cachedir = os.path.join(cfg.voc_root, 'annotations_cache')
aps = []
for i, cls in enumerate(cfg.classes):
if cls == '__background__':
continue
filename = files[i]
rec, prec, ap = voc_eval(filename, annopath, imagesetfile, cls, cachedir,
ovthresh=0.5, use_07_metric=True)
aps += [ap]
print('AP for {} = {:.4f}'.format(cls, ap))
print('Mean AP = {:.4f}'.format(np.mean(aps)))
mAP = np.mean(aps).item()
return mAP
if __name__ == '__main__':
output_dir = get_output_dir(cfg.output_folder, 'eval')
filename = output_dir + '/all.txt'
files = [output_dir + '/' + name + '.txt' for name in cfg.classes]
if not os.path.isfile(filename):
pred_file(filename, files)
map = test_map(files)
scores = dets[:, 0].cpu().numpy()
cls_dets = np.c_[boxes.cpu().numpy(), scores]
all_boxes[j][i] = cls_dets
print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, img_num,
detect_time))
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
return all_boxes
if __name__ == '__main__':
# load net
net = build_ssd('test', bone='vgg')
net.load_state_dict(torch.load(cfg.trained_model))
net.eval()
if cfg.test_cuda:
net = net.cuda()
# load dataset
dataset = VOCDetection(cfg.voc_root, [('2007', 'test')], BaseTransform(),
AnnotationTransform())
# save file
output_dir = get_output_dir(cfg.output_folder, 'eval')
det_file = os.path.join(output_dir, 'detections.pkl')
# generate boxes
print('predict boxes, and save to .pkl')
box_list = generate_boxes(dataset, net, det_file)
print('Evaluating detections')
write_voc_results_file(box_list, dataset)
do_python_eval(get_output_dir(cfg.save_folder, 'eval'))
with open(filename, mode='a') as f:
f.write('Predictions: ' + '\n')
score = y[0, i, j, 0].item()
label_name = labelmap[i - 1]
pt = (y[0, i, j, 1:] * scale).cpu().numpy()
coords = (pt[0], pt[1], pt[2], pt[3])
pred_num += 1
with open(filename, mode='a') as f:
f.write(
str(pred_num) + ' label: ' +
' || '.join('{:4.1f}'.format(c)
for c in coords) + ' || ' + label_name +
' || ' + '{:4.2f}'.format(score) + '\n')
j += 1
if __name__ == '__main__':
file = get_output_dir(cfg.output_folder, 'test')
filename = file + '/test.txt'
# load network
net = build_ssd('test', cfg.bone)
net.load_state_dict(torch.load(cfg.trained_model))
net.eval()
print('Finished loading model !')
if cfg.test_cuda:
net = net.cuda()
# load dataset
testset = VOCDetection(cfg.voc_root, [('2007', 'test')], None,
AnnotationTransform())
test(net, testset, filename)