Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-02-10 21:06:22 +01:00
parent 9d12a162f8
commit 97909df1a6
1 changed files with 48 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

81
detect.py
View File

@ -8,10 +8,18 @@ from utils.utils import *
from utils import torch_utils from utils import torch_utils
def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3, nms_thres=0.45, save_txt=False, def detect(
save_images=True): cfg,
weights,
images,
output='output',
img_size=416,
conf_thres=0.3,
nms_thres=0.45,
save_txt=False,
save_images=True
):
device = torch_utils.select_device() device = torch_utils.select_device()
os.system('rm -rf ' + output) os.system('rm -rf ' + output)
os.makedirs(output, exist_ok=True) os.makedirs(output, exist_ok=True)
@ -39,43 +47,42 @@ def detect(cfg, weights, images, output='output', img_size=416, conf_thres=0.3,
t = time.time() t = time.time()
# Get detections # Get detections
with torch.no_grad(): img = torch.from_numpy(img).unsqueeze(0).to(device)
img = torch.from_numpy(img).unsqueeze(0).to(device) if ONNX_EXPORT:
if ONNX_EXPORT: pred = torch.onnx._export(model, img, 'weights/model.onnx', verbose=True)
pred = torch.onnx._export(model, img, 'weights/model.onnx', verbose=True) return # ONNX export
return # ONNX export pred = model(img)
pred = model(img) pred = pred[pred[:, :, 4] > conf_thres]
pred = pred[pred[:, :, 4] > conf_thres]
if len(pred) > 0: if len(pred) > 0:
detections = non_max_suppression(pred.unsqueeze(0), conf_thres, nms_thres)[0] detections = non_max_suppression(pred.unsqueeze(0), conf_thres, nms_thres)[0]
# Draw bounding boxes and labels of detections # Draw bounding boxes and labels of detections
if detections is not None: if detections is not None:
save_img_path = os.path.join(output, path.split('/')[-1]) save_img_path = os.path.join(output, path.split('/')[-1])
save_txt_path = save_img_path + '.txt' save_txt_path = save_img_path + '.txt'
# Rescale boxes from 416 to true image size # Rescale boxes from 416 to true image size
detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape) detections[:, :4] = scale_coords(img_size, detections[:, :4], im0.shape)
unique_classes = detections[:, -1].cpu().unique() unique_classes = detections[:, -1].cpu().unique()
for i in unique_classes: for i in unique_classes:
n = (detections[:, -1].cpu() == i).sum() n = (detections[:, -1].cpu() == i).sum()
print('%g %ss' % (n, classes[int(i)]), end=', ') print('%g %ss' % (n, classes[int(i)]), end=', ')
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections: for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
if save_txt: # Write to file if save_txt: # Write to file
with open(save_txt_path, 'a') as file: with open(save_txt_path, 'a') as file:
file.write('%g %g %g %g %g %g\n' % (x1, y1, x2, y2, cls_pred, cls_conf * conf)) file.write('%g %g %g %g %g %g\n' % (x1, y1, x2, y2, cls_pred, cls_conf * conf))
if save_images: # Add bbox to the image if save_images: # Add bbox to the image
label = '%s %.2f' % (classes[int(cls_pred)], conf) label = '%s %.2f' % (classes[int(cls_pred)], conf)
plot_one_box([x1, y1, x2, y2], im0, label=label, color=colors[int(cls_pred)]) plot_one_box([x1, y1, x2, y2], im0, label=label, color=colors[int(cls_pred)])
if save_images: # Save generated image with detections if save_images: # Save generated image with detections
cv2.imwrite(save_img_path, im0) cv2.imwrite(save_img_path, im0)
print(' Done. (%.3fs)' % (time.time() - t)) print(' Done. (%.3fs)' % (time.time() - t))
if platform == 'darwin': # MacOS if platform == 'darwin': # MacOS
os.system('open ' + output + '&& open ' + save_img_path) os.system('open ' + output + '&& open ' + save_img_path)
@ -92,4 +99,12 @@ if __name__ == '__main__':
opt = parser.parse_args() opt = parser.parse_args()
print(opt) print(opt)
detect(opt.cfg, opt.weights, opt.images, img_size=opt.img_size, conf_thres=opt.conf_thres, nms_thres=opt.nms_thres) with torch.no_grad():
detect(
opt.cfg,
opt.weights,
opt.images,
img_size=opt.img_size,
conf_thres=opt.conf_thres,
nms_thres=opt.nms_thres
)