From 4d75c933542a21c456e284dfd6179df021b207b1 Mon Sep 17 00:00:00 2001 From: YassineYousfi Date: Tue, 8 Dec 2020 19:15:08 -0500 Subject: [PATCH] Fixes #2 IAAPerspective corrupts the mask if not in {0;1} thanks @geohot --- retriever.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/retriever.py b/retriever.py index 5efe907..8d2317c 100644 --- 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) - \ No newline at end of file +