training/detectors/ae_detector.py

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)