MNISTファイルから画像抽出

ちょっと所要で手書きの数字を認識させたい今日この頃。
手書きの数字といえばMNIST
これをtesstrainを利用してTesseract用の辞書にするため、画像ファイルとラベルファイルに変換したい。
VisualStudioで適当なC#コンソールアプリを作ったので、ベロっとソース貼っておきます。
Tesseractと戯れてる人はそこに時間を掛けてる場合じゃないと思うので、ベロっとソースが必要なんだ。
誰かに使ってもらえたらそれで良い。
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;
            }
        }
    }
}

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

返信を残す

メールアドレスが公開されることはありません。

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