[Object Detection] YOLO(v3)를 PyTorch로 바닥부터 구현하기 - Part 4

 이 포스팅은 공부 목적으로 아래 게시물을 번역한 글입니다.

 

How to implement a YOLO (v3) object detector from scratch in PyTorch: Part 3

 

 파이토치로 YOLO v3 detector를 구현하는 튜토리얼의 part 4 입니다. 지난 part에서 우리는 신경망의 순전파를 구현했습니다. 이번 part에서, detections를 object confidence로 threshold하고 non-maximum suppression을 할 것입니다.

 

 이번 튜토리얼에 대한 코드는 Python 3.5와 PyTorch 0.4로 구동되도록 설계되었습니다. 전체 코드는 여기에서 확인하실 수 있습니다.

 

 이 튜토리얼은 5가지 Part로 나뉘어져 있습니다.

 

1. Part 1 : YOLO가 어떻게 작동하는지 이해하기

2. Part 2 : 신경망 구조의 계층 생성하기

3. Part 3 : 신경망의 순전파 구현하기

4. Part 4 : (현재) 비-최대 억제(Non-maximum suppression)와 객체 점수 임계값

5. Part 5 : 입력값과 출력값 

 

---

 이전 parts에서 주어진 입력 이미지로 여러개의 object detections를 출력하는 model을 구축했습니다. 좀더 정확하게, 출력은 B x 10647 x 85 shape의 tensor 입니다. B는 batch에서 이미지의 수이고, 10647은 이미지당 예측된 bounding boxes의 수이고, 85는 bounding box 속성 입니다.

 

 하지만 Part 1에서 설명한 것 처럼, output을 objectness score thresholding과 Non-maximum suppression을 적용해야 합니다. 이것을 하기 위해 util.py 파일 안에 write_results 함수를 작성해야 합니다.

 

def write_results(prediction, confidence, num_classes, nms_conf = 0.4):

 

 이 함수는 입력으로 prediction, confidence (objectness score threshold), num_classes (80), nmn_conf (NMS IoU threshold) 4가지를 취합니다.

 

---

Object Confidence Thresholding

 prediction tensor는 B x 10647 bounding boxes에 대한 정보를 포함합니다. threshold 아래에 있는 objectness score를 갖고 있는 각각의 bounding box에 대하여, 모든 속성 값들을 0으로 설정합니다.

 

    # threshold 보다 낮은 objectness score를 갖은 bounding box의 속성을 0으로 설정
    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
    prediction = prediction * conf_mask

 

---

Non-maximum Suppression 수행하기

 bounding box의 중앙 좌표와 높이, 넓이를 설명하는 bounding box 속성을 갖고 있습니다. 하지만 각 box의 대각 모서리 쌍의 좌표를 사용하여 두 boxes의 IoU를 계산하는 것이 쉽습니다. 따라서 다음과 같이 중심점을 좌측 상단, 우측 하단 모서리 좌표로 변경합니다.

 

    # bounding box의 중심 점을 좌측 상단, 우측 하단 모서리 좌표로 변환하기
    box_corner = prediction.new(prediction.shape)
    box_corner[:,:,0] = (prediction[:,:,0] - prediction[:,:,2]/2)
    box_corner[:,:,1] = (prediction[:,:,1] - prediction[:,:,3]/2)
    box_corner[:,:,2] = (prediction[:,:,0] + prediction[:,:,2]/2) 
    box_corner[:,:,3] = (prediction[:,:,1] + prediction[:,:,3]/2)
    prediction[:,:,:4] = box_corner[:,:,:4]

 

 모든 이미지에서 true detections의 수는 아마 다를 것 입니다. 예를 들어, 이미지 1, 2, 3은 각각 5, 2, 4개의 true detections를 갖고 있습니다. 그러므로 confidence thresholding과 NMS는 한 번에 하나의 이미지로 수행되어야 합니다. 이것은 연산히 포함된 vectorise가 불가능하고, prediction(이미지의 index를 포함한)의 첫 번째 차원을 반복해야 합니다.

 

    # 한번에 하나의 이미지에 대하여 수행
    for ind in range(batch_size):
        image_pred = prediction[ind]

 

 이전에 설명했던 것 처럼, write flag는 output을 초기화하지 않은 것을 나타냅니다. 그리고 전체 batch에 걸쳐서 true detections를 수집하기 위해 사용할 tensor 입니다.

 

 각 bounding box row는 85개 속성을 갖고 있으며, 이 중 80개는 class score입니다. 이 시점에서, maximum value를 지닌 class score만 관심을 두겠습니다. 그러므로 각 row로부터 80개의 class socre를 제거합니다. 그리고 가장 높은 값을 갖은 class의 index와 그 class의 class score를 추가합니다.

 

        # 가장 높은 값을 가진 class score를 제외하고 모두 삭제
        max_conf, max_conf_score = torch.max(image_pred[:, 5:5 + num_classes], 1)
        max_conf = max_conf.float().unsqueeze(1)
        max_conf_score = max_conf_score.float().unsqueeze(1)
        seq = (image_pred[:,:5], max_conf, max_conf_score)
        image_pred = torch.cat(seq, 1)

 

 threshold보다 낮은 object confidence를 가진 bounding bow rows를 0으로 설정했던 것을 기억 하시나요? 그것들을 제거하겠습니다.

 

        # threshold보다 낮은 object confidence를 지닌 bounding box rows를 0으로 설정한 것을 제거
        non_zero_ind =  (torch.nonzero(image_pred[:,4]))
        try:
            image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
        except:
            continue
        
        # PyToch 0.4 호환성
        # scalar가 PyTorch 0.4에서 지원되기 때문에 no detection에 대한 
        # not raise exception 코드입니다.
        if image_pred_.shape[0] == 0:
            continue 

 

 try-exept block은 no detections를 얻었을 때 상황을 다루기 위한 것 입니다. 이 경우에 이 이미지에 대하여 loop를 건너 뛰기 위해 continue를 사용합니다.

 

 이제 이미지에서 검출된 class를 얻겠습니다.

 

        # 이미지에서 검출된 다양한 classes를 얻기
        img_classes = unique(image_pred_[:,-1]) # -1 index는 class index를 지니고 있습니다.

 

 동일한 클래스에 다수의 true detections가 있을 수 있기 때문에, 주어진 이미지에서 나타내는 클래스를 얻기 위한 unique 함수를 사용하겠습니다.

 

def unique(tensor):
    tensor_np = tensor.cpu().numpy()
    unique_np = np.unique(tensor_np)
    unique_tensor = torch.from_numpy(unique_np)
    
    tensor_res = tensor.new(unique_tensor.shape)
    tensor_res.copy_(unique_tensor)
    return tensor_res

 

 그리고나서 NMS를 수행하겠습니다.

 

        for cls in img_classes:
            # NMS 실행하기

 

 한번 loop가 작동되면 첫번째로 할 것은 특정 클래스(cls 변수로 표기된)의 detections를 추출하는 것입니다.

 

            # 특정 클래스에 대한 detections 얻기
            cls_mask = image_pred_*(image_pred_[:,-1] == cls).float().unsqueeze(1)
            class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
            image_pred_class = image_pred_[class_mask_ind].view(-1,7)
            
            # 가장 높은 objectness를 지닌 detections 순으로 정렬하기
            # confidence는 가장 위에 있습니다.
            conf_sort_index = torch.sort(image_pred_class[:,4], descending = True )[1]
            image_pred_class = image_pred_class[conf_sort_index]
            idx = image_pred_class.size(0)   #detections의 수

 

 이제 NMS를 실행하겠습니다.

 

            for i in range(idx):
                # 모든 박스에 대해 하나하나 IOU 얻기
                try:
                    ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])
                except ValueError:
                    break

                except IndexError:
                    break
                
                # IoU > threshhold인 detections를 0으로 만들기
                iou_mask = (ious < nms_conf).float().unsqueeze(1)
                image_pred_class[i+1:] *= iou_mask       

                # non-zero 항목 제거하기
                non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                image_pred_class = image_pred_class[non_zero_ind].view(-1,7)

 

 여기서, bbox_iou 함수를 사용하겠습니다. 첫 번째 입력값은 loop에서 변수 i에 의해 인덱싱된 bounding box 행 입니다.

 

 bbox_iou에 두 번째 입력은 bounding boxes의 tensor입니다. bbox_iou 함수의 출력값은 bounding box의 IoU를 포함하는 tensor 입니다. 

 

 

 만약 동일한 클래스에서  threshold보다 큰 IoU를 지닌 두 개의 bounding boxes를 갖고 있다면, 낮은 class confidence를 지닌 것이 제거됩니다. 이미 bounding boxes들을 confidence 순으로 정렬 했습니다.

 

 loop의 중간에서 다음 코드는 i 인덱스를 갖고 있는 box의 IoU와 i보다 큰 인덱스를 지닌 bounding boxes를 줍니다.

 

                    #  i 인덱스를 갖고 있는 box의 IoU와 i보다 큰 인덱스를 지닌 bounding boxes 얻기
                    ious = bbox_iou(image_pred_class[i].unsqueeze(0), image_pred_class[i+1:])

 

 모든 반복에서, 만약 i 보다 큰 인덱스를 지닌 bounding boxes가 threshold nms_thresh 보다 큰 IoU를 갖고 있으면, 그 box는 제거됩니다.

 

                # IoU > threshhold인 detections를 0으로 만들기
                iou_mask = (ious < nms_conf).float().unsqueeze(1)
                image_pred_class[i+1:] *= iou_mask       

                # non-zero 항목 제거하기
                non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                image_pred_class = image_pred_class[non_zero_ind].view(-1,7)

 

 try-catch 블락에서 ious를 계산하는 코드가 있습니다. 이것은 loop가 idx iterations(image_pred_class에서 row의 수)을 실행하도록 설계되었기 때문입니다. 하지만, loop에서 진행했던 것 처럼, bounding boxes의 수는 image_pred_class로 제거됩니다. 이것은 image_pred_class에 의해 값이 하나라도 제거 됬으면, idx iterations를 가질 수 없다는 것을 의미합니다. 따라서, 경계가 벗어난 값(IndexError)을 인덱싱하거나 image_pred_class[i+1:] 슬라이싱은 빈 tensor를 반환하여 ValueError를 트리거할 수 있습니다. 이 시점에서, NMS가 추가적인 bounding boxes를 제거할 수 없음을 확인할 수 있고 loop를 중지합니다.

 

IoU 계산하기

 이것은 bbox_iou 함수입니다.

 

# iou 계산하기
def bbox_iou(box1, box2):
    # 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]
    
    # 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 영역
    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 영역
    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)
    
    return iou    

 

예측값 저장하기

 write_results 함수는 Dx8 크기의 tensor를 출력합니다. 여기서 D는 모든 이미지에서 true detections이고, 각 행으로 나타납니다. 각 detections는 8개의 속성을 갖고 있습니다. 즉 배치에서 이미지의 인덱스, 4개의 꼭지점 좌표, objectness score, 가장 높은 클래스 score, 그 class의 인덱스

 

 이전과 마찬가지로, 출력 tensor에 할당할 detection을 갖고 있지 않는 한, 출력 tensor를 초기화하지 않습니다. 이것이 한번 초기화되면, 이것에 후속의 detections를 연결합니다. tensor가 초기화 됬는지 않됬는지 나타내기 위해 write flag를 사용합니다. 클래스에 대한 반복 loop 끝에서, 최종 detections를 tensor output에 추가합니다.

 

            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)      
            # 이미지에 있는 class의 detections 만큼 batch_id를 반복합니다.
            seq = batch_ind, image_pred_class

            if not write:
                output = torch.cat(seq,1)
                write = True
            else:
                out = torch.cat(seq,1)
                output = torch.cat((output,out))

 

 함수의 끝에서, output이 초기화 도었는지 확인합니다. batch에 어떤 이미지에서 detection이 하나라도 검출되지 않았다면 0을 반환합니다.

 

    try:
        return output
    except:
        return 0

 

  이 포스팅은 여기까지 입니다. 포스팅의 마지막에서, 최종적으로 각 prediction을 묶은 tensor의 형태로 예측값을 갖게 됩니다. 마지막으로 하나 남은 것은 이미지를 읽기 위한 입력 pipeline을 생성하는 것입니다. 그리고 prediction을 계산하고 이미지에 bounding boxes를 그립니다. 그리고나서 이 이미지들을 출력합니다. 이것들이 다음 part에서 할 것들입니다.