Segment Anything으로 labeling 없는 데이터를 fine tuning하기

 

 

 

어떤 task에 대해 labeling이 되어있지 않는(ground-truth가 없는) 데이터셋을 segmentation해야 할 일이 있었습니다.

이 때, 어떻게 segmentation해야 할지 애매한 부분이 있었습니다. 데이터셋은 충분히 존재하고 정작 labeling이 하나도 되어 있지 않아서 이것을 일일이 labeling하기에는 시간, 노동이 너무 비효율적이라고 판단하였습니다.

 

SAM(Segment Anything Model)이 zero-shot task로서 이 부분에서 강점이 있다고 판단하여 해당 데이터셋에 labeling을 해주는 것을 SAM에게 할당하였습니다. 아래 설명드리겠지만 pseudo labeling과 같은 방식으로 처리했습니다.

 

1. labeling이 되지 않은 데이터셋들을 SAM이 inference(예측)합니다.
2. inference한 결과물을 labeling을 합니다.
3. 이후 학습 시, 해당 ground-truth를 바탕으로 fine tuning을 진행하여 데이터셋에 대한 최적화 학습을 진행합니다.

 

어느 정도 최적화를 이루면 아래 이미지와 같은 양상을 나타낼 수 있습니다.

Predicted mask : pseudo labeling하여 fine tuning을 진행한 모델의 inference(예측)

Ground truth mask : 처음 SAM이 예측한 inference(정확한 ground-truth는 아니지만 이걸 우선 사용하여 pseudo labeling하게 됩니다.)

모델 최적화가 잘 되어서 처음의 SAM의 inference보다 실질적으로 왼쪽과 같이 해당 object 영역을 더 잘 잡고 있습니다.

 

 

당연히 모든 데이터셋에 최적화되는지는 장담할 수 없습니다. 그러나 충분히 먹힐 만한 이유가 이미 매우 큰 데이터셋으로 SAM이 학습되어 있어서 어느 정도 pseudo labeling에 대한 결과를 보증할 수 있기 때문입니다. 그리고 labeling할 시간이 부족한 상황에서 유용합니다.

 

fine tuning이 필요한 상황에서는 이렇게 진행하시면 되는데 그 전에 3가지 pretrained model(vit-h, vit-l, vit-b)에 대해 먼저 쓸 만한지 확인하고 진행하는 것을 추천합니다.(vit-h의 성능이 꽤 좋은 편입니다.)