Object Detection - YOLO v3 Pytorch 구현 (1)

[Machine Learning/기타] - Object Detection - YOLO v3

---

이번 포스트에서는 YOLO v3 Pytorch 구현을 해보고자 합니다.

Datasets

PASCAL VOC 2007 / 2012 데이터셋을 먼저 다운받습니다. YOLO 공식 홈페이지에 접속하면 PASCAL VOC 데이터 처리 방법에 대해 상세히 나와 있습니다. wget 명령어를 이용하여 VOC 데이터셋 tar 파일을 다운받습니다. 구글 코랩을 이용할 경우 '/content' 경로에 다운받아지며 리눅스 CLI로 동작시키기 위해서는 느낌표를 명령어 앞에 붙여야 합니다.

!wget https://pjreddie.com/media/files/VOCtrainval_11-May-2012.tar
!wget https://pjreddie.com/media/files/VOCtrainval_06-Nov-2007.tar
!wget https://pjreddie.com/media/files/VOCtest_06-Nov-2007.tar

import os
os.getcwd()

다운받은 이후 tar 명령어로 압축을 풀어줍니다. 압축을 풀어주면 '/content/VOCdevkit' 경로에 데이터가 담기게 됩니다.

!tar xf VOCtrainval_11-May-2012.tar
!tar xf VOCtrainval_06-Nov-2007.tar
!tar xf VOCtest_06-Nov-2007.tar

이제는 각 이미지 별로 "<object-class> <x> <y> <width> <height>" 라벨 정보가 담긴 (총 5개) 텍스트 파일을 (*.txt) 만들어주저야 합니다. "voc_label.py" 파일을 wget 명령어를 이용해 받고 이를 실행시키면 다음과 같이 각 파일의 경로가 담긴 텍스트 파일이 생성됩니다. (2007_train.txt, 2007_val.txt, 2012_train.txt, 2012_test.txt)

우리는 PASCAL VOC 2007 train, val / 2012 train, val 데이터셋을 훈련 데이터로 사용하고 2007 test 데이터셋을 validation 으로 사용할 것이므로 2007_train.txt, 2007_val.txt, 2012_train.txt, 2012_val.txt 파일의 내용을 train.txt 라는 하나의 파일로 합쳐줍니다.

with open('train.txt', 'w') as f:
    linelist = []
    ## 2007_train.txt
    with open('2007_train.txt', 'r') as f1:
        linelist.extend(f1.readlines())
    ## 2007_val.txt
    with open('2007_val.txt', 'r') as f2:
        linelist.extend(f2.readlines())
    ## 2012_train.txt
    with open('2012_train.txt', 'r') as f3:
        linelist.extend(f3.readlines())
    ## 2012_val.txt
    with open('2012_val.txt', 'r') as f4:
        linelist.extend(f4.readlines())
    for line in linelist:
        f.write(line)

---

다음으로 Pytorch 의 Dataset 모듈을 이용하여 데이터셋을 불러오는 클래스를 구현해줍니다.

import os
import cv2
import torch

import numpy as np

from torch.utils.data import Dataset, DataLoader

class ListDataset(Dataset):
    def __init__(self, list_path, img_size=416):
        with open(list_path, 'r') as file:
            self.img_files = file.readlines()
        self.label_files = [path.replace('JPEGImages', 'labels').replace('.jpg', '.txt') for path in self.img_files]
        self.img_shape = (img_size, img_size)
        self.max_objects = 50

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

    def __getitem__(self, index):

        #---------
        #  Image
        #---------

        img_path = self.img_files[index % len(self.img_files)].rstrip()
        img = cv2.imread(img_path)
        h, w, _ = img.shape
        dim_diff = np.abs(h - w)
        # Upper (left) and lower (right) padding
        pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
        # Determine padding
        pad = ((pad1, pad2), (0, 0), (0, 0)) if h <= w else ((0, 0), (pad1, pad2), (0, 0))
        # Add padding
        pad_img = np.pad(img, pad, 'constant', constant_values=128)

        padded_h, padded_w, _ = pad_img.shape

        # Resize
        pad_img = cv2.resize(pad_img, self.img_shape)
        # Channels-first
        input_img = pad_img[:, :, ::-1].transpose((2, 0, 1)).copy()
        # As pytorch tensor
        input_img = torch.from_numpy(input_img).float().div(255.0)

        #---------
        #  Label
        #---------

        label_path = self.label_files[index % len(self.img_files)].rstrip()

        labels = None
        if os.path.exists(label_path):
            labels = np.loadtxt(label_path).reshape(-1, 5)
            # Extract coordinates for unpadded + unscaled image
            x1 = w * (labels[:, 1] - labels[:, 3]/2) # Upper left x
            y1 = h * (labels[:, 2] - labels[:, 4]/2) # Upper left y
            x2 = w * (labels[:, 1] + labels[:, 3]/2) # Lower right x
            y2 = h * (labels[:, 2] + labels[:, 4]/2) # Lower right y
            # Adjust for added padding
            x1 += pad[1][0]
            y1 += pad[0][0]
            x2 += pad[1][0]
            y2 += pad[0][0]
            # Calculate ratios from coordinates
            labels[:, 1] = ((x1 + x2) / 2) / padded_w # Box center x
            labels[:, 2] = ((y1 + y2) / 2) / padded_h # Box center y
            labels[:, 3] *= w / padded_w # width
            labels[:, 4] *= h / padded_h # height
        # Fill matrix
        filled_labels = np.zeros((self.max_objects, 5))
        if labels is not None:
            filled_labels[range(len(labels))[:self.max_objects]] = labels[:self.max_objects]

        filled_labels = torch.from_numpy(filled_labels)

        sample = {'input_img': input_img, 'orig_img': pad_img, 'label': filled_labels, 'path': img_path}

        return sample

• __init__() 함수에서 이미지 파일경로와 라벨 파일경로를 읽습니다. 라벨 파일경로는 이미지 파일경로에서 "JPEGImage' 를 'labels', 'jpg' 를 'txt' 로 바꿔주면 얻을 수 있습니다. 
• 라벨은 이미지에서 객체의 중심, 높이, 너비의 상대적인 비율로 구성되어 있습니다. YOLO v3는 416x416 크기의 입력 이미지를 사용하므로 이미지를 정사각형으로 패딩하고 리사이즈해주기 때문에 라벨도 이에 마찬가지로 다시 계산되어야 합니다. 
• __getitem__() 함수에서 인덱스에 따른 이미지, 라벨 처리과정을 구현합니다. 이미지는 cv2 라이브러리로 읽으며 먼저 패딩을 통해 정사각형으로 만들고 416 크기로 리사이즈한 후 255를 나눠주어 [0,1]로 정규화합니다.
• 라벨에 대해서는 패딩되기 전의 원본 이미지 상의 객체 중심점, 높이, 너비를 구한 이후 패딩과 리사이즈에 따른 효과를 계산합니다.
• 마지막으로 한 이미지가 담을 수 있는 최대 객체 수를 max_object=50 으로 설정합니다. 따라서 이후 dataloader 로부터 라벨 데이터를 읽으면 [batch size, 50, 5] 형태가 됩니다. 

---

이후에는 Pytorch 의 DataLoader 모듈을 이용하여 배치를 뽑아내는 데이터로더를 구현합니다. 입력 이미지는 [batch size, 3, 416, 416] 형태가 되어야 하며 라벨은 [batch size, 50, 5] 형태가 되어야 합니다.

batch_size = 8
inp_dim = 416

dataloaders = {'train': DataLoader(ListDataset(train_path, img_size=inp_dim), batch_size=batch_size, shuffle=True, pin_memory=True, drop_last=False),
                'val': DataLoader(ListDataset(val_path, img_size=inp_dim), batch_size=batch_size, shuffle=False, pin_memory=True)}

for i_batch, sample_batch in enumerate(dataloaders['train']):
    input_images_batch, orig_images_batch, label_batch, path_batch = sample_batch['input_img'], sample_batch['orig_img'], sample_batch['label'], sample_batch['path']

    print(i_batch, input_images_batch.shape, orig_images_batch.shape, label_batch.shape)
    
    print('\t', label_batch[0,0])
    if i_batch == 3:
        break

 

Loss calculation

데이터 별 loss 를 계산하기 위해서는 네트워크의 output feature map 상에서 어떤 셀의 앵커 박스가 해당 객체에 대한 책임이 존재하는지 알아야 합니다. 즉, 라벨 정보를 네트워크의 출력에 맞게 재구성하는 과정이 필요합니다.

def build_targets(target, anchors, grid_size, num_anchors = 3, num_classes = 20):
    '''
    Arguments:
    - target: [batch size, max object, 5]
    '''

    nB = target.size(0)
    nA = num_anchors
    nC = num_classes
    nG = grid_size
    mask = torch.zeros(nB, nA, nG, nG) # [batch, num_anchors, grid, grid]
    tx = torch.zeros(nB, nA, nG, nG)
    ty = torch.zeros(nB, nA, nG, nG)
    tw = torch.zeros(nB, nA, nG, nG)
    th = torch.zeros(nB, nA, nG, nG)
    tconf = torch.zeros(nB, nA, nG, nG)
    tcls = torch.zeros(nB, nA, nG, nG, nC) # [batch, num_anchors, grid, grid, num_classes]

    for b in range(nB):  # for each image
        for t in range(target.shape[1]):  # for each object
            if target[b, t].sum() == 0:  # if the row is empty
                continue
            # Convert to object label data to feature map
            gx = target[b, t, 1] * nG
            gy = target[b, t, 2] * nG
            gw = target[b, t, 3] * nG
            gh = target[b, t, 4] * nG
            # Get grid box indices
            gi = int(gx)
            gj = int(gy)
            # Get shape of gt box
            gt_box = torch.FloatTensor(np.array([0, 0, gw, gh])).unsqueeze(0)  # 1 x 4
            # Get shape of anchor box
            anchor_shapes = torch.FloatTensor(
                np.concatenate((np.zeros((len(anchors), 2)), np.array(anchors)), 1))
            # Calculate iou between gt and anchor shapes
            anch_ious = bbox_iou(gt_box, anchor_shapes)
            # Find the best matching anchor box
            best_n = np.argmax(anch_ious)
            # Masks
            mask[b, best_n, gj, gi] = 1
            # Coordinates
            tx[b, best_n, gj, gi] = gx - gi
            ty[b, best_n, gj, gi] = gy - gj
            # Width and height
            tw[b, best_n, gj, gi] = math.log(gw / anchors[best_n][0] + 1e-16)
            th[b, best_n, gj, gi] = math.log(gh / anchors[best_n][1] + 1e-16)
            # One-hot encoding of label
            target_label = int(target[b, t, 0]) # Float 2 int
            tcls[b, best_n, gj, gi, target_label] = 1
            tconf[b, best_n, gj, gi] = 1

    return mask, tx, ty, tw, th, tconf, tcls

• mask 는 [batch size, num anchors, grid, grid] 형태이며, 객체 예측에 할당된 박스에 대해서는 1, 아니면 0으로 구성됩니다.
• tx, ty, tw, th 는 mask 와 같은 형태를 가지는 텐서이고 YOLO v3의 Figure 2와 같이 tx / ty 는 해당 셀 안에서의 오프셋, tw / th 는 K-Means 알고리즘으로 미리 정의된 앵커박스 크기 비율의 로그 변환으로 정의합니다.
• Output feature map 상에서 셀의 오프셋을 gi, gj 변수에 할당합니다. gx, gy, gw, gh 는 output feature map 상에서의 위치, 크기를 나타냅니다.
• 각 셀별로 ground-truth 박스와의 IOU가 제일 높은 앵커박스를 mask 텐서의 알맞인 인덱싱을 통해 1로 설정합니다. 이 앵커박스에 대해 tx, ty, tw, th 텐서 또한 해당되는 인덱스에 대해 값을 대입합니다.
• tconf 는 mask 와 동일하며 tcls 는 classification 을 위한 텐서로 [batch size, num anchors, grid, grid, num classes] 형태이며 해당하는 클래스 인덱스에 1을 할당합니다.

---

Bounding 박스 간 IOU를 계산하는 코드는 다음과 같습니다.

def bbox_iou(box1, box2, x1y1x2y2=True):
    """
    Returns the IoU of two bounding boxes
    Two scenarios： 1. box is represented by center and width 2. box is represented by cooridnates of left up and right bottom
    """
    if not x1y1x2y2:
        # Transform from center and width to exact coordinates
        b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
        b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
        b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
        b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
    else:
        # Get the coordinates of bounding boxes
        b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
        b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

    # get the corrdinates of the intersection rectangle
    inter_rect_x1 = torch.max(b1_x1, b2_x1)
    inter_rect_y1 = torch.max(b1_y1, b2_y1)
    inter_rect_x2 = torch.min(b1_x2, b2_x2)
    inter_rect_y2 = torch.min(b1_y2, b2_y2)
    # Intersection area
    inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1 + 1, min=0) * torch.clamp(
        inter_rect_y2 - inter_rect_y1 + 1, min=0
    )
    # Union Area
    b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
    b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

    iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

    return iou

---

마지막으로 Loss 함수를 정의합니다. 네트워크의 출력을 [batch size, num_anchors, grid, grid, num_class] 형태로 재구성하고 x, y, w, h 에는 MSE를 (Mean Squared Error) 적용하고 conf, cls 에 대해서는 BCE를 (Binary Cross Entropy) 적용합니다.

def Loss(input, target, anchors, inp_dim, num_anchors = 3, num_classes = 20):

    nA = num_anchors  # number of anchors
    nB = input.size(0)  # number of batches
    nG = input.size(2)  # number of grid size
    nC = num_classes
    stride = inp_dim / nG

    # Tensors for cuda support
    FloatTensor = torch.cuda.FloatTensor if input.is_cuda else torch.FloatTensor
    ByteTensor = torch.cuda.ByteTensor if input.is_cuda else torch.ByteTensor

    prediction = input.view(nB, nA, 5 + nC, nG, nG).permute(0, 1, 3, 4, 2).contiguous()  # reshape the output data

    # Get outputs
    x = torch.sigmoid(prediction[..., 0])  # Center x
    y = torch.sigmoid(prediction[..., 1])  # Center y
    w = prediction[..., 2]  # Width
    h = prediction[..., 3]  # Height
    pred_conf = torch.sigmoid(prediction[..., 4])  # Conf
    pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred

    # Calculate offsets for each grid
    grid_x = torch.arange(nG).repeat(nG, 1).view([1, 1, nG, nG]).type(FloatTensor)
    grid_y = torch.arange(nG).repeat(nG, 1).t().view([1, 1, nG, nG]).type(FloatTensor)
    scaled_anchors = FloatTensor([(a_w / stride, a_h / stride) for a_w, a_h in anchors])
    anchor_w = scaled_anchors[:, 0:1].view((1, nA, 1, 1))
    anchor_h = scaled_anchors[:, 1:2].view((1, nA, 1, 1))

    # Add offset and scale with anchors
    pred_boxes = FloatTensor(prediction[..., :4].shape)
    pred_boxes[..., 0] = x.data + grid_x
    pred_boxes[..., 1] = y.data + grid_y
    pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
    pred_boxes[..., 3] = torch.exp(h.data) * anchor_h

    mask, tx, ty, tw, th, tconf, tcls = build_targets(
        target=target.cpu().data,
        anchors=scaled_anchors.cpu().data,
        grid_size=nG,
        num_anchors=nA,
        num_classes=num_classes)

    # Handle target variables
    tx, ty = tx.type(FloatTensor), ty.type(FloatTensor)
    tw, th = tw.type(FloatTensor), th.type(FloatTensor)
    tconf, tcls = tconf.type(FloatTensor), tcls.type(FloatTensor)
    mask = mask.type(ByteTensor)

    mse_loss = nn.MSELoss(reduction='sum')  # Coordinate loss
    bce_loss = nn.BCELoss(reduction='sum')  # Confidence loss
    loss_x = mse_loss(x[mask], tx[mask])
    loss_y = mse_loss(y[mask], ty[mask])
    loss_w = mse_loss(w[mask], tw[mask])
    loss_h = mse_loss(h[mask], th[mask])
    loss_conf = bce_loss(pred_conf, tconf)
    loss_cls = bce_loss(pred_cls[mask], tcls[mask])
    loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls

    return (loss, loss_x, loss_y, loss_w, loss_h, loss_conf, loss_cls)

• YOLO v3 에서는 multi-scale prediction 을 사용하는데, 다음과 같은 앵커박스를 사용합니다. 이 앵커박스 크기는 입력 이미지 상에서의 크기로 output feature map 상에서 계산을 해주어야 하기 때문에 (입력크기/feature map 크기)를 나누어 스케일링합니다.(build_targets 함수에 스케일된 앵커박스 사이즈가 입력으로 들어갑니다.)

  anchors = ([(10, 13), (16, 30), (33, 23)], 
             [(30, 61), (62, 45), (59, 119)], 
             [(116, 90), (156, 198), (373, 326)])​

 

---

다음 포스트

[Machine Learning/기타] - Object Detection - YOLO v3 Pytorch 구현 (2)