added changes to use grayscale images during training

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
 

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