Faster-RCNN

목적

패션 데이터셋을 활용하여 모자, 상의, 하의 등의 객체를 탐지하는 코드를 실행합니다.

데이터 로드

로드하려는 csv 파일 안에 패션에 대한 이미지 URL, 종류, 영역 좌표의 정보가 포함되어 있습니다.

이 정보를 로드하여 시각화하고 확인해봅니다.

import os
import json
import pandas as pd
import numpy as np
from urllib.parse import *
from requests.utils import requote_uri
import requests
from tqdm.notebook import tqdm

BASE_DIR = os.getcwd()
annotations_path = os.path.join(BASE_DIR, 'annotations')
images_path = os.path.join(BASE_DIR, 'images')
test_annotations_path = os.path.join(BASE_DIR, 'test_annotations')
test_images_path = os.path.join(BASE_DIR, 'test_images')

dirs = [annotations_path, images_path, test_annotations_path, test_images_path]
list(map(lambda x : os.makedirs(x, exist_ok=True), dirs))

origin_labels = pd.read_csv('sample.csv')
origin_labels.head()

첫번째 라인에는 모자 두번째 라인에는 바지에 관한 정보가 들어있는 것을 확인해볼 수 있습니다.

	Unamed	filename	width	height	class	xmin	ymin	xmax	ymax
0	1	https://sample.com/image_source/test1.jpg	591	1137	Hat	247	61	474	284
1	2	https://sample.com/image_source/test2.jpg	591	1137	Pants	219	367	590	848

이미지 종류에 대해서만 그룹핑을 해보도록 합니다. 모자, 바지 등 8종류가 있습니다.

classes = origin_labels.groupby('class')
class_indexes = list(classes.groups.keys())
print(class_indexes)
class_indexes.index('Pants')

['Hat', 'Pants', 'One Piece', 'T-Shirt', 'Shoes', 'Jaket', 'Skirt', 'Coat']

1

 

데이터 전처리

하나의 이미지에 여러가지 영역과 이미지 종류가 있기때문에 이미지 URL 정보로 그룹핑을 합니다.

def get_bbox(group):
    new_bbox = []
    for item in group.values:
        new_bbox.append((class_indexes.index(item[4]), [float(item[5]), float(item[6]), float(item[7]), float(item[8])]))
    return np.asarray(new_bbox)

labels는 이미지 URL을 키로하여 dictionary 형태로 되어 있고 값은

[(이미지 종류, 박스 영역), (이미지 종류, 박스 영역)] 형태의 데이터 구조로 되어있는 것을 알 수 있습니다.

labels = origin_labels.groupby('filename').apply(get_bbox)
print(labels.keys())
print(labels[labels.keys()[0]])

['https://sample.com/image_source/test1.jpg']

[(0, [247, 61, 474, 284]), (1, [0, 10, 50, 100])]

 

전체 데이터에서 생성할 훈련 데이터, 평가 데이터를 8:2의 비율로 나눕니다.

total_values = np.unique(labels.index.values)
total_len = len(total_values)

train_ratio = int(total_len * 0.8)
train_image_ids = total_values[0:train_ratio]
val_image_ids = total_values[train_ratio:]

이미지를 읽어서 전부 메모리에 가지고 있으면 부하가 커지기때문에 전부 다운로드화하여 파일화 시킵니다.

def parse_url(url):
    p = urlparse(url, 'http')
    if p.netloc:
        netloc = p.netloc
        path = p.path
    else:
        netloc = p.path
        path = ''
    p = p._replace(netloc=netloc, path=path)
    return p.geturl()

def download_image_from_url(image_url, source='train'):  
    try:
        target_url = parse_url(image_url)      
        response = requests.get(requote_uri(target_url), stream=True)
        filename = image_url.split('/')[-1] 
        if source == 'train':
            filepath = os.path.join(images_path, filename)
        else:
            filepath = os.path.join(test_images_path, filename)
        
        if os.path.exists(filepath):
            return True
            
        if response.status_code == 200:
            with open(filepath, 'wb') as f:
                f.write(response.content)
                return True        
    except Exception as e:
        pass
    return False

def write_annotation(image_url, source='train'):
    basename = image_url.split('/')[-1] 
    basename, _ = os.path.splitext(basename)
    filename = basename + '.json'
    
    if source == 'train':
        filepath = os.path.join(annotations_path, filename)
    else:
        filepath = os.path.join(test_annotations_path, filename)            
        
    if os.path.exists(filepath):
        return True        
        
    annotations = labels[image_url]    
    new_annotations = []
    for item in annotations:   
        new_annotations.append({
            'class_index': item[0],
            'annotation' : item[1]
        })
    with open(filepath, 'w') as f:
        f.write(json.dumps(new_annotations))

훈련 및 평가 데이터 파일을 생성합니다.

make_dataset(train_image_ids, source='train')

make_dataset(val_image_ids, source='test')

데이터 파이프 라인

전부 데이터를 가져와서 메모리에 갖고 있기에는 부하가 있기때문에 훈련시 배치 사이즈만큼 가져와서 가공할 수 있도록 데이터 파이프 라인을 만들어 보기로 합니다.

import os
import numpy as np
import matplotlib.patches as patches
import matplotlib.pyplot as plt
from bs4 import BeautifulSoup
from PIL import Image
import torchvision
from torchvision import transforms, datasets, models
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
import time
import torch

def generate_target(file): 
    with open(file, encoding='utf-8') as f:
        data = json.loads(f.read())
        boxes = []
        labels = []
        for item in data:
            boxes.append(item['annotation'])
            labels.append(item['class_index'])

        boxes = torch.as_tensor(boxes, dtype=torch.float32) 
        labels = torch.as_tensor(labels, dtype=torch.int64) 
        
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        
        return target

class Dataset(object):
    def __init__(self, transforms, path):
        self.transforms = transforms
        self.path = os.path.join(BASE_DIR, path)
        self.imgs = list(sorted(os.listdir(self.path)))

    def __getitem__(self, idx): #special method
        file_image = self.imgs[idx]
        file_label = self.imgs[idx][:-3] + 'json'
        img_path = os.path.join(self.path, file_image)
        
        if 'test' in self.path:
            label_path = os.path.join(BASE_DIR, "test_annotations/", file_label)
        else:
            label_path = os.path.join(BASE_DIR, "annotations/", file_label)

        img = Image.open(img_path).convert("RGB")
        target = generate_target(label_path)
        
        if self.transforms is not None:
            img = self.transforms(img)

        return img, target

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

data_transform = transforms.Compose([transforms.ToTensor()])

def collate_fn(batch):
    return tuple(zip(*batch))

훈련 및 평가 데이터 파이프라인을 생성합니다.

dataset = Dataset(data_transform, 'images')
test_dataset = Dataset(data_transform, 'test_images')
data_loader = torch.utils.data.DataLoader(dataset, batch_size=2, collate_fn=collate_fn)
test_data_loader = torch.utils.data.DataLoader(test_dataset, batch_size=5, collate_fn=collate_fn)

모델 로드

faster rcnn 모델을 로드합니다.

def get_model_instance_segmentation(num_classes):  
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model

model = get_model_instance_segmentation(len(classes))
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') 
model.to(device)

FasterRCNN(
  (transform): GeneralizedRCNNTransform(
      Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
      Resize(min_size=(800,), max_size=1333, mode='bilinear')
  )
  ... <생략>

torch.cuda.is_available()

True

num_epochs = 100
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
                                momentum=0.9, weight_decay=0.0005)

훈련을 시작합니다.

print('----------------------train start--------------------------')
for epoch in range(num_epochs):
    start = time.time()
    model.train()
    i = 0    
    epoch_loss = 0
    for imgs, annotations in data_loader:
        i += 1
        imgs = list(img.to(device) for img in imgs)
        annotations = [{k: v.to(device) for k, v in t.items()} for t in annotations]
        loss_dict = model(imgs, annotations) 
        losses = sum(loss for loss in loss_dict.values())        

        optimizer.zero_grad()
        losses.backward()
        optimizer.step() 
        epoch_loss += losses
    print(f'epoch : {epoch+1}, Loss : {epoch_loss}, time : {time.time() - start}')

epoch : 1, Loss : 246.2340850830078, time : 283.55497002601624

epoch : 2, Loss : 202.09490966796875, time : 284.3149654865265

epoch : 3, Loss : 167.4278564453125, time : 284.43429923057556

epoch : 4, Loss : 147.18411254882812, time : 284.58269715309143

epoch : 5, Loss : 135.09730529785156, time : 284.5432515144348

...<생략>

 

모델을 저장합니다.

torch.save(model.state_dict(), os.path.join(BASE_DIR, f'model_{num_epochs}.pt'))

예측

저장한 모델을 로드해봅니다.

model.load_state_dict(torch.load(f'model_{num_epochs}.pt'))

def make_prediction(model, img, threshold):
    model.eval()
    preds = model(img)
    for id in range(len(preds)) :
        idx_list = []
        for idx, score in enumerate(preds[id]['scores']) :
            if score > threshold :
                idx_list.append(idx)
        preds[id]['boxes'] = preds[id]['boxes'][idx_list]
        preds[id]['labels'] = preds[id]['labels'][idx_list]
        preds[id]['scores'] = preds[id]['scores'][idx_list]
    return preds

시각화 할 함수를 만듭니다.

cmap = plt.get_cmap('jet', len(class_indexes))
colors = [cmap(i) for i in np.linspace(0, 1, 50)]

def plot_image_from_output(img, annotation):    
    img = img.cpu().permute(1,2,0)
    
    fig,ax = plt.subplots(1)
    ax.imshow(img)    
    for idx in range(len(annotation["boxes"])):
        xmin, ymin, xmax, ymax = annotation["boxes"][idx].cpu()
        class_index = annotation['labels'][idx].cpu().numpy() 
        class_name = class_indexes[class_index]
        color = colors[class_index]        
        rect = patches.Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin), linewidth=1, edgecolor=color, facecolor='none', label=class_name)        
        ax.add_patch(rect)
        plt.text(xmin, ymin, s=class_name, 
            color='white', verticalalignment='top', bbox={'color': color, 'pad': 0})

    plt.show()

샘플링을 하기 위하여 평가 데이터셋에서 한번의 배치 사이즈만 로드해봅니다.

with torch.no_grad(): 
    for imgs, annotations in test_data_loader:
        imgs = list(img.to(device) for img in imgs)
        pred = make_prediction(model, imgs, 0.5)
        print(pred)
        break

실제 데이터와 예측한 결과를 비교해봅니다.

_idx = 0
print("Target : ", annotations[_idx]['labels'])
plot_image_from_output(imgs[_idx], annotations[_idx])
print("Prediction : ", pred[_idx]['labels'])
plot_image_from_output(imgs[_idx], pred[_idx])