Source code for ewoksndreg.tasks.reg2d_preeval

from typing import Any
from typing import Dict
from typing import List
from typing import Optional

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

from ..evaluation.evaluation import pre_evaluation
from ..io.input_stack import InputDataType
from ..io.input_stack import input_context




class Inputs(BaseInputModel):
    image_stacks: InputDataType = Field(
        description="Image stacks as a dictionary of numpy arrays or list of HDF5 dataset URI's.",
        examples=[
            {
                "stack1": "/path/to/file.h5::/entry/process/results/parameters/Ca-K",
                "stack2": "/path/to/file.h5::/entry/process/results/parameters/Fe-K",
            },
            {"stack1": [[0, 0, 0], [1, 1, 1], [2, 2, 2]]},
        ],
    )
    reference_stack: Optional[str] = Field(
        default=None,
        description="Force select stack.",
        examples=["stack1", "stack2"],
    )
    output_configuration: Optional[Dict[str, Any]] = Field(
        default=None,
        description="Registration configuration parameters to be saved.",
        examples=[{"param1": 0, "param2": 1}],
    )
    skip: bool = Field(
        False,
        description="Do not rank the stacks.",
    )





class Outputs(BaseOutputModel):
    image_stacks: InputDataType = Field(
        description="Dictionary of image stacks in memory or URIs."
    )
    reference_stack: Optional[str] = Field(
        description="Transformations of all stacks is based on the image registration of this stack.",
        examples=["stack1", "stack2"],
    )
    ranked_stack_names: Optional[List[str]] = Field(
        description="`image_stacks` keys ordered from best to worst."
    )
    output_configuration: Optional[Dict[str, Any]] = Field(
        description="Registration configuration parameters to be saved.",
        examples=[{"param1": 0, "param2": 1}],
    )





class Reg2DPreEvaluation(Task, input_model=Inputs, output_model=Outputs):
    """Given several stacks of images requiring the same alignment,
    determine the stack which is most suitable to get a correct alignment.
    """



    def run(self):
        reference_stack = self.inputs.reference_stack
        image_stacks = self.inputs.image_stacks

        if self.inputs.skip:
            self.outputs.reference_stack = reference_stack
            self.outputs.ranked_stack_names = None
        elif reference_stack:
            if reference_stack not in image_stacks:
                raise ValueError(f"{reference_stack=} must be in {list(image_stacks)}")
            self.outputs.reference_stack = reference_stack
            self.outputs.ranked_stack_names = None
        else:
            with input_context(image_stacks) as stacks:
                pre_eval_rank = pre_evaluation(stacks)
            self.outputs.reference_stack = pre_eval_rank[0]
            self.outputs.ranked_stack_names = pre_eval_rank
        self.outputs.image_stacks = image_stacks

        output_configuration = self.get_input_value("output_configuration") or dict()
        output_configuration["reference_stack"] = reference_stack
        self.outputs.output_configuration = output_configuration