Rustでpythonコードを呼び出してトラッキング
概要
仕様するモデルがonnxであればortを使って呼び出してrustのみで物体検出とかできるんですが、例えばトラッキングの場合は物体検出モデルを呼び出した後にカルマンフィルタなどを呼び出してIDを追跡するので、全てをRustで完結させるのは難しいです。
そこで、rustからpythonのコードを呼び出してトラッキングを行い、rust側で整形してクラウドと通信すれば良いんじゃないかと思い、rustとpythonを組み合わせてトラッキングをしてみました。
リポジトリは下記に公開しています。
Mac上では、READMEの手順に従って実現できていることを確認しています。ただ、python側のメモリについてはrustの範囲外なので、大きな動画サイズのトラッキングをしてメモリが足りなくなった場合にパニックを起こしてしまいます(ので、python側でよろしくなるなどの対応が必要なはず)。
解説
Dockerfileについて
今回はrustとpythonを共存させたdocker環境を用意する必要があるため、下記のようなdockerfileを用意します。
FROM python:3.11-slim AS base
# The Stage to install rust.
FROM base AS rust-install
RUN apt-get update && \
apt-get install -y curl \
build-essential g++ cmake pkg-config libssl-dev libopencv-dev && \
curl https://sh.rustup.rs -sSf | sh -s -- -y && \
. "$HOME/.cargo/env" && \
rustc --version && cargo --version \
apt-get clean
# Final stage.
FROM rust-install AS final
# Set path.
ENV PATH="/root/.cargo/bin:${PATH}"
WORKDIR /app
# Copy only Cargo.toml and Cargo.lock first, then build the dependent libraries (to make use of Docker's cache)
COPY Cargo.toml Cargo.lock ./
RUN mkdir src && echo 'fn main() {}' > src/main.rs && \
cargo build --release && \
rm -rf src
# Copy source codes.
COPY . .
# Install python library.
RUN pip install --upgrade pip
RUN pip install poetry
COPY . .
RUN poetry config virtualenvs.create false \
&& poetry install --without ml --no-interaction --no-ansi
RUN rm pyproject.toml poetry.lock
# Build rust codes.
RUN cargo build --release
docker環境の立ち上げは、docker-compose.ymlと組み合わせて下記のコマンドで完結します。ただし、下記を実行する前にpoetry lockでpoetry.lockファイルを作成してください。
# 環境変数を読み込み
$ cp .env{.example,}
# ビルドしてコンテナに入る
$ DOCKER_BUILDKIT=1 docker build . -t rust-track-image
$ docker compose run --rm rust-track-demo
python側のコード
pythonコードはsrc/tracking.pyに置かれていて、ultralyticsでトラッキングをした結果をフレームごとの結果のリストとして取得しています。
from ultralytics import YOLO
def run_tracking(video_path: str, model_hash: str, conf_threshold: float):
# 簡易的なトラッキングコード例
model = YOLO(model_hash)
results = model.track(
source=video_path,
persist=True,
verbose=False,
conf=conf_threshold,
stream=True,
)
outputs = []
for result in results:
output = {}
if result.boxes is not None:
for track in result.boxes:
track_id = int(track.id.cpu().item()) if track.id is not None else -1
xywhn = track.xywhn.cpu().numpy().tolist()
xyxy = track.xyxy.cpu().numpy().tolist()
cls = int(track.cls.cpu().item())
conf = float(track.conf.cpu().item())
output[track_id] = {
"xywhn": xywhn,
"xyxy": xyxy,
"cls": cls,
"conf": conf,
}
outputs.append(output)
return outputs
rust側のコード
必要なライブラリ
Cargo.toml は下記のようになっています。今回はrustからpythonコードを呼び出すのでpyo3をインストールしています。また、jsonファイルの書き込みができるようにserde, clapのderive APIでコマンドラインの引数を設計できるようにしています。
[package]
name = "track"
version = "0.1.0"
edition = "2021"
[dependencies]
clap = { version = "4.5.32", features = ["derive"] }
pyo3 = { version = "0.21", features = ["auto-initialize"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
パラメータ設定
まず、pythonのコードを呼び出す際に異なるモデルを呼び出したりconfidenceを変更できるように、下記の引数を設計します。今回は、動画のファイルパスとモデルパス、confidenceを指定できるようにしています。
ソースコードはsrc/args.rsに対応しています。
use clap::Parser;
#[derive(Parser, Clone)]
#[command(name = "track", version = "0.1.0", author = "chiikawa", about = "yolo tracking")]
pub struct Args {
/// ONNX model path
#[arg(short = 'm', long = "model", value_name = "MODEL", default_value = "yolov8n", help = "ONNX model path")]
pub model: String,
/// input path
#[arg(value_name = "SOURCE", help = "Input data.", required = true)]
pub source: String,
/// Confidence threshold.
#[arg(short='c', long="conf-threshold", default_value="0.7")]
pub conf: f32,
}
pythonの呼び出し部分
ソースコードはリポジトリのsrc/lib.rsに対応しています。
まず、pythonコードを認識できるように、sys_path.insertで./srcを追加します。その後、tracking.pyを呼び出して関数を取得します。
今回は複数の引数を指定する必要があるため、PyTupleで複数の引数を渡せるようにしてから関数を呼び出しています。tracking.pyの結果はList[Dict[str, Any]]なので、PyListとして取得できるようにしています。
Python::with_gil(|py| {
let sys_path: &PyList = py.import("sys")?.getattr("path")?.downcast()?;
sys_path.insert(0, "./src")?; // カレントディレクトリ
let tracking_module = PyModule::import_bound(py, "tracking")?;
// Python関数を取得
let run_tracking = tracking_module.getattr("run_tracking")?;
// 引数をPythonタプルとして渡す
let args = PyTuple::new_bound(
py,
&[
config.source.clone().into_py(py),
config.model_hash.clone().into_py(py),
config.conf_threshold.into_py(py),
],
);
// 関数を呼び出し、結果をPyListとして取得
let result = run_tracking.call1(args)?;
let result_dict = result.downcast::<PyList>()?;
json形式に整形
最後にjsonに整形して保存できるようにしています。
// Pythonのjson.dumpsを使ってstrに変換
let json_module = PyModule::import_bound(py, "json")?;
let json_str_obj = json_module.getattr("dumps")?.call1((result_dict,))?;
let json_str: &str = json_str_obj.extract()?;
// Rust側でserde_json::Valueとしてパース
let result_json: Value = serde_json::from_str(json_str)
.map_err(|e| PyValueError::new_err(e.to_string()))?;
// JSONファイルに保存
let file = File::create("tracking_result.json").expect("Failed to create file");
to_writer_pretty(file, &result_json).expect("Failed to write JSON");
実行
では、最後に実行してみます。dockerのコンテナに入ったら、下記を実行してみてください。
$ cargo run -- -m yolov8n.pt <動画ファイル>
root@312b84ee0b3c:~/workspace# cargo run -- -m yolov8n.pt sample.mkv
Compiling track v0.1.0 (/root/workspace)
Finished `dev` profile [unoptimized + debuginfo] target(s) in 2.65s
Running `target/debug/track -m yolov8n.pt sample.mkv`
Tracking result saved to tracking_result.json
まとめ
今回はrustからpythonを呼び出してトラッキングするところまで実施してみました。
将来的には、トラッキングのアルゴリズムも含めて、rustでやれるところまでやってみても面白いかもしれません。
Discussion