RMBG-1.4 / MyConfig.py
OriLib's picture
integrate with transformers (#21)
33c41df verified
raw
history blame
326 Bytes
from transformers import PretrainedConfig
from typing import List
class RMBGConfig(PretrainedConfig):
model_type = "SegformerForSemanticSegmentation"
def __init__(
self,
in_ch=3,
out_ch=1,
**kwargs):
self.in_ch = in_ch
self.out_ch = out_ch
super().__init__(**kwargs)