update sc_correction_dataset.py: define classes

xing-lab-pitt · Mar 22, 2024 · 5d98628 · 5d98628
1 parent d0ef61f
commit 5d98628
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/livecellx/model_zoo/segmentation/sc_correction_dataset.py b/livecellx/model_zoo/segmentation/sc_correction_dataset.py
@@ -34,6 +34,11 @@
 class CorrectSegNetDataset(torch.utils.data.Dataset):
     """Dataset for training CorrectSegNetDatasert"""
 
+    OVERSEG_ONEHOT = [1, 0, 0, 0]
+    UNDERSEG_ONEHOT = [0, 1, 0, 0]
+    DROPOUT_ONEHOT = [0, 0, 1, 0]
+    CORRECT_ONEHOT = [0, 0, 0, 1]
+
     def __init__(
         self,
         raw_img_paths: List[str],