From 7a9454fa6f2bfc990c5cace0f935b208a69edca4 Mon Sep 17 00:00:00 2001
From: Gouvernathor <44340603+Gouvernathor@users.noreply.github.com>
Date: Thu, 14 Apr 2022 15:17:59 +0200
Subject: [PATCH 1/3] Allow for full-RGB mask channels

---
 utils/data_loading.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/utils/data_loading.py b/utils/data_loading.py
index 24893b02db..7091fabb60 100644
--- a/utils/data_loading.py
+++ b/utils/data_loading.py
@@ -33,7 +33,10 @@ def preprocess(pil_img, scale, is_mask):
         pil_img = pil_img.resize((newW, newH), resample=Image.NEAREST if is_mask else Image.BICUBIC)
         img_ndarray = np.asarray(pil_img)
 
-        if not is_mask:
+        if is_mask:
+            if img_ndarray.ndim > 2:
+                img_ndarray = img_ndarray.sum(axis=-1, dtype=img_ndarray.dtype)
+        else:
             if img_ndarray.ndim == 2:
                 img_ndarray = img_ndarray[np.newaxis, ...]
             else:

From 278483977d3949f8767f44e61727e5b28f0c7a1b Mon Sep 17 00:00:00 2001
From: Gouvernathor <44340603+Gouvernathor@users.noreply.github.com>
Date: Thu, 14 Apr 2022 17:58:06 +0200
Subject: [PATCH 2/3] raise appropriate exception

---
 utils/data_loading.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/utils/data_loading.py b/utils/data_loading.py
index 7091fabb60..1e18817681 100644
--- a/utils/data_loading.py
+++ b/utils/data_loading.py
@@ -35,7 +35,8 @@ def preprocess(pil_img, scale, is_mask):
 
         if is_mask:
             if img_ndarray.ndim > 2:
-                img_ndarray = img_ndarray.sum(axis=-1, dtype=img_ndarray.dtype)
+                # customize this function if you want it to support RGB masks
+                raise RuntimeError("Only black-and-white images are supported as masks.")
         else:
             if img_ndarray.ndim == 2:
                 img_ndarray = img_ndarray[np.newaxis, ...]

From bc310abad99233e7e685d496a9cea3a978fe5171 Mon Sep 17 00:00:00 2001
From: Gouvernathor <44340603+Gouvernathor@users.noreply.github.com>
Date: Fri, 15 Apr 2022 04:31:24 +0200
Subject: [PATCH 3/3] Add warning in the readme

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 350fac9418..be82acc678 100644
--- a/README.md
+++ b/README.md
@@ -177,7 +177,7 @@ bash scripts/download_data.sh
 
 The input images and target masks should be in the `data/imgs` and `data/masks` folders respectively (note that the `imgs` and `masks` folder should not contain any sub-folder or any other files, due to the greedy data-loader). For Carvana, images are RGB and masks are black and white.
 
-You can use your own dataset as long as you make sure it is loaded properly in `utils/data_loading.py`.
+You can use your own dataset as long as you make sure it is loaded properly in `utils/data_loading.py`. For example, loading RGB-coded black and white masks will not work with the data-loading code as-is.
 
 
 ---