runner.py

import copy
import json
import logging
import math
import os
import random
import time
from collections import defaultdict
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

import albumentations as A
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import segmentation_models_pytorch as smp
import torch
import torch.nn as nn
import torch.optim as optim
import yaml
from IPython import get_ipython
from timm.utils import ModelEmaV3
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader
from torchvision.transforms import v2

from components import (
    BackboneParamGroupsBuilder,
    BaseTransform,
    CombinedLoss,
    CombinedTransforms,
    CSIROBioMassDataset,
    CSVLogger,
    GrassCloverSegmentationDataset,
    InverseBiomassEvaluator,
    LogZScoreTransform,
    ModelFreezeSchedule,
    SegmentationEvaluator,
    SwitchTrainPipelineScehdule,
    WeightedSmoothL1Loss,
    ZScoreTransform,
    build_transforms_albu,
    build_transforms_torchvision,
    evaluator,
    hooks,
    modules,
    set_seed,
    setup_logger,
)
from components.utils import visualize_backbone_attention

# import tqdm dynamically
try:
    from IPython import get_ipython

    shell = get_ipython().__class__.__name__
    if shell == "ZMQInteractiveShell":
        from tqdm.notebook import tqdm
    else:
        from tqdm import tqdm

except Exception as error:
    from tqdm import tqdm


class BaseRunner:
    @staticmethod
    def get_default_transforms(train=True):
        if train:
            return [
                {"type": "Resize", "size": [1024, 2048]},
                {"type": "RandomHorizontalFlip", "p": 0.5},
                {"type": "RandomVerticalFlip", "p": 0.5},
                {"type": "RandomAffine", "degrees": 20, "scale": [0.9, 1.1]},
                {"type": "ColorJitter", "brightness": 0.2, "contrast": 0.2, "saturation": 0.2, "hue": 0.0},
                {"type": "RandomAdjustSharpness", "sharpness_factor": 2},
                {"type": "GaussianBlur", "kernel_size": [3, 3], "sigma": [0.1, 1.0]},
                {"type": "ToImage"},
                {"type": "ToDtype", "dtype": torch.float32, "scale": True},
                {"type": "Normalize", "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
            ]
        else:
            return [
                {"type": "Resize", "size": [1024, 2048]},
                {"type": "ToImage"},
                {"type": "ToDtype", "dtype": torch.float32, "scale": True},
                {"type": "Normalize", "mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]},
            ]

    def build_transforms(self, transforms_config: Dict):
        config_ = copy.deepcopy(transforms_config)
        transform_type = config_.pop("type").lower()
        if transform_type == "albu":
            return build_transforms_albu(config_)
        elif transform_type == "torchvision":
            return build_transforms_torchvision(config_)
        else:
            raise ValueError(f"Unknown aug framework {transform_type}")

    @staticmethod
    def build_config(cfg: str | Path | Dict[str, Any]):
        if isinstance(cfg, (str, Path)):
            path = Path(cfg)
            with path.open("r", encoding="utf-8") as file:
                if path.suffix in {".yml", ".yaml"}:
                    return yaml.safe_load(file)
                elif path.suffix == ".json":
                    return json.load(file)
                else:
                    raise ValueError(f"Unsupported config file type: {path.suffix}")

        elif isinstance(cfg, (list, dict)):
            return cfg

        else:
            raise ValueError("Unsupported config data type")

    def _build_model(self, model_kwargs: Optional[Dict] = None) -> nn.Module:
        if model_kwargs is None:
            model_kwargs = self.cfg["model_kwargs"]
        model_kwargs_cpy = copy.deepcopy(model_kwargs)
        model_type = model_kwargs_cpy.pop("type", "BioMassCNNRegressor")
        constructor = getattr(modules, model_type)
        return constructor(**model_kwargs_cpy)

    def _build_data_preprocess(self) -> BaseTransform:
        data_info = copy.deepcopy(self.cfg["data_info"])
        proc_type = data_info.pop("proc_type", "z-score")
        if proc_type == "z-score":
            self.target_names = data_info["targets"]
            return ZScoreTransform(**data_info)
        elif proc_type == "log-z-score":
            self.target_names = data_info["targets"]
            return LogZScoreTransform(**data_info)
        elif proc_type == "combined":
            self.target_names = data_info["biomass_stats"]["targets"]
            self.aux_target_names = data_info["aux_stats"]["targets"]
            return CombinedTransforms(**data_info)
        else:
            raise ValueError("Unknown data processor type")

    def _build_dataloader(self, train: bool, transforms: Optional[Any] = None) -> DataLoader:
        dataset_kwargs = copy.deepcopy(self.cfg["dataset_kwargs"])
        dataset_kwargs = dataset_kwargs["train_kwargs"] if train == True else dataset_kwargs["val_kwargs"]

        # build dataset
        if transforms is None:
            transforms = dataset_kwargs.pop("transforms", "default")
            if transforms == "default":
                transforms = self.get_default_transforms(train)
            dataset_kwargs["transforms"] = self.build_transforms(transforms)
        else:
            dataset_kwargs["transforms"] = transforms
        dataset = CSIROBioMassDataset(**dataset_kwargs)

        # build dataloader
        dataloader_kwargs = self.cfg["dataloader_kwargs"]
        if self.device == "cpu":
            dataloader_kwargs["pin_memory"] = False

        if train:

            def seed_worker(worker_id: int) -> None:
                worker_seed = torch.initial_seed() % 2**32
                np.random.seed(worker_seed)
                random.seed(worker_seed)

            g = torch.Generator()
            g.manual_seed(self.cfg["seed"])
            return DataLoader(dataset, shuffle=True, generator=g, worker_init_fn=seed_worker, **dataloader_kwargs)
        else:
            return DataLoader(dataset, shuffle=False, **dataloader_kwargs)


class BioMassRunner(BaseRunner):
    def __init__(
        self,
        work_dir: str | Path,
        cfg: str | Path | Dict[str, Any],
        track_loss_metrics: List[str] = ["loss", "biomass_loss"],
        track_acc_metrics: List[str] = ["r2"],
        save_every_epoch: bool = False,
        logger_suffix: str = "",
        device: Optional[str] = None,
    ) -> None:
        # setup work directories and paths
        self.work_dir = Path(work_dir)
        self.data_sample_path = self.work_dir / "data_samples"
        self.checkpoint_dir = self.work_dir / "checkpoints"
        self.log_path = self.work_dir / "logs.log"
        self.csv_path = self.work_dir / "logs_metrics.csv"
        os.makedirs(self.checkpoint_dir, exist_ok=True)
        os.makedirs(self.data_sample_path, exist_ok=True)

        # Logging Setup
        self.cfg = self.build_config(cfg)
        self.logger = self._build_logger(suffix=logger_suffix)
        self.csv_writer = CSVLogger(csv_path=self.csv_path)
        self.visualize = True
        self.use_ema = False
        self.save_every_epoch = save_every_epoch

        # select device
        self.device = device
        if self.device is None:
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        if self.device == "cuda" and self.cfg["use_amp"] == True:
            self.scaler = GradScaler()
            self.logger.info("Using Automatic Mixed Precision (AMP) training.")

        # model state tracker
        self.track_loss_metrics = track_loss_metrics
        self.track_acc_metrics = track_acc_metrics
        self.loss_metrics_records = {metric_name: float("inf") for metric_name in track_loss_metrics}
        self.acc_metrics_records = {metric_name: -float("inf") for metric_name in track_acc_metrics}
        self.least_loss_checkpoints = {
            metric_name: self.checkpoint_dir / f"least_val_{metric_name}.pth" for metric_name in track_loss_metrics
        }
        self.best_acc_checkpoints = {
            metric_name: self.checkpoint_dir / f"best_val_{metric_name}.pth" for metric_name in track_acc_metrics
        }

    def _build_logger(self, suffix: str = ""):
        return setup_logger(name=f"BioMassRunner{suffix}", log_path=self.log_path)

    def _build_components(self):
        self.model = self._build_model().to(self.device)
        self.data_proc = self._build_data_preprocess()
        self.train_dataloader = self._build_dataloader(train=True)
        self.val_dataloader = self._build_dataloader(train=False)
        self.optimizer = self._build_optimizer()
        self.schedular = self._build_lr_schedular()
        self.loss_fn = self._build_loss()
        self.post_epoch_hooks = self._build_pipeline_schedular()
        self.post_epoch_hooks += self._build_hooks()

        # init ema settings
        ema_cfg = copy.deepcopy(self.cfg.get("ema"))
        if ema_cfg:
            start_epoch = ema_cfg.pop("start_after")
            start_iter = start_epoch * len(self.train_dataloader)
            if self.model.backbone_eval_mode == True and hasattr(self.model, "heads"):
                self.ema_model = ModelEmaV3(
                    self.model.heads, update_after_step=start_iter, device=self.device, **ema_cfg
                )
            else:
                self.ema_model = ModelEmaV3(self.model, update_after_step=start_iter, device=self.device, **ema_cfg)

            self.use_ema = True

    def _build_optimizer(self) -> optim.Optimizer:
        optim_kwargs = copy.deepcopy(self.cfg["optim_kwargs"])
        backbone_param_kwargs = optim_kwargs["backbone_param_kwargs"]
        base_lr = optim_kwargs["base_lr"]

        # build parameters group for backbone
        parameter_groups = []
        param_type = backbone_param_kwargs.pop("type")
        param_fn = getattr(BackboneParamGroupsBuilder, f"get_{param_type}_param_groups")
        parameter_groups += param_fn(backbone=self.model.backbone, base_lr=base_lr, **backbone_param_kwargs)

        # build param group for attention_blocks
        attn_parameters = []
        if hasattr(self.model, "attention_blocks"):
            attn_parameters += list(self.model.attention_blocks.parameters())
        if hasattr(self.model, "entry_attention"):
            attn_parameters += list(self.model.entry_attention.parameters())
        if hasattr(self.model, "backbone_attention"):
            attn_parameters += list(self.model.backbone_attention.parameters())
        if len(attn_parameters) >= 1:
            attn_lr = base_lr * optim_kwargs["attn_mmltiplier"]
            parameter_groups.append({"params": attn_parameters, "lr": attn_lr, "name": "attention_layers"})

        # create param groups for rest of the componenets
        rest_components = []
        for name, module in self.model.named_children():
            if name not in ("backbone", "segm_model", "attention_blocks", "entry_attention", "backbone_attention"):
                rest_components.extend(list(module.parameters()))

        parameter_groups.append({"params": rest_components, "lr": base_lr, "name": "rest_layers"})
        return optim.AdamW(parameter_groups, weight_decay=optim_kwargs["weight_decay"])

    def _build_loss(self) -> nn.Module:
        loss_kwargs = copy.deepcopy(self.cfg["loss_kwargs"])
        loss_type = loss_kwargs.pop("type", "smoothL1")
        if loss_type == "smoothL1":
            loss = torch.nn.SmoothL1Loss(**loss_kwargs)
        elif loss_type == "weighted_smoothL1":
            loss = WeightedSmoothL1Loss(**loss_kwargs)
        elif loss_type == "combined":
            return CombinedLoss(**loss_kwargs)
        else:
            raise ValueError("Unknown loss encountered")
        return loss

    def _build_freeze_scehdular(self) -> ModelFreezeSchedule:
        freeze_kwargs = self.cfg["freeze_schedule_kwargs"]
        return ModelFreezeSchedule(self.logger, self.model, **freeze_kwargs)

    def _build_hooks(self) -> List[Callable]:
        hooks_cfgs = copy.deepcopy(self.cfg.get("hooks", []))

        hooks_funcs = []
        for hook_cfg in hooks_cfgs:
            hook_type = hook_cfg.pop("type")
            constructor = getattr(hooks, hook_type)
            hooks_funcs.append(constructor(**hook_cfg))

        return hooks_funcs

    def _build_pipeline_schedular(self) -> SwitchTrainPipelineScehdule:
        swith_pipeline_hooks = []
        pipeline_cfg = copy.deepcopy(self.cfg.get("pipeline_schedules", {}))
        for num_epoch, pipeline_cfg in pipeline_cfg.items():
            swith_pipeline_hooks.append(SwitchTrainPipelineScehdule(switch_at=num_epoch, pipeline_cfg=pipeline_cfg))
        return swith_pipeline_hooks

    def _build_lr_schedular(self) -> optim.lr_scheduler.LambdaLR:
        lr_schedule_kwargs = self.cfg["lr_schedule_kwargs"]
        warmup_epochs = lr_schedule_kwargs["warmup_epochs"] + 1
        anneal_after = lr_schedule_kwargs["anneal_after"]
        total_epochs = lr_schedule_kwargs["total_epochs"]
        mir_lr_factor = lr_schedule_kwargs["mir_lr_factor"]
        T_max = total_epochs - anneal_after

        def lr_lambda_func(current_epoch: int) -> float:
            # linear warmup
            current_epoch += 1
            if current_epoch < warmup_epochs:
                return float(current_epoch) / float(max(1, warmup_epochs))

            # constant phase
            if current_epoch < anneal_after:
                return 1.0

            # start cosine annealing
            progress = min(1.0, (current_epoch - anneal_after) / T_max)
            cosine_out = 0.5 * (1.0 + math.cos(math.pi * progress))
            return mir_lr_factor + (1.0 - mir_lr_factor) * cosine_out

        return optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lr_lambda_func)

    def _build_evaluator(self) -> InverseBiomassEvaluator:
        evaluator_kwargs = copy.deepcopy(
            self.cfg["evaluator_kwargs"],
        )
        evaluator_type = evaluator_kwargs.pop("type", "InverseBiomassEvaluator")
        constructor = getattr(evaluator, evaluator_type)
        return constructor(**evaluator_kwargs)

    def visualize_batch(
        self,
        epoch: int,
        batch_idx: int,
        images: torch.Tensor,
        targets: torch.Tensor,
        attn_feats: Dict[str, torch.Tensor],
        show: bool = True,
    ):
        # fetch mean and stds for image
        image_mean, image_std = None, None
        data_kwargs = self.cfg["dataset_kwargs"]["train_kwargs"]
        for key, value in data_kwargs["transforms"].items():
            if "transform" not in key:
                continue
            for transform in value:
                if transform["type"].lower() == "normalize":
                    image_mean = transform["mean"]
                    image_std = transform["std"]

        if image_mean is not None or image_std is not None:
            mean = torch.tensor(image_mean).view(1, 3, 1, 1)
            std = torch.tensor(image_std).view(1, 3, 1, 1)

        # denormalize attention map
        img_batches = copy.deepcopy(images).cpu()
        if len(img_batches.shape) == 5:
            img_batches = self.train_dataloader.dataset.transforms.inverse_tiles(img_batches, is_batch=True)
        if image_mean is not None or image_std is not None:
            img_batches = img_batches * std + mean
            img_batches = np.clip(img_batches, 0, 1)

        # create figure
        fig, axes = plt.subplots(nrows=images.shape[0], ncols=1, figsize=(12, 20))
        for i in range(len(images)):
            img = img_batches[i].permute(1, 2, 0).numpy()
            axes[i].imshow(img)

            # Create label string
            label_str = " | ".join([f"{name}: {val:.2f}" for name, val in zip(self.target_names, targets[i])])
            axes[i].set_title(f"Sample {i} - {label_str}", fontsize=10)
            axes[i].axis("off")

        plt.tight_layout()
        fig_path = self.data_sample_path / f"train_epoch{epoch}_batch_{batch_idx}.jpeg"
        fig.savefig(fig_path)
        if show:
            plt.show()
        else:
            plt.close(fig)

        # handle the attention map
        if attn_feats is not None:
            attn_fig_path = self.data_sample_path / f"train_epoch{epoch}_batch_{batch_idx}_attn.jpeg"
            visualize_backbone_attention(
                conv_feats=attn_feats["pre_attn"],
                attention_feats=attn_feats["post_attn"],
                image_tensor=img_batches,
                out_dir=attn_fig_path,
                alpha=0.5,
                show=show,
            )

    def plot_training_logs(self, show: bool = True):
        df = pd.read_csv(self.csv_path)
        epochs = df["epoch"]

        # 1. Identify Keys Dynamically
        # Loss keys: find unique base names (e.g., 'loss' from 'train/loss' and 'val/loss')
        all_cols = df.columns.tolist()
        loss_base_keys = sorted(list(set(c.replace("train/", "").replace("val/", "") for c in all_cols if "loss" in c)))

        lr_cols = [c for c in all_cols if c.startswith("train/") and c.endswith("lr")]
        val_metric_cols = [c for c in all_cols if c.startswith("val/") and "loss" not in c]

        plt.style.use("seaborn-v0_8-muted")

        # --- FIGURE 1: LOSSES AND LEARNING RATES ---
        num_loss_plots = len(loss_base_keys) + (1 if lr_cols else 0)
        cols_per_row = 3
        rows = math.ceil(num_loss_plots / cols_per_row)

        fig_loss, axes_loss = plt.subplots(rows, cols_per_row, figsize=(18, 5 * rows))
        axes_loss = axes_loss.flatten() if num_loss_plots > 1 else [axes_loss]

        # Plot Losses
        for i, base_key in enumerate(loss_base_keys):
            ax = axes_loss[i]
            train_key = f"train/{base_key}"
            val_key = f"val/{base_key}"

            if train_key in df.columns:
                ax.plot(epochs, df[train_key], label="Train", marker="o", markersize=3)
            if val_key in df.columns:
                ax.plot(epochs, df[val_key], label="Val", linestyle="--", marker="s", markersize=3)

            ax.set_title(f"Loss: {base_key.replace('_', ' ').title()}", fontweight="bold")
            ax.set_xlabel("Epoch")
            ax.set_ylabel("Loss Value")
            ax.legend()
            ax.grid(True, alpha=0.3)

        # Plot LRs in the last loss slot
        if lr_cols:
            ax_lr = axes_loss[num_loss_plots - 1]
            for lr_col in lr_cols:
                label = lr_col.replace("train/", "").replace("_", " ").title()
                ax_lr.plot(epochs, df[lr_col], label=label, linewidth=2)
            ax_lr.set_title("Learning Rates", fontweight="bold")
            ax_lr.set_xlabel("Epoch")
            ax_lr.set_ylabel("LR Value")
            ax_lr.legend()
            ax_lr.grid(True, alpha=0.3)

        # Hide unused subplots
        for j in range(num_loss_plots, len(axes_loss)):
            axes_loss[j].axis("off")

        fig_loss.suptitle("Training Progress: Losses & Optimization", fontsize=16, fontweight="bold")
        fig_loss.tight_layout(rect=[0, 0.03, 1, 0.95])
        fig_loss.savefig(self.work_dir / "loss_lr_dashboard.png", dpi=300)

        # --- FIGURE 2: VALIDATION METRICS ---
        if val_metric_cols:
            num_metrics = len(val_metric_cols)
            rows_m = math.ceil(num_metrics / cols_per_row)
            fig_met, axes_met = plt.subplots(rows_m, cols_per_row, figsize=(18, 5 * rows_m))
            axes_met = axes_met.flatten() if num_metrics > 1 else [axes_met]

            for i, col in enumerate(val_metric_cols):
                ax = axes_met[i]
                label = col.replace("val/", "").replace("_", " ").upper()
                ax.plot(epochs, df[col], label=label, color="green", marker="x")
                ax.set_title(label, fontweight="bold")
                ax.set_xlabel("Epoch")
                ax.set_ylabel("Value")
                ax.grid(True, alpha=0.3)

            # Hide unused metric subplots
            for j in range(num_metrics, len(axes_met)):
                axes_met[j].axis("off")

            fig_met.suptitle("Validation Metrics Performance", fontsize=16, fontweight="bold")
            fig_met.tight_layout(rect=[0, 0.03, 1, 0.95])
            fig_met.savefig(self.work_dir / "metrics_dashboard.png", dpi=300)

        # Handle Show/Close
        if show:
            plt.show()
        else:
            plt.close("all")

    def train_one_epoch(self, epoch: int, epochs: int) -> float:
        self.model.train()
        loss_records = defaultdict(lambda: 0)
        data_time = 0.0
        tqdm_desc = f"Epoch [{epoch}/{epochs}] Train "
        with tqdm(total=len(self.train_dataloader), desc=tqdm_desc, leave=False, unit="batch") as pbar:
            previous = time.time()
            for batch_idx, batch_data in enumerate(self.train_dataloader):
                data_time = time.time() - previous

                # get the data and move to the accelerator data
                images = batch_data["img"].to(self.device)
                targets = batch_data["targets"]
                if isinstance(targets, dict):
                    for key in targets.keys():
                        targets[key] = targets[key].to(self.device)
                else:
                    targets = targets.to(self.device)
                targets_norm = self.data_proc.normalize_data(targets)

                # visualize initial batches for epochs
                if self.visualize:
                    attn_feats = None
                    if hasattr(self.model, "get_attention_feats"):
                        attn_feats = self.model.get_attention_feats(images)
                        for key in attn_feats.keys():
                            attn_feats[key] = attn_feats[key][:4, ...]

                    self.visualize_batch(
                        epoch=epoch,
                        batch_idx=batch_idx,
                        images=images[:4, ...],
                        targets=(targets["biomass"][:4, ...] if isinstance(targets, dict) else targets[:4, ...]),
                        attn_feats=attn_feats,
                        show=False,
                    )
                    if batch_idx == 3:
                        self.visualize = False

                # forward pass and step the loss
                self.optimizer.zero_grad()
                if hasattr(self, "scaler"):
                    with autocast():
                        preds = self.model(images)
                        loss, loss_values = self.loss_fn(preds, targets_norm)
                    self.scaler.scale(loss).backward()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    preds = self.model(images)
                    loss, loss_values = self.loss_fn(preds, targets_norm)
                    loss.backward()
                    self.optimizer.step()

                # update ema
                if self.use_ema:
                    if self.model.backbone_eval_mode == True and hasattr(self.model, "heads"):
                        self.ema_model.update(
                            self.model.heads, step=(epoch - 1) * len(self.train_dataloader) + batch_idx
                        )
                    else:
                        self.ema_model.update(self.model, step=(epoch - 1) * len(self.train_dataloader) + batch_idx)

                # clip gradients
                if self.cfg["clip_grad"] is not None:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.cfg["clip_grad"])

                # post step operations
                blr = [param_grp["lr"] for param_grp in self.optimizer.param_groups if "backbone" in param_grp["name"]]
                blr = blr[0] if len(blr) >= 1 else 0.0
                for key, value in loss_values.items():
                    loss_records[key] += value

                # set the values on progress bar
                pbar.set_postfix(
                    {
                        "data_time": data_time,
                        **{k: f"{v:.4f}" for k, v in loss_values.items()},
                        "loss": f"{loss.item():.4f}",
                        "blr": f"{blr:.3e}",
                        "lr": f"{self.optimizer.param_groups[-1]['lr']:.3e}",
                    }
                )
                pbar.update()
                previous = time.time()

            # log the last state of progress bar
            self.logger.info(pbar.format_meter(**pbar.format_dict), extra={"file_only": True})

        # End of Training Epoch
        train_details = {}
        train_details["train/lr"] = self.optimizer.param_groups[-1]["lr"]
        train_details["train/blr"] = blr
        for key, value in loss_records.items():
            train_details[f"train/{key}"] = value / len(self.train_dataloader)

        # step the scehdular
        self.schedular.step()
        return train_details

    @torch.no_grad()
    def validate(self, dataloader: DataLoader, prefix: str = "Val") -> Dict[str, float]:
        self.model.eval()
        loss_records = defaultdict(lambda: 0)
        evaluator = self._build_evaluator()
        tqdm_desc = 16 * " " + f"{prefix} "
        with tqdm(total=len(dataloader), desc=tqdm_desc, leave=False, unit="batch") as pbar:
            previous = time.time()
            for batch_data in dataloader:
                data_time = time.time() - previous
                images = batch_data["img"].to(self.device)
                targets = batch_data["targets"]
                if isinstance(targets, dict):
                    for key in targets.keys():
                        targets[key] = targets[key].to(self.device)
                else:
                    targets = targets.to(self.device)
                targets_norm = self.data_proc.normalize_data(targets)

                # predict on the model
                if hasattr(self, "scaler"):
                    with autocast():
                        if self.use_ema:
                            if self.model.backbone_eval_mode and hasattr(self.model, "heads"):
                                back_outs = self.model.get_backbone_feats(images)
                                preds = self.model._forward_heads_staticmethod_for_ema(self.ema_model.module, back_outs)
                            else:
                                preds = self.ema_model(images)
                        else:
                            preds = self.model(images)

                        _, loss_values = self.loss_fn(preds, targets_norm)

                # add records
                for key, value in loss_values.items():
                    loss_records[key] += value

                # denorm the predicted data add to evaluator
                preds = self.data_proc.denormalize_data(preds)
                if isinstance(targets, dict):
                    targets = self.data_proc.denormalize_data(targets_norm)
                    preds["species"] = preds["species"].sigmoid()
                evaluator.add(y_true=targets, y_pred=preds)

                # update the progress bar
                pbar.set_postfix(
                    {
                        "data_time": data_time,
                        **{k: f"{v:.4f}" for k, v in loss_values.items()},
                    }
                )
                pbar.update()
                previous = time.time()

            # log the last state of progress bar
            self.logger.info(pbar.format_meter(**pbar.format_dict), extra={"file_only": True})

        # complete the evaluation
        val_details = {}
        for key, value in loss_records.items():
            val_details[f"val/{key}"] = value / len(dataloader)

        val_metrics = (
            evaluator.evaluate(targets_names=self.target_names, aux_names=self.aux_target_names)
            if hasattr(self, "aux_target_names")
            else evaluator.evaluate(targets_names=self.target_names)
        )
        report_str = val_metrics.pop("report_str", None)
        if report_str is not None:
            self.logger.info(f"Evaluation Report is as follos \n {report_str}")

        for name, value in val_metrics.items():
            val_details[f"{prefix.lower()}/{name}"] = value
        return val_details

    def run(self, show: bool = False) -> None:
        # dump the config in work directory
        cfg_path = self.work_dir / f"{self.work_dir.stem}_config.yaml"
        with cfg_path.open("w", encoding="utf-8") as file:
            yaml.safe_dump(self.cfg, file)

        # seed using the random seed
        set_seed(self.cfg["seed"])

        # build all the componsnet
        self._build_components()

        # freeze model heads
        if hasattr(self.model, "freeze_backbone"):
            self.freeze_schedule = self._build_freeze_scehdular()

        # start the traininig
        self.logger.info(f"Starting training on {self.device}...")
        for epoch in range(1, self.cfg["epochs"] + 1):
            train_details = self.train_one_epoch(epoch, epochs=self.cfg["epochs"])
            val_details = self.validate(self.val_dataloader)
            epoch_details = dict(epoch=epoch, **train_details, **val_details)

            # save the best, every epoch checkpoints
            states = dict(
                epoch=epoch,
                details=epoch_details,
                weights=(self.ema_model.module.state_dict() if self.use_ema else self.model.state_dict()),
            )
            if self.save_every_epoch:
                epoch_save_path = self.checkpoint_dir / f"Epoch{epoch}.pth"
                torch.save(states, epoch_save_path)
            for loss_metric in self.track_loss_metrics:
                if epoch_details[f"val/{loss_metric}"] <= self.loss_metrics_records[loss_metric]:
                    self.loss_metrics_records[loss_metric] = epoch_details[f"val/{loss_metric}"]
                    torch.save(states, self.least_loss_checkpoints[loss_metric])
                    self.logger.info(
                        f"Saved least loss Model (Val {loss_metric}: {self.loss_metrics_records[loss_metric]})"
                    )
            for acc_metric in self.track_acc_metrics:
                if epoch_details[f"val/{acc_metric}"] >= self.acc_metrics_records[acc_metric]:
                    self.acc_metrics_records[acc_metric] = epoch_details[f"val/{acc_metric}"]
                    torch.save(states, self.best_acc_checkpoints[acc_metric])
                    self.logger.info(f"Saved Best acc Model (Val {acc_metric}: {self.acc_metrics_records[acc_metric]})")
            if self.use_ema:
                self.logger.info(f"Saved EMA model weights at epoch {epoch} as stated above")

            # log the epoch details in logger and csv
            self.csv_writer.write_rows(epoch_details)
            message = f"Epoch [{epoch}/{self.cfg['epochs']}] Details - \n"
            if self.use_ema:
                self.logger.info(f"Metrics computed on EMA module weights . . . .")
            for key, value in epoch_details.items():
                message += f"{key}: {value} \n"
            self.logger.info(message)

            # apply freeze schedule
            if hasattr(self.model, "freeze_backbone"):
                self.freeze_schedule(current_epoch=epoch)

            # apply hooks
            for hook_obj in self.post_epoch_hooks:
                hook_obj(epoch, self)

        self.logger.info("Training Completed . . . .")
        self.plot_training_logs(show=show)
        return self.csv_path


class BioMassKFoldRunner(BaseRunner):
    def __init__(
        self, base_work_dir: str | Path, folds_cfg: str | Path | Dict[str, Any], base_cfg: str | Path | Dict[str, Any]
    ):
        self.base_work_dir = Path(base_work_dir)
        self.folds_cfg = BioMassRunner.build_config(folds_cfg)
        self.base_cfg = BioMassRunner.build_config(base_cfg)

        log_path = self.base_work_dir / "fold_logs.log"
        self.logger = setup_logger(name="BioMassKFoldRunner", log_path=log_path)

    def _override_config(self, base_dict: Dict, override_dict: Dict) -> Dict:
        for key, value in override_dict.items():
            if isinstance(value, dict) and key in base_dict and isinstance(base_dict[key], dict):
                self._override_config(base_dict[key], value)
            else:
                base_dict[key] = copy.deepcopy(value)
        return base_dict

    def summarize_results(self, csv_paths: List):
        self.logger.info("#" * 100)
        self.logger.info(" CROSS-VALIDATION SUMMARY . . . .")
        self.logger.info("#" * 100)

        # collect best checkpoint from each fold
        fold_summaries = []
        for fold_cfg, csv_path in zip(self.folds_cfg["folds"], csv_paths):
            df = pd.read_csv(csv_path)
            best_idx = df["val/r2"].idxmax()
            best_row = df.loc[best_idx].to_dict()
            best_row["fold"] = fold_cfg["name"]
            fold_summaries.append(best_row)
            self.logger.info(
                f"{fold_cfg['name']} | Best Epoch: {int(best_row['epoch'])} | "
                f"R2: {best_row['val/r2']:.4f} | MAE: {best_row['val/mae']:.4f}"
            )

        # calculate global statistics
        self.logger.info("-" * 30)
        stats_output = {}
        summary_df = pd.DataFrame(fold_summaries)
        for metric in ["val/loss", "val/mae", "val/rmse", "val/r2"]:
            mean_val = summary_df[metric].mean()
            std_val = summary_df[metric].std()
            stats_output[metric] = {"mean": mean_val, "std": std_val}
            self.logger.info(f"Global {metric}: {mean_val:.4f} ± {std_val:.4f}")

        # save summary to a JSON in the work_dir
        with open(self.base_work_dir / "cv_summary.json", "w") as f:
            json.dump(stats_output, f, indent=4)

        self.logger.info("=" * 30 + "\n")
        return summary_df

    def run(self):
        csv_paths = []
        for fold in self.folds_cfg["folds"]:
            # log the fold information
            self.logger.info(f"{'#'*100}")
            self.logger.info(f" Preparing: {fold['name']} ")
            self.logger.info(f"{'#'*100}\n")

            # Setup fold firectory and override fold config
            fold_work_dir = self.base_work_dir / fold["name"]
            fold_work_dir.mkdir(parents=True, exist_ok=True)
            current_cfg = copy.deepcopy(self.base_cfg)
            current_cfg = self._override_config(current_cfg, fold["override"])

            # select runner class
            if self.folds_cfg["runner"] == "BioMassRunner":
                runner_class = BioMassRunner

            # Initialize the fold runner and attach it's logger
            suffix = f".{fold['name']}"
            runner = runner_class(work_dir=fold_work_dir, cfg=current_cfg, logger_suffix=suffix)
            for handler in self.logger.handlers:
                if isinstance(handler, logging.FileHandler):
                    runner.logger.addHandler(handler)

            # run the training on the fold
            logs_csv = runner.run()
            csv_paths.append(logs_csv)

        # summarise metrics across the folds
        summary_df_path = self.base_work_dir / "summary.csv"
        summary_df = self.summarize_results(csv_paths)
        summary_df.to_csv(summary_df_path, index=False)


class GrassCloverSemanticRunner(BioMassRunner):
    def _build_components(self):
        self.model = self._build_model().to(self.device)
        self.train_dataloader = self._build_dataloader(train=True)
        self.val_dataloader = self._build_dataloader(train=False)
        self.optimizer = self._build_optimizer()
        self.schedular = self._build_lr_schedular()
        self.loss_fn = self._build_loss()

    def _build_logger(self, suffix: str = ""):
        return setup_logger(name=f"GrassCloverSemanticRuner{suffix}", log_path=self.log_path)

    def _build_dataloader(self, train: bool) -> DataLoader:
        dataset_kwargs = copy.deepcopy(self.cfg["dataset_kwargs"])
        dataset_kwargs = dataset_kwargs["train_kwargs"] if train == True else dataset_kwargs["val_kwargs"]

        # build dataset
        transforms = dataset_kwargs.pop("transforms", "default")
        if transforms == "default":
            transforms = self.get_default_transforms(train)
        transforms = self.build_transforms(transforms)
        dataset = GrassCloverSegmentationDataset(transforms=transforms, **dataset_kwargs)

        # build dataloader
        dataloader_kwargs = self.cfg["dataloader_kwargs"]
        if self.device == "cpu":
            dataloader_kwargs["pin_memory"] = False
        return DataLoader(dataset, shuffle=train, **dataloader_kwargs)

    def _build_model(self) -> nn.Module:
        model_kwargs = copy.deepcopy(self.cfg["model_kwargs"])
        model_type = model_kwargs.pop("type", "Unet")
        model_constructor = getattr(smp, model_type)
        return model_constructor(**model_kwargs)

    def _build_optimizer(self) -> optim.Optimizer:
        optim_kwargs = copy.deepcopy(self.cfg["optim_kwargs"])
        backbone_param_kwargs = optim_kwargs.pop("optim_param_kwargs", {})

        # collect the scale params
        base_lr = optim_kwargs.pop("lr")
        encoder_stem_scale = backbone_param_kwargs.pop("encoder_stem_scale", 0.0)
        encoder_scale = backbone_param_kwargs.pop("encoder_scale", 1.0)
        decoder_scale = backbone_param_kwargs.pop("decoder_scale", 1.0)
        head_scale = backbone_param_kwargs.pop("head_scale", 1.0)

        # create param groups
        stem_params = []
        stem_layers = ["stem_0", "stem_1"]
        for layer_name in stem_layers:
            layer = getattr(self.model.encoder.model, layer_name)
            stem_params.extend(list(layer.parameters()))

        rest_encoder_params = []
        for layer_name, layer in self.model.encoder.model.named_children():
            if layer_name in stem_layers:
                continue
            rest_encoder_params.extend(list(layer.parameters()))

        parameter_groups = [
            {"params": stem_params, "lr": base_lr * encoder_stem_scale, "name": "stem_params"},
            {"params": rest_encoder_params, "lr": base_lr * encoder_scale, "name": "encoder_params"},
            {"params": self.model.decoder.parameters(), "lr": base_lr * decoder_scale, "name": "decoder_params"},
            {"params": self.model.segmentation_head.parameters(), "lr": base_lr * head_scale, "name": "head_params"},
        ]

        # build the optimizer
        optim_type = optim_kwargs.pop("type", "AdamW")
        optim_constructor = getattr(optim, optim_type)
        return optim_constructor(parameter_groups, **optim_kwargs)

    def _build_loss(self) -> nn.Module:
        loss_kwargs = self.cfg["loss_kwargs"]
        loss_type = loss_kwargs.pop("type", "focal")
        if loss_type == "cross_entrophy":
            loss_fn = torch.nn.CrossEntropyLoss(**loss_kwargs)
        elif loss_type == "focal":
            loss_fn = smp.losses.FocalLoss(**loss_kwargs)
        elif loss_type == "dice":
            loss_fn = smp.losses.DiceLoss(**loss_kwargs)
        else:
            raise ValueError("Unknown loss encountered")
        return loss_fn

    def _build_evaluator(self) -> SegmentationEvaluator:
        evaluator_kwargs = self.cfg["evaluator_kwargs"]
        return SegmentationEvaluator(**evaluator_kwargs)

    def visualize_batch(
        self, batch_idx: int, images: torch.Tensor, masks: torch.Tensor, alpha: float = 0.4, show: bool = False
    ):
        if len(images.shape) == 3:
            images = images[None, ...]
            masks = masks[None, ...]

        # Create subplots
        dpi = 100
        num_images = images.shape[0]
        img_side_inches = 512 / dpi
        fig_width = (img_side_inches * 3) + 2.5
        fig_height = img_side_inches * num_images
        fig, axes = plt.subplots(
            num_images,
            3,
            dpi=dpi,
            figsize=(fig_width, fig_height),
        )
        axes = np.array(axes).reshape(num_images, 3)

        # extract the image means and stds
        image_means = None
        image_stds = None
        for transfrom_cfg in self.cfg["dataset_kwargs"]["train_kwargs"]["transforms"]:
            if transfrom_cfg["type"].lower().startswith("normalize"):
                image_means = torch.tensor(transfrom_cfg["mean"])[:, None, None]
                image_stds = torch.tensor(transfrom_cfg["std"])[:, None, None]

        for idx in range(num_images):
            image_vis = images[idx]
            mask = masks[idx]

            # Denormalize Image (Tensor -> Numpy RGB) if required
            if image_means is not None or image_stds is not None:
                image_vis = image_vis * image_stds + image_means
            image_vis = image_vis.permute(1, 2, 0).numpy()
            image_vis = np.clip(image_vis, 0, 1)

            # create overlay image
            cmap = plt.get_cmap("tab20")
            mask_colored = cmap(mask / 19.0)
            mask_rgb = mask_colored[..., :3]  # Drop Alpha
            overlay = (1 - alpha) * image_vis + alpha * mask_rgb

            # plot 1 Original image
            axes[idx][0].imshow(image_vis)
            axes[idx][0].set_title(f"Original Image (Index: {idx})")
            axes[idx][0].axis("off")

            # plot 2 Raw Segmentation Map
            im = axes[idx][1].imshow(mask, cmap="tab20", vmin=0, vmax=len(self.target_names) - 1)
            axes[idx][1].set_title("Segmentation Map")
            axes[idx][1].axis("off")

            # Add colorbar for Map
            class_ids = range(len(self.target_names))
            cbar = plt.colorbar(im, ax=axes[idx][1], ticks=class_ids, fraction=0.046, pad=0.04)
            cbar.ax.set_yticklabels([f"{i}: {self.target_names[i]}" for i in class_ids])
            cbar.set_label("Class ID")

            # plot 3 Overlay with Legend
            axes[idx][2].imshow(overlay)
            axes[idx][2].set_title("Overlay (Image + Mask)")
            axes[idx][2].axis("off")

            # Dynamic Legend (Only present classes)
            present_classes = np.unique(mask)
            patches = []
            for cls_id in present_classes:
                if cls_id in self.target_names:  # Safety check
                    name = self.target_names[cls_id]
                    color = cmap(cls_id / 19.0)
                    patch = mpatches.Patch(color=color, label=f"{cls_id}: {name}")
                    patches.append(patch)

            axes[idx][2].legend(handles=patches, bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0.0)

        plt.tight_layout()
        fig_path = self.work_dir / f"train_batch_{batch_idx}.jpeg"
        fig.savefig(fig_path)
        if show:
            plt.show()
        else:
            plt.close(fig)

    def train_one_epoch(self, epoch: int, epochs: int) -> float:
        self.model.train()
        total_loss = 0.0
        data_time = 0.0
        tqdm_desc = f"Epoch [{epoch}/{epochs}] Train "
        with tqdm(total=len(self.train_dataloader), desc=tqdm_desc, leave=False, unit="batch") as pbar:
            previous = time.time()
            for batch_idx, batch_data in enumerate(self.train_dataloader):
                data_time = time.time() - previous
                images = batch_data["image"].to(self.device)
                segm_maps = batch_data["segm_map"].to(self.device)

                # visualize initial batches
                if epoch == 1 and batch_idx in (0, 1, 2):
                    self.visualize_batch(
                        batch_idx=batch_idx, images=images[:4, ...].cpu(), masks=segm_maps[:4, ...].cpu(), show=False
                    )

                # forward pass and step the loss
                self.optimizer.zero_grad()
                if hasattr(self, "scaler"):
                    with autocast():
                        preds = self.model(images)
                        loss = self.loss_fn(preds, segm_maps)
                    self.scaler.scale(loss).backward()
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                else:
                    preds = self.model(images)
                    loss = self.loss_fn(preds, segm_maps)
                    loss.backward()
                    self.optimizer.step()

                # post step operations
                total_loss += loss.item()
                pbar.set_postfix(
                    {
                        "data_time": data_time,
                        "loss": f"{loss.item():.4f}",
                        "lr": f"{self.optimizer.param_groups[-1]['lr']:.3e}",
                    }
                )
                pbar.update()
                previous = time.time()

            # log the last state of progress bar
            self.logger.info(pbar.format_meter(**pbar.format_dict), extra={"file_only": True})

        # End of Training Epoch
        train_details = {}
        train_details["train/lr"] = self.optimizer.param_groups[-1]["lr"]
        train_details["train/loss"] = total_loss / len(self.train_dataloader)

        # step the scehdular
        self.schedular.step()
        return train_details

    @torch.no_grad()
    def validate(self, dataloader: DataLoader, prefix: str = "Val") -> Dict[str, float]:
        self.model.eval()
        total_loss = 0.0
        evaluator = self._build_evaluator()
        tqdm_desc = 16 * " " + f"{prefix} "
        with tqdm(total=len(dataloader), desc=tqdm_desc, leave=False, unit="batch") as pbar:
            previous = time.time()
            for batch_data in dataloader:
                data_time = time.time() - previous
                images = batch_data["image"].to(self.device)
                segm_maps = batch_data["segm_map"].to(self.device)

                # predict on the model
                preds = self.model(images)
                loss = self.loss_fn(preds, segm_maps)
                total_loss += loss.item()

                # denorm the predicted data add to evaluator
                pred_segm_map = preds.softmax(dim=1)
                evaluator.add(y_true=segm_maps, y_pred=pred_segm_map)
                pbar.set_postfix({"data_time": data_time, "loss": f"{loss.item():.4f}"})
                pbar.update()
                previous = time.time()

            # log the last state of progress bar
            self.logger.info(pbar.format_meter(**pbar.format_dict), extra={"file_only": True})

        # complete the evaluation
        val_details = {}
        val_details[f"{prefix.lower()}/loss"] = total_loss / len(self.val_dataloader)
        val_metrics = evaluator.evaluate()
        for name, value in val_metrics.items():
            val_details[f"{prefix.lower()}/{name}"] = value
        return val_details