From 63ee8ffad1fd65c1d435cc8aa4a6fb10e67c3bd0 Mon Sep 17 00:00:00 2001
From: alex-l-kong <31424707+alex-l-kong@users.noreply.github.com>
Date: Fri, 23 Feb 2024 16:14:51 -0800
Subject: [PATCH] Ensure nuclear masks don't appear in cell table with updated
 ezSeg pipeline if nuclear_counts=False (#1107)

---
 src/ark/segmentation/marker_quantification.py |  9 +++++++--
 .../marker_quantification_test.py             | 19 ++++++++++++-------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/ark/segmentation/marker_quantification.py b/src/ark/segmentation/marker_quantification.py
index 533de3e32..5f393ca35 100644
--- a/src/ark/segmentation/marker_quantification.py
+++ b/src/ark/segmentation/marker_quantification.py
@@ -529,10 +529,15 @@ def generate_cell_table(segmentation_dir, tiff_dir, img_sub_folder="TIFs",
         whole_cell_file = fov_name + '_whole_cell.tiff'
         nuclear_file = fov_name + '_nuclear.tiff'
 
-        # for each label given in the argument. read in that mask for the fov, and proceed with label and table appending
+        # for each label given in the argument, read in that mask for the fov, and proceed with
+        # label and table appending
         mask_files = io_utils.list_files(segmentation_dir, substrs=fov_name)
         mask_types = get_existing_mask_types(fov_names=fovs, mask_names=mask_files)
-        
+ 
+        # remove nuclear from mask_types if nuclear_counts False
+        if not nuclear_counts and "nuclear" in mask_types:
+            mask_types.remove("nuclear")
+
         for mask_type in mask_types:
             # load the segmentation labels in
             fov_mask_name = fov_name + '_' + mask_type + ".tiff"
diff --git a/tests/segmentation/marker_quantification_test.py b/tests/segmentation/marker_quantification_test.py
index f2281b2cd..c3424892f 100644
--- a/tests/segmentation/marker_quantification_test.py
+++ b/tests/segmentation/marker_quantification_test.py
@@ -735,18 +735,20 @@ def test_generate_cell_table_tree_loading():
             img_sub_folder=img_sub_folder, is_mibitiff=False, fovs=fovs_subset,
             nuclear_counts=True)
 
-        assert norm_data_nuc.shape[0] == norm_data_fov_sub.shape[0]
+        # setting nuclear_counts True generates data for both whole_cell and nuclear
+        # so there should be double the number of rows
+        assert norm_data_nuc.shape[0] == norm_data_fov_sub.shape[0] * 2
         assert norm_data_nuc.shape[1] == norm_data_fov_sub.shape[1] * 2
         misc_utils.verify_in_list(
             nuclear_col='nc_ratio',
             nuc_cell_table_cols=norm_data_nuc.columns.values
         )
 
-        assert arcsinh_data_nuc.shape[0] == arcsinh_data_fov_sub.shape[0]
+        assert arcsinh_data_nuc.shape[0] == arcsinh_data_fov_sub.shape[0] * 2
         assert arcsinh_data_nuc.shape[1] == norm_data_fov_sub.shape[1] * 2
         misc_utils.verify_in_list(
             nuclear_col='nc_ratio',
-            nuc_cell_table_cols=norm_data_nuc.columns.values
+            nuc_cell_table_cols=arcsinh_data_nuc.columns.values
         )
 
 
@@ -820,18 +822,20 @@ def test_generate_cell_table_mibitiff_loading():
             segmentation_dir=seg_dir, tiff_dir=tiff_dir, is_mibitiff=True, fovs=fovs_subset,
             nuclear_counts=True)
 
-        assert norm_data_nuc.shape[0] == norm_data_fov_sub.shape[0]
+        # setting nuclear_counts True generates data for both whole_cell and nuclear
+        # so there should be double the number of rows
+        assert norm_data_nuc.shape[0] == norm_data_fov_sub.shape[0] * 2
         assert norm_data_nuc.shape[1] == norm_data_fov_sub.shape[1] * 2
         misc_utils.verify_in_list(
             nuclear_col='nc_ratio',
             nuc_cell_table_cols=norm_data_nuc.columns.values
         )
 
-        assert arcsinh_data_nuc.shape[0] == arcsinh_data_fov_sub.shape[0]
+        assert arcsinh_data_nuc.shape[0] == arcsinh_data_fov_sub.shape[0] * 2
         assert arcsinh_data_nuc.shape[1] == norm_data_fov_sub.shape[1] * 2
         misc_utils.verify_in_list(
             nuclear_col='nc_ratio',
-            nuc_cell_table_cols=norm_data_nuc.columns.values
+            nuc_cell_table_cols=arcsinh_data_nuc.columns.values
         )
 
 
@@ -871,7 +875,8 @@ def test_generate_cell_table_extractions():
 
         default_norm_data, _ = marker_quantification.generate_cell_table(
             segmentation_dir=temp_dir, tiff_dir=tiff_dir,
-            img_sub_folder=img_sub_folder, is_mibitiff=False
+            img_sub_folder=img_sub_folder, is_mibitiff=False,
+            nuclear_counts=True
         )
 
         # verify total intensity extraction, same for whole_cell and nuclear mask types