Upload main.py

main.py

@@ -0,0 +1,369 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, models, transforms
+from torch.utils.data import DataLoader, random_split, Dataset
+from torch.optim import lr_scheduler
+import pandas as pd
+import numpy as np
+import time
+import copy
+import os
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import confusion_matrix, classification_report, roc_auc_score
+from sklearn.preprocessing import label_binarize
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+root_dir = r"path"
+data_dir = os.path.join(root_dir, 'Training')
+save_dir = "./improved_results"
+os.makedirs(save_dir, exist_ok=True)
+
+CONFIG = {
+    'model_name': 'ResNet50_Improved',
+    'batch_size': 32,
+    'lr': 0.001,
+    'epochs': 25,
+    'scheduler_step': 7,
+    'gamma': 0.1,
+    'weight_decay': 5e-4, 
+    'dropout_rate': 0.6,   
+    'early_stopping_patience': 5,
+    'early_stopping_min_delta': 0.001
+}
+
+train_transforms = transforms.Compose([
+    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),  
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomRotation(20),  
+    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2), 
+    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),  
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+test_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+
+class EarlyStopping:
+    def __init__(self, patience=5, min_delta=0.001, mode='max'):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.mode = mode
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+        
+    def __call__(self, score):
+        if self.best_score is None:
+            self.best_score = score
+            return False
+            
+        if self.mode == 'max':
+            if score > self.best_score + self.min_delta:
+                self.best_score = score
+                self.counter = 0
+            else:
+                self.counter += 1
+        else:
+            if score < self.best_score - self.min_delta:
+                self.best_score = score
+                self.counter = 0
+            else:
+                self.counter += 1
+                
+        if self.counter >= self.patience:
+            self.early_stop = True
+            print(f"\nEarly stopping triggered! No improvement for {self.patience} epochs.")
+            return True
+        return False
+
+
+class TransformedSubset(Dataset):
+    def __init__(self, subset, transform=None):
+        self.subset = subset
+        self.transform = transform
+
+    def __getitem__(self, index):
+        x, y = self.subset[index]
+        if self.transform:
+            x = self.transform(x)
+        return x, y
+
+    def __len__(self):
+        return len(self.subset)
+
+
+base_dataset = datasets.ImageFolder(root=data_dir)
+class_names = base_dataset.classes
+num_classes = len(class_names)
+
+train_size = int(0.8 * len(base_dataset))
+test_size = len(base_dataset) - train_size
+train_indices, test_indices = random_split(base_dataset, [train_size, test_size])
+
+train_dataset = TransformedSubset(train_indices, transform=train_transforms)
+test_dataset = TransformedSubset(test_indices, transform=test_transforms)
+
+dataloaders = {
+    'train': DataLoader(train_dataset, batch_size=CONFIG['batch_size'], shuffle=True, num_workers=0),
+    'test': DataLoader(test_dataset, batch_size=CONFIG['batch_size'], shuffle=False, num_workers=0)
+}
+dataset_sizes = {'train': train_size, 'test': test_size}
+
+
+def get_model():
+    model = models.resnet50(pretrained=True)
+    num_ftrs = model.fc.in_features
+    model.fc = nn.Sequential(
+        nn.Dropout(CONFIG['dropout_rate']),
+        nn.Linear(num_ftrs, num_classes)
+    )
+    return model.to(device)
+
+
+model = get_model()
+criterion = nn.CrossEntropyLoss()
+optimizer = optim.Adam(
+    model.parameters(), 
+    lr=CONFIG['lr'], 
+    weight_decay=CONFIG['weight_decay']
+)
+exp_lr_scheduler = lr_scheduler.StepLR(
+    optimizer, 
+    step_size=CONFIG['scheduler_step'], 
+    gamma=CONFIG['gamma']
+)
+
+
+def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
+    since = time.time()
+    best_model_wts = copy.deepcopy(model.state_dict())
+    best_acc = 0.0
+    history = []
+    
+    early_stopping = EarlyStopping(
+        patience=CONFIG['early_stopping_patience'],
+        min_delta=CONFIG['early_stopping_min_delta'],
+        mode='max'
+    )
+
+    for epoch in range(num_epochs):
+        print(f'\n{"="*50}')
+        print(f'Epoch {epoch+1}/{num_epochs}')
+        print("="*50)
+        
+        epoch_stats = {'Epoch': epoch+1}
+
+        for phase in ['train', 'test']:
+            if phase == 'train':
+                model.train()
+            else:
+                model.eval()
+
+            running_loss = 0.0
+            running_corrects = 0
+
+            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)
+
+                    if phase == 'train':
+                        loss.backward()
+                        optimizer.step()
+
+                running_loss += loss.item() * inputs.size(0)
+                running_corrects += torch.sum(preds == labels.data)
+
+            if phase == 'train':
+                scheduler.step()
+
+            epoch_loss = running_loss / dataset_sizes[phase]
+            epoch_acc = running_corrects.double() / dataset_sizes[phase]
+
+            print(f'{phase.upper():5s} | Loss: {epoch_loss:.4f} | Acc: {epoch_acc:.4f} ({epoch_acc*100:.2f}%)')
+            
+            epoch_stats[f'{phase}_loss'] = epoch_loss
+            epoch_stats[f'{phase}_acc'] = epoch_acc.item()
+
+            if phase == 'test':
+                if epoch_acc > best_acc:
+                    best_acc = epoch_acc
+                    best_model_wts = copy.deepcopy(model.state_dict())
+                    torch.save(model.state_dict(), os.path.join(save_dir, 'best_model.pth'))
+                    print(f"✅ New Record! Test Acc: {best_acc:.4f}")
+                
+                if early_stopping(epoch_acc.item()):
+                    print(f"\nTraining stopped (Epoch {epoch+1})")
+                    model.load_state_dict(best_model_wts)
+                    df = pd.DataFrame(history)
+                    df.to_csv(os.path.join(save_dir, 'training_logs.csv'), index=False)
+                    return model, df
+
+        history.append(epoch_stats)
+
+    time_elapsed = time.time() - since
+    print(f'\n{"="*50}')
+    print(f'Training completed: {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
+    print(f'Best Test Accuracy: {best_acc:.4f} ({best_acc*100:.2f}%)')
+    print("="*50)
+
+    model.load_state_dict(best_model_wts)
+    df = pd.DataFrame(history)
+    df.to_csv(os.path.join(save_dir, 'training_logs.csv'), index=False)
+    
+    return model, df
+
+
+def evaluate_model(model, dataloader, class_names):
+    model.eval()
+    all_preds = []
+    all_labels = []
+    all_probs = []
+    
+    with torch.no_grad():
+        for inputs, labels in dataloader:
+            inputs = inputs.to(device)
+            outputs = model(inputs)
+            probs = torch.softmax(outputs, dim=1)
+            _, preds = torch.max(outputs, 1)
+            
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.numpy())
+            all_probs.extend(probs.cpu().numpy())
+    
+    all_preds = np.array(all_preds)
+    all_labels = np.array(all_labels)
+    all_probs = np.array(all_probs)
+    
+    cm = confusion_matrix(all_labels, all_preds)
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                xticklabels=class_names, yticklabels=class_names,
+                cbar_kws={'label': 'Count'})
+    plt.title('Confusion Matrix', fontsize=16, fontweight='bold')
+    plt.ylabel('True Class', fontsize=12)
+    plt.xlabel('Predicted Class', fontsize=12)
+    plt.tight_layout()
+    plt.savefig(os.path.join(save_dir, 'confusion_matrix.png'), dpi=300)
+    plt.show()
+    
+    print("\n" + "="*60)
+    print("DETAILED PERFORMANCE REPORT")
+    print("="*60)
+    report = classification_report(all_labels, all_preds, 
+                                   target_names=class_names, 
+                                   digits=4)
+    print(report)
+    
+    report_dict = classification_report(all_labels, all_preds, 
+                                       target_names=class_names, 
+                                       output_dict=True)
+    
+    metrics = ['precision', 'recall', 'f1-score']
+    class_metrics = {metric: [] for metric in metrics}
+    
+    for class_name in class_names:
+        for metric in metrics:
+            class_metrics[metric].append(report_dict[class_name][metric])
+    
+    fig, ax = plt.subplots(figsize=(12, 6))
+    x = np.arange(len(class_names))
+    width = 0.25
+    
+    for i, metric in enumerate(metrics):
+        ax.bar(x + i*width, class_metrics[metric], width, 
+               label=metric.capitalize(), alpha=0.8)
+    
+    ax.set_xlabel('Classes', fontsize=12)
+    ax.set_ylabel('Score', fontsize=12)
+    ax.set_title('Per-Class Performance Metrics', fontsize=14, fontweight='bold')
+    ax.set_xticks(x + width)
+    ax.set_xticklabels(class_names, rotation=45, ha='right')
+    ax.legend()
+    ax.set_ylim([0, 1.05])
+    ax.grid(axis='y', alpha=0.3)
+    plt.tight_layout()
+    plt.savefig(os.path.join(save_dir, 'class_metrics.png'), dpi=300)
+    plt.show()
+    
+    try:
+        y_bin = label_binarize(all_labels, classes=range(num_classes))
+        auc_scores = []
+        for i in range(num_classes):
+            auc = roc_auc_score(y_bin[:, i], all_probs[:, i])
+            auc_scores.append(auc)
+            print(f"ROC-AUC ({class_names[i]}): {auc:.4f}")
+        print(f"Mean ROC-AUC: {np.mean(auc_scores):.4f}")
+    except:
+        print("ROC-AUC could not be calculated")
+    
+    return cm, report
+
+
+def plot_training_results(df):
+    sns.set_style("whitegrid")
+    fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+    
+    axes[0].plot(df['Epoch'], df['train_loss'], 'o-', label='Train Loss', linewidth=2, markersize=6)
+    axes[0].plot(df['Epoch'], df['test_loss'], 's-', label='Test Loss', linewidth=2, markersize=6)
+    axes[0].set_title('Loss Evolution', fontsize=14, fontweight='bold')
+    axes[0].set_xlabel('Epoch', fontsize=12)
+    axes[0].set_ylabel('Loss', fontsize=12)
+    axes[0].legend(fontsize=11)
+    axes[0].grid(True, alpha=0.3)
+    
+    axes[1].plot(df['Epoch'], df['train_acc'], 'o-', label='Train Acc', linewidth=2, markersize=6, color='green')
+    axes[1].plot(df['Epoch'], df['test_acc'], 's-', label='Test Acc', linewidth=2, markersize=6, color='orange')
+    axes[1].set_title('Accuracy Evolution', fontsize=14, fontweight='bold')
+    axes[1].set_xlabel('Epoch', fontsize=12)
+    axes[1].set_ylabel('Accuracy', fontsize=12)
+    axes[1].legend(fontsize=11)
+    axes[1].grid(True, alpha=0.3)
+    axes[1].set_ylim([0, 1.05])
+    
+    plt.tight_layout()
+    plt.savefig(os.path.join(save_dir, 'training_curves.png'), dpi=300)
+    plt.show()
+    
+    df['overfit_gap'] = df['train_acc'] - df['test_acc']
+    print(f"\nOverfitting Analysis:")
+    print(f"Mean Train-Test Gap: {df['overfit_gap'].mean():.4f}")
+    print(f"Max Gap: {df['overfit_gap'].max():.4f} (Epoch {df.loc[df['overfit_gap'].idxmax(), 'Epoch']:.0f})")
+
+
+print("\nStarting training...\n")
+model_ft, logs = train_model(
+    model, criterion, optimizer, exp_lr_scheduler, 
+    num_epochs=CONFIG['epochs']
+)
+
+print("\nVisualizing results...")
+plot_training_results(logs)
+
+print("\nPerforming detailed evaluation...")
+cm, report = evaluate_model(model_ft, dataloaders['test'], class_names)
+
+print("\n" + "="*60)
+print("SUMMARY REPORT")
+print("="*60)
+print(f"Model: {CONFIG['model_name']}")
+print(f"Total Epochs: {len(logs)}")
+print(f"Best Test Accuracy: {logs['test_acc'].max():.4f} ({logs['test_acc'].max()*100:.2f}%)")
+print(f"Final Test Accuracy: {logs['test_acc'].iloc[-1]:.4f}")
+print(f"Final Train Accuracy: {logs['train_acc'].iloc[-1]:.4f}")
+print(f"Overfitting Gap: {logs['train_acc'].iloc[-1] - logs['test_acc'].iloc[-1]:.4f}")
+print(f"\nAll results saved to '{save_dir}'")
+print("="*60)