[DeepSeg] MNIST 데이터셋 YOLO 형식으로 변환

숫자 이미지 전처리를 위해 OpenCV로 ROI를 추출하던 중, YOLOv를 이용해 숫자 객체를 탐지한 후 이를 MNIST 모델을 이용해 분류하는 것이 낫겠다고 생각했다. 이 포스트에서는 MNIST 데이터셋을 YOLO용으로 변환해 학습시키는 과정을 정리하고, 다음 포스트에서 변환된 데이터셋으로 YOLOv5 모델을 학습시키는 과정을 다루도록 한다.

1. MNIST 데이터셋 다운로드 

import torchvision.datasets as datasets

# MNIST 데이터 다운로드 (train + test 세트)
datasets.MNIST(root="../mnist_data", train=True, download=True)
datasets.MNIST(root="../mnist_data", train=False, download=True)

파이썬 코드를 실행해 데이터셋을 다운로드 받는다.

 

2. MNIST 데이터셋을 YOLO 데이터셋으로 변환

다음 코드를 실행해서 MNIST 데이터셋을 .jpg로 변환하고, YOLO학습을 위해 YOLO 데이터셋 형식의 .txt 레이블 파일을 생성한다.

#include <opencv2/opencv.hpp>
#include <fstream>
#include <iostream>

using namespace cv;
using namespace std;

// MNIST 데이터 경로 설정
#define IMAGE_FILE "../data/mnist_data/MNIST/raw/train-images-idx3-ubyte"
#define LABEL_FILE "../data/mnist_data/MNIST/raw/train-labels-idx1-ubyte"

// 이미지 저장 및 YOLO 라벨 생성
void saveImageAndLabel(const Mat& img, int label, int index) {
    string imgPath = "../data/yolo_data/images/" + to_string(index) + ".jpg";
    imwrite(imgPath, img);

    float cx = 0.5, cy = 0.5, w = 1.0, h = 1.0;
    string labelPath = "../data/yolo_data/labels/" + to_string(index) + ".txt";
    ofstream labelFile(labelPath);
    labelFile << label << " " << cx << " " << cy << " " << w << " " << h << endl;
    labelFile.close();
}

int main() {
    ifstream imageFile(IMAGE_FILE, ios::binary);
    ifstream labelFile(LABEL_FILE, ios::binary);
    
    if (!imageFile.is_open() || !labelFile.is_open()) {
        cerr << "Could not open MNIST dataset!" << endl;
        return -1;
    }

    // MNIST 데이터셋 헤더 건너뛰기
    imageFile.seekg(16, ios::beg);
    labelFile.seekg(8, ios::beg);

    for (int i = 0; i < 60000; i++) {  // 60,000개의 MNIST 학습 데이터
        Mat img(28, 28, CV_8U);
        char label;
        imageFile.read((char*)img.data, 28 * 28);
        labelFile.read(&label, 1);

        resize(img, img, Size(640, 640));
        saveImageAndLabel(img, label, i);
    }

    imageFile.close();
    labelFile.close();
    return 0;
}

MNIST 데이터는 28×28 크기의 그레이스케일 숫자 이미지로 구성된다. 따라서 이 데이터를 640×640 크기의 YOLO 입력 이미지와 라벨 형식으로 변환한다.

 

#define IMAGE_FILE "../data/mnist_data/MNIST/raw/train-images-idx3-ubyte"
#define LABEL_FILE "../data/mnist_data/MNIST/raw/train-labels-idx1-ubyte"

    ifstream imageFile(IMAGE_FILE, ios::binary);
    ifstream labelFile(LABEL_FILE, ios::binary);

 

train-images-idx3-ubyte는 28x28 픽셀의 이미지 데이터이고, train-labels-idx1-ubyte는 0~9의 정답 레이블이 저장된 바이너리 파일이다. 이미지 파일의 앞 16byte는 헤더 부분으로 데이터에 대한 정보가 저장되고, 이후 60000개의 28x28 크기 이미지 데이터가 존재한다. 레이블 파일의 앞 8byte는 마찬가지로 헤더 부분이며, 이후 60000개의 정답 레이블이 존재한다.

 

    imageFile.seekg(16, ios::beg);
    labelFile.seekg(8, ios::beg);

따라서 위와 같이 헤더 부분을 건너뛰고 데이터 부분만 읽도록 설정한다. 

 

    for (int i = 0; i < 60000; i++) {  // 60,000개의 MNIST 학습 데이터
        Mat img(28, 28, CV_8U);
        char label;
        imageFile.read((char*)img.data, 28 * 28);
        labelFile.read(&label, 1);

        resize(img, img, Size(640, 640));
        saveImageAndLabel(img, label, i);
    }

 

MNIST 데이터를 읽고 처리한다. 28×28 크기의 숫자 이미지를 OpenCV Mat 형태로 읽어 처리하고, 정답 레이블은 1byte 크기로 읽는다. YOLO 모델은 일반적으로 640x640의 입력을 받으므로, 읽은 이미지 데이터를 640x640으로 리사이즈한다.

 

void saveImageAndLabel(const Mat& img, int label, int index) {
    string imgPath = "../data/yolo_data/images/" + to_string(index) + ".jpg";
    imwrite(imgPath, img);

    float cx = 0.5, cy = 0.5, w = 1.0, h = 1.0;
    string labelPath = "../data/yolo_data/labels/" + to_string(index) + ".txt";
    ofstream labelFile(labelPath);
    labelFile << label << " " << cx << " " << cy << " " << w << " " << h << endl;
    labelFile.close();
}

 읽은 데이터를 YOLO 데이터셋 형식으로 변환한다. YOLO 데이터셋은 다음과 같은 형식으로 구성된다.

<class_id> <x_center> <y_center> <width> <height>

 

class_id : 객체 레이블

x_center: 객체 중심의 X 좌표 (0~1)

y_center: 객체 중심의 Y 좌표 (0~1)

width: 객체의 폭 (0~1)

height: 객체의 높이 (0~1) 

 

MNIST 데이터셋은 배경 없이 숫자가 이미지를 채우므로, 이미지 전체를 객체로 간주하여 이미지의 정중앙을 객체 위치로 설정한다. 위와 같이 단순하게 설정할 수도 있지만, 다음 코드처럼 실제 숫자 객체의 위치를 기준으로 계산할 수도 있다. 

    // 이진화 처리
    Mat bin;
    threshold(img, bin, 10, 255, THRESH_BINARY); // 픽셀 값 10 이상을 흰색(255)으로 설정

    Rect bbox = boundingRect(bin); // 이진화된 이미지에서 숫자 영역 찾기

    //  YOLO 형식 좌표 변환
    float x = (bbox.x + bbox.width / 2.0) / 640.0;
    float y = (bbox.y + bbox.height / 2.0) / 640.0;
    float w = bbox.width / 640.0;
    float h = bbox.height / 640.0;

여기서는 실제 바운딩박스를 기준으로 숫자 객체의 중심 좌표를 계산한다. boundingRect는 입력 점들의 집합, 즉 외곽선을 입력받아 입력 점을 감싸는 바운딩 박스를 생성한다. MNIST 이미지는 그레이스케일 값이므로 이진화 처리를 해준다.

 

따라서 boundingRect(bin)은 숫자 윤곽선을 감싸는 최소 크기의 바운딩 박스를 반환하여 숫자가 있는 부분만 감싸는 바운딩박스를 얻을 수 있다. 

 

bbox.x는 바운딩박스의 왼쪽 상단 x좌표, bbox.width는 바운딩박스의 가로 길이이다. 따라서 bbox.x + bbox.width / 2.0는 바운딩박스 중심의 x좌표이다. 이 좌표를 이미지의 크기인 640으로 나누어 정규화한다. 이와 같은 방법으로 나머지 좌표들도 생성해주면 각 이미지 데이터에서 숫자의 위치를 중심으로 YOLO 데이터셋을 생성할 수 있다. 

 

#include <opencv2/opencv.hpp>
#include <fstream>
#include <iostream>

using namespace cv;
using namespace std;

// MNIST 데이터 경로 설정
#define IMAGE_FILE "../data/mnist_data/MNIST/raw/train-images-idx3-ubyte"
#define LABEL_FILE "../data/mnist_data/MNIST/raw/train-labels-idx1-ubyte"

// 이미지 저장 및 YOLO 라벨 생성
void saveImageAndLabel(const Mat &img, int label, int index)
{
    string imgPath = "../data/yolo_data/images/" + to_string(index) + ".jpg";
    imwrite(imgPath, img);

    // 이진화 처리
    Mat bin;
    threshold(img, bin, 10, 255, THRESH_BINARY); // 픽셀 값 10 이상을 흰색(255)으로 설정

    Rect bbox = boundingRect(bin); // 이진화된 이미지에서 숫자 영역 찾기

    //  YOLO 형식 좌표 변환
    float x = (bbox.x + bbox.width / 2.0) / 640.0;
    float y = (bbox.y + bbox.height / 2.0) / 640.0;
    float w = bbox.width / 640.0;
    float h = bbox.height / 640.0;

    // YOLO 레이블 저장
    string labelPath = "../data/yolo_data/labels/" + to_string(index) + ".txt";
    ofstream labelFile(labelPath);
    labelFile << label << " ";
    labelFile << fixed << setprecision(6) << x << " " << y << " " << w << " " << h << endl;
    labelFile.close();
}

int main()
{
    ifstream imageFile(IMAGE_FILE, ios::binary);
    ifstream labelFile(LABEL_FILE, ios::binary);

    if (!imageFile.is_open() || !labelFile.is_open())
    {
        cerr << "Could not open MNIST dataset!" << endl;
        return -1;
    }

    // MNIST 데이터셋 헤더 건너뛰기
    imageFile.seekg(16, ios::beg);
    labelFile.seekg(8, ios::beg);

    for (int i = 0; i < 60000; i++)
    { // 60,000개의 MNIST 학습 데이터
        Mat img(28, 28, CV_8U);
        char label;
        imageFile.read((char *)img.data, 28 * 28);
        labelFile.read(&label, 1);

        resize(img, img, Size(640, 640));
        saveImageAndLabel(img, label, i);
    }

    imageFile.close();
    labelFile.close();
    return 0;
}

이제 위 코드를 실행하면 다음과 같은 디렉터리 구조로 데이터가 생성된다. 레이블 저장 코드에서 labelFile << fixed << setprecision(6)는 소수점 이하 6번째 자리까지 출력하도록 한다. 

 

zerogod@WIN-KICA6EUHAVM:~/project1$ tree data
data
├── mnist_data
│   └── MNIST
│       └── raw
└── yolo_data
    ├── images
    │	├── 0.jpg
    │	├── 1.jpg
    │	├── 2.jpg
    │	├── ...
    └── labels
    	├── 0.txt
        ├── 1.txt
        ├── 2.txt
        ├── ...

 

생성된 레이블들은 다음과 같이 저장된다.

5 0.500000 0.535937 0.743750 0.746875
0 0.519531 0.500000 0.635938 0.746875
4 0.460938 0.535937 0.737500 0.746875