From 82b97d5d4cb17165404dcab0b954852eb8b8a732 Mon Sep 17 00:00:00 2001 From: Clea Parcerisas Date: Thu, 24 Aug 2023 14:42:29 +0200 Subject: [PATCH] changed data/loaders and engine/predictor so they accept 1 channel images for predictions --- ultralytics/data/loaders.py | 2 +- ultralytics/engine/predictor.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/ultralytics/data/loaders.py b/ultralytics/data/loaders.py index 1534a50..0089944 100644 --- a/ultralytics/data/loaders.py +++ b/ultralytics/data/loaders.py @@ -249,7 +249,7 @@ class LoadImages: else: # Read image self.count += 1 - im0 = cv2.imread(path) # BGR + im0 = cv2.imread(path, cv2.IMREAD_GRAYSCALE) # BGR if im0 is None: raise FileNotFoundError(f'Image Not Found {path}') s = f'image {self.count}/{self.nf} {path}: ' diff --git a/ultralytics/engine/predictor.py b/ultralytics/engine/predictor.py index c42c1b5..3dc0191 100644 --- a/ultralytics/engine/predictor.py +++ b/ultralytics/engine/predictor.py @@ -122,7 +122,8 @@ class BasePredictor: not_tensor = not isinstance(im, torch.Tensor) if not_tensor: im = np.stack(self.pre_transform(im)) - im = im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) + im = np.expand_dims(im, -1) + im = im[..., ::].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) im = np.ascontiguousarray(im) # contiguous im = torch.from_numpy(im)