diff --git a/mantidimaging/core/data/dataset.py b/mantidimaging/core/data/dataset.py index 8d71c6ab082..993c78a4426 100644 --- a/mantidimaging/core/data/dataset.py +++ b/mantidimaging/core/data/dataset.py @@ -2,7 +2,6 @@ # SPDX - License - Identifier: GPL-3.0-or-later from __future__ import annotations import uuid -from dataclasses import dataclass import numpy as np @@ -25,7 +24,17 @@ def remove_nones(image_stacks: list[ImageStack | None]) -> list[ImageStack]: class BaseDataset: - def __init__(self, *, name: str = "", stacks: list[ImageStack] | None = None) -> None: + def __init__( + self, + *, + name: str = "", + stacks: list[ImageStack] | None = None, + sample: ImageStack | None = None, + flat_before: ImageStack | None = None, + flat_after: ImageStack | None = None, + dark_before: ImageStack | None = None, + dark_after: ImageStack | None = None, + ) -> None: self._id: uuid.UUID = uuid.uuid4() self.name = name @@ -34,6 +43,15 @@ def __init__(self, *, name: str = "", stacks: list[ImageStack] | None = None) -> stacks = [] if stacks is None else stacks self._stacks: list[ImageStack] = stacks + self.sample = sample + self.flat_before = flat_before + self.flat_after = flat_after + self.dark_before = dark_before + self.dark_after = dark_after + + if self.name == "" and sample is not None: + self.name = sample.name + @property def id(self) -> uuid.UUID: return self._id @@ -89,31 +107,7 @@ class MixedDataset(BaseDataset): pass -@dataclass class StrictDataset(BaseDataset): - sample: ImageStack - flat_before: ImageStack | None = None - flat_after: ImageStack | None = None - dark_before: ImageStack | None = None - dark_after: ImageStack | None = None - - def __init__(self, - *, - sample: ImageStack, - flat_before: ImageStack | None = None, - flat_after: ImageStack | None = None, - dark_before: ImageStack | None = None, - dark_after: ImageStack | None = None, - name: str = ""): - super().__init__(name=name) - self.sample = sample - self.flat_before = flat_before - self.flat_after = flat_after - self.dark_before = dark_before - self.dark_after = dark_after - - if self.name == "": - self.name = sample.name @property def all(self) -> list[ImageStack]: @@ -165,6 +159,8 @@ def proj180deg(self) -> ImageStack | None: @proj180deg.setter def proj180deg(self, proj180deg: ImageStack | None) -> None: + if self.sample is None: + raise RuntimeError("Can't set a 180 projection without a sample") self.sample.proj180deg = proj180deg def delete_stack(self, images_id: uuid.UUID) -> None: @@ -179,6 +175,7 @@ def delete_stack(self, images_id: uuid.UUID) -> None: elif isinstance(self.dark_after, ImageStack) and self.dark_after.id == images_id: self.dark_after = None elif isinstance(self.proj180deg, ImageStack) and self.proj180deg.id == images_id: + assert self.sample is not None self.sample.clear_proj180deg() elif isinstance(self.sinograms, ImageStack) and self.sinograms.id == images_id: self.sinograms = None @@ -221,7 +218,7 @@ def _get_stack_data_type(stack_id: uuid.UUID, dataset: BaseDataset) -> str: if stack_id in [stack.id for stack in dataset._stacks]: return "Images" if isinstance(dataset, StrictDataset): - if stack_id == dataset.sample.id: + if dataset.sample is not None and stack_id == dataset.sample.id: return "Sample" if dataset.flat_before is not None and stack_id == dataset.flat_before.id: return "Flat Before" diff --git a/mantidimaging/core/io/saver.py b/mantidimaging/core/io/saver.py index c7f790d52df..041394f75a5 100644 --- a/mantidimaging/core/io/saver.py +++ b/mantidimaging/core/io/saver.py @@ -247,6 +247,7 @@ def _nexus_save(nexus_file: h5py.File, dataset: StrictDataset, sample_name: str, data["image_key"] = detector["image_key"] for recon in dataset.recons: + assert dataset.sample is not None assert dataset.sample.filenames is not None _save_recon_to_nexus(nexus_file, recon, dataset.sample.filenames[0]) diff --git a/mantidimaging/gui/windows/main/model.py b/mantidimaging/gui/windows/main/model.py index 5f3bc7c3110..91479f9fa0e 100644 --- a/mantidimaging/gui/windows/main/model.py +++ b/mantidimaging/gui/windows/main/model.py @@ -185,8 +185,10 @@ def remove_container(self, container_id: uuid.UUID) -> list[uuid.UUID]: proj_180_id = None # If we're deleting a sample from a StrictDataset then any linked 180 projection will also be # deleted - if isinstance(dataset, StrictDataset) and dataset.proj180deg and dataset.sample.id == container_id: - proj_180_id = dataset.proj180deg.id + if isinstance(dataset, StrictDataset) and dataset.proj180deg: + assert dataset.sample is not None + if dataset.sample.id == container_id: + proj_180_id = dataset.proj180deg.id dataset.delete_stack(container_id) return [container_id, proj_180_id] if proj_180_id else [container_id] if container_id == dataset.recons.id: diff --git a/mantidimaging/gui/windows/main/presenter.py b/mantidimaging/gui/windows/main/presenter.py index 7c65df284fc..1c5e3296d0d 100644 --- a/mantidimaging/gui/windows/main/presenter.py +++ b/mantidimaging/gui/windows/main/presenter.py @@ -241,6 +241,7 @@ def add_alternative_180_if_required(self, dataset: StrictDataset) -> None: Checks if the dataset has a 180 projection and tries to find an alternative if one is missing. :param dataset: The loaded dataset. """ + assert dataset.sample is not None if dataset.sample.has_proj180deg() and dataset.sample.proj180deg.filenames: # type: ignore return else: @@ -260,6 +261,7 @@ def create_strict_dataset_stack_windows(self, dataset: StrictDataset) -> StackVi :param dataset: The loaded dataset. :return: The stack widget for the sample. """ + assert dataset.sample is not None sample_stack_vis = self._create_lone_stack_window(dataset.sample) self._tabify_stack_window(sample_stack_vis) @@ -357,6 +359,7 @@ def create_strict_dataset_tree_view_items(self, dataset: StrictDataset) -> None: Creates the tree view items for a strict dataset. :param dataset: The loaded dataset. """ + assert dataset.sample is not None dataset_tree_item = self.view.create_dataset_tree_widget_item(dataset.name, dataset.id) self.view.create_child_tree_item(dataset_tree_item, dataset.sample.id, "Projections")