def demo(img_list, save_path=None):
net = build_yolo3('test', cfg)
net.load_state_dict(torch.load(cfg.trained_model))
if cfg.test_cuda: net = net.cuda()
net.eval()
transform = BaseTransform(size=416, mean=(0, 0, 0), scale=True)
for img in img_list:
if use_cv2:
image = cv2.imread(img)
h, w, _ = image.shape
else:
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)
def pred_file(filename, files, office):
# Note: delete all the exists files
[os.remove(file) for file in files if os.path.isfile(file)]
print('Loading trained network ...')
net = build_yolo('test', cfg, eval=True)
net.load_state_dict(torch.load(cfg.trained_model))
net.eval()
net = net.cuda() if cfg.test_cuda else net
print('Loading VOC dataset ...')
if office:
mean = (0, 0, 0)
else:
mean = (104, 117, 123) if use_cv2 else (123, 117, 104)
transform = BaseTransform(size=416, mean=mean, scale=True)
dataset = VOCDetection(cfg.voc_root, [('2007', 'test')], transform,
AnnotationTransform())
num_images = len(dataset)
_t = {'im_detect': Timer(), 'misc': Timer()}
x = torch.randn((1, 3, 416, 416))
x = x.cuda() if cfg.test_cuda else x
for i in range(num_images):
print('Testing image {:d}/{:d}....'.format(i + 1, num_images))
im, gt, h, w = dataset.pull_item(i)
idx = dataset.ids[i]
x.copy_(im.unsqueeze(0))
_t['im_detect'].tic()
with torch.no_grad():
y = net(x)
detect_time = _t['im_detect'].toc(avg=False)
# TODO: unfinish
for k in range(0, y.size(1)):
dets = y[0, k, :]
mask = dets[:, 0].gt(0.).expand(5, dets.size(0)).t()
dets = torch.masked_select(dets, mask).view(-1, 5)
if dets.size(0) == 0:
continue
boxes = dets[:, 1:]
boxes[:, 0::2] *= w
boxes[:, 1::2] *= h
scores = dets[:, 0].cpu().numpy()
cls_dets = np.c_[boxes.cpu().numpy(), scores]
for j in range(cls_dets.shape[0]):
with open(filename, mode='a') as f:
f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
idx, cls_dets[j, 4], cls_dets[j, 0], cls_dets[j, 1],
cls_dets[j, 2], cls_dets[j, 3]))
with open(files[k], mode='a') as f:
f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
idx, cls_dets[j, 4], cls_dets[j, 0], cls_dets[j, 1],
cls_dets[j, 2], cls_dets[j, 3]))
def test(net, testset, filename, transform=BaseTransform()):
num_images = len(testset)
for i in range(num_images):
print('Testing image {:d}/{:d}...'.format(i + 1, num_images))
img = testset.pull_image(i)
img_id, annotation = testset.pull_anno(i)
x, _, _ = transform(img)
x = torch.from_numpy(x).permute(2, 0, 1).unsqueeze(0)
with open(filename, mode='a') as f:
f.write('\n ground truth for: ' + img_id + '\n')
for box in annotation:
f.write('label: ' + ' || '.join(str(b) for b in box[:-1]) +
' || ' + labelmap[box[-1]] + '\n')
if cfg.test_cuda:
x = x.cuda()
with torch.no_grad():
y = net(x)
if use_cv2:
scale = torch.Tensor(
[img.shape[1], img.shape[0], img.shape[1], img.shape[0]])
else:
scale = torch.Tensor(
[img.size[0], img.size[1], img.size[0], img.size[1]])
pred_num = 0
for i in range(1, y.size(1)):
j = 0
while y[0, i, j, 0] >= 0.6:
if pred_num == 0:
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
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'))
score = y[0, i, j, 0].item()
label_name = labelmap[i]
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'
open(filename, 'w').close() # clean the txt file
# load network
net = build_yolo('test')
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())
transform = BaseTransform(size=416, mean=(0, 0, 0), scale=True)
test(net, testset, filename, transform)
net = net.cuda() if cfg.cuda else net
net.train()
loc_loss, conf_loss, prob_loss = 0, 0, 0
epoch = 0
step_index = 0
img_path = '../dataset/000004.jpg'
anno_path = '../dataset/000004.xml'
print("Training YOLO on Single Image")
if use_cv2:
image = cv2.imread(img_path)
h, w, _ = image.shape
img, _, _ = BaseTransform(416, mean=(0, 0, 0), scale=True)(image)
img = img[:, :, (2, 1, 0)]
else:
image = Image.open(img_path)
w, h = image.size
img, _, _ = BaseTransform(416, mean=(0, 0, 0), scale=True)(image)
anno = ET.parse(anno_path).getroot()
x = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0)
target = AnnotationTransform()(anno, w, h)
target = [torch.Tensor(target)]
x.requires_grad = True
if cfg.cuda:
x = x.cuda()
target = [t.cuda() for t in target]
with open(files[k], mode='a') as f:
f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'.format(
idx, cls_dets[j, 4], cls_dets[j, 0], cls_dets[j, 1],
cls_dets[j, 2], cls_dets[j, 3]))
if __name__ == '__main__':
# load net
net = build_yolo('test', cfg, eval=True)
net.load_state_dict(torch.load(cfg.trained_model))
net.eval()
if cfg.test_cuda:
net = net.cuda()
# load dataset
if cfg.use_office:
mean = (0, 0, 0)
else:
mean = (104, 117, 123) if use_cv2 else (123, 117, 104)
transform = BaseTransform(size=416, mean=mean, scale=True)
dataset = VOCDetection(cfg.voc_root, [('2007', 'test')], transform,
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.output_folder, 'eval'))