CornViT - A Multi-stage CVT Framework

README.md

@@ -1,3 +1,97 @@
 ---
-license: apache-2.0
+language: en
+license: mit
+tags:
+  - keras
+  - tensorflow
+  - computer-vision
+  - image-processing
+  - corn-kernel-classification
+pipeline_tag: image-classification
+library_name: keras
 ---
+
+# CornViT
+
+A Multi-Stage Convolutional Vision Transformer Framework for Corn Kernel Analysis
+
+## Overview
+
+Three-stage hierarchical classification pipeline for automated corn kernel quality assessment:
+
+- **Stage 1**: Purity detection (Pure vs Impure)
+- **Stage 2**: Shape classification (Flat vs Round)
+- **Stage 3**: Embryo orientation (Up vs Down)
+
+## Architecture
+
+- **Model**: CvT-13 (384×384) with ImageNet-22k pretraining
+- **Framework**: PyTorch + Microsoft CvT
+- **Test Accuracy**: 93.8% (Stage 1), 94.1% (Stage 2), 91.1% (Stage 3)
+
+## Setup
+
+```bash
+# Clone repository
+git clone https://github.com/microsoft/CvT.git
+
+# Install dependencies
+pip install -r requirements.txt
+```
+
+## Training
+
+Each stage has independent training scripts:
+
+```bash
+python stage1/train_cvt13.py  # Purity classification
+python stage2/train_cvt13.py  # Shape classification
+python stage3/train_cvt13.py  # Embryo orientation
+```
+
+## Inference
+
+```bash
+python stage1/inference_cvt13.py
+python stage2/inference_cvt13.py
+python stage3/inference_cvt13.py
+```
+
+## Baselines
+
+ResNet50 and DenseNet121 baselines available in `baselines/`.
+
+## Structure
+
+```
+├── stage1/          # Purity classification
+├── stage2/          # Shape classification
+├── stage3/          # Embryo orientation
+└── preprocess/      # Data preprocessing scripts
+```
+
+## Requirements
+
+- Python 3.13+
+- PyTorch 2.9+
+- CUDA (optional, for GPU training)
+
+---
+
+## Citation
+If you use this code, models, or catalog in your research, please cite:
+
+```bibtex
+@Article{computers15010002,
+  AUTHOR = {Erukude, Sai Teja and Mascarenhas, Jane and Shamir, Lior},
+  TITLE = {CornViT: A Multi-Stage Convolutional Vision Transformer Framework for Hierarchical Corn Kernel Analysis},
+  JOURNAL = {Computers},
+  VOLUME = {15},
+  YEAR = {2026},
+  NUMBER = {1},
+  ARTICLE-NUMBER = {2},
+  URL = {https://www.mdpi.com/2073-431X/15/1/2},
+  ISSN = {2073-431X},
+  DOI = {10.3390/computers15010002}
+}
+```

config.json

@@ -0,0 +1,52 @@
+{
+    "model_type": "cvt",
+    "architectures": [
+        "CvTForImageClassification"
+    ],
+    "paper": {
+        "title": "CornViT: A Multi-Stage Convolutional Vision Transformer Framework for Hierarchical Corn Kernel Analysis",
+        "year": 2025,
+        "doi": "https://doi.org/10.3390/computers15010002"
+    },
+    "pipeline_tag": "image-classification",
+    "library_name": "keras",
+    "framework": "pytorch",
+    "image_size": 384,
+    "num_channels": 3,
+    "pretrained": "imagenet-22k",
+    "backbone": "cvt-13",
+    "hierarchical_pipeline": true,
+    "stages": [
+        {
+            "stage": 1,
+            "name": "purity_detection",
+            "labels": {
+                "0": "pure",
+                "1": "impure"
+            },
+            "num_labels": 2,
+            "accuracy": 0.938
+        },
+        {
+            "stage": 2,
+            "name": "shape_classification",
+            "labels": {
+                "0": "flat",
+                "1": "round"
+            },
+            "num_labels": 2,
+            "accuracy": 0.941
+        },
+        {
+            "stage": 3,
+            "name": "embryo_orientation",
+            "labels": {
+                "0": "up",
+                "1": "down"
+            },
+            "num_labels": 2,
+            "accuracy": 0.911
+        }
+    ],
+    "training_framework": "pytorch"
+}

preprocess/delete.py

@@ -0,0 +1,120 @@
+import os
+import random
+from   tqdm         import tqdm
+
+
+######################
+#
+######################
+def delete_files(dir1: str, dir2: str) -> bool:
+    """
+    Desc: 
+        This method compares two directories and deletes files if present.
+    Args:
+        dir1 (str): Path to the directory 1.
+        dir2 (str): Path to the directory 2.
+    Returns:
+        True, if the deletion was complete, otherwise False.
+    """
+    try:
+        if not os.path.isdir(dir1) or not os.path.isdir(dir2):
+            return False
+        
+        dir1_files = set(os.listdir(dir1))
+        dir2_files = set(os.listdir(dir2))
+        
+        for idx, file in enumerate(dir1_files):
+            print(f"Processing file {idx}...")
+            
+            file_path = os.path.join(dir1, file)
+            if os.path.isfile(file_path):
+                
+                if file in dir2_files:
+                    # Delete the file is it is present in dir2
+                    os.remove(file_path)
+                    
+        return True
+        
+    except Exception as delete_ex:
+        print(f"Deletion error: {delete_ex}.")
+        return False
+    
+
+######################
+#
+######################
+def delete_n_random_files(dir: str, n: int) -> bool: 
+    """
+    Desc: 
+        This method deletes 'n' random files from the provided directory.
+    Args:
+        dir (str): Path to the directory.
+        n (int): The number of random files to be deleted.
+    Returns:
+        True, if the deletion was complete, otherwise False.
+    """
+    try:
+        if not os.path.isdir(dir):
+            print(f"Directory '{dir}' does not exist.")
+            return False
+        
+        # Get all files (not directories) in the specified directory
+        all_files = [f for f in os.listdir(dir) if os.path.isfile(os.path.join(dir, f))]
+        
+        if len(all_files) < n:
+            print(f"Cannot delete '{n}' files, directory only contains '{len(all_files)}' files.")
+            return False
+        
+        files_to_delete = random.sample(all_files, n)
+        
+        for file in tqdm(files_to_delete):
+            file_path = os.path.join(dir, file)
+            os.remove(file_path)
+        
+        return True
+    
+    except Exception as delete_ex:
+        print(f"Error occurred while deleting: {str(delete_ex)}.")
+        return False
+
+
+######################
+#
+######################
+def delete_files_name_contains(dir: str, word: str) -> bool:
+    """
+    Desc:
+        Deletes all files in a directory whose filenames contain a specific word (case-insensitive).
+    Parameters:
+        dir (str): The directory to search for files.
+        word (str): Substring to search for in filenames.
+    Returns:
+        bool: True if deletion completes (even if no files matched), False if an error occurred.
+    """
+    try:
+        if not os.path.isdir(dir):
+            print(f"Directory '{dir}' does not exist.")
+            return False
+
+        all_files = [f for f in os.listdir(dir) if os.path.isfile(os.path.join(dir, f))]
+
+        for file in tqdm(all_files, desc="Deleting files"):
+            if word.lower() in file.lower():
+                file_path = os.path.join(dir, file)
+                os.remove(file_path)
+
+        return True
+
+    except Exception as delete_ex:
+        print(f"Error occurred while deleting: {str(delete_ex)}.")
+        return False
+
+
+######################
+#
+######################
+if __name__ == "__main__":
+    
+    dir1 = ""
+    dir2 = ""
+    delete_files(dir1, dir2)

preprocess/move.py

@@ -0,0 +1,153 @@
+import os
+import random
+import shutil
+from   tqdm         import tqdm
+
+
+######################
+#
+######################
+def move_n_random_files(dir1: str, dir2: str, n: int) -> bool: 
+    """
+    Desc: 
+        This method moves 'n' random files from the source directory (dir1) to destination directory (dir2).
+    Args:
+        dir1 (str): Path to the source directory.
+        dir2 (str): Path to the destination directory.
+        n (int): The number of random files to be moved.
+    Returns:
+        True, if the operation was successful, otherwise False.
+    """
+    try:
+        # Check if the source and destination directory exists
+        if not os.path.isdir(dir1):
+            print(f"Directory '{dir1}' does not exist.")
+            return False
+        
+        if not os.path.isdir(dir2):
+            print(f"Directory '{dir2}' does not exist. Creating it.")
+            os.makedirs(dir2)
+        
+        # Get all files (not directories) in the specified directory
+        all_files = [f for f in os.listdir(dir1) if os.path.isfile(os.path.join(dir1, f))]
+        
+        if len(all_files) < n:
+            print(f"Cannot move '{n}' files, directory only contains '{len(all_files)}' files.")
+            return False
+        
+        files_to_move = random.sample(all_files, n)
+        
+        for file in tqdm(files_to_move):
+            source_path      = os.path.join(dir1, file)
+            destination_path = os.path.join(dir2, file)
+
+            # Move the file to the destination directory
+            shutil.move(source_path, destination_path)
+        
+        return True
+    
+    except Exception as move_ex:
+        print(f"Error occurred while moving files: {str(move_ex)}.")
+        return False
+
+
+def copy_n_random_files(dir1: str, dir2: str, n: int) -> bool:
+    """
+    Desc:
+        Randomly select and copy 'n' files from one directory to another.
+    Args:
+        dir1 (str): Path to the source directory.
+        dir2 (str): Path to the destination directory. Will be created if it doesn't exist.
+        n (int): Number of files to randomly copy.
+    Returns:
+        bool: True if the operation was successful, False otherwise.
+    """
+    try:
+        # Check if the source directory exists
+        if not os.path.isdir(dir1):
+            print(f"Directory '{dir1}' does not exist.")
+            return False
+
+        # Ensure destination directory exists
+        if not os.path.isdir(dir2):
+            print(f"Directory '{dir2}' does not exist. Creating it.")
+            os.makedirs(dir2)
+
+        # Get list of all files (not directories) in source
+        all_files = [f for f in os.listdir(dir1) if os.path.isfile(os.path.join(dir1, f))]
+
+        if len(all_files) < n:
+            n = len(all_files)
+
+        print(f"Copying '{n}' files to '{dir2}'...")
+        files_to_copy = random.sample(all_files, n)
+
+        for file in tqdm(files_to_copy, desc="Copying files"):
+            source_path = os.path.join(dir1, file)
+            destination_path = os.path.join(dir2, file)
+            shutil.copy(source_path, destination_path)
+
+        return True
+
+    except Exception as copy_ex:
+        print(f"Error occurred while copying files: {str(copy_ex)}.")
+        return False
+    
+
+######################
+#
+######################
+def copy_n_unique_files(dir1, dir2, output_dir, n):
+    """
+    Desc: 
+        This method iterates files in dir1 and checks if they are not present in dir2. If not present, copies the file to output_dir. Moves 'n' files in total.
+    Args:
+        dir1 (str): Path to directory 1.
+        dir2 (str): Path to directory 2.
+        output_dir (str): Path to the destination directory.
+        n (int): The number of random files to be moved.
+    Returns:
+        True, if the operation was successful, otherwise False.
+    """
+    try:
+        # List all files in dir1 and dir2
+        dir1_files = [f for f in os.listdir(dir1) if os.path.isfile(os.path.join(dir1, f))]
+        dir2_files = [f for f in os.listdir(dir2) if os.path.isfile(os.path.join(dir2, f))]
+
+        # Filter out files that already exist in dir2
+        unique_files = [f for f in dir1_files if f not in dir2_files]
+
+        # If no unique files are found
+        if not unique_files:
+            print("No unique files to move.")
+            return False
+
+        # Randomly select 'n' files to copy (make sure we don't select more than we have)
+        files_to_copy = random.sample(unique_files, min(n, len(unique_files)))
+
+        # Copy selected files to output_dir
+        files_copied = 0
+        for file in files_to_copy:
+            src_path  = os.path.join(dir1, file)
+            dest_path = os.path.join(output_dir, file)
+            shutil.copy(src_path, dest_path)
+            files_copied += 1
+            print(f"Copied: {file}")
+
+        print(f"Total files copied: {files_copied}")
+        return True
+    
+    except Exception as copy_ex:
+        print(f"An error occurred while copying: {copy_ex}")
+        return False
+ 
+
+######################
+#
+######################
+if __name__ == "__main__":
+    dir1 = ""
+    dir2 = ""
+    n    = 0
+    
+    move_n_random_files(dir1, dir2, n)

preprocess/rename.py

@@ -0,0 +1,51 @@
+import os
+from   tqdm     import tqdm
+
+
+def rename_files(input_dir: str, word: str, new_word: str = '') -> bool:
+    """
+    Desc: 
+        This method renames the files that has {word} in the filename
+    Args:
+        input_dir (str): Path to the input directory 
+        word (str): A word to look for in the filename
+        new_word (str): A new word that is used to replace
+    Returns:
+        True, if renaming operation is success, else False.
+    """
+    try:        
+        if not os.path.isdir(input_dir):
+            print(f"The directory {input_dir} does not exist.")
+            return False
+        
+        for file in tqdm(os.listdir(input_dir)):
+            basename, ext = os.path.splitext(file)
+
+            if word in basename:
+                new_name = f"{basename.replace(word, new_word)}{ext}"
+                
+                # Construct full paths for renaming
+                old_file_path = os.path.join(input_dir, file)
+                new_file_path = os.path.join(input_dir, new_name)
+                
+                try:
+                    # Rename the file
+                    os.rename(old_file_path, new_file_path)
+                except:
+                    continue
+        
+        return True
+    
+    except Exception as rename_ex:
+        print(f"An error occurred while renaming: {rename_ex}")
+        return False
+
+   
+
+
+if __name__ == "__main__":
+    inp_dir = ""
+    word = ""
+    new_word = ""
+    rename_files(inp_dir, word, new_word)
+    

stage1/inference_cvt13.py

@@ -0,0 +1,637 @@
+import torch
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import sys
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+import json
+from datetime import datetime
+from sklearn.metrics import (
+    accuracy_score, precision_score, recall_score, f1_score,
+    confusion_matrix, classification_report, roc_curve, auc,
+    precision_recall_curve, average_precision_score, roc_auc_score
+)
+import seaborn as sns
+import pandas as pd
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+BASE_DIR        = "path_to_CornViT"
+
+# Path to the Microsoft CvT repository
+CVT_REPO_PATH = f"{BASE_DIR}/CvT"
+
+# Model configuration
+IMG_SIZE = 384
+NUM_CLASSES = 2
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+RUN = "cvt13_run_2025xxxx_xxxxxx"
+
+# Path to trained model
+MODEL_PATH = f"metrics/{RUN}/train/best_model.pth"
+
+# Test data folder (should have subfolders for each class like train/val structure)
+TEST_DATA_DIR = f"{BASE_DIR}/stage1/data/test"
+
+# Class names (update these to match your dataset)
+CLASS_NAMES = ["Pure", "Impure"]
+
+# Output directory for evaluation results (within the same metrics folder)
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+EVAL_OUTPUT_DIR = f"metrics/{RUN}/evals/eval_{timestamp}"
+os.makedirs(EVAL_OUTPUT_DIR, exist_ok=True)
+
+
+# ============================================================
+# SETUP: Import CvT model
+# ============================================================
+
+# Fix torch._six compatibility
+cls_cvt_path = os.path.join(CVT_REPO_PATH, "lib", "models", "cls_cvt.py")
+if os.path.exists(cls_cvt_path):
+    with open(cls_cvt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    if "from torch._six import container_abcs" in content:
+        content = content.replace(
+            "from torch._six import container_abcs",
+            "import collections.abc as container_abcs"
+        )
+        content = content.replace(
+            "or pretrained_layers[0] is '*'",
+            "or pretrained_layers[0] == '*'"
+        )
+        with open(cls_cvt_path, 'w', encoding='utf-8') as f:
+            f.write(content)
+
+sys.path.insert(0, CVT_REPO_PATH)
+
+import warnings
+warnings.filterwarnings('ignore', category=SyntaxWarning)
+
+from lib.models import cls_cvt
+from lib.config import config, update_config
+
+
+# ============================================================
+# MODEL LOADING
+# ============================================================
+
+def load_model(model_path, config_path=None):
+    """Load the trained CvT model"""
+    
+    # Load config
+    if config_path is None:
+        config_path = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt", "cvt-13-384x384.yaml")
+    
+    config.defrost()
+    config.merge_from_file(config_path)
+    config.MODEL.NUM_CLASSES = NUM_CLASSES
+    config.MODEL.PRETRAINED = ''
+    config.freeze()
+    
+    # Create model
+    model = cls_cvt.get_cls_model(config)
+    
+    # Load trained weights
+    checkpoint = torch.load(model_path, map_location=DEVICE)
+    if 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+    
+    model = model.to(DEVICE)
+    model.eval()
+    
+    print(f"✅ Model loaded from: {model_path}")
+    return model
+
+
+# ============================================================
+# DATA LOADING
+# ============================================================
+
+def get_test_dataloader(test_dir, batch_size=32):
+    """Create test dataloader"""
+    test_transforms = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                           [0.229, 0.224, 0.225])
+    ])
+    
+    test_dataset = datasets.ImageFolder(test_dir, transform=test_transforms)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, 
+                            shuffle=False, num_workers=0, pin_memory=True)
+    
+    print(f"✅ Test dataset loaded: {len(test_dataset)} images")
+    print(f"   Classes: {test_dataset.classes}")
+    return test_loader, test_dataset
+
+
+# ============================================================
+# EVALUATION FUNCTIONS
+# ============================================================
+
+def evaluate_model(model, test_loader, test_dataset):
+    """
+    Evaluate model with single image predictions
+    
+    Returns:
+        all_preds: Predicted class labels
+        all_labels: Ground truth labels
+        all_probs: Predicted probabilities for all classes
+        all_confidences: Confidence scores
+        image_paths: List of image paths
+    """
+    model.eval()
+    
+    all_preds = []
+    all_labels = []
+    all_probs = []
+    all_confidences = []
+    image_paths = []
+    
+    print("\n🔍 Running single-image inference on test set...")
+    
+    # Process each image individually
+    total_images = len(test_dataset)
+    
+    for idx in range(total_images):
+        # Get single image and label
+        image, label = test_dataset[idx]
+        img_path, _ = test_dataset.samples[idx]
+        
+        # Add batch dimension and move to device
+        image = image.unsqueeze(0).to(DEVICE)
+        
+        with torch.no_grad():
+            # Forward pass
+            output = model(image)
+            
+            # Ensure output has correct shape
+            if output.dim() == 1:
+                output = output.unsqueeze(0)
+                
+            probabilities = torch.softmax(output, dim=1)
+            confidence, predicted = torch.max(probabilities, 1)
+            
+            # Collect results
+            all_preds.append(predicted.item())
+            all_labels.append(label)
+            all_probs.append(probabilities.cpu().numpy()[0])
+            all_confidences.append(confidence.item())
+            image_paths.append(img_path)
+        
+        # Progress update
+        if (idx + 1) % 50 == 0 or (idx + 1) == total_images:
+            print(f"   Processed {idx + 1}/{total_images} images...")
+    
+    print(f"✅ Inference complete: {len(all_preds)} predictions")
+    
+    return (np.array(all_preds), np.array(all_labels), np.array(all_probs), 
+            np.array(all_confidences), image_paths)
+
+
+# ============================================================
+# METRICS CALCULATION
+# ============================================================
+
+def calculate_metrics(y_true, y_pred, y_probs):
+    """Calculate all classification metrics"""
+    
+    metrics = {}
+    
+    # Basic metrics
+    metrics['accuracy'] = accuracy_score(y_true, y_pred)
+    metrics['precision_macro'] = precision_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['precision_weighted'] = precision_score(y_true, y_pred, average='weighted', zero_division=0)
+    metrics['recall_macro'] = recall_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['recall_weighted'] = recall_score(y_true, y_pred, average='weighted', zero_division=0)
+    metrics['f1_macro'] = f1_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['f1_weighted'] = f1_score(y_true, y_pred, average='weighted', zero_division=0)
+    
+    # Per-class metrics
+    precision_per_class = precision_score(y_true, y_pred, average=None, zero_division=0)
+    recall_per_class = recall_score(y_true, y_pred, average=None, zero_division=0)
+    f1_per_class = f1_score(y_true, y_pred, average=None, zero_division=0)
+    
+    metrics['per_class'] = {}
+    for i, class_name in enumerate(CLASS_NAMES):
+        metrics['per_class'][class_name] = {
+            'precision': float(precision_per_class[i]),
+            'recall': float(recall_per_class[i]),
+            'f1_score': float(f1_per_class[i])
+        }
+    
+    # ROC-AUC (for binary and multi-class)
+    if NUM_CLASSES == 2:
+        metrics['roc_auc'] = roc_auc_score(y_true, y_probs[:, 1])
+        metrics['average_precision'] = average_precision_score(y_true, y_probs[:, 1])
+    else:
+        metrics['roc_auc_ovr'] = roc_auc_score(y_true, y_probs, multi_class='ovr', average='macro')
+        metrics['roc_auc_ovo'] = roc_auc_score(y_true, y_probs, multi_class='ovo', average='macro')
+    
+    return metrics
+
+
+# ============================================================
+# VISUALIZATION FUNCTIONS
+# ============================================================
+
+def plot_confusion_matrix(y_true, y_pred, save_path):
+    """Plot and save confusion matrix"""
+    cm = confusion_matrix(y_true, y_pred)
+    
+    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
+    
+    # Raw counts
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                ax=axes[0], cbar_kws={'label': 'Count'})
+    axes[0].set_xlabel('Predicted Label', fontsize=12)
+    axes[0].set_ylabel('True Label', fontsize=12)
+    axes[0].set_title('Confusion Matrix (Counts)', fontsize=14, fontweight='bold')
+    
+    # Normalized
+    cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
+    sns.heatmap(cm_normalized, annot=True, fmt='.2%', cmap='Blues',
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                ax=axes[1], cbar_kws={'label': 'Percentage'})
+    axes[1].set_xlabel('Predicted Label', fontsize=12)
+    axes[1].set_ylabel('True Label', fontsize=12)
+    axes[1].set_title('Confusion Matrix (Normalized)', fontsize=14, fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Confusion matrix saved to: {save_path}")
+    plt.close()
+    
+    return cm
+
+
+def plot_roc_curve(y_true, y_probs, save_path):
+    """Plot ROC curve"""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    
+    if NUM_CLASSES == 2:
+        # Binary classification
+        fpr, tpr, _ = roc_curve(y_true, y_probs[:, 1])
+        roc_auc = auc(fpr, tpr)
+        
+        ax.plot(fpr, tpr, color='darkorange', lw=2,
+                label=f'ROC curve (AUC = {roc_auc:.3f})')
+        ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Classifier')
+        
+    else:
+        # Multi-class (one-vs-rest)
+        for i, class_name in enumerate(CLASS_NAMES):
+            y_true_binary = (y_true == i).astype(int)
+            fpr, tpr, _ = roc_curve(y_true_binary, y_probs[:, i])
+            roc_auc = auc(fpr, tpr)
+            ax.plot(fpr, tpr, lw=2, label=f'{class_name} (AUC = {roc_auc:.3f})')
+        
+        ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Classifier')
+    
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.05])
+    ax.set_xlabel('False Positive Rate', fontsize=12)
+    ax.set_ylabel('True Positive Rate', fontsize=12)
+    ax.set_title('Receiver Operating Characteristic (ROC) Curve', fontsize=14, fontweight='bold')
+    ax.legend(loc="lower right", fontsize=10)
+    ax.grid(alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 ROC curve saved to: {save_path}")
+    plt.close()
+
+
+def plot_precision_recall_curve(y_true, y_probs, save_path):
+    """Plot Precision-Recall curve"""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    
+    if NUM_CLASSES == 2:
+        # Binary classification
+        precision, recall, _ = precision_recall_curve(y_true, y_probs[:, 1])
+        avg_precision = average_precision_score(y_true, y_probs[:, 1])
+        
+        ax.plot(recall, precision, color='darkorange', lw=2,
+                label=f'PR curve (AP = {avg_precision:.3f})')
+        
+    else:
+        # Multi-class
+        for i, class_name in enumerate(CLASS_NAMES):
+            y_true_binary = (y_true == i).astype(int)
+            precision, recall, _ = precision_recall_curve(y_true_binary, y_probs[:, i])
+            avg_precision = average_precision_score(y_true_binary, y_probs[:, i])
+            ax.plot(recall, precision, lw=2, 
+                   label=f'{class_name} (AP = {avg_precision:.3f})')
+    
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.05])
+    ax.set_xlabel('Recall', fontsize=12)
+    ax.set_ylabel('Precision', fontsize=12)
+    ax.set_title('Precision-Recall Curve', fontsize=14, fontweight='bold')
+    ax.legend(loc="lower left", fontsize=10)
+    ax.grid(alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Precision-Recall curve saved to: {save_path}")
+    plt.close()
+
+
+def plot_class_distribution(y_true, y_pred, save_path):
+    """Plot class distribution comparison"""
+    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
+    
+    # True distribution
+    true_counts = [np.sum(y_true == i) for i in range(NUM_CLASSES)]
+    axes[0].bar(CLASS_NAMES, true_counts, color='steelblue', alpha=0.7)
+    axes[0].set_ylabel('Count', fontsize=12)
+    axes[0].set_title('True Label Distribution', fontsize=14, fontweight='bold')
+    axes[0].grid(axis='y', alpha=0.3)
+    for i, count in enumerate(true_counts):
+        axes[0].text(i, count + max(true_counts)*0.01, str(count), 
+                    ha='center', va='bottom', fontweight='bold')
+    
+    # Predicted distribution
+    pred_counts = [np.sum(y_pred == i) for i in range(NUM_CLASSES)]
+    axes[1].bar(CLASS_NAMES, pred_counts, color='coral', alpha=0.7)
+    axes[1].set_ylabel('Count', fontsize=12)
+    axes[1].set_title('Predicted Label Distribution', fontsize=14, fontweight='bold')
+    axes[1].grid(axis='y', alpha=0.3)
+    for i, count in enumerate(pred_counts):
+        axes[1].text(i, count + max(pred_counts)*0.01, str(count), 
+                    ha='center', va='bottom', fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Class distribution saved to: {save_path}")
+    plt.close()
+
+
+def plot_per_class_metrics(metrics, save_path):
+    """Plot per-class performance metrics"""
+    classes = list(metrics['per_class'].keys())
+    precision_vals = [metrics['per_class'][c]['precision'] for c in classes]
+    recall_vals = [metrics['per_class'][c]['recall'] for c in classes]
+    f1_vals = [metrics['per_class'][c]['f1_score'] for c in classes]
+    
+    x = np.arange(len(classes))
+    width = 0.25
+    
+    fig, ax = plt.subplots(figsize=(12, 7))
+    
+    bars1 = ax.bar(x - width, precision_vals, width, label='Precision', color='steelblue', alpha=0.8)
+    bars2 = ax.bar(x, recall_vals, width, label='Recall', color='coral', alpha=0.8)
+    bars3 = ax.bar(x + width, f1_vals, width, label='F1-Score', color='lightgreen', alpha=0.8)
+    
+    ax.set_ylabel('Score', fontsize=12)
+    ax.set_title('Per-Class Performance Metrics', fontsize=14, fontweight='bold')
+    ax.set_xticks(x)
+    ax.set_xticklabels(classes)
+    ax.legend(fontsize=11)
+    ax.set_ylim([0, 1.1])
+    ax.grid(axis='y', alpha=0.3)
+    
+    # Add value labels on bars
+    def autolabel(bars):
+        for bar in bars:
+            height = bar.get_height()
+            ax.text(bar.get_x() + bar.get_width()/2., height + 0.02,
+                   f'{height:.3f}', ha='center', va='bottom', fontsize=9)
+    
+    autolabel(bars1)
+    autolabel(bars2)
+    autolabel(bars3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Per-class metrics saved to: {save_path}")
+    plt.close()
+
+
+def plot_confidence_distribution(y_true, y_pred, confidences, save_path):
+    """Plot confidence score distribution for correct vs incorrect predictions"""
+    # Confidence scores are already extracted
+    correct = (y_true == y_pred)
+    
+    fig, axes = plt.subplots(2, 1, figsize=(12, 10))
+    
+    # Histogram
+    axes[0].hist(confidences[correct], bins=50, alpha=0.7, label='Correct', 
+                color='green', edgecolor='black')
+    axes[0].hist(confidences[~correct], bins=50, alpha=0.7, label='Incorrect', 
+                color='red', edgecolor='black')
+    axes[0].set_xlabel('Confidence Score', fontsize=12)
+    axes[0].set_ylabel('Frequency', fontsize=12)
+    axes[0].set_title('Confidence Distribution: Correct vs Incorrect Predictions', 
+                     fontsize=14, fontweight='bold')
+    axes[0].legend(fontsize=11)
+    axes[0].grid(alpha=0.3)
+    
+    # Box plot
+    data_to_plot = [confidences[correct], confidences[~correct]]
+    box = axes[1].boxplot(data_to_plot, labels=['Correct', 'Incorrect'], 
+                          patch_artist=True, showmeans=True)
+    box['boxes'][0].set_facecolor('lightgreen')
+    box['boxes'][1].set_facecolor('lightcoral')
+    axes[1].set_ylabel('Confidence Score', fontsize=12)
+    axes[1].set_title('Confidence Score Box Plot', fontsize=14, fontweight='bold')
+    axes[1].grid(axis='y', alpha=0.3)
+    
+    # Add statistics
+    correct_mean = np.mean(confidences[correct])
+    incorrect_mean = np.mean(confidences[~correct]) if (~correct).sum() > 0 else 0
+    axes[1].text(1, correct_mean, f'μ={correct_mean:.3f}', 
+                ha='right', va='center', fontweight='bold', fontsize=10)
+    if (~correct).sum() > 0:
+        axes[1].text(2, incorrect_mean, f'μ={incorrect_mean:.3f}', 
+                    ha='left', va='center', fontweight='bold', fontsize=10)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Confidence distribution saved to: {save_path}")
+    plt.close()
+
+
+# ============================================================
+# RESULTS SAVING
+# ============================================================
+
+def save_predictions_to_csv(image_paths, y_true, y_pred, y_probs, confidences, save_path):
+    """Save detailed predictions to CSV"""
+    results = []
+    
+    for img_path, true_label, pred, probs, conf in zip(image_paths, y_true, y_pred, y_probs, confidences):
+        result = {
+            'image_path': img_path,
+            'image_name': os.path.basename(img_path),
+            'true_label': CLASS_NAMES[true_label],
+            'true_label_idx': true_label,
+            'predicted_label': CLASS_NAMES[pred],
+            'predicted_label_idx': pred,
+            'confidence': conf,
+            'correct': pred == true_label
+        }
+        
+        # Add probabilities for each class
+        for i, class_name in enumerate(CLASS_NAMES):
+            result[f'prob_{class_name}'] = probs[i]
+        
+        results.append(result)
+    
+    df = pd.DataFrame(results)
+    df.to_csv(save_path, index=False)
+    print(f"💾 Predictions saved to: {save_path}")
+    
+    # Print some statistics
+    print(f"\n📊 Prediction Statistics:")
+    print(f"   Total images: {len(df)}")
+    print(f"   Correct predictions: {df['correct'].sum()} ({df['correct'].sum()/len(df)*100:.2f}%)")
+    print(f"   Incorrect predictions: {(~df['correct']).sum()} ({(~df['correct']).sum()/len(df)*100:.2f}%)")
+    print(f"   Average confidence: {df['confidence'].mean():.4f}")
+    print(f"   Confidence on correct: {df[df['correct']]['confidence'].mean():.4f}")
+    print(f"   Confidence on incorrect: {df[~df['correct']]['confidence'].mean():.4f}" if (~df['correct']).sum() > 0 else "")
+    
+    return df
+
+
+def save_metrics_json(metrics, save_path):
+    """Save metrics to JSON file"""
+    with open(save_path, 'w') as f:
+        json.dump(metrics, f, indent=4)
+    print(f"💾 Metrics saved to: {save_path}")
+
+
+def generate_classification_report_file(y_true, y_pred, save_path):
+    """Generate and save sklearn classification report"""
+    report = classification_report(y_true, y_pred, target_names=CLASS_NAMES, digits=4)
+    
+    with open(save_path, 'w') as f:
+        f.write("="*60 + "\n")
+        f.write("CLASSIFICATION REPORT\n")
+        f.write("="*60 + "\n\n")
+        f.write(report)
+    
+    print(f"📄 Classification report saved to: {save_path}")
+
+
+# ============================================================
+# MAIN EVALUATION PIPELINE
+# ============================================================
+
+def main():
+    """Main evaluation pipeline"""
+    
+    print("\n" + "="*60)
+    print("CvT-13 MODEL EVALUATION PIPELINE")
+    print("Single Image Prediction Mode")
+    print("="*60 + "\n")
+    
+    # Load model
+    print("📦 Loading model...")
+    model = load_model(MODEL_PATH)
+    
+    # Load test data
+    print("\n📂 Loading test data...")
+    test_loader, test_dataset = get_test_dataloader(TEST_DATA_DIR, batch_size=1)
+    
+    # Run evaluation with single image predictions
+    print("\n🔍 Evaluating model (single image predictions)...")
+    y_pred, y_true, y_probs, confidences, image_paths = evaluate_model(model, test_loader, test_dataset)
+    
+    # Calculate metrics
+    print("\n📊 Calculating metrics...")
+    metrics = calculate_metrics(y_true, y_pred, y_probs)
+    
+    # Print key metrics
+    print("\n" + "="*60)
+    print("EVALUATION RESULTS")
+    print("="*60)
+    print(f"Total Images Evaluated: {len(y_pred)}")
+    print(f"Accuracy:           {metrics['accuracy']*100:.2f}%")
+    print(f"Precision (Macro):  {metrics['precision_macro']*100:.2f}%")
+    print(f"Recall (Macro):     {metrics['recall_macro']*100:.2f}%")
+    print(f"F1-Score (Macro):   {metrics['f1_macro']*100:.2f}%")
+    if 'roc_auc' in metrics:
+        print(f"ROC-AUC:            {metrics['roc_auc']:.4f}")
+    print("\nPer-Class Metrics:")
+    for class_name, class_metrics in metrics['per_class'].items():
+        print(f"  {class_name}:")
+        print(f"    Precision: {class_metrics['precision']*100:.2f}%")
+        print(f"    Recall:    {class_metrics['recall']*100:.2f}%")
+        print(f"    F1-Score:  {class_metrics['f1_score']*100:.2f}%")
+    print("="*60)
+    
+    # Generate all visualizations
+    print("\n📊 Generating visualizations...")
+    plot_confusion_matrix(y_true, y_pred, 
+                         os.path.join(EVAL_OUTPUT_DIR, "confusion_matrix.png"))
+    plot_roc_curve(y_true, y_probs, 
+                  os.path.join(EVAL_OUTPUT_DIR, "roc_curve.png"))
+    plot_precision_recall_curve(y_true, y_probs, 
+                               os.path.join(EVAL_OUTPUT_DIR, "precision_recall_curve.png"))
+    plot_class_distribution(y_true, y_pred, 
+                          os.path.join(EVAL_OUTPUT_DIR, "class_distribution.png"))
+    plot_per_class_metrics(metrics, 
+                          os.path.join(EVAL_OUTPUT_DIR, "per_class_metrics.png"))
+    plot_confidence_distribution(y_true, y_pred, confidences,
+                                os.path.join(EVAL_OUTPUT_DIR, "confidence_distribution.png"))
+    
+    # Save results
+    print("\n💾 Saving results...")
+    df = save_predictions_to_csv(image_paths, y_true, y_pred, y_probs, confidences,
+                                os.path.join(EVAL_OUTPUT_DIR, "predictions.csv"))
+    save_metrics_json(metrics, 
+                     os.path.join(EVAL_OUTPUT_DIR, "metrics.json"))
+    generate_classification_report_file(y_true, y_pred,
+                                       os.path.join(EVAL_OUTPUT_DIR, "classification_report.txt"))
+    
+    # Save misclassified images list
+    misclassified = df[~df['correct']]
+    if len(misclassified) > 0:
+        misclassified_path = os.path.join(EVAL_OUTPUT_DIR, "misclassified_images.csv")
+        misclassified.to_csv(misclassified_path, index=False)
+        print(f"⚠️  Misclassified images saved to: {misclassified_path}")
+        print(f"   Total misclassified: {len(misclassified)}")
+    
+    # Save low confidence predictions
+    low_conf_threshold = 0.7
+    low_confidence = df[df['confidence'] < low_conf_threshold]
+    if len(low_confidence) > 0:
+        low_conf_path = os.path.join(EVAL_OUTPUT_DIR, "low_confidence_predictions.csv")
+        low_confidence.to_csv(low_conf_path, index=False)
+        print(f"⚠️  Low confidence predictions saved to: {low_conf_path}")
+        print(f"   Total with confidence < {low_conf_threshold}: {len(low_confidence)}")
+    
+    print("\n" + "="*60)
+    print(f"✅ Evaluation complete!")
+    print(f"📁 All results saved to: {EVAL_OUTPUT_DIR}")
+    print("="*60 + "\n")
+    
+    print("Generated files:")
+    print("  📊 confusion_matrix.png - Confusion matrix visualization")
+    print("  📊 roc_curve.png - ROC curve")
+    print("  📊 precision_recall_curve.png - Precision-Recall curve")
+    print("  📊 class_distribution.png - Class distribution comparison")
+    print("  📊 per_class_metrics.png - Per-class performance")
+    print("  📊 confidence_distribution.png - Confidence analysis")
+    print("  💾 predictions.csv - Detailed predictions for each image")
+    print("  💾 misclassified_images.csv - List of incorrectly classified images")
+    print("  💾 low_confidence_predictions.csv - Predictions with low confidence")
+    print("  💾 metrics.json - All metrics in JSON format")
+    print("  📄 classification_report.txt - Sklearn classification report")
+
+
+if __name__ == '__main__':
+    main()

stage1/stage1_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:445c5a5b94b86649cab12ef3c2fe4df9461f9879864c43d52a7cc9560204fcc3
+size 78733538

stage1/train_cvt13.py

@@ -0,0 +1,457 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import sys
+import os
+from timm.loss import SoftTargetCrossEntropy
+from timm.scheduler import CosineLRScheduler
+from timm.utils import accuracy
+import matplotlib.pyplot as plt
+import json
+from datetime import datetime
+
+
+# ============================================================
+# SETUP: Clone and import from Microsoft CvT repository
+# ============================================================
+"""
+First, clone the Microsoft CvT repository:
+    git clone https://github.com/microsoft/CvT.git
+    cd CvT
+    pip install -r requirements.txt
+"""
+
+BASE_DIR = "path_to_CornViT"
+
+# Add the CvT repo to Python path
+CVT_REPO_PATH = f"{BASE_DIR}/CvT"
+
+if not os.path.exists(CVT_REPO_PATH):
+    print(f"❌ CvT repository not found at {CVT_REPO_PATH}")
+    print("Please clone it: git clone https://github.com/microsoft/CvT.git")
+    sys.exit(1)
+
+# Fix torch._six compatibility BEFORE importing
+print("Applying compatibility fixes for newer PyTorch versions...")
+cls_cvt_path = os.path.join(CVT_REPO_PATH, "lib", "models", "cls_cvt.py")
+
+if os.path.exists(cls_cvt_path):
+    with open(cls_cvt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    # Fix 1: Replace torch._six import
+    if "from torch._six import container_abcs" in content:
+        content = content.replace(
+            "from torch._six import container_abcs",
+            "import collections.abc as container_abcs"
+        )
+        
+        # Fix 2: Replace 'is' with '==' for string comparison
+        content = content.replace(
+            "or pretrained_layers[0] is '*'",
+            "or pretrained_layers[0] == '*'"
+        )
+        
+        with open(cls_cvt_path, 'w', encoding='utf-8') as f:
+            f.write(content)
+        print("✅ Applied compatibility patches to cls_cvt.py")
+    else:
+        print("✅ Compatibility patches already applied")
+else:
+    print(f"❌ Could not find cls_cvt.py at {cls_cvt_path}")
+    sys.exit(1)
+
+# Now import
+sys.path.insert(0, CVT_REPO_PATH)
+
+# Suppress the SyntaxWarning
+import warnings
+warnings.filterwarnings('ignore', category=SyntaxWarning)
+
+from lib.models import cls_cvt
+from lib.config import config, update_config
+print("✅ Successfully imported Microsoft CvT models")
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+DATA_DIR        = f"{BASE_DIR}/stage1/data"
+BATCH_SIZE      = 32
+IMG_SIZE        = 384
+NUM_CLASSES     = 2
+NUM_EPOCHS      = 100
+DEVICE          = "cuda" if torch.cuda.is_available() else "cpu"
+PRETRAINED_PATH = f"{BASE_DIR}/CvT-13-384x384-IN-22k.pth"
+
+# Create output directory for saving results
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+OUTPUT_DIR = f"metrics/cvt13_run_{timestamp}"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+print(f"Metrics will be saved to: {OUTPUT_DIR}")
+
+
+# ============================================================
+# DATASET & AUGMENTATION
+# ============================================================
+
+train_transforms = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomVerticalFlip(),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+    transforms.RandomRotation(15),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+val_transforms = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+train_dataset = datasets.ImageFolder(f"{DATA_DIR}/train", transform=train_transforms)
+val_dataset = datasets.ImageFolder(f"{DATA_DIR}/val", transform=val_transforms)
+
+train_loader = DataLoader(
+    train_dataset, 
+    batch_size=BATCH_SIZE, 
+    shuffle=True, 
+    num_workers=0, 
+    pin_memory=True,
+    drop_last=True
+)
+val_loader = DataLoader(
+    val_dataset, 
+    batch_size=BATCH_SIZE,
+    shuffle=False, 
+    num_workers=0, 
+    pin_memory=True,
+    drop_last=True
+)
+
+
+# ============================================================
+# MODEL SETUP - Using Microsoft CvT Implementation
+# ============================================================
+
+# Load the CvT-13 config from the repository
+cvt_config_path = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt", "cvt-13-384x384.yaml")
+
+if not os.path.exists(cvt_config_path):
+    print(f"⚠️ Config file not found at {cvt_config_path}")
+    print("Available configs:")
+    config_dir = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt")
+    if os.path.exists(config_dir):
+        for f in os.listdir(config_dir):
+            if f.endswith('.yaml'):
+                print(f"  - {f}")
+    sys.exit(1)
+
+print(f"Loading config from: {cvt_config_path}")
+
+# Load config directly using merge_from_file
+config.defrost()
+config.merge_from_file(cvt_config_path)
+
+# Update the number of classes for our task
+config.MODEL.NUM_CLASSES = NUM_CLASSES
+config.MODEL.PRETRAINED = ''  # We'll load weights manually
+config.freeze()
+
+print("Creating CvT-13 model...")
+# Create model using the official CvT architecture
+model = cls_cvt.get_cls_model(config)
+model = model.to(DEVICE)
+
+# Load pretrained weights
+if os.path.exists(PRETRAINED_PATH):
+    print(f"Loading pretrained weights from {PRETRAINED_PATH}")
+    try:
+        checkpoint = torch.load(PRETRAINED_PATH, map_location=DEVICE)
+        
+        # Handle different checkpoint formats
+        if 'model' in checkpoint:
+            state_dict = checkpoint['model']
+        elif 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+        
+        # Remove 'module.' prefix if present
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            name = k.replace("module.", "")
+            new_state_dict[name] = v
+        
+        # Remove head layers from pretrained weights (they have different dimensions)
+        filtered_state_dict = {k: v for k, v in new_state_dict.items() if 'head' not in k}
+        
+        # Load weights - strict=False will only load matching layers
+        missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False)
+        
+        # Count how many weights were actually loaded
+        loaded_keys = [k for k in filtered_state_dict.keys() if k in model.state_dict()]
+        print(f"✅ Loaded pretrained weights: {len(loaded_keys)} layers from backbone")
+        print(f"   Head layer initialized randomly for {NUM_CLASSES} classes")
+        
+        # Show what's missing (should only be head-related)
+        head_missing = [k for k in missing_keys if 'head' in k]
+        other_missing = [k for k in missing_keys if 'head' not in k]
+        
+        if other_missing:
+            print(f"⚠️  Warning - Missing non-head keys: {other_missing}")
+        if unexpected_keys:
+            print(f"⚠️  Unexpected keys: {unexpected_keys}")
+            
+    except Exception as e:
+        print(f"⚠️ Error loading pretrained weights: {e}")
+        import traceback
+        traceback.print_exc()
+        print("Continuing with random initialization...")
+else:
+    print(f"⚠️ Pretrained weights not found at {PRETRAINED_PATH}")
+    print("Training from scratch...")
+
+# Freeze backbone - only train the head for faster training and less overfitting
+print("Freezing backbone layers (keeping only head trainable)...")
+for name, param in model.named_parameters():
+    if "head" not in name:
+        param.requires_grad = False
+
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Trainable parameters: {trainable_params:,} / {total_params:,}")
+print(f"Frozen parameters: {total_params - trainable_params:,}")
+
+
+# ============================================================
+# OPTIMIZER AND LOSS
+# ============================================================
+
+optimizer = optim.AdamW(
+    filter(lambda p: p.requires_grad, model.parameters()), 
+    lr=1e-4, 
+    weight_decay=0.05
+)
+
+criterion = SoftTargetCrossEntropy()
+
+lr_scheduler = CosineLRScheduler(
+    optimizer,
+    t_initial=NUM_EPOCHS,
+    lr_min=1e-6,
+    warmup_t=5,
+    warmup_lr_init=1e-5,
+)
+
+
+# ============================================================
+# TRAINING & VALIDATION LOOP
+# ============================================================
+
+def train_one_epoch(epoch, history):
+    model.train()
+    total_loss, total_acc = 0, 0
+    
+    for images, targets in train_loader:
+        images, targets = images.to(DEVICE), targets.to(DEVICE)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+
+        acc1, _ = accuracy(outputs, targets.argmax(dim=1), topk=(1, 5))
+        total_loss += loss.item()
+        total_acc += acc1.item()
+
+    avg_loss = total_loss / len(train_loader)
+    avg_acc = total_acc / len(train_loader)
+    
+    history['train_loss'].append(avg_loss)
+    history['train_acc'].append(avg_acc)
+    history['learning_rate'].append(optimizer.param_groups[0]['lr'])
+    
+    print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] | Train Loss: {avg_loss:.4f} | Train Acc: {avg_acc:.2f}% | LR: {optimizer.param_groups[0]['lr']:.6f}")
+    return avg_loss, avg_acc
+
+
+def validate(epoch, history):
+    model.eval()
+    total_loss, total_acc = 0, 0
+    
+    with torch.no_grad():
+        for images, targets in val_loader:
+            images, targets = images.to(DEVICE), targets.to(DEVICE)
+            outputs = model(images)
+            loss = nn.CrossEntropyLoss()(outputs, targets)
+            acc1, _ = accuracy(outputs, targets, topk=(1, 5))
+            
+            total_loss += loss.item()
+            total_acc += acc1.item()
+    
+    avg_loss = total_loss / len(val_loader)
+    avg_acc = total_acc / len(val_loader)
+    
+    history['val_loss'].append(avg_loss)
+    history['val_acc'].append(avg_acc)
+    
+    print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] | Val Loss: {avg_loss:.4f} | Val Acc: {avg_acc:.2f}%")
+    return avg_acc
+
+
+def plot_training_history(history, save_path):
+    """Plot and save training metrics"""
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+    
+    epochs = range(1, len(history['train_loss']) + 1)
+    
+    # Plot 1: Loss
+    axes[0, 0].plot(epochs, history['train_loss'], 'b-', label='Train Loss', linewidth=2)
+    axes[0, 0].plot(epochs, history['val_loss'], 'r-', label='Val Loss', linewidth=2)
+    axes[0, 0].set_xlabel('Epoch', fontsize=12)
+    axes[0, 0].set_ylabel('Loss', fontsize=12)
+    axes[0, 0].set_title('Training and Validation Loss', fontsize=14, fontweight='bold')
+    axes[0, 0].legend()
+    axes[0, 0].grid(True, alpha=0.3)
+    
+    # Plot 2: Accuracy
+    axes[0, 1].plot(epochs, history['train_acc'], 'b-', label='Train Acc', linewidth=2)
+    axes[0, 1].plot(epochs, history['val_acc'], 'r-', label='Val Acc', linewidth=2)
+    axes[0, 1].set_xlabel('Epoch', fontsize=12)
+    axes[0, 1].set_ylabel('Accuracy (%)', fontsize=12)
+    axes[0, 1].set_title('Training and Validation Accuracy', fontsize=14, fontweight='bold')
+    axes[0, 1].legend()
+    axes[0, 1].grid(True, alpha=0.3)
+    
+    # Plot 3: Learning Rate
+    axes[1, 0].plot(epochs, history['learning_rate'], 'g-', linewidth=2)
+    axes[1, 0].set_xlabel('Epoch', fontsize=12)
+    axes[1, 0].set_ylabel('Learning Rate', fontsize=12)
+    axes[1, 0].set_title('Learning Rate Schedule', fontsize=14, fontweight='bold')
+    axes[1, 0].set_yscale('log')
+    axes[1, 0].grid(True, alpha=0.3)
+    
+    # Plot 4: Val Acc vs Train Acc (Overfitting check)
+    axes[1, 1].plot(epochs, history['train_acc'], 'b-', label='Train Acc', linewidth=2)
+    axes[1, 1].plot(epochs, history['val_acc'], 'r-', label='Val Acc', linewidth=2)
+    gap = [t - v for t, v in zip(history['train_acc'], history['val_acc'])]
+    axes[1, 1].fill_between(epochs, history['val_acc'], history['train_acc'], 
+                            alpha=0.3, color='orange', label='Overfitting Gap')
+    axes[1, 1].set_xlabel('Epoch', fontsize=12)
+    axes[1, 1].set_ylabel('Accuracy (%)', fontsize=12)
+    axes[1, 1].set_title('Overfitting Analysis', fontsize=14, fontweight='bold')
+    axes[1, 1].legend()
+    axes[1, 1].grid(True, alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Training plots saved to: {save_path}")
+    plt.close()
+
+
+def save_training_summary(history, best_acc, save_path):
+    """Save training summary as JSON"""
+    summary = {
+        'config': {
+            'model': 'CvT-13',
+            'batch_size': BATCH_SIZE,
+            'img_size': IMG_SIZE,
+            'num_classes': NUM_CLASSES,
+            'num_epochs': NUM_EPOCHS,
+            'device': DEVICE,
+            'pretrained': PRETRAINED_PATH,
+        },
+        'final_metrics': {
+            'best_val_accuracy': best_acc,
+            'final_train_loss': history['train_loss'][-1],
+            'final_train_acc': history['train_acc'][-1],
+            'final_val_loss': history['val_loss'][-1],
+            'final_val_acc': history['val_acc'][-1],
+        },
+        'history': history
+    }
+    
+    with open(save_path, 'w') as f:
+        json.dump(summary, f, indent=4)
+    
+    print(f"💾 Training summary saved to: {save_path}")
+
+
+# ============================================================
+# MAIN TRAINING LOOP
+# ============================================================
+
+if __name__ == '__main__':
+    print("\n" + "="*60)
+    print("STARTING TRAINING")
+    print("="*60 + "\n")
+
+    # Initialize history tracking
+    history = {
+        'train_loss': [],
+        'train_acc': [],
+        'val_loss': [],
+        'val_acc': [],
+        'learning_rate': []
+    }
+    
+    best_acc = 0.0
+    best_epoch = 0
+
+    for epoch in range(NUM_EPOCHS):
+        train_loss, train_acc = train_one_epoch(epoch, history)
+        val_acc = validate(epoch, history)
+        lr_scheduler.step(epoch + 1)
+
+        # Save best model
+        if val_acc > best_acc:
+            best_acc = val_acc
+            best_epoch = epoch + 1
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'best_acc': best_acc,
+                'history': history,
+            }, os.path.join(OUTPUT_DIR, "best_model.pth"))
+            print(f"✅ Saved best model at epoch {epoch+1} with val acc {best_acc:.2f}%\n")
+        
+        # Save checkpoint every 10 epochs
+        if (epoch + 1) % 10 == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_acc,
+                'history': history,
+            }, os.path.join(OUTPUT_DIR, f"checkpoint_epoch_{epoch+1}.pth"))
+            print(f"💾 Checkpoint saved at epoch {epoch+1}\n")
+        
+        # Plot and save metrics every 5 epochs
+        if (epoch + 1) % 5 == 0 or epoch == NUM_EPOCHS - 1:
+            plot_training_history(history, os.path.join(OUTPUT_DIR, "training_metrics.png"))
+
+    # Final summary
+    print("="*60)
+    print(f"🎉 Training complete!")
+    print(f"Best validation accuracy: {best_acc:.2f}% at epoch {best_epoch}")
+    print(f"Final train accuracy: {history['train_acc'][-1]:.2f}%")
+    print(f"Final val accuracy: {history['val_acc'][-1]:.2f}%")
+    print("="*60)
+    
+    # Save final training summary
+    save_training_summary(history, best_acc, os.path.join(OUTPUT_DIR, "training_summary.json"))
+    
+    # Save final plot
+    plot_training_history(history, os.path.join(OUTPUT_DIR, "final_training_metrics.png"))
+    
+    print(f"\n📁 All outputs saved to: {OUTPUT_DIR}")

stage2/inference_cvt13.py

@@ -0,0 +1,637 @@
+import torch
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import sys
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+import json
+from datetime import datetime
+from sklearn.metrics import (
+    accuracy_score, precision_score, recall_score, f1_score,
+    confusion_matrix, classification_report, roc_curve, auc,
+    precision_recall_curve, average_precision_score, roc_auc_score
+)
+import seaborn as sns
+import pandas as pd
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+BASE_DIR        = "path_to_CornViT"
+
+# Path to the Microsoft CvT repository
+CVT_REPO_PATH = f"{BASE_DIR}/CvT"
+
+# Model configuration
+IMG_SIZE = 384
+NUM_CLASSES = 2
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+RUN = "cvt13_run_2025xxxx_xxxxxx"
+
+# Path to trained model
+MODEL_PATH = f"metrics/{RUN}/train/best_model.pth"
+
+# Test data folder (should have subfolders for each class like train/val structure)
+TEST_DATA_DIR = f"{BASE_DIR}/stage2/data/test"
+
+# Class names (update these to match your dataset)
+CLASS_NAMES = ["Flat", "Round"]
+
+# Output directory for evaluation results (within the same metrics folder)
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+EVAL_OUTPUT_DIR = f"metrics/{RUN}/evals/eval_{timestamp}"
+os.makedirs(EVAL_OUTPUT_DIR, exist_ok=True)
+
+
+# ============================================================
+# SETUP: Import CvT model
+# ============================================================
+
+# Fix torch._six compatibility
+cls_cvt_path = os.path.join(CVT_REPO_PATH, "lib", "models", "cls_cvt.py")
+if os.path.exists(cls_cvt_path):
+    with open(cls_cvt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    if "from torch._six import container_abcs" in content:
+        content = content.replace(
+            "from torch._six import container_abcs",
+            "import collections.abc as container_abcs"
+        )
+        content = content.replace(
+            "or pretrained_layers[0] is '*'",
+            "or pretrained_layers[0] == '*'"
+        )
+        with open(cls_cvt_path, 'w', encoding='utf-8') as f:
+            f.write(content)
+
+sys.path.insert(0, CVT_REPO_PATH)
+
+import warnings
+warnings.filterwarnings('ignore', category=SyntaxWarning)
+
+from lib.models import cls_cvt
+from lib.config import config, update_config
+
+
+# ============================================================
+# MODEL LOADING
+# ============================================================
+
+def load_model(model_path, config_path=None):
+    """Load the trained CvT model"""
+    
+    # Load config
+    if config_path is None:
+        config_path = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt", "cvt-13-384x384.yaml")
+    
+    config.defrost()
+    config.merge_from_file(config_path)
+    config.MODEL.NUM_CLASSES = NUM_CLASSES
+    config.MODEL.PRETRAINED = ''
+    config.freeze()
+    
+    # Create model
+    model = cls_cvt.get_cls_model(config)
+    
+    # Load trained weights
+    checkpoint = torch.load(model_path, map_location=DEVICE)
+    if 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+    
+    model = model.to(DEVICE)
+    model.eval()
+    
+    print(f"✅ Model loaded from: {model_path}")
+    return model
+
+
+# ============================================================
+# DATA LOADING
+# ============================================================
+
+def get_test_dataloader(test_dir, batch_size=32):
+    """Create test dataloader"""
+    test_transforms = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                           [0.229, 0.224, 0.225])
+    ])
+    
+    test_dataset = datasets.ImageFolder(test_dir, transform=test_transforms)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, 
+                            shuffle=False, num_workers=0, pin_memory=True)
+    
+    print(f"✅ Test dataset loaded: {len(test_dataset)} images")
+    print(f"   Classes: {test_dataset.classes}")
+    return test_loader, test_dataset
+
+
+# ============================================================
+# EVALUATION FUNCTIONS
+# ============================================================
+
+def evaluate_model(model, test_loader, test_dataset):
+    """
+    Evaluate model with single image predictions
+    
+    Returns:
+        all_preds: Predicted class labels
+        all_labels: Ground truth labels
+        all_probs: Predicted probabilities for all classes
+        all_confidences: Confidence scores
+        image_paths: List of image paths
+    """
+    model.eval()
+    
+    all_preds = []
+    all_labels = []
+    all_probs = []
+    all_confidences = []
+    image_paths = []
+    
+    print("\n🔍 Running single-image inference on test set...")
+    
+    # Process each image individually
+    total_images = len(test_dataset)
+    
+    for idx in range(total_images):
+        # Get single image and label
+        image, label = test_dataset[idx]
+        img_path, _ = test_dataset.samples[idx]
+        
+        # Add batch dimension and move to device
+        image = image.unsqueeze(0).to(DEVICE)
+        
+        with torch.no_grad():
+            # Forward pass
+            output = model(image)
+            
+            # Ensure output has correct shape
+            if output.dim() == 1:
+                output = output.unsqueeze(0)
+                
+            probabilities = torch.softmax(output, dim=1)
+            confidence, predicted = torch.max(probabilities, 1)
+            
+            # Collect results
+            all_preds.append(predicted.item())
+            all_labels.append(label)
+            all_probs.append(probabilities.cpu().numpy()[0])
+            all_confidences.append(confidence.item())
+            image_paths.append(img_path)
+        
+        # Progress update
+        if (idx + 1) % 50 == 0 or (idx + 1) == total_images:
+            print(f"   Processed {idx + 1}/{total_images} images...")
+    
+    print(f"✅ Inference complete: {len(all_preds)} predictions")
+    
+    return (np.array(all_preds), np.array(all_labels), np.array(all_probs), 
+            np.array(all_confidences), image_paths)
+
+
+# ============================================================
+# METRICS CALCULATION
+# ============================================================
+
+def calculate_metrics(y_true, y_pred, y_probs):
+    """Calculate all classification metrics"""
+    
+    metrics = {}
+    
+    # Basic metrics
+    metrics['accuracy'] = accuracy_score(y_true, y_pred)
+    metrics['precision_macro'] = precision_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['precision_weighted'] = precision_score(y_true, y_pred, average='weighted', zero_division=0)
+    metrics['recall_macro'] = recall_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['recall_weighted'] = recall_score(y_true, y_pred, average='weighted', zero_division=0)
+    metrics['f1_macro'] = f1_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['f1_weighted'] = f1_score(y_true, y_pred, average='weighted', zero_division=0)
+    
+    # Per-class metrics
+    precision_per_class = precision_score(y_true, y_pred, average=None, zero_division=0)
+    recall_per_class = recall_score(y_true, y_pred, average=None, zero_division=0)
+    f1_per_class = f1_score(y_true, y_pred, average=None, zero_division=0)
+    
+    metrics['per_class'] = {}
+    for i, class_name in enumerate(CLASS_NAMES):
+        metrics['per_class'][class_name] = {
+            'precision': float(precision_per_class[i]),
+            'recall': float(recall_per_class[i]),
+            'f1_score': float(f1_per_class[i])
+        }
+    
+    # ROC-AUC (for binary and multi-class)
+    if NUM_CLASSES == 2:
+        metrics['roc_auc'] = roc_auc_score(y_true, y_probs[:, 1])
+        metrics['average_precision'] = average_precision_score(y_true, y_probs[:, 1])
+    else:
+        metrics['roc_auc_ovr'] = roc_auc_score(y_true, y_probs, multi_class='ovr', average='macro')
+        metrics['roc_auc_ovo'] = roc_auc_score(y_true, y_probs, multi_class='ovo', average='macro')
+    
+    return metrics
+
+
+# ============================================================
+# VISUALIZATION FUNCTIONS
+# ============================================================
+
+def plot_confusion_matrix(y_true, y_pred, save_path):
+    """Plot and save confusion matrix"""
+    cm = confusion_matrix(y_true, y_pred)
+    
+    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
+    
+    # Raw counts
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                ax=axes[0], cbar_kws={'label': 'Count'})
+    axes[0].set_xlabel('Predicted Label', fontsize=12)
+    axes[0].set_ylabel('True Label', fontsize=12)
+    axes[0].set_title('Confusion Matrix (Counts)', fontsize=14, fontweight='bold')
+    
+    # Normalized
+    cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
+    sns.heatmap(cm_normalized, annot=True, fmt='.2%', cmap='Blues',
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                ax=axes[1], cbar_kws={'label': 'Percentage'})
+    axes[1].set_xlabel('Predicted Label', fontsize=12)
+    axes[1].set_ylabel('True Label', fontsize=12)
+    axes[1].set_title('Confusion Matrix (Normalized)', fontsize=14, fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Confusion matrix saved to: {save_path}")
+    plt.close()
+    
+    return cm
+
+
+def plot_roc_curve(y_true, y_probs, save_path):
+    """Plot ROC curve"""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    
+    if NUM_CLASSES == 2:
+        # Binary classification
+        fpr, tpr, _ = roc_curve(y_true, y_probs[:, 1])
+        roc_auc = auc(fpr, tpr)
+        
+        ax.plot(fpr, tpr, color='darkorange', lw=2,
+                label=f'ROC curve (AUC = {roc_auc:.3f})')
+        ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Classifier')
+        
+    else:
+        # Multi-class (one-vs-rest)
+        for i, class_name in enumerate(CLASS_NAMES):
+            y_true_binary = (y_true == i).astype(int)
+            fpr, tpr, _ = roc_curve(y_true_binary, y_probs[:, i])
+            roc_auc = auc(fpr, tpr)
+            ax.plot(fpr, tpr, lw=2, label=f'{class_name} (AUC = {roc_auc:.3f})')
+        
+        ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Classifier')
+    
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.05])
+    ax.set_xlabel('False Positive Rate', fontsize=12)
+    ax.set_ylabel('True Positive Rate', fontsize=12)
+    ax.set_title('Receiver Operating Characteristic (ROC) Curve', fontsize=14, fontweight='bold')
+    ax.legend(loc="lower right", fontsize=10)
+    ax.grid(alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 ROC curve saved to: {save_path}")
+    plt.close()
+
+
+def plot_precision_recall_curve(y_true, y_probs, save_path):
+    """Plot Precision-Recall curve"""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    
+    if NUM_CLASSES == 2:
+        # Binary classification
+        precision, recall, _ = precision_recall_curve(y_true, y_probs[:, 1])
+        avg_precision = average_precision_score(y_true, y_probs[:, 1])
+        
+        ax.plot(recall, precision, color='darkorange', lw=2,
+                label=f'PR curve (AP = {avg_precision:.3f})')
+        
+    else:
+        # Multi-class
+        for i, class_name in enumerate(CLASS_NAMES):
+            y_true_binary = (y_true == i).astype(int)
+            precision, recall, _ = precision_recall_curve(y_true_binary, y_probs[:, i])
+            avg_precision = average_precision_score(y_true_binary, y_probs[:, i])
+            ax.plot(recall, precision, lw=2, 
+                   label=f'{class_name} (AP = {avg_precision:.3f})')
+    
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.05])
+    ax.set_xlabel('Recall', fontsize=12)
+    ax.set_ylabel('Precision', fontsize=12)
+    ax.set_title('Precision-Recall Curve', fontsize=14, fontweight='bold')
+    ax.legend(loc="lower left", fontsize=10)
+    ax.grid(alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Precision-Recall curve saved to: {save_path}")
+    plt.close()
+
+
+def plot_class_distribution(y_true, y_pred, save_path):
+    """Plot class distribution comparison"""
+    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
+    
+    # True distribution
+    true_counts = [np.sum(y_true == i) for i in range(NUM_CLASSES)]
+    axes[0].bar(CLASS_NAMES, true_counts, color='steelblue', alpha=0.7)
+    axes[0].set_ylabel('Count', fontsize=12)
+    axes[0].set_title('True Label Distribution', fontsize=14, fontweight='bold')
+    axes[0].grid(axis='y', alpha=0.3)
+    for i, count in enumerate(true_counts):
+        axes[0].text(i, count + max(true_counts)*0.01, str(count), 
+                    ha='center', va='bottom', fontweight='bold')
+    
+    # Predicted distribution
+    pred_counts = [np.sum(y_pred == i) for i in range(NUM_CLASSES)]
+    axes[1].bar(CLASS_NAMES, pred_counts, color='coral', alpha=0.7)
+    axes[1].set_ylabel('Count', fontsize=12)
+    axes[1].set_title('Predicted Label Distribution', fontsize=14, fontweight='bold')
+    axes[1].grid(axis='y', alpha=0.3)
+    for i, count in enumerate(pred_counts):
+        axes[1].text(i, count + max(pred_counts)*0.01, str(count), 
+                    ha='center', va='bottom', fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Class distribution saved to: {save_path}")
+    plt.close()
+
+
+def plot_per_class_metrics(metrics, save_path):
+    """Plot per-class performance metrics"""
+    classes = list(metrics['per_class'].keys())
+    precision_vals = [metrics['per_class'][c]['precision'] for c in classes]
+    recall_vals = [metrics['per_class'][c]['recall'] for c in classes]
+    f1_vals = [metrics['per_class'][c]['f1_score'] for c in classes]
+    
+    x = np.arange(len(classes))
+    width = 0.25
+    
+    fig, ax = plt.subplots(figsize=(12, 7))
+    
+    bars1 = ax.bar(x - width, precision_vals, width, label='Precision', color='steelblue', alpha=0.8)
+    bars2 = ax.bar(x, recall_vals, width, label='Recall', color='coral', alpha=0.8)
+    bars3 = ax.bar(x + width, f1_vals, width, label='F1-Score', color='lightgreen', alpha=0.8)
+    
+    ax.set_ylabel('Score', fontsize=12)
+    ax.set_title('Per-Class Performance Metrics', fontsize=14, fontweight='bold')
+    ax.set_xticks(x)
+    ax.set_xticklabels(classes)
+    ax.legend(fontsize=11)
+    ax.set_ylim([0, 1.1])
+    ax.grid(axis='y', alpha=0.3)
+    
+    # Add value labels on bars
+    def autolabel(bars):
+        for bar in bars:
+            height = bar.get_height()
+            ax.text(bar.get_x() + bar.get_width()/2., height + 0.02,
+                   f'{height:.3f}', ha='center', va='bottom', fontsize=9)
+    
+    autolabel(bars1)
+    autolabel(bars2)
+    autolabel(bars3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Per-class metrics saved to: {save_path}")
+    plt.close()
+
+
+def plot_confidence_distribution(y_true, y_pred, confidences, save_path):
+    """Plot confidence score distribution for correct vs incorrect predictions"""
+    # Confidence scores are already extracted
+    correct = (y_true == y_pred)
+    
+    fig, axes = plt.subplots(2, 1, figsize=(12, 10))
+    
+    # Histogram
+    axes[0].hist(confidences[correct], bins=50, alpha=0.7, label='Correct', 
+                color='green', edgecolor='black')
+    axes[0].hist(confidences[~correct], bins=50, alpha=0.7, label='Incorrect', 
+                color='red', edgecolor='black')
+    axes[0].set_xlabel('Confidence Score', fontsize=12)
+    axes[0].set_ylabel('Frequency', fontsize=12)
+    axes[0].set_title('Confidence Distribution: Correct vs Incorrect Predictions', 
+                     fontsize=14, fontweight='bold')
+    axes[0].legend(fontsize=11)
+    axes[0].grid(alpha=0.3)
+    
+    # Box plot
+    data_to_plot = [confidences[correct], confidences[~correct]]
+    box = axes[1].boxplot(data_to_plot, labels=['Correct', 'Incorrect'], 
+                          patch_artist=True, showmeans=True)
+    box['boxes'][0].set_facecolor('lightgreen')
+    box['boxes'][1].set_facecolor('lightcoral')
+    axes[1].set_ylabel('Confidence Score', fontsize=12)
+    axes[1].set_title('Confidence Score Box Plot', fontsize=14, fontweight='bold')
+    axes[1].grid(axis='y', alpha=0.3)
+    
+    # Add statistics
+    correct_mean = np.mean(confidences[correct])
+    incorrect_mean = np.mean(confidences[~correct]) if (~correct).sum() > 0 else 0
+    axes[1].text(1, correct_mean, f'μ={correct_mean:.3f}', 
+                ha='right', va='center', fontweight='bold', fontsize=10)
+    if (~correct).sum() > 0:
+        axes[1].text(2, incorrect_mean, f'μ={incorrect_mean:.3f}', 
+                    ha='left', va='center', fontweight='bold', fontsize=10)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Confidence distribution saved to: {save_path}")
+    plt.close()
+
+
+# ============================================================
+# RESULTS SAVING
+# ============================================================
+
+def save_predictions_to_csv(image_paths, y_true, y_pred, y_probs, confidences, save_path):
+    """Save detailed predictions to CSV"""
+    results = []
+    
+    for img_path, true_label, pred, probs, conf in zip(image_paths, y_true, y_pred, y_probs, confidences):
+        result = {
+            'image_path': img_path,
+            'image_name': os.path.basename(img_path),
+            'true_label': CLASS_NAMES[true_label],
+            'true_label_idx': true_label,
+            'predicted_label': CLASS_NAMES[pred],
+            'predicted_label_idx': pred,
+            'confidence': conf,
+            'correct': pred == true_label
+        }
+        
+        # Add probabilities for each class
+        for i, class_name in enumerate(CLASS_NAMES):
+            result[f'prob_{class_name}'] = probs[i]
+        
+        results.append(result)
+    
+    df = pd.DataFrame(results)
+    df.to_csv(save_path, index=False)
+    print(f"💾 Predictions saved to: {save_path}")
+    
+    # Print some statistics
+    print(f"\n📊 Prediction Statistics:")
+    print(f"   Total images: {len(df)}")
+    print(f"   Correct predictions: {df['correct'].sum()} ({df['correct'].sum()/len(df)*100:.2f}%)")
+    print(f"   Incorrect predictions: {(~df['correct']).sum()} ({(~df['correct']).sum()/len(df)*100:.2f}%)")
+    print(f"   Average confidence: {df['confidence'].mean():.4f}")
+    print(f"   Confidence on correct: {df[df['correct']]['confidence'].mean():.4f}")
+    print(f"   Confidence on incorrect: {df[~df['correct']]['confidence'].mean():.4f}" if (~df['correct']).sum() > 0 else "")
+    
+    return df
+
+
+def save_metrics_json(metrics, save_path):
+    """Save metrics to JSON file"""
+    with open(save_path, 'w') as f:
+        json.dump(metrics, f, indent=4)
+    print(f"💾 Metrics saved to: {save_path}")
+
+
+def generate_classification_report_file(y_true, y_pred, save_path):
+    """Generate and save sklearn classification report"""
+    report = classification_report(y_true, y_pred, target_names=CLASS_NAMES, digits=4)
+    
+    with open(save_path, 'w') as f:
+        f.write("="*60 + "\n")
+        f.write("CLASSIFICATION REPORT\n")
+        f.write("="*60 + "\n\n")
+        f.write(report)
+    
+    print(f"📄 Classification report saved to: {save_path}")
+
+
+# ============================================================
+# MAIN EVALUATION PIPELINE
+# ============================================================
+
+def main():
+    """Main evaluation pipeline"""
+    
+    print("\n" + "="*60)
+    print("CvT-13 MODEL EVALUATION PIPELINE")
+    print("Single Image Prediction Mode")
+    print("="*60 + "\n")
+    
+    # Load model
+    print("📦 Loading model...")
+    model = load_model(MODEL_PATH)
+    
+    # Load test data
+    print("\n📂 Loading test data...")
+    test_loader, test_dataset = get_test_dataloader(TEST_DATA_DIR, batch_size=1)
+    
+    # Run evaluation with single image predictions
+    print("\n🔍 Evaluating model (single image predictions)...")
+    y_pred, y_true, y_probs, confidences, image_paths = evaluate_model(model, test_loader, test_dataset)
+    
+    # Calculate metrics
+    print("\n📊 Calculating metrics...")
+    metrics = calculate_metrics(y_true, y_pred, y_probs)
+    
+    # Print key metrics
+    print("\n" + "="*60)
+    print("EVALUATION RESULTS")
+    print("="*60)
+    print(f"Total Images Evaluated: {len(y_pred)}")
+    print(f"Accuracy:           {metrics['accuracy']*100:.2f}%")
+    print(f"Precision (Macro):  {metrics['precision_macro']*100:.2f}%")
+    print(f"Recall (Macro):     {metrics['recall_macro']*100:.2f}%")
+    print(f"F1-Score (Macro):   {metrics['f1_macro']*100:.2f}%")
+    if 'roc_auc' in metrics:
+        print(f"ROC-AUC:            {metrics['roc_auc']:.4f}")
+    print("\nPer-Class Metrics:")
+    for class_name, class_metrics in metrics['per_class'].items():
+        print(f"  {class_name}:")
+        print(f"    Precision: {class_metrics['precision']*100:.2f}%")
+        print(f"    Recall:    {class_metrics['recall']*100:.2f}%")
+        print(f"    F1-Score:  {class_metrics['f1_score']*100:.2f}%")
+    print("="*60)
+    
+    # Generate all visualizations
+    print("\n📊 Generating visualizations...")
+    plot_confusion_matrix(y_true, y_pred, 
+                         os.path.join(EVAL_OUTPUT_DIR, "confusion_matrix.png"))
+    plot_roc_curve(y_true, y_probs, 
+                  os.path.join(EVAL_OUTPUT_DIR, "roc_curve.png"))
+    plot_precision_recall_curve(y_true, y_probs, 
+                               os.path.join(EVAL_OUTPUT_DIR, "precision_recall_curve.png"))
+    plot_class_distribution(y_true, y_pred, 
+                          os.path.join(EVAL_OUTPUT_DIR, "class_distribution.png"))
+    plot_per_class_metrics(metrics, 
+                          os.path.join(EVAL_OUTPUT_DIR, "per_class_metrics.png"))
+    plot_confidence_distribution(y_true, y_pred, confidences,
+                                os.path.join(EVAL_OUTPUT_DIR, "confidence_distribution.png"))
+    
+    # Save results
+    print("\n💾 Saving results...")
+    df = save_predictions_to_csv(image_paths, y_true, y_pred, y_probs, confidences,
+                                os.path.join(EVAL_OUTPUT_DIR, "predictions.csv"))
+    save_metrics_json(metrics, 
+                     os.path.join(EVAL_OUTPUT_DIR, "metrics.json"))
+    generate_classification_report_file(y_true, y_pred,
+                                       os.path.join(EVAL_OUTPUT_DIR, "classification_report.txt"))
+    
+    # Save misclassified images list
+    misclassified = df[~df['correct']]
+    if len(misclassified) > 0:
+        misclassified_path = os.path.join(EVAL_OUTPUT_DIR, "misclassified_images.csv")
+        misclassified.to_csv(misclassified_path, index=False)
+        print(f"⚠️  Misclassified images saved to: {misclassified_path}")
+        print(f"   Total misclassified: {len(misclassified)}")
+    
+    # Save low confidence predictions
+    low_conf_threshold = 0.7
+    low_confidence = df[df['confidence'] < low_conf_threshold]
+    if len(low_confidence) > 0:
+        low_conf_path = os.path.join(EVAL_OUTPUT_DIR, "low_confidence_predictions.csv")
+        low_confidence.to_csv(low_conf_path, index=False)
+        print(f"⚠️  Low confidence predictions saved to: {low_conf_path}")
+        print(f"   Total with confidence < {low_conf_threshold}: {len(low_confidence)}")
+    
+    print("\n" + "="*60)
+    print(f"✅ Evaluation complete!")
+    print(f"📁 All results saved to: {EVAL_OUTPUT_DIR}")
+    print("="*60 + "\n")
+    
+    print("Generated files:")
+    print("  📊 confusion_matrix.png - Confusion matrix visualization")
+    print("  📊 roc_curve.png - ROC curve")
+    print("  📊 precision_recall_curve.png - Precision-Recall curve")
+    print("  📊 class_distribution.png - Class distribution comparison")
+    print("  📊 per_class_metrics.png - Per-class performance")
+    print("  📊 confidence_distribution.png - Confidence analysis")
+    print("  💾 predictions.csv - Detailed predictions for each image")
+    print("  💾 misclassified_images.csv - List of incorrectly classified images")
+    print("  💾 low_confidence_predictions.csv - Predictions with low confidence")
+    print("  💾 metrics.json - All metrics in JSON format")
+    print("  📄 classification_report.txt - Sklearn classification report")
+
+
+if __name__ == '__main__':
+    main()

stage2/stage2_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:445c5a5b94b86649cab12ef3c2fe4df9461f9879864c43d52a7cc9560204fcc3
+size 78733538

stage2/train_cvt13.py

@@ -0,0 +1,458 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import sys
+import os
+from timm.loss import SoftTargetCrossEntropy
+from timm.scheduler import CosineLRScheduler
+from timm.utils import accuracy
+import matplotlib.pyplot as plt
+import json
+from datetime import datetime
+
+
+# ============================================================
+# SETUP: Clone and import from Microsoft CvT repository
+# ============================================================
+"""
+First, clone the Microsoft CvT repository:
+    git clone https://github.com/microsoft/CvT.git
+    cd CvT
+    pip install -r requirements.txt
+"""
+
+BASE_DIR = "path_to_CornViT"
+
+# Add the CvT repo to Python path
+CVT_REPO_PATH = f"{BASE_DIR}/CvT"
+
+if not os.path.exists(CVT_REPO_PATH):
+    print(f"❌ CvT repository not found at {CVT_REPO_PATH}")
+    print("Please clone it: git clone https://github.com/microsoft/CvT.git")
+    sys.exit(1)
+
+# Fix torch._six compatibility BEFORE importing
+print("Applying compatibility fixes for newer PyTorch versions...")
+cls_cvt_path = os.path.join(CVT_REPO_PATH, "lib", "models", "cls_cvt.py")
+
+if os.path.exists(cls_cvt_path):
+    with open(cls_cvt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    # Fix 1: Replace torch._six import
+    if "from torch._six import container_abcs" in content:
+        content = content.replace(
+            "from torch._six import container_abcs",
+            "import collections.abc as container_abcs"
+        )
+        
+        # Fix 2: Replace 'is' with '==' for string comparison
+        content = content.replace(
+            "or pretrained_layers[0] is '*'",
+            "or pretrained_layers[0] == '*'"
+        )
+        
+        with open(cls_cvt_path, 'w', encoding='utf-8') as f:
+            f.write(content)
+        print("✅ Applied compatibility patches to cls_cvt.py")
+    else:
+        print("✅ Compatibility patches already applied")
+else:
+    print(f"❌ Could not find cls_cvt.py at {cls_cvt_path}")
+    sys.exit(1)
+
+# Now import
+sys.path.insert(0, CVT_REPO_PATH)
+
+# Suppress the SyntaxWarning
+import warnings
+warnings.filterwarnings('ignore', category=SyntaxWarning)
+
+from lib.models import cls_cvt
+from lib.config import config, update_config
+print("✅ Successfully imported Microsoft CvT models")
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+DATA_DIR        = f"{BASE_DIR}/stage2/data"
+BATCH_SIZE      = 32
+IMG_SIZE        = 384
+NUM_CLASSES     = 2
+NUM_EPOCHS      = 100
+DEVICE          = "cuda" if torch.cuda.is_available() else "cpu"
+PRETRAINED_PATH = f"{BASE_DIR}/CvT-13-384x384-IN-22k.pth"
+
+# Create output directory for saving results
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+OUTPUT_DIR = f"metrics/cvt13_run_{timestamp}"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+print(f"Metrics will be saved to: {OUTPUT_DIR}")
+
+
+# ============================================================
+# DATASET & AUGMENTATION
+# ============================================================
+
+train_transforms = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomVerticalFlip(),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+    transforms.RandomRotation(15),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+val_transforms = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+train_dataset = datasets.ImageFolder(f"{DATA_DIR}/train", transform=train_transforms)
+val_dataset = datasets.ImageFolder(f"{DATA_DIR}/val", transform=val_transforms)
+
+train_loader = DataLoader(
+    train_dataset, 
+    batch_size=BATCH_SIZE, 
+    shuffle=True, 
+    num_workers=0, 
+    pin_memory=True,
+    drop_last=True
+)
+val_loader = DataLoader(
+    val_dataset, 
+    batch_size=BATCH_SIZE,
+    shuffle=False, 
+    num_workers=0, 
+    pin_memory=True,
+    drop_last=True
+)
+
+
+# ============================================================
+# MODEL SETUP - Using Microsoft CvT Implementation
+# ============================================================
+
+# Load the CvT-13 config from the repository
+cvt_config_path = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt", "cvt-13-384x384.yaml")
+
+if not os.path.exists(cvt_config_path):
+    print(f"⚠️ Config file not found at {cvt_config_path}")
+    print("Available configs:")
+    config_dir = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt")
+    if os.path.exists(config_dir):
+        for f in os.listdir(config_dir):
+            if f.endswith('.yaml'):
+                print(f"  - {f}")
+    sys.exit(1)
+
+print(f"Loading config from: {cvt_config_path}")
+
+# Load config directly using merge_from_file
+config.defrost()
+config.merge_from_file(cvt_config_path)
+
+# Update the number of classes for our task
+config.MODEL.NUM_CLASSES = NUM_CLASSES
+config.MODEL.PRETRAINED = ''  # We'll load weights manually
+config.freeze()
+
+print("Creating CvT-13 model...")
+# Create model using the official CvT architecture
+model = cls_cvt.get_cls_model(config)
+model = model.to(DEVICE)
+
+# Load pretrained weights
+if os.path.exists(PRETRAINED_PATH):
+    print(f"Loading pretrained weights from {PRETRAINED_PATH}")
+    try:
+        checkpoint = torch.load(PRETRAINED_PATH, map_location=DEVICE)
+        
+        # Handle different checkpoint formats
+        if 'model' in checkpoint:
+            state_dict = checkpoint['model']
+        elif 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+        
+        # Remove 'module.' prefix if present
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            name = k.replace("module.", "")
+            new_state_dict[name] = v
+        
+        # Remove head layers from pretrained weights (they have different dimensions)
+        filtered_state_dict = {k: v for k, v in new_state_dict.items() if 'head' not in k}
+        
+        # Load weights - strict=False will only load matching layers
+        missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False)
+        
+        # Count how many weights were actually loaded
+        loaded_keys = [k for k in filtered_state_dict.keys() if k in model.state_dict()]
+        print(f"✅ Loaded pretrained weights: {len(loaded_keys)} layers from backbone")
+        print(f"   Head layer initialized randomly for {NUM_CLASSES} classes")
+        
+        # Show what's missing (should only be head-related)
+        head_missing = [k for k in missing_keys if 'head' in k]
+        other_missing = [k for k in missing_keys if 'head' not in k]
+        
+        if other_missing:
+            print(f"⚠️  Warning - Missing non-head keys: {other_missing}")
+        if unexpected_keys:
+            print(f"⚠️  Unexpected keys: {unexpected_keys}")
+            
+    except Exception as e:
+        print(f"⚠️ Error loading pretrained weights: {e}")
+        import traceback
+        traceback.print_exc()
+        print("Continuing with random initialization...")
+else:
+    print(f"⚠️ Pretrained weights not found at {PRETRAINED_PATH}")
+    print("Training from scratch...")
+
+# Freeze backbone - only train the head for faster training and less overfitting
+print("Freezing backbone layers (keeping only head trainable)...")
+for name, param in model.named_parameters():
+    if "head" not in name:
+        param.requires_grad = False
+
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Trainable parameters: {trainable_params:,} / {total_params:,}")
+print(f"Frozen parameters: {total_params - trainable_params:,}")
+
+
+# ============================================================
+# OPTIMIZER AND LOSS
+# ============================================================
+
+optimizer = optim.AdamW(
+    filter(lambda p: p.requires_grad, model.parameters()), 
+    lr=1e-4, 
+    weight_decay=0.05
+)
+
+criterion = SoftTargetCrossEntropy()
+
+lr_scheduler = CosineLRScheduler(
+    optimizer,
+    t_initial=NUM_EPOCHS,
+    lr_min=1e-6,
+    warmup_t=5,
+    warmup_lr_init=1e-5,
+)
+
+
+# ============================================================
+# TRAINING & VALIDATION LOOP
+# ============================================================
+
+def train_one_epoch(epoch, history):
+    model.train()
+    total_loss, total_acc = 0, 0
+    
+    for images, targets in train_loader:
+        images, targets = images.to(DEVICE), targets.to(DEVICE)
+        images, targets = mixup_fn(images, targets)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+
+        acc1, _ = accuracy(outputs, targets.argmax(dim=1), topk=(1, 5))
+        total_loss += loss.item()
+        total_acc += acc1.item()
+
+    avg_loss = total_loss / len(train_loader)
+    avg_acc = total_acc / len(train_loader)
+    
+    history['train_loss'].append(avg_loss)
+    history['train_acc'].append(avg_acc)
+    history['learning_rate'].append(optimizer.param_groups[0]['lr'])
+    
+    print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] | Train Loss: {avg_loss:.4f} | Train Acc: {avg_acc:.2f}% | LR: {optimizer.param_groups[0]['lr']:.6f}")
+    return avg_loss, avg_acc
+
+
+def validate(epoch, history):
+    model.eval()
+    total_loss, total_acc = 0, 0
+    
+    with torch.no_grad():
+        for images, targets in val_loader:
+            images, targets = images.to(DEVICE), targets.to(DEVICE)
+            outputs = model(images)
+            loss = nn.CrossEntropyLoss()(outputs, targets)
+            acc1, _ = accuracy(outputs, targets, topk=(1, 5))
+            
+            total_loss += loss.item()
+            total_acc += acc1.item()
+    
+    avg_loss = total_loss / len(val_loader)
+    avg_acc = total_acc / len(val_loader)
+    
+    history['val_loss'].append(avg_loss)
+    history['val_acc'].append(avg_acc)
+    
+    print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] | Val Loss: {avg_loss:.4f} | Val Acc: {avg_acc:.2f}%")
+    return avg_acc
+
+
+def plot_training_history(history, save_path):
+    """Plot and save training metrics"""
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+    
+    epochs = range(1, len(history['train_loss']) + 1)
+    
+    # Plot 1: Loss
+    axes[0, 0].plot(epochs, history['train_loss'], 'b-', label='Train Loss', linewidth=2)
+    axes[0, 0].plot(epochs, history['val_loss'], 'r-', label='Val Loss', linewidth=2)
+    axes[0, 0].set_xlabel('Epoch', fontsize=12)
+    axes[0, 0].set_ylabel('Loss', fontsize=12)
+    axes[0, 0].set_title('Training and Validation Loss', fontsize=14, fontweight='bold')
+    axes[0, 0].legend()
+    axes[0, 0].grid(True, alpha=0.3)
+    
+    # Plot 2: Accuracy
+    axes[0, 1].plot(epochs, history['train_acc'], 'b-', label='Train Acc', linewidth=2)
+    axes[0, 1].plot(epochs, history['val_acc'], 'r-', label='Val Acc', linewidth=2)
+    axes[0, 1].set_xlabel('Epoch', fontsize=12)
+    axes[0, 1].set_ylabel('Accuracy (%)', fontsize=12)
+    axes[0, 1].set_title('Training and Validation Accuracy', fontsize=14, fontweight='bold')
+    axes[0, 1].legend()
+    axes[0, 1].grid(True, alpha=0.3)
+    
+    # Plot 3: Learning Rate
+    axes[1, 0].plot(epochs, history['learning_rate'], 'g-', linewidth=2)
+    axes[1, 0].set_xlabel('Epoch', fontsize=12)
+    axes[1, 0].set_ylabel('Learning Rate', fontsize=12)
+    axes[1, 0].set_title('Learning Rate Schedule', fontsize=14, fontweight='bold')
+    axes[1, 0].set_yscale('log')
+    axes[1, 0].grid(True, alpha=0.3)
+    
+    # Plot 4: Val Acc vs Train Acc (Overfitting check)
+    axes[1, 1].plot(epochs, history['train_acc'], 'b-', label='Train Acc', linewidth=2)
+    axes[1, 1].plot(epochs, history['val_acc'], 'r-', label='Val Acc', linewidth=2)
+    gap = [t - v for t, v in zip(history['train_acc'], history['val_acc'])]
+    axes[1, 1].fill_between(epochs, history['val_acc'], history['train_acc'], 
+                            alpha=0.3, color='orange', label='Overfitting Gap')
+    axes[1, 1].set_xlabel('Epoch', fontsize=12)
+    axes[1, 1].set_ylabel('Accuracy (%)', fontsize=12)
+    axes[1, 1].set_title('Overfitting Analysis', fontsize=14, fontweight='bold')
+    axes[1, 1].legend()
+    axes[1, 1].grid(True, alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Training plots saved to: {save_path}")
+    plt.close()
+
+
+def save_training_summary(history, best_acc, save_path):
+    """Save training summary as JSON"""
+    summary = {
+        'config': {
+            'model': 'CvT-13',
+            'batch_size': BATCH_SIZE,
+            'img_size': IMG_SIZE,
+            'num_classes': NUM_CLASSES,
+            'num_epochs': NUM_EPOCHS,
+            'device': DEVICE,
+            'pretrained': PRETRAINED_PATH,
+        },
+        'final_metrics': {
+            'best_val_accuracy': best_acc,
+            'final_train_loss': history['train_loss'][-1],
+            'final_train_acc': history['train_acc'][-1],
+            'final_val_loss': history['val_loss'][-1],
+            'final_val_acc': history['val_acc'][-1],
+        },
+        'history': history
+    }
+    
+    with open(save_path, 'w') as f:
+        json.dump(summary, f, indent=4)
+    
+    print(f"💾 Training summary saved to: {save_path}")
+
+
+# ============================================================
+# MAIN TRAINING LOOP
+# ============================================================
+
+if __name__ == '__main__':
+    print("\n" + "="*60)
+    print("STARTING TRAINING")
+    print("="*60 + "\n")
+
+    # Initialize history tracking
+    history = {
+        'train_loss': [],
+        'train_acc': [],
+        'val_loss': [],
+        'val_acc': [],
+        'learning_rate': []
+    }
+    
+    best_acc = 0.0
+    best_epoch = 0
+
+    for epoch in range(NUM_EPOCHS):
+        train_loss, train_acc = train_one_epoch(epoch, history)
+        val_acc = validate(epoch, history)
+        lr_scheduler.step(epoch + 1)
+
+        # Save best model
+        if val_acc > best_acc:
+            best_acc = val_acc
+            best_epoch = epoch + 1
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'best_acc': best_acc,
+                'history': history,
+            }, os.path.join(OUTPUT_DIR, "best_model.pth"))
+            print(f"✅ Saved best model at epoch {epoch+1} with val acc {best_acc:.2f}%\n")
+        
+        # Save checkpoint every 10 epochs
+        if (epoch + 1) % 10 == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_acc,
+                'history': history,
+            }, os.path.join(OUTPUT_DIR, f"checkpoint_epoch_{epoch+1}.pth"))
+            print(f"💾 Checkpoint saved at epoch {epoch+1}\n")
+        
+        # Plot and save metrics every 5 epochs
+        if (epoch + 1) % 5 == 0 or epoch == NUM_EPOCHS - 1:
+            plot_training_history(history, os.path.join(OUTPUT_DIR, "training_metrics.png"))
+
+    # Final summary
+    print("="*60)
+    print(f"🎉 Training complete!")
+    print(f"Best validation accuracy: {best_acc:.2f}% at epoch {best_epoch}")
+    print(f"Final train accuracy: {history['train_acc'][-1]:.2f}%")
+    print(f"Final val accuracy: {history['val_acc'][-1]:.2f}%")
+    print("="*60)
+    
+    # Save final training summary
+    save_training_summary(history, best_acc, os.path.join(OUTPUT_DIR, "training_summary.json"))
+    
+    # Save final plot
+    plot_training_history(history, os.path.join(OUTPUT_DIR, "final_training_metrics.png"))
+    
+    print(f"\n📁 All outputs saved to: {OUTPUT_DIR}")

stage3/inference_cvt13.py

@@ -0,0 +1,637 @@
+import torch
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import sys
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+import json
+from datetime import datetime
+from sklearn.metrics import (
+    accuracy_score, precision_score, recall_score, f1_score,
+    confusion_matrix, classification_report, roc_curve, auc,
+    precision_recall_curve, average_precision_score, roc_auc_score
+)
+import seaborn as sns
+import pandas as pd
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+BASE_DIR        = "path_to_CornViT"
+
+# Path to the Microsoft CvT repository
+CVT_REPO_PATH = f"{BASE_DIR}/CvT"
+
+# Model configuration
+IMG_SIZE = 384
+NUM_CLASSES = 2
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+RUN = "cvt13_run_2025xxxx_xxxxxx"
+
+# Path to trained model
+MODEL_PATH = f"metrics/{RUN}/train/best_model.pth"
+
+# Test data folder (should have subfolders for each class like train/val structure)
+TEST_DATA_DIR = f"{BASE_DIR}/stage3/data/test"
+
+# Class names (update these to match your dataset)
+CLASS_NAMES = ["Up", "Down"]
+
+# Output directory for evaluation results (within the same metrics folder)
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+EVAL_OUTPUT_DIR = f"metrics/{RUN}/evals/eval_{timestamp}"
+os.makedirs(EVAL_OUTPUT_DIR, exist_ok=True)
+
+
+# ============================================================
+# SETUP: Import CvT model
+# ============================================================
+
+# Fix torch._six compatibility
+cls_cvt_path = os.path.join(CVT_REPO_PATH, "lib", "models", "cls_cvt.py")
+if os.path.exists(cls_cvt_path):
+    with open(cls_cvt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    if "from torch._six import container_abcs" in content:
+        content = content.replace(
+            "from torch._six import container_abcs",
+            "import collections.abc as container_abcs"
+        )
+        content = content.replace(
+            "or pretrained_layers[0] is '*'",
+            "or pretrained_layers[0] == '*'"
+        )
+        with open(cls_cvt_path, 'w', encoding='utf-8') as f:
+            f.write(content)
+
+sys.path.insert(0, CVT_REPO_PATH)
+
+import warnings
+warnings.filterwarnings('ignore', category=SyntaxWarning)
+
+from lib.models import cls_cvt
+from lib.config import config, update_config
+
+
+# ============================================================
+# MODEL LOADING
+# ============================================================
+
+def load_model(model_path, config_path=None):
+    """Load the trained CvT model"""
+    
+    # Load config
+    if config_path is None:
+        config_path = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt", "cvt-13-384x384.yaml")
+    
+    config.defrost()
+    config.merge_from_file(config_path)
+    config.MODEL.NUM_CLASSES = NUM_CLASSES
+    config.MODEL.PRETRAINED = ''
+    config.freeze()
+    
+    # Create model
+    model = cls_cvt.get_cls_model(config)
+    
+    # Load trained weights
+    checkpoint = torch.load(model_path, map_location=DEVICE)
+    if 'model_state_dict' in checkpoint:
+        model.load_state_dict(checkpoint['model_state_dict'])
+    else:
+        model.load_state_dict(checkpoint)
+    
+    model = model.to(DEVICE)
+    model.eval()
+    
+    print(f"✅ Model loaded from: {model_path}")
+    return model
+
+
+# ============================================================
+# DATA LOADING
+# ============================================================
+
+def get_test_dataloader(test_dir, batch_size=32):
+    """Create test dataloader"""
+    test_transforms = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406],
+                           [0.229, 0.224, 0.225])
+    ])
+    
+    test_dataset = datasets.ImageFolder(test_dir, transform=test_transforms)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, 
+                            shuffle=False, num_workers=0, pin_memory=True)
+    
+    print(f"✅ Test dataset loaded: {len(test_dataset)} images")
+    print(f"   Classes: {test_dataset.classes}")
+    return test_loader, test_dataset
+
+
+# ============================================================
+# EVALUATION FUNCTIONS
+# ============================================================
+
+def evaluate_model(model, test_loader, test_dataset):
+    """
+    Evaluate model with single image predictions
+    
+    Returns:
+        all_preds: Predicted class labels
+        all_labels: Ground truth labels
+        all_probs: Predicted probabilities for all classes
+        all_confidences: Confidence scores
+        image_paths: List of image paths
+    """
+    model.eval()
+    
+    all_preds = []
+    all_labels = []
+    all_probs = []
+    all_confidences = []
+    image_paths = []
+    
+    print("\n🔍 Running single-image inference on test set...")
+    
+    # Process each image individually
+    total_images = len(test_dataset)
+    
+    for idx in range(total_images):
+        # Get single image and label
+        image, label = test_dataset[idx]
+        img_path, _ = test_dataset.samples[idx]
+        
+        # Add batch dimension and move to device
+        image = image.unsqueeze(0).to(DEVICE)
+        
+        with torch.no_grad():
+            # Forward pass
+            output = model(image)
+            
+            # Ensure output has correct shape
+            if output.dim() == 1:
+                output = output.unsqueeze(0)
+                
+            probabilities = torch.softmax(output, dim=1)
+            confidence, predicted = torch.max(probabilities, 1)
+            
+            # Collect results
+            all_preds.append(predicted.item())
+            all_labels.append(label)
+            all_probs.append(probabilities.cpu().numpy()[0])
+            all_confidences.append(confidence.item())
+            image_paths.append(img_path)
+        
+        # Progress update
+        if (idx + 1) % 50 == 0 or (idx + 1) == total_images:
+            print(f"   Processed {idx + 1}/{total_images} images...")
+    
+    print(f"✅ Inference complete: {len(all_preds)} predictions")
+    
+    return (np.array(all_preds), np.array(all_labels), np.array(all_probs), 
+            np.array(all_confidences), image_paths)
+
+
+# ============================================================
+# METRICS CALCULATION
+# ============================================================
+
+def calculate_metrics(y_true, y_pred, y_probs):
+    """Calculate all classification metrics"""
+    
+    metrics = {}
+    
+    # Basic metrics
+    metrics['accuracy'] = accuracy_score(y_true, y_pred)
+    metrics['precision_macro'] = precision_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['precision_weighted'] = precision_score(y_true, y_pred, average='weighted', zero_division=0)
+    metrics['recall_macro'] = recall_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['recall_weighted'] = recall_score(y_true, y_pred, average='weighted', zero_division=0)
+    metrics['f1_macro'] = f1_score(y_true, y_pred, average='macro', zero_division=0)
+    metrics['f1_weighted'] = f1_score(y_true, y_pred, average='weighted', zero_division=0)
+    
+    # Per-class metrics
+    precision_per_class = precision_score(y_true, y_pred, average=None, zero_division=0)
+    recall_per_class = recall_score(y_true, y_pred, average=None, zero_division=0)
+    f1_per_class = f1_score(y_true, y_pred, average=None, zero_division=0)
+    
+    metrics['per_class'] = {}
+    for i, class_name in enumerate(CLASS_NAMES):
+        metrics['per_class'][class_name] = {
+            'precision': float(precision_per_class[i]),
+            'recall': float(recall_per_class[i]),
+            'f1_score': float(f1_per_class[i])
+        }
+    
+    # ROC-AUC (for binary and multi-class)
+    if NUM_CLASSES == 2:
+        metrics['roc_auc'] = roc_auc_score(y_true, y_probs[:, 1])
+        metrics['average_precision'] = average_precision_score(y_true, y_probs[:, 1])
+    else:
+        metrics['roc_auc_ovr'] = roc_auc_score(y_true, y_probs, multi_class='ovr', average='macro')
+        metrics['roc_auc_ovo'] = roc_auc_score(y_true, y_probs, multi_class='ovo', average='macro')
+    
+    return metrics
+
+
+# ============================================================
+# VISUALIZATION FUNCTIONS
+# ============================================================
+
+def plot_confusion_matrix(y_true, y_pred, save_path):
+    """Plot and save confusion matrix"""
+    cm = confusion_matrix(y_true, y_pred)
+    
+    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
+    
+    # Raw counts
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                ax=axes[0], cbar_kws={'label': 'Count'})
+    axes[0].set_xlabel('Predicted Label', fontsize=12)
+    axes[0].set_ylabel('True Label', fontsize=12)
+    axes[0].set_title('Confusion Matrix (Counts)', fontsize=14, fontweight='bold')
+    
+    # Normalized
+    cm_normalized = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
+    sns.heatmap(cm_normalized, annot=True, fmt='.2%', cmap='Blues',
+                xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES,
+                ax=axes[1], cbar_kws={'label': 'Percentage'})
+    axes[1].set_xlabel('Predicted Label', fontsize=12)
+    axes[1].set_ylabel('True Label', fontsize=12)
+    axes[1].set_title('Confusion Matrix (Normalized)', fontsize=14, fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Confusion matrix saved to: {save_path}")
+    plt.close()
+    
+    return cm
+
+
+def plot_roc_curve(y_true, y_probs, save_path):
+    """Plot ROC curve"""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    
+    if NUM_CLASSES == 2:
+        # Binary classification
+        fpr, tpr, _ = roc_curve(y_true, y_probs[:, 1])
+        roc_auc = auc(fpr, tpr)
+        
+        ax.plot(fpr, tpr, color='darkorange', lw=2,
+                label=f'ROC curve (AUC = {roc_auc:.3f})')
+        ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Classifier')
+        
+    else:
+        # Multi-class (one-vs-rest)
+        for i, class_name in enumerate(CLASS_NAMES):
+            y_true_binary = (y_true == i).astype(int)
+            fpr, tpr, _ = roc_curve(y_true_binary, y_probs[:, i])
+            roc_auc = auc(fpr, tpr)
+            ax.plot(fpr, tpr, lw=2, label=f'{class_name} (AUC = {roc_auc:.3f})')
+        
+        ax.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random Classifier')
+    
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.05])
+    ax.set_xlabel('False Positive Rate', fontsize=12)
+    ax.set_ylabel('True Positive Rate', fontsize=12)
+    ax.set_title('Receiver Operating Characteristic (ROC) Curve', fontsize=14, fontweight='bold')
+    ax.legend(loc="lower right", fontsize=10)
+    ax.grid(alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 ROC curve saved to: {save_path}")
+    plt.close()
+
+
+def plot_precision_recall_curve(y_true, y_probs, save_path):
+    """Plot Precision-Recall curve"""
+    fig, ax = plt.subplots(figsize=(10, 8))
+    
+    if NUM_CLASSES == 2:
+        # Binary classification
+        precision, recall, _ = precision_recall_curve(y_true, y_probs[:, 1])
+        avg_precision = average_precision_score(y_true, y_probs[:, 1])
+        
+        ax.plot(recall, precision, color='darkorange', lw=2,
+                label=f'PR curve (AP = {avg_precision:.3f})')
+        
+    else:
+        # Multi-class
+        for i, class_name in enumerate(CLASS_NAMES):
+            y_true_binary = (y_true == i).astype(int)
+            precision, recall, _ = precision_recall_curve(y_true_binary, y_probs[:, i])
+            avg_precision = average_precision_score(y_true_binary, y_probs[:, i])
+            ax.plot(recall, precision, lw=2, 
+                   label=f'{class_name} (AP = {avg_precision:.3f})')
+    
+    ax.set_xlim([0.0, 1.0])
+    ax.set_ylim([0.0, 1.05])
+    ax.set_xlabel('Recall', fontsize=12)
+    ax.set_ylabel('Precision', fontsize=12)
+    ax.set_title('Precision-Recall Curve', fontsize=14, fontweight='bold')
+    ax.legend(loc="lower left", fontsize=10)
+    ax.grid(alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Precision-Recall curve saved to: {save_path}")
+    plt.close()
+
+
+def plot_class_distribution(y_true, y_pred, save_path):
+    """Plot class distribution comparison"""
+    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
+    
+    # True distribution
+    true_counts = [np.sum(y_true == i) for i in range(NUM_CLASSES)]
+    axes[0].bar(CLASS_NAMES, true_counts, color='steelblue', alpha=0.7)
+    axes[0].set_ylabel('Count', fontsize=12)
+    axes[0].set_title('True Label Distribution', fontsize=14, fontweight='bold')
+    axes[0].grid(axis='y', alpha=0.3)
+    for i, count in enumerate(true_counts):
+        axes[0].text(i, count + max(true_counts)*0.01, str(count), 
+                    ha='center', va='bottom', fontweight='bold')
+    
+    # Predicted distribution
+    pred_counts = [np.sum(y_pred == i) for i in range(NUM_CLASSES)]
+    axes[1].bar(CLASS_NAMES, pred_counts, color='coral', alpha=0.7)
+    axes[1].set_ylabel('Count', fontsize=12)
+    axes[1].set_title('Predicted Label Distribution', fontsize=14, fontweight='bold')
+    axes[1].grid(axis='y', alpha=0.3)
+    for i, count in enumerate(pred_counts):
+        axes[1].text(i, count + max(pred_counts)*0.01, str(count), 
+                    ha='center', va='bottom', fontweight='bold')
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Class distribution saved to: {save_path}")
+    plt.close()
+
+
+def plot_per_class_metrics(metrics, save_path):
+    """Plot per-class performance metrics"""
+    classes = list(metrics['per_class'].keys())
+    precision_vals = [metrics['per_class'][c]['precision'] for c in classes]
+    recall_vals = [metrics['per_class'][c]['recall'] for c in classes]
+    f1_vals = [metrics['per_class'][c]['f1_score'] for c in classes]
+    
+    x = np.arange(len(classes))
+    width = 0.25
+    
+    fig, ax = plt.subplots(figsize=(12, 7))
+    
+    bars1 = ax.bar(x - width, precision_vals, width, label='Precision', color='steelblue', alpha=0.8)
+    bars2 = ax.bar(x, recall_vals, width, label='Recall', color='coral', alpha=0.8)
+    bars3 = ax.bar(x + width, f1_vals, width, label='F1-Score', color='lightgreen', alpha=0.8)
+    
+    ax.set_ylabel('Score', fontsize=12)
+    ax.set_title('Per-Class Performance Metrics', fontsize=14, fontweight='bold')
+    ax.set_xticks(x)
+    ax.set_xticklabels(classes)
+    ax.legend(fontsize=11)
+    ax.set_ylim([0, 1.1])
+    ax.grid(axis='y', alpha=0.3)
+    
+    # Add value labels on bars
+    def autolabel(bars):
+        for bar in bars:
+            height = bar.get_height()
+            ax.text(bar.get_x() + bar.get_width()/2., height + 0.02,
+                   f'{height:.3f}', ha='center', va='bottom', fontsize=9)
+    
+    autolabel(bars1)
+    autolabel(bars2)
+    autolabel(bars3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Per-class metrics saved to: {save_path}")
+    plt.close()
+
+
+def plot_confidence_distribution(y_true, y_pred, confidences, save_path):
+    """Plot confidence score distribution for correct vs incorrect predictions"""
+    # Confidence scores are already extracted
+    correct = (y_true == y_pred)
+    
+    fig, axes = plt.subplots(2, 1, figsize=(12, 10))
+    
+    # Histogram
+    axes[0].hist(confidences[correct], bins=50, alpha=0.7, label='Correct', 
+                color='green', edgecolor='black')
+    axes[0].hist(confidences[~correct], bins=50, alpha=0.7, label='Incorrect', 
+                color='red', edgecolor='black')
+    axes[0].set_xlabel('Confidence Score', fontsize=12)
+    axes[0].set_ylabel('Frequency', fontsize=12)
+    axes[0].set_title('Confidence Distribution: Correct vs Incorrect Predictions', 
+                     fontsize=14, fontweight='bold')
+    axes[0].legend(fontsize=11)
+    axes[0].grid(alpha=0.3)
+    
+    # Box plot
+    data_to_plot = [confidences[correct], confidences[~correct]]
+    box = axes[1].boxplot(data_to_plot, labels=['Correct', 'Incorrect'], 
+                          patch_artist=True, showmeans=True)
+    box['boxes'][0].set_facecolor('lightgreen')
+    box['boxes'][1].set_facecolor('lightcoral')
+    axes[1].set_ylabel('Confidence Score', fontsize=12)
+    axes[1].set_title('Confidence Score Box Plot', fontsize=14, fontweight='bold')
+    axes[1].grid(axis='y', alpha=0.3)
+    
+    # Add statistics
+    correct_mean = np.mean(confidences[correct])
+    incorrect_mean = np.mean(confidences[~correct]) if (~correct).sum() > 0 else 0
+    axes[1].text(1, correct_mean, f'μ={correct_mean:.3f}', 
+                ha='right', va='center', fontweight='bold', fontsize=10)
+    if (~correct).sum() > 0:
+        axes[1].text(2, incorrect_mean, f'μ={incorrect_mean:.3f}', 
+                    ha='left', va='center', fontweight='bold', fontsize=10)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Confidence distribution saved to: {save_path}")
+    plt.close()
+
+
+# ============================================================
+# RESULTS SAVING
+# ============================================================
+
+def save_predictions_to_csv(image_paths, y_true, y_pred, y_probs, confidences, save_path):
+    """Save detailed predictions to CSV"""
+    results = []
+    
+    for img_path, true_label, pred, probs, conf in zip(image_paths, y_true, y_pred, y_probs, confidences):
+        result = {
+            'image_path': img_path,
+            'image_name': os.path.basename(img_path),
+            'true_label': CLASS_NAMES[true_label],
+            'true_label_idx': true_label,
+            'predicted_label': CLASS_NAMES[pred],
+            'predicted_label_idx': pred,
+            'confidence': conf,
+            'correct': pred == true_label
+        }
+        
+        # Add probabilities for each class
+        for i, class_name in enumerate(CLASS_NAMES):
+            result[f'prob_{class_name}'] = probs[i]
+        
+        results.append(result)
+    
+    df = pd.DataFrame(results)
+    df.to_csv(save_path, index=False)
+    print(f"💾 Predictions saved to: {save_path}")
+    
+    # Print some statistics
+    print(f"\n📊 Prediction Statistics:")
+    print(f"   Total images: {len(df)}")
+    print(f"   Correct predictions: {df['correct'].sum()} ({df['correct'].sum()/len(df)*100:.2f}%)")
+    print(f"   Incorrect predictions: {(~df['correct']).sum()} ({(~df['correct']).sum()/len(df)*100:.2f}%)")
+    print(f"   Average confidence: {df['confidence'].mean():.4f}")
+    print(f"   Confidence on correct: {df[df['correct']]['confidence'].mean():.4f}")
+    print(f"   Confidence on incorrect: {df[~df['correct']]['confidence'].mean():.4f}" if (~df['correct']).sum() > 0 else "")
+    
+    return df
+
+
+def save_metrics_json(metrics, save_path):
+    """Save metrics to JSON file"""
+    with open(save_path, 'w') as f:
+        json.dump(metrics, f, indent=4)
+    print(f"💾 Metrics saved to: {save_path}")
+
+
+def generate_classification_report_file(y_true, y_pred, save_path):
+    """Generate and save sklearn classification report"""
+    report = classification_report(y_true, y_pred, target_names=CLASS_NAMES, digits=4)
+    
+    with open(save_path, 'w') as f:
+        f.write("="*60 + "\n")
+        f.write("CLASSIFICATION REPORT\n")
+        f.write("="*60 + "\n\n")
+        f.write(report)
+    
+    print(f"📄 Classification report saved to: {save_path}")
+
+
+# ============================================================
+# MAIN EVALUATION PIPELINE
+# ============================================================
+
+def main():
+    """Main evaluation pipeline"""
+    
+    print("\n" + "="*60)
+    print("CvT-13 MODEL EVALUATION PIPELINE")
+    print("Single Image Prediction Mode")
+    print("="*60 + "\n")
+    
+    # Load model
+    print("📦 Loading model...")
+    model = load_model(MODEL_PATH)
+    
+    # Load test data
+    print("\n📂 Loading test data...")
+    test_loader, test_dataset = get_test_dataloader(TEST_DATA_DIR, batch_size=1)
+    
+    # Run evaluation with single image predictions
+    print("\n🔍 Evaluating model (single image predictions)...")
+    y_pred, y_true, y_probs, confidences, image_paths = evaluate_model(model, test_loader, test_dataset)
+    
+    # Calculate metrics
+    print("\n📊 Calculating metrics...")
+    metrics = calculate_metrics(y_true, y_pred, y_probs)
+    
+    # Print key metrics
+    print("\n" + "="*60)
+    print("EVALUATION RESULTS")
+    print("="*60)
+    print(f"Total Images Evaluated: {len(y_pred)}")
+    print(f"Accuracy:           {metrics['accuracy']*100:.2f}%")
+    print(f"Precision (Macro):  {metrics['precision_macro']*100:.2f}%")
+    print(f"Recall (Macro):     {metrics['recall_macro']*100:.2f}%")
+    print(f"F1-Score (Macro):   {metrics['f1_macro']*100:.2f}%")
+    if 'roc_auc' in metrics:
+        print(f"ROC-AUC:            {metrics['roc_auc']:.4f}")
+    print("\nPer-Class Metrics:")
+    for class_name, class_metrics in metrics['per_class'].items():
+        print(f"  {class_name}:")
+        print(f"    Precision: {class_metrics['precision']*100:.2f}%")
+        print(f"    Recall:    {class_metrics['recall']*100:.2f}%")
+        print(f"    F1-Score:  {class_metrics['f1_score']*100:.2f}%")
+    print("="*60)
+    
+    # Generate all visualizations
+    print("\n📊 Generating visualizations...")
+    plot_confusion_matrix(y_true, y_pred, 
+                         os.path.join(EVAL_OUTPUT_DIR, "confusion_matrix.png"))
+    plot_roc_curve(y_true, y_probs, 
+                  os.path.join(EVAL_OUTPUT_DIR, "roc_curve.png"))
+    plot_precision_recall_curve(y_true, y_probs, 
+                               os.path.join(EVAL_OUTPUT_DIR, "precision_recall_curve.png"))
+    plot_class_distribution(y_true, y_pred, 
+                          os.path.join(EVAL_OUTPUT_DIR, "class_distribution.png"))
+    plot_per_class_metrics(metrics, 
+                          os.path.join(EVAL_OUTPUT_DIR, "per_class_metrics.png"))
+    plot_confidence_distribution(y_true, y_pred, confidences,
+                                os.path.join(EVAL_OUTPUT_DIR, "confidence_distribution.png"))
+    
+    # Save results
+    print("\n💾 Saving results...")
+    df = save_predictions_to_csv(image_paths, y_true, y_pred, y_probs, confidences,
+                                os.path.join(EVAL_OUTPUT_DIR, "predictions.csv"))
+    save_metrics_json(metrics, 
+                     os.path.join(EVAL_OUTPUT_DIR, "metrics.json"))
+    generate_classification_report_file(y_true, y_pred,
+                                       os.path.join(EVAL_OUTPUT_DIR, "classification_report.txt"))
+    
+    # Save misclassified images list
+    misclassified = df[~df['correct']]
+    if len(misclassified) > 0:
+        misclassified_path = os.path.join(EVAL_OUTPUT_DIR, "misclassified_images.csv")
+        misclassified.to_csv(misclassified_path, index=False)
+        print(f"⚠️  Misclassified images saved to: {misclassified_path}")
+        print(f"   Total misclassified: {len(misclassified)}")
+    
+    # Save low confidence predictions
+    low_conf_threshold = 0.7
+    low_confidence = df[df['confidence'] < low_conf_threshold]
+    if len(low_confidence) > 0:
+        low_conf_path = os.path.join(EVAL_OUTPUT_DIR, "low_confidence_predictions.csv")
+        low_confidence.to_csv(low_conf_path, index=False)
+        print(f"⚠️  Low confidence predictions saved to: {low_conf_path}")
+        print(f"   Total with confidence < {low_conf_threshold}: {len(low_confidence)}")
+    
+    print("\n" + "="*60)
+    print(f"✅ Evaluation complete!")
+    print(f"📁 All results saved to: {EVAL_OUTPUT_DIR}")
+    print("="*60 + "\n")
+    
+    print("Generated files:")
+    print("  📊 confusion_matrix.png - Confusion matrix visualization")
+    print("  📊 roc_curve.png - ROC curve")
+    print("  📊 precision_recall_curve.png - Precision-Recall curve")
+    print("  📊 class_distribution.png - Class distribution comparison")
+    print("  📊 per_class_metrics.png - Per-class performance")
+    print("  📊 confidence_distribution.png - Confidence analysis")
+    print("  💾 predictions.csv - Detailed predictions for each image")
+    print("  💾 misclassified_images.csv - List of incorrectly classified images")
+    print("  💾 low_confidence_predictions.csv - Predictions with low confidence")
+    print("  💾 metrics.json - All metrics in JSON format")
+    print("  📄 classification_report.txt - Sklearn classification report")
+
+
+if __name__ == '__main__':
+    main()

stage3/stage3_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:445c5a5b94b86649cab12ef3c2fe4df9461f9879864c43d52a7cc9560204fcc3
+size 78733538

stage3/train_cvt13.py

@@ -0,0 +1,457 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+import sys
+import os
+from timm.loss import SoftTargetCrossEntropy
+from timm.scheduler import CosineLRScheduler
+from timm.utils import accuracy
+import matplotlib.pyplot as plt
+import json
+from datetime import datetime
+
+
+# ============================================================
+# SETUP: Clone and import from Microsoft CvT repository
+# ============================================================
+"""
+First, clone the Microsoft CvT repository:
+    git clone https://github.com/microsoft/CvT.git
+    cd CvT
+    pip install -r requirements.txt
+"""
+
+BASE_DIR = "path_to_CornViT"
+
+# Add the CvT repo to Python path
+CVT_REPO_PATH = f"{BASE_DIR}/CvT"
+
+if not os.path.exists(CVT_REPO_PATH):
+    print(f"❌ CvT repository not found at {CVT_REPO_PATH}")
+    print("Please clone it: git clone https://github.com/microsoft/CvT.git")
+    sys.exit(1)
+
+# Fix torch._six compatibility BEFORE importing
+print("Applying compatibility fixes for newer PyTorch versions...")
+cls_cvt_path = os.path.join(CVT_REPO_PATH, "lib", "models", "cls_cvt.py")
+
+if os.path.exists(cls_cvt_path):
+    with open(cls_cvt_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    
+    # Fix 1: Replace torch._six import
+    if "from torch._six import container_abcs" in content:
+        content = content.replace(
+            "from torch._six import container_abcs",
+            "import collections.abc as container_abcs"
+        )
+        
+        # Fix 2: Replace 'is' with '==' for string comparison
+        content = content.replace(
+            "or pretrained_layers[0] is '*'",
+            "or pretrained_layers[0] == '*'"
+        )
+        
+        with open(cls_cvt_path, 'w', encoding='utf-8') as f:
+            f.write(content)
+        print("✅ Applied compatibility patches to cls_cvt.py")
+    else:
+        print("✅ Compatibility patches already applied")
+else:
+    print(f"❌ Could not find cls_cvt.py at {cls_cvt_path}")
+    sys.exit(1)
+
+# Now import
+sys.path.insert(0, CVT_REPO_PATH)
+
+# Suppress the SyntaxWarning
+import warnings
+warnings.filterwarnings('ignore', category=SyntaxWarning)
+
+from lib.models import cls_cvt
+from lib.config import config, update_config
+print("✅ Successfully imported Microsoft CvT models")
+
+
+# ============================================================
+# CONFIGURATION
+# ============================================================
+
+DATA_DIR        = f"{BASE_DIR}/stage3/data"
+BATCH_SIZE      = 32
+IMG_SIZE        = 384
+NUM_CLASSES     = 2
+NUM_EPOCHS      = 100
+DEVICE          = "cuda" if torch.cuda.is_available() else "cpu"
+PRETRAINED_PATH = f"{BASE_DIR}/CvT-13-384x384-IN-22k.pth"
+
+# Create output directory for saving results
+timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+OUTPUT_DIR = f"metrics/cvt13_run_{timestamp}"
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+print(f"Metrics will be saved to: {OUTPUT_DIR}")
+
+
+# ============================================================
+# DATASET & AUGMENTATION
+# ============================================================
+
+train_transforms = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.RandomHorizontalFlip(),
+    transforms.RandomVerticalFlip(),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+    transforms.RandomRotation(15),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+val_transforms = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406],
+                         [0.229, 0.224, 0.225])
+])
+
+train_dataset = datasets.ImageFolder(f"{DATA_DIR}/train", transform=train_transforms)
+val_dataset = datasets.ImageFolder(f"{DATA_DIR}/val", transform=val_transforms)
+
+train_loader = DataLoader(
+    train_dataset, 
+    batch_size=BATCH_SIZE, 
+    shuffle=True, 
+    num_workers=0, 
+    pin_memory=True,
+    drop_last=True
+)
+val_loader = DataLoader(
+    val_dataset, 
+    batch_size=BATCH_SIZE,
+    shuffle=False, 
+    num_workers=0, 
+    pin_memory=True,
+    drop_last=True
+)
+
+
+# ============================================================
+# MODEL SETUP - Using Microsoft CvT Implementation
+# ============================================================
+
+# Load the CvT-13 config from the repository
+cvt_config_path = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt", "cvt-13-384x384.yaml")
+
+if not os.path.exists(cvt_config_path):
+    print(f"⚠️ Config file not found at {cvt_config_path}")
+    print("Available configs:")
+    config_dir = os.path.join(CVT_REPO_PATH, "experiments", "imagenet", "cvt")
+    if os.path.exists(config_dir):
+        for f in os.listdir(config_dir):
+            if f.endswith('.yaml'):
+                print(f"  - {f}")
+    sys.exit(1)
+
+print(f"Loading config from: {cvt_config_path}")
+
+# Load config directly using merge_from_file
+config.defrost()
+config.merge_from_file(cvt_config_path)
+
+# Update the number of classes for our task
+config.MODEL.NUM_CLASSES = NUM_CLASSES
+config.MODEL.PRETRAINED = ''  # We'll load weights manually
+config.freeze()
+
+print("Creating CvT-13 model...")
+# Create model using the official CvT architecture
+model = cls_cvt.get_cls_model(config)
+model = model.to(DEVICE)
+
+# Load pretrained weights
+if os.path.exists(PRETRAINED_PATH):
+    print(f"Loading pretrained weights from {PRETRAINED_PATH}")
+    try:
+        checkpoint = torch.load(PRETRAINED_PATH, map_location=DEVICE)
+        
+        # Handle different checkpoint formats
+        if 'model' in checkpoint:
+            state_dict = checkpoint['model']
+        elif 'state_dict' in checkpoint:
+            state_dict = checkpoint['state_dict']
+        else:
+            state_dict = checkpoint
+        
+        # Remove 'module.' prefix if present
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            name = k.replace("module.", "")
+            new_state_dict[name] = v
+        
+        # Remove head layers from pretrained weights (they have different dimensions)
+        filtered_state_dict = {k: v for k, v in new_state_dict.items() if 'head' not in k}
+        
+        # Load weights - strict=False will only load matching layers
+        missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False)
+        
+        # Count how many weights were actually loaded
+        loaded_keys = [k for k in filtered_state_dict.keys() if k in model.state_dict()]
+        print(f"✅ Loaded pretrained weights: {len(loaded_keys)} layers from backbone")
+        print(f"   Head layer initialized randomly for {NUM_CLASSES} classes")
+        
+        # Show what's missing (should only be head-related)
+        head_missing = [k for k in missing_keys if 'head' in k]
+        other_missing = [k for k in missing_keys if 'head' not in k]
+        
+        if other_missing:
+            print(f"⚠️  Warning - Missing non-head keys: {other_missing}")
+        if unexpected_keys:
+            print(f"⚠️  Unexpected keys: {unexpected_keys}")
+            
+    except Exception as e:
+        print(f"⚠️ Error loading pretrained weights: {e}")
+        import traceback
+        traceback.print_exc()
+        print("Continuing with random initialization...")
+else:
+    print(f"⚠️ Pretrained weights not found at {PRETRAINED_PATH}")
+    print("Training from scratch...")
+
+# Freeze backbone - only train the head for faster training and less overfitting
+print("Freezing backbone layers (keeping only head trainable)...")
+for name, param in model.named_parameters():
+    if "head" not in name:
+        param.requires_grad = False
+
+trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Trainable parameters: {trainable_params:,} / {total_params:,}")
+print(f"Frozen parameters: {total_params - trainable_params:,}")
+
+
+# ============================================================
+# OPTIMIZER AND LOSS
+# ============================================================
+
+optimizer = optim.AdamW(
+    filter(lambda p: p.requires_grad, model.parameters()), 
+    lr=1e-4, 
+    weight_decay=0.05
+)
+
+criterion = SoftTargetCrossEntropy()
+
+lr_scheduler = CosineLRScheduler(
+    optimizer,
+    t_initial=NUM_EPOCHS,
+    lr_min=1e-6,
+    warmup_t=5,
+    warmup_lr_init=1e-5,
+)
+
+
+# ============================================================
+# TRAINING & VALIDATION LOOP
+# ============================================================
+
+def train_one_epoch(epoch, history):
+    model.train()
+    total_loss, total_acc = 0, 0
+    
+    for images, targets in train_loader:
+        images, targets = images.to(DEVICE), targets.to(DEVICE)
+
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, targets)
+        loss.backward()
+        optimizer.step()
+
+        acc1, _ = accuracy(outputs, targets.argmax(dim=1), topk=(1, 5))
+        total_loss += loss.item()
+        total_acc += acc1.item()
+
+    avg_loss = total_loss / len(train_loader)
+    avg_acc = total_acc / len(train_loader)
+    
+    history['train_loss'].append(avg_loss)
+    history['train_acc'].append(avg_acc)
+    history['learning_rate'].append(optimizer.param_groups[0]['lr'])
+    
+    print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] | Train Loss: {avg_loss:.4f} | Train Acc: {avg_acc:.2f}% | LR: {optimizer.param_groups[0]['lr']:.6f}")
+    return avg_loss, avg_acc
+
+
+def validate(epoch, history):
+    model.eval()
+    total_loss, total_acc = 0, 0
+    
+    with torch.no_grad():
+        for images, targets in val_loader:
+            images, targets = images.to(DEVICE), targets.to(DEVICE)
+            outputs = model(images)
+            loss = nn.CrossEntropyLoss()(outputs, targets)
+            acc1, _ = accuracy(outputs, targets, topk=(1, 5))
+            
+            total_loss += loss.item()
+            total_acc += acc1.item()
+    
+    avg_loss = total_loss / len(val_loader)
+    avg_acc = total_acc / len(val_loader)
+    
+    history['val_loss'].append(avg_loss)
+    history['val_acc'].append(avg_acc)
+    
+    print(f"Epoch [{epoch+1}/{NUM_EPOCHS}] | Val Loss: {avg_loss:.4f} | Val Acc: {avg_acc:.2f}%")
+    return avg_acc
+
+
+def plot_training_history(history, save_path):
+    """Plot and save training metrics"""
+    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
+    
+    epochs = range(1, len(history['train_loss']) + 1)
+    
+    # Plot 1: Loss
+    axes[0, 0].plot(epochs, history['train_loss'], 'b-', label='Train Loss', linewidth=2)
+    axes[0, 0].plot(epochs, history['val_loss'], 'r-', label='Val Loss', linewidth=2)
+    axes[0, 0].set_xlabel('Epoch', fontsize=12)
+    axes[0, 0].set_ylabel('Loss', fontsize=12)
+    axes[0, 0].set_title('Training and Validation Loss', fontsize=14, fontweight='bold')
+    axes[0, 0].legend()
+    axes[0, 0].grid(True, alpha=0.3)
+    
+    # Plot 2: Accuracy
+    axes[0, 1].plot(epochs, history['train_acc'], 'b-', label='Train Acc', linewidth=2)
+    axes[0, 1].plot(epochs, history['val_acc'], 'r-', label='Val Acc', linewidth=2)
+    axes[0, 1].set_xlabel('Epoch', fontsize=12)
+    axes[0, 1].set_ylabel('Accuracy (%)', fontsize=12)
+    axes[0, 1].set_title('Training and Validation Accuracy', fontsize=14, fontweight='bold')
+    axes[0, 1].legend()
+    axes[0, 1].grid(True, alpha=0.3)
+    
+    # Plot 3: Learning Rate
+    axes[1, 0].plot(epochs, history['learning_rate'], 'g-', linewidth=2)
+    axes[1, 0].set_xlabel('Epoch', fontsize=12)
+    axes[1, 0].set_ylabel('Learning Rate', fontsize=12)
+    axes[1, 0].set_title('Learning Rate Schedule', fontsize=14, fontweight='bold')
+    axes[1, 0].set_yscale('log')
+    axes[1, 0].grid(True, alpha=0.3)
+    
+    # Plot 4: Val Acc vs Train Acc (Overfitting check)
+    axes[1, 1].plot(epochs, history['train_acc'], 'b-', label='Train Acc', linewidth=2)
+    axes[1, 1].plot(epochs, history['val_acc'], 'r-', label='Val Acc', linewidth=2)
+    gap = [t - v for t, v in zip(history['train_acc'], history['val_acc'])]
+    axes[1, 1].fill_between(epochs, history['val_acc'], history['train_acc'], 
+                            alpha=0.3, color='orange', label='Overfitting Gap')
+    axes[1, 1].set_xlabel('Epoch', fontsize=12)
+    axes[1, 1].set_ylabel('Accuracy (%)', fontsize=12)
+    axes[1, 1].set_title('Overfitting Analysis', fontsize=14, fontweight='bold')
+    axes[1, 1].legend()
+    axes[1, 1].grid(True, alpha=0.3)
+    
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches='tight')
+    print(f"📊 Training plots saved to: {save_path}")
+    plt.close()
+
+
+def save_training_summary(history, best_acc, save_path):
+    """Save training summary as JSON"""
+    summary = {
+        'config': {
+            'model': 'CvT-13',
+            'batch_size': BATCH_SIZE,
+            'img_size': IMG_SIZE,
+            'num_classes': NUM_CLASSES,
+            'num_epochs': NUM_EPOCHS,
+            'device': DEVICE,
+            'pretrained': PRETRAINED_PATH,
+        },
+        'final_metrics': {
+            'best_val_accuracy': best_acc,
+            'final_train_loss': history['train_loss'][-1],
+            'final_train_acc': history['train_acc'][-1],
+            'final_val_loss': history['val_loss'][-1],
+            'final_val_acc': history['val_acc'][-1],
+        },
+        'history': history
+    }
+    
+    with open(save_path, 'w') as f:
+        json.dump(summary, f, indent=4)
+    
+    print(f"💾 Training summary saved to: {save_path}")
+
+
+# ============================================================
+# MAIN TRAINING LOOP
+# ============================================================
+
+if __name__ == '__main__':
+    print("\n" + "="*60)
+    print("STARTING TRAINING")
+    print("="*60 + "\n")
+
+    # Initialize history tracking
+    history = {
+        'train_loss': [],
+        'train_acc': [],
+        'val_loss': [],
+        'val_acc': [],
+        'learning_rate': []
+    }
+    
+    best_acc = 0.0
+    best_epoch = 0
+
+    for epoch in range(NUM_EPOCHS):
+        train_loss, train_acc = train_one_epoch(epoch, history)
+        val_acc = validate(epoch, history)
+        lr_scheduler.step(epoch + 1)
+
+        # Save best model
+        if val_acc > best_acc:
+            best_acc = val_acc
+            best_epoch = epoch + 1
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'best_acc': best_acc,
+                'history': history,
+            }, os.path.join(OUTPUT_DIR, "best_model.pth"))
+            print(f"✅ Saved best model at epoch {epoch+1} with val acc {best_acc:.2f}%\n")
+        
+        # Save checkpoint every 10 epochs
+        if (epoch + 1) % 10 == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_acc': val_acc,
+                'history': history,
+            }, os.path.join(OUTPUT_DIR, f"checkpoint_epoch_{epoch+1}.pth"))
+            print(f"💾 Checkpoint saved at epoch {epoch+1}\n")
+        
+        # Plot and save metrics every 5 epochs
+        if (epoch + 1) % 5 == 0 or epoch == NUM_EPOCHS - 1:
+            plot_training_history(history, os.path.join(OUTPUT_DIR, "training_metrics.png"))
+
+    # Final summary
+    print("="*60)
+    print(f"🎉 Training complete!")
+    print(f"Best validation accuracy: {best_acc:.2f}% at epoch {best_epoch}")
+    print(f"Final train accuracy: {history['train_acc'][-1]:.2f}%")
+    print(f"Final val accuracy: {history['val_acc'][-1]:.2f}%")
+    print("="*60)
+    
+    # Save final training summary
+    save_training_summary(history, best_acc, os.path.join(OUTPUT_DIR, "training_summary.json"))
+    
+    # Save final plot
+    plot_training_history(history, os.path.join(OUTPUT_DIR, "final_training_metrics.png"))
+    
+    print(f"\n📁 All outputs saved to: {OUTPUT_DIR}")