tan200224/Synthetic-CT-Scan_VAE_Conditional

1"""2VAE model definition.3Input: (B, 4, 256, 256) — 4 slices (3D CT label/mask).4Output: decode(z) -> (B, 4, 256, 256) — 4 reconstructed slices.5"""6import torch7import torch.nn as nn8 9 10class Conv(nn.Module):11    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0):12        super().__init__()13        self.conv = nn.Sequential(14            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),15            nn.BatchNorm2d(out_channels),16            nn.LeakyReLU(inplace=True),17        )18 19    def forward(self, x: torch.Tensor) -> torch.Tensor:20        return self.conv(x)21 22 23class ConvTranspose(nn.Module):24    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0):25        super().__init__()26        self.conv = nn.Sequential(27            nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),28            nn.BatchNorm2d(out_channels),29            nn.LeakyReLU(inplace=True),30        )31 32    def forward(self, x: torch.Tensor) -> torch.Tensor:33        return self.conv(x)34 35 36class VAE(nn.Module):37    """VAE: 4-channel input (4 slices) -> latent -> 4-channel output (4 slices)."""38 39    def __init__(self, base: int = 64):40        super().__init__()41        self.base = base42 43        self.encoder = nn.Sequential(44            Conv(4, base, 3, stride=2, padding=1),45            Conv(base, 2 * base, 3, padding=1),46            Conv(2 * base, 2 * base, 3, stride=2, padding=1),47            Conv(2 * base, 2 * base, 3, padding=1),48            Conv(2 * base, 2 * base, 3, stride=2, padding=1),49            Conv(2 * base, 4 * base, 3, padding=1),50            Conv(4 * base, 4 * base, 3, stride=2, padding=1),51            Conv(4 * base, 4 * base, 3, padding=1),52            Conv(4 * base, 4 * base, 3, stride=2, padding=1),53            nn.Conv2d(4 * base, 64 * base, 8),54            nn.LeakyReLU(inplace=True),55        )56 57        self.encoder_mu = nn.Conv2d(64 * base, 32 * base, 1)58        self.encoder_logvar = nn.Conv2d(64 * base, 32 * base, 1)59 60        self.decoder = nn.Sequential(61            nn.Conv2d(32 * base, 64 * base, 1),62            ConvTranspose(64 * base, 4 * base, 8),63            Conv(4 * base, 4 * base, 3, padding=1),64            ConvTranspose(4 * base, 4 * base, 4, stride=2, padding=1),65            Conv(4 * base, 4 * base, 3, padding=1),66            ConvTranspose(4 * base, 4 * base, 4, stride=2, padding=1),67            Conv(4 * base, 2 * base, 3, padding=1),68            ConvTranspose(2 * base, 2 * base, 4, stride=2, padding=1),69            Conv(2 * base, 2 * base, 3, padding=1),70            ConvTranspose(2 * base, 2 * base, 4, stride=2, padding=1),71            Conv(2 * base, base, 3, padding=1),72            ConvTranspose(base, base, 4, stride=2, padding=1),73            nn.Conv2d(base, 4, 3, padding=1),74            nn.Sigmoid(),75        )76 77    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:78        x = self.encoder(x)79        return self.encoder_mu(x), self.encoder_logvar(x)80 81    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:82        """Standard VAE reparameterization: z = mu + std * eps. For inference (eval mode), return mu for deterministic output."""83        if not self.training:84            return mu85        std = torch.exp(0.5 * logvar)86        eps = torch.randn_like(std)87        return mu + std * eps88 89    def decode(self, z: torch.Tensor) -> torch.Tensor:90        return self.decoder(z)91 92    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:93        mu, logvar = self.encode(x)94        z = self.reparameterize(mu, logvar)95        return self.decode(z), mu, logvar96