remove dask vestiges

rosalindfranklininstitute · Oct 20, 2023 · 480d557 · 480d557
1 parent ba1d721
commit 480d557
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 7 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -24,7 +24,6 @@ torch = ">=1.7.1"
 segmentation-models-pytorch = "^0.2.1"
 termplotlib = "^0.3.6"
 imagecodecs = "> 2022.2.22"
-dask ="*"
 
 [tool.poetry.dev-dependencies]
 opencv-python-headless = "4.8.0.74"

diff --git a/volume_segmantics/model/operations/vol_seg_2d_predictor.py b/volume_segmantics/model/operations/vol_seg_2d_predictor.py
@@ -11,8 +11,6 @@
 from volume_segmantics.data.dataloaders import get_2d_prediction_dataloader
 from volume_segmantics.model.model_2d import create_model_from_file
 from volume_segmantics.utilities.base_data_utils import Axis
-from dask import array as da
-
 
 class VolSeg2dPredictor:
     """Class that performs U-Net prediction operations. Does not interact with disk."""
@@ -186,7 +184,6 @@ def _predict_3_ways_max_probs(self, data_vol):
         )
         logging.info("Merging max of XY and ZX volumes with ZY volume.")
         self._merge_vols_in_mem(prob_container, label_container)
-        logging.debug("Nor using dask")
         return label_container[0], prob_container[0]
 
 
@@ -221,14 +218,12 @@ def _predict_12_ways_max_probs(self, data_vol):
             )
             self._merge_vols_in_mem(prob_container, label_container)
 
-        logging.debug("Not using dask")
         return label_container[0], prob_container[0]
 
     def _predict_single_axis_to_one_hot(self, data_vol, axis=Axis.Z):
         prediction, _ = self._predict_single_axis(data_vol, axis=axis)
         return utils.one_hot_encode_array(prediction, self.num_labels)
 
-    #TODO: implement dask for large volumes that fail to predict
     def _predict_3_ways_one_hot(self, data_vol):
         logging.debug("_predict_3_ways_one_hot()")
         one_hot_out = self._predict_single_axis_to_one_hot(data_vol)

diff --git a/volume_segmantics/scripts/predict_2d_model.py b/volume_segmantics/scripts/predict_2d_model.py
@@ -36,7 +36,7 @@ def main():
     settings = get_settings_data(settings_path)
     # Create prediction manager and predict
     #pred_manager = VolSeg2DPredictionManager(model_file_path, data_vol_path, settings)
-    pred_manager = VolSeg2DPredictionManager(model_file_path, data_vol_path, settings, True) #version that uses dask
+    pred_manager = VolSeg2DPredictionManager(model_file_path, data_vol_path, settings)
     pred_manager.predict_volume_to_path(output_path)
 
     logging.info(f"Prediction completed, saved to file: {output_path}")