SegFormer란?

Min Jae Cho·2024년 12월 3일

오늘은 SegFormer 모델을 활용하여 균열 감지를 위한 세그멘테이션 작업을 진행했습니다. 프로젝트의 주요 목표는 다양한 이미지에서 균열을 정확히 탐지하는 모델을 학습시키는 것이었으며, 이를 위해 SegFormer라는 세그멘테이션 모델을 사용했습니다. 이번 프로젝트에서 사용한 기술과 참고할 만한 내용들을 정리해보겠습니다.

SegFormer 모델 개요

SegFormer는 최근 각광받고 있는 세그멘테이션 모델로, 단순하면서도 효율적인 성능을 보여줍니다. 이번 작업에서는 간단한 U-Net 구조와 유사하게 SegFormer를 구현하여 균열 감지 문제에 적용해 보았습니다. 기본적으로 Encoder-Decoder 형태로 구성된 이 모델은 다음과 같은 주요 구성 요소를 포함합니다:

Encoder: 이미지에서 특징을 추출하는 역할로, 컨볼루션 계층과 ReLU 활성화 함수를 사용하여 구성했습니다.

Decoder: Encoder에서 추출한 특징을 바탕으로 원본 이미지의 해상도로 복원하여 각 픽셀의 클래스를 예측하는 역할을 합니다.

데이터 전처리와 마스크 생성

특히 이번 프로젝트에서는 레이블 데이터가 JSON 파일로 주어졌고, 이 JSON 파일에서 균열 정보를 기반으로 마스크 이미지를 생성하는 과정을 거쳤습니다. 이를 위해 다음과 같은 단계를 진행했습니다:

JSON 파일 파싱 및 마스크 생성: 각 JSON 파일에는 'annotations'라는 키가 있으며, 이 안에 균열 정보가 포함되어 있습니다. 'polygon' 또는 'polyline' 형태로 저장된 좌표를 이용하여 균열이 있는 영역만 마스크 이미지로 변환했습니다.

마스크 이미지 생성 후 저장: 균열이 있는 경우에만 마스크 이미지를 생성하고, 빈 마스크는 생성하지 않도록 하여 효율적으로 데이터를 관리했습니다.

CSV 파일 생성: 균열이 있는 이미지 파일명을 따로 CSV 파일에 저장하여 후속 작업에서 참고할 수 있도록 했습니다.

모델 학습

SegFormer 모델 학습은 아래의 과정으로 진행되었습니다:

데이터셋 로딩: 이미지와 마스크를 불러와 학습에 사용할 수 있도록 데이터셋을 정의했습니다. 이미지와 마스크 간 개수가 일치하지 않는 경우 경고를 출력하고, 일치하는 경우에만 학습에 사용하도록 했습니다.

손실 함수 및 옵티마이저: CrossEntropyLoss를 사용하여 모델의 예측 결과와 실제 마스크 간의 손실을 계산하고, AdamW 옵티마이저를 사용하여 학습을 진행했습니다.

학습 루프: 모델 학습 중 매 에포크마다 학습 손실과 검증 손실을 출력하여 모델 성능을 모니터링했습니다. 또한 검증 손실이 개선될 경우 모델을 저장했습니다.

오늘의 성과와 느낀 점

JSON 기반 마스크 생성: 레이블링 데이터에서 유효한 정보를 추출하고 이를 마스크로 변환하는 과정이 특히 흥미로웠습니다. 자동으로 마스크를 생성하여 균열이 있는 이미지만 선별할 수 있었던 점이 효율적이었습니다.

모델 성능 개선 가능성: 아직 초기 학습 단계라서 모델 성능이 완벽하지는 않지만, 데이터 전처리의 중요성을 다시 한번 느꼈습니다. 특히 균열이 없는 이미지를 잘 걸러내는 것이 학습 효율을 높이는 데 도움이 될 것으로 보입니다.

참고할 만한 내용

데이터 전처리 효율화: JSON 파일에서 필요 없는 데이터를 걸러내는 방식으로 데이터셋을 효율화하는 작업은 이후 다른 프로젝트에서도 사용할 수 있을 것 같습니다. 특히 대용량 데이터에서 불필요한 연산을 줄이는 것이 중요했습니다.

SegFormer 구조 변경: SegFormer의 기본 구조를 간단하게 구현했지만, 추후 Transformer 기반으로 encoder를 개선하는 것도 고려해볼 만합니다.

다음 단계

모델 구조 개선: U-Net과 유사한 단순한 구조에서 벗어나 SegFormer의 원래 구조에 가까운 Transformer 기반의 인코더로 업그레이드할 계획입니다.

데이터셋 추가 확보 및 증강: 현재 사용하고 있는 데이터셋 외에도 추가적인 균열 데이터를 확보하고, 데이터 증강 기법을 적용하여 모델의 일반화 성능을 높일 예정입니다.

오늘 진행한 프로젝트는 균열 탐지 문제를 다루는 데 있어 데이터 전처리의 중요성과 SegFormer 모델의 가능성을 보여주는 좋은 경험이었습니다.

profile
A.I. Engineer

0개의 댓글