import os import datetime import logging import numpy as np from sklearn import metrics from typing import Union from collections import defaultdict import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.nn import DataParallel from torch.utils.tensorboard import SummaryWriter from torchvision.models import vgg16 import torchvision.transforms.functional as F_tv from metrics.base_metrics_class import calculate_metrics_for_train, calculate_acc_for_train from .base_detector import AbstractDetector from detectors import DETECTOR from networks import BACKBONE from loss import LOSSFUNC logger = logging.getLogger(__name__) # --------------------------------------------------- class Swish(nn.Module): """Swish activation commonly used in LDM""" def forward(self, x): return x * torch.sigmoid(x) def weights_init(m): """Initialize weights (for GAN and AE)""" classname = m.__class__.__name__ if classname.find('Conv') != -1: nn.init.normal_(m.weight.data, 0.0, 0.02) elif classname.find('BatchNorm') != -1 or classname.find('GroupNorm') != -1: nn.init.normal_(m.weight.data, 1.0, 0.02) nn.init.constant_(m.bias.data, 0) class PerceptualLoss(nn.Module): """LDM-style perceptual loss (VGG16 intermediate-layer features) - fixed for batch normalization issues""" def __init__(self): super().__init__() # Load pretrained VGG16 and freeze its parameters vgg = vgg16(pretrained=True).features[:10] for param in vgg.parameters(): param.requires_grad = False self.vgg = vgg # Store normalization parameters (use functional ops instead of instantiating transforms.Normalize) self.mean = torch.tensor([0.485, 0.456, 0.406]) self.std = torch.tensor([0.229, 0.224, 0.225]) self.mse = nn.MSELoss() def forward(self, x, x_recon): """ Fix: support normalization for batched data ([B,3,H,W]) x: original image ([B,3,H,W], ∈[-1,1]) x_recon: reconstructed image ([B,3,H,W], ∈[-1,1]) """ # 1. [-1,1] → [0,1] x = (x + 1) / 2 x_recon = (x_recon + 1) / 2 # 2. Move mean/std to the same device as x (GPU/CPU) mean = self.mean.to(x.device, dtype=x.dtype) std = self.std.to(x.device, dtype=x.dtype) # 3. Use torchvision.functional.normalize (supports batched data [B,C,H,W]) # Note: mean/std must be expanded to [1,3,1,1] to match the batched tensor dimensions x_norm = F_tv.normalize(x, mean=mean.view(1, 3, 1, 1), std=std.view(1, 3, 1, 1)) x_recon_norm = F_tv.normalize(x_recon, mean=mean.view(1, 3, 1, 1), std=std.view(1, 3, 1, 1)) # 4. Extract VGG features and compute MSE feat_x = self.vgg(x_norm) feat_recon = self.vgg(x_recon_norm) return self.mse(feat_x, feat_recon) class LDMEncoder(nn.Module): """LDM-style VAE encoder""" def __init__(self, in_channels=3, latent_channels=4, base_channels=64): super().__init__() self.encoder = nn.Sequential( nn.Conv2d(in_channels, base_channels, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels), Swish(), nn.Conv2d(base_channels, base_channels*2, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*2), Swish(), nn.Conv2d(base_channels*2, base_channels*4, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*4), Swish(), nn.Conv2d(base_channels*4, base_channels*8, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*8), Swish(), nn.Conv2d(base_channels*8, latent_channels, kernel_size=3, stride=1, padding=1) ) def forward(self, x): return self.encoder(x) class LDMDecoder(nn.Module): """LDM-style VAE decoder""" def __init__(self, out_channels=3, latent_channels=4, base_channels=64): super().__init__() self.decoder = nn.Sequential( nn.Conv2d(latent_channels, base_channels*8, kernel_size=3, stride=1, padding=1), nn.GroupNorm(32, base_channels*8), Swish(), nn.ConvTranspose2d(base_channels*8, base_channels*4, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*4), Swish(), nn.ConvTranspose2d(base_channels*4, base_channels*2, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*2), Swish(), nn.ConvTranspose2d(base_channels*2, base_channels, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels), Swish(), nn.ConvTranspose2d(base_channels, out_channels, kernel_size=4, stride=2, padding=1), nn.Tanh() ) def forward(self, z): return self.decoder(z) class PatchGANDiscriminator(nn.Module): """PatchGAN discriminator (used for adversarial loss)""" def __init__(self, in_channels=3, base_channels=64): super().__init__() self.discriminator = nn.Sequential( nn.Conv2d(in_channels, base_channels, kernel_size=4, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels, base_channels*2, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*2), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels*2, base_channels*4, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*4), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels*4, base_channels*8, kernel_size=4, stride=2, padding=1), nn.GroupNorm(32, base_channels*8), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels*8, 1, kernel_size=4, stride=1, padding=1) ) def forward(self, x): return self.discriminator(x) class AETotalLoss(nn.Module): """Combined loss: perceptual loss + MSE + adversarial loss + classification loss + LDM KL loss""" def __init__(self, lambda_perceptual=0, lambda_mse=0, lambda_adv=0, lambda_cls=1.0, lambda_kl=0): super().__init__() self.lambda_perceptual = lambda_perceptual self.lambda_mse = lambda_mse self.lambda_adv = lambda_adv self.lambda_cls = lambda_cls self.lambda_kl = lambda_kl # initial KL weight self.perceptual_loss = PerceptualLoss() self.mse = nn.MSELoss() self.ce = nn.CrossEntropyLoss() def forward(self, x, x_recon, pred, labels, z, d_fake): """ Generator loss: excludes the discriminator real branch (discriminator is computed separately) Args: d_fake: discriminator output for the reconstructed image (already detached or from an independent computation graph) """ # 1. perceptual loss loss_perceptual = self.perceptual_loss(x, x_recon) # 2. MSE loss loss_mse = self.mse(x, x_recon) # 3. adversarial loss (generator part only; d_fake is already detached) loss_adv_gen = self.mse(d_fake, torch.ones_like(d_fake)) # 4. classification loss loss_cls = self.ce(pred, labels) # 5. simplified LDM KL loss (z is the mean μ) loss_kl = 0.5 * torch.sum(z ** 2, dim=[1,2,3]).mean() # total generator loss total_loss = ( self.lambda_perceptual * loss_perceptual + self.lambda_mse * loss_mse + self.lambda_adv * loss_adv_gen + self.lambda_cls * loss_cls + self.lambda_kl * loss_kl ) return { 'gen_total': loss_cls, 'loss_perceptual': loss_perceptual, 'loss_mse': loss_mse, 'loss_adv_gen': loss_adv_gen, 'loss_cls': loss_cls, 'loss_kl': loss_kl } class DiscriminatorLoss(nn.Module): """Standalone discriminator loss (computed separately to avoid conflicts with the generator graph)""" def __init__(self): super().__init__() self.mse = nn.MSELoss() def forward(self, d_real, d_fake): """ discriminator loss (LSGAN) Args: d_real: discriminator output for the original image d_fake: discriminator output for the reconstructed image (detach the generator graph) """ loss_real = self.mse(d_real, torch.ones_like(d_real)) loss_fake = self.mse(d_fake.detach(), torch.zeros_like(d_fake)) return 0.5 * (loss_real + loss_fake) # -------------------------- Bench-adapted detector (core class) -------------------------- @DETECTOR.register_module(module_name='ae_detector') class LDM_AE_Classify_Detector(AbstractDetector): def __init__(self, config): super().__init__() self.config = config self.num_classes = config['backbone_config']['num_classes'] self.latent_channels = config.get('latent_channels', 1024) self.base_channels = config.get('base_channels', 64) self.in_channels = config.get('in_channels', 3) self.out_channels = config.get('out_channels', 3) # KL annealing configuration self.kl_annealing_steps = config.get('kl_annealing_steps', 2000) self.max_lambda_kl = config.get('lambda_kl', 0.001) self.current_step = 0 # global iteration step # Build network components self.encoder, self.decoder, self.classifier_head, self.discriminator = self.build_backbone(config) # Build the loss function (separate generator and discriminator losses) self.gen_loss_func, self.disc_loss_func = self.build_loss(config) # Logging configuration logger.info(f"LDM-AE Detector initialized successfully:") logger.info(f" - number of classes: {self.num_classes}") logger.info(f" - latent-space channel count: {self.latent_channels}") logger.info(f" - KL annealing steps: {self.kl_annealing_steps}") logger.info(f" - maximum KL weight: {self.max_lambda_kl}") def build_backbone(self, config): """Build the AE encoder, decoder, classification head, and discriminator""" # 1. LDM encoder encoder = LDMEncoder( in_channels=self.in_channels, latent_channels=self.latent_channels, base_channels=self.base_channels ) # 2. LDM decoder decoder = LDMDecoder( out_channels=self.out_channels, latent_channels=self.latent_channels, base_channels=self.base_channels ) # 3. Latent-space classification head (renamed to classifier_head to avoid conflict with the abstract classifier method) classifier_head = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(self.latent_channels, 256), nn.LayerNorm(256), Swish(), nn.Dropout(0.5), nn.Linear(256, self.num_classes) ) # 4. PatchGAN discriminator discriminator = PatchGANDiscriminator(in_channels=self.in_channels) # Initialize weights encoder.apply(weights_init) decoder.apply(weights_init) classifier_head.apply(weights_init) discriminator.apply(weights_init) return encoder, decoder, classifier_head, discriminator def build_loss(self, config): """Build separate loss functions for the generator and discriminator""" # generator loss gen_loss_kwargs = { 'lambda_perceptual': config.get('lambda_perceptual', 0.8), 'lambda_mse': config.get('lambda_mse', 0.2), 'lambda_adv': config.get('lambda_adv', 0.1), 'lambda_cls': config.get('lambda_cls', 1.0), 'lambda_kl': config.get('lambda_kl', 0.001) } gen_loss_func = AETotalLoss(**gen_loss_kwargs) # discriminator loss (independent) disc_loss_func = DiscriminatorLoss() return gen_loss_func, disc_loss_func def features(self, data_dict: dict) -> tuple: """Extract features: encode + decode, and return latent features and the reconstructed image""" x = data_dict['image'] # Input: [B,3,H,W] ∈[-1,1] z = self.encoder(x) # Encode: [B,4,H/16,W/16] x_recon = self.decoder(z) # Decode: [B,3,H,W] ∈[-1,1] return z, x_recon def classifier(self, z: torch.Tensor) -> torch.Tensor: """Implement the abstract classifier method (matching the AbstractDetector interface)""" return self.classifier_head(z) def get_gen_losses(self, data_dict: dict, pred_dict: dict) -> dict: """Compute generator loss (independent of the discriminator)""" # 1. Extract inputs and predictions x = data_dict['image'] # original image labels = data_dict['label'] # ground-truth labels x_recon = pred_dict['x_recon'] # reconstructed image pred = pred_dict['cls'] # classification prediction z = pred_dict['z'] # latent features d_fake = pred_dict['d_fake'] # discriminator output for the reconstructed image (already detached) # 2. KL annealing: dynamically adjust the KL weight if self.current_step < self.kl_annealing_steps: self.gen_loss_func.lambda_kl = self.max_lambda_kl * (self.current_step / self.kl_annealing_steps) else: self.gen_loss_func.lambda_kl = self.max_lambda_kl self.current_step += 1 # 3. Compute generator loss loss_dict = self.gen_loss_func(x, x_recon, pred, labels, z, d_fake) # 4. Assemble the loss dictionary return { 'overall': loss_dict['gen_total'], # total generator loss 'loss_perceptual': loss_dict['loss_perceptual'].item(), 'loss_mse': loss_dict['loss_mse'].item(), 'loss_adv_gen': loss_dict['loss_adv_gen'].item(), 'loss_cls': loss_dict['loss_cls'].item(), 'loss_kl': loss_dict['loss_kl'].item(), 'lambda_kl': self.gen_loss_func.lambda_kl } def get_disc_losses(self, data_dict: dict, pred_dict: dict) -> torch.Tensor: """Compute discriminator loss (fully independent)""" d_real = pred_dict['d_real'] # discriminator output for the original image d_fake = pred_dict['d_fake'] # discriminator output for the reconstructed image return self.disc_loss_func(d_real, d_fake) def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict: """Compute training metrics (adapted to the Bench interface)""" labels = data_dict['label'].detach() pred = pred_dict['cls'].detach() acc, mAP = calculate_acc_for_train(labels, pred, self.num_classes) return {'acc': acc, 'mAP': mAP} def forward(self, data_dict: dict, inference=False) -> dict: """Forward pass (separate generator and discriminator computation graphs)""" x = data_dict['image'] # 1. Generator forward pass (encode + decode + classify) z, x_recon = self.features(data_dict) pred = self.classifier(z) prob = torch.softmax(pred, dim=1) # 2. Discriminator forward pass (computed independently to avoid coupling with the generator graph) with torch.no_grad(): # Do not compute discriminator gradients during generator inference d_real = self.discriminator(x) # discriminate the original image d_fake = self.discriminator(x_recon) # discriminate the reconstructed image # 3. Build the prediction dictionary (including discriminator outputs for loss computation) pred_dict = { 'cls': pred, # classification logits 'prob': prob, # classification probabilities 'z': z, # latent features 'x_recon': x_recon, # reconstructed image 'feat': z.mean(dim=[2,3]), # global mean of latent features 'd_real': d_real, # discriminator real output 'd_fake': d_fake # discriminator fake output } # Visualization during inference # if inference and hasattr(self, 'writer'): # self.writer.add_images('recon/image_gt', (x + 1)/2, global_step=self.current_step, dataformats='NCHW') # self.writer.add_images('recon/image_pred', (x_recon + 1)/2, global_step=self.current_step, dataformats='NCHW') return pred_dict def get_discriminator(self): """Provide an accessor for the discriminator (the Bench training loop needs to optimize it separately)""" return self.discriminator def get_losses(self, data_dict: dict, pred_dict: dict) -> dict: """Bench-compatible interface: return generator loss (the discriminator loss is obtained separately via get_disc_losses)""" return self.get_gen_losses(data_dict, pred_dict)