[AI] Semantic Segmentation을 해보자. feat.Custom dataset

어찌저찌 회사에서 업무를 하다보니

Semantic Segmentation을 진행하게 되었다. 그것도 Custom Dataset으로,,,

중요한건 데이터셋을 외부에서 직접 공수해야한다는 점이고, 그말은 즉 데이터 수집부터 라벨링까지 전부 스스로 해결해야 한다는 의미이다.

이에 그 기록을 남겨보자.

 

1. Semantic Segmentation?

여기서 다룰 내용이므로 간략한 설명은 아래의 포스트를 참고하자

https://medium.com/hyunjulie/1%ED%8E%B8-semantic-segmentation-%EC%B2%AB%EA%B1%B8%EC%9D%8C-4180367ec9cb

1편: Semantic Segmentation 첫걸음!

 

간단하게 생각하면 "물체 판별" 하는 기법! 이라고 이해하면 쉽다.

 

2. 데이터 수집

이 포스트를 쓰는 이유이다.

kaggle이나, VOC dataset 등 원한다면 데이터셋을 쉽게 구할 수 있다. ( 라벨링까지 완료된 )

이건 아주 나이스한 경우이고,,,, 나와같이 딱 맞는 데이터가 없다면 구글링으로 이미지를 다운받아야 한다.

구글에 본인이 찾고자 하는 이미지를 넉넉하게 다른 이름으로 저장 을 통해서 수집하자.

파일명은 0001~0xxx 이런식으로 만들어주는게 관리하기 쉽다.

 

3. 데이터 라벨링 

커스텀 데이터셋을 위해선 라벨링작업이 필수적이다.

다행히 라벨링 툴은 git에서 pull 받아 사용할 수 있으며, 내가 사용한 툴은 labelme라는 툴이다.

https://github.com/wkentaro/labelme

GitHub - wkentaro/labelme: Image Polygonal Annotation with Python (polygon, rectangle, circle, line, point and image-level flag

풀 받아서 README에 적힌대로 따라가면, 쉽게 사용이 가능하고, 

수집한 이미지 데이터 폴더를 읽거나, 파일을 읽으면 아래와 같이 로드된다.

그리고 객체를 한 땀 한 땀 클릭해서 폴리곤을 그리자. 보면 새 한마리에 폴리곤을 그리고 label명을 bird로 설정했다.

 

4. VOC 변환

우리의 친절한 labelme씨는 폴리곤을 그리면 동일한 파일명의 json이 생성해준다.

semantic segmentation을 하기 위해선, 객체에 해당하는 부분만 그려져있는 PNG 파일이 필요한데,

labelme에서 이 기능을 사용할 수 있다.

labelme-master/examples/semantic_segmentation/labelme2voc.py 

인데, 우선 주요 파라미터부터 숙지하자.

positional arguments:
  input_dir        input annotated directory
  output_dir       output dataset directory

optional arguments:
  -h, --help       show this help message and exit
  --labels LABELS  labels file (default: None)
  --noviz          no visualization (default: False)

input_dir, output_dir이 필수적이고, labels가 optional 인데, 이 부분까지 추가해서 만들어보자.

input_dir : 라벨링이 완료된 디렉터리명

output_dir : VOC format의 데이터셋이 담길 디렉터리명

labels : 지정한 라벨들 ( _backgroud_,__ignore__, 이 필수적이며, 사용자가 추가한 라벨들은 이 다음에 적어준다 )

ex) 
#labels.txt
_background_,
__ignore__,
your_label1,
your_label2,
,,,

이렇게 만들었다면, VOC 포멧으로 데이터를 변환하자.

labelme 내부의 디렉터리 구조이고, labelme 디렉터리에서 실행시켰다.

labelme와 동일한 depth에 crawling_image 디렉터리에 이미지와 라벨링한 json이 포함되어있다.

 

 

python examples/semantic_segmentation/labelme2voc.py ../crawling_image ./transfer --labels labels.txt

정상적으로 완료된다면, 아래와 같은 폴더가 생성된다.

원본 이미지, npy, 라벨링된 영역 이미지, 반투명(?) 된 이미지 등이 저장되고, 이 데이터셋으로 학습을 진행하면 된다.

 

5. 학습

사실 이 부분에서 시간을 많이 썼는데, kaggle에서 공개된 코드를 수정해서 진행했다.

https://www.kaggle.com/code/balraj98/unet-resnet50-frontend-road-segmentation-pytorch/notebook

UNet (ResNet50 frontend) Road Segmentation PyTorch

 

아래와 같이 구성하면 되는데, x_test_dir, x_val_dir을 적절하게 나눈 다음 추가해주면 된다.

아직 수집과 라벨링중이라 train_dir밖에 없다,,, ㅠ

BASE_DIR = '../transfer/'

x_train_dir = os.path.join(BASE_DIR, 'JPEGImages')
y_train_dir = os.path.join(BASE_DIR, 'SegmentationClassPNG')

그 외의 부분들은 본인의 상황에 맞게 수정하면 될 듯 하다.