Pipeline Components: LearnTransform and ApplyTransform

Makefile

@@ -54,7 +54,7 @@ help-unit:
 ### DOCS #####################################################
 #
 docs-%:
-	make -C docs $*
+	make -C docs clean && make -C docs $*
 
 help-docs:
 	$(call print_help, docs-TASK, alias for 'make TASK' in the docs subdirectory)

docs/source/api/image/index.rst

@@ -10,8 +10,5 @@ Image Manipulation
 .. toctree::
    filtering/index.rst
 
-.. toctree::
-   registration/index.rst
-
 .. toctree::
    segmentation/index.rst

docs/source/index.rst

@@ -59,7 +59,6 @@ To see the code or report a bug, please visit the `github repository
   <div class="col-md-3">
   <h2>Features</h2>
 
-* Registration: :ref:`API <registration>`
 * Filtering: :ref:`API <filtering>`
 * Spot-Finding: :ref:`API <detection>`
 * Decoding: :ref:`API <decoding>`

docs/source/usage/iss/iss_pipeline.py

@@ -1,7 +1,7 @@
 import os
 
 import starfish
-from starfish.image import Filter, Registration, Segmentation
+from starfish.image import ApplyTransform, Filter, LearnTransform, Segmentation
 from starfish.spots import SpotFinder, TargetAssignment
 from starfish.types import Axes
 
@@ -11,12 +11,12 @@
 def iss_pipeline(fov, codebook):
     primary_image = fov.get_image(starfish.FieldOfView.PRIMARY_IMAGES)
 
-    # register the raw images
-    registration = Registration.FourierShiftRegistration(
-        upsampling=1000,
-        reference_stack=fov.get_image('dots')
-    )
-    registered = registration.run(primary_image, in_place=False)
+    # register the raw image
+    learn_translation = LearnTransform.Translation(reference_stack=fov.get_image('dots'),
+                                                   axis=Axes.ROUND, upsampling=100)
+    transforms_list = learn_translation.run(primary_image.max_proj(Axes.CH, Axes.ZPLANE))
+    warp = ApplyTransform.Warp(transforms_list=transforms_list)
+    registered = warp.run(primary_image, in_place=False, verbose=True)
 
     # filter raw data
     masking_radius = 15

notebooks/ISS_Pipeline_-_Breast_-_1_FOV.ipynb

@@ -209,13 +209,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from starfish.image import Registration\n",
+    "from starfish.image import ApplyTransform, LearnTransform\n",
     "\n",
-    "registration = Registration.FourierShiftRegistration(\n",
-    "    upsampling=1000,\n",
-    "    reference_stack=dots,\n",
-    "    verbose=True)\n",
-    "registered_image = registration.run(primary_image, in_place=False)"
+    "learn_translation = LearnTransform.Translation(reference_stack=dots, axis=Axes.ROUND, upsampling=1000)\n",
+    "transforms_list = learn_translation.run(primary_image.max_proj(Axes.CH, Axes.ZPLANE))\n",
+    "warp = ApplyTransform.Warp(transforms_list=transforms_list)\n",
+    "registered_image = warp.run(primary_image, in_place=False, verbose=True)"
    ]
   },
   {

notebooks/assay_comparison.ipynb

@@ -482,4 +482,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

notebooks/py/ISS_Pipeline_-_Breast_-_1_FOV.py

@@ -129,13 +129,12 @@
 # EPY: END markdown
 
 # EPY: START code
-from starfish.image import Registration
+from starfish.image import ApplyTransform, LearnTransform
 
-registration = Registration.FourierShiftRegistration(
-    upsampling=1000,
-    reference_stack=dots,
-    verbose=True)
-registered_image = registration.run(primary_image, in_place=False)
+learn_translation = LearnTransform.Translation(reference_stack=dots, axis=Axes.ROUND, upsampling=1000)
+transforms_list = learn_translation.run(primary_image.max_proj(Axes.CH, Axes.ZPLANE))
+warp = ApplyTransform.Warp(transforms_list=transforms_list)
+registered_image = warp.run(primary_image, in_place=False, verbose=True)
 # EPY: END code
 
 # EPY: START markdown

starfish/image/__init__.py

@@ -1,3 +1,4 @@
+from ._apply_transform import ApplyTransform
 from ._filter import Filter
-from ._registration import Registration
+from ._learn_transform import LearnTransform
 from ._segmentation import Segmentation

starfish/image/_registration/__init__.py → starfish/image/_apply_transform/__init__.py

@@ -7,29 +7,28 @@
 import_all_submodules(__file__, __package__)
 
 
-class Registration(PipelineComponent):
+class ApplyTransform(PipelineComponent):
 
     @classmethod
     def _get_algorithm_base_class(cls) -> Type[AlgorithmBase]:
-        return _base.RegistrationAlgorithmBase
+        return _base.ApplyTransformBase
 
     @classmethod
-    def _cli_run(cls, ctx, instance):
+    def _cli_run(cls, ctx, instance, *args, **kwargs):
         output = ctx.obj["output"]
         stack = ctx.obj["stack"]
-        instance.run(stack)
-        stack.export(output)
+        transformed = instance.run(stack)
+        transformed.export(output)
 
     @staticmethod
-    @click.group("registration")
+    @click.group("apply_transform")
     @click.option("-i", "--input", type=click.Path(exists=True))
     @click.option("-o", "--output", required=True)
     @click.pass_context
     def _cli(ctx, input, output):
-        """translation correction of image stacks"""
-        print("Registering...")
+        print("Applying Transform to images...")
         ctx.obj = dict(
-            component=Registration,
+            component=ApplyTransform,
             input=input,
             output=output,
             stack=ImageStack.from_path_or_url(input),

starfish/image/_registration/_base.py → starfish/image/_apply_transform/_base.py

@@ -1,10 +1,8 @@
-from typing import Optional
-
 from starfish.imagestack.imagestack import ImageStack
 from starfish.pipeline.algorithmbase import AlgorithmBase
 
 
-class RegistrationAlgorithmBase(AlgorithmBase):
-    def run(self, stack) -> Optional[ImageStack]:
+class ApplyTransformBase(AlgorithmBase):
+    def run(self, stack) -> ImageStack:
         """Performs registration on the stack provided."""
         raise NotImplementedError()

starfish/image/_apply_transform/warp.py

@@ -0,0 +1,106 @@
+from copy import deepcopy
+from typing import Union
+
+import numpy as np
+import xarray as xr
+from skimage import transform
+from skimage.transform._geometric import GeometricTransform
+from tqdm import tqdm
+
+from starfish.config import StarfishConfig
+from starfish.image._apply_transform._base import ApplyTransformBase
+from starfish.image._learn_transform.transforms_list import TransformsList
+from starfish.imagestack.imagestack import ImageStack
+from starfish.types import Axes
+from starfish.util import click
+
+
+class Warp(ApplyTransformBase):
+    """Class that applies a list of arbitrary skimage GeometricTransforms to an ImageStack
+     using skimage.transform.warp"""
+
+    def __init__(self, transforms_list: Union[str, TransformsList]):
+        """
+        Parameters
+        ----------
+        transforms_list: TransformsList
+            A list of tuples that describe a subset of an ImageStack axis and a GeometricTransform
+            to apply to it.
+        """
+        if isinstance(transforms_list, TransformsList):
+            self.transforms_list = transforms_list
+        else:
+            self.transforms_list = TransformsList.from_json(filename=transforms_list)
+
+    def run(
+            self, stack: ImageStack,
+            in_place: bool=False, verbose: bool=False, **kwargs) -> ImageStack:
+        """Applies a list of transformations to an ImageStack
+
+        Parameters
+        ----------
+        stack : ImageStack
+            Stack to be transformed.
+        in_place : bool
+            if True, process ImageStack in-place, otherwise return a new stack
+        verbose : bool
+            if True, report on filtering progress (default = False)
+
+        Returns
+        -------
+        ImageStack :
+            If in-place is False, return the results of the transforms as a new stack.
+            Otherwise return the original stack.
+        """
+        if not in_place:
+            # create a copy of the ImageStack, call apply on that stack with in_place=True
+            image_stack = deepcopy(stack)
+            return self.run(
+                image_stack,
+                in_place=True,
+                **kwargs
+            )
+        if verbose and StarfishConfig().verbose:
+            self.transforms_list.transforms = tqdm(self.transforms_list.transforms)
+        all_axes = {Axes.ROUND, Axes.CH, Axes.ZPLANE}
+        for selector, transformation_object in self.transforms_list.transforms:
+            other_axes = all_axes - set(selector.keys())
+            # iterate through remaining axes
+            for axes in stack._iter_axes(other_axes):
+                # combine all axes data to select one tile
+                selector.update(axes)  # type: ignore
+                selected_image, _ = stack.get_slice(selector)
+                warped_image = warp(selected_image, transformation_object, **kwargs
+                                    ).astype(np.float32)
+                stack.set_slice(selector, warped_image)
+        return stack
+
+    @staticmethod
+    @click.command("Warp")
+    @click.option("--transformation-list", required=True, type=click.Path(exists=True),
+                  help="The list of transformations to apply to the ImageStack.")
+    @click.pass_context
+    def _cli(ctx, transformation_list):
+        ctx.obj["component"]._cli_run(ctx, Warp(transformation_list))
+
+
+def warp(image: Union[xr.DataArray, np.ndarray],
+         transformation_object: GeometricTransform,
+         **kwargs
+         ) -> np.ndarray:
+    """ Wrapper around skimage.transform.warp. Warps an image according to a
+    given coordinate transformation.
+
+    Parameters
+    ----------
+    image: np.ndarray
+        The image to be transformed
+    transformation_object: skimage.transform.GeometricTransform
+        The transformation object to apply.
+
+    Returns
+    -------
+    np.ndarray:
+        the warped image.
+    """
+    return transform.warp(image, transformation_object, **kwargs)

starfish/image/_learn_transform/__init__.py

@@ -0,0 +1,36 @@
+from typing import Type
+
+from starfish.imagestack.imagestack import ImageStack
+from starfish.pipeline import AlgorithmBase, import_all_submodules, PipelineComponent
+from starfish.util import click
+from . import _base
+import_all_submodules(__file__, __package__)
+
+
+class LearnTransform(PipelineComponent):
+
+    @classmethod
+    def _get_algorithm_base_class(cls) -> Type[AlgorithmBase]:
+        return _base.LearnTransformBase
+
+    @classmethod
+    def _cli_run(cls, ctx, instance, *args, **kwargs):
+        output = ctx.obj["output"]
+        stack = ctx.obj["stack"]
+        transformation_list = instance.run(stack)
+        transformation_list.to_json(output)
+
+    @staticmethod
+    @click.group("learn_transform")
+    @click.option("-i", "--input", type=click.Path(exists=True))
+    @click.option("-o", "--output", required=True)
+    @click.pass_context
+    def _cli(ctx, input, output):
+        """learn a set of transforms for an ImageStack."""
+        print("Learning Transforms for images...")
+        ctx.obj = dict(
+            component=LearnTransform,
+            input=input,
+            output=output,
+            stack=ImageStack.from_path_or_url(input),
+        )

starfish/image/_learn_transform/_base.py

@@ -0,0 +1,8 @@
+from starfish.image._learn_transform.transforms_list import TransformsList
+from starfish.pipeline.algorithmbase import AlgorithmBase
+
+
+class LearnTransformBase(AlgorithmBase):
+    def run(self, stack) -> TransformsList:
+        """Learns Transforms for a given stack."""
+        raise NotImplementedError()