updates
This commit is contained in:
parent
0a725a4bad
commit
80516dd758
22
detect.py
22
detect.py
|
@ -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)
|
||||||
|
|
|
@ -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='')
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue