Add inference codes for 3-ph category

Co-authored-by: Copilot <copilot@github.com>

src/datasets.py

@@ -272,5 +272,39 @@ class triplePhotonDataset(Dataset):
         sample = torch.tensor(sample, dtype=torch.float32).unsqueeze(0)
         label = torch.tensor(label, dtype=torch.float32)
         return sample, label
+    def __len__(self):
+        return self.length
+
+class triplePhotonInferenceDataset(Dataset):
+    def __init__(self, sampleList, sampleRatio, datasetName):
+        self.sampleFileList = sampleList
+        self.sampleRatio = sampleRatio
+        self.datasetName = datasetName
+        all_samples = []
+        all_ref_pts = []
+        for idx, sampleFile in enumerate(self.sampleFileList):
+            if '.npz' in sampleFile:
+                data = np.load(sampleFile)
+                all_samples.append(data['samples'])
+                all_ref_pts.append(data['referencePoint'])
+            elif '.h5' in sampleFile:
+                import h5py
+                with h5py.File(sampleFile, 'r') as f:
+                    samples = f['clusters'][:]
+                    ref_pts = f['referencePoint'][:]
+                    all_samples.append(samples)
+                    all_ref_pts.append(ref_pts)
+        self.samples = np.concatenate(all_samples, axis=0) if all_samples else None
+        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 __getitem__(self, index):
+        sample = self.samples[index]
+        # sample[sample == 0] += np.random.normal(loc=0.0, scale=0.13, size=sample[sample == 0].shape)  ### add noise to zero pixels
+        sample = torch.tensor(sample, dtype=torch.float32).unsqueeze(0)
+        dummy_label = np.zeros((3, 4), dtype=np.float32)  ### dummy label for 3 photons
+        return sample, torch.tensor(dummy_label, dtype=torch.float32)
     def __len__(self):
         return self.length