Pytorch를 이용한 YOLO v3 논문 구현 #4 - NMS 구현과 objectness score

※본 포스팅은 아래 블로그를 참조해 번역하고 공부한 것입니다.

https://blog.paperspace.com/how-to-implement-a-yolo-object-detector-in-pytorch/

Tutorial on implementing YOLO v3 from scratch in PyTorch

output 값인 5x(B+C)를 NMS와 objectness score에 적용을 해야합니다. 이번 챕터에서는 util.py를 마저 정리해보겠습니다.

---

Object Confidence thresholding

def write_results(prediction, confidence, num_classes, nms_conf = 0.4):
    conf_mask = (prediction[:,:,4] > confidence).float().unsqueeze(2)
    prediction = prediction*conf_mask

이 함수는 output 바운딩 박스의 값에 대한 정보를 포함한다. threshold 아래에 있는 objectness score를 갖고 있는 각각의 바운딩 박스에 대하여, 모든 속성 값을 0으로 설정한다.

---

NMS (Non-maximum Suppression)

    #바운딩 박스의 중심 점을 좌측 상단, 우측 하단 모서리 좌표로 변환
    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]

각 바운딩 박스의 대각 모서리 쌍의 좌표를 사용하여 IoU를 계산하는 것이 코드 작성에 용이하다. 따라서 위 코드는 중심점을 모서리 좌표로 변환하기 위한 로직이다.

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

이미지마다 정답 영역과 정답 객체의 수는 다를 것이다. 그러므로 confidence thresholding과 NMS는 한번에 하나의 이미로 수행해야 한다. 이는 이미지의 index를 포함한 prediction의 첫 번째 차원을 반복해야 한다.

       #NMS
       #가장 높은 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)

maximum value를 지닌 class score만 살리고, 나머지는 삭제한다. 그리고 그 class의 index 값과 class score를 추가한다.

        non_zero_ind =  (torch.nonzero(image_pred[:,4]))
        try:
            image_pred_ = image_pred[non_zero_ind.squeeze(),:].view(-1,7)
        except:
            continue
        
        if image_pred_.shape[0] == 0:
            continue

앞에서 threshold보다 낮은 objectness score를 가진 바운딩 박스를 0으로 설정했는데, 이를 제거해준다. 만약 검출이 이뤄지지 않았다면, 그 이미지에 대해 무한 반복이 수행될 수 있기 때문에 continue를 추가해준다.

        #이미지에서 검출된 다양한 class 얻기
        img_classes = unique(image_pred_[:,-1])  # -1: class index

이미지에서 검출된 여러가지 class를 얻기 위한 코드다. 

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

동일한 클래스에 다수의 검출이 이뤄질 수 있기 때문에, 주어진 이미지에서 나타내는 클래스만 얻기 위해 unique 함수를 선헌한다.

        for cls in img_classes:
            #NMS
        
            #특정 클래스에 대한 detection
            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를 가진 detection 순으로 정렬
            #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를 수행하는 로직. 한번 반복 수행이 시작되면 가장 먼저 특정 클래스의 detecton을 추출한다.

            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 > theshhold인 detection -> 0
                iou_mask = (ious < nms_conf).float().unsqueeze(1)
                image_pred_class[i+1:] *= iou_mask       
            
                #0이 아닌 항목 제거
                non_zero_ind = torch.nonzero(image_pred_class[:,4]).squeeze()
                image_pred_class = image_pred_class[non_zero_ind].view(-1,7)

본격적인 NMS 로직. 일반적으로 알려져있는 NMS의 공식을 수행하기 위한 코드이다.

1. i 인덱스를 갖고 있는 box의 IoU와 i보다 큰 인덱스를 가진 바운딩 박스를 얻는다.
2. 모든 반복 수행중에서, 만약 i보다 큰 인덱스를 지닌 바운딩 박스가 nms_thresh보다 큰 값을 가지면 그 박스는 제거된다.
3. image_pred_class에 의해 값이 하나라도 제거 됐으면, idx iterations를 가질 수 없다. 따라서 NMS가 바운딩 박스를 추가적으로 제거할 수 없음을 확인하고 반복 수행을 중단한다.

---

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

문자 그래도 iou를 계산하기 위한 함수 선언이다. iou에 대한 자세한 수식은 다음 포스팅을 참조하자.

https://it-the-hunter.tistory.com/29?category=1036021 

[딥러닝]IOU에 대해서 이해해보자

---

예측값 저장

            batch_ind = image_pred_class.new(image_pred_class.size(0), 1).fill_(ind)     
            #이미지에 있는 class의 detection만큼 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))

write_resuls 함수는 d*8 크기의 tensor를 출력한다. 이전과 마찬가지로, 출력 tensor에 할당할 객체를 갖고 있지 않으면 출력 tensor를 초기화하지 않는다. 그러나 객체가 존재해서 한번이라도 초기화되면, 후속 객체를 연결하고 tensor가 초기화 됐는지 확인하기 위해 write flag를 사용한다.

    try:
        return output
    except:
        return 0

최종적으로 output이 초기화됐는지 확인 후 0을 반환한다. 각 prediction을 묶은 tensor의 형태로 예측값을 가진다.

---

util.py도 마무리입니다. 다음 챕터에서는 이미지를 읽기 위해 입력 파이프라인을 행성하고, 예측을 진행해 바운딩 박스를 그리는 로직을 수행할 것입니다.