Fixes #2

IAAPerspective corrupts the mask if not in {0;1} thanks @geohot
YassineYousfi · Dec 9, 2020 · 4d75c93 · 4d75c93
1 parent c1d5c11
commit 4d75c93
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/retriever.py b/retriever.py
@@ -78,7 +78,7 @@ def __getitem__(self, index: int):
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
         mask = cv2.imread(str(self.data_path/self.masks_folder/image_name), 0)
-        mask = np.stack([(mask == v) for v in self.class_values], axis=-1).astype('float')
+        mask = np.stack([(mask == v) for v in self.class_values], axis=-1).astype('uint8')
 
         if self.transforms:
             sample = self.transforms(image=image, mask=mask)
@@ -95,4 +95,4 @@ def __getitem__(self, index: int):
     def __len__(self) -> int:
         return len(self.image_names)
 
-
+