From e99942c7c325663ac7d6f86d3164bb15bf68985a Mon Sep 17 00:00:00 2001 From: "xiangyu.xie" Date: Wed, 10 Jun 2026 11:15:58 +0200 Subject: [PATCH] Simplify datasets, move coordinates into models; remove old 2ph models --- src/datasets.py | 85 +++------------ src/models.py | 277 ------------------------------------------------ 2 files changed, 15 insertions(+), 347 deletions(-) diff --git a/src/datasets.py b/src/datasets.py index 9a181ee..dae0fd4 100644 --- a/src/datasets.py +++ b/src/datasets.py @@ -114,80 +114,36 @@ class singlePhotonDataset(Dataset): return self.effectiveLength class doublePhotonDataset(Dataset): - def __init__(self, sampleList, sampleRatio, datasetName, reuselFactor=1, noiseKeV=0, nSize=6, noiseThreshold=0, normalize=False): + def __init__(self, sampleList, sampleRatio, datasetName): self.sampleFileList = sampleList self.sampleRatio = sampleRatio self.datasetName = datasetName - self.noiseKeV = noiseKeV - self.nSize = nSize - self.normalize = normalize - self._init_coords() all_samples = [] all_labels = [] for idx, sampleFile in enumerate(self.sampleFileList): - if '.npz' in sampleFile: - data = np.load(sampleFile) - all_samples.append(data['samples']) - all_labels.append(data['labels']) - elif '.h5' in sampleFile: - import h5py - with h5py.File(sampleFile, 'r') as f: - samples = f['clusters'][:] - labels = f['labels'][:] - all_samples.append(samples) - all_labels.append(labels) + data = np.load(sampleFile) + all_samples.append(data['samples']) + all_labels.append(data['labels']) self.samples = np.concatenate(all_samples, axis=0) - 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) - self.samples = self.samples + noise - if noiseThreshold != 0: - print(f'[{self.datasetName} dataset] \t Setting values below noise threshold ({noiseThreshold} keV) to zero') - self.samples[self.samples < noiseThreshold] = 0 ### set values below threshold to zero - 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 * 30. # normalize each sample by its total charge + # self.samples = torch.tensor(self.samples, dtype=torch.float32).unsqueeze(1) ### B,1,nSize,nSize + self.samples = torch.from_numpy(self.samples).float().unsqueeze(1) ### B,1,nSize,nSize + self.labels = np.concatenate(all_labels, axis=0) + self.labels_ph0 = self.labels[:, 0, :] ### B,4 (x1,y1,z1,e1) + self.labels_ph0[:, :2] -= self.samples.shape[-1] / 2. ### adjust labels to be centered at sample center + self.labels_ph1 = self.labels[:, 1, :] ### B,4 (x2,y2,z2,e2) + self.labels_ph1[:, :2] -= self.samples.shape[-1] / 2. ### adjust labels to be centered at sample center + self.labels = np.concatenate((self.labels_ph0, self.labels_ph1), axis=1) ### B,8 (x1,y1,z1,e1,x2,y2,z2,e2) + self.labels = torch.from_numpy(self.labels).float() ### B,8 ### total number of samples - self.length = int(self.samples.shape[0] * self.sampleRatio) // 2 * reuselFactor + self.length = int(self.samples.shape[0] * self.sampleRatio) 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(-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((self.nSize+2, self.nSize+2), dtype=np.float32) - sample = np.random.normal(loc=0.0, scale=self.noiseKeV, size=(self.nSize+2, self.nSize+2)) ### add noise to the whole sample - idx1 = np.random.randint(0, self.samples.shape[0]) - idx2 = np.random.randint(0, self.samples.shape[0]) - photon1 = self.samples[idx1] - photon2 = self.samples[idx2] - singlePhotonSize = photon1.shape[0] - - ### random position for photons in - pos_x1 = np.random.randint(1, 4) - pos_y1 = np.random.randint(1, 4) - sample[pos_y1:pos_y1+singlePhotonSize, pos_x1:pos_x1+singlePhotonSize] += photon1 - 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: 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 - - 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) + return self.samples[index], self.labels[index] def __len__(self): return self.length @@ -199,7 +155,6 @@ class doublePhotonInferenceDataset(Dataset): self.sampleRatio = sampleRatio self.datasetName = datasetName self.nSize = nSize - self._init_coords() all_samples = [] all_ref_pts = [] @@ -221,20 +176,10 @@ class doublePhotonInferenceDataset(Dataset): 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(-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] - # 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) - sample = torch.cat((sample, self.x_grid, self.y_grid), dim=0) ### concatenate coordinate channels dummy_label = np.zeros((8,), dtype=np.float32) ### dummy label return sample, torch.tensor(dummy_label, dtype=torch.float32) diff --git a/src/models.py b/src/models.py index 65aa189..9bf2f2a 100644 --- a/src/models.py +++ b/src/models.py @@ -147,283 +147,6 @@ class singlePhotonNet_260511(nn.Module): coords = self.fc(flat_feat) # [B, 2] return coords -class doublePhotonNet_250909(nn.Module): - def __init__(self): - super(doublePhotonNet_250909, self).__init__() - self.conv1 = nn.Conv2d(1, 3, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(3, 5, kernel_size=3, padding=1) - self.conv3 = nn.Conv2d(5, 5, kernel_size=3, padding=1) - self.fc1 = nn.Linear(5*6*6, 4) - # self.fc2 = nn.Linear(50, 4) - # 初始化更稳一些 - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, nonlinearity="relu") - if isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - # x = F.relu(self.fc1(x)) - # x = self.fc2(x) - x = self.fc1(x) - return x -class doublePhotonNet_250910(nn.Module): - def __init__(self): - super(doublePhotonNet_250910, self).__init__() - ### x shape: (B, 1, 6, 6) - self.conv1 = nn.Conv2d(1, 5, kernel_size=5, padding=2) # (B,5,6,6) - self.conv2 = nn.Conv2d(5, 10, kernel_size=5, padding=2) # (B,10,6,6) - self.conv3 = nn.Conv2d(10, 20, kernel_size=3, padding=0) # (B,20,4,4) - self.fc1 = nn.Linear(20*4*4, 4) - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, nonlinearity="relu") - if isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - x = F.relu(self.conv1(x)) - x = F.relu(self.conv2(x)) - x = F.relu(self.conv3(x)) - x = x.view(x.size(0), -1) - # x = F.relu(self.fc1(x)) - # x = self.fc2(x) - x = self.fc1(x) * 6 - return x - -class doublePhotonNet_251001(nn.Module): - def __init__(self): - super().__init__() - # 保持空间分辨率：使用小卷积核 + 无池化 - self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1) # 6x6 -> 6x6 - self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) # 6x6 -> 6x6 - self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # 6x6 -> 6x6 - - # 全局特征提取（替代全连接层） - self.global_avg_pool = nn.AdaptiveAvgPool2d((1,1)) # 64x1x1 - self.global_max_pool = nn.AdaptiveMaxPool2d((1,1)) # 64x1x1 - - # 回归头：输出4个坐标 (x1,y1,x2,y2) - self.fc = nn.Sequential( - nn.Linear(64, 128), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(128, 4), # 直接输出坐标 - # nn.Sigmoid() # sigmoid leads to overfitting - ) - # for m in self.modules(): - # if isinstance(m, nn.Conv2d): - # nn.init.kaiming_normal_(m.weight, nonlinearity="relu") - # if isinstance(m, nn.Linear): - # nn.init.xavier_uniform_(m.weight) - # nn.init.zeros_(m.bias) - - def forward(self, x): - x = torch.relu(self.conv1(x)) - x = torch.relu(self.conv2(x)) - x = torch.relu(self.conv3(x)) - # x = self.global_avg_pool(x).view(x.size(0), -1) - x = self.global_max_pool(x).view(x.size(0), -1) - coords = self.fc(x) - return coords # shape: [B, 4] - -class doublePhotonNet_251001_2(nn.Module): - def __init__(self): - super().__init__() - # Backbone: deeper + residual-like blocks - self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) - self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) - - # Spatial Attention Module (轻量但有效) - self.spatial_attn = nn.Sequential( - nn.Conv2d(128, 1, kernel_size=1), - nn.Sigmoid() - ) - - # Multi-scale feature fusion (optional but helpful) - self.reduce1 = nn.Conv2d(32, 32, kernel_size=1) # from conv1 - self.reduce2 = nn.Conv2d(64, 32, kernel_size=1) # from conv2 - self.fuse = nn.Conv2d(32*3, 128, kernel_size=1) - - # Global context with both Max and Avg pooling (better than GAP alone) - self.global_max_pool = nn.AdaptiveMaxPool2d((1,1)) - self.global_avg_pool = nn.AdaptiveAvgPool2d((1,1)) - - # Enhanced regression head - self.fc = nn.Sequential( - nn.Linear(128 * 2, 256), # concat max + avg - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(256, 128), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(128, 4), - # nn.Sigmoid() # output in [0,1] - ) - - self._init_weights() - - def _init_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - # Feature extraction - c1 = F.relu(self.conv1(x)) # [B, 32, 6, 6] - c2 = F.relu(self.conv2(c1)) # [B, 64, 6, 6] - c3 = F.relu(self.conv3(c2)) # [B,128, 6, 6] - - # Spatial attention: highlight photon peaks - attn = self.spatial_attn(c3) # [B, 1, 6, 6] - c3 = c3 * attn # reweight features - - # (Optional) Multi-scale fusion — uncomment if needed - # r1 = F.interpolate(self.reduce1(c1), size=(6,6), mode='nearest') - # r2 = self.reduce2(c2) - # fused = torch.cat([r1, r2, c3], dim=1) - # c3 = self.fuse(fused) - - # Global context: MaxPool better captures peaks, Avg for context - g_max = self.global_max_pool(c3).flatten(1) # [B, 128] - g_avg = self.global_avg_pool(c3).flatten(1) # [B, 128] - global_feat = torch.cat([g_max, g_avg], dim=1) # [B, 256] - - # Regression - coords = self.fc(global_feat) - return coords # [B, 4] - -class doublePhotonNet_251124(nn.Module): ### adapted for 7x7 input from 251001_2 - def __init__(self): - super().__init__() - # Backbone: deeper + residual-like blocks - self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) ### 7x7 - self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) ### 7x7 - self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) ### 7x7 - - # Spatial Attention Module (轻量但有效) - self.spatial_attn = nn.Sequential( - nn.Conv2d(128, 1, kernel_size=1), - nn.Sigmoid() - ) - - # Multi-scale feature fusion (optional but helpful) - self.reduce1 = nn.Conv2d(32, 32, kernel_size=1) # from conv1 - self.reduce2 = nn.Conv2d(64, 32, kernel_size=1) # from conv2 - self.fuse = nn.Conv2d(32*3, 128, kernel_size=1) - - # Global context with both Max and Avg pooling (better than GAP alone) - self.global_max_pool = nn.AdaptiveMaxPool2d((1,1)) - self.global_avg_pool = nn.AdaptiveAvgPool2d((1,1)) - - # Enhanced regression head - self.fc = nn.Sequential( - nn.Linear(128 * 2, 256), # concat max + avg - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(256, 128), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(128, 4), - # nn.Sigmoid() # output in [0,1] - ) - - self._init_weights() - - def _init_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - # Feature extraction - c1 = F.relu(self.conv1(x)) # [B, 32, 7, 7] - c2 = F.relu(self.conv2(c1)) # [B, 64, 7, 7] - c3 = F.relu(self.conv3(c2)) # [B,128, 7, 7] - - # Spatial attention: highlight photon peaks - attn = self.spatial_attn(c3) # [B, 1, 7, 7] - c3 = c3 * attn # reweight features - - # (Optional) Multi-scale fusion — uncomment if needed - # r1 = F.interpolate(self.reduce1(c1), size=(7,7), mode='nearest') - # r2 = self.reduce2(c2) - # fused = torch.cat([r1, r2, c3], dim=1) - # c3 = self.fuse(fused) - - # Global context: MaxPool better captures peaks, Avg for context - g_max = self.global_max_pool(c3).flatten(1) # [B, 128] - g_avg = self.global_avg_pool(c3).flatten(1) # [B, 128] - global_feat = torch.cat([g_max, g_avg], dim=1) # [B, 256] - - # Regression - coords = self.fc(global_feat) - return coords # [B, 4] - -class doublePhotonNet_260507(nn.Module): ## adapted from 251124, removed max pooling, added more capacity in FC layers - def __init__(self): - super().__init__() - # Backbone 保持不变 - self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) - self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) - self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) - - # 空间注意力模块 (Spatial Attention Module) - self.spatial_attn = nn.Sequential( - nn.Conv2d(128, 1, kernel_size=1), - nn.Sigmoid() - ) - - # 我们移除了全局池化层 (Global Pooling)。 - # 7x7 的空间特征图有 128 个通道，展平后是 128 * 7 * 7 = 6272 个特征。 - # 这使得全连接层能够直接“看到”精确的空间分布规律。 - self.fc = nn.Sequential( - nn.Linear(128 * 7 * 7, 512), # 增加了中间层的维度以适配展平后的输入 - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(512, 128), - nn.ReLU(), - nn.Dropout(0.3), - nn.Linear(128, 4) - ) - - self._init_weights() - - def _init_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') - elif isinstance(m, nn.Linear): - nn.init.xavier_uniform_(m.weight) - nn.init.zeros_(m.bias) - - def forward(self, x): - c1 = F.relu(self.conv1(x)) # [B, 32, 7, 7] - c2 = F.relu(self.conv2(c1)) # [B, 64, 7, 7] - c3 = F.relu(self.conv3(c2)) # [B, 128, 7, 7] - - attn = self.spatial_attn(c3) # [B, 1, 7, 7] - c3 = c3 * attn # [B, 128, 7, 7] - - # 直接展平 (Flatten) 而不是池化 - flat_feat = c3.view(c3.size(0), -1) # [B, 6272] - - coords = self.fc(flat_feat) # [B, 4] - return coords - class doublePhotonNet_260608(nn.Module): ## adapted from 260507, removed padding and dropout def __init__(self, nSize=6): super().__init__()