main.py

from lightning.pytorch import Trainer, seed_everything
from lightning.pytorch.callbacks import ModelCheckpoint
import hydra
from omegaconf import DictConfig , OmegaConf
from hydra.utils import instantiate
import torch
import os
from pathlib import Path
import warnings
from pydantic._internal._generate_schema import UnsupportedFieldAttributeWarning
from omegaconf import OmegaConf
from datetime import datetime


warnings.filterwarnings(
    "ignore", category=UnsupportedFieldAttributeWarning
)
warnings.filterwarnings("ignore", message=".*resume.*ignored.*offline.*")
warnings.filterwarnings("ignore", message=".*isinstance.*LeafSpec.*deprecated.*")

def generate_run_name(model_name: str) -> str:
    """Generate a meaningful run name with timestamp."""
    now = datetime.now()
    timestamp = now.strftime("%Y_%m_%d_%H_%M_%S")
    return f"{model_name}_{timestamp}"

@hydra.main(version_base=None, config_path='./configs', config_name='experiment')
def main(cfg: DictConfig):
    # Speed optimization: Use Tensor Cores efficiently on NVIDIA GPUs
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision('medium')  # 'medium' or 'high' for speed

    seed_everything(cfg.seed, workers=True)
    OmegaConf.register_new_resolver("torch.tensor", lambda x: torch.tensor(x))
    data_module = instantiate(cfg.dataset)

    if cfg.mode == "tune":
        import matplotlib.pyplot as plt
        from lightning.pytorch.tuner import Tuner
        model = instantiate(cfg.model)
        setattr(model, "lr", cfg.model.cfg.optimizer.lr)
        trainer = Trainer(
            accelerator=cfg.accelerator,
            devices=cfg.devices,
            precision=cfg.precision,
            logger=False,
            enable_checkpointing=False,
            enable_progress_bar=True,
        )
        tuner = Tuner(trainer)
        lr_finder = tuner.lr_find(model, datamodule=data_module)
        fig = lr_finder.plot(suggest=True)

        # Organized tune output directory (same naming as logs)
        model_name = cfg.model._target_.split('.')[-1]
        tune_run_name = generate_run_name(model_name)
        output_dir = Path("logs/tune") / tune_run_name
        output_dir.mkdir(parents=True, exist_ok=True)

        plot_path = output_dir / "lr_finder_plot.png"
        fig.savefig(plot_path)

        print(f"\n{'='*70}")
        print(f"📊 Learning Rate Finder Results")
        print(f"{'='*70}")
        print(f"  Run Name:     {tune_run_name}")
        print(f"  Suggested LR: {lr_finder.suggestion()}")
        print(f"  Plot saved:   {plot_path}")
        print(f"{'='*70}\n")

    elif cfg.mode in ('train' , 'test'):
        model = instantiate(cfg.model)

        # Calculate and print actual FLOPs (optional, non-critical)
        try:
            from fvcore.nn import FlopCountAnalysis, flop_count_table

            # Use fvcore instead of thop (more reliable)
            device = 'cuda' if torch.cuda.is_available() else 'cpu'
            model_temp = model.to(device)
            dummy_input = torch.randn(2, 1, 3, 256, 256).to(device)

            # Calculate FLOPs
            flops = FlopCountAnalysis(model_temp, dummy_input)
            total_flops = flops.total()

            # Format the number
            if total_flops >= 1e9:
                flops_str = f"{total_flops/1e9:.3f}G"
            elif total_flops >= 1e6:
                flops_str = f"{total_flops/1e6:.3f}M"
            else:
                flops_str = f"{total_flops/1e3:.3f}K"

            print(f"\n{'='*70}")
            print(f"📊 Actual Model Statistics")
            print(f"{'='*70}")
            print(f"  Parameters: 290.3K")
            print(f"  FLOPs: {flops_str}")
            print(f"{'='*70}\n")

            del dummy_input
            torch.cuda.empty_cache() if torch.cuda.is_available() else None

        except ImportError:
            # fvcore not available, skip FLOPs calculation
            pass
        except Exception as e:
            # FLOPs calculation failed, but training can continue
            pass

        # Ensure model is in train mode
        model.train()

        # Extract model name from config path
        model_name = cfg.model._target_.split('.')[-1]  # e.g., 'TinyCDEnhanced'
        run_name = generate_run_name(model_name)

        # Handle both single logger and multiple loggers
        if 'loggers' in cfg.logger:
            # Multiple loggers (e.g., WandB + CSV)
            loggers = [instantiate(logger_cfg) for logger_cfg in cfg.logger.loggers]
            logger = loggers
            primary_logger = loggers[0]  # Use first logger for special operations
        else:
            # Single logger (backward compatible)
            logger = instantiate(cfg.logger)
            primary_logger = logger

        # Override logger name with meaningful run name
        if hasattr(primary_logger, '_name'):
            primary_logger._name = run_name
        elif hasattr(primary_logger, 'experiment') and hasattr(primary_logger.experiment, 'name'):
            primary_logger.experiment.name = run_name

        global_hparams = {
            "seed": cfg.seed,
            "accumulate_grad_batches": cfg.gradient_accumulation,
            "precision": cfg.precision,
            "gradient_clip_val": cfg.gradient_clip_val,
            "devices": cfg.devices,
            "epochs": cfg.epochs,
            "model": cfg.model,
            "dataset": cfg.dataset,
        }

        # Log hyperparams to all loggers
        if isinstance(logger, list):
            for log in logger:
                log.log_hyperparams(global_hparams)
        else:
            logger.log_hyperparams(global_hparams)

        # Watch model (only for WandB logger)
        if cfg.grad_track:
            if hasattr(primary_logger, 'watch'):
                primary_logger.watch(model, log="all", log_freq=cfg.grad_track)

        # Organized checkpoint directory structure
        checkpoint_dir = Path("logs/checkpoints") / run_name

        # Set best metrics path for the model
        model.best_metrics_path = str(checkpoint_dir)

        checkpoint_callback = ModelCheckpoint(
            dirpath=checkpoint_dir,
            monitor="val_loss",
            mode="min",
            save_top_k=-1,
            filename="{epoch:02d}",
            save_weights_only=True,
        )

        # Organized logging directory structure
        log_dir = Path("logs/csv") / run_name

        # Create config backup directory for this run
        config_backup_dir = checkpoint_dir.parent / f"{run_name}_configs"
        config_backup_dir.mkdir(parents=True, exist_ok=True)

        # Save complete configuration to YAML file
        full_config_path = config_backup_dir / "full_config.yaml"
        with open(full_config_path, 'w') as f:
            f.write("# Full Configuration for Run: " + run_name + "\n")
            f.write("# Generated: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n\n")
            OmegaConf.save(cfg, f)

        # Save individual config files for easy reference
        import shutil
        config_src_dir = Path("configs")
        for config_file in config_src_dir.rglob("*.yaml"):
            rel_path = config_file.relative_to(config_src_dir)
            dest_file = config_backup_dir / rel_path
            dest_file.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(config_file, dest_file)

        # Save a human-readable summary
        summary_path = config_backup_dir / "run_summary.txt"
        with open(summary_path, 'w') as f:
            f.write(f"{'='*70}\n")
            f.write(f"RUN SUMMARY: {run_name}\n")
            f.write(f"{'='*70}\n\n")
            f.write(f"Date/Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"Model: {cfg.model._target_}\n")
            f.write(f"Dataset: {cfg.dataset._target_}\n")
            f.write(f"Mode: {cfg.mode}\n\n")

            f.write(f"TRAINING SETTINGS:\n")
            f.write(f"  Epochs: {cfg.epochs}\n")
            f.write(f"  Batch Size: {cfg.dataset.cfg.batch_size}\n")
            f.write(f"  Learning Rate: {cfg.model.cfg.optimizer.lr}\n")
            f.write(f"  Precision: {cfg.precision}\n")
            f.write(f"  Gradient Accumulation: {cfg.gradient_accumulation}\n")
            f.write(f"  Gradient Clipping: {cfg.gradient_clip_val}\n")
            f.write(f"  Devices: {cfg.devices}\n")
            f.write(f"  Seed: {cfg.seed}\n\n")

            f.write(f"DATA SETTINGS:\n")
            f.write(f"  Root: {cfg.dataset.cfg.root}\n")
            f.write(f"  Num Workers: {cfg.dataset.cfg.num_workers}\n")
            f.write(f"  Pin Memory: {cfg.dataset.cfg.pin_memory}\n")
            if 'prefetch_factor' in cfg.dataset.cfg:
                f.write(f"  Prefetch Factor: {cfg.dataset.cfg.prefetch_factor}\n")
            f.write(f"\n")

            f.write(f"MODEL SETTINGS:\n")
            f.write(f"  Loss Function: {cfg.model.cfg.loss_fn._target_}\n")
            f.write(f"  Optimizer: {cfg.model.cfg.optimizer._target_}\n")
            if 'scheduler' in cfg.model.cfg and cfg.model.cfg.scheduler is not None:
                f.write(f"  Scheduler: {cfg.model.cfg.scheduler._target_}\n")
            else:
                f.write(f"  Scheduler: None\n")
            f.write(f"\n")

            f.write(f"PATHS:\n")
            f.write(f"  Checkpoints: {checkpoint_dir}\n")
            f.write(f"  Logs: {log_dir}\n")
            f.write(f"  Configs: {config_backup_dir}\n")
            f.write(f"{'='*70}\n")

        # Print organized directory structure
        print(f"\n{'='*70}")
        print(f"📁 Organized Logging Structure")
        print(f"{'='*70}")
        print(f"  Run Name:     {run_name}")
        print(f"  Checkpoints:  {checkpoint_dir}")
        print(f"  Logs:         {log_dir}")
        print(f"  Configs:      {config_backup_dir}")
        print(f"{'='*70}\n")

        # Build Trainer kwargs conditionally
        trainer_kwargs = {
            "accelerator": cfg.accelerator,
            "devices": cfg.devices,
            "max_epochs": cfg.epochs,
            "precision": cfg.precision,
            "logger": logger,
            "default_root_dir": str(log_dir),
            "enable_checkpointing": True,
            "enable_progress_bar": True,
            "callbacks": [checkpoint_callback],
            "accumulate_grad_batches": cfg.gradient_accumulation,
        }

        # Only add gradient_clip_val if it's not None
        if cfg.gradient_clip_val is not None:
            trainer_kwargs["gradient_clip_val"] = cfg.gradient_clip_val

        trainer = Trainer(**trainer_kwargs)
        if cfg.mode == "train":
            trainer.fit(model, datamodule=data_module, ckpt_path=cfg.ckpt_path)
            trainer.test(model, datamodule=data_module, ckpt_path=cfg.ckpt_path)
        elif cfg.mode == "test":
            trainer.test(model, datamodule=data_module, ckpt_path=cfg.ckpt_path)

        # Unwatch model
        if cfg.grad_track:
            if hasattr(primary_logger, 'experiment') and hasattr(primary_logger.experiment, 'unwatch'):
                primary_logger.experiment.unwatch(model)
            

    elif cfg.mode == "infer":
        #### Note that you need to change the model class when using different models here (for example I use SiamDiff here), this part of the code generates predictions and saves them to manually inspect them.
        from models.tinycd_enhanced import TinyCDEnhanced
        import matplotlib.pyplot as plt

        ckpt_path = cfg.ckpt_path

        print(f"Loading model from checkpoint: {ckpt_path}")

        model = TinyCDEnhanced.load_from_checkpoint(
            ckpt_path,
            cfg=cfg.model.cfg,
            map_location="cuda" if torch.cuda.is_available() else "cpu",
        )
        model.eval()

        data_module = instantiate(cfg.dataset)
        data_module.setup("test")
        test_loader = data_module.test_dataloader()

        # Organized inference output directory (extract run name from checkpoint if available)
        if cfg.ckpt_path and 'checkpoints' in str(cfg.ckpt_path):
            # Extract run name from checkpoint path
            # e.g., logs/checkpoints/TinyCDEnhanced_2025_12_26_00_15_06/epoch=199.ckpt
            ckpt_path = Path(cfg.ckpt_path)
            infer_run_name = ckpt_path.parent.name  # Get the run name directory
        else:
            # No checkpoint or unrecognized path, generate new name
            model_name = cfg.model._target_.split('.')[-1]
            infer_run_name = generate_run_name(model_name)

        save_dir = Path("logs/inference") / infer_run_name
        save_dir.mkdir(parents=True, exist_ok=True)

        print(f"\n{'='*70}")
        print(f"📁 Inference Output Directory")
        print(f"{'='*70}")
        print(f"  Run Name:   {infer_run_name}")
        print(f"  Checkpoint: {cfg.ckpt_path}")
        print(f"  Output:     {save_dir}")
        print(f"{'='*70}\n")

        for i, batch in enumerate(test_loader):
            pre = batch["pre"].to(model.device)
            post = batch["post"].to(model.device)
            mask = batch["mask"].cpu().numpy()  

            with torch.no_grad():
                preds = model((pre,post)) 
                pred_mask = torch.argmax(torch.softmax(preds, dim=1), dim=1).cpu().numpy()  

            for j in range(pre.shape[0]):  

                current_mask = mask[j].squeeze()  
                current_pred_mask = pred_mask[j] 

                fig, axes = plt.subplots(1, 2, figsize=(10, 5))
                axes[0].imshow(current_mask, cmap="gray")
                axes[0].set_title("Ground Truth")
                axes[0].axis("off")
                axes[1].imshow(current_pred_mask, cmap="gray")
                axes[1].set_title("Prediction")
                axes[1].axis("off")

                fig.savefig(save_dir / f"pred_{i:03d}_{j:03d}.png")
                plt.close(fig)

        print(f"Inference complete. Results saved in {save_dir}")
if __name__ == "__main__":
    main()