Step by step object detection, STEP 1. 데이터셋 준비하기

---

객체 검출 학습 모델 구현을 단계별로 정리하는 포스트를 작성해보려고 한다. 

은근 자잘하게 신경쓸게 많아서 정리해두면 나중에 활용하기에 좋을 것 같다.

 

먼저 이 과정은 sgrvinod 깃허브의 'Deep Tutorials for PyTorch' 튜토리얼을 참고해서 진행되었다. 

내가 사용한 객체 검출 데이터셋은 AI HUB 의 '건강관리를 위한 음식 이미지 데이터' 이다. 

 

이 데이터셋에는 음식과 관련된 500여개 카테고리가 존재했는데, 컴퓨팅적 리소스를 고려하여, 과일과 채소 대분류 중 일부 카테고리만을 선정하여 객체 검출을 진행하기로 했다. 

그리하여 선정된 카테고리는 다음과 같다.

• 과일: 망고(1), 멜론(2), 딸기(3), 블루베리(4), 사과(5), 자몽(6), 자두(7), 천도복숭아(8), 청포도(9), 체리(10)
• 채소: 군고구마(11), 군밤(12), 고추(13), 치커리(14), 콜라비(15), 파프리카(16), 토마토(17), 표고버섯(18), 단호박(19), 피망(20)

 

STEP 1. 데이터셋 준비하기

 

1) 랜덤 샘플링

각 카테고리 별로 raw data 를 다운받아 보면 한 카테고리 안에도 이미지가 너무 많다. 그래서 학습 데이터로 400장, validation 데이터로 50장, 테스트 데이터로 50장을 뽑아서 활용할 것이다.

이를 위해 랜덤 샘플링을 500개 만큼 진행하였다.

모든 카테고리는 다음과 같은 구성으로 이미지와 라벨이 저장되어있다. 이곳 (source) 에서 500장을 랜덤 샘플링해서 복사하고 외부 폴더로 붙여넣기를 할 것이다. 따라서 외부에 블루베리라는 폴더를 만들고 하위 폴더로 images와 labels 를 만든다.

< 기존 데이터셋 - 블루베리 폴더 >

📦 [원천] 음식204_Tra
└─ 블루베리
   ├─ A220121XX_03592.jpg
   ├─ A220121XX_03593.jpg
   └─ (이하 생략)

📦 [라벨] 음식204_Tra
└─ 블루베리
   ├─ A220121XX_03592.json
   ├─ A220121XX_03593.json
   └─ (이하 생략)

< 새로 만든 폴더 >

📦 블루베리
├─ images
└─ labels

랜덤 샘플링하여 복사 붙여넣기를 하는 코드는 다음과 같다. (500_random_sampling.ipynb)

source_img = "/Users/jeongin/Desktop/장어구이 image"
destin_img = "/Users/jeongin/Desktop/장어구이/images"
source_json = "/Users/jeongin/Desktop/장어구이 json"
destin_json = "/Users/jeongin/Desktop/장어구이/labels"


file_list = os.listdir(source_img)
image_list = [f for f in file_list if f.endswith('.jpg')]

# 이미지 파일 500개 랜덤 샘플링
random_images = random.sample(image_list, 500)
print("random",random_images)

# 이미지 파일 복사 붙여넣기
copied_files = []
for image in random_images:
    copy_this = os.path.join(source_img, image)
    paste_here = os.path.join(destin_img, image)
    shutil.copy(copy_this, paste_here)
    copied_files.append(image)
print("copied", copied_files)

json_list = []

for img in copied_files:
    new = img.replace('.jpg', '.json')
    json_list.append(new)
print(json_list)

for json in json_list:
    copy_from = os.path.join(source_json, json)
    paste_to = os.path.join(destin_json, json)
    shutil.copy(copy_from, paste_to)​

 

 

이 작업을 완료하고 폴더를 우클릭해 '정보 가져오기'를 하면 501개의 파일이 들어있다고 뜨는데, 이 하나의 정체는 '.DS_Store' 라는 파일로 macOS 운영 체제에서 사용되는 숨겨진 시스템 파일이다. 이 파일은 특정 폴더의 디렉토리 구조와 폴더 내부의 파일 및 폴더의 레이아웃을 저장하는 데 사용되며 무시해도 된다. 

크기: 501개의 항목이 있다고 뜬다.

위의 과정을 선정한 20개의 카테고리에 대해 반복하며 총 500x20=10,000장의 사진과 라벨 파일이 저장된다.

랜덤으로 뽑은 500개의 사진과 라벨중 50, 50 개를 떼어내서 test, validation 에 나눠주면 된다.

이제 다시 새로운 폴더를 만들어 전체 카테고리에 대해 train, validation, test 셋으로 분리시킬 것이다. 

📦 all
├─ train (8000)
│  ├─ images
│  └─ labels
├─ validation (1000)
│  ├─ images
│  └─ labels
└─ test (1000)
   ├─ images
   └─ labels

 

2) 이미지 resize & json 파일 수정(label, resolution 키 추가)

모든 이미지가 같은 사이즈를 가지도록 리사이징 해준다. 리사이징 해상도는 300x300 이다.

원래 image 파일은 각기 다른 해상도를 가지고 있다. 이에 대한 정보를 resolution 으로 넣어준다.

또한 label_map 으로 매핑될 label 숫자를 넣어준다.

코드: 

import os
import random
import shutil
import json
import cv2
source_dir = "/Users/jeongin/Desktop/all/training/images"
destin_dir = "/Users/jeongin/Desktop/all/training/destin_images"

def resize_and_update_json(image, json_data, target_size=(300,300)):
    resized_image = cv2.resize(image, target_size)
    updated_json = []
    for obj in json_data:
        obj["label"] = 20 # 바꿔줄 부분
        updated_json.append(obj)
    updated_json = [{'label': obj['label'], **obj} for obj in updated_json]
    resolution = {
        'width': image.shape[1],
        'height': image.shape[0]
    }
    return resized_image, updated_json, resolution
    
    
    
for file in source_dir:
    # 이미지 파일 복사
    source_img = os.path.join(file)
    destin_img = os.path.join(file)
    
    # JSON 파일 복사
    json_file = os.path.splitext(file)[0] + ".json"
    source_json = os.path.join(source_dir, "labels", json_file)
    destin_json = os.path.join(destin_dir, "labels", json_file)
    
    # Load the image and JSON data
    image = cv2.imread(source_img)
    with open(source_json, 'r') as json_f:
        json_data = json.load(json_f)
    
    # Resize image and adjust bounding boxes
    resized_image, updated_json, resolution = resize_and_update_json(image, json_data)
    
    # Save the resized image and updated JSON data
    cv2.imwrite(destin_img, resized_image)
    
    # Group the updated object annotations and resolution into a single dictionary
    output_data = {
        'objects': updated_json,
        'resolution': resolution
    }
    
    with open(destin_json, 'w') as json_f:
        json.dump(output_data, json_f, indent=2)

 

업데이트된 json 파일:

 

3) 절대좌표로 변환해준다. 

 

원래 json 파일 안에는 중심좌표 (c_x, c_y)와 폭 (width), 높이(height) 가 상대좌표로 들어있다.

이것을 추출해서 절대좌표로 넣어줄 것이다. 절대좌표로 변환하는 코드는 아래와 같다. 

코드: 

import os
import json

from PIL import Image, ImageDraw
import torch
import torchvision.transforms.functional as FT

def convert(json_dir):

    for filename in os.listdir(json_dir):
        if filename.endswith(".json"):  

            json_file_path = os.path.join(json_dir, filename)  

            with open(json_file_path, 'r') as json_file:
                data = json.load(json_file)

            for item in data["objects"]:
                w = float(item["W"])
                h = float(item["H"])
                x, y = map(float, item["Point(x,y)"].split(','))

                resized_width = 300
                resized_height = 300

                x_min = (x - w / 2) * resized_width
                y_min = (y - h / 2) * resized_height
                x_max = (x + w / 2) * resized_width
                y_max = (y + h / 2) * resized_height

                new_xmin = x_min
                new_ymin = y_min
                new_xmax = x_max
                new_ymax = y_max

                print(new_xmin, new_ymin, new_xmax, new_ymax)

                item["boxes"] = [new_xmin, new_ymin, new_xmax, new_ymax]

            with open(json_file_path, 'w') as updated_json_file:
                json.dump(data, updated_json_file, indent=4)

            # 업데이트된 JSON 파일 경로 출력
            print(f"업데이트된 JSON 파일: {json_file_path}")
            
 # JSON 파일이 들어 있는 디렉토리 경로
json_dir = "/Users/jeongin/Desktop/all/train/labels"

# JSON 파일 변환 함수 호출
convert(json_dir)

최종 json 파일 형태

파란색 글씨 부분이 새롭게 업데이트된 부분이다.

---

4) train_images.json, train_objects.json, test_images.json, test_objects.json 만들기

 

학습을 할 때 데이터를 로드해와서 학습 함수에 입력해주려면 CustomDataset 클래스 정의와 DataLoader 가 필요하다. 

CustomDataset 에서 이미지와 레이블을 불러와야 한다. 이때 필요한 레이블은 json 파일 내부에서 "label" 값과 "boxes" 뿐이니, 해당하는 key들만 가져와서 새로운 'train_objects.json' 파일에 리스트로 저장할 것이다.

 

또한, 모든 이미지의 절대 경로도 새로운 'train_images.json' 파일을 만들어 리스트로 저장하려고 한다. 이렇게 하면 이미지의 경로와, 그 이미지에 대한 label&boxes 가 순서대로 저장되어 인덱스로 불러올 수 있을 것이다.

 

< train_images.json 과 train_objects.json 만들기 >

이하 모든 코드는 레포지토리의 utils.py 내부에 존재하는 코드이다.

 

1) label_map 지정

먼저 label map 을 정의해야 한다. 20개의 카테고리를 1부터 20까지의 레이블로 지정할 것이다.

레이블 0은 사진 속에 아무런 객체가 없는 경우로 'background' 클래스로 지정할 것이다. 

food_labels = ('mango', 'melon', 'strawberry', 'blueberry', 'apple', 'grapefruit',
               'plum', 'peach', 'grape', 'cherry', 'yam', 'chestnuts', 'pepper',
               'chicory', 'Kohlrabi', 'Paprika', 'Tomato', 'Mushroom', 'Pumpkin',
               'Pimento')
label_map = {k: v+1 for v, k in enumerate(food_labels)}

label_map['background'] = 0
print(len(label_map))
rev_label_map = {v: k for k, v in label_map.items()}  # inverse mapping

distinct_colors = ['#e6194b', '#3cb44b', '#ffe119', '#0082c8', '#f58231', 
		'#911eb4', '#46f0f0', '#f032e6', '#d2f53c', '#fabebe', 
                '#008080', '#000080', '#aa6e28', '#fffac8', '#800000', 
                '#aaffc3', '#808000', '#ffd8b1', '#e6beff', '#808080', '#FFFFFF']
label_color_map = {k: distinct_colors[i] for i, k in enumerate(label_map.keys())}

print("label map:",label_map) # 'mango':

21
label_map: {
            'mango': 1,
            'melon': 2,
            'strawberry': 3,
            'blueberry': 4,
            'apple': 5,
            'grapefruit': 6,
            'plum': 7,
            'peach': 8,
            'grape': 9,
            'cherry': 10,
            'yam': 11,
            'chestnuts': 12,
            'pepper': 13,
            'chicory': 14,
            'Kohlrabi': 15,
            'Paprika': 16,
            'Tomato': 17,
            'Mushroom': 18,
            'Pumpkin': 19,
            'Pimento': 20,
            'background': 0
            }

2) json 파일의 label과 boxes 키에 대한 value 가져오기

def parse_annotation(json_path):
    with open(json_path, 'r') as json_file:
        data = json.load(json_file)
    objects = data.get('objects', [])
    boxes = []
    labels = []
    for obj in objects:
        label = obj.get('label') + 1
        box = obj.get('boxes')
        if label is not None and box is not None:
            xmin, ymin, xmax, ymax = box
            boxes.append([xmin, ymin, xmax, ymax])
            labels.append(label)
    return {'labels': labels, 'boxes': boxes}

json 파일 안에는 레이블이 0~19 까지의 매핑대로 라벨링되어 있으니 train_objects.json 에 가져올 때는 label에 1을 더해서 가져와야 한다. boxes는 xmin, ymin, xmax, ymax 로 그대로 가져온다. 

이제 가져온 값들을 새로운 json 파일에 write 해서 저장하면 된다. 이미지의 경로를 가져오기 위해서 이미지의 모든 파일 제목이 담겨있는 train.txt 와 custom_path (레포지토리 절대 경로) 를 활용할 것이다. 

def create_data_lists(custom_path, output_folder):
    train_images = list()
    train_objects = list()
    n_objects = 0
    # Training data
    for path in [custom_path]:
        with open(os.path.join(path, 'train.txt')) as f:
            ids = f.read().splitlines()
        for id in ids:
            objects = parse_annotation(os.path.join(path, 'train/labels', id + '.json'))
            if len(objects['boxes']) == 0:
                continue
            n_objects += len(objects['labels'])  # Use objects['labels'] here
            train_objects.append(objects)
            train_images.append(os.path.join(path, 'train/images', id + '.jpg'))
    assert len(train_objects) == len(train_images)

    # Save to file
    with open(os.path.join(output_folder, 'train_images.json'), 'w') as j:
        json.dump(train_images, j)
    with open(os.path.join(output_folder, 'train_objects.json'), 'w') as j:
        json.dump(train_objects, j)
    print('\nThere are %d test images containing a total of %d objects. '
          'Files have been saved to %s.' % (
              len(train_images), n_objects, os.path.abspath(output_folder)))
    print("length of test_objects", len(train_objects))
    print("length of test_images", len(train_images))

이제 이 함수를 실행하기 위해 custom_path 에 train과 test 폴더를 포함하고 있는 현재 레포지토리의 경로를 넣어주고, output_folder에는 json 파일들을 저장할 경로를 넣어준다. 

from utils import create_data_lists

if __name__ == '__main__':
    create_data_lists(custom_path = '/Users/jeongin/PycharmProjects/21class_food_detection',
                      output_folder='/Users/jeongin/PycharmProjects/21class_food_detection/train/json_files')

test 데이터에 대해서도 동일하게 create_data_list 함수를 적용시켜 주면 된다.