from transformers import PretrainedConfig

_ATR_LABELS = [
    "Background",
    "Hat",
    "Hair",
    "Sunglasses",
    "Upper-clothes",
    "Skirt",
    "Pants",
    "Dress",
    "Belt",
    "Left-shoe",
    "Right-shoe",
    "Face",
    "Left-leg",
    "Right-leg",
    "Left-arm",
    "Right-arm",
    "Bag",
    "Scarf",
]


class SCHPConfig(PretrainedConfig):
    r"""
    Configuration for **Self-Correction-Human-Parsing (SCHP)**.

    Args:
        num_labels (`int`, *optional*, defaults to 18):
            Number of segmentation classes (18 for ATR dataset).
        input_size (`int`, *optional*, defaults to 512):
            Spatial resolution the model expects (height = width).
        backbone (`str`, *optional*, defaults to `"resnet101"`):
            Backbone architecture name. Only `"resnet101"` is supported.
    """

    model_type = "schp"

    def __init__(
        self,
        num_labels: int = 18,
        input_size: int = 512,
        backbone: str = "resnet101",
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.num_labels = num_labels
        self.input_size = input_size
        self.backbone = backbone

        if "id2label" not in kwargs:
            self.id2label = {
                str(i): lbl for i, lbl in enumerate(_ATR_LABELS[:num_labels])
            }
        if "label2id" not in kwargs:
            self.label2id = {
                lbl: str(i) for i, lbl in enumerate(_ATR_LABELS[:num_labels])
            }