Yolo Segmentation 커스텀모델 Training, JsonToYolo 포함

Json파일을 Yolo txt로 변환하여, Yolo Segmentation 모델을 만드는 방식을 정리해보려고 한다.

 

보통 Yolo 라벨링은 LabelImg라는 라벨링 툴을 사용하지만,

LabelImg는 직사각형만을 지원하기 때문에 Segmentation 라벨링으로는 적합하지 않다.

 

그래서 Labelme라는 프로그램을 통해서 다각형 라벨링을 하고,

JsonToYolo 변환을 하여 Yolo Segmentation 커스텀 모델을 만들었다.

 

0. Yolo Segmentation

YOLO Segmentation은 YOLO(You Only Look Once)의 물체 검출(Object Detection) 능력에
픽셀 단위의 세분화(segmentation) 기능을 추가한 모델로,
한 번의 추론으로 물체 탐지와 동시에 물체의 정확한 형태를 마스크(mask)로 추출한다.

• 일반 YOLO는 물체 위치를 박스(box) 로 탐지한다.
• YOLO Segmentation은 정확한 형태(픽셀 수준의 마스크) 까지 추가로 예측한다.

모델의 구성은 아래와 같다.

YOLO Segmentation 모델
    │
    ├── Backbone (특징 추출기)
    │      └── CSPDarknet53 등
    │
    ├── Neck (특징 강화 및 융합)
    │      └── PANet, FPN
    │
    └── Head (검출 + 세그먼테이션 예측)
           └── Detection + Mask Prediction

 

1. Segmentation Yolo txt 형식

class_id x1 y1 x2 y2 x3 y3 ... xn yn
class_id x1 y1 x2 y2 x3 y3 ... xn yn
class_id x1 y1 x2 y2 x3 y3 ... xn yn

 

최소 3개 이상의 점으로 구성되어야 하며, 0과 1 사이의 값으로 정규화되어야 한다.

 

2. JSONToYOLO

대략적인 흐름만 코드로 보여준다면, 아래와 같다.

for idx, pt in enumerate(tpoints):
    if idx == 0:
        txt_outfile.write(cls) # class_id 입력
    x, y = pt
    bb = convert_coor((width, height), [x, y])
    txt_outfile.write(" " + " ".join([str(a) for a in bb])) # 좌표 입력

txt_outfile.write("\n")

여기서 convert_coor 함수를 통해, 0과 1 사이의 값으로 정규화한다.

 

이렇게 Yolo txt로 변환한다면, 데이터 준비는 끝났다.

 

3. 환경 설정

가상환경을 만들고, 아래 라이브러리들을 설치하면 된다.

ultralytics
torch # 자신의 cuda 버전에 맞춰서 설치
torchaudio # 자신의 cuda 버전에 맞춰서 설치
torchvision # 자신의 cuda 버전에 맞춰서 설치

 

torch, torchaudio, torchvision 3가지는 pytorch에서 자신의 cuda 버전에 맞는 라이브러리를 찾아서 직접 설치해야한다.

https://pytorch.org/get-started/locally/

Get Started

본인의 환경에 맞춰 명령어를 터미널에 입력하면 된다.

혹시 CUDA 11.8 이전 버전이라면, 

https://pytorch.org/get-started/previous-versions/

Previous PyTorch Versions

여기서 찾아서 명령어를 입력하면 된다.

 

4. Yaml 파일 생성

# .yaml 파일 생성
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../yolo/data/datasets # dataset root dir
train: carplate_segment/images  # train images (relative to 'path') 128 images
val: carplate_segment/validation  # val images (relative to 'path') 128 images
test: carplate_segment/test # test images (optional)


# Classes
names:
  0: None
  1: car
  2: truck
  3: bus
  4: motorcycle
  5: bicycle
  6: human
  7: helmet
  8: kickboard
  9: plate

 

각 데이터 path와 클래스를 정리하여 yaml 파일을 만들면 된다.

 

5. Train

자 그럼 모든 준비는 끝났고, 가상환경을 활성화하여 명령어를 입력하면 된다.

# epochs : 총 학습할 에폭 수, imgsz : 이미지 사이즈, patience : 얼리스타핑을 결정하는데 사용되는 파라미터
yolo segment train model=yolov8l-seg.pt data=/.../dataset-seg.yaml
yolo task=segment mode=train model=yolov8l-seg.pt data=mydata.yaml epochs=100
yolo segment train model=yolov8l-seg.pt data=/.../dataset-seg.yaml imgsz=512
yolo segment train model=yolov8l-seg.pt data=/.../dataset-seg.yaml patience=5

 

model에는 yolo segmentation 모델의 pre-trained된 파일을 입력하면 된다.

 

모델의 파일명의 의미는 다음과 같다.

yolov8   +   l   +  seg    + .pt
└───┬───┘   └┬┘   └──┬───┘   └─┬─┘
 모델 버전    크기    작업유형   파일형식
 (YOLOv8)  (large) (segmentation) (.pt: PyTorch weights)

 

모델의 크기는

n : nano 

s : small

m : medium

l : large

x : extra-large

이다. 

 

만들고자 하는 모델의 성능에 맞는 pretrained 모델과 파라미터들을 설정하여, 명령어를 입력하면 된다.

 

추가적인 train 세팅은 아래 공식문서에서 확인하면 된다.

https://docs.ultralytics.com/modes/train/#train-settings

Train

 

validation 결과 출력

 

 

만약 pt모델을 다른 형식으로 변환하여 export하고 싶다면,

yolo task=detect mode=export model="runs/detect/train5/weights/best.pt" format="onnx"
yolo task=detect mode=export model="runs/detect/train5/weights/best.pt" format="engine"

 

원하는 형식에 맞춰 명령어를 입력해주면 된다.

https://docs.ultralytics.com/ko/modes/export/#export-formats

내보내기