update single_cell.py and tutorial_general_usage_test.py: add test fo…

…r get bin mask and label mask

livecellx/core/single_cell.py

@@ -258,7 +258,7 @@ def get_img(self):
     def get_mask(self, dtype=bool):
         if isinstance(self.mask_dataset, SingleImageDataset) or isinstance(self.mask_dataset, LiveCellImageDataset):
             try:
-                return self.mask_dataset.get_img_by_time()
+                return self.mask_dataset.get_img_by_time(time=self.timeframe)
             except Exception as e:
                 main_exception(
                     f"Error getting mask for single cell {self.id} at timeframe {self.timeframe} from presented mask dataset. Using contour mask instead."

tests/tutorial_general_usage_test.py

@@ -64,6 +64,22 @@ def test_loading_single_cells(self):
         assert len(single_cells_by_time[1]) == 14
         assert len(single_cells_by_time[2]) == 15
 
+    def test_sc_label_mask_retrieval(self):
+        sc = self.single_cells[0]
+        mask = sc.get_label_mask()
+
+        assert mask is not None
+        assert mask.shape == sc.get_img().shape
+        print("length of mask", len(np.unique(mask)))
+        assert len(np.unique(mask)) >= 3, "Mask should have at least three unique values"
+
+    def test_sc_bin_mask_retrieval(self):
+        sc = self.single_cells[0]
+        mask = sc.get_mask().astype(bool)
+        assert mask is not None
+        assert mask.shape == sc.get_img().shape
+        assert len(np.unique(mask)) == 2, "Mask should have exactly two unique values"
+
     def test_visualize_single_cell(self):
         sc = self.single_cells[0]
         assert isinstance(sc, SingleCellStatic)