13. Image 데이터 확장(augmentation) - pineland/object-tracking GitHub Wiki

여기에서는 Python 라이브러리 imgaug를 사용하여 Image augmentation하는 것을 설명한다.

1. imgaug 설치

pip3 install imgaug

2. Image와 Annotation 정보 읽기

다음은 Python에서 xml모듈을 이용하여 xml 구조를 파싱하는 코드이다. 한 이미지에 대한 Annotation(PASCAL VOC)으로부터 해당하는 이미지 filename과 이미지 내의 모든 Bounding-Box 정보를 읽어서 반환한다.

import xml.etree.ElementTree as ET

def read_anntation(xml_file: str):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    bounding_box_list = []

    file_name = root.find('filename').text
    for obj in root.iter('object'):

        object_label = obj.find("name").text
        for box in obj.findall("bndbox"):
            x_min = int(box.find("xmin").text)
            y_min = int(box.find("ymin").text)
            x_max = int(box.find("xmax").text)
            y_max = int(box.find("ymax").text)

        bounding_box = [object_label, x_min, y_min, x_max, y_max]
        bounding_box_list.append(bounding_box)

    return bounding_box_list, file_name

read_train_dataset 함수는 image 파일들과 각 이미지의 annotation xml 파일들이 함께 있는 디렉토리의 경로를 인자로 전달했을 때, 4차원(N:이미지 수, W:이미지 너비, H:이미지 높이, D:RGB) nparray로 변환된 이미지들과 3개의 항(bounding-box 리스트, xml 파일명, 관련 이미지 파일명)을 갖는 Annotation 튜플의 리스트를 반환한다.

from os import listdir
import cv2
import numpy as np

def read_train_dataset(dir):
    images = []
    annotations = []

    for file in listdir(dir):
        if 'jpg' in file.lower() or 'png' in file.lower():
            images.append(cv2.imread(dir + file, 1))
            annotation_file = file.replace(file.split('.')[-1], 'xml')
            bounding_box_list, file_name = read_anntation(dir + annotation_file)
            annotations.append((bounding_box_list, annotation_file, file_name))

    images = np.array(images)

    return images, annotations

3. Image augmentation

A simple example의 코드를 아래와 같이 수정하여, 가지고 있는 이미지와 해당 Annotation 정보에 적용한다.

import imgaug as ia
from imgaug import augmenters as iaa
from files import *

ia.seed(1)

dir = 'train/'
images, annotations = read_train_dataset(dir)

for idx in range(len(images)):
    image = images[idx]
    boxes = annotations[idx][0]

    ia_bounding_boxes = []
    for box in boxes:
        ia_bounding_boxes.append(ia.BoundingBox(x1=box[1], y1=box[2], x2=box[3], y2=box[4]))
    bbs = ia.BoundingBoxesOnImage(ia_bounding_boxes, shape=image.shape)

    seq = iaa.Sequential([         # 이미지를 어떻게 augment하는지를 결정
        iaa.Multiply((1.2, 1.5)),
        iaa.Affine(
            translate_px={"x": 40, "y": 60},
            scale=(0.5, 0.7),
            rotate=45
        )
    ])

    seq_det = seq.to_deterministic()

    image_aug = seq_det.augment_images([image])[0]
    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]

    for i in range(len(bbs.bounding_boxes)):
        before = bbs.bounding_boxes[i]
        after = bbs_aug.bounding_boxes[i]
        print("BB %d: (%.4f, %.4f, %.4f, %.4f) -> (%.4f, %.4f, %.4f, %.4f)" % (
            i,
            before.x1, before.y1, before.x2, before.y2,
            after.x1, after.y1, after.x2, after.y2)
        )

    image_before = bbs.draw_on_image(image, thickness=20)
    image_after = bbs_aug.draw_on_image(image_aug, thickness=20, color=[0, 0, 255])

    cv2.imshow('image_before', cv2.resize(image_before, (380, 640)))
    cv2.imshow('image_after', cv2.resize(image_after, (380, 640)))

    ret = cv2.waitKey(0)

    if ret == 27:  # ESC키를 누르면,
        break
    else:
        continue

콘솔 출력결과

BB 0: (615.0000, 815.0000, 1573.0000, 1073.0000) -> (905.7255, 1211.4296, 1468.9925, 1363.1236)
BB 1: (466.0000, 1062.0000, 1529.0000, 1304.0000) -> (818.1193, 1356.6560, 1443.1222, 1498.9427)
BB 2: (1006.0000, 1287.0000, 1594.0000, 1434.0000) -> (1135.6184, 1488.9473, 1481.3397, 1575.3776)
BB 3: (471.0000, 1285.0000, 969.0000, 1446.0000) -> (821.0591, 1487.7714, 1113.8638, 1582.4332)
BB 4: (632.0000, 2541.0000, 1162.0000, 2725.0000) -> (915.7209, 2226.2509, 1227.3404

원본 이미지를 어떠한 방식으로 augmentation 할 것인지는 iaa.Sequential 부분을 수정함으로써 결정할 수 있다. 보다 다양한 옵션들은 Examples: Basics에서 확인 할 수 있다.

seq = iaa.Sequential([
    iaa.Multiply((1.2, 1.5)),
    iaa.Affine(
        translate_px={"x": 40, "y": 60},
        scale=(0.5, 0.7),
        rotate=45  # 이미지 회전 추가
    )
])

4. images and annotations 저장

새로운 이미지의 Annotation 정보를 담는 xml 문서는 pascal-voc-writer 모듈을 기반으로 생성한다.
pascal-voc-writer 모듈은 다음과 같이 pip를 이용하여 설치할 수 있다.

pip3 install pascal-voc-writer

다음은 imgaug와 pascal-voc-writer를 이용한 이미지 생성 및 저장하는 코드 예제이다. 동일 디렉토리내에 'after_'가 붙은 이미지와 annotation 파일을 생성한다.

import imgaug as ia from imgaug import augmenters as iaa from files import * from pascal_voc_writer import Writer

ia.seed(1)

dir = 'train/' images, annotations = read_train_dataset(dir)

for idx in range(len(images)): image = images[idx] boxes = annotations[idx][0]

  ia_bounding_boxes = []
  for box in boxes:
      ia_bounding_boxes.append(ia.BoundingBox(x1=box[1], y1=box[2], x2=box[3], y2=box[4]))

  bbs = ia.BoundingBoxesOnImage(ia_bounding_boxes, shape=image.shape)

  seq = iaa.Sequential([
      iaa.Multiply((1.2, 1.5)),
      iaa.Affine(
          translate_px={"x": 40, "y": 60},
          scale=(0.5, 0.7),
          rotate=45
      )
  ])

  seq_det = seq.to_deterministic()

  image_aug = seq_det.augment_images([image])[0]
  bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]

  new_image_file = dir + 'after_' + annotations[idx][2]
  cv2.imwrite(new_image_file, image_aug)

  h, w = np.shape(image_aug)[0:2]
  voc_writer = Writer(new_image_file, w, h)

  for i in range(len(bbs_aug.bounding_boxes)):
      bb_box = bbs_aug.bounding_boxes[i]
      voc_writer.addObject(boxes[i][0], int(bb_box.x1), int(bb_box.y1), int(bb_box.x2), int(bb_box.y2))

  voc_writer.save(dir + 'after_' + annotations[idx][1])