RendeiroLab/MetPredict-lung-structure-segmentation

1"""HuggingFace `PreTrainedModel` + `PretrainedConfig` wrapper around `DPT`.2 3Lets consumers do `AutoModel.from_pretrained(repo_id, trust_remote_code=True)`4without importing the local `DPT` class. The `auto_map` field on the config5tells HF to bundle `hf_model.py` + `dpt.py` with the uploaded weights so the6classes are reconstructable in a clean env.7"""8from __future__ import annotations9 10from typing import Any, Literal, Optional, cast11 12import torch13import torch.nn.functional as F14from transformers import (15    AutoConfig,16    AutoModel,17    PretrainedConfig,18    PreTrainedModel,19)20from transformers.modeling_outputs import SemanticSegmenterOutput21 22from .dpt import DPT23 24 25class DPTConfig(PretrainedConfig):26    model_type = "metpredict_dpt"27 28    def __init__(29        self,30        n_classes: int = 4,31        class_names: Optional[list[str]] = None,32        backbone: str = "hf-hub:bioptimus/H-optimus-0",33        encoder_depth: int = 4,34        decoder_intermediate_channels: tuple[int, ...] = (224, 448, 896, 896),35        decoder_fusion_channels: int = 224,36        decoder_readout: str = "cat",37        activation: Optional[str] = None,38        in_channels: int = 3,39        **kwargs,40    ):41        super().__init__(**kwargs)42        self.n_classes = n_classes43        self.class_names = list(class_names) if class_names else []44        self.backbone = backbone45        self.encoder_depth = encoder_depth46        self.decoder_intermediate_channels = list(decoder_intermediate_channels)47        self.decoder_fusion_channels = decoder_fusion_channels48        self.decoder_readout = decoder_readout49        self.activation = activation50        self.in_channels = in_channels51        # `auto_map` makes the repo loadable as AutoModel without local imports.52        self.auto_map = {53            "AutoConfig": "hf_model.DPTConfig",54            "AutoModel": "hf_model.DPTForSegmentation",55        }56 57 58class DPTForSegmentation(PreTrainedModel):59    config_class = DPTConfig60    base_model_prefix = "dpt"61    main_input_name = "pixel_values"62    all_tied_weights_keys: dict = {}  # To be compatible with transformers 4.x and 5.x63 64    def __init__(self, config: DPTConfig):65        super().__init__(config)66        # `decoder_readout` is Literal-typed in DPT — cast since pydantic-loaded67        # value is a plain str.68        readout = cast(Literal["ignore", "add", "cat"], config.decoder_readout)69        dpt_kwargs: dict[str, Any] = dict(70            encoder_name=config.backbone,71            encoder_depth=config.encoder_depth,72            decoder_readout=readout,73            decoder_intermediate_channels=tuple(config.decoder_intermediate_channels),74            decoder_fusion_channels=config.decoder_fusion_channels,75            in_channels=config.in_channels,76            classes=config.n_classes,77            activation=config.activation,78        )79        self.dpt = DPT(**dpt_kwargs)80        # Skip post_init weight init — DPT.initialize() already ran inside DPT.__init__.81 82    def forward(83        self,84        pixel_values: torch.Tensor,85        labels: Optional[torch.Tensor] = None,86        return_dict: bool = True,87    ):88        logits = self.dpt(pixel_values)89        loss: Optional[torch.Tensor] = None90        if labels is not None:91            loss = F.cross_entropy(logits, labels.long())92        if not return_dict:93            return (loss, logits) if loss is not None else (logits,)94        # SemanticSegmenterOutput expects FloatTensor — cast suppresses Pylance.95        return SemanticSegmenterOutput(96            loss=cast(Any, loss),97            logits=cast(Any, logits),98        )99 100 101def _register() -> None:102    try:103        AutoConfig.register("metpredict_dpt", DPTConfig)104        AutoModel.register(DPTConfig, DPTForSegmentation)105    except ValueError:106        # Already registered (re-import).107        pass108 109 110_register()111