Matting/tools/model.py

import os
import sys

import paddle
import paddleseg
from paddleseg.utils import get_sys_env, logger

from ppmatting.transforms import Compose

import ppmatting
from ppmatting.core import load
from ppmatting.utils import Config, MatBuilder

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: paddle.nn.layer.Layer
    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 load_model():
    m = modelDict.get("ppmattingv2")
    global _model
    _model = m

    model_path=get_rel_path("configs/quick_start/ppmattingv2-stdc1-human_512.yml")
    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