This commit is contained in:
Glenn Jocher 2019-08-31 19:22:53 +02:00
parent 0a725a4bad
commit 80516dd758
2 changed files with 12 additions and 12 deletions

View File

@ -7,7 +7,7 @@ from utils.datasets import *
from utils.utils import * from utils.utils import *
def detect(save_txt=False, save_images=True): def detect(save_txt=False, save_img=True):
img_size = (320, 192) if ONNX_EXPORT else opt.img_size # (320, 192) or (416, 256) or (608, 352) for (height, width) img_size = (320, 192) if ONNX_EXPORT else opt.img_size # (320, 192) or (416, 256) or (608, 352) for (height, width)
out = opt.output out = opt.output
@ -47,7 +47,7 @@ def detect(save_txt=False, save_images=True):
# Set Dataloader # Set Dataloader
vid_path, vid_writer = None, None vid_path, vid_writer = None, None
if opt.webcam: if opt.webcam:
save_images = False save_img = False
dataloader = LoadWebcam(img_size=img_size, half=opt.half) dataloader = LoadWebcam(img_size=img_size, half=opt.half)
else: else:
dataloader = LoadImages(opt.input, img_size=img_size, half=opt.half) dataloader = LoadImages(opt.input, img_size=img_size, half=opt.half)
@ -67,8 +67,8 @@ def detect(save_txt=False, save_images=True):
pred, _ = model(img) pred, _ = model(img)
det = non_max_suppression(pred.float(), opt.conf_thres, opt.nms_thres)[0] det = non_max_suppression(pred.float(), opt.conf_thres, opt.nms_thres)[0]
if det is not None and len(det) > 0: if det is not None and len(det):
# Rescale boxes from 416 to true image size # Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round() det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results to screen # Print results to screen
@ -77,22 +77,22 @@ def detect(save_txt=False, save_images=True):
n = (det[:, -1] == c).sum() n = (det[:, -1] == c).sum()
print('%g %ss' % (n, classes[int(c)]), end=', ') print('%g %ss' % (n, classes[int(c)]), end=', ')
# Draw bounding boxes and labels of detections # Write results
for *xyxy, conf, cls_conf, cls in det: for *xyxy, conf, _, cls in det:
if save_txt: # Write to file if save_txt: # Write to file
with open(save_path + '.txt', 'a') as file: with open(save_path + '.txt', 'a') as file:
file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf)) file.write(('%g ' * 6 + '\n') % (*xyxy, cls, conf))
# Add bbox to the image if save_img: # Add bbox to image
label = '%s %.2f' % (classes[int(cls)], conf) label = '%s %.2f' % (classes[int(cls)], conf)
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)]) plot_one_box(xyxy, im0, label=label, color=colors[int(cls)])
print('Done. (%.3fs)' % (time.time() - t)) print('Done. (%.3fs)' % (time.time() - t))
if opt.webcam: # Show live webcam if opt.webcam: # Show live webcam
cv2.imshow(opt.weights, im0) cv2.imshow(opt.weights, im0)
if save_images: # Save image with detections if save_img: # Save image with detections
if dataloader.mode == 'images': if dataloader.mode == 'images':
cv2.imwrite(save_path, im0) cv2.imwrite(save_path, im0)
else: else:
@ -107,7 +107,7 @@ def detect(save_txt=False, save_images=True):
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (width, height)) vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (width, height))
vid_writer.write(im0) vid_writer.write(im0)
if save_images: if save_img:
print('Results saved to %s' % os.getcwd() + os.sep + out) print('Results saved to %s' % os.getcwd() + os.sep + out)
if platform == 'darwin': # MacOS if platform == 'darwin': # MacOS
os.system('open ' + out + ' ' + save_path) os.system('open ' + out + ' ' + save_path)

View File

@ -162,7 +162,7 @@ class LoadWebcam: # for inference
break break
# Print # Print
assert ret_val, 'Webcam Error' assert ret_val, 'Camera Error %s' % self.pipe
img_path = 'webcam_%g.jpg' % self.count img_path = 'webcam_%g.jpg' % self.count
print('webcam %g: ' % self.count, end='') print('webcam %g: ' % self.count, end='')