Matting/tools/model.py

import os
import sys

import paddle
import paddleseg
from paddleseg.utils import get_sys_env

import ppmatting
from ppmatting.core import load
from ppmatting.utils import Config, MatBuilder
import configs
current_path = os.path.abspath(os.path.dirname(__file__))


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


_model: configs.MattingModel


def load_model():
    m = configs.get_model()
    global _model
    _model = m

    model_path = m.config
    cfg = Config(model_path)
    builder = MatBuilder(cfg)

    paddleseg.utils.show_env_info()
    paddleseg.utils.show_cfg_info(cfg)
    paddleseg.utils.set_device("cpu")

    m.model = builder.model
    m.transforms = ppmatting.transforms.Compose(builder.val_transforms)

    load(m.model, m.path)
    m.init = True


def get_model():
    return _model