train_DevNet.py

import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from torchvision.datasets import ImageFolder
from sklearn.utils.class_weight import compute_class_weight
from sklearn.metrics import classification_report, confusion_matrix, f1_score, accuracy_score, precision_score, recall_score
import matplotlib.pyplot as plt
import argparse
import json
import time
from PIL import Image
from model.DevNet import *
from sklearn.metrics import roc_curve
from tqdm import tqdm
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# Argument parser
parser = argparse.ArgumentParser(description='Train or test a CNN model.')
parser.add_argument('--mode', type=str, choices=['train', 'test'], required=True, help='Choose to train or test the model.')
parser.add_argument('--h', type=int, default=128, help='Height of the input image')
parser.add_argument('--w', type=int, default=128, help='Width of the input image')
parser.add_argument('--model', type=str, choices=['DevNet',], required=False, help='Specify the model architecture.')
parser.add_argument('--dataset', type=str, default=r'K:\QYPAD\dataset\palm\train',required=False, help='Specify the dataset name.')
parser.add_argument('--threshold', type=float, default=2.50, help='Specify the threshold for classifying predictions (default: 2.50)')

    
class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.samples = []
        
        # Walk through all subdirectories (two levels deep)
        for dirpath, dirnames, filenames in os.walk(root_dir):
            for filename in filenames:
                if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  # Adjust file types as needed
                    label = dirpath.split(os.path.sep)[-3]  # Get the label from the second to last folder (A/B)
                    self.samples.append((os.path.join(dirpath, filename), label))
                    # print((os.path.join(dirpath, filename), label))
        
        # self.samples = self.samples[::10]
        self.classes = ['spoof', 'live']  # Adjust this based on your actual labels
        print(f"root_dir: {root_dir}")
        print(f"total: {len(self.samples)}")
        print("type:")
        for cls in self.classes:
            count = sum(1 for _, label in self.samples if label == cls)
            print(f"  {cls}: {count} ({count/len(self.samples)*100:.1f}%)")
        
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        image = Image.open(img_path).convert('RGB')
        label = 0 if label == 'spoof' else 1  # Convert labels to 0/1 if binary classification
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

def generate_class_weights(data_dir):
    dataset = ImageFolder(data_dir)
    labels = [label for _, label in dataset.samples]
    class_weights = compute_class_weight(class_weight='balanced', classes=np.unique(labels), y=labels)
    return torch.FloatTensor(class_weights)

def get_model(model_name):
    if model_name == 'DevNet':
        model = SemiADNet()
    else:
        raise ValueError(f"Model {model_name} not supported")
    
    print(model_name)
    # exit()

    return model

def find_optimal_threshold(model, val_loader, device):

    all_outputs = []
    all_labels = []
    
    model.eval()
    print("Finding optimal threshold on validation set...")
    with torch.no_grad():
        for inputs, labels in tqdm(val_loader):
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs).squeeze()
            all_outputs.extend(outputs.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    all_outputs = np.array(all_outputs)
    all_labels = np.array(all_labels)
    
    live_scores = all_outputs[all_labels == 1]
    spoof_scores = all_outputs[all_labels == 0]

    live_min, live_max = np.min(live_scores), np.max(live_scores)
    spoof_min, spoof_max = np.min(spoof_scores), np.max(spoof_scores)

    threshold_min = max(0, min(live_min, spoof_min) * 0.9)  
    threshold_max = max(live_max, spoof_max) * 1.1  
    
    print(f"Score statistics - Live: min={live_min:.4f}, max={live_max:.4f}")
    print(f"Score statistics - Spoof: min={spoof_min:.4f}, max={spoof_max:.4f}")
    print(f"Setting threshold search range: [{threshold_min:.4f}, {threshold_max:.4f}]")

    plt.figure(figsize=(10, 6))
    plt.hist(live_scores, bins=50, alpha=0.5, label='Live')
    plt.hist(spoof_scores, bins=50, alpha=0.5, label='Spoof')
    plt.xlabel('Score')
    plt.ylabel('Frequency')
    plt.legend()
    plt.savefig(os.path.join('checkpoint/DevNet', 'score_distribution.png'))
    
    thresholds = np.linspace(threshold_min, threshold_max, 100)

    min_eer_diff = float('inf')  
    optimal_threshold_eer = 0.5  
    best_eer = 1.0  

    min_acer = float('inf')
    optimal_threshold_acer = 0.5

    best_metrics = {}
    best_metrics_eer = {}
    
    
    for threshold in thresholds:

        predictions = (all_outputs >= threshold).astype(int)

        real_indices = (all_labels == 1)  
        fake_indices = (all_labels == 0)  
        
        if np.sum(real_indices) > 0 and np.sum(fake_indices) > 0:
            
            bpcer = np.sum((predictions == 0) & real_indices) / np.sum(real_indices)

            apcer = np.sum((predictions == 1) & fake_indices) / np.sum(fake_indices)
            
            acer = (apcer + bpcer) / 2

            eer_diff = np.abs(apcer - bpcer)

            metrics = {
                "threshold": float(threshold),
                "apcer": float(apcer),
                "bpcer": float(bpcer),
                "acer": float(acer),
                "eer_diff": float(eer_diff)
            }

            print(f"Threshold: {threshold:.4f}\t, APCER: {apcer:.4f}, BPCER: {bpcer:.4f}\t, |APCER-BPCER|: {eer_diff:.4f}\t, ACER: {acer:.4f}")

            if eer_diff < min_eer_diff:
                min_eer_diff = eer_diff
                optimal_threshold_eer = threshold
                best_eer = (apcer + bpcer) / 2  
                best_metrics_eer = metrics.copy()
                best_metrics_eer["eer"] = best_eer

            if acer < min_acer:
                min_acer = acer
                optimal_threshold_acer = threshold
                best_metrics = metrics.copy()

    print("\nEqual Error Rate (EER) Statistics:")
    print(f"Optimal Threshold for EER: {optimal_threshold_eer:.4f}")
    print(f"At this threshold - APCER: {best_metrics_eer['apcer']:.4f}, BPCER: {best_metrics_eer['bpcer']:.4f}")
    print(f"EER: {best_eer:.4f}, |APCER-BPCER|: {min_eer_diff:.6f}")

    print("\nMinimal ACER Statistics:")
    print(f"Optimal Threshold for min ACER: {optimal_threshold_acer:.4f}")
    print(f"At this threshold - APCER: {best_metrics['apcer']:.4f}, BPCER: {best_metrics['bpcer']:.4f}")
    print(f"ACER: {min_acer:.4f}")

    return optimal_threshold_eer, best_metrics_eer


def train_model(dataset, model_choice):
    # Data directories
    data_dir = f'{dataset}/'
    train_dir = os.path.join(data_dir, 'train')
    val_dir = os.path.join(data_dir, 'validation')
    
    # Data transforms
    train_transform = transforms.Compose([
        transforms.Resize((h, w)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(20),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize((h, w)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    # class_to_label = {'real': 1, 'fake': 0}
    train_dataset = CustomDataset(root_dir=train_dir, transform=train_transform)
    val_dataset = CustomDataset(root_dir=val_dir, transform=val_transform)
    
    # Create data loaders
    train_loader = DataLoader(train_dataset, batch_size=2048, shuffle=True, num_workers=16)
    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False, num_workers=4)
    
    # Get model
    model = get_model(model_choice)
    model = model.to(device)

    # Training loop
    best_metric = float('inf')  
    patience = 5
    patience_counter = 0
    epochs = 50

    # Loss function and optimizer
    criterion = DeviationLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
    
    # Create model directory
    model_dir = f'./checkpoint/{model_choice}/'
    os.makedirs(model_dir, exist_ok=True)
    

    train_losses = []
    train_accs = []
    val_losses = []
    val_accs = []
    val_acers = []
    
    for epoch in range(epochs):
        model.train()
        train_loss = 0
        train_correct = 0
        train_total = 0
        
        for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} - Training"):
            inputs = inputs.to(device)
            labels = labels.float().to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs).squeeze()
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            predicted = (outputs >= args.threshold).float()
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()

        train_acc = 100 * train_correct / train_total
        train_losses.append(train_loss / len(train_loader))
        train_accs.append(train_acc)
        
        # Validation
        model.eval()
        val_loss = 0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            
            for inputs, labels in tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} - Validation"):
                inputs = inputs.to(device)
                labels = labels.float().to(device)
                
                outputs = model(inputs).squeeze()
                loss = criterion(outputs, labels)
                
                val_loss += loss.item()
                predicted = (outputs >= args.threshold).float()
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()
        
        val_acc = 100 * val_correct / val_total
        val_losses.append(val_loss / len(val_loader))
        val_accs.append(val_acc)

        curr_threshold, curr_val_metrics = find_optimal_threshold(model, val_loader, device)
        curr_acer = curr_val_metrics["acer"]
        val_acers.append(curr_acer)

        print(f'Epoch [{epoch+1}/{epochs}]: '
              f'Train Loss: {train_losses[-1]:.4f}, Train Acc: {train_acc:.2f}%, '
              f'Val Loss: {val_losses[-1]:.4f}, Val Acc: {val_acc:.2f}%, '
              f'Val ACER: {curr_acer:.4f}, Threshold: {curr_threshold:.4f}')

        scheduler.step(curr_acer)

        if curr_acer < best_metric:
            best_metric = curr_acer
            print(f"New best model found with ACER: {best_metric:.4f}")
            torch.save(model.state_dict(), os.path.join(model_dir, 'model.pth'))

            with open(os.path.join(model_dir, 'threshold.json'), 'w') as f:
                json.dump(curr_val_metrics, f, indent=4)
                
            patience_counter = 0
        else:
            patience_counter += 1

        if patience_counter >= patience:
            print(f"Early stopping triggered after {epoch+1} epochs")
            break

    plt.figure(figsize=(15, 5))

    plt.subplot(1, 3, 1)
    plt.plot(train_accs, label='Training Accuracy')
    plt.plot(val_accs, label='Validation Accuracy')
    plt.title('Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.legend()

    plt.subplot(1, 3, 2)
    plt.plot(train_losses, label='Training Loss')
    plt.plot(val_losses, label='Validation Loss')
    plt.title('Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()

    plt.subplot(1, 3, 3)
    plt.plot(val_acers, label='Validation ACER', color='red')
    plt.title('ACER (Lower is Better)')
    plt.xlabel('Epoch')
    plt.ylabel('ACER')
    plt.legend()
    
    plt.tight_layout()
    plt.savefig(os.path.join(model_dir, 'training_history.png'))
    
    print(f"Training completed. Best validation ACER: {best_metric:.4f}")
    print(f"Model saved to {os.path.join(model_dir, 'model.pth')}")
    print(f"Threshold and metrics saved to {os.path.join(model_dir, 'threshold.json')}")

def test_model(dataset, model_choice):
    # Data directories and transforms
    data_dir = f'{dataset}/'
    test_dir = os.path.join(data_dir, 'test')
    
    test_transform = transforms.Compose([
        transforms.Resize((h, w)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    # Create test dataset and loader
    # class_to_label = {'real': 0, 'fake': 1}
    test_dataset = CustomDataset(root_dir=test_dir, transform=test_transform)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)
    
    # Load model
    model = get_model(model_choice)
    model_dir = f'./checkpoint/{model_choice}/'
    model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pth')))
    model = model.to(device)
    model.eval()

    try:
        with open(os.path.join(model_dir, 'threshold.json'), 'r') as f:
            val_metrics = json.load(f)
        optimal_threshold = val_metrics["threshold"]
        print(f"Loaded optimal threshold from validation: {optimal_threshold}")
    except FileNotFoundError:
        print("Warning: No threshold.json found. Using default threshold of 0.5")
        optimal_threshold = 0.5
        
    # Test loop
    all_predictions = []
    all_labels = []
    
    print("Testing model on test set...")
    with torch.no_grad():
        for inputs, labels in tqdm(test_loader):
            inputs = inputs.to(device)
            outputs = model(inputs).squeeze()
            predictions = outputs.cpu().numpy()
            all_predictions.extend(predictions)
            all_labels.extend(labels.numpy())
    
    all_predictions = np.array(all_predictions)
    all_labels = np.array(all_labels)
    print(f"Test set distribution - Spoof: {np.sum(all_labels==0)}, Live: {np.sum(all_labels==1)}")

    y_pred = (all_predictions >= optimal_threshold).astype(int)

    accuracy = accuracy_score(all_labels, y_pred)
    f1 = f1_score(all_labels, y_pred, average='weighted')
    precision = precision_score(all_labels, y_pred, average='weighted')
    recall = recall_score(all_labels, y_pred, average='weighted')
    # conf_matrix = confusion_matrix(all_labels, y_pred)
    conf_matrix = confusion_matrix(all_labels, y_pred, labels=[0, 1])

    tn, fp, fn, tp = conf_matrix.ravel()
    apcer = fp / (fp + tn) if (fp + tn) > 0 else 0  
    bpcer = fn / (fn + tp) if (fn + tp) > 0 else 0  
    acer = (apcer + bpcer) / 2  

    test_results = {
        "model": model_choice,
        "threshold": float(optimal_threshold),
        "accuracy": float(accuracy),
        "f1_score": float(f1),
        "precision": float(precision),
        "recall": float(recall),
        "confusion_matrix": conf_matrix.tolist(),
        "apcer": float(apcer),
        "bpcer": float(bpcer),
        "acer": float(acer)
    }

    results_path = os.path.join(model_dir, 'test_results.json')
    with open(results_path, 'w') as results_file:
        json.dump(test_results, results_file, indent=4)
    
    print(f"Test completed. Results saved to {results_path}")
    print(f"Test metrics: APCER={apcer:.4f}, BPCER={bpcer:.4f}, ACER={acer:.4f}")
    print(f"Accuracy={accuracy:.4f}, F1-score={f1:.4f}")

    fpr, tpr, thresholds = roc_curve(all_labels, all_predictions)
    fnr = 1 - tpr
    eer_idx = np.argmin(np.abs(fnr - fpr))
    eer = fpr[eer_idx]  

    plt.figure(figsize=(10, 8))
    plt.plot(fpr, tpr, label=f'ROC curve (EER = {eer:.4f})')
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'ROC Curve - {model_choice}')
    plt.legend()
    plt.savefig(os.path.join(model_dir, 'roc_curve.png'))

    eer_info = {
        "eer": float(eer),
        "eer_threshold": float(thresholds[eer_idx])
    }
    with open(os.path.join(model_dir, 'eer_info.json'), 'w') as f:
        json.dump(eer_info, f, indent=4)
    
    print(f"EER on test set: {eer:.4f} at threshold {thresholds[eer_idx]:.4f}")
    print('Test completed successfully!')  

if __name__ == '__main__':
    args = parser.parse_args()
    # Device configuration
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    h, w = args.h, args.w
    if args.mode == 'train':
        selected_datasets = [args.dataset] if args.dataset else ['data_1', 'data_2']
        selected_models = [args.model] if args.model else [
            'convnext_large', 'densenet201', 'efficientnet_b7', 'efficientnet_v2_l',
            'mobilenet_v3_large', 'regnet_x_32gf', 'resnet152', 'vgg19',
        ]
        
        for dataset in selected_datasets:
            for model_choice in selected_models:
                model_dir = f"./checkpoint/{dataset}/{model_choice}/"
                if not os.path.exists(model_dir):
                    print(f"Training {model_choice} on {dataset}")
                    train_model(dataset, model_choice)
                else:
                    print(f"Model directory {model_dir} already exists. Skipping training.")
    
    elif args.mode == 'test':
        selected_datasets = [args.dataset] if args.dataset else ['data_1', 'data_2']
        selected_models = [args.model] if args.model else [
            'densenet201', 'efficientnet_b7', 
            'mobilenet_v3_large', 'vgg19'
        ]
        
        for dataset in selected_datasets:
            for model_choice in selected_models:
                test_model(dataset, model_choice)