cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-09 21:12:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

Rename HoloDataset to DirectCDIDataset and update related methods for clarity

Browse Source
This commit is contained in:
gnzng
2025-03-11 15:46:42 -07:00
parent b4eee52aed
commit fbc23af55b
1 changed files with 20 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/datasets/holo_dataset.py
+20 -15
View File
@@ -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