Simplify datasets, move coordinates into models; remove old 2ph models
This commit is contained in:
+15
-70
@@ -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)
|
||||
|
||||
|
||||
-277
@@ -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__()
|
||||
|
||||
Reference in New Issue
Block a user