From eeca8b7ad9c11361fef0cd1a0f086cf8b76bd8cf Mon Sep 17 00:00:00 2001
From: "Hannes F. Kuchelmeister" <hannes@kuchelmeister.org>
Date: Wed, 14 Jul 2021 11:49:06 +0200
Subject: [PATCH] add type hints to registrations

---
 situr/registration/channel_registration.py     | 11 +++++++----
 situr/registration/registration.py             | 11 +++++++----
 situr/registration/round_registration.py       |  9 +++++----
 situr/transformation/__init__.py               |  1 +
 situr/transformation/channel_transformation.py |  4 +++-
 situr/transformation/round_transformation.py   |  3 ++-
 situr/transformation/transformation.py         |  2 ++
 7 files changed, 27 insertions(+), 14 deletions(-)
 create mode 100644 situr/transformation/transformation.py

diff --git a/situr/registration/channel_registration.py b/situr/registration/channel_registration.py
index 35fdd8a..8e356ba 100644
--- a/situr/registration/channel_registration.py
+++ b/situr/registration/channel_registration.py
@@ -1,11 +1,14 @@
-from situr.registration import Registration, FilterregRegistrationFunction
-from situr.transformation import ScaleRotateTranslateChannelTransform
+from situr.image.situ_image import SituImage
+from situr.transformation.channel_transformation import ChannelTransform
+from situr.registration import Registration, RegistrationFunction, FilterregRegistrationFunction
+from situr.transformation import ChannelTransform, ScaleRotateTranslateChannelTransform
 
 
 class ChannelRegistration(Registration):
-    def __init__(self, registration_function=FilterregRegistrationFunction(ScaleRotateTranslateChannelTransform)):
+    def __init__(self, registration_function: RegistrationFunction[ChannelTransform] = FilterregRegistrationFunction(ScaleRotateTranslateChannelTransform)):
         super().__init__(registration_function)
-    def do_channel_registration(self, situ_img, reference_channel=0):
+
+    def do_channel_registration(self, situ_img: SituImage, reference_channel: int = 0):
         # For each channel (except nucleus) compute transform compared to reference_channel
         # Add Channel transformation to Channel
         reference_peaks = situ_img.get_channel_peaks(reference_channel)
diff --git a/situr/registration/registration.py b/situr/registration/registration.py
index a5c4626..ae3737b 100644
--- a/situr/registration/registration.py
+++ b/situr/registration/registration.py
@@ -1,14 +1,17 @@
 import abc
+from situr.transformation.channel_transformation import ChannelTransform
+from situr.transformation.round_transformation import RoundTransform
 import open3d as o3
 from probreg import filterreg
+import numpy as np
 
 from situr.image import extend_dim
-
+from situr.transformation import Transform
 
 class RegistrationFunction:
     __metaclass__ = abc.ABCMeta
 
-    def __init__(self, transormation_type):
+    def __init__(self, transormation_type: Transform):
         self.transormation_type = transormation_type
 
     @abc.abstractmethod
@@ -17,7 +20,7 @@ class RegistrationFunction:
 
 
 class FilterregRegistrationFunction(RegistrationFunction):
-    def do_registration(self, data_peaks, reference_peaks):
+    def do_registration(self, data_peaks: np.ndarray, reference_peaks: np.ndarray) -> Transform:
         source = o3.geometry.PointCloud()
         source.points = o3.utility.Vector3dVector(extend_dim(data_peaks))
         target = o3.geometry.PointCloud()
@@ -30,5 +33,5 @@ class FilterregRegistrationFunction(RegistrationFunction):
 
 
 class Registration:
-    def __init__(self, registration_function):
+    def __init__(self, registration_function: RegistrationFunction):
         self.registration_function = registration_function
diff --git a/situr/registration/round_registration.py b/situr/registration/round_registration.py
index cf78f60..38cccb5 100644
--- a/situr/registration/round_registration.py
+++ b/situr/registration/round_registration.py
@@ -1,12 +1,13 @@
-from situr.registration import Registration, FilterregRegistrationFunction
-from situr.transformation import ScaleRotateTranslateRoundTransform
+from situr.registration import Registration, RegistrationFunction, FilterregRegistrationFunction
+from situr.transformation import RoundTransform, ScaleRotateTranslateRoundTransform
+from situr.image import Tile
 
 
 class RoundRegistration(Registration):
-    def __init__(self, registration_function=FilterregRegistrationFunction(ScaleRotateTranslateRoundTransform)):
+    def __init__(self, registration_function: RegistrationFunction[RoundTransform] = FilterregRegistrationFunction(ScaleRotateTranslateRoundTransform)):
         super().__init__(registration_function)
 
-    def do_round_registration(self, situ_tile, reference_round=0, reference_channel=0):
+    def do_round_registration(self, situ_tile: Tile, reference_round: int = 0, reference_channel: int = 0):
         """This method generates a round registration transformation for a tile and saves it in the tile. 
 
         Args:
diff --git a/situr/transformation/__init__.py b/situr/transformation/__init__.py
index 505f82e..899e19a 100644
--- a/situr/transformation/__init__.py
+++ b/situr/transformation/__init__.py
@@ -1,2 +1,3 @@
 from .channel_transformation import ChannelTransform, IdentityChannelTransform, ScaleRotateTranslateChannelTransform
 from .round_transformation import RoundTransform, IdentityRoundTransform, ScaleRotateTranslateRoundTransform
+from .transformation import Transform
diff --git a/situr/transformation/channel_transformation.py b/situr/transformation/channel_transformation.py
index a12fb19..7c9b783 100644
--- a/situr/transformation/channel_transformation.py
+++ b/situr/transformation/channel_transformation.py
@@ -2,8 +2,10 @@ import abc
 import numpy as np
 import scipy
 
+from situr.transformation import Transform
 
-class ChannelTransform:
+
+class ChannelTransform(Transform):
     __metaclass__ = abc.ABCMeta
 
     @abc.abstractmethod
diff --git a/situr/transformation/round_transformation.py b/situr/transformation/round_transformation.py
index 9e52b1b..14f6852 100644
--- a/situr/transformation/round_transformation.py
+++ b/situr/transformation/round_transformation.py
@@ -2,9 +2,10 @@ import abc
 import scipy
 import numpy as np
 from situr.image import situ_image
+from situr.transformation import Transform
 
 
-class RoundTransform:
+class RoundTransform(Transform):
     __metaclass__ = abc.ABCMeta
 
     @abc.abstractmethod
diff --git a/situr/transformation/transformation.py b/situr/transformation/transformation.py
new file mode 100644
index 0000000..ec20a6b
--- /dev/null
+++ b/situr/transformation/transformation.py
@@ -0,0 +1,2 @@
+class Transform:
+    pass