说明文档

Informative-Drawings-ONNX

本仓库包含来自论文 C. Chan et al., "Informative Drawings: Learning to generate line drawings that convey geometry and semantics", CVPR 2021 的 ONNX 模型。

环境要求

本仓库已在以下环境中测试：

Ubuntu 22.04
Python 3.11
ONNXRuntime 1.19

使用方法

安装依赖：

pip install numpy onnxruntime Pillow

保存以下脚本并运行：

import argparse
import time
from contextlib import contextmanager

import numpy as np
import onnxruntime
from PIL import Image


@contextmanager
def timer(name: str):
  start = time.perf_counter()
  yield
  end = time.perf_counter()
  print(f\"{name}: {end - start:.3f} [sec]\")


def preprocess(image: Image.Image, size: int) -> np.ndarray:
  image = image.resize((size, size))
  image = np.float32(image) / 255.
  image = np.transpose(image, (2, 0, 1))  # (w, h, c) -> (c, w, h)
  image = np.expand_dims(image, axis=0)  # (c, w, h) -> (b, c, w, h)
  return image


def postprocess(output: np.ndarray, original_size: tuple[int, int] = None) -> Image.Image:
  output_image = np.squeeze(output[0, 0] * 255).astype(np.uint8)
  output_image = Image.fromarray(output_image).convert(\"L\")
  if original_size:
    output_image = output_image.resize(original_size)
  return output_image


def main():

    parser = argparse.ArgumentParser()
    parser.add_argument(\"--image\", type=str, required=True, help=\"Path to input image\")
    parser.add_argument(\"--model\", type=str, required=True, help=\"Path to ONNX model\")
    args = parser.parse_args()

    input_image = Image.open(args.image).convert(\"RGB\")
    input_image.thumbnail((1536, 1536))

    # Load input image and the model
    with timer(f\"Load model\"):
        session = onnxruntime.InferenceSession(args.model)

    size = session.get_inputs()[0].shape[2]  # width

    # Inference
    input_data = preprocess(input_image, size)
    with timer(f\"Inference\"):
        output_data = session.run(None, {\"inputs\": input_data})[0]
    output_image = postprocess(output_data, input_image.size)
    output_image.show()


if __name__ == \"__main__\":
   main()