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

updates

This commit is contained in:
Glenn Jocher 2019-09-19 18:43:29 +02:00
parent 0f3f6c03e7
commit b5db03827f
2 changed files with 10 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
detect.py
View File

@ -9,8 +9,7 @@ from utils.utils import *
def detect(save_txt=False, save_img=False, view_img=False): def detect(save_txt=False, save_img=False, view_img=False):
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, source, weights, half = opt.output, opt.source, opt.weights, opt.half out, source, weights, half = opt.output, opt.source, opt.weights, opt.half
webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')
streams = 'streams' in source and source.endswith('.txt')
# Initialize # Initialize
device = torch_utils.select_device(device='cpu' if ONNX_EXPORT else opt.device) device = torch_utils.select_device(device='cpu' if ONNX_EXPORT else opt.device)
@ -47,13 +46,10 @@ def detect(save_txt=False, save_img=False, view_img=False):
# Set Dataloader # Set Dataloader
vid_path, vid_writer = None, None vid_path, vid_writer = None, None
if streams: if webcam:
view_img = False view_img = True
torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=img_size, half=half) dataset = LoadStreams(source, img_size=img_size, half=half)
elif webcam:
view_img = True
dataset = LoadWebcam(source, img_size=img_size, half=half)
else: else:
save_img = True save_img = True
dataset = LoadImages(source, img_size=img_size, half=half) dataset = LoadImages(source, img_size=img_size, half=half)
@ -74,7 +70,7 @@ def detect(save_txt=False, save_img=False, view_img=False):
pred, _ = model(img) pred, _ = model(img)
for i, det in enumerate(non_max_suppression(pred, opt.conf_thres, opt.nms_thres)): # detections per image for i, det in enumerate(non_max_suppression(pred, opt.conf_thres, opt.nms_thres)): # detections per image
if streams: # batch_size > 1 if webcam: # batch_size >= 1
p, s, im0 = path[i], '%g: ' % i, im0s[i] p, s, im0 = path[i], '%g: ' % i, im0s[i]
else: else:
p, s, im0 = path, '', im0s p, s, im0 = path, '', im0s

8
utils/datasets.py
View File

@ -186,10 +186,12 @@ class LoadWebcam: # for inference
class LoadStreams: # multiple IP or RTSP cameras class LoadStreams: # multiple IP or RTSP cameras
def __init__(self, path='streams.txt', img_size=416, half=False): def __init__(self, sources='streams.txt', img_size=416, half=False):
self.img_size = img_size self.img_size = img_size
self.half = half # half precision fp16 images self.half = half # half precision fp16 images
with open(path, 'r') as f:
if os.path.isfile(sources):
with open(sources, 'r') as f:
sources = [x.strip() for x in f.read().splitlines() if len(x.strip())] sources = [x.strip() for x in f.read().splitlines() if len(x.strip())]
n = len(sources) n = len(sources)
@ -208,7 +210,7 @@ class LoadStreams: # multiple IP or RTSP cameras
thread.start() thread.start()
print('') # newline print('') # newline
time.sleep(0.5) time.sleep(0.5) # allow connections to start
def update(self, index, cap): def update(self, index, cap):
# Read next stream frame in a daemon thread # Read next stream frame in a daemon thread