[core][dataset] Fix typo when only saving part of a dataset

darfix/core/dataset.py

@@ -26,7 +26,7 @@
 
 __authors__ = ["J. Garriga"]
 __license__ = "MIT"
-__date__ = "05/07/2021"
+__date__ = "09/07/2021"
 
 import copy
 import glob
@@ -345,7 +345,6 @@ class Dataset():
         self.__dims.remove_dim(axis)
 
     def zsum(self, indices=None, dimension=None):
-
         data = self.get_data(indices, dimension)
         return data.sum(axis=0)
 
@@ -718,7 +717,7 @@ class Dataset():
         self.running_data = self.get_data(indices)
         if self._in_memory:
             new_data = apply_3D_ROI(self.running_data, origin, size, center).view(Data)
-            new_data.save(roi_dir + "/data.hdf5")
+            new_data.save(roi_dir + "/data.hdf5", new_shape=new_data.shape)
         else:
             shape = numpy.append([self.running_data.shape[0]], apply_2D_ROI(self.running_data[0], origin, size, center).shape)
             urls = self.running_data.apply_funcs([(apply_2D_ROI, [origin, size, center])],
@@ -790,7 +789,7 @@ class Dataset():
         _file = h5py.File(_dir + '/data.hdf5', 'a')
         dataset_name = "dataset"
         if "dataset" in _file:
-            _file.create_dataset("update_dataset", data.shape, dtype=data.dtype)
+            _file.create_dataset("update_dataset", data=_file["dataset"])
             dataset_name = "update_dataset"
         else:
             _file.create_dataset("dataset", data.shape, dtype=data.dtype)
@@ -861,7 +860,6 @@ class Dataset():
             del _file["update_dataset"]
 
         _file.close()
-
         data = Data(new_urls.reshape(self.data.urls.shape), self.data.metadata, in_memory=self._in_memory)
         return Dataset(_dir=_dir, data=data, dims=self.__dims, transformation=self.transformation,
                        in_memory=self._in_memory)
@@ -1470,8 +1468,6 @@ class Data(numpy.ndarray):
         """
         if indices is None:
             indices = range(len(self))
-        if new_shape is None:
-            new_shape = self.shape
         if isinstance(indices, int):
             indices = [indices]
         urls = []
@@ -1486,9 +1482,13 @@ class Data(numpy.ndarray):
         _file = h5py.File(save, 'a')
         dataset_name = "dataset"
         if "dataset" in _file:
-            _file.create_dataset("update_dataset", new_shape, dtype=self.dtype)
+            if new_shape is not None:
+                _file.create_dataset("update_dataset", new_shape, dtype=self.dtype)
+            else:
+                _file.create_dataset("update_dataset", data=_file["dataset"])
             dataset_name = "update_dataset"
         else:
+            new_shape = self.shape if new_shape is None else new_shape
             _file.create_dataset("dataset", new_shape, dtype=self.dtype)
 
         for i in indices:
@@ -1525,7 +1525,7 @@ class Data(numpy.ndarray):
         _file.close()
         return numpy.array(urls)
 
-    def save(self, path, indices=None):
+    def save(self, path, indices=None, new_shape=None):
         """
         Save the data into `path` folder and replace Data urls.
         TODO: check if urls already exist and, if so, modify only urls[indices].
@@ -1543,9 +1543,13 @@ class Data(numpy.ndarray):
             data = self[indices]
 
         _file = h5py.File(path, 'a')
-        if "dataset" in _file:
+
+        if "dataset" not in _file:
+            _file.create_dataset("dataset", self.shape, dtype=self.dtype)
+        elif new_shape is not None:
             del _file["dataset"]
-        _file.create_dataset("dataset", self.shape, dtype=self.dtype)
+            _file.create_dataset("dataset", new_shape, dtype=self.dtype)
+
 
         for i, img in enumerate(data):
             _file["dataset"][i] = img