MLXID/DeepLearning
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add configurable cluster size for double phtoon sample

Browse Source
This commit is contained in:
xie_x1
2025-11-25 15:25:07 +01:00
parent 1c4f03a308
commit eef6a87f06
1 changed files with 21 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
src/datasets.py
View File

@@ -101,11 +101,12 @@ class singlePhotonDataset(Dataset):
return self.effectiveLength
class doublePhotonDataset(Dataset):
def __init__(self, sampleList, sampleRatio, datasetName, reuselFactor=1, noiseKeV=0):
def __init__(self, sampleList, sampleRatio, datasetName, reuselFactor=1, noiseKeV=0, nSize=6):
self.sampleFileList = sampleList
self.sampleRatio = sampleRatio
self.datasetName = datasetName
self.noiseKeV = noiseKeV
self.nSize = nSize
self._init_coords()
all_samples = []
@@ -118,6 +119,8 @@ class doublePhotonDataset(Dataset):
if self.noiseKeV != 0:
print(f'Adding Gaussian noise with sigma = {self.noiseKeV} keV to samples in {self.datasetName} dataset')
noise = np.random.normal(loc=0.0, scale=self.noiseKeV, size=self.samples.shape)
#### add noise only to pixels that not zero
noise[self.samples == 0] = 0
self.samples = self.samples + noise
self.labels = np.concatenate(all_labels, axis=0)
@@ -127,14 +130,14 @@ class doublePhotonDataset(Dataset):
def _init_coords(self):
# Create a coordinate grid for 3x3 input
x = np.linspace(0, 5, 6)
y = np.linspace(0, 5, 6)
x_grid, y_grid = np.meshgrid(x, y, indexing='ij') # (6,6), (6,6)
self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous() # (1, 6, 6)
self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous() # (1, 6, 6)
x = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize)
y = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize)
x_grid, y_grid = np.meshgrid(x, y, indexing='ij') # (nSize,nSize), (nSize,nSize)
self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous() # (1, nSize, nSize)
self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous() # (1, nSize, nSize)
def __getitem__(self, index):
sample = np.zeros((8, 8), dtype=np.float32)
sample = np.zeros((self.nSize+2, self.nSize+2), dtype=np.float32)
idx1 = np.random.randint(0, self.samples.shape[0])
idx2 = np.random.randint(0, self.samples.shape[0])
photon1 = self.samples[idx1]
@@ -148,13 +151,12 @@ class doublePhotonDataset(Dataset):
pos_x2 = np.random.randint(1, 4)
pos_y2 = np.random.randint(1, 4)
sample[pos_y2:pos_y2+singlePhotonSize, pos_x2:pos_x2+singlePhotonSize] += photon2
sample = sample[1:-1, 1:-1] ### sample size: 6x6
sample = sample[1:-1, 1:-1] ### sample size: nSize x nSize
sample = torch.tensor(sample, dtype=torch.float32).unsqueeze(0)
sample = torch.cat((sample, self.x_grid, self.y_grid), dim=0) ### concatenate coordinate channels
doublePhotonSize = 6
label1 = self.labels[idx1] + np.array([pos_x1-1-doublePhotonSize//2, pos_y1-1-doublePhotonSize//2, 0, 0])
label2 = self.labels[idx2] + np.array([pos_x2-1-doublePhotonSize//2, pos_y2-1-doublePhotonSize//2, 0, 0])
label1 = self.labels[idx1] + np.array([pos_x1-1-self.nSize/2., pos_y1-1-self.nSize/2., 0, 0])
label2 = self.labels[idx2] + np.array([pos_x2-1-self.nSize/2., pos_y2-1-self.nSize/2., 0, 0])
label = np.concatenate((label1, label2), axis=0)
return sample, torch.tensor(label, dtype=torch.float32)
@@ -163,10 +165,11 @@ class doublePhotonDataset(Dataset):
class doublePhotonInferenceDataset(Dataset):
def __init__(self, sampleList, sampleRatio, datasetName):
def __init__(self, sampleList, sampleRatio, datasetName, nSize=6):
self.sampleFileList = sampleList
self.sampleRatio = sampleRatio
self.datasetName = datasetName
self.nSize = nSize
self._init_coords()
all_samples = []
@@ -187,15 +190,16 @@ class doublePhotonInferenceDataset(Dataset):
self.referencePoint = np.concatenate(all_ref_pts, axis=0) if all_ref_pts else None
### total number of samples
self.length = int(self.samples.shape[0] * self.sampleRatio)
self.referencePoint = self.referencePoint[:self.length]
print(f"[{self.datasetName} dataset] \t Total number of samples: {self.length}")
def _init_coords(self):
# Create a coordinate grid for 3x3 input
x = np.linspace(0, 5, 6)
y = np.linspace(0, 5, 6)
x_grid, y_grid = np.meshgrid(x, y, indexing='ij') # (6,6), (6,6)
self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous() # (1, 6, 6)
self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous() # (1, 6, 6)
x = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize)
y = np.linspace(-self.nSize/2. + 0.5, self.nSize/2. - 0.5, self.nSize)
x_grid, y_grid = np.meshgrid(x, y, indexing='ij') # (nSize,nSize), (nSize,nSize)
self.x_grid = torch.tensor(np.expand_dims(x_grid, axis=0)).float().contiguous() # (1, nSize, nSize)
self.y_grid = torch.tensor(np.expand_dims(y_grid, axis=0)).float().contiguous() # (1, nSize, nSize)
def __getitem__(self, index):
sample = self.samples[index]