C++와 TensorFlow를 이용한 쓰레기 이미지 분류

안녕하세요, 여러분! 오늘은 C++과 TensorFlow를 이용하여 쓰레기 이미지를 분석하고 분류하는 방법에 대해 알아보려고 합니다. 이 과정에서는 사전 학습된 CNN 모델을 사용하여 이미지가 유리, 플라스틱, 알루미늄, 종이, 또는 재활용할 수 없는 재질인지 분류합니다.

TensorFlow 모델 다운로드

TensorFlow 허브는 다양한 사전 학습된 모델을 제공하는 저장소입니다. 이 중 MobileNet V2 모델을 사용하여 이미지를 분류해 보겠습니다. MobileNet V2 모델은 경량화된 신경망으로, 모바일과 임베디드 애플리케이션에 적합합니다.

모델 다운로드 링크: https://www.kaggle.com/models/google/mobilenet-v2/

준비 사항

1. TensorFlow 모델 다운로드: 사전 학습된 MobileNet V2 모델을 TensorFlow 허브에서 다운로드합니다.
2. OpenCV 설치: OpenCV가 설치되어 있어야 하며, DNN 모듈을 포함하여 빌드해야 합니다.
3. TensorFlow C++ API 설치: TensorFlow C++ API를 설치하고 빌드합니다.

C++ 코드 작성

이제 C++와 TensorFlow를 사용하여 사전 학습된 모델을 로드하고 이미지를 분류하는 코드를 작성해보겠습니다.

1. 헤더 파일 포함 및 네임스페이스 사용

먼저, 필요한 헤더 파일을 포함하고 네임스페이스를 설정합니다.

#include <tensorflow/cc/client/client_session.h>
#include <tensorflow/cc/ops/standard_ops.h>
#include <tensorflow/core/framework/tensor.h>
#include <tensorflow/core/platform/env.h>
#include <tensorflow/core/public/session.h>
#include <tensorflow/core/util/command_line_flags.h>
#include <opencv2/opencv.hpp>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>

using namespace tensorflow;
using namespace cv;
using namespace std;

2. 레이블 로드 함수

다음으로, 클래스 레이블이 포함된 파일을 읽어 벡터에 저장하는 함수를 작성합니다.

vector<string> readLabels(const string& labelFile) {
    vector<string> labels;
    ifstream file(labelFile);
    string line;
    while (getline(file, line)) {
        labels.push_back(line);
    }
    return labels;
}

3. 텐서 변환 함수

이미지를 TensorFlow 텐서로 변환하는 함수를 작성합니다.

Tensor MatToTensor(const Mat& img) {
    Tensor input_tensor(DT_FLOAT, TensorShape({1, img.rows, img.cols, img.channels()}));
    auto input_tensor_mapped = input_tensor.tensor<float, 4>();
    
    for (int y = 0; y < img.rows; ++y) {
        for (int x = 0; x < img.cols; ++x) {
            for (int c = 0; c < img.channels(); ++c) {
                input_tensor_mapped(0, y, x, c) = img.at<Vec3f>(y, x)[c];
            }
        }
    }
    return input_tensor;
}

4. 메인 함수

마지막으로, 모델과 구성을 로드하고 이미지를 분류하는 메인 함수를 작성합니다.

int main(int argc, char* argv[]) {
    // 모델 파일 경로
    string modelFile = "path/to/saved_model.pb";
    string labelFile = "path/to/labels.txt";
    string imageFile = "path/to/image.jpg";

    // 레이블 로드
    vector<string> labels = readLabels(labelFile);

    // 이미지 로드 및 전처리
    Mat img = imread(imageFile);
    if (img.empty()) {
        cerr << "Could not read input image..." << endl;
        return -1;
    }
    resize(img, img, Size(224, 224));
    img.convertTo(img, CV_32FC3);
    img = (img / 127.5) - 1.0;

    // 텐서로 변환
    Tensor input_tensor = MatToTensor(img);

    // 세션 생성 및 모델 로드
    Session* session;
    Status status = NewSession(SessionOptions(), &session);
    if (!status.ok()) {
        cerr << status.ToString() << endl;
        return -1;
    }
    GraphDef graph_def;
    status = ReadBinaryProto(Env::Default(), modelFile, &graph_def);
    if (!status.ok()) {
        cerr << status.ToString() << endl;
        return -1;
    }
    status = session->Create(graph_def);
    if (!status.ok()) {
        cerr << status.ToString() << endl;
        return -1;
    }

    // 입력 및 출력 텐서 이름 설정
    vector<pair<string, Tensor>> inputs = {{"input", input_tensor}};
    vector<Tensor> outputs;
    status = session->Run(inputs, {"MobilenetV2/Predictions/Reshape_1"}, {}, &outputs);
    if (!status.ok()) {
        cerr << status.ToString() << endl;
        return -1;
    }

    // 결과 분석
    Tensor output = outputs[0];
    auto scores = output.flat<float>();
    float max_score = scores(0);
    int max_index = 0;
    for (int i = 1; i < scores.size(); ++i) {
        if (scores(i) > max_score) {
            max_score = scores(i);
            max_index = i;
        }
    }

    // 결과 출력
    string label = labels[max_index];
    cout << "Class: " << label << ", Confidence: " << max_score << endl;

    // 결과 이미지 저장
    putText(img, format("%s: %.2f", label.c_str(), max_score), Point(20, 20), FONT_HERSHEY_SIMPLEX, 0.7, Scalar(0, 255, 0), 2);
    imwrite("result.jpg", img);

    // 세션 종료
    session->Close();
    delete session;

    return 0;
}

주요 단계 설명

1. 헤더 파일 포함 및 네임스페이스 사용: TensorFlow와 OpenCV의 헤더 파일을 포함하고, 네임스페이스를 설정합니다.
2. 레이블 로드 함수: 클래스 레이블이 포함된 파일을 읽어 벡터에 저장하는 함수입니다.
3. 텐서 변환 함수: OpenCV의 Mat 객체를 TensorFlow 텐서로 변환하는 함수입니다.
4. 메인 함수: 모델, 레이블 및 이미지를 로드하고 전처리합니다. 이미지를 텐서로 변환하고 TensorFlow 세션을 생성합니다. 모델을 실행하여 예측 결과를 얻고 가장 높은 확률을 가진 클래스를 선택합니다. 결과를 콘솔에 출력하고, 이미지에 레이블을 추가하여 저장합니다.

이 코드 예제는 TensorFlow와 OpenCV를 사용하여 이미지를 분류하는 기본적인 예제입니다. 필요에 따라 모델 파일, 구성 파일, 레이블 파일 경로를 설정하고, 추가적인 전처리 및 후처리 과정을 포함할 수 있습니다. 이 과정을 통해 여러분도 쓰레기 이미지를 분류하는 시스템을 구축할 수 있습니다.

그럼 오늘도 즐거운 코딩 되세요!