Dataset 구현

Dataset

---

• Pytorch C++ API (libtorch)로 작성한 데이터셋 클래스: Dataset

• yes와 no라는 두 개의 클래스로 구성된 이미지 데이터를 로드

• 이를 파이토치의 데이터셋 형식으로 변환해주는 역할

• 주요 구성 요소

  1. 생성자
  2. 데이터 가져오기
  3. 데이터셋 크기 반환
  4. 이미지 정규화

1. 생성자

---

Dataset::Dataset(const std::string& data_dir, const std::vector<float>& mean, const std::vector<float>& std)
		
		// 정규화에 필요한 평균, 표준편차를 각각 저장
		// 이는 이미지 데이터를 정규화하는 데 사용함
    : mean_(mean), std_(std) {
    // no: 0, yes: 1로 라벨링 -> 라벨과 클래스 디렉토리 연결
    for (const auto& label : { "no", "yes" }) {
        int label_value = (label == std::string("yes")) ? 1 : 0;
        fs::path sub_dir = fs::path(data_dir) / label;
				
				// 이미지 파일 탐색
        for (const auto& entry : fs::directory_iterator(sub_dir)) {
            if (entry.is_regular_file()) {
                image_paths_.push_back(entry.path().string()); // 이미지 파일 경로 저장
                labels_.push_back(label_value); // 이미지 라벨 저장 
            }
        }
    }
}

• 주어진 디렉토리 경로를 탐색하여 클래스에 해당하는 이미지 파일을 수집
• 각 파일 경로와 라벨을 리스트에 저장하여 데이터셋 구성하는 것이 목적

2. 데이터 가져오기

---

// torch::data::Example<>: 파이토치에서 샘플 하나를 표현하는 객체, 입력 데이터와 라벨의 쌍으로 구성
torch::data::Example<> Dataset::get(size_t index) {
    // 이미지 로드
    cv::Mat img = cv::imread(image_paths_[index], cv::IMREAD_COLOR);
    // 오류 처리: 이미지 파일을 찾을 수 없거나 로드 실패 시
    if (img.empty()) {
        throw std::runtime_error("Failed to load image: " + image_paths_[index]);
    }

    // 이미지 정규화 및 텐서 반환
    // normalize(): OpenCV의 cv::Mat -> torch::Tensor
    // 정규화 함수는 따로 정의해줌 (밑에서)
    torch::Tensor img_tensor = normalize(img, mean_, std_);

    // 라벨 텐서 생성
    torch::Tensor label_tensor = torch::tensor(labels_[index], torch::kInt64);

    return { img_tensor, label_tensor };
}

• 특정 인덱스에 해당하는 샘플(이미지와 라벨)을 가져오는 함수
• 해당 함수는 데이터셋에서 하나의 샘플을 불러오고 이를 정규화하여 텐서 형식을 반환하는 것이 목표

3. 데이터셋 크기 반환

---

torch::optional<size_t> Dataset::size() const {
    return image_paths_.size();
}

4. 이미지 정규화

---

torch::Tensor Dataset::normalize(const cv::Mat& img, const std::vector<float>& mean, const std::vector<float>& std) {
		// 이미지 크기 조정
    cv::Mat resized_img;
    cv::resize(img, resized_img, cv::Size(224, 224)); // Resize for CNN input
    // 픽셀 값 정규화
    resized_img.convertTo(resized_img, CV_32F, 1.0 / 255.0);
		
		// OpenCV 데이터를 Pytorch 텐서로 변환
		// 1: 배치 차원
		// 3: RGB 채널
    auto img_tensor = torch::from_blob(resized_img.data, { 1, resized_img.rows, resized_img.cols, 3 });
    
    // 텐서의 차원 순서 변경
    // OpenCV의 데이터 순서가 (Batch, Height, Width, Channels)이므로, PyTorch에서 사용하는 (Batch, Channels, Height, Width)로 변환
    img_tensor = img_tensor.permute({ 0, 3, 1, 2 });
    
    // 채널별 정규화
    // 각 채널(RGB)에 대해 수행
    for (int i = 0; i < 3; i++) {
        img_tensor[0][i] = (img_tensor[0][i] - mean[i]) / std[i];
    }
		
		// 텐서 복사 후 반환
		// 메모리 공유 가능성 방지
    return img_tensor.clone(); // Ensure the data is not deallocated
}

• resized_img.convertTo(resized_img, CV_32F, 1.0 / 255.0)
  • CV_32F
    • 기본적으로 이미지의 픽셀 값은 0~255 범위를 가짐
    • 딥러닝 모델 학습에서는 이미지 데이터를 소수점으로 처리해야 하기 때문에 CV_32F 형식으로 반환해주어야 함
  • 1.0 / 255.0
    • 픽셀 값을 255로 나누어서 0~1 사이로 스케일링을 수행
    • 니누어진 값은 converTo() 내부에서 자동으로 resized_img에 곱해짐

→ 해당 코드는 데이터셋 준비와 전처리를 간단하게 하여 모델 학습에 바로 사용할 수 있게 함