diff --git a/models/common.py b/models/common.py index 9911b207d060..6b8b4e4cb42f 100644 --- a/models/common.py +++ b/models/common.py @@ -259,7 +259,7 @@ def forward(self, imgs, size=640, augment=False, profile=False): files.append(Path(f).with_suffix('.jpg').name) if im.shape[0] < 5: # image in CHW im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1) - im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input + im = im[..., :3] if im.ndim == 3 else np.tile(im[..., None], 3) # enforce 3ch input s = im.shape[:2] # HWC shape0.append(s) # image shape g = (size / max(s)) # gain diff --git a/utils/datasets.py b/utils/datasets.py index 5c76a908c559..5a3fbefa28b7 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -218,7 +218,7 @@ def __next__(self): img = letterbox(img0, self.img_size, stride=self.stride)[0] # Convert - img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB and HWC to CHW + img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB img = np.ascontiguousarray(img) return path, img, img0, self.cap @@ -264,7 +264,7 @@ def __next__(self): img = letterbox(img0, self.img_size, stride=self.stride)[0] # Convert - img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB and HWC to CHW + img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB img = np.ascontiguousarray(img) return img_path, img, img0, None @@ -345,7 +345,7 @@ def __next__(self): img = np.stack(img, 0) # Convert - img = img[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB and BHWC to BCHW + img = img[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW img = np.ascontiguousarray(img) return self.sources, img, img0, None @@ -526,7 +526,6 @@ def __getitem__(self, index): if random.random() < hyp['mixup']: img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1))) - else: # Load image img, (h0, w0), (h, w) = load_image(self, index) @@ -579,7 +578,7 @@ def __getitem__(self, index): labels_out[:, 1:] = torch.from_numpy(labels) # Convert - img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3 x img_height x img_width + img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB img = np.ascontiguousarray(img) return torch.from_numpy(img), labels_out, self.img_files[index], shapes