Source code for ewoksndreg.tasks.example2d_stacks

from typing import Dict
from typing import List
from typing import Optional
from typing import Sequence

from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from ewokscore.task import Task
from pydantic import Field

from ..io import data_for_registration
from ..io.output_stack import OutputDataTypeForInput
from ..transformation.base import Transformation
from ..transformation.numpy_backend import NumpyHomography
from ..transformation.types import TransformationType

try:
    from ..transformation.simpleitk_backend import SimpleITKTransformation
except ImportError:
    SimpleITKTransformation = None




class Inputs(BaseInputModel):
    name: str = Field(description="Name of the image", examples=["astronaut"])
    transformation_type: TransformationType = Field(
        description="Transformation type", examples=["translation", "rigid"]
    )
    shape: Optional[Sequence[int]] = Field(
        default=None,
        description="Image shape",
        examples=[(200, 220)],
        min_length=2,
        max_length=2,
    )
    nimages: Optional[int] = Field(
        default=None, description="Number of images per stack"
    )
    nstacks: Optional[int] = Field(default=None, description="Number of image stacks")
    noise: Optional[int] = Field(
        default=None, description="Number of image stacks", examples=["s&p", "uniform"]
    )





class Outputs(BaseOutputModel):
    image_stacks: OutputDataTypeForInput = Field(
        description="Dictionary of image stacks in memory or URIs."
    )
    transformations: Dict[str, List[Transformation]] = Field(
        description="Transformation between the images of each stack."
    )





class Example2DStacks(Task, input_model=Inputs, output_model=Outputs):
    """Generate one or more stacks of transformed images to test registration methods."""



    def run(self):
        image = data_for_registration.generate_image(name=self.inputs.name)
        image_stacks, _, passive_matrices = data_for_registration.generate_image_stack(
            image,
            self.inputs.transformation_type,
            shape=self.inputs.shape,
            nimages=self.inputs.nimages,
        )
        image_stacks = data_for_registration.generate_image_stacks(
            image_stacks, self.inputs.nstacks, noise=self.inputs.noise
        )

        if self.inputs.transformation_type in ["displacement_field", "bspline"]:
            if SimpleITKTransformation is None:
                raise ValueError(
                    "displacement field transforms cannot be generated without SimpleITK"
                )
            transformations = [
                SimpleITKTransformation(displacement_field=d) for d in passive_matrices
            ]
        else:
            transformations = [NumpyHomography(M) for M in passive_matrices]

        self.outputs.image_stacks = image_stacks
        self.outputs.transformations = {name: transformations for name in image_stacks}