From fbc23af55b87708996fe89c76c6f18dc49656d8a Mon Sep 17 00:00:00 2001
From: gnzng <damian@guenzing.de>
Date: Mon, 10 Mar 2025 14:21:19 -0700
Subject: [PATCH] Rename HoloDataset to DirectCDIDataset and update related
 methods for clarity

---
 src/cdtools/datasets/holo_dataset.py | 35 ++++++++++++++++------------
 1 file changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/cdtools/datasets/holo_dataset.py b/src/cdtools/datasets/holo_dataset.py
index 374d6d5..4bcbc3f 100644
--- a/src/cdtools/datasets/holo_dataset.py
+++ b/src/cdtools/datasets/holo_dataset.py
@@ -6,14 +6,15 @@ import pathlib
 from PIL import Image
 
 
-class HoloDataset(torchDataset):
+class DirectCDIDataset(torchDataset):
     def __init__(self, data, transform=None):
         self.data = data
         self.transform = transform
 
     @classmethod
-    def from_h5(cls, h5_file, path_within_h5, axes_to_average=None, device=None, replace_dead_pixels=True):
-        """Generates a new HoloDataset from a .h5 file directly
+    def from_h5(cls, h5_file, path_within_h5, axes_to_average=None,
+                device=None, replace_dead_pixels=True):
+        """Generates a new DirectCDIDataset from a .h5 file directly
 
         Parameters
         ----------
@@ -28,16 +29,17 @@ class HoloDataset(torchDataset):
 
         Returns
         -------
-        dataset : HoloDataset
+        dataset : DirectCDIDataset
             The constructed dataset object
         """
-        
         # If a bare string is passed
         if isinstance(h5_file, str) or isinstance(h5_file, pathlib.Path):
-            with h5py.File(h5_file,'r') as f:
+            with h5py.File(h5_file, 'r') as f:
                 return cls.from_h5(f, path_within_h5, axes_to_average, device)
-        
-        data = t.tensor(h5_file[path_within_h5][:,:,:,:], device=device, dtype=t.float32)
+
+        data = t.tensor(h5_file[path_within_h5][:, :, :, :],
+                        device=device,
+                        dtype=t.float32)
         if axes_to_average is not None:
             data = t.mean(data, dim=axes_to_average)
 
@@ -45,14 +47,19 @@ class HoloDataset(torchDataset):
         if replace_dead_pixels:
             dead_pixels = list(t.argwhere(data == 0))
             for n in dead_pixels:
-                data[n[0],n[1]] = t.mean(t.tensor([data[n[0]+1,n[1]+1], data[n[0]+1,n[1]], data[n[0],n[1]+1], data[n[0]-1,n[1]-1], data[n[0]-1,n[1]], data[n[0],n[1]-1]]))
-        
+                data[n[0], n[1]] = t.mean(t.tensor([data[n[0]+1, n[1]+1],
+                                                    data[n[0]+1, n[1]],
+                                                    data[n[0], n[1]+1],
+                                                    data[n[0]-1, n[1]-1],
+                                                    data[n[0]-1, n[1]],
+                                                    data[n[0], n[1]-1]]))
+
         return cls(data)
 
     @classmethod
     def from_tif(cls, tif_file):
-        """Generates a new HoloDataset from a .tif file directly. It assumes that the .tif file is 
-        already stitched, edited and is a single pattern.
+        """Generates a new DirectCDIDataset from a .tif file directly. It assumes
+        that the .tif file is already stitched, edited and is a single pattern.
 
         Parameters
         ----------
@@ -61,7 +68,7 @@ class HoloDataset(torchDataset):
 
         Returns
         -------
-        dataset : HoloDataset
+        dataset : DirectCDIDataset
             The constructed dataset object
         """
         data = t.tensor(np.array(Image.open(tif_file)), dtype=t.float32)
@@ -75,5 +82,3 @@ class HoloDataset(torchDataset):
         if self.transform:
             data = self.transform(data)
         return data
-    
-    
\ No newline at end of file