move all delayed array creation into DaskImageDataStack

Author: MikeSullivan7

mantidimaging/eyes_tests/live_viewer_window_test.py

@@ -7,7 +7,7 @@
 import numpy as np
 import os
 from mantidimaging.core.operations.loader import load_filter_packages
-from mantidimaging.gui.windows.live_viewer.model import Image_Data
+from mantidimaging.gui.windows.live_viewer.model import Image_Data, DaskImageDataStack
 from mantidimaging.test_helpers.unit_test_helper import FakeFSTestCase
 from pathlib import Path
 from mantidimaging.eyes_tests.base_eyes import BaseEyesTest
@@ -61,31 +61,34 @@ def test_live_view_opens_without_data(self, _mock_time, _mock_image_watcher):
     @mock.patch("time.time", return_value=4000.0)
     def test_live_view_opens_with_data(self, _mock_time, _mock_image_watcher, mock_load_image):
         file_list = self._make_simple_dir(self.live_directory)
-        image_list = [Image_Data(path, create_delayed_array=False) for path in file_list]
+        image_list = [Image_Data(path) for path in file_list]
+        dask_image_stack = DaskImageDataStack(image_list, create_delayed_array=False)
         mock_load_image.return_value = self._generate_image()
         self.imaging.show_live_viewer(self.live_directory)
-        self.imaging.live_viewer.presenter.model._handle_image_changed_in_list(image_list)
+        self.imaging.live_viewer.presenter.model._handle_image_changed_in_list(image_list, dask_image_stack)
         self.check_target(widget=self.imaging.live_viewer)
 
     @mock.patch('mantidimaging.gui.windows.live_viewer.presenter.LiveViewerWindowPresenter.load_image')
     @mock.patch('mantidimaging.gui.windows.live_viewer.model.ImageWatcher')
     @mock.patch("time.time", return_value=4000.0)
     def test_live_view_opens_with_bad_data(self, _mock_time, _mock_image_watcher, mock_load_image):
         file_list = self._make_simple_dir(self.live_directory)
-        image_list = [Image_Data(path, create_delayed_array=False) for path in file_list]
+        image_list = [Image_Data(path) for path in file_list]
+        dask_image_stack = DaskImageDataStack(image_list, create_delayed_array=False)
         mock_load_image.side_effect = ValueError
         self.imaging.show_live_viewer(self.live_directory)
-        self.imaging.live_viewer.presenter.model._handle_image_changed_in_list(image_list)
+        self.imaging.live_viewer.presenter.model._handle_image_changed_in_list(image_list, dask_image_stack)
         self.check_target(widget=self.imaging.live_viewer)
 
     @mock.patch('mantidimaging.gui.windows.live_viewer.presenter.LiveViewerWindowPresenter.load_image')
     @mock.patch('mantidimaging.gui.windows.live_viewer.model.ImageWatcher')
     @mock.patch("time.time", return_value=4000.0)
     def test_rotate_operation_rotates_image(self, _mock_time, _mock_image_watcher, mock_load_image):
         file_list = self._make_simple_dir(self.live_directory)
-        image_list = [Image_Data(path, create_delayed_array=False) for path in file_list]
+        image_list = [Image_Data(path) for path in file_list]
+        dask_image_stack = DaskImageDataStack(image_list, create_delayed_array=False)
         mock_load_image.return_value = self._generate_image()
         self.imaging.show_live_viewer(self.live_directory)
-        self.imaging.live_viewer.presenter.model._handle_image_changed_in_list(image_list)
+        self.imaging.live_viewer.presenter.model._handle_image_changed_in_list(image_list, dask_image_stack)
         self.imaging.live_viewer.rotate_angles_group.actions()[1].trigger()
         self.check_target(widget=self.imaging.live_viewer)

mantidimaging/gui/windows/live_viewer/model.py

@@ -24,25 +24,41 @@ class DaskImageDataStack:
     """
     A Dask Image Data Stack Class to hold a delayed array of all the images in the Live Viewer Path
     """
-    delayed_stack: dask.array.Array | None = None
+    delayed_stack: dask.array.Array
+    image_list: list[Image_Data]
 
-    def __init__(self, image_list: list[Image_Data] | None):
-        if image_list:
-            if image_list[0].create_delayed_array:
+    def __init__(self, image_list: list[Image_Data], create_delayed_array: bool = True):
+        self.image_list = image_list
+
+        if image_list and create_delayed_array:
+            arrays = self.get_delayed_arrays()
+            if arrays:
                 if image_list[0].image_path.suffix.lower() in [".tif", ".tiff"]:
-                    arrays = [image_data.delayed_array for image_data in image_list]
                     self.delayed_stack = dask.array.stack(dask.array.array(arrays))
                 elif image_list[0].image_path.suffix.lower() in [".fits"]:
                     with fits.open(image_list[0].image_path.__str__()) as fit:
                         sample = fit[0].data
-                    arrays = [image_data.delayed_array for image_data in image_list]
                     lazy_arrays = [dask.array.from_delayed(x, shape=sample.shape, dtype=sample.dtype) for x in arrays]
                     self.delayed_stack = dask.array.stack(lazy_arrays)
 
     @property
     def shape(self):
         return self.delayed_stack.shape
 
+    def get_delayed_arrays(self) -> list[dask.array.Array] | None:
+        if self.image_list[0].image_path.suffix.lower() in [".tif", ".tiff"]:
+            return [dask_image.imread.imread(image_data.image_path)[0] for image_data in self.image_list]
+        elif self.image_list[0].image_path.suffix.lower() == ".fits":
+            return [dask.delayed(fits.open)(image_data.image_path)[0].data for image_data in self.image_list]
+        else:
+            return None
+
+    def get_delayed_image(self, index) -> dask.array.Array:
+        return self.delayed_stack[index]
+
+    def get_image_data(self, index) -> Image_Data:
+        return self.image_list[index]
+
 
 class Image_Data:
     """
@@ -66,7 +82,7 @@ class Image_Data:
     delayed_array: dask.array.Array
     create_delayed_array: bool
 
-    def __init__(self, image_path: Path, create_delayed_array: bool = True):
+    def __init__(self, image_path: Path):
         """
         Constructor for Image_Data class.
 
@@ -78,9 +94,6 @@ def __init__(self, image_path: Path, create_delayed_array: bool = True):
         self.image_path = image_path
         self.image_name = image_path.name
         self._stat = image_path.stat()
-        self.create_delayed_array = create_delayed_array
-        if self.create_delayed_array:
-            self.set_delayed_array()
 
     @property
     def stat(self) -> stat_result:
@@ -96,12 +109,6 @@ def image_modified_time_stamp(self) -> str:
         """Return the image modified time as a string"""
         return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.image_modified_time))
 
-    def set_delayed_array(self) -> None:
-        if self.image_path.suffix.lower() in [".tif", ".tiff"]:
-            self.delayed_array = dask_image.imread.imread(self.image_path)[0]
-        elif self.image_path.suffix.lower() == ".fits":
-            self.delayed_array = dask.delayed(fits.open)(self.image_path)[0].data
-
 
 class SubDirectory:
 
@@ -145,7 +152,7 @@ def __init__(self, presenter: LiveViewerWindowPresenter):
         self._dataset_path: Path | None = None
         self.image_watcher: ImageWatcher | None = None
         self.images: list[Image_Data] = []
-        self.image_stack: DaskImageDataStack | None
+        self.image_stack: DaskImageDataStack
 
     @property
     def path(self) -> Path | None:
@@ -159,9 +166,8 @@ def path(self, path: Path) -> None:
         self.image_watcher.recent_image_changed.connect(self.handle_image_modified)
         self.image_watcher._handle_notified_of_directry_change(str(path))
 
-    def _handle_image_changed_in_list(self,
-                                      image_files: list[Image_Data],
-                                      dask_image_stack: DaskImageDataStack | None = None) -> None:
+    def _handle_image_changed_in_list(self, image_files: list[Image_Data],
+                                      dask_image_stack: DaskImageDataStack) -> None:
         """
         Handle an image changed event. Update the image in the view.
         This method is called when the image_watcher detects a change
@@ -243,7 +249,7 @@ def find_images(self, directory: Path) -> list[Image_Data]:
         for file_path in directory.iterdir():
             if self._is_image_file(file_path.name):
                 try:
-                    image_obj = Image_Data(file_path, create_delayed_array=self.create_delayed_array)
+                    image_obj = Image_Data(file_path)
                     image_files.append(image_obj)
                 except FileNotFoundError:
                     continue
@@ -315,7 +321,7 @@ def _handle_directory_change(self) -> None:
                 break
 
         images = self.sort_images_by_modified_time(images)
-        dask_image_stack = DaskImageDataStack(images)
+        dask_image_stack = DaskImageDataStack(images, create_delayed_array=self.create_delayed_array)
         self.update_recent_watcher(images[-1:])
         self.image_changed.emit(images, dask_image_stack)
 

mantidimaging/gui/windows/live_viewer/presenter.py

@@ -7,6 +7,7 @@
 from collections.abc import Callable
 from logging import getLogger
 import numpy as np
+import dask.array
 
 from imagecodecs._deflate import DeflateError
 
@@ -40,6 +41,8 @@ def __init__(self, view: LiveViewerWindowView, main_window: MainWindowView):
         self.main_window = main_window
         self.model = LiveViewerWindowModel(self)
         self.selected_image: Image_Data | None = None
+        self.selected_delayed_image: dask.array.Array
+
         self.filters = {f.filter_name: f for f in load_filter_packages()}
 
     def close(self) -> None:
@@ -73,20 +76,21 @@ def update_image_list(self, images_list: list[Image_Data]) -> None:
         self.view.set_image_index(len(images_list) - 1)
 
     def select_image(self, index: int) -> None:
-        if not self.model.images:
+        self.selected_image = self.model.image_stack.get_image_data(index)
+        self.selected_delayed_image = self.model.image_stack.get_delayed_image(index)
+        if not self.selected_image:
             return
-        self.selected_image = self.model.images[index]
         image_timestamp = self.selected_image.image_modified_time_stamp
         self.view.label_active_filename.setText(f"{self.selected_image.image_name} - {image_timestamp}")
 
-        self.display_image(self.selected_image)
+        self.display_image(self.selected_image, self.selected_delayed_image)
 
-    def display_image(self, image_data_obj: Image_Data) -> None:
+    def display_image(self, image_data_obj: Image_Data, delayed_image: dask.array.Array) -> None:
         """
         Display image in the view after validating contents
         """
         try:
-            image_data = self.load_image(image_data_obj)
+            image_data = self.load_image(delayed_image)
         except (OSError, KeyError, ValueError, DeflateError) as error:
             message = f"{type(error).__name__} reading image: {image_data_obj.image_path}: {error}"
             logger.error(message)
@@ -104,28 +108,27 @@ def display_image(self, image_data_obj: Image_Data) -> None:
         self.view.live_viewer.show_error(None)
 
     @staticmethod
-    def load_image(image_data_obj: Image_Data) -> np.ndarray:
+    def load_image(delayed_image: dask.array.Array) -> np.ndarray:
         """
         Load a .Tif, .Tiff or .Fits file only if it exists
         and returns as an ndarray
         """
-        if image_data_obj.image_path.suffix.lower() in [".tif", ".tiff", ".fits"]:
-            image_data = image_data_obj.delayed_array.compute()
+        image_data = delayed_image.compute()
         return image_data
 
     def update_image_modified(self, image_path: Path) -> None:
         """
         Update the displayed image when the file is modified
         """
         if self.selected_image and image_path == self.selected_image.image_path:
-            self.display_image(self.selected_image)
+            self.display_image(self.selected_image, self.selected_delayed_image)
 
     def update_image_operation(self) -> None:
         """
         Reload the current image if an operation has been performed on the current image
         """
         if self.selected_image is not None:
-            self.display_image(self.selected_image)
+            self.display_image(self.selected_image, self.selected_delayed_image)
 
     def convert_image_to_imagestack(self, image_data) -> ImageStack:
         """