diff --git a/utils/datasets.py b/utils/datasets.py index 29ee4b051e85..38ddc0af4b8b 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -107,7 +107,7 @@ def __iter__(self): class LoadImages: # for inference - def __init__(self, path, img_size=640): + def __init__(self, path, img_size=640, auto=True): p = str(Path(path)) # os-agnostic p = os.path.abspath(p) # absolute path if '*' in p: @@ -128,6 +128,7 @@ def __init__(self, path, img_size=640): self.nf = ni + nv # number of files self.video_flag = [False] * ni + [True] * nv self.mode = 'images' + self.auto = auto if any(videos): self.new_video(videos[0]) # new video else: @@ -169,7 +170,7 @@ def __next__(self): print('image %g/%g %s: ' % (self.count, self.nf, path), end='') # Padded resize - img = letterbox(img0, new_shape=self.img_size)[0] + img = letterbox(img0, new_shape=self.img_size, auto=self.auto)[0] # Convert img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416 @@ -253,7 +254,7 @@ def __len__(self): class LoadStreams: # multiple IP or RTSP cameras - def __init__(self, sources='streams.txt', img_size=640): + def __init__(self, sources='streams.txt', img_size=640, auto=True): self.mode = 'images' self.img_size = img_size @@ -266,6 +267,7 @@ def __init__(self, sources='streams.txt', img_size=640): n = len(sources) self.imgs = [None] * n self.sources = sources + self.auto = auto for i, s in enumerate(sources): # Start the thread to read frames from the video stream print('%g/%g: %s... ' % (i + 1, n, s), end='') @@ -281,7 +283,7 @@ def __init__(self, sources='streams.txt', img_size=640): print('') # newline # check for common shapes - s = np.stack([letterbox(x, new_shape=self.img_size)[0].shape for x in self.imgs], 0) # inference shapes + s = np.stack([letterbox(x, new_shape=self.img_size, auto=self.auto)[0].shape for x in self.imgs], 0) # inference shapes self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal if not self.rect: print('WARNING: Different stream shapes detected. For optimal performance supply similarly-shaped streams.') @@ -310,7 +312,7 @@ def __next__(self): raise StopIteration # Letterbox - img = [letterbox(x, new_shape=self.img_size, auto=self.rect)[0] for x in img0] + img = [letterbox(x, new_shape=self.img_size, auto=self.rect and self.auto)[0] for x in img0] # Stack img = np.stack(img, 0)