main.py

"""Main entry point for NanoImage."""

from __future__ import annotations

import random

import hydra
import numpy as np
from omegaconf import DictConfig
from rich.console import Console
from rich.table import Table

from config.config_utils import print_config_table
from nanoimage.attacks import evaluate_robustness
from nanoimage.data import load_mnist
from nanoimage.model import NeuralNetwork
from nanoimage.trainer import train_standard


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)


def print_results(results: dict) -> None:
    console = Console()
    table = Table(title="Robustness Report", show_header=True, header_style="bold magenta")
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="green")
    table.add_row("Clean Accuracy", f"{results['clean_accuracy']*100:.2f}%")
    table.add_row("Adversarial Accuracy", f"{results['adv_accuracy']*100:.2f}%")
    table.add_row("Samples Evaluated", str(int(results['n_samples'])))
    console.print(table)


@hydra.main(version_base=None, config_path="config", config_name="defaults")
def main(cfg: DictConfig) -> None:
    """Main function: train -> evaluate -> attack -> report."""
    print_config_table(cfg, style="tree")

    set_seed(cfg.seed)

    X_train, y_train, X_test, y_test = load_mnist()  # noqa: N806

    model = NeuralNetwork(
        layers=list(cfg.model.layers),
        activation=cfg.model.activation,
        seed=cfg.seed,
    )

    train_standard(
        model,
        X_train,
        y_train,
        lr=cfg.training.learning_rate,
        epochs=cfg.training.epochs,
        sampling=cfg.training.sampling,
    )

    # Evaluate on a subset for speed (first 1000 test images)
    results = evaluate_robustness(model, X_test[:1000], y_test[:1000], cfg.attack.epsilon)
    print_results(results)


if __name__ == "__main__":
    main()