[core][dataset] Add transformation parameter to dataset

darfix/core/dataset.py

@@ -26,7 +26,7 @@
 
 __authors__ = ["J. Garriga"]
 __license__ = "MIT"
-__date__ = "04/06/2021"
+__date__ = "07/06/2021"
 
 import copy
 import glob
@@ -100,7 +100,7 @@ class Dataset():
     """
 
     def __init__(self, _dir, data=None, first_filename=None, filenames=None,
-                 dims=None, treated="treated_data", in_memory=True):
+                 dims=None, treated="treated_data", transformation=None, in_memory=True):
 
         self._data = None
         self._frames_intensity = []
@@ -110,7 +110,7 @@ class Dataset():
         self.running_data = None
         self.moments_dims = {}
         self.operations_state = numpy.zeros(len(Operation))
-        self._transformation = None
+        self._transformation = transformation
         if not os.path.isdir(self.dir + "/" + treated):
             try:
                 os.mkdir(self.dir + "/" + treated)
@@ -591,7 +591,8 @@ class Dataset():
             new_data = Data(urls.reshape(self.data.urls.shape), self.data.metadata,
                             self._in_memory)
 
-        return Dataset(_dir=bs_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)
+        return Dataset(_dir=bs_dir, data=new_data, dims=self.__dims, transformation=self.transformation,
+                       in_memory=self._in_memory)
 
     def apply_hot_pixel_removal(self, kernel=3, indices=None):
         """
@@ -635,7 +636,8 @@ class Dataset():
             new_data = Data(urls.reshape(self.data.urls.shape), self.data.metadata,
                             self._in_memory)
 
-        return Dataset(_dir=_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)
+        return Dataset(_dir=_dir, data=new_data, dims=self.__dims, transformation=self.transformation,
+                       in_memory=self._in_memory)
 
     def apply_threshold_removal(self, bottom=None, top=None, indices=None):
         """
@@ -679,7 +681,8 @@ class Dataset():
             new_data = Data(urls.reshape(self.data.urls.shape), self.data.metadata,
                             self._in_memory)
 
-        return Dataset(_dir=_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)
+        return Dataset(_dir=_dir, data=new_data, dims=self.__dims, transformation=self.transformation,
+                       in_memory=self._in_memory)
 
     def apply_roi(self, origin=None, size=None, center=None, indices=None, roi_dir=None):
         """
@@ -722,7 +725,8 @@ class Dataset():
             shape.append(new_data.shape[-2])
             shape.append(new_data.shape[-1])
             new_data = new_data.reshape(shape)
-        return Dataset(_dir=roi_dir, data=new_data, dims=self.__dims, in_memory=self._in_memory)
+        return Dataset(_dir=roi_dir, data=new_data, dims=self.__dims, transformation=transformation,
+                       in_memory=self._in_memory)
 
     def find_shift(self, dimension=None, h_max=0.5, h_step=0.01, indices=None):
         """
@@ -828,7 +832,8 @@ class Dataset():
         self._lock.release()
 
         data = Data(new_urls.reshape(self.data.urls.shape), self.data.metadata, in_memory=self._in_memory)
-        return Dataset(_dir=shift_dir, data=data, dims=self.__dims, in_memory=self._in_memory)
+        return Dataset(_dir=shift_dir, data=data, dims=self.__dims, transformation=self.transformation,
+                       in_memory=self._in_memory)
 
     def find_and_apply_shift(self, dimension=None, h_max=0.5, h_step=0.01, shift_approach="fft",
                              indices=None, callback=None):
@@ -1257,7 +1262,7 @@ class Dataset():
         else:
             new_urls = numpy.array(urls)
         data = Data(new_urls.reshape(self.data.urls.shape), self.data.metadata, in_memory=True)  # to modify
-        return Dataset(_dir=_dir, data=data, dims=self.__dims, in_memory=True)
+        return Dataset(_dir=_dir, data=data, dims=self.__dims, transformation=self.transformation, in_memory=True)
 
     def compute_transformation(self, d, rotate=False):