codingutf8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license collection of utils to be used by backbones and their components import enum import inspect from typing import iterable list optional tuple union class backbonetypeenum enum timm timm transformers transformers def verifyoutfeaturesoutindices outfeatures optionaliterablestr outindices optionaliterableint stagenames optionaliterablestr if stagenames is none raise valueerrorstagenames must be set for transformers backbones if outfeatures is not none if not isinstanceoutfeatures list raise valueerrorfoutfeatures must be a list typeoutfeatures if anyfeat not in stagenames for feat in outfeatures raise valueerrorfoutfeatures must be a subset of stagenames stagenames got outfeatures if outindices is not none if not isinstanceoutindices list tuple raise valueerrorfoutindices must be a list or tuple got typeoutindices if anyidx lenstagenames for idx in outindices raise valueerrorfoutindices must be valid indices for stagenames stagenames got outindices if outfeatures is not none and outindices is not none if lenoutfeatures lenoutindices raise valueerroroutfeatures and outindices should have the same length if both are set if outfeatures stagenamesidx for idx in outindices raise valueerroroutfeatures and outindices should correspond to the same stages if both are set def alignoutputfeaturesoutputindices outfeatures optionalliststr outindices optionalunionlistint tupleint stagenames liststr if outindices is none and outfeatures is none outindices lenstagenames 1 outfeatures stagenames1 elif outindices is none and outfeatures is not none outindices stagenames indexlayer for layer in outfeatures elif outfeatures is none and outindices is not none outfeatures stagenamesidx for idx in outindices return outfeatures outindices def getalignedoutputfeaturesoutputindices outfeatures optionalliststr outindices optionalunionlistint tupleint stagenames liststr tupleliststr listint first verify that the outfeatures and outindices are valid verifyoutfeaturesoutindicesoutfeaturesoutfeatures outindicesoutindices stagenamesstagenames outputfeatures outputindices alignoutputfeaturesoutputindices outfeaturesoutfeatures outindicesoutindices stagenamesstagenames verify that the aligned outfeatures and outindices are valid verifyoutfeaturesoutindicesoutfeaturesoutputfeatures outindicesoutputindices stagenamesstagenames return outputfeatures outputindices class backbonemixin backbonetype optionalbackbonetype none def inittimmbackboneself config none if getattrself backbone none is none raise valueerrorself backbone must be set before calling inittimmbackbone these will diagree with the defaults for the transformers models e g for resnet50 the transformer model has outfeatures stem stage1 stage2 stage3 stage4 the timm model has outfeatures act layer1 layer2 layer3 layer4 self stagenames stagemodule for stage in self backbone featureinfo info self numfeatures stagenumchs for stage in self backbone featureinfo info outindices self backbone featureinfo outindices outfeatures self backbone featureinfo modulename we verify the out indices and out features are valid verifyoutfeaturesoutindices outfeaturesoutfeatures outindicesoutindices stagenamesself stagenames self outfeatures self outindices outfeatures outindices def inittransformersbackboneself config none stagenames getattrconfig stagenames outfeatures getattrconfig outfeatures none outindices getattrconfig outindices none self stagenames stagenames self outfeatures self outindices getalignedoutputfeaturesoutputindices outfeaturesoutfeatures outindicesoutindices stagenamesstagenames number of channels for each stage this is set in the transformer backbone model init self numfeatures none def initbackboneself config none self config config self usetimmbackbone getattrconfig usetimmbackbone false self backbonetype backbonetype timm if self usetimmbackbone else backbonetype transformers if self backbonetype backbonetype timm self inittimmbackboneconfig elif self backbonetype backbonetype transformers self inittransformersbackboneconfig else raise valueerrorfbackbonetype self backbonetype not supported property def outfeaturesself return self outfeatures outfeatures setter def outfeaturesself outfeatures liststr self outfeatures self outindices getalignedoutputfeaturesoutputindices outfeaturesoutfeatures outindicesnone stagenamesself stagenames property def outindicesself return self outindices outindices setter def outindicesself outindices uniontupleint listint self outfeatures self outindices getalignedoutputfeaturesoutputindices outfeaturesnone outindicesoutindices stagenamesself stagenames property def outfeaturechannelsself the current backbones will output the number of channels for each stage even if that stage is not in the outfeatures list return stage self numfeaturesi for i stage in enumerateself stagenames property def channelsself return self outfeaturechannelsname for name in self outfeatures def forwardwithfilteredkwargsself args kwargs signature dictinspect signatureself forward parameters filteredkwargs k v for k v in kwargs items if k in signature return selfargs filteredkwargs def forward self pixelvalues outputhiddenstates optionalbool none outputattentions optionalbool none returndict optionalbool none raise notimplementederrorthis method should be implemented by the derived class def todictself output super todict outputoutfeatures output popoutfeatures outputoutindices output popoutindices return output class backboneconfigmixin property def outfeaturesself return self outfeatures outfeatures setter def outfeaturesself outfeatures liststr self outfeatures self outindices getalignedoutputfeaturesoutputindices outfeaturesoutfeatures outindicesnone stagenamesself stagenames property def outindicesself return self outindices outindices setter def outindicesself outindices uniontupleint listint self outfeatures self outindices getalignedoutputfeaturesoutputindices outfeaturesnone outindicesoutindices stagenamesself stagenames def todictself output super todict outputoutfeatures output popoutfeatures outputoutindices output popoutindices return output coding utf 8 2023 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license collection of utils to be used by backbones and their components verify that out_indices and out_features are valid for the given stage_names finds the corresponding out_features and out_indices for the given stage_names the logic is as follows out_features not set out_indices set out_features is set to the out_features corresponding to the out_indices out_indices not set out_features set out_indices is set to the out_indices corresponding to the out_features out_indices and out_features not set out_indices and out_features are set to the last stage out_indices and out_features set input out_indices and out_features are returned args out_features list str the names of the features for the backbone to output out_indices list int or tuple int the indices of the features for the backbone to output stage_names list str the names of the stages of the backbone get the out_features and out_indices so that they are aligned the logic is as follows out_features not set out_indices set out_features is set to the out_features corresponding to the out_indices out_indices not set out_features set out_indices is set to the out_indices corresponding to the out_features out_indices and out_features not set out_indices and out_features are set to the last stage out_indices and out_features set they are verified to be aligned args out_features list str the names of the features for the backbone to output out_indices list int or tuple int the indices of the features for the backbone to output stage_names list str the names of the stages of the backbone first verify that the out_features and out_indices are valid verify that the aligned out_features and out_indices are valid initialize the backbone model from timm the backbone must already be loaded to self _backbone these will diagree with the defaults for the transformers models e g for resnet50 the transformer model has out_features stem stage1 stage2 stage3 stage4 the timm model has out_features act layer1 layer2 layer3 layer4 we verify the out indices and out features are valid number of channels for each stage this is set in the transformer backbone model init method to initialize the backbone this method is called by the constructor of the base class after the pretrained model weights have been loaded set the out_features attribute this will also update the out_indices attribute to match the new out_features set the out_indices attribute this will also update the out_features attribute to match the new out_indices the current backbones will output the number of channels for each stage even if that stage is not in the out_features list serializes this instance to a python dictionary override the default to_dict from pretrainedconfig to include the out_features and out_indices attributes a mixin to support handling the out_features and out_indices attributes for the backbone configurations set the out_features attribute this will also update the out_indices attribute to match the new out_features set the out_indices attribute this will also update the out_features attribute to match the new out_indices serializes this instance to a python dictionary override the default to_dict from pretrainedconfig to include the out_features and out_indices attributes
import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union


class BackboneType(enum.Enum):
    TIMM = "timm"
    TRANSFORMERS = "transformers"


def verify_out_features_out_indices(
    out_features: Optional[Iterable[str]], out_indices: Optional[Iterable[int]], stage_names: Optional[Iterable[str]]
):
    
    if stage_names is None:
        raise ValueError("Stage_names must be set for transformers backbones")

    if out_features is not None:
        if not isinstance(out_features, (list,)):
            raise ValueError(f"out_features must be a list {type(out_features)}")
        if any(feat not in stage_names for feat in out_features):
            raise ValueError(f"out_features must be a subset of stage_names: {stage_names} got {out_features}")

    if out_indices is not None:
        if not isinstance(out_indices, (list, tuple)):
            raise ValueError(f"out_indices must be a list or tuple, got {type(out_indices)}")
        if any(idx >= len(stage_names) for idx in out_indices):
            raise ValueError(f"out_indices must be valid indices for stage_names {stage_names}, got {out_indices}")

    if out_features is not None and out_indices is not None:
        if len(out_features) != len(out_indices):
            raise ValueError("out_features and out_indices should have the same length if both are set")
        if out_features != [stage_names[idx] for idx in out_indices]:
            raise ValueError("out_features and out_indices should correspond to the same stages if both are set")


def _align_output_features_output_indices(
    out_features: Optional[List[str]],
    out_indices: Optional[Union[List[int], Tuple[int]]],
    stage_names: List[str],
):
    
    if out_indices is None and out_features is None:
        out_indices = [len(stage_names) - 1]
        out_features = [stage_names[-1]]
    elif out_indices is None and out_features is not None:
        out_indices = [stage_names.index(layer) for layer in out_features]
    elif out_features is None and out_indices is not None:
        out_features = [stage_names[idx] for idx in out_indices]
    return out_features, out_indices


def get_aligned_output_features_output_indices(
    out_features: Optional[List[str]],
    out_indices: Optional[Union[List[int], Tuple[int]]],
    stage_names: List[str],
) -> Tuple[List[str], List[int]]:
    
    
    verify_out_features_out_indices(out_features=out_features, out_indices=out_indices, stage_names=stage_names)
    output_features, output_indices = _align_output_features_output_indices(
        out_features=out_features, out_indices=out_indices, stage_names=stage_names
    )
    
    verify_out_features_out_indices(out_features=output_features, out_indices=output_indices, stage_names=stage_names)
    return output_features, output_indices


class BackboneMixin:
    backbone_type: Optional[BackboneType] = None

    def _init_timm_backbone(self, config) -> None:
        
        if getattr(self, "_backbone", None) is None:
            raise ValueError("self._backbone must be set before calling _init_timm_backbone")

        
        
        
        self.stage_names = [stage["module"] for stage in self._backbone.feature_info.info]
        self.num_features = [stage["num_chs"] for stage in self._backbone.feature_info.info]
        out_indices = self._backbone.feature_info.out_indices
        out_features = self._backbone.feature_info.module_name()

        
        verify_out_features_out_indices(
            out_features=out_features, out_indices=out_indices, stage_names=self.stage_names
        )
        self._out_features, self._out_indices = out_features, out_indices

    def _init_transformers_backbone(self, config) -> None:
        stage_names = getattr(config, "stage_names")
        out_features = getattr(config, "out_features", None)
        out_indices = getattr(config, "out_indices", None)

        self.stage_names = stage_names
        self._out_features, self._out_indices = get_aligned_output_features_output_indices(
            out_features=out_features, out_indices=out_indices, stage_names=stage_names
        )
        
        self.num_features = None

    def _init_backbone(self, config) -> None:
        
        self.config = config

        self.use_timm_backbone = getattr(config, "use_timm_backbone", False)
        self.backbone_type = BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

        if self.backbone_type == BackboneType.TIMM:
            self._init_timm_backbone(config)
        elif self.backbone_type == BackboneType.TRANSFORMERS:
            self._init_transformers_backbone(config)
        else:
            raise ValueError(f"backbone_type {self.backbone_type} not supported.")

    @property
    def out_features(self):
        return self._out_features

    @out_features.setter
    def out_features(self, out_features: List[str]):
        
        self._out_features, self._out_indices = get_aligned_output_features_output_indices(
            out_features=out_features, out_indices=None, stage_names=self.stage_names
        )

    @property
    def out_indices(self):
        return self._out_indices

    @out_indices.setter
    def out_indices(self, out_indices: Union[Tuple[int], List[int]]):
        
        self._out_features, self._out_indices = get_aligned_output_features_output_indices(
            out_features=None, out_indices=out_indices, stage_names=self.stage_names
        )

    @property
    def out_feature_channels(self):
        
        
        return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names)}

    @property
    def channels(self):
        return [self.out_feature_channels[name] for name in self.out_features]

    def forward_with_filtered_kwargs(self, *args, **kwargs):
        signature = dict(inspect.signature(self.forward).parameters)
        filtered_kwargs = {k: v for k, v in kwargs.items() if k in signature}
        return self(*args, **filtered_kwargs)

    def forward(
        self,
        pixel_values,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        raise NotImplementedError("This method should be implemented by the derived class.")

    def to_dict(self):
        
        output = super().to_dict()
        output["out_features"] = output.pop("_out_features")
        output["out_indices"] = output.pop("_out_indices")
        return output


class BackboneConfigMixin:
    

    @property
    def out_features(self):
        return self._out_features

    @out_features.setter
    def out_features(self, out_features: List[str]):
        
        self._out_features, self._out_indices = get_aligned_output_features_output_indices(
            out_features=out_features, out_indices=None, stage_names=self.stage_names
        )

    @property
    def out_indices(self):
        return self._out_indices

    @out_indices.setter
    def out_indices(self, out_indices: Union[Tuple[int], List[int]]):
        
        self._out_features, self._out_indices = get_aligned_output_features_output_indices(
            out_features=None, out_indices=out_indices, stage_names=self.stage_names
        )

    def to_dict(self):
        
        output = super().to_dict()
        output["out_features"] = output.pop("_out_features")
        output["out_indices"] = output.pop("_out_indices")
        return output