Skip to content

Commit

Permalink
added ability to read fits files via delayed dask functions
Browse files Browse the repository at this point in the history
  • Loading branch information
MikeSullivan7 committed Aug 8, 2024
1 parent e7b03b9 commit aeace47
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 8 deletions.
18 changes: 15 additions & 3 deletions mantidimaging/gui/windows/live_viewer/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from PyQt5.QtCore import QFileSystemWatcher, QObject, pyqtSignal, QTimer

import dask_image.imread
from astropy.io import fits

if TYPE_CHECKING:
from os import stat_result
Expand All @@ -25,10 +26,18 @@ class DaskImageDataStack:
"""
delayed_stack: dask.array.Array | None = None

def __init__(self, image_list: list[Image_Data]):
def __init__(self, image_list: list[Image_Data] | None):
if image_list:
if image_list[0].create_delayed_array:
self.delayed_stack = dask.array.concatenate([image_data.delayed_array for image_data in image_list])
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):
Expand Down Expand Up @@ -83,7 +92,10 @@ def image_modified_time(self) -> float:
return self._stat.st_mtime

def set_delayed_array(self) -> None:
self.delayed_array = dask_image.imread.imread(self.image_path)
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:
Expand Down
7 changes: 2 additions & 5 deletions mantidimaging/gui/windows/live_viewer/presenter.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,11 +109,8 @@ def load_image(image_data_obj: Image_Data) -> 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"]:
image_data = image_data_obj.delayed_array.compute()[0]
elif image_data_obj.image_path.suffix.lower() == ".fits":
with fits.open(image_data_obj.image_path.__str__()) as fit:
image_data = fit[0].data
if image_data_obj.image_path.suffix.lower() in [".tif", ".tiff", ".fits"]:
image_data = image_data_obj.delayed_array.compute()
return image_data

def update_image_modified(self, image_path: Path) -> None:
Expand Down

0 comments on commit aeace47

Please sign in to comment.