From 833d8646893d956274ad1031d9d250a7608b061b Mon Sep 17 00:00:00 2001 From: Clea Parcerisas Date: Mon, 21 Aug 2023 10:51:22 +0200 Subject: [PATCH] added changes to use grayscale images during training --- ultralytics/cfg/models/v8/yolov8.yaml | 1 + ultralytics/data/augment.py | 10 ++++++++-- ultralytics/data/base.py | 2 +- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/ultralytics/cfg/models/v8/yolov8.yaml b/ultralytics/cfg/models/v8/yolov8.yaml index 2255450..6238958 100644 --- a/ultralytics/cfg/models/v8/yolov8.yaml +++ b/ultralytics/cfg/models/v8/yolov8.yaml @@ -3,6 +3,7 @@ # Parameters nc: 80 # number of classes +ch: 1 # number of channels scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' # [depth, width, max_channels] n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPs diff --git a/ultralytics/data/augment.py b/ultralytics/data/augment.py index 39866ae..f0e5e44 100644 --- a/ultralytics/data/augment.py +++ b/ultralytics/data/augment.py @@ -741,9 +741,15 @@ class Format: def _format_img(self, img): """Format the image for YOLOv5 from Numpy array to PyTorch tensor.""" + # if len(img.shape) < 3: + # img = np.expand_dims(img, -1) + # img = np.ascontiguousarray(img.transpose(2, 0, 1)[::-1]) + # img = torch.from_numpy(img) + # return img + if len(img.shape) < 3: - img = np.expand_dims(img, -1) - img = np.ascontiguousarray(img.transpose(2, 0, 1)[::-1]) + img = img.reshape([1, *img.shape]) + img = np.ascontiguousarray(img) img = torch.from_numpy(img) return img diff --git a/ultralytics/data/base.py b/ultralytics/data/base.py index bfc3cc1..a506954 100644 --- a/ultralytics/data/base.py +++ b/ultralytics/data/base.py @@ -148,7 +148,7 @@ class BaseDataset(Dataset): if fn.exists(): # load npy im = np.load(fn) else: # read image - im = cv2.imread(f) # BGR + im = cv2.imread(f, cv2.IMREAD_GRAYSCALE) # BGR if im is None: raise FileNotFoundError(f'Image Not Found {f}') h0, w0 = im.shape[:2] # orig hw