abdulwaqar/Synthetic_Data_Generation

1import torch2import torch.nn as nn3 4LATENT_DIM = 1005IMAGE_SIZE = 646CHANNELS = 37 8 9class Generator(nn.Module):10    def __init__(self):11        super().__init__()12 13        self.model = nn.Sequential(14            nn.ConvTranspose2d(15                LATENT_DIM, 512, kernel_size=4, stride=1,16                padding=0, bias=False17            ),18            nn.BatchNorm2d(512),19            nn.ReLU(True),20 21            nn.ConvTranspose2d(22                512, 256, kernel_size=4, stride=2,23                padding=1, bias=False24            ),25            nn.BatchNorm2d(256),26            nn.ReLU(True),27 28            nn.ConvTranspose2d(29                256, 128, kernel_size=4, stride=2,30                padding=1, bias=False31            ),32            nn.BatchNorm2d(128),33            nn.ReLU(True),34 35            nn.ConvTranspose2d(36                128, 64, kernel_size=4, stride=2,37                padding=1, bias=False38            ),39            nn.BatchNorm2d(64),40            nn.ReLU(True),41 42            nn.ConvTranspose2d(43                64, CHANNELS, kernel_size=4, stride=2,44                padding=1, bias=False45            ),46            nn.Tanh()47        )48 49    def forward(self, x):50        return self.model(x)51 52 53def load_generator(checkpoint_path="generator.pth", device=None):54    if device is None:55        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")56 57    model = Generator().to(device)58 59    checkpoint = torch.load(60        checkpoint_path,61        map_location=device,62        weights_only=True63    )64 65    # Supports either a raw state_dict or a checkpoint dictionary.66    if isinstance(checkpoint, dict) and "state_dict" in checkpoint:67        checkpoint = checkpoint["state_dict"]68 69    # Remove DataParallel prefix if the model was saved with it.70    checkpoint = {71        key.replace("module.", "", 1) if key.startswith("module.") else key: value72        for key, value in checkpoint.items()73    }74 75    model.load_state_dict(checkpoint)76    model.eval()77 78    return model, device79