LoginSignup
2
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

Qiitanがほしい人の一人アドカレAdvent Calendar 2024

Day 23
@nikawamikan(Mikan Nikawa)inTeam shahu

PythonからRust製モジュールを呼び出して高速化したかった

Posted at

よみとばしてください

限界派遣SESのnikawamikanです。どうも。
アドベントカレンダーもあと少しですね。
ひとりアドカレ23日目です。

また遅れました。ネタが思いつかなくて・・・。

RustからPythonを呼び出せるPyO3

PyO3はRustでPythonの拡張モジュールを書くためのライブラリです。
これを使えば、ちょっと遅いPythonくんの処理を高速でイケイケなRustくんにおもたーい処理を任せることができます。

そんなわけで、FastAPIで作ったエンドポイントの処理をRustで書いた拡張モジュールを使って簡単な画像の合成処理を高速化してみたいと思います。

ちなみに筆者はRustに慣れてないので、かなり苦戦しました。

環境構築

今回はRustとPythonの環境が必要なため、DevContainerでDockerfileを書いて環境を作っていきます。

FROM mcr.microsoft.com/devcontainers/rust:1-1-bookworm

# 必要なパッケージをインストール
RUN apt update -y && apt install -y \
    wget \
    build-essential \
    libssl-dev \
    zlib1g-dev \
    libncurses5-dev \
    libncursesw5-dev \
    libreadline-dev \
    libsqlite3-dev \
    libgdbm-dev \
    libdb5.3-dev \
    libbz2-dev \
    libexpat1-dev \
    liblzma-dev \
    tk-dev \
    libffi-dev \
    libgdbm-compat-dev

RUN wget https://www.python.org/ftp/python/3.12.0/Python-3.12.0.tgz && \
    tar xzf Python-3.12.0.tgz && \
    cd Python-3.12.0 && \
    ./configure --enable-optimizations && \
    make -j$(nproc) && \
    make altinstall && \
    cd .. && \
    rm -rf Python-3.12.0 Python-3.12.0.tgz

RUN ln -s /usr/local/bin/python3.12 /usr/bin/python && \
    ln -s /usr/local/bin/pip3.12 /usr/bin/pip

apt installでPythonのインストールをしてもいいのですが、バージョンが古めになってしまうため、ソースからビルドしています。
ビルドする場合は色々とパッケージが必要になってしまうので、それらをインストールしています。

またdevcontainer.jsonは以下のようになります。

devcontainer.json
{
	"name": "pyo3",
	"build": {
		"dockerfile": "Dockerfile",
	},
	"customizations": {
		"vscode": {
			"extensions": [
				"ms-python.vscode-pylance",
				"ms-python.python"
			]
		}
	},
    "postCreateCommand": "source .venv/bin/activate && pip install -r requirements.txt",
}

Pythonの拡張機能のms-python.pythonとms-python.vscode-pylanceをインストールしています。
また、postCreateCommandで仮想環境の有効化と必要なライブラリのインストールをしています。
それらについては後述します。

次にvenvを作成して、プロジェクトを作成します。

python -m venv .venv
source .venv/bin/activate

次にPythonのライブラリをインストールします。

ライブラリ 用途
FastAPI APIサーバーを作成するため
uvicorn FastAPIのサーバーを起動するため
maturin PyO3のビルドツール 
pip install fastapi "uvicorn[standard]" maturin

最後にpip freeze > requirements.txtでライブラリのリストを保存します。

DevContainer自体が仮想環境なのですが、venvで仮想環境を作成しているのは、maturinがvenvを前提としているためです。

Rustのプロジェクトを作成

次にRustのプロジェクトを作成します。

cargo new --lib rust_image

Cargo.tomlにpyo3と、画像生成に必要なクレートを追加します。

Cargo.toml
[dependencies]
pyo3 = { version = "0.23.3", features = ["extension-module"] }
image = "0.25.5"
imageproc = "0.25.0"
rusttype = "0.9.2"
ab_glyph = "0.2.29"

なお、私は画像生成に必要なライブラリの互換性の問題で1時間ほど悩みましたが、cargo treeで依存関係を確認すれば問題ないという学びを得ました。

次にsrc/lib.rsに画像生成の処理を書いていきます。

とりあえず全体像は以下のようになります。

src/lib.rs
use ab_glyph::PxScale;
use imageproc::drawing::{draw_text_mut, text_size};
use pyo3::prelude::*;

#[pyclass]
struct ImageGenerator {
    font_data: Vec<u8>,
}

#[pymethods]
impl ImageGenerator {
    #[new]
    pub fn new(font_path: &str) -> PyResult<Self> {
        let font_data = std::fs::read(font_path)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;
        Ok(Self { font_data })
    }

    pub fn generate_image(&self, image_path: &str, text: &str) -> PyResult<Vec<u8>> {
        let mut img = image::open(image_path)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?
            .to_rgb8();
        let font = ab_glyph::FontRef::try_from_slice(&self.font_data)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

        // スケールを設定
        let scale = PxScale::from(20.0);

        let (w, h) = img.dimensions();

        // テキストのサイズを取得
        let (text_width, _) = text_size(scale, &font, text);

        // テキスト幅と画像の幅が一致するようにスケールを調整
        let scale = PxScale::from((w as f32 / text_width as f32) * 20.0);

        // テキストのサイズを再取得
        let (text_width, text_height) = text_size(scale, &font, text);

        // テキストを中央に配置
        let center_x = (w as f32 - text_width as f32) / 2.0;
        let center_y = (h as f32 - text_height as f32) / 2.0;

        draw_text_mut(
            &mut img,
            image::Rgb([255u8, 255u8, 255u8]),
            center_x as i32,
            center_y as i32,
            scale,
            &font,
            &text.replace('_', " "),
        );

        // 画像をバイト配列として返す
        let mut buf = Vec::new();
        let mut cursor = std::io::Cursor::new(&mut buf);
        img.write_to(&mut cursor, image::ImageFormat::Png)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(buf)
    }
}

#[pymodule]
fn rust_image(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<ImageGenerator>()?;
    Ok(())
}

順に説明します。

エントリーポイント

#[pymodule]
fn rust_image(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<ImageGenerator>()?;
    Ok(())
}

#[pymodule]でモジュールを定義しています。
このm.add_class::<ImageGenerator>()?;でPythonから呼び出せるクラスを追加しています。

struct ImageGenerator

#[pyclass]
struct ImageGenerator {
    font_data: Vec<u8>,
}

ImageGeneratorという構造体を定義しています。
#[pyclass]でPythonから呼び出せるクラスとして定義しています。

impl ImageGenerator::new

#[pymethods]
impl ImageGenerator {
    #[new]
    pub fn new(font_path: &str) -> PyResult<Self> {
        let font_data = std::fs::read(font_path)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?;
        Ok(Self { font_data })
    }
}

#[pymethods]でメソッドし、#[new]でコンストラクタを定義しています。

これはPythonでいう

class ImageGenerator:
    def __init__(self, font_path: str):
        self.font_data = open(font_path, "rb").read()

と同じです。

impl ImageGenerator::generate_image

    pub fn generate_image(&self, image_path: &str, text: &str) -> PyResult<Vec<u8>> {
        let mut img = image::open(image_path)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))?
            .to_rgb8();
        let font = ab_glyph::FontRef::try_from_slice(&self.font_data)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;

        // スケールを設定
        let scale = PxScale::from(20.0);

        let (w, h) = img.dimensions();

        // テキストのサイズを取得
        let (text_width, _) = text_size(scale, &font, text);

        // テキスト幅と画像の幅が一致するようにスケールを調整
        let scale = PxScale::from((w as f32 / text_width as f32) * 20.0);

        // テキストのサイズを再取得
        let (text_width, text_height) = text_size(scale, &font, text);

        // テキストを中央に配置
        let center_x = (w as f32 - text_width as f32) / 2.0;
        let center_y = (h as f32 - text_height as f32) / 2.0;

        draw_text_mut(
            &mut img,
            image::Rgb([255u8, 255u8, 255u8]),
            center_x as i32,
            center_y as i32,
            scale,
            &font,
            &text.replace('_', " "),
        );

        // 画像をバイト配列として返す
        let mut buf = Vec::new();
        let mut cursor = std::io::Cursor::new(&mut buf);
        img.write_to(&mut cursor, image::ImageFormat::Png)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
        Ok(buf)
    }

画像をこねこねしています。
この、画像をこねこねする作業ですが、Rustのライブラリではimageとimageprocとab_glyphを使うという仕様で、それぞれの依存関係があるためCargo.tomlでは喧嘩しないように一つのターゲットに合わせたバージョンを指定する必要があります。

Pythonから呼び出された際の引数として、画像のパスとテキストを受け取り、テキストを画像に描画してバイト配列として返しています。

モジュールのビルド

次にモジュールをビルドします。

maturin develop

これでtargetディレクトリにlibrust_image.soが生成されます。
また、.venvのsite-packagesにrust_imageというディレクトリが生成され、その中に__init__.pyとrust_image.cpython-39-x86_64-linux-gnu.soが生成されます。
そのため、Python側からはimport rust_imageでモジュールを読み込むことができます。

FastAPIでAPIを作成

次にFastAPIでAPIを作成します。

from io import BytesIO
from fastapi.responses import StreamingResponse
import rust_image
from fastapi import FastAPI


app = FastAPI()


@app.get("/")
def generate_image(text: str):
    image_generator = rust_image.ImageGenerator("./assets/Makinas-4-Flat.otf")

    bin_image = image_generator.generate_image("./assets/bancan.png", text)
    
    image = BytesIO(bin_image)

    image.seek(0)

    return StreamingResponse(image, media_type="image/png")

コード量が少なくて安心しますね（？）

rust_imageは先ほどビルドしたモジュールです。
ImageGeneratorクラスをインスタンス化して、generate_imageメソッドを呼び出して画像を生成しています。

byte配列が返ってくるので、BytesIOでバッファに書き込んで、StreamingResponseで画像を返しています。

ちなみにMakinasはフリーフォントでめちゃくちゃかわいいのでみんな使いましょう。

サーバーの起動

最後にサーバーを起動します。

uvicorn main:app --reload

http://localhost:8000/docsにアクセスして、テキストを入力して画像を生成してみましょう。

ほげーとしてますね。

image.png

気になる点

Python側のコードではrust_image配下のクラスなどの型がわからないため、IDEが補完をしてくれません。
pyiを作成して型を定義することで補完が効くようになるそうですが、自動で生成してくれないとこれは不便かなぁーと思いました。

image-1.png

まとめ

PyO3を使ってRustでPythonの拡張モジュールを作成して、FastAPIでAPIを作成してみました。
正直、Rust単体でAPIまで作成しまっても問題ない気がしますが、Pythonのライブラリを使いたい場合などは画像生成部分のみをRustで書くことにより高速化できるかもしれません。(未検証です。すみません。時間がなかったので・・・。おそらくそのうち検証します。)

また、Rustのライブラリの互換性についても学べたので、良い経験になりました。

参考

2
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
2
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?