Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from torch.nn.utils.parametrizations import weight_norm | |
| class AttentionFusion(nn.Module): | |
| def __init__(self, feature_dim: int, attention_dim: int): | |
| super().__init__() | |
| self.attention = nn.Linear(feature_dim, attention_dim) | |
| self.context_vector = nn.Linear(attention_dim, 1, bias=False) | |
| def forward(self, features: torch.Tensor) -> torch.Tensor: | |
| # features: [batch_size, feature_dim] | |
| attention_scores = self.context_vector(torch.tanh(self.attention(features))) # [batch_size, 1] | |
| attention_weights = torch.softmax(attention_scores, dim=1) # [batch_size, 1] | |
| weighted_features = features * attention_weights # [batch_size, feature_dim] | |
| return weighted_features | |
| class ParkinsonMultimodalEncoder(nn.Module): | |
| def __init__(self, | |
| mri_channels=1, | |
| clinical_dim=8, | |
| datscan_dim=4, | |
| latent_dim=64): # Final latent embedding dimension | |
| super().__init__() | |
| # -------- MRI Pathway (input: [B, 1, D, H, W]) -------- | |
| self.mri_encoder = nn.Sequential( | |
| nn.Conv3d(mri_channels, 16, kernel_size=3, padding=1), # [B, 16, D, H, W] | |
| nn.BatchNorm3d(16), | |
| nn.ReLU(), | |
| nn.MaxPool3d(2), # [B, 16, D/2, H/2, W/2] | |
| self._make_res_block3d(16, 32, stride=2), # [B, 32, D/4, H/4, W/4] | |
| self._make_res_block3d(32, 64, stride=1), # [B, 64, D/4, H/4, W/4] | |
| nn.Conv3d(64, 64, kernel_size=1), # [B, 64, D/4, H/4, W/4] | |
| nn.Sigmoid(), | |
| nn.AdaptiveAvgPool3d(1) # [B, 64, 1, 1, 1] | |
| ) | |
| # -------- Clinical Pathway (input: [B, T, 8]) -------- | |
| self.clinical_norm = nn.LayerNorm(clinical_dim) | |
| self.clinical_conv = weight_norm(nn.Conv1d(8, 16, kernel_size=3, padding=1)) # [B, 16, T] | |
| self.clinical_motor_branch = nn.Sequential(nn.Linear(16, 32), nn.ReLU()) # For time step 0 | |
| self.clinical_nonmotor_branch = nn.Sequential(nn.Linear(16, 32), nn.ReLU()) # For time step 1 | |
| # -------- DaTscan Pathway (input: [B, 4]) -------- | |
| self.datscan_base = nn.Sequential( | |
| nn.Linear(datscan_dim, 16), nn.ReLU(), # [B, 16] | |
| nn.Linear(16, 8), nn.Tanh() # [B, 8] | |
| ) | |
| self.putamen_head = nn.Sequential(nn.Linear(8, 4), nn.ReLU()) # [B, 4] | |
| self.caudate_head = nn.Sequential(nn.Linear(8, 4), nn.ReLU()) # [B, 4] | |
| self.ratio_head = nn.Linear(8, 4) # [B, 4] | |
| # -------- Fusion -------- | |
| fusion_input_dim = 64 + 64 + 12 # MRI + Clinical + DaTscan = 140 | |
| self.attention_fusion = AttentionFusion(feature_dim=fusion_input_dim, attention_dim=64) | |
| self.fusion = nn.Sequential( | |
| nn.Linear(fusion_input_dim, 256), # [B, 256] | |
| nn.LeakyReLU(0.2), | |
| nn.Dropout(0.3), | |
| nn.Linear(256, 64), # [B, 64] | |
| nn.Sigmoid(), | |
| nn.Linear(64, latent_dim) # [B, latent_dim] | |
| ) | |
| self._init_weights() | |
| def _make_res_block3d(self, in_channels, out_channels, stride=1): | |
| return nn.Sequential( | |
| nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1), | |
| nn.BatchNorm3d(out_channels), | |
| nn.ReLU(), | |
| nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1), | |
| nn.BatchNorm3d(out_channels) | |
| ) | |
| def _init_weights(self): | |
| for m in self.modules(): | |
| if isinstance(m, nn.Conv3d): | |
| nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') | |
| if m.bias is not None: | |
| nn.init.constant_(m.bias, 0) | |
| elif isinstance(m, nn.Linear): | |
| nn.init.xavier_uniform_(m.weight) | |
| if m.bias is not None: | |
| nn.init.constant_(m.bias, 0.01) | |
| # ----- Process each modality individually ----- | |
| def process_mri(self, mri: torch.Tensor) -> torch.Tensor: | |
| x = self.mri_encoder(mri) # [B, 64, 1, 1, 1] | |
| return x.view(x.size(0), -1) # [B, 64] | |
| def process_clinical(self, clinical: torch.Tensor) -> torch.Tensor: | |
| clinical = self.clinical_norm(clinical) # [B, T, 8] | |
| clinical = clinical.permute(0, 2, 1) # [B, 8, T] | |
| clinical = self.clinical_conv(clinical) # [B, 16, T] | |
| clinical = clinical.permute(0, 2, 1) # [B, T, 16] | |
| motor = self.clinical_motor_branch(clinical[:, 0, :]) # [B, 32] | |
| nonmotor = self.clinical_nonmotor_branch(clinical[:, 1, :]) # [B, 32] | |
| return torch.cat([motor, nonmotor], dim=1) # [B, 64] | |
| def process_datscan(self, datscan: torch.Tensor) -> torch.Tensor: | |
| base = self.datscan_base(datscan) # [B, 8] | |
| putamen = self.putamen_head(base) # [B, 4] | |
| caudate = self.caudate_head(base) # [B, 4] | |
| ratio = self.ratio_head(base) # [B, 4] | |
| return torch.cat([putamen, caudate, ratio], dim=1) # [B, 12] | |
| def forward(self, mri, clinical, datscan): | |
| # Each returns [B, X] | |
| mri_feat = self.process_mri(mri) # [B, 64] | |
| clinical_feat = self.process_clinical(clinical) # [B, 64] | |
| dat_feat = self.process_datscan(datscan) # [B, 12] | |
| combined = torch.cat([mri_feat, clinical_feat, dat_feat], dim=1) # [B, 140] | |
| weighted = self.attention_fusion(combined) # [B, 140] | |
| latent = self.fusion(weighted) # [B, latent_dim] | |
| return latent |