BioMike/clipsegmulticlass

1from transformers import PretrainedConfig
2from transformers.utils import logging
3
4logger = logging.get_logger(__name__)
5
6class ClipSegMultiClassConfig(PretrainedConfig):
7    model_type = "clipseg-multiclass"
8    is_composition = False
9
10    def __init__(
11        self,
12        class_labels=None,
13        label2color=None,
14        model="CIDAS/clipseg-rd64-refined",
15        image_size=352,
16        **kwargs
17    ):
18        super().__init__(**kwargs)
19
20        self.class_labels = class_labels or []
21        self.num_classes = len(self.class_labels)
22
23        self.label2color = label2color or {
24            i: [
25                int(255 * (i / max(1, self.num_classes - 1))),
26                0,
27                255 - int(255 * (i / max(1, self.num_classes - 1)))
28            ]
29            for i in range(self.num_classes)
30        }
31
32        self.model = model
33        self.image_size = image_size
34