AI 허브 데이터 YOLO format으로 변환 과정

예시로 사용한 AI 허브 데이터: 생활 폐기물 이미지

Original Json File(예시)

15_X001_C011_1006_0.json

{
  "FILE NAME": "15_X001_C011_1006_0.jpg",
  "COLLECTION METHOD": "제작자료",
  "FORM": "이미지",
  "DATE": "2020-10-06 14:40:02",
  "GPS": "37.5212093997222/126.715606899722",
  "ID CODE": "C011",
  "RESOLUTION": "1920*1440",
  "focus distance": "4.32mm",
  "exposure time": "1/278sec",
  "Aperture values": "2.52",
  "Sensitivity iso": "50",
  "exposure method": "normal program",
  "MAKE": "samsung",
  "Camera Model Name": "SM-G977N",
  "Software": "G977NKSU4CTG1",
  "File Size": "589983",
  "DAY/NIGHT": "주간",
  "PLACE": "실외",
  "PROJECT SORTING": "생활폐기물",
  "BoundingCount": "1",
  "Bounding": [
    {
      "CLASS": "비닐류",
      "DETAILS": "과자봉지",
      "DAMAGE": "일부훼손",
      "TRANSPARENCY": "불투명",
      "Color": "176/137/113",
      "Shape": "평면형",
      "Texture": "부드러움",
      "Object Size": "소",
      "Direction": "원거리",
      "Drawing": "BOX",
      "x1": "540",
      "y1": "765",
      "x2": "827",
      "y2": "930"
    }
  ]
}

1. 이 중 class와 bbox만을 필요로 하는 YOLO format에서는 사실상 나머지 정보는 필요가 없고, "Bounding" 안의 "CLASS"(혹은 "DETAILS"), "x1", "x2", "y1", "y2"가 필요할 뿐이다. class 중에서도 자신의 필요에 맞는 class만 남기고 삭제하는 작업이 필요하다.
2. YOLO format의 bbox가 nomalize된 값인 것을 고려하면 "RESOLUTION"(사진의 width와 height) 정보도 필요하나, 이는 이미지를 직접 불러와 읽어들이도록 한다.
3. "Drawing"이 "BOX"가 아니라 "POLYGON"인 것은 제외한다. 이 경우 bbox annotation을 제공하지 않은 것과 같으므로... 어쩔 수 없이 직접 bbox를 라벨링 할 수 밖에 없다.

1. labels와 images 각각 원하는 디렉토리에 옮기기

import os
import shutil

orig_label_dir = 'path/to/생활 폐기물 이미지/Validation/[V라벨링]라벨링데이터'
label_dir = 'path/to/labels'

for curDir, dirs, files in os.walk(orig_label_dir):
	for f in files:
    		shutil.move(os.path.join(curDir, f), os.path.join(label_dir, f))
            
orig_image_dir = 'path/to/생활 폐기물 이미지/Validation'
image_dir = 'path/to/images'

for curDir, dirs, files in os.walk(orig_image_dir):
    for f in files:
        shutil.move(os.path.join(curDir, f), os.path.join(image_dir, f))

• os.walk()의 출력
  • curDir: 파일이 존재하는 모든 디렉토리 경로(str, 파일명은 미포함)
  • dirs: 모든 디렉토리 명(list)
  • files: 모든 디렉토리별로 존재하는 모든 파일명(list)

2. 필요한 class들 정의하기

-> image는 AI hub에서 부분적으로는 골라서 받을 수 있지만, label은 내 task에 필요하지 않은 것들도 한번에 받도록 되어 있으므로, 맞는 image가 없는 label data는 먼저 삭제한다.

filenames = os.listdir(image_dir)
names = [name[:-4] for name in filenames]    # '파일명.jpg'에서 '파일명'만 남김

for curDir, dirs, files in os.walk(label_dir):
    for f in files:
        if f[:-5] not in names:    # '파일명.Json'에서 '파일명'으로 판단하고 없으면 삭제
            os.remove(os.path.join(curDir, f))

-> 남은 label들이 어떤 클래스를 갖고 있는지 먼저 추리고 파악한다.

import json
labels = []
for file in names:
   json_path = os.path.join(label_dir, file+'.Json')
   
   # json 파일 내용에 한글이 포함되어 있으므로 encoding 추가
   with open(json_path, 'r', encoding='utf-8') as f:    
       json_data = json.load(f)
   for obj in json_data['Bounding']:
       labels.append(obj['DETAILS'])
label_set = set(labels)    # "DETAILS"를 모두 저장한 다음 set 자료형으로 중복 제거
print(label_set)

-> task 정의에 따라 label_set들을 class를 0부터 시작하는 정수 번호로 지정해주는 함수를 정의한다.

def class_choose(cls):
	bottle = ['기타술병', '맥주병', '물병', '박카스병', '소주병', '음료수병', '페트병']
	can = ['맥주캔', '음료수캔', '통조림캔', '커피캔', '참기름캔', '스팸류']
    if cls in bottle:
        return 0
    elif cls in can:
        return 1
    else:
    	return 2

3. YOLO format annotation txt 파일 만들기

import json
import cv2
for file in names:
    json_path = os.path.join(label_dir, file+'.Json')
    img_path = os.path.join(image_dir, file+'.jpg')
    image = cv2.imread(img_path)
    h,w,_ = image.shape    # 이미지 처리에서는 이미지 shape를 height, width, channal 순으로 정의한다.
    with open(json_path, 'r', encoding='utf-8') as f:
        json_data = json.load(f)
    objs = []
    for obj in json_data['Bounding']:
        if obj['Drawing'] != 'BOX':    # bbox가 아닌 polygon 등 다른 방식의 annotation 제외
            break
        cls = obj['DETAILS']
        cls_num = class_choose(cls)
        x1 = int(obj['x1'])
        x2 = int(obj['x2'])
        y1 = int(obj['y1'])
        y2 = int(obj['y2'])
        center_x = (x1 + x2) / (2*w)
        center_y = (y1 + y2) / (2*h)
        width = (x2 - x1) / w
        height = (y2 - y1) / h    # YOLO format의 normalized bbox
        objs.append(f'{cls_num} {center_x} {center_y} {width} {height}\n')
    with open('path/to/yolo/labels/{file}.txt', 'w') as f:
        f.writelines(objs)

이후 적절히 annotation file과 image file의 경로를 맞추고, 라벨링 상태를 확인한 다음 YOLO를 훈련시키면 된다.