befozg/PortraitTransfer

1import torchvision2import io3import numpy as np4import torch.onnx5import onnx6from tools import Inference, Matting, log, extract_matte, harmonize, css, execute_onnx_model7from omegaconf import OmegaConf8import os9import sys10import torch11import numpy as np12import torchvision.transforms.functional as tf13from PIL import Image14import cv2 as cv15from onnxruntime import InferenceSession16 17args = OmegaConf.load(os.path.join(f"./config/demo.yaml"))18 19log("Model loading")20phnet = Inference(**args)21stylematte = Matting(**args)22log("Model loaded")23model = stylematte.model24 25x = torch.randn((1, 3, 720, 1280))26mask = torch.ones((1, 1, 512, 512))27path = 'checkpoints/stylematte-test.onnx'28 29# Export30torch.onnx.export(model, x, path, opset_version=16)31 32# Validation33onnx_model = onnx.load(path)34onnx.checker.check_model(onnx_model)35# execute_onnx_model(x, onnx_model)36