Faster R-CNN을 이용한 여드름 탐지 - 데이터 전처리

Emile Perron on Unsplash

 

이전 글 'Roboflow를 이용한 data annotation' 에서는 여드름 이미지 데이터를 준비하고, bounding box annotation을 작성하여 데이터셋을 준비했습니다. 하지만, 데이터셋을 준비했다고 해서 곧바로 학습을 시작할 수 있는 것은 아닙니다. 데이터를 모델이 수용할 수 있는 구조로 변화시키고, 데이터를 증강하는 등의 전처리가 필요합니다. 이번 글에서는 코드 위주로, 학습 이전의 전처리 단계들을 진행해 보겠습니다. 

 

 

 

 

 

---

 

 

 

 

 

 

0. Import

import numpy as np
import pandas as pd
from PIL import Image
import torch
import os
import torchvision
import torchvision.transforms as T
import torch.optim as optim
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from torchvision import models
from torch.utils.data import DataLoader, Dataset
from torchvision.io import read_image
import random

from torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FasterRCNN_ResNet50_FPN_Weights
from sklearn.model_selection import train_test_split
from torchvision.models.detection.rpn import AnchorGenerator

필요한 라이브러리들과 함수들을 import 합니다.

 

 

 

 

 

1. Dataframe 준비

1-1. converter 함수

데이터셋 활용을 용이하게 하기 위해, 데이터셋에 대한 정보들 (이미지 이름, 확장자, bounding box)을 하나의 dataframe에 정리합니다. 이때 .txt 형식으로 저장되어 있는 bounding box 데이터를 모델이 읽어낼 수 있는 형태로 변환하는 과정이 필요합니다. 이를 converter 함수로 정의해 봅시다. 

def converter(labels_path,img):
  img_w, img_h = img.size
  d = {'boxes' : [], 'labels' : []}

  with open(labels_path) as f:
    lines_txt = f.readlines()
    for lines in lines_txt:
      line = []
      line.append([int(lines.split()[0])]+[round(float(el),5) for el in lines.split()[1:]])
      line = line[0]

      d['boxes'].append([(line[1]-line[3]/2)*img_w,(line[2]-line[4]/2)*img_h,(line[1]+line[3]/2)*img_w,(line[2]+line[4]/2)*img_h])
      d['labels'].append(line[0])

  return d

converter 함수는 label이 저장되어 있는 경로(label_path)와, image를 입력받습니다. 입력받은 경로의 .txt 파일에는 bounding box의 (center x, center y, width, height) 좌표가 (0,1) scale로 저장되어 있습니다. Faster R-CNN이 요구하는 box의 좌표 형식은 (x1, y1, x2, y2) 이므로, 해당 조건에 맞게 좌표를 변환해 줘야 합니다. 

 

또한, Faster R-CNN이 요구하는 target의 형식이 list of dictionary 이기 때문에, 'boxes'와 'labels'를 key로 갖는 dictionary를 산출합니다. 

 

 

1-2. dataframe 작성

#train,test dataframe

train_df = pd.DataFrame()

%cd '/content/drive/MyDrive/projects/acne_detection/data/train/images'
names = []
extentions = []
for i in os.listdir():
  name, ext = os.path.splitext(i)
  names.append(name)
  extentions.append(ext)
train_df['names'] = names
train_df['extentions'] = extentions

targets = []
for idx,name in enumerate(names):
  imgname = names[idx] + extentions[idx]
  img = Image.open(os.path.join('/content/drive/MyDrive/projects/acne_detection/data/train/images',imgname))
  targets.append(converter(os.path.join('/content/drive/MyDrive/projects/acne_detection/data/train/labels',name+'.txt'),img))
train_df['targets'] = targets

test_df = pd.DataFrame()

%cd '/content/drive/MyDrive/projects/acne_detection/data/test'
names = []
extentions = []
for i in os.listdir():
  name, ext = os.path.splitext(i)
  names.append(name)
  extentions.append(ext)
test_df['names'] = names
test_df['extentions'] = extentions

이제, 데이터셋에 대한 정보 (image name, extentions, target)을 담고 있는 dataframe을 작성합니다. 이때, test dataframe에 대해서는 target 정보를 작성하지 않습니다. 

image의 name, extention을 분리하기 위해 os.path.splitext() 함수를 사용했고, target을 기록하기 위해 앞서 정의한 converter 함수를 사용했습니다. 

 

Dataframe은 다음과 같이 만들어졌습니다. 

 

 

 

 

 

 

---

 

 

 

 

 

2. Dataset, Dataloader

2-1. custom dataset

위에서 작성한 Dataframe을 바탕으로, 학습에 사용할 custom dataset을 정의합니다.

class acnedataset(Dataset):
  def __init__(self,root,df,indices):
    self.root = root
    self.names = df['names'].values
    self.exts = df['extentions'].values
    self.target = None
    self.indices = indices
    if 'targets' in df.keys():
      self.target = df['targets'].values

  def __len__(self):
    return len(self.indices)

  def __getitem__(self,idx):
    idx = self.indices[idx]
    img_path = os.path.join(self.root,self.names[idx]+self.exts[idx])
    image = read_image(img_path)
    if self.target is not None:
      target = self.target[idx]
      return image, target
    else:
      return image

우리의 acnedataset은 root (source dataset의 위치), dataframe, indices를 입력으로 받습니다. indices의 역할은 train data와 validation data를 분리하는 기능을 합니다. 

 

pytorch의 custom dataset을 작성하는 규칙에 맞춰, __init__()과 __len()__, 그리고 __getitem()__ 함수를 구현해 줍니다. 

 

 

2-2. collate_fn

def my_collate(batch):
  img_resize = T.Resize(350,antialias=None)
  img_list = []
  target_list = []
  train = False
  if len(batch[0]) == 2: train = True

  #resize
  for data in batch:
    #train
    if train:
      img,target = data
      resized_target = {}
      resized_target['labels'] = torch.tensor(target['labels'])
      resized_target['boxes'] = torch.zeros(torch.tensor(target['boxes']).shape)

      _,h,w = img.shape
      img = img_resize(img)
      _,re_h,re_w = img.shape
      resize_ratio = re_h/h
      img_list.append(img/255)

      for idx,coefficients in enumerate(target['boxes']):
        resized_target['boxes'][idx] = torch.tensor([coef*resize_ratio for coef in coefficients])
      target_list.append(resized_target)

    #test
    else:
      img = img_resize(data)
      img_list.append(img/255)

  if not train: return img_list

  return img_list,target_list

custom dataset을 Dataloader에 전달하기 앞서, 데이터를 전처리하는 custom collate_fn 함수를 정의합니다. pytorch에서 디폴트로 사용되는 collate_fn을 사용하지 않는 이유는 다음과 같습니다. 

 

• batch 내의 image size가 서로 다르기 때문입니다. 디폴트 collate_fn에서는 이것을 허용하지 않아서 error를 발생시킵니다. (Faster R-CNN 모델은 이를 허용합니다.)
• 이미지를 350으로 Resize할 때, bounding box도 비율에 맞게 조정해야 하기 때문입니다. (이 과정은 augmentation 단계에서 수행해도 될 것 같습니다.)

 

 

2-3. dataloaders

#dataloader
train_inds, val_inds = train_test_split(range(len(train_df['names'])), test_size = 30)
train_root = '/content/drive/MyDrive/projects/acne_detection/data/train/images'
train_dataset = acnedataset(train_root,train_df,train_inds)
train_dataloader = DataLoader(train_dataset, batch_size = 4, shuffle = True, num_workers = 2, collate_fn = my_collate)

val_dataset = acnedataset(train_root,train_df,val_inds)
val_dataloader = DataLoader(val_dataset, batch_size = 4, shuffle = True, num_workers = 2, collate_fn = my_collate)

test_root = '/content/drive/MyDrive/projects/acne_detection/data/test'
test_inds = list(range(len(test_df['names'])))
test_dataset = acnedataset(test_root,test_df,test_inds)
test_dataloader = DataLoader(test_dataset, batch_size = 4, shuffle = True, num_workers = 2, collate_fn = my_collate, drop_last = True)

sklearn 라이브러리의 train_test_split 함수를 이용해 train indices와 validation indices를 분리합니다. 총 135개의 데이터 중, 30개의 데이터를 validation set으로 설정하였으며, 모든 dataloader의 batch size는 4로 설정하였습니다. 

 

 

2-4. Visualize

img, target = next(iter(train_dataloader))
plt.figure(figsize = (12,12))
for it, (img,target) in enumerate(zip(img[:4],target[:4])):
  ax = plt.subplot(2,2,it+1)
  ax.imshow(img.permute(1,2,0))
  plt.axis('off')

  for box,label in zip(target['boxes'],target['labels']):
    if label == 0:
      rect = patches.Rectangle((box[0],box[1]),box[2]-box[0],box[3]-box[1],linewidth=1, edgecolor= 'red', facecolor='none')
    else:
      rect = patches.Rectangle((box[0],box[1]),box[2]-box[0],box[3]-box[1],linewidth=1, edgecolor='purple', facecolor='none')
    ax.add_patch(rect)

제작한 dataloader가 의도한 대로 작동하는지 확인하기 위해, matplotlib을 이용하여 visualize해 보았습니다. bounding box를 표현하는 데에는 matplotlib의 patches 라이브러리를 사용하였습니다. 시각화된 1 batch의 데이터셋은 아래와 같습니다. 

 

 

 

 

 

 

---

 

 

 

 

 

마치며

이번 글은 본격적인 모델 학습을 위한 전처리 등의 준비 단계에 관한 글이었습니다. 특별히 어려운 점은 없었지만, 이 모든 과정을 처음부터 직접 작성해 본 것은 처음이라 꽤나 큰 공부가 되었던 것 같습니다. 특히 collate_fn이나 dataframe을 제작하는 과정에서 다양한 형태의 데이터를 원하는 형태로 가공하는 연습이 많이 되었습니다. 

 

다음 글에서는 본격적인 모델 학습을 진행해 보고, 한정된 데이터에서 가장 준수한 성과를 얻어내는 모델 학습 방법을 연구해 보려고 합니다. 긴 글을 읽어 주셔서 감사합니다!