experiments/data_scaling_study/models.py

"""
Thin re-exports of the U-Net and SegFormer-B0 model classes used in the
existing pv_panel_models/ baselines, so the data-scaling study trains
exactly the same architectures.

We load the source files via importlib without putting their parent
directories on sys.path — otherwise their per-model `dataset.py` (with a
different class name) would shadow this experiment's `dataset.py`.
"""
import importlib.util
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]
PV_DIR = REPO_ROOT / "pv_panel_models"


def _load(module_name: str, file_path: Path):
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    if spec is None or spec.loader is None:
        raise ImportError(f"could not load {file_path}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


_unet_mod = _load("_pv_unet_model", PV_DIR / "unet_model" / "unet_model.py")
_segformer_mod = _load("_pv_segformer_model", PV_DIR / "vit_model" / "segformer_model.py")

UNet = _unet_mod.UNet
UNetLoss = _unet_mod.CombinedLoss
SegformerModel = _segformer_mod.SegformerModel
SegformerLoss = _segformer_mod.CombinedLoss


def build_unet():
    return UNet(in_channels=3, out_channels=1), UNetLoss(bce_weight=0.5)


def build_segformer_b0():
    return SegformerModel(pretrained_name="nvidia/mit-b0", num_classes=1), SegformerLoss(bce_weight=0.5)


MODEL_REGISTRY = {
    "unet": build_unet,
    "segformer_b0": build_segformer_b0,
}