Added sub_dir option for save_fov_images (#580)

* added sub_dir option for save_fov_imagse

* Update ark/utils/data_utils.py

Co-authored-by: Noah F. Greenwald <noahfgreenwald@gmail.com>

* changed test sub_dir directory

* changed the correct sub_dir

Co-authored-by: Noah F. Greenwald <noahfgreenwald@gmail.com>

ark/utils/data_utils.py

@@ -10,7 +10,7 @@
 from ark.utils.misc_utils import verify_in_list
 
 
-def save_fov_images(fovs, data_dir, img_xr, name_suffix=''):
+def save_fov_images(fovs, data_dir, img_xr, sub_dir=None, name_suffix=''):
     """Given an xarray of images per fov, saves each image separately
 
     Args:
@@ -20,8 +20,12 @@ def save_fov_images(fovs, data_dir, img_xr, name_suffix=''):
             The directory to save the images
         img_xr (xarray.DataArray):
             The array of images per fov
+        sub_dir (Optional[str]):
+            The subdirectory to save the images in. If specified images are saved to
+            "data_dir/sub_dir". If `sub_dir = None` the images are saved to "data_dir". Defaults
+            to None.
         name_suffix (str):
-            Specify what to append at the end of every fov
+            Specify what to append at the end of every fov.
     """
 
     if not os.path.exists(data_dir):
@@ -33,6 +37,15 @@ def save_fov_images(fovs, data_dir, img_xr, name_suffix=''):
         img_xr_fovs=img_xr.fovs.values
     )
 
+    if sub_dir is not None:
+        # Save the fovs in the directory `data_dir/sub_dir/`
+        save_dir = os.path.join(data_dir, sub_dir)
+        if not os.path.exists(save_dir):
+            os.makedirs(save_dir)
+    else:
+        # Save the fovs in the directory `data_dir`
+        save_dir = data_dir
+
     for fov in fovs:
         # retrieve the image for the fov
         fov_img_data = img_xr.loc[fov, ...].values
@@ -41,7 +54,7 @@ def save_fov_images(fovs, data_dir, img_xr, name_suffix=''):
         fov_file = fov + name_suffix + '.tiff'
 
         # save the image to data_dir
-        io.imsave(os.path.join(data_dir, fov_file), fov_img_data, check_contrast=False)
+        io.imsave(os.path.join(save_dir, fov_file), fov_img_data, check_contrast=False)
 
 
 def label_cells_by_cluster(fovs, all_data, label_maps, fov_col=settings.FOV_ID,

ark/utils/data_utils_test.py

@@ -50,12 +50,21 @@ def test_save_fov_images():
         for fov in fovs_sub:
             assert os.path.exists(os.path.join(temp_dir, fov + '.tiff'))
 
-    # test 3: name suffix provided
+    # test 3: name suffix provided, along with sub_dir provided
     with tempfile.TemporaryDirectory() as temp_dir:
-        data_utils.save_fov_images(fovs_sub, temp_dir, sample_img_xr, name_suffix='_test')
+        data_utils.save_fov_images(fovs_sub, temp_dir, sample_img_xr, sub_dir="sub_directory",
+                                   name_suffix='_test')
 
         for fov in fovs_sub:
-            assert os.path.exists(os.path.join(temp_dir, fov + '_test.tiff'))
+            assert os.path.exists(os.path.join(os.path.join(temp_dir, 'sub_directory'),
+                                               fov + '_test.tiff'))
+
+    # test 4: name suffix not provided, sub_dir provided
+    with tempfile.TemporaryDirectory() as temp_dir:
+        data_utils.save_fov_images(fovs_sub, temp_dir, sample_img_xr, sub_dir="sub_directory")
+        save_dir = os.path.join(temp_dir, "sub_directory")
+        for fov in fovs_sub:
+            assert os.path.exists(os.path.join(save_dir, fov + ".tiff"))
 
 
 def test_generate_deepcell_input():

templates_ark/example_cell_clustering.ipynb

@@ -684,6 +684,7 @@
     "        cell_fovs,\n",
     "        os.path.join(base_dir, cell_output_dir),\n",
     "        cell_cluster_masks,\n",
+    "        sub_dir='cell_masks',\n",
     "        name_suffix='_cell_mask'\n",
     "    )"
    ]

templates_ark/example_pixel_clustering.ipynb

@@ -660,6 +660,7 @@
     "        pixel_fovs,\n",
     "        os.path.join(base_dir, pixel_output_dir),\n",
     "        pixel_cluster_masks,\n",
+    "        sub_dir='pixel_masks',\n",
     "        name_suffix='_pixel_mask'\n",
     "    )"
    ]