train.py

"""
Training script for surgical instrument classification
"""

import os
import pickle
import cv2
import pandas as pd
import numpy as np
from utils.utils import extract_features_from_image, fit_pca_transformer, train_svm_model, augment_image
from utils.utils import extract_features_from_image, fit_pca_transformer, augment_image
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, classification_report

def train_and_save_model(base_path, images_folder, gt_csv, save_dir, n_components=100):
    """
    Complete training pipeline that saves everything needed for submission
    
    Args:
        base_path: Base directory path
        images_folder: Folder name containing images
        gt_csv: Ground truth CSV filename
        save_dir: Directory to save model artifacts
        n_components: Number of PCA components
    """
    
    print("="*80)
    print("TRAINING SURGICAL INSTRUMENT CLASSIFIER")
    print("="*80)
    
    # Setup paths
    PATH_TO_GT = os.path.join(base_path, gt_csv)
    PATH_TO_IMAGES = os.path.join(base_path, images_folder)
    
    print(f"\nConfiguration:")
    print(f"  Ground Truth: {PATH_TO_GT}")
    print(f"  Images: {PATH_TO_IMAGES}")
    print(f"  PCA Components: {n_components}")
    print(f"  Save directory: {save_dir}")
    
    # Load ground truth
    df = pd.read_csv(PATH_TO_GT)
    print(f"\nLoaded {len(df)} training samples")
    print(f"\nLabel distribution:")
    print(df['category_id'].value_counts().sort_index())
    
    # Extract features
    print(f"\n{'='*80}")
    print("STEP 1: FEATURE EXTRACTION WITH AUGMENTATION")
    print("="*80)
    
    # Augmentation configuration
    AUGMENTATIONS_PER_IMAGE = 2  # Conservative: 3x total dataset
    
    print(f"\nAugmentation settings:")
    print(f"  Augmentations per image: {AUGMENTATIONS_PER_IMAGE}")
    print(f"  Rotation range: -10° to +10°")
    print(f"  Brightness range: 0.9 to 1.1")
    print(f"  Horizontal flip: Yes")
    print(f"  Gaussian noise: σ=3")
    print(f"  Expected total samples: {len(df) * (1 + AUGMENTATIONS_PER_IMAGE)}")
    
    features = []
    labels = []
    
    for i in range(len(df)):
        if i % 500 == 0:
            print(f"  Processing {i}/{len(df)* (1 + AUGMENTATIONS_PER_IMAGE)}...")
        
        image_name = df.iloc[i]["file_name"]
        label = df.iloc[i]["category_id"]
        
        path_to_image = os.path.join(PATH_TO_IMAGES, image_name)
        
        try:
            image = cv2.imread(path_to_image)
            if image is None:
                print(f"  Warning: Could not read {image_name}, skipping...")
                continue
            
            # ORIGINAL IMAGE
            original_features = extract_features_from_image(image)
            features.append(original_features)
            labels.append(label)
            
            # AUGMENTED VERSIONS
            for aug_idx in range(AUGMENTATIONS_PER_IMAGE):
                # Generate augmented image
                aug_image = augment_image(
                    image,
                    rotation_range=(-10, 10),
                    brightness_range=(0.9, 1.1),
                    add_noise=True,
                    noise_sigma=3,
                    horizontal_flip=(aug_idx == 0)  # Only flip on first augmentation
                )
                
                # Extract features from augmented image
                aug_features = extract_features_from_image(aug_image)
                features.append(aug_features)
                labels.append(label)
            
        except Exception as e:
            print(f"  Error processing {image_name}: {e}")
            continue
    
    features_array = np.array(features)
    labels_array = np.array(labels)
    
    print(f"\nFeature extraction complete!")
    print(f"  Original samples: {len(df)}")
    print(f"  Total samples (with augmentation): {len(features)}")
    print(f"  Features shape: {features_array.shape}")
    print(f"  Labels shape: {labels_array.shape}")
    print(f"  Feature dimension: {features_array.shape[1]}")
    
    # Apply PCA
    print(f"\n{'='*80}")
    print("STEP 2: DIMENSIONALITY REDUCTION (PCA)")
    print("="*80)
    
    pca_params, features_reduced = fit_pca_transformer(features_array, n_components)
    
    print(f"  Reduced from {features_array.shape[1]} to {n_components} dimensions")
    print(f"  Explained variance: {pca_params['cumulative_variance'][-1]:.4f}")
    
    # Train SVM with Grid Search
    print(f"\n{'='*80}")
    print("STEP 3: TRAINING SVM CLASSIFIER WITH GRID SEARCH")
    print("="*80)
    
    # Split data for training and testing
    X_train, X_test, y_train, y_test = train_test_split(
        features_reduced, 
        labels_array, 
        test_size=0.2, 
        random_state=42,
        stratify=labels_array
    )
    
    print(f"\nData split:")
    print(f"  Training samples: {len(X_train)}")
    print(f"  Test samples: {len(X_test)}")
    
    # Define parameter grid
    param_grid = {
        'C': [1, 10, 50, 100],
        'gamma': ['scale', 0.001, 0.01, 0.1],
        'kernel': ['rbf']
    }
    
    print(f"\nGrid Search parameters:")
    print(f"  C values: {param_grid['C']}")
    print(f"  Gamma values: {param_grid['gamma']}")
    print(f"  Kernel: {param_grid['kernel']}")
    print(f"  Total combinations: {len(param_grid['C']) * len(param_grid['gamma'])}")
    print(f"  Cross-validation folds: 3")
    print(f"\nThis will take 15-30 minutes...")
    
    # Perform Grid Search
    grid_search = GridSearchCV(
        SVC(),
        param_grid,
        cv=3,
        scoring='f1_macro',
        n_jobs=-1,
        verbose=2
    )
    
    print("\nStarting Grid Search...")
    grid_search.fit(X_train, y_train)
    
    # Get best model
    svm_model = grid_search.best_estimator_
    
    print(f"\n{'='*80}")
    print("GRID SEARCH COMPLETE!")
    print("="*80)
    print(f"\nBest parameters found:")
    print(f"  C: {grid_search.best_params_['C']}")
    print(f"  Gamma: {grid_search.best_params_['gamma']}")
    print(f"  Kernel: {grid_search.best_params_['kernel']}")
    print(f"\nBest cross-validation F1-score: {grid_search.best_score_:.4f}")
    
    # Train final model and evaluate
    print(f"\n{'='*80}")
    print("FINAL MODEL EVALUATION")
    print("="*80)
    
    # Training set performance
    y_train_pred = svm_model.predict(X_train)
    train_accuracy = accuracy_score(y_train, y_train_pred)
    train_f1 = f1_score(y_train, y_train_pred, average='macro')
    train_precision = precision_score(y_train, y_train_pred, average='macro')
    train_recall = recall_score(y_train, y_train_pred, average='macro')
    
    # Test set performance
    y_test_pred = svm_model.predict(X_test)
    test_accuracy = accuracy_score(y_test, y_test_pred)
    test_f1 = f1_score(y_test, y_test_pred, average='macro')
    test_precision = precision_score(y_test, y_test_pred, average='macro')
    test_recall = recall_score(y_test, y_test_pred, average='macro')
    
    print(f"\nTraining Set Performance:")
    print(f"  Accuracy:  {train_accuracy:.4f}")
    print(f"  Precision: {train_precision:.4f}")
    print(f"  Recall:    {train_recall:.4f}")
    print(f"  F1-score:  {train_f1:.4f}")
    
    print(f"\nTest Set Performance:")
    print(f"  Accuracy:  {test_accuracy:.4f}")
    print(f"  Precision: {test_precision:.4f}")
    print(f"  Recall:    {test_recall:.4f}")
    print(f"  F1-score:  {test_f1:.4f}")
    
    print(f"\nDetailed Classification Report:")
    print(classification_report(y_test, y_test_pred, 
                                target_names=[f'Class {i}' for i in sorted(np.unique(labels_array))]))
    
    print(f"\nModel Details:")
    print(f"  Support vectors: {len(svm_model.support_)}")
    print(f"  Support vectors per class: {svm_model.n_support_}")

    # Save SVM model
    model_path = os.path.join(save_dir, "multiclass_model.pkl")
    with open(model_path, "wb") as f:
        pickle.dump(svm_model, f)
    print(f"  ✓ Saved SVM model: {model_path}")
    
    # Save PCA parameters
    pca_path = os.path.join(save_dir, "pca_params.pkl")
    with open(pca_path, "wb") as f:
        pickle.dump(pca_params, f)
    print(f"  ✓ Saved PCA params: {pca_path}")
    

    print(f"\n{'='*80}")
    print("TRAINING COMPLETE!")
    print("="*80)
    print(f"\nFinal Optimized Results:")
    print(f"  Best Parameters: C={grid_search.best_params_['C']}, gamma={grid_search.best_params_['gamma']}")
    print(f"  CV F1-score: {grid_search.best_score_:.4f}")
    print(f"  Test F1-score: {test_f1:.4f}")
    print(f"  Test Precision: {test_precision:.4f}")
    print(f"  Test Recall: {test_recall:.4f}")
    print(f"\nFiles saved to: {save_dir}")
    print(f"\nNext steps:")
    print(f"  1. Create a 'utils' folder in your HuggingFace repository")
    print(f"  2. Copy utils.py into the 'utils' folder")
    print(f"  3. Copy script.py, multiclass_model.pkl, and pca_params.pkl to the repository root")
    print(f"  4. Create an empty __init__.py file in the 'utils' folder")
    print(f"  5. Submit to competition!")


if __name__ == "__main__":
    
    BASE_PATH = "C:/Users/anna2/ISM/ANNA/phase1a-data-augmentation"
    IMAGES_FOLDER = "C:/Users/anna2/ISM/Images"
    GT_CSV = "C:/Users/anna2/ISM/Baselines/phase_1a/gt_for_classification_multiclass_from_filenames_0_index.csv"
    
    SAVE_DIR = "C:/Users/anna2/ISM/ANNA/phase1a-data-augmentation"
    
    # Number of PCA components
    N_COMPONENTS = 250 #can be adjusted 
    
    # Train and save
    train_and_save_model(BASE_PATH, IMAGES_FOLDER, GT_CSV, SAVE_DIR, N_COMPONENTS)