線形代数学+Rustで画像圧縮のアルゴリズムを実装する

こんにちは👋
長く暑い夏が終わろうとしている今ですが、筆者は秋の季節を満喫しております。

LabBaseでは線形代数学の基礎を使って検索エンジンを構築していますが、レコメンド、検索アルゴリズムによく使われる王道の手法について記事を書くことにしました。

概要

線形代数学の特異値分解（SVD）の知識を活かして、原始的な画像圧縮アルゴリズムをRustで実装します。

SVDとは？

SVDは、線形代数学でよく使われる行列の分解です。行列の分解は、同じマトリックスを他のマトリックスに分けて表現することです。SVDの他に、LU三角分解、QR分解などがあります。

SVDは、あるAというマトリックスの列空間と行空間の固有ベクトルを計算して、それぞれをUとVというマトリックスに収めます。さらに、Σという対角行列に、固有値の平方根を入れます。Vの転置行列をV'と定義しますが、以下の分解になります。

A = UΣV'

Σの体格行列の値が、大きさ順に並んでいますが、この意味合いはUの列とVの行がそれぞれにAを表現するのにどれほど重要なのかを表しています。

これがなぜ有意義なのかというと、この原理を利用することによって、マトリックスAを、より少ないデータで推測することができるからです。

例えば、Uの１から5の列、Σの５平方マトリックス、V'の1から5までの行を掛けると、Aに近いマトリックスが出来上がるのです。

今回の記事はこの現象を画像圧縮で応用しますが、複雑で膨大なデータのどの部分が大きな影響力を持っているのか、つまり傾向を見抜けるのに有効な手法であり、機械学習の基礎的な原理でもあります。

Rustのimage crateで画像を開く

本題のさてRustでの実装ですが、まず、画像を読み込みたいので、imageというクレートを使います。

そして、使う画像ですが、愛する地元の観光推進に使われている画像を借ります。

こちらです。

この画像をダウンロードして新しいRustのプロジェクトダイレクトリに入れてください。筆者は./images/PureMichigan-Background-EbanIceCaves2_0.jpgに入れています。

それからこの画像を開くロジックを書きます。

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let img = image::open("images/PureMichigan-Background-EbanIceCaves2_0.jpg")?;

    let grayscale_matrix = img.grayscale().to_luma8();

    todo!();
}

灰色の0から255の値で表したマトリックスを作りたいのでluma8を使います。

画像をマトリックスに移す

上記読み込んだ画像を線形代数学の計算ができるライブラリーに移したいのですが、そのライブラリーはnalgebraと言います。

こちらをプロジェクトに追加します。

それから、マトリックスの型を定義します。

use nalgebra::DMatrix;

type ImageMatrix = DMatrix<f64>;

f64が値の動的マトリックスになります。
なぜf64を使うかというと、画像の0-255の値を推測するような計算になるので、小数点を使わないとエラーが大きくなるからです。
それと、なぜ動的マトリックスを使うかというと、画像の幅と高さが動的だからです。上記の画像は形がわかっているから静的マトリックスも定義できます。nalgebraで静的マトリックスを定義すると、間違ったマトリックス計算をしたら、コンパイル時に指摘してもらえるのです。なので、マトリックスのサイズが決まっているならSMatrixを使った方がいいです。

imageのマトリックスをnalgebraのマトリックスにする

imageで読み込んだマトリックスをnalgebraのマトリックスにしますが、筆者は強引な手法を使いますのでもっといい方法をぜひお考えください。

    let mut matrix = ImageMatrix::zeros(
        grayscale_matrix.height() as usize,
        grayscale_matrix.width() as usize,
    );

    for (row_i, row) in grayscale_matrix.rows().enumerate() {
        for (col_i, entry) in row.enumerate() {
            let mat_entry = matrix.get_mut((row_i, col_i)).expect("must have entry");

            *mat_entry = entry[0] as f64;
        }
    }

綺麗ではないのですが、nalgebraのDMatrix::from_***の大量の関数を解読するのが億劫でした😂。

nalgebraでSVDを計算する

次、上記のマトリックスをSVDによりUΣV'の分解にします。

    let svd = nalgebra::SVD::new(matrix, true, true);

    let u = svd.u.unwrap();
    let v_t = svd.v_t.unwrap();

    let mut svd_s_diag = ImageMatrix::zeros(u.nrows(), v_t.ncols());
    svd_s_diag.set_diagonal(&svd.singular_values);

SVDにtrue,trueの引数を入れているのは、UとV'の計算を省かないように指示するためです。Uだけを計算したい、もしくはV'だけを計算したい時がありますのでこの機能があります。今回は行空間も列空間も必須なので計算してもらいます。

それで、UとV'をuとv_tに開梱しますが、ここはSomeであることが上記の引数でわかるのでunwrapでOKです。

svd_s_diagですが、なぜかnalgebraはΣを対角行列ではなく、一つの行として返すのです。筆者は線形代数学の魔法使いのレベル１にも満たないので、慣れ親しんでいる対角行列に戻して計算を考えたいのです。それで上記のsvd_s_diagという平方行列を作ることにしました。
ちなみに、理論上、assert!(u.nrows() != v_t.ncols())です。元のマトリックス（今回は画像）が平方だった場合のみ、u.nrows() == v_t.ncols()がtrueです。

SVDを使ってランク別の画像圧縮を試してみる

これから、4種類の圧縮した画像を、UとΣとV'を掛けて作ります。Uのどれくらいの列を、V'のどれくらいの行を使うか、これをランクと言います。

ランクが高ければ高いほど元のマトリックス（画像）に近づいていきます。今回の場合は画像の幅が1080pxなので、ランク1080まで行けば、小数点のエラーを除けば元の画像が出るはずです。
逆に、ランクを低くするというのは、画像のマトリックスを推測するのに使うデータを少なくするということです。なので、ランクが低ければ低いほど、画像が荒くなります。

まずコードを共有します。

    for rank in [5, 20, 100, 200] {
        let u_ncols = u.columns(0, rank);
        let nxn_sigma = svd_s_diag.view((0, 0), (rank, rank));
        let v_t_nrows = v_t.rows(0, rank);

        let img_approx_matrix = u_ncols * nxn_sigma * v_t_nrows;

        let mut img_r_1 =
            image::DynamicImage::new_luma8(grayscale_matrix.width(), grayscale_matrix.height());

        for (row_i, pixels) in img_r_1.as_mut_luma8().unwrap().rows_mut().enumerate() {
            let svd_row = img_approx_matrix.row(row_i);

            for (col_i, pixel) in pixels.enumerate() {
                let entry = svd_row.get(col_i).expect("must have entry");

                pixel[0] = *entry as u8;
            }
        }

        img_r_1
            .save(std::path::Path::new(&format!("./output/rank_{}.png", rank)))
            .unwrap();
    }

保存するのに、またimageの画像マトリックスに戻すのですが、こちらも強引にやってしまいました😂。もうちょい腕力のある人ならより綺麗に書くことでしょう。

結果画像

上記のコードを実行すると、それぞれのランクの画像を出力します。

注意
cargo runを実行すると、最適化されていないマシーンコードがコンパイルされます。通常、これはどうということでもないのですが、今回の場合は大きく変わるので、cargo run --profile releaseを使うことを推奨します。

元の画像

ランク5

ランク20

ランク100

ランク200

元画像のおおよそ2割のデータ量です。

まとめ

線形代数学の理論を応用して、画像圧縮アルゴリズムを作りましたが、いかがでしたでしょうか？

筆者は、まず、Rustでこうして数学実験ができることを喜ばしく思いました。PythonやMATLABでやるのが遥かに簡単でしょうが、Rustで実験してもいいくらい簡単に思いました。

それから、理論上の数学を教科書で読んだのをそのまま、コードに落としただけでこれだけの結果が出ることに素直に驚きました。鳥肌が立ったと言ったら過言かもしれませんが、数学は素晴らしいなと思いました。

ちなみに、JPEGのような画像圧縮にはSVDではなく、離散フーリエ変換（Fourier Transform）が使われているそうです。

今回の記事は画像圧縮に特化したものですが、実はSVDはさまざまなデータマトリックスに対して応用できるものです。何千行もあるユーザーのクリック履歴を分析したり、画像認識にも使える手法です。顔認識はまさにこのSVDを使って行われているのです。

筆者が働くLabBaseではこうした基礎的な分野においても追求しております。機械学習のはやりだけでなく、数学の理論を活かしてパフォーマンスを追求した機械学習を内製化しているのです。高みを目指したい仲間がいるなら歓迎します。