From ffdb58b0e07d964eb2d148a6814d22a4a26d47cc Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 2 Sep 2022 14:12:10 +0200
Subject: [PATCH] Move `cudnn.benchmarks(True)` to LoadStreams (#9258)
* Move cudnn.benchmarks(True) to LoadStreams
* Update dataloaders.py
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
* Move cudnn.benchmarks(True) to LoadStreams
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---
classify/predict.py | 2 --
detect.py | 2 --
utils/dataloaders.py | 54 ++++----------------------------------------
3 files changed, 4 insertions(+), 54 deletions(-)
diff --git a/classify/predict.py b/classify/predict.py
index 937704d0f..76115c750 100644
--- a/classify/predict.py
+++ b/classify/predict.py
@@ -31,7 +31,6 @@ import sys
from pathlib import Path
import torch
-import torch.backends.cudnn as cudnn
import torch.nn.functional as F
FILE = Path(__file__).resolve()
@@ -89,7 +88,6 @@ def run(
# 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, transforms=classify_transforms(imgsz[0]))
bs = len(dataset) # batch_size
else:
diff --git a/detect.py b/detect.py
index 60a821b59..cf75d0f11 100644
--- a/detect.py
+++ b/detect.py
@@ -31,7 +31,6 @@ import sys
from pathlib import Path
import torch
-import torch.backends.cudnn as cudnn
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
@@ -97,7 +96,6 @@ def run(
# 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:
diff --git a/utils/dataloaders.py b/utils/dataloaders.py
index 5f86f8378..38ae3399c 100755
--- a/utils/dataloaders.py
+++ b/utils/dataloaders.py
@@ -283,62 +283,17 @@ class LoadImages:
return self.nf # number of files
-class LoadWebcam: # for inference
- # YOLOv5 local webcam dataloader, i.e. `python detect.py --source 0`
- def __init__(self, pipe='0', img_size=640, stride=32):
- self.img_size = img_size
- self.stride = stride
- self.pipe = eval(pipe) if pipe.isnumeric() else pipe
- self.cap = cv2.VideoCapture(self.pipe) # video capture object
- self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 3) # set buffer size
-
- def __iter__(self):
- self.count = -1
- return self
-
- def __next__(self):
- self.count += 1
- if cv2.waitKey(1) == ord('q'): # q to quit
- self.cap.release()
- cv2.destroyAllWindows()
- raise StopIteration
-
- # Read frame
- ret_val, im0 = self.cap.read()
- im0 = cv2.flip(im0, 1) # flip left-right
-
- # Print
- assert ret_val, f'Camera Error {self.pipe}'
- img_path = 'webcam.jpg'
- s = f'webcam {self.count}: '
-
- # Process
- im = letterbox(im0, self.img_size, stride=self.stride)[0] # resize
- im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
- im = np.ascontiguousarray(im) # contiguous
-
- return img_path, im, im0, None, s
-
- def __len__(self):
- return 0
-
-
class LoadStreams:
# YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams`
def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True, transforms=None):
+ torch.backends.cudnn.benchmark = True # faster for fixed-size inference
self.mode = 'stream'
self.img_size = img_size
self.stride = stride
-
- if os.path.isfile(sources):
- with open(sources) as f:
- sources = [x.strip() for x in f.read().strip().splitlines() if len(x.strip())]
- else:
- sources = [sources]
-
+ sources = Path(sources).read_text().rsplit() if Path(sources).is_file() else [sources]
n = len(sources)
- self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
self.sources = [clean_str(x) for x in sources] # clean source names for later
+ self.imgs, self.fps, self.frames, self.threads = [None] * n, [0] * n, [0] * n, [None] * n
for i, s in enumerate(sources): # index, source
# Start thread to read frames from video stream
st = f'{i + 1}/{n}: {s}... '
@@ -377,8 +332,7 @@ class LoadStreams:
n, f, read = 0, self.frames[i], 1 # frame number, frame array, inference every 'read' frame
while cap.isOpened() and n < f:
n += 1
- # _, self.imgs[index] = cap.read()
- cap.grab()
+ cap.grab() # .read() = .grab() followed by .retrieve()
if n % read == 0:
success, im = cap.retrieve()
if success: