build.py

from __future__ import annotations
from pathlib import Path
from functools import partial
import numpy as np
import torch
from torch.utils.data import DataLoader
from vit_pytorch import SimpleViT

from src.data import DatasetConfig, BkDatasetPatchTo1D
from src.networks.cross_modal import CustomModel1

# external utils preserved from your project
from utils import (
    get_balanced_batch_sampler, get_criterion,
    IsoMaxPlusLossFirstPart, CoTeachingLoss,
)

# optional nets still referenced in your script
from src.networks.audioclip import AudioCLIP
from src.networks.inception_1d import InceptionModel


def get_dataset(opt, set_name: str, file_resolver, specified_min_inv=None, specified_metadata=None, extra_metadata=None):
    # determine min_inv for this split
    min_inv = specified_min_inv if specified_min_inv is not None else (opt.min_inv if set_name == "train" else opt.min_inv_eval)

    cfg = DatasetConfig(
        metadata_csv=opt.metadata_csv,   # <-- ARG, no hardcoded path
        fold_idx=opt.fold_idx,
        set_name=set_name.split("_")[0],
        min_involvement=min_inv,
        min_gleason=opt.min_gs,
        drop_ids=getattr(opt, "drop_ids", None),
    )

    # For now we instantiate only the used pipeline: 1D+2D
    ds = BkDatasetPatchTo1D(
        config=cfg,
        file_resolver=file_resolver,
        seq_len=opt.seq_len,
        aug_1d=opt.aug_1d,
        patch_size=opt.patch_dim[0],
        input_size=tuple(opt.input_size) if isinstance(opt.input_size, (list, tuple)) else (opt.input_size, opt.input_size),
    )
    return ds


def set_loaders(opt, file_resolver, verbose: bool = True):
    datasets, loaders = {}, {}
    for set_name in ["train", "val", "test"]:
        datasets[set_name] = get_dataset(opt, set_name, file_resolver=file_resolver)
        if verbose:
            x, y = datasets[set_name][0]
            print(set_name, (x[0] if isinstance(x, (tuple, list)) else x).shape, y)

        if set_name == "train":
            labels = np.array(datasets[set_name].metadata.TrueLabel).astype("uint8")
            sampler = get_balanced_batch_sampler(labels)
            dl = DataLoader(datasets[set_name], batch_size=opt.batch_size, sampler=sampler, drop_last=True, num_workers=opt.workers, pin_memory=False)
        else:
            dl = DataLoader(datasets[set_name], batch_size=opt.batch_size, shuffle=False, num_workers=opt.workers, pin_memory=False)
        loaders[set_name] = dl
    return loaders


def set_model(opt, device: torch.device, verbose: bool = True):
    if "_1d" in opt.dataset_name:
        if opt.model_name == "clip":
            backbone = AudioCLIP(nmb_prototypes=opt.nmb_prototypes, d_model=opt.out_channels, cfg=opt, multilabel=False, seq_len=opt.seq_len)
            model = CustomModel1(backbone)
        elif opt.model_name == "inception":
            model = InceptionModel(
                opt.num_blocks,
                opt.in_channel,
                out_channels=opt.out_channels,
                stride=opt.stride,
                bottleneck_channels=opt.bottleneck_channels,
                kernel_sizes=opt.kernel_sizes,
                input_length=opt.seq_len,
                use_residuals="default",
                num_pred_classes=opt.num_classes,
                self_train=False, num_positions=0
            )
        elif opt.model_name == "timesnet":
            from src.networks.timesnet import Model
            model = Model("classification", opt.seq_len, 1, 0, 3, 1024, 1, 3, 512, 256, "fixed", "h", 0.0, 0, 2)
        elif opt.model_name == "transformer":
            from src.networks.Transformer import Model
            model = Model(opt)
        else:
            raise NotImplementedError(opt.model_name)
    elif "_2d" in opt.dataset_name or "teus" in opt.dataset_name:
        from src.networks.resnet import resnet34
        backbone, emb_dim = resnet34(num_channels=1)
        head = torch.nn.Linear(emb_dim, opt.num_classes)
        model = torch.nn.Sequential(backbone, head)
    else:
        raise NotImplementedError(opt.dataset_name)

    # checkpoint loading (path is an argument; no defaults)
    if getattr(opt, "ckpt", None):
        ckpt = torch.load(opt.ckpt, map_location="cpu")
        try:
            model.load_state_dict(ckpt, strict=True)
            print("Loaded checkpoint (strict).")
        except Exception as e:
            print("Partial load:", e)
            model.load_state_dict(ckpt, strict=False)

    # optional IsoMax head reset when not resuming
    if not getattr(opt, "resumed", False) and opt.criterion == "isomax":
        if opt.model_name == "inception":
            num_features = model.classifier.linear01[0].in_features
            model.classifier = IsoMaxPlusLossFirstPart(num_features, opt.num_classes)
        elif hasattr(model, "projection"):
            num_features = model.projection.in_features
            model.projection = IsoMaxPlusLossFirstPart(num_features, opt.num_classes)
            model.classifier = model.projection
        else:
            num_features = list(model.children())[-1].in_features
            head = IsoMaxPlusLossFirstPart(num_features, opt.num_classes)
            backbone = torch.nn.Sequential(*list(model.children())[:-1])
            model = torch.nn.Sequential(backbone, head)

    # data parallel / device move
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        model = torch.nn.DataParallel(model)
    model = model.to(device)

    if getattr(opt, "freeze_backbone", False):
        backbone = torch.nn.Sequential(*list(model.children())[:-1])
        for p in backbone.parameters(): p.requires_grad = False
        backbone.eval()

    if verbose:
        tot = sum(p.numel() for p in model.parameters()) / 1e6
        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
        print(f"Trainable params: {trainable}/{tot:.2f}M")

    return model