Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from diffusers import DDIMScheduler | |
| from einops import rearrange | |
| # Create sinusoidal time embeddings as used in diffusion models | |
| def get_sinusoidal_embeddings(timesteps, embedding_dim): | |
| assert len(timesteps.shape) == 1 | |
| half_dim = embedding_dim // 2 | |
| emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1) | |
| emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb).to(timesteps.device) | |
| emb = timesteps.float()[:, None] * emb[None, :] | |
| emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) | |
| if embedding_dim % 2 == 1: | |
| emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) # zero padding if odd | |
| return emb | |
| # 3D Residual Block with time conditioning | |
| class ResidualBlock3D(nn.Module): | |
| def __init__(self, in_channels, out_channels, time_emb_dim, dropout): | |
| super().__init__() | |
| self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1) | |
| self.bn1 = nn.BatchNorm3d(out_channels) | |
| self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1) | |
| self.bn2 = nn.BatchNorm3d(out_channels) | |
| self.time_mlp = nn.Sequential( | |
| nn.SiLU(), | |
| nn.Linear(time_emb_dim, out_channels) | |
| ) | |
| self.residual_conv = nn.Conv3d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity() | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, t): | |
| h = F.relu(self.bn1(self.conv1(x))) | |
| time_emb = self.time_mlp(t).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) | |
| h = h + time_emb | |
| h = self.dropout(h) | |
| h = F.relu(self.bn2(self.conv2(h))) | |
| return h + self.residual_conv(x) | |
| # Downsampling block with residual block and max pooling | |
| class DownBlock3D(nn.Module): | |
| def __init__(self, in_channels, out_channels, time_emb_dim, dropout): | |
| super().__init__() | |
| self.res = ResidualBlock3D(in_channels, out_channels, time_emb_dim, dropout) | |
| self.pool = nn.MaxPool3d(2) | |
| def forward(self, x, t): | |
| x = self.res(x, t) | |
| return self.pool(x), x # skip connection is returned | |
| # Upsampling block with skip connection | |
| class UpBlock3D(nn.Module): | |
| def __init__(self, in_channels, out_channels, time_emb_dim, dropout): | |
| super().__init__() | |
| self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True) | |
| self.res = ResidualBlock3D(in_channels * 2, out_channels, time_emb_dim, dropout) | |
| def forward(self, x, t, skip): | |
| x = self.up(x) | |
| x = torch.cat([x, skip], dim=1) # concatenate skip connection | |
| return self.res(x, t) | |
| # U-Net 3D architecture with time and conditioning embeddings | |
| class UNet3D(nn.Module): | |
| def __init__(self, in_channels=1, out_channels=1, time_emb_dim=128, base_dim=32, dropout=0.1, cond_dim=256): | |
| super().__init__() | |
| self.time_mlp = nn.Sequential( | |
| nn.Linear(time_emb_dim, time_emb_dim * 4), | |
| nn.SiLU(), | |
| nn.Linear(time_emb_dim * 4, time_emb_dim) | |
| ) | |
| self.cond_proj = nn.Linear(cond_dim, time_emb_dim) | |
| self.init_conv = nn.Conv3d(in_channels, base_dim, kernel_size=3, padding=1) | |
| self.down1 = DownBlock3D(base_dim, base_dim * 2, time_emb_dim, dropout) | |
| self.down2 = DownBlock3D(base_dim * 2, base_dim * 4, time_emb_dim, dropout) | |
| self.bottleneck = ResidualBlock3D(base_dim * 4, base_dim * 4, time_emb_dim, dropout) | |
| self.up1 = UpBlock3D(base_dim * 4, base_dim * 2, time_emb_dim, dropout) | |
| self.up2 = UpBlock3D(base_dim * 2, base_dim, time_emb_dim, dropout) | |
| self.out_conv = nn.Conv3d(base_dim, out_channels, kernel_size=1) | |
| def forward(self, x, timesteps, cond): | |
| # Build and project time embeddings | |
| t_emb = get_sinusoidal_embeddings(timesteps, self.time_mlp[0].in_features) | |
| t_emb = self.time_mlp(t_emb.half()) | |
| cond_emb = self.cond_proj(cond) | |
| t_emb = t_emb + cond_emb | |
| # Apply U-Net | |
| x = self.init_conv(x) | |
| x, skip1 = self.down1(x, t_emb) | |
| x, skip2 = self.down2(x, t_emb) | |
| x = self.bottleneck(x, t_emb) | |
| x = self.up1(x, t_emb, skip2) | |
| x = self.up2(x, t_emb, skip1) | |
| return self.out_conv(x) | |
| # Full PATCDF model | |
| class PATCDF(nn.Module): | |
| def __init__(self, region_coords=None, num_braak_stages=0): | |
| super().__init__() | |
| # Multimodal encoder to embed input features | |
| self.encoder = ParkinsonMultimodalEncoder( | |
| mri_channels=1, | |
| clinical_dim=8, | |
| datscan_dim=4, | |
| latent_dim=128 | |
| ) | |
| # Time-conditioned disease progression embedding | |
| self.conditioner = ParkinsonProgressionConditioner( | |
| time_dim=64, | |
| clinical_dim=32, | |
| biomarker_dim=32, | |
| latent_dim=128, | |
| max_time_interval=10 | |
| ) | |
| # Regularization component for PPSE, symptoms, etc. | |
| self.pareg = ParkinsonAwareRegularization( | |
| latent_dim=128, | |
| cond_dim=128, | |
| num_braak_stages=num_braak_stages, | |
| region_coords=region_coords, | |
| modality_dims={'mri': 64, 'clinical': 64, 'datscan': 12} | |
| ) | |
| # 3D diffusion model and noise scheduler | |
| self.unet = UNet3D(in_channels=1, out_channels=1, cond_dim=256) | |
| self.scheduler = DDIMScheduler(num_train_timesteps=1000) | |
| def forward(self, baseline_mri, clinical, datscan, delta_t, clinical_delta, biomarker_delta, current_stage, targets=None, return_reg_losses=False): | |
| batch_size = baseline_mri.shape[0] | |
| device = baseline_mri.device | |
| # Encode each modality | |
| mri_feat = self.encoder.process_mri(baseline_mri) | |
| clinical_feat = self.encoder.process_clinical(clinical) | |
| datscan_feat = self.encoder.process_datscan(datscan) | |
| latent_base = self.encoder(baseline_mri, clinical, datscan) | |
| modality_features = { | |
| 'mri': mri_feat, | |
| 'clinical': clinical_feat, | |
| 'datscan': datscan_feat | |
| } | |
| # Generate progression-aware conditioning vector | |
| cond_vector = self.conditioner(delta_t, clinical_delta, biomarker_delta, current_stage) | |
| full_cond = torch.cat([latent_base, cond_vector], dim=-1) | |
| # Sample timestep and apply noise | |
| timesteps = torch.randint(0, self.scheduler.config.num_train_timesteps, (batch_size,), device=device, dtype=torch.long) | |
| noise = torch.randn_like(baseline_mri) | |
| noisy_mri = self.scheduler.add_noise(baseline_mri, noise, timesteps) | |
| # Predict noise from the model | |
| pred_noise = self.unet(noisy_mri, timesteps, full_cond) | |
| # Apply multimodal regularization | |
| pareg_outs = self.pareg( | |
| latent=latent_base, | |
| cond_vector=cond_vector, | |
| scan_gen=pred_noise, | |
| scan_base=baseline_mri, | |
| modality_features=modality_features, | |
| return_losses=return_reg_losses, | |
| targets=targets | |
| ) | |
| return pred_noise, pareg_outs | |
| # Inference pipeline to generate MRI predictions | |
| def generate(self, baseline_mri, clinical, datscan, delta_t, clinical_delta, biomarker_delta, current_stage, num_inference_steps=50): | |
| self.unet.eval() | |
| self.scheduler.set_timesteps(num_inference_steps, device=baseline_mri.device) | |
| batch_size = baseline_mri.shape[0] | |
| device = baseline_mri.device | |
| # Encode conditioning information | |
| latent_base = self.encoder(baseline_mri, clinical, datscan) | |
| cond_vector = self.conditioner(delta_t, clinical_delta, biomarker_delta, current_stage) | |
| full_cond = torch.cat([latent_base, cond_vector], dim=-1) | |
| sample = torch.randn_like(baseline_mri) | |
| # Iteratively denoise using the U-Net and scheduler | |
| for t in self.scheduler.timesteps: | |
| with torch.no_grad(): | |
| timesteps = torch.full((batch_size,), t, device=device, dtype=torch.long) | |
| pred_noise = self.unet(sample, timesteps, full_cond) | |
| sample = self.scheduler.step(pred_noise, t, sample).prev_sample | |
| return sample |