def __init__(self):
self.conf = self.DETECTION_THRESHOLD
self.device = select_device()
self.model = DetectMultiBackend(paths.yolo_model, device=self.device)
self.model.model.half()
self.model.warmup(imgsz=(1, 3, *self.IMG_SIZE), half=True)
class YoloDetector(PedestrianDetector):
# TODO: Store model parameters somewhere else
NMS_IOU_THRESHOLD = 0.6
DETECTION_THRESHOLD = 0.451
IMG_SIZE = [640, 640]
def __init__(self):
self.conf = self.DETECTION_THRESHOLD
self.device = select_device()
self.model = DetectMultiBackend(paths.yolo_model, device=self.device)
self.model.model.half()
self.model.warmup(imgsz=(1, 3, *self.IMG_SIZE), half=True)
@torch.no_grad()
def detect_pedestrians(self, camera_frame: np.ndarray) -> List[BoundingBox]:
# Pre-process
im = letterbox(camera_frame, self.IMG_SIZE, stride=self.model.stride)[0]
im = np.ascontiguousarray(im.transpose(2, 0, 1))
im = torch.from_numpy(np.expand_dims(im, 0)).to(self.device)
im = im.half()
im /= 255
# Inference
pred = self.model(im)
# NMS
pred = non_max_suppression(
pred, conf_thres=self.conf, iou_thres=self.NMS_IOU_THRESHOLD
)[0]
# Scale predictions to original image coordinates
pred_boxes = (
scale_coords(im.shape[2:], pred[:, :4], camera_frame.shape)
.round()
.cpu()
.numpy()
.astype(int)
)
return [self.detection_to_bounding_box(det) for det in pred_boxes]
def detection_to_bounding_box(self, detection: np.ndarray) -> BoundingBox:
x_min, y_min, x_max, y_max = detection
return BoundingBox(x_min=x_min, y_min=y_min, x_max=x_max, y_max=y_max)
def _create(
name,
pretrained=True,
channels=3,
classes=80,
autoshape=True,
verbose=True,
device=None,
):
"""Creates or loads a YOLOv5 model
Arguments:
name (str): model name 'yolov5s' or path 'path/to/best.pt'
pretrained (bool): load pretrained weights into the model
channels (int): number of input channels
classes (int): number of model classes
autoshape (bool): apply YOLOv5 .autoshape() wrapper to model
verbose (bool): print all information to screen
device (str, torch.device, None): device to use for model parameters
Returns:
YOLOv5 model
"""
from pathlib import Path
from yolov5.models.common import AutoShape, DetectMultiBackend
from yolov5.models.yolo import Model
from yolov5.utils.downloads import attempt_download
from yolov5.utils.general import (
LOGGER,
check_requirements,
intersect_dicts,
logging,
)
from yolov5.utils.torch_utils import select_device
if not verbose:
LOGGER.setLevel(logging.WARNING)
check_requirements(exclude=("tensorboard", "thop", "opencv-python"))
name = Path(name)
path = name.with_suffix(
".pt") if name.suffix == "" else name # checkpoint path
try:
device = select_device(("0" if torch.cuda.is_available() else "cpu"
) if device is None else device)
if pretrained and channels == 3 and classes == 80:
model = DetectMultiBackend(
path, device=device) # download/load FP32 model
# model = models.experimental.attempt_load(path, map_location=device) # download/load FP32 model
else:
cfg = list(
(Path(__file__).parent /
"models").rglob(f"{path.stem}.yaml"))[0] # model.yaml path
model = Model(cfg, channels, classes) # create model
if pretrained:
ckpt = torch.load(attempt_download(path),
map_location=device) # load
csd = (ckpt["model"].float().state_dict()
) # checkpoint state_dict as FP32
csd = intersect_dicts(csd,
model.state_dict(),
exclude=["anchors"]) # intersect
model.load_state_dict(csd, strict=False) # load
if len(ckpt["model"].names) == classes:
model.names = ckpt[
"model"].names # set class names attribute
if autoshape:
model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS
return model.to(device)
except Exception as e:
help_url = "https://github.com/ultralytics/yolov5/issues/36"
s = f"{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help."
raise Exception(s) from e
def run(
weights=ROOT / 'yolov5s.pt', # model.pt path(s)
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
):
source = str(source)
save_img = not nosave and not source.endswith(
'.txt') # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
is_url = source.lower().startswith(
('rtsp://', 'rtmp://', 'http://', 'https://'))
webcam = source.isnumeric() or source.endswith('.txt') or (is_url
and not is_file)
if is_url and is_file:
source = check_file(source) # download
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= (
pt or jit or onnx or engine
) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()
# Dataloader
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
bs = len(dataset) # batch_size
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Run inference
model.warmup(imgsz=(1 if pt else bs, 3, *imgsz), half=half) # warmup
dt, seen = [0.0, 0.0, 0.0], 0
for path, im, im0s, vid_cap, s in dataset:
t1 = time_sync()
im = torch.from_numpy(im).to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path).stem,
mkdir=True) if visualize else False
pred = model(im, augment=augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# NMS
pred = non_max_suppression(pred,
conf_thres,
iou_thres,
classes,
agnostic_nms,
max_det=max_det)
dt[2] += time_sync() - t3
# Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
# Process predictions
for i, det in enumerate(pred): # per image
seen += 1
if webcam: # batch_size >= 1
p, im0, frame = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg
txt_path = str(save_dir / 'labels' / p.stem) + (
'' if dataset.mode == 'image' else f'_{frame}') # im.txt
s += '%gx%g ' % im.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1,
0]] # normalization gain whwh
imc = im0.copy() if save_crop else im0 # for save_crop
annotator = Annotator(im0,
line_width=line_thickness,
example=str(names))
if len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(im.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
for *xyxy, conf, cls in reversed(det):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if save_conf else (cls,
*xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class
label = None if hide_labels else (
names[c]
if hide_conf else f'{names[c]} {conf:.2f}')
annotator.box_label(xyxy, label, color=colors(c, True))
if save_crop:
save_one_box(xyxy,
imc,
file=save_dir / 'crops' / names[c] /
f'{p.stem}.jpg',
BGR=True)
# Stream results
im0 = annotator.result()
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream'
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release(
) # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path = str(Path(save_path).with_suffix(
'.mp4')) # force *.mp4 suffix on results videos
vid_writer[i] = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps,
(w, h))
vid_writer[i].write(im0)
# Print time (inference-only)
LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}'
% t)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
strip_optimizer(weights) # update model (to fix SourceChangeWarning)
def rtsp_to_mongodb():
with open("/home/asyed/airflow/dags/parameters.json") as f:
parms = json.load(f)
agnostic_nms = parms["agnostic_nms"]
augment = parms["augment"]
classes = parms["classes"]
conf_thres = parms["conf_thres"]
config_deepsort = parms["config_deepsort"]
deep_sort_model = parms["deep_sort_model"]
device = parms["device"]
dnn = False
evaluate = parms["evaluate"]
exist_ok = parms["exist_ok"]
fourcc = parms["fourcc"]
half = False
print(device)
imgsz = parms["imgsz"]
iou_thres = parms["iou_thres"]
max_det = parms["max_det"]
name = parms["name"]
# save_vid = parms["save_vid"]
#show_vid = parms["show_vid"]
source = parms["source"]
visualize = parms["visualize"]
yolo_model = parms["yolo_model"]
webcam = parms["webcam"]
save_txt = parms["save_txt"]
homography = np.array(parms["homography"])
url = "mongodb://localhost:27017"
client = MongoClient(url)
db = client.trajectory_database
today_date = date.today().strftime("%m-%d-%y")
new = "file_image_coordinates_" + today_date
collection = db[new]
cfg = get_config()
cfg.merge_from_file(config_deepsort)
deepsort = DeepSort(deep_sort_model,
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
device = select_device(device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
# make new output folder
# Load model
device = select_device(device)
model = DetectMultiBackend(yolo_model, device=device, dnn=dnn)
stride, names, pt, jit, _ = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source,
img_size=imgsz,
stride=stride,
auto=pt and not jit)
bs = len(dataset) # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
if pt and device.type != 'cpu':
model(
torch.zeros(1, 3, *imgsz).to(device).type_as(
next(model.model.parameters()))) # warmup
# global framess_im2
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
# arr = None
past = []
for frame_idx, (path, img, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
img = torch.from_numpy(img).to(device)
# print("raw_frame",img.shape)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t2 = time_sync()
dt[0] += t2 - t1
pred = model(img, augment=augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
pred = non_max_suppression(pred,
conf_thres,
iou_thres,
classes,
agnostic_nms,
max_det=max_det)
dt[2] += time_sync() - t3
# Process detections
# dets_per_img = []
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # batch_size >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
annotator = Annotator(im0, line_width=2, pil=not ascii)
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort
t4 = time_sync()
outputs = deepsort.update(xywhs.cpu(), confs.cpu(),
clss.cpu(), im0)
t5 = time_sync()
dt[3] += t5 - t4
if len(outputs) > 0:
for j, (output, conf) in enumerate(zip(outputs,
confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
c = int(cls) # integer class
label = f'{id} {names[c]} {conf:.2f}'
annotator.box_label(bboxes,
label,
color=colors(c, True))
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# bbox_left = bbox_left + bbox_h
bbox_top = bbox_top + bbox_h
agent_data = {
'frame': int(frame_idx + 1),
'agent_id': int(id),
"labels": str(names[c]),
"x": int(bbox_left),
"y": int(bbox_top)
}
print("agent", agent_data)
collection.insert_one(agent_data)
#db.object_detection.insert_one(agent_data)
#db.pedestrian_detection_15_june.insert_one(agent_data)
#db.test_21_july.insert_one(agent_data)
LOGGER.info(
f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)'
)
else:
deepsort.increment_ages()
LOGGER.info('No detections')
im0 = annotator.result()
def run(
data,
weights=None, # model.pt path(s)
batch_size=32, # batch size
imgsz=640, # inference size (pixels)
conf_thres=0.001, # confidence threshold
iou_thres=0.6, # NMS IoU threshold
task='val', # train, val, test, speed or study
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
workers=8, # max dataloader workers (per RANK in DDP mode)
single_cls=False, # treat as single-class dataset
augment=False, # augmented inference
verbose=False, # verbose output
save_txt=False, # save results to *.txt
save_hybrid=False, # save label+prediction hybrid results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_json=False, # save a COCO-JSON results file
project=ROOT / 'runs/val', # save to project/name
name='exp', # save to project/name
exist_ok=False, # existing project/name ok, do not increment
half=True, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
model=None,
dataloader=None,
save_dir=Path(''),
plots=True,
callbacks=Callbacks(),
compute_loss=None,
):
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device, pt, jit, engine = next(model.parameters(
)).device, True, False, False # get model device, PyTorch model
half &= device.type != 'cpu' # half precision only supported on CUDA
model.half() if half else model.float()
else: # called directly
device = select_device(device, batch_size=batch_size)
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data)
stride, pt, jit, onnx, engine = model.stride, model.pt, model.jit, model.onnx, model.engine
imgsz = check_img_size(imgsz, s=stride) # check image size
half &= (
pt or jit or onnx or engine
) and device.type != 'cpu' # FP16 supported on limited backends with CUDA
if pt or jit:
model.model.half() if half else model.model.float()
elif engine:
batch_size = model.batch_size
else:
half = False
batch_size = 1 # export.py models default to batch-size 1
device = torch.device('cpu')
LOGGER.info(
f'Forcing --batch-size 1 square inference shape(1,3,{imgsz},{imgsz}) for non-PyTorch backends'
)
# Data
data = check_dataset(data) # check
# Configure
model.eval()
is_coco = isinstance(data.get('val'), str) and data['val'].endswith(
'coco/val2017.txt') # COCO dataset
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training:
model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz),
half=half) # warmup
pad = 0.0 if task == 'speed' else 0.5
task = task if task in (
'train', 'val', 'test') else 'val' # path to train/val/test images
dataloader = create_dataloader(data[task],
imgsz,
batch_size,
stride,
single_cls,
pad=pad,
rect=pt,
workers=workers,
prefix=colorstr(f'{task}: '))[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R',
'[email protected]', '[email protected]:.95')
dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0,
0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
pbar = tqdm(dataloader,
desc=s,
bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
t1 = time_sync()
if pt or jit or engine:
im = im.to(device, non_blocking=True)
targets = targets.to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
nb, _, height, width = im.shape # batch size, channels, height, width
t2 = time_sync()
dt[0] += t2 - t1
# Inference
out, train_out = model(im) if training else model(
im, augment=augment, val=True) # inference, loss outputs
dt[1] += time_sync() - t2
# Loss
if compute_loss:
loss += compute_loss([x.float() for x in train_out],
targets)[1] # box, obj, cls
# NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:]
for i in range(nb)] if save_hybrid else [] # for autolabelling
t3 = time_sync()
out = non_max_suppression(out,
conf_thres,
iou_thres,
labels=lb,
multi_label=True,
agnostic=single_cls)
dt[2] += time_sync() - t3
# Metrics
for si, pred in enumerate(out):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path, shape = Path(paths[si]), shapes[si][0]
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
if single_cls:
pred[:, 5] = 0
predn = pred.clone()
scale_coords(im[si].shape[1:], predn[:, :4], shape,
shapes[si][1]) # native-space pred
# Evaluate
if nl:
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
scale_coords(im[si].shape[1:], tbox, shape,
shapes[si][1]) # native-space labels
labelsn = torch.cat((labels[:, 0:1], tbox),
1) # native-space labels
correct = process_batch(predn, labelsn, iouv)
if plots:
confusion_matrix.process_batch(predn, labelsn)
else:
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(),
tcls)) # (correct, conf, pcls, tcls)
# Save/log
if save_txt:
save_one_txt(predn,
save_conf,
shape,
file=save_dir / 'labels' / (path.stem + '.txt'))
if save_json:
save_one_json(predn, jdict, path,
class_map) # append to COCO-JSON dictionary
callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
# Plot images
if plots and batch_i < 3:
f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images,
args=(im, targets, paths, f, names),
daemon=True).start()
f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images,
args=(im, output_to_target(out), paths, f, names),
daemon=True).start()
# Compute metrics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
ap50, ap = ap[:, 0], ap.mean(1) # [email protected], [email protected]:0.95
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
LOGGER.info(pf %
(names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
if not training:
shape = (batch_size, 3, imgsz, imgsz)
LOGGER.info(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
callbacks.run('on_val_end')
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights
).stem if weights is not None else '' # weights
anno_json = str(
Path(data.get('path', '../coco')) /
'annotations/instances_val2017.json') # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
check_requirements(['pycocotools'])
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [
int(Path(x).stem) for x in dataloader.dataset.img_files
] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
LOGGER.info(f'pycocotools unable to run: {e}')
# Return results
model.float() # for training
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
def detect(opt):
memory = {}
counter = 0
out, source, yolo_model, deep_sort_model, show_vid, save_vid, save_txt, imgsz, evaluate, half, project, name, exist_ok= \
opt.output, opt.source, opt.yolo_model, opt.deep_sort_model, opt.show_vid, opt.save_vid, \
opt.save_txt, opt.imgsz, opt.evaluate, opt.half, opt.project, opt.name, opt.exist_ok
webcam = source == '0' or source.startswith('rtsp') or source.startswith(
'http') or source.endswith('.txt')
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(deep_sort_model,
torch.device("cpu"),
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET)
# Initialize
device = select_device(opt.device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
if not evaluate:
if os.path.exists(out):
pass
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
save_dir.mkdir(parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model = DetectMultiBackend(yolo_model, device=device, dnn=opt.dnn)
stride, names, pt, jit, _ = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
if show_vid:
show_vid = check_imshow()
# Dataloader
if webcam:
show_vid = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source,
img_size=imgsz,
stride=stride,
auto=pt and not jit)
bs = len(dataset) # batch_size
else:
dataset = LoadImages(source,
img_size=imgsz,
stride=stride,
auto=pt and not jit)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# extract what is in between the last '/' and last '.'
txt_file_name = source.split('/')[-1].split('.')[0]
txt_path = str(Path(save_dir)) + '/' + txt_file_name + '.txt'
if pt and device.type != 'cpu':
model(
torch.zeros(1, 3, *imgsz).to(device).type_as(
next(model.model.parameters()))) # warmup
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
regionid = set()
for frame_idx, (path, img, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path).stem,
mkdir=True) if opt.visualize else False
pred = model(img, augment=opt.augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# Apply NMS
pred = non_max_suppression(pred,
opt.conf_thres,
opt.iou_thres,
opt.classes,
opt.agnostic_nms,
max_det=opt.max_det)
dt[2] += time_sync() - t3
# Process detections
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # batch_size >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
s += '%gx%g ' % img.shape[2:] # print string
annotator = Annotator(im0,
line_width=2,
font='Arial.ttf',
pil=not ascii)
if det is not None and len(det):
tboxes = []
indexIDs = []
previous = memory.copy()
memory = {}
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4],
im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort
t4 = time_sync()
outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss.cpu(),
im0)
t5 = time_sync()
dt[3] += t5 - t4
# draw boxes for visualization
if len(outputs) > 0:
for j, (output, conf) in enumerate(zip(outputs, confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
roi = [(0, 0), (640, 0), (640, 380), (0, 380)]
(x, y) = (int(bboxes[0]), int(bboxes[1]))
(w, h) = (int(bboxes[2]), int(bboxes[3]))
inside = cv2.pointPolygonTest(np.array(roi), (x, h),
False)
if inside > 0:
regionid.add(id)
c = int(cls) # integer class
label = f' {names[c]} {conf:.2f}'
cv2.putText(im0, "count =" + str(len(regionid)),
(20, 50), 0, 1, (100, 200, 0), 2)
annotator.box_label(bboxes,
label,
color=colors(c, True))
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# Write MOT compliant results to file
with open(txt_path, 'a') as f:
f.write(('%g ' * 10 + '\n') % (
frame_idx + 1,
id,
bbox_left, # MOT format
bbox_top,
bbox_w,
bbox_h,
-1,
-1,
-1,
-1))
LOGGER.info(
f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)'
)
LOGGER.info(f'counter = {len(regionid)}')
else:
deepsort.increment_ages()
LOGGER.info('No detections')
# Stream results
im0 = annotator.result()
if show_vid:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_vid:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
vid_writer = cv2.VideoWriter(
save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps,
(w, h))
vid_writer.write(im0)
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms deep sort update \
per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_vid:
print('Results saved to %s' % save_path)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
def detect(opt):
out, source, yolo_model, deep_sort_model, show_vid, save_vid, save_txt, imgsz, evaluate, half, \
project, exist_ok, update, save_crop = \
opt.output, opt.source, opt.yolo_model, opt.deep_sort_model, opt.show_vid, opt.save_vid, \
opt.save_txt, opt.imgsz, opt.evaluate, opt.half, opt.project, opt.exist_ok, opt.update, opt.save_crop
webcam = source == '0' or source.startswith(
'rtsp') or source.startswith('http') or source.endswith('.txt')
# Initialize
device = select_device(opt.device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
if not evaluate:
if os.path.exists(out):
pass
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Directories
if type(yolo_model) is str: # single yolo model
exp_name = yolo_model.split(".")[0]
elif type(yolo_model) is list and len(yolo_model) == 1: # single models after --yolo_model
exp_name = yolo_model[0].split(".")[0]
else: # multiple models after --yolo_model
exp_name = "ensemble"
exp_name = exp_name + "_" + deep_sort_model.split('/')[-1].split('.')[0]
save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok) # increment run if project name exists
(save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Load model
model = DetectMultiBackend(yolo_model, device=device, dnn=opt.dnn)
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
if show_vid:
show_vid = check_imshow()
# Dataloader
if webcam:
show_vid = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
nr_sources = len(dataset)
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
nr_sources = 1
vid_path, vid_writer, txt_path = [None] * nr_sources, [None] * nr_sources, [None] * nr_sources
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
# Create as many trackers as there are video sources
deepsort_list = []
for i in range(nr_sources):
deepsort_list.append(
DeepSort(
deep_sort_model,
device,
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
)
)
outputs = [None] * nr_sources
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# Run tracking
model.warmup(imgsz=(1 if pt else nr_sources, 3, *imgsz)) # warmup
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
im = torch.from_numpy(im).to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255.0 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path[0]).stem, mkdir=True) if opt.visualize else False
pred = model(im, augment=opt.augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms, max_det=opt.max_det)
dt[2] += time_sync() - t3
# Process detections
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # nr_sources >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
p = Path(p) # to Path
s += f'{i}: '
txt_file_name = p.name
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
# video file
if source.endswith(VID_FORMATS):
txt_file_name = p.stem
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
# folder with imgs
else:
txt_file_name = p.parent.name # get folder name containing current img
save_path = str(save_dir / p.parent.name) # im.jpg, vid.mp4, ...
txt_path = str(save_dir / 'tracks' / txt_file_name) # im.txt
s += '%gx%g ' % im.shape[2:] # print string
imc = im0.copy() if save_crop else im0 # for save_crop
annotator = Annotator(im0, line_width=2, pil=not ascii)
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort
t4 = time_sync()
outputs[i] = deepsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
t5 = time_sync()
dt[3] += t5 - t4
# draw boxes for visualization
if len(outputs[i]) > 0:
for j, (output, conf) in enumerate(zip(outputs[i], confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# Write MOT compliant results to file
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * 10 + '\n') % (frame_idx + 1, id, bbox_left, # MOT format
bbox_top, bbox_w, bbox_h, -1, -1, -1, i))
if save_vid or save_crop or show_vid: # Add bbox to image
c = int(cls) # integer class
label = f'{id} {names[c]} {conf:.2f}'
annotator.box_label(bboxes, label, color=colors(c, True))
if save_crop:
txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''
save_one_box(bboxes, imc, file=save_dir / 'crops' / txt_file_name / names[c] / f'{id}' / f'{p.stem}.jpg', BGR=True)
LOGGER.info(f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s)')
else:
deepsort_list[i].increment_ages()
LOGGER.info('No detections')
# Stream results
im0 = annotator.result()
if show_vid:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
if save_vid:
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer[i].write(im0)
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms deep sort update \
per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_vid:
s = f"\n{len(list(save_dir.glob('tracks/*.txt')))} tracks saved to {save_dir / 'tracks'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
strip_optimizer(yolo_model) # update model (to fix SourceChangeWarning)
def evaluate(
data: Path,
weights: Path,
conf_threshold: float = None,
batch_size: int = 256,
img_size: int = 640,
iou_thres: float = 0.6,
) -> None:
"""Evaluate model on dataset
Args:
data_dir: Dataset dir path
weights: Path to model.pt
conf_threshold: Confidence threshold
batch_size: Batch size
img_size: Inference size (pixels)
iou_thres: Non-max supression IoU threshold
"""
device = select_device(batch_size=batch_size) # cuda | cpu
save_dir = config_paths.temp_dir_path / "yolo" / "eval" / str(int(time()))
save_dir.mkdir(parents=True)
model = DetectMultiBackend(weights, device)
model.model.float()
model.eval()
img_size = check_img_size(img_size, s=model.stride) # type:ignore
model.warmup(imgsz=(1, 3, img_size, img_size))
metadata = pd.read_csv(data / "metadata.csv")
dataloader = create_dataloader(
data,
img_size,
batch_size,
model.stride,
single_cls=False,
pad=0.5,
rect=True,
)[0]
# Configure
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector 0.5:0.95
seen = 0
stats: List[ImgRes] = []
augmented_res = {}
for batch_i, (im, targets, paths, shapes) in enumerate(tqdm(dataloader)):
# Preprocess
im = im.to(device, non_blocking=True)
im = im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
*_, height, width = im.shape # batch size, channels, height, width
# [Total targetx x 6] = Tensor[[img_index, class, ...xcycwh-rel]]
targets = targets.to(device)
# Inference
out, _ = model(im, val=True) # inference, loss outputs
# NMS
# Out format (after NMS), sligthly unclear before NMS...
# [Images x Predictions x [...xyxy-abs, conf, class]]
# List[Tensor[Tensor[]]]
# Looks like bboxes are in xyxy-abs format
out = non_max_suppression(prediction=out,
conf_thres=0.001,
iou_thres=iou_thres)
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(
device) # to pixels i.e. Tensor[img_index, class, ...xcycwh-abs]
# Metrics
# Pred: Tensor[Predictions x [...xyxy-abs, conf class]]
for si, pred in enumerate(out): # for each image
seen += 1
path = Path(paths[si]) # path to img,
shape = shapes[si][0] # original img dimensions
# Labels: Each target where image_index matches current index excluding image_index col
# i.e. all targets in current image [Targets x [class, ...xcyxcwh-abs]]
labels = targets[targets[:, 0] == si, 1:]
# target class, one entry for each labeled object
labelsn = torch.Tensor()
if len(labels):
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes [xyxy-abs]
scale_coords(im[si].shape[1:], tbox, shape,
shapes[si][1]) # native-space labels
labelsn = torch.cat((labels[:, 0:1], tbox),
1) # native-space labels
predn = pred.clone()
if len(pred):
# Predictions
# native-space pred
# bbox was predicted on rescaled image (i.e. multiple of 32)
# now we scale the predicted bbox to match the orignal image shape
scale_coords(im[si].shape[1:], predn[:, :4], shape,
shapes[si][1])
img_res = evaluate_single_img(labelsn, predn, iouv)
if img_res is not None:
stats.append(img_res)
append_augmented_img_res(augmented_res, path, metadata, labelsn,
predn, img_res)
# Plot prediction examples
if batch_i < 3:
plot_batch_predictions(batch_i, save_dir, im, targets, out, paths)
augmented_res_df = pd.DataFrame.from_dict(augmented_res, orient="index")
augmented_res_df.sort_values(by=["run_id", "frame_index"], inplace=True)
augmented_res_df.reset_index(inplace=True, drop=True)
augmented_res_df.to_pickle(save_dir / f"{data.stem}-augmented-res.pkl")
results_by_slice(augmented_res_df, save_dir, conf_threshold)
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE,
n_init=cfg.DEEPSORT.N_INIT,
nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
device = select_device(device)
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
# make new output folder
# Load model
device = select_device(device)
model = DetectMultiBackend(yolo_model, device=device, dnn=dnn)
stride, names, pt, jit, _ = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source,