[리뷰] TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-captured Scenarios

이번에는 VisDrone 2021 ICCV workshop에 발표된 논문인 TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-captured Scenarios를 읽고, 리뷰해보고자 합니다.

Index
1. Background
    1.1. YOLO
    1.2. Transformer
    1.3. Convolutional Block Attention Model
2. Abstract
3. Introduction
4. Related Work
5. Method
    5.1. Overview of YOLOv5
    5.2. TPH-YOLOv5
6. Experiment
7. Conclusion 

1. Background

1.1. YOLO

[리뷰] You Only Look Once: Unified, Real-Time Object Detection

1.2. Transformer

[리뷰] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

1.3. Convoultional Block Attention Model

[번외] Metric Learning을 위한 딥러닝 - CBAM(Conv Block Attention Module)

 

2. Abstract

• 최근 드론으로 촬영된 영상에서의 object detection은 인기 있음
• 그러나, 아래와 같은 문제점 때문에 어려움이 있음 
  • 드론은 비행 고도가 빈번하게 바뀌기 때문에, object의 scale 변화가 심하여 네트워크 최적화가 어려움
  • 드론의 고속 및 저고도 비행은 object들이 촘촘하게 있는 장면에서, motion blur를 발생시켜, object를 구별하는 것이 어려움

 

3. Introduction

• 앞서 언급한 문제들을 해결하기 위해 본 논문이 제안됨
  • YOLOv5를 기반으로, 다양한 scale의 object를 detection하기 위해 prediction head를 하나 더 추가
  • 기존의 prediction head를 transformer prediction head(TPH)로 교체하여 self-attention으로 예측 가능성을 탐색
  • convolutional block attention model(CBAM)을 통합하여 object가 밀집되어 있는 장면에서 attention region을 찾음
  • data augmentation, multi-scale test, multi-model integration, self-trained classifier와 같은 bag of tricks를 이용

 

4. Related Work

• Data Augmentation
  • 모델의 성능을 향상시키고, 다른 환경에 대해서도 robust하게 작동할 수 있도록 하기 위함
  • scaling, cropping, translation, shearing, rotation 등이 있음
    • 이 방법들이 조합된 MixUp, CutMix, Mosaic 등의 방법 또한 존재
  • 본 논문에서는 MixUp, Mosaic 및 기존 방식을 조합하여 이용
• Multi-Model Ensemble Method in Object Detection
  
  • 모델의 분산을 줄이는 방법으로, 단일 모델 대신 여러 모델을 훈련 시킨 후, 모델들의 예측을 결합하는 방법이 있음
  • object detection 모델의 bbox 예측 결과를 ensemble 하는 방법으로는 3가지가 있음
    • NMS 방법에서는 여러 bbox 중, 가장 confidence score가 높은 하나의 bbox만을 남기고, IoU가 특정 임계값(보통 0.5 이용)보다 높고 으면 동일한 object에 속하는 것으로 간주하므로 모델 성능에 IoU가 큰 영향을 미침
    • Soft-NMS는 NMS에 약간 변화를 준 방법으로, confidence score를 0으로 설정하고 삭제하는 대신 IoU 값을 기반으로 인접 경계 상자의 신뢰도에 대한 감쇄 기능을 설정
    • WBF 방법은 NMS와 다르게 모든 상자를 병합하여 최종 결과를 형성하여 모델의 모든 부정확한 예측을 해결할 수 있음
  • 본 논문에서는 최종 모델을 앙상블하기 위해 WBF를 이용
• Object Detection
  • 다양한 CNN기반의 object detector가 존재하며, 최근에는 backbone과 head 사이에 layer를 삽입하는 경우가 많음(neck이라 부름)
  • backbone
    • 분류 및 기타 문제에 대해 강력한 feature 추출 기능을 가지고 있음
    • VGG, ResNet, DenseNet, MobileNet, EfficientNet, CSPDarknet, Swin Transformer 등이 있음
  • neck
    • backbone에서 추출한 feature를 reprocess하고 비율을 조정하여 사용함으로써, 더욱 잘 활용하도록 설계
    • 일반적으로 bottom-up path와 top-down path로 구성
    • FPN, PANet, NAS-FPN, BiFPN, ASFF, SFAM 등이 있음
    • SPP, ASPP, RFB, CBAM 등과 같이 추가적으로 사용되는 것도 있음
  • head
    • classification network로서, object의 위치 및 카테고리를 판단하는 역할 수행

 

5. Method

5.1. Overview of YOLOv5

• YOLOv5에는 YOLOv5, YOLOv5m, YOLOv5l, YOLOv5x 등 4가지 모델이 있음
• 일반적으로 SPP 계층을 backbone으로 CSPDarknet53, neck으로 PANet을 사용
• 모델을 더욱 최적화하기 위해 bag of freebies와 specials 사용
• 본 논문에서는 계산 cost는 높지만, 성능이 가장 높은 YOLOv5x를 사용

5.2. TPH-YOLOv5

• Overview

• Prediction head for tiny objects
  
  • tiny object detection을 위해 prediction head를 하나 더 추가 (Figure 3의 TPH 1번)
    • object의 scale변화에서 발생하는 negative한 영향을 완화
    • 계산 cost는 증가
• Transformer encoder block
  • ViT에서 영감을 받아 Conv block과 CSP bottleneck block을 transformer encoder block으로 교체
    • 기존의 CSP bottleneck block에 비해 global information과 context를 잘 추출함
    • self-atention 매커니즘으로 feature representation potential을 탐색할 수 있음
    • 2개의 sub-layer를 포함하는데, 하나는 multi-head attention layer이며, 남은 하나는 FC layer임
      • 이들을 residual connection을 이용하여 연결
    • 복잡한 장면에서 occlusion에 대해 더 나은 성능을 가짐
  • head 부분에만 transformer encoder block을 적용하여 TPH(Transformer Prediction Head)와 backbone의 끝을 형성
    • backbone의 끝에서 추출되는 feature map의 해상도가 낮아서 계산 cost를 낮출 수 있기 때문
  • 입력 이미지의 해상도를 확대할 때, 일부 TPH block을 선택적으로 제거하여 학습을 수행할 수 있도록 함

• Convolutional block attention module (CBAM)
  • 간단하고 가볍지만 효과적인 CNN 구조에 통합가능한 attention module
  • end-to-end 방식으로 학습 가능
  • feature map이 주어진 상태에서 channel과 spatial 공간을 따라 feature map을 순차적으로 inference한 다음, 주어진 feature map과 attention map을 곱하여 adaptive한 feature refinement 수행
  • CBAM을 이용하면 confusing information에 저항하고, 유의미한 target object에 집중하도록 attention map을 추출할 수 있음

• Ms-testing and model ensemble
  • 모델 ensemble을 위해 5개의 모델을 훈련하며, inference 단계에서는 단일 모델에 대해 ms-test 전략 수행
  • ms-testing 전략의 구현 세부 사항
    • 테스트 이미지 크기를 1.3배로 조정
    • 이미지를 각각 1배, 0.83배, 0.67배로 줄임
    • 이미지를 수평으로 뒤집음
    • 6개의 서로 다른 scale의 이미지를 TPH-YOLOv5에 공급하고, NMS를 사용하여 prediction을 합침
    • 서로 다른 모델에서 동일한 ms-testing을 수행하고, WBF를 사용하여 5개의 prediction을 융합하여 최종 결과를 얻음
• Self-trained classifier
  • VisDrone2021에 대해 테스트한 결과, localization은 우수하지만, classification 능력이 떨어진다는 것을 발견하여, 이를 해결하기 위해 추가적으로 self-trained classifier를 제안
    • ground-truth의 bbox만을 추출하여 각 이미지 patch의 크기를 64x64로 조정하여 학습
    • ResNet18을 이용하여 분류 수행

 

6. Experiment

• VisDrone2021 dataset의 성능 비교

• 사용한 요소에 따른 성능 차이

• 입력 이미지의 크기와 각 카테고리의 label 수에 따른 성능 차이 및 모델 앙상블의 효과

• 데모 영상

 

7. Conclusion 

• YOLOv5에 transformer encoder block, CBAM, trick 등을 추가하여 특히 드론으로 촬영된 영상에서 object detection에 탁월한 TPH-YOLOv5 제안
• VisDrone2021 dataset에서 SOTA 달성

 

 

논문 링크

https://arxiv.org/abs/2108.11539

https://github.com/cv516Buaa/tph-yolov5