File size: 1,484 Bytes
a54fa65 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | 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])
}
|