MNISTファイルから画像抽出

ちょっと所要で手書きの数字を認識させたい今日この頃。
手書きの数字といえばMNIST
これをtesstrainを利用してTesseract用の辞書にするため、画像ファイルとラベルファイルに変換した

VisualStudioで適当なC#コンソールアプリを作ったので、ベロっとソース貼っておきます。
Tesseractと戯れてる人はそこに時間を掛けてる場合じゃないと思うので、ベロっとソースが必要なんだ。
誰かに使ってもらえたらそれで良い。

ダウンロードしたtrain-images-idx3-ubyte.gzとtrain-labels-idx1-ubyte.gzは、解凍しておいてください。
C#のコンソールアプリプロジェクトを作ったら、System.Drawingの参照を追加して、下記ソースをベロっと貼ります。
Main関数の先頭ローカル変数4つを適当に修正して、実行してください。
実行完了まで数分かかります。

using System;
using System.Drawing;
using System.Drawing.Imaging;
using System.IO;

namespace MNISTfilesConverter
{
    class Program
    {
        private const int MNIST_IMAGE_COUNT = 60000;    // MNISTの画像ファイル数
        private const int MNIST_SIZE = 28;              // MNISTの画像ファイルの1辺の長さ

        static void Main(string[] args)
        {
            string input_dir = @"C:\Users\hogehoge\Documents\MNIST";
            string output_dir = @"\\wsl$\Ubuntu\home\hogehoge\ocrd-train\data\mnist-ground-truth";
            int lineLen = 32;   // 1画像あたりの文字数
            int scale = 1;      // 拡大率

            // フォルダの用意(前のデータは削除)
            if (Directory.Exists(output_dir))
            {
                Directory.Delete(output_dir, true);
            }
            Directory.CreateDirectory(output_dir);

            // 画像情報を読み取る
            ImageInfo[] trainImages = LoadData(input_dir + @"\train-images.idx3-ubyte", input_dir + @"\train-labels.idx1-ubyte");

            // OCR-D用にtifファイルとgt.txtファイルを出力
            for (int i = 0; 0 <= i; i++)
            {
                Bitmap ret = null;
                string text = "";
                string name = @"\" + i.ToString("D5");
                for (int j = 0; j < lineLen; j++)
                {
                    int idx = i * lineLen + j;
                    if (trainImages.Length <= idx) break;

                    if (ret == null)
                    {
                        ret = MakeBitmap(trainImages[idx], scale);
                    }
                    else
                    {
                        ret = ChainBitmap(ret, MakeBitmap(trainImages[idx], scale));
                    }
                    text += trainImages[idx].Label;
                }

                if (ret == null) break;

                ret.Save(output_dir + name + ".tif", ImageFormat.Tiff);
                StreamWriter sw = new StreamWriter(output_dir + name + ".gt.txt");
                sw.Write(text + "\n");
                sw.Close();
            }
        }

        /// <summary>
        /// 画像を横につなげる
        /// </summary>
        /// <param name="img1">元データ</param>
        /// <param name="img2">連携つデータ</param>
        /// <returns>連結後の画像</returns>
        private static Bitmap ChainBitmap(Bitmap img1, Bitmap img2)
        {
            Bitmap ret = new Bitmap(img1.Width + img2.Width, img1.Height > img2.Height ? img1.Height : img2.Height);
            Graphics g = Graphics.FromImage(ret);
            g.DrawImage(img1, new Point(0, 0));
            g.DrawImage(img2, new Point(img1.Width, 0));
            g.Dispose();
            return ret;
        }

        /// <summary>
        /// 画像に変換
        /// </summary>
        /// <param name="info">画像情報</param>
        /// <param name="scale">拡大率(1で等倍)</param>
        /// <returns>画像</returns>
        private static Bitmap MakeBitmap(ImageInfo info, int scale)
        {
            Bitmap result = new Bitmap(MNIST_SIZE * scale, MNIST_SIZE * scale);
            Graphics gr = Graphics.FromImage(result);
            for (int i = 0; i < MNIST_SIZE; ++i)
            {
                for (int j = 0; j < MNIST_SIZE; ++j)
                {
                    int pixelColor = 255 - info.Pixels[i, j];
                    Color color = Color.FromArgb(pixelColor, pixelColor, pixelColor);
                    SolidBrush sb = new SolidBrush(color);
                    gr.FillRectangle(sb, j * scale, i * scale, scale, scale);
                }
            }
            return result;
        }

        /// <summary>
        /// MNISTのファイルを画像情報に変換
        /// </summary>
        /// <param name="pixelFile">train-images.idx3-ubyteファイルパス</param>
        /// <param name="labelFile">train-labels.idx1-ubyteファイルパス</param>
        /// <returns>画像情報配列</returns>
        private static ImageInfo[] LoadData(string pixelFile, string labelFile)
        {
            ImageInfo[] ret = new ImageInfo[MNIST_IMAGE_COUNT];
            byte[,] pixels = new byte[MNIST_SIZE, MNIST_SIZE];
            FileStream fsPixels = new FileStream(pixelFile, FileMode.Open);
            FileStream fsLabels = new FileStream(labelFile, FileMode.Open);
            BinaryReader brImages = new BinaryReader(fsPixels);
            BinaryReader brLabels = new BinaryReader(fsLabels);

            // 画像データは先頭16byteは識別子、画像データ数、行数、列数なので読み飛ばす
            brImages.ReadBytes(16);
            // ラベルデータは先頭8byteは識別子とデータ数なので読み飛ばす
            brLabels.ReadBytes(8);

            for (int i = 0; i < MNIST_IMAGE_COUNT; i++)
            {
                for (int x = 0; x < MNIST_SIZE; x++)
                {
                    for (int y = 0; y < MNIST_SIZE; y++)
                    {
                        pixels[x, y] = brImages.ReadByte();
                    }
                }
                ret[i] = new ImageInfo(pixels, brLabels.ReadByte());
            }
            fsPixels.Close();
            brImages.Close();
            fsLabels.Close();
            brLabels.Close();

            return ret;
        }
        /// <summary>
        /// 画像情報
        /// </summary>
        private class ImageInfo
        {
            // 画像データ
            public byte[,] Pixels = new byte[MNIST_SIZE, MNIST_SIZE];
            // ラベル
            public byte Label;
            public ImageInfo(byte[,] pixels, byte label)
            {
                Array.Copy(pixels, this.Pixels, pixels.Length);
                this.Label = label;
            }
        }
    }
}

ローカル変数output_dirに指定したフォルダにtifファイルとtxtファイルができていれば完成です。
たぶん。
下記はlineLen = 1で実行した結果。

ちなみに、1画像40文字にすると、txtファイルは40文字を1行におさめたファイルを生成します。
Tesseractで学習する際、こんなエラーメッセージが出るためです。

nohup: ignoring input
find -L data/mnist-ground-truth -name '*.gt.txt' | xargs paste -s > "data/mnist/all-gt"
unicharset_extractor --output_unicharset "data/mnist/unicharset" --norm_mode 2 "data/mnist/all-gt"
Extracting unicharset from box file data/mnist/all-gt
Wrote unicharset file data/mnist/unicharset
PYTHONIOENCODING=utf-8 python3 generate_line_box.py -i "data/mnist-ground-truth/00136.tif" -t "data/mnist-ground-truth/00136.gt.txt" > "data/mnist-ground-truth/00136.box"
Traceback (most recent call last):
  File "/root/tess/ocrd-train/generate_line_box.py", line 32, in <module>
    raise ValueError("ERROR: %s: Ground truth text file should contain exactly one line, not %s" % (args.txt, len(lines)))
ValueError: ERROR: data/mnist-ground-truth/00136.gt.txt: Ground truth text file should contain exactly one line, not 40
make: *** [Makefile:218: data/mnist-ground-truth/00136.box] Error 1
Command exited with non-zero status 2
Run time = 0:00.12

こちらも、1行のプレーンテキストにすべしと仰っている。

Transcriptions must be single-line plain text and have the same name as the line image but with the image extension replaced by .gt.txt

text2image形式で吐き出す

ただ、BOXファイルとしては正しくないよな?
ベースになる辞書に再学習させようとするとBOXファイル形式がおかしいと怒られるので、text2image形式でも吐き出したい(Ubuntu18.04なら怒られなかった…まぁ、無駄)

Main関数を下記に変更します。

        static void Main(string[] args)
        {
            string input_dir = @"C:\Users\hogehoge\Documents\MNIST";
            string output_dir = @"\\wsl$\Ubuntu\home\hogehoge\ocrd-train\data\mnist-ground-truth";
            int lineLen = 32;   // 1画像あたりの文字数
            int scale = 1;      // 拡大率

            // フォルダの用意(前のデータは削除)
            if (Directory.Exists(output_dir))
            {
                Directory.Delete(output_dir, true);
            }
            Directory.CreateDirectory(output_dir);

            // 画像情報を読み取る
            ImageInfo[] trainImages = LoadData(input_dir + @"\train-images.idx3-ubyte", input_dir + @"\train-labels.idx1-ubyte");

            // OCR-D用にtifファイルとgt.txtファイルを出力
            for (int i = 0; 0 <= i; i++)
            {
                Bitmap ret = null;
                string name = @"\" + i.ToString("D5");
                string text = "";
                int left = 0;
                for (int j = 0; j < lineLen; j++)
                {
                    int idx = i * lineLen + j;
                    if (trainImages.Length <= idx) break;

                    if (ret == null)
                    {
                        ret = MakeBitmap(trainImages[idx], scale);
                    }
                    else
                    {
                        ret = ChainBitmap(ret, MakeBitmap(trainImages[idx], scale));
                    }

                    if(text.Length != 0 )
                    {
                        text += "\n";
                    }
                    text += trainImages[idx].Label + " " 
                            + left.ToString() + " 0 " + (left + MNIST_SIZE).ToString() + " " + MNIST_SIZE.ToString() + " 0";
                    left += MNIST_SIZE;
                }

                if (ret == null) break;

                ret.Save(output_dir + name + ".tif", ImageFormat.Tiff);
                StreamWriter sw = new StreamWriter(output_dir + name + ".gt.txt");
                sw.Write(text);
                sw.Close();
            }
        }

EMNISTから抽出

微妙にデータが違うので、よく似てるけど別のコードを用意します。


他にもニッチなIT関連要素をまとめていますので、よければ一覧記事もご覧ください。

返信を残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

日本語が含まれない投稿は無視されますのでご注意ください。(スパム対策)