import os
import ppmatting

current_path = os.path.abspath(os.path.dirname(__file__))


def get_rel_path(path: str):
    return os.path.join(current_path, '..', path)


class MattingModel:
    path = ""
    config = ""
    model = None
    transforms: ppmatting.transforms.Compose
    init = False

    def __init__(self, p, c):
        self.path = get_rel_path(p)
        self.config = get_rel_path(c)


_model: MattingModel

modelDict = {
    "ppmattingv2": MattingModel("models/ppmattingv2-stdc1-human_512.pdparams",
                                "configs/ppmattingv2/ppmattingv2-stdc1-human_512.yml"),

    "ppmatting": MattingModel("models/ppmatting-hrnet_w18-human_512.pdparams",
                              "configs/quick_start/ppmattingv2-stdc1-human_512.yml")
}


def get_model():
    return modelDict.get("ppmattingv2")