src/train_vision_model.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader

def train_model():
    print("Configuration du Device...")
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"Utilisation de l'appareil: {device}")

    # Transformations obligatoires
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ColorJitter(brightness=0.3, contrast=0.3),
            transforms.RandomRotation(15),
            transforms.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0)),
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    data_dir = 'data/images'
    
    print("Préparation des DataLoader...")
    # datasets.ImageFolder prend l'arborescence et associe fake/ et real/ à des classes 0 et 1.
    image_datasets = {
        x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
        for x in ['train', 'test']
    }
    
    dataloaders = {
        x: DataLoader(image_datasets[x], batch_size=16, shuffle=(True if x == 'train' else False), num_workers=0)
        for x in ['train', 'test']
    }
    
    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'test']}
    class_names = image_datasets['train'].classes
    
    print(f"Classes détectées : {class_names}")
    print(f"Tailles des sets : Train={dataset_sizes['train']}, Test={dataset_sizes['test']}")

    # Chargement du modèle ResNet18 préentraîné
    print("Initialisation du modèle torchvision.models.resnet18(pretrained=True)...")
    # Pour éviter le warning deprecation, on peut utiliser weights=ResNet18_Weights.DEFAULT,
    # mais pretrained=True est explicitement demandé et supporté.
    model = models.resnet18(pretrained=True)

    # Geler les paramètres de base
    for param in model.parameters():
        param.requires_grad = False

    # Remplacement de la dernière couche linéaire adaptée à 2 sorties
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, 2)
    
    model = model.to(device)

    # Définition de l'optimiseur (uniquement sur les paramètres de la couche dense)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.fc.parameters(), lr=0.001)

    num_epochs = 10
    print(f"\nDémarrage de l'entraînement sur {num_epochs} epochs...")

    for epoch in range(num_epochs):
        print(f"Epoch {epoch+1}/{num_epochs}")
        print("-" * 10)

        for phase in ['train', 'test']:
            if phase == 'train':
                model.train()  # Phase d'entraînement
            else:
                model.eval()   # Phase d'évaluation

            running_loss = 0.0
            running_corrects = 0

            # Itérations sur les données
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # Backpropagation et optimisation seulement dans train
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # Statistiques
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f"{phase.capitalize()} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}")

    # Évaluation finale (Accuracy Finale)
    print("\nÉvaluation finale terminée.")
    model.eval()
    final_corrects = 0
    with torch.no_grad():
        for inputs, labels in dataloaders['test']:
            inputs = inputs.to(device)
            labels = labels.to(device)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            final_corrects += torch.sum(preds == labels.data)
            
    final_acc = final_corrects.double() / dataset_sizes['test']
    print("=" * 40)
    print(f"RÉSULTATS DE L'ÉVALUATION (Test)")
    print("=" * 40)
    print(f"Accuracy finale sur le jet de test : {final_acc * 100:.2f} %")
    print("=" * 40)

    print("\nSauvegarde du modèle...\n")
    os.makedirs('models', exist_ok=True)
    save_path = 'models/robust_vision_model.pth'
    
    # On sauvegarde les poids du modèle (recommandé en PyTorch)
    torch.save(model.state_dict(), save_path)
    print(f"Modèle sauvegardé dans : {save_path}")
    print("Terminé avec succès !")

if __name__ == '__main__':
    train_model()