Add normalize, noiseThreshold options; add rms x/y outputs

src/datasets.py

@@ -3,11 +3,13 @@ import torch
 import numpy as np
 
 class singlePhotonDataset(Dataset):
-    def __init__(self, sampleList, sampleRatio, datasetName, noiseKeV=0, numberOfAugOps=1):
+    def __init__(self, sampleList, sampleRatio, datasetName, noiseKeV=0, numberOfAugOps=1, normalize=False, noiseThreshold=0):
         self.sampleFileList = sampleList
         self.sampleRatio = sampleRatio
         self.datasetName = datasetName
         self.numberOfAugOps = numberOfAugOps
+        self.normalize = normalize
+        self.noiseThreshold = noiseThreshold
         self._init_coords()
 
         all_samples = []
@@ -45,12 +47,23 @@ class singlePhotonDataset(Dataset):
             print(f'Adding Gaussian noise with sigma = {noiseKeV} keV to samples in {self.datasetName} dataset')
             noise = np.random.normal(loc=0.0, scale=noiseKeV, size=self.samples.shape)
             self.samples = self.samples + noise
+        if self.noiseThreshold != 0 and noiseKeV != 0:
+            print(f'[{self.datasetName} dataset] \t Setting values below noise threshold ({self.noiseThreshold} keV) to zero')
+            self.samples[self.samples < self.noiseThreshold] = 0  ### set values below threshold to zero
         self.labels = np.concatenate(all_labels, axis=0)
         self.referencePoint = np.concatenate(all_ref_pts, axis=0) if all_ref_pts else None
+
+        if self.normalize:
+            print(f'Normalizing samples in {self.datasetName} dataset by total charge')
+            total_charge = np.sum(self.samples, axis=(1,2), keepdims=True)  # (B, 1, 1)
+            total_charge[total_charge == 0] = 1  # avoid division by zero
+            self.samples = self.samples / total_charge * 15.  # normalize each sample by its total charge
         
         if self.samples.shape[1] == 5: ### if sample size is 5x5, remove border pixels to make it 3x3
             self.samples = self.samples[:, 1:-1, 1:-1]  ### remove border pixels
             self.labels = self.labels - np.array([1, 1, 0, 0])  ### adjust labels accordingly
+            if self.referencePoint is not None:
+                self.referencePoint = self.referencePoint + np.array([1, 1])  ### adjust reference points accordingly
         self.samples = np.expand_dims(self.samples, axis=1)
         self.labels -= np.array([self.samples.shape[-1]/2., self.samples.shape[-1]/2., 0, 0])  ### B,D,3,3 adjust labels to be centered at (0,0)
         self.labels[:, 2] /= 650. ### normalize z coordinate (depth) to [0, 1]