add type hints to registrations

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)

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

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:

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

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

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

situr/transformation/transformation.py

@@ -0,0 +1,2 @@
+class Transform:
+    pass