YOLO : You Only Look Once: Unified, Real-Time Object Detection 논문 요약
Abstract
객체 탐지의 이전 연구들에서는 분류기에서도 detection을 수행할 수 있도록 했다. 본 연구에서는 하나의 단일 신경망을 사용해서 여러 bounding box와 class probabilities를 예측하는 객체 탐지를 regression 문제로 취급했다. 본 연구의 YOLO 모델은 객체 탐지 파이프라인 전체가 하나의 네트워크로 이루어져있어서 실시간 이미지를 초당 45프레임으로 처리할만큼 굉장히 빠르다. 다른 네트워크와 비교했을 때 localization에 대한 오류가 존재하긴하지만 배경에 대해서는 예측을 잘 수행하고 다른 도메인에 대해서도 RCNN, DPM보다 일반화가 잘 된 예측을 수행한다.
Details
Introduction
- 최근 detection networks들은 분류기가 탐지를 수행할 수 있도록 변형한다. 분류기는 탐지 역할을 하기 위해서 다양한 위치와 크기의 이미지들에서 객체를 탐지하고 평가한다. DPM 형태의 모델에서는 sliding window 개념을 활용하여 전체 이미지에 대해 균일한 간격으로 분류를 수행한다.
- 보다 최근인 RCNN에서는 물체가 있을법한 곳에 bounding box를 생성하고 분류기는 이 bounding box에 대해 분류를 수행한다. 분류 후에는 bounding box 내의 객체가 중복되면 이를 제거하고, box를 다른 object들을 기반으로 rescoring 한다. RCNN의 경우 각각의 tasks들이 별개로 이루어져있기때문에 복잡하고 예측에 많은 시간이 소요된다는 단점을 가지고있다.
- YOLO(You Only Look Once)에서는 단일 네트워크로 CNN에서 여러 bbox와 각 bbox 마다 존재하는 class의 확률을 예측한다. 아래 Figure 1. 이미지를 보면, YOLO는 이미지를 고정된 448 x 448 사이즈로 리사이징 하고, CNN을 통해 class 분류 및 bbox를 생성하고, confidence를 기반으로 NMS을 수행해 객체당 하나의 bbox를 남긴다.
- YOLO의 장점은 첫째, YOLO는 detection을 하나의 regression problem으로 취급하기때문에 파이프라인이 복잡하지 않다. 속도가 매우 빠르면서 다른 실시간 탐지 시스템에 비해 두배 이상의 mAP를 기록했다.
- 둘째, YOLO는 sliding window나 region proposals 기법과 달리 훈련 및 테스트 시간 동안 고정된 grid cell로 전체 이미지를 보기때문에 클래스에 대한 context 정보나 class의 형태 등과 같은 정보를 암시적으로 인코딩할 수 있다. Fast R-CNN의 경우 large context 정보가 제대로 반영되지 않기때문에 background patches에 예측력이 좋지않다. YOLO는 Fast R-CNN의 절반 정도의 background 예측 에러를 갖고있다.
- 셋째, YOLO는 natural한 이미지를 학습했고, 예술 이미지에 대해 test했을 때 일반화 능력이 뛰어났다. DPM, Fast R-CNN보다 새로운 도메인에 대해서도 예측을 잘 수행했다.
- 속도가 매우 빠르지만 아직 localization과 작은 객체 탐지에 대해 어려움을 겪고 있고 여러 실험을 통해 속도와 정확도간의 trade-off에 대해 살펴본다.
Unified Detection
YOLO network는 한 이미지에 존재하는 모든 클래스에 걸쳐 모든 bbox를 동시에 예측한다. 즉, 네트워크는 전체 이미지와 이미지에 포함된 객체를 전역적으로 예측한다는 것이다. YOLO는 이미지를 SxS의 그리드로 나누는데 만약 객체의 중심이 그리드셀 내에 있다면 해당 그리드셀에서는 반드시 객체가 탐지되어야 할 것이다. 각각의 그리드 셀마다 B개의 bbox와 bbox 내의 confidence score(물체 존재 확신도)를 예측한다. confidence를 수식화하면 $Pr(Object) * IOU(truth|pred)$로 표현한다. 만약, 그리드셀 내에 객체가 없다면 confidence score는 0이 되어야한다. 물체가 존재한다면 confidence score는 IOU와 같아지길 원한다.
Bounding box는 x,y,w,h,confidence score 5개의 예측값을 가진다. x,y는 grid cell 내에서의 중심점 좌표고 w,h는 전체 이미지 내의 넓이와 높이다. Confidence score는 예측 box와 어떤 ground-truth box간의 IOU를 대변한다.
또한, 각각의 그리드 셀은 C(conditonal class probabilities)라는 클래스 확률을 예측한다.
\[C = Pr(Class_{i}|Object)\]이 조건부 확률은 그리드 셀안에 객체가 있다는 조건 하에 해당 객체가 어떤 class인지에 대한 확률이다. 그리드 셀 내의 Bounding box의 갯수와 상관없이 한 그리드 셀마다 하나의 class probabilities만 예측한다.
Test 단계에서는 C와 cofidence score를 곱해 box내에 특정 class가 존재할 확률과 bounding box가 얼마나 이 객체에 맞게 잘 형성되었는지를 파악한다.
\[Pr(Class_i|Object) * Pr(Object) * IOU(truth|pred) = Pr(Class_i) * IOU(truth|pred)\]아래 Figure 2.에서 YOLO가 어떻게 detection을 수행하는지 잘 설명했다. Box당 4개의 좌표와 물체 존재 확신도, class probobilities가 필요하고 이것을 그리드마다 수행하게되므로 예측에는 총 S x S x (5B + C)의 텐서가 필요하다.
1. Network Design
- PASCAL VOC 데이터셋에서 평가를 진행한다.
- CNN에서 이미지의 특징을 추출하고 FC층에서 class 확률과 bbox 좌표를 예측한다.
- ImageNet 데이터를 활용하여 classification task로 pretrain했고, detection에서는 224x224크기의 이미지를 두배로 늘려서 사용했다.
- 아키텍처는 GoogLeNet과 유사한데 인셉션 모듈대신 1x1, 3x3 합성곱을 사용했다(DarkNet).
- Fast YOLO는 24 layer 대신 9 layer만 사용해서 속도를 빠르게 했다.
2. Training
- 사전훈련한 모델은 ImageNet 2012에서 GoogLeNet과 비슷한 수준의 성능을 기록했고, 모든 훈련과 추론에 DarkNet 프레임워크를 사용했다.
- 객체 탐지를 수행하기 위해 4개의 convolutional layer와 랜덤하게 가중치가 초기화된 2개의 완전결합층을 추가해서 모델을 조금 수정했다. 또한, detection에는 더 세밀한 작업이 필요해 resolution을 2배로 높였다(448 x 448).
- 최종 layer에서는 class 확률과 bbox 좌표를 예측하는데 bbox는 w,h를 정규화해서 0과 1사이가 되도록했고, x,y 좌표도 특정 그리드 셀 위의 오프셋으로 변환하여 0과 1사이로 변환했다.
- 최종 layer만 linear activation을 사용했고 모든 다른 layer에서는 Leaky ReLU를 사용했다.
- MSE를 사용하면, 손실함수가 간단하지만 mAP를 극대화하려는 목표와는 거리가 있다. 또한, grid cell에 물체가 포함되지 않은 경우가 많다면 물체가 있는데도 confidence score가 0으로 되어 예측 자체가 수행되지 않은 경우가 생길 수 있기때문에 학습이 불안정할 수도 있다.
- 이를 해결하기 위해, bbox 좌표 예측의 손실을 높이고, 객체가 포함되지 않은 box에 대한 confidence의 예측 손실을 줄이는 방법을 사용했다. localization과 classification 중 localization의 가중치를 더 증가시키고 객체가 없는 confidence loss의 가중치를 있는 가중치보다 더 감소시키는 방법이다(객체가 없는 경우가 훨씬 많기때문에 객체가 있을 때의 loss가 더 중요함). 이는 $\lambda_{coord} = 5$, $\lambda_{noobj} = 0.5$를 설정하여 해결할 수 있다(coordinate는 가중치가 정수라서 원래보다 커지고, noobj는 소수라서 작아짐).
- MSE는 bbox가 크든, 작든 동일한 가중치를 부여한다는 단점이 있다. 오차의 관점에서 크기가 큰 bbox는 작은 bbox보다 오차에 많은 영향을 미친다. 크기가 큰만큼 오차가 상대적으로 더 클 것이기때문이다. 또한, 작은 bbox는 큰 bbox보다 편차에 민감하다. 예를 들어, 큰 객체를 감싸고 있는 bbox가 0.5 움직여도 여전히 객체를 감싸고 있을 수도 있지만, 작은 객체를 감싸고 있는 bbox가 0.5 움직이면 객체에서 벗어날 수도 있기때문이다. 따라서 큰 bbox의 loss 때문에 작은 bbox가 영향을 받는다면 최적화가 잘 이루어지지 않을 수 있다. 이 문제를 해결하기 위해 넓이와 높이를 직접적으로 예측하지 않고, 제곱근을 예측하는 방법을 사용했다. 제곱근을 예측하면 box의 크기가 클수록 증가율이 낮아지기때문에 bbox가 크더라도 제곱근으로 인해 작은 bbox가 큰 영향을 받지 않게되고 loss를 줄일 수 있다. 결과적으로 Box가 클수록 증가율이 작아져 IOU에 적은 영향을 끼치게된다.
- 위의 손실함수 관련 내용을 수식으로 표현하면 다음과 같다.
- $1_{ij}^{obj}$는 grid cell내에 class가 존재한다는 것을 의미하고 존재하면 1, 아니면 0으로 표현한다.
- $1_{ij}^{noobj}$는 객체가 존재하지 않을 때 confidence score loss를 계산하기 위해 사용된다.
- $1_{ij}^{obj}$ 는 grid cell $i$의 $j$번째 bbox predictor가 사용되는지의 여부이다. 위의 수식을 5단계로 표현하면
- 먼저, 객체가 존재하고 그리드 셀 $i$의 bbox predictor $j$에 대해 x,y loss를 구한다.
- 두번째, 객체가 존재하고 그리드 셀 $i$의 bbox predictor $j$에 대해 w,h loss를 구한다.(큰 bbox의 증가율이 커지지 않도록 제곱근을 예측)
- 세번째, 객체가 존재하고 그리드 셀 $i$의 bbox predictor $j$에 대해 confidence loss를 구한다(물체가 존재하기때문에 $C_i=1$)
- 네번째, 객체가 존재하지않을때 그리드 셀 $i$의 bbox predictor $j$에 대해 confidence loss를 구한다.($C_i=0$)
- 다섯번째, 객체가 존재하고 그리드 셀 $i$의 bbox predictor $j$에 대해 class probabilities loss를 구한다.(class가 맞으면 $p_i(C)=1$ 아니면 0)
- $\lambda_{coord}$ x,y,w,h의 좌표 loss와 다른 loss간의 밸런스를 위한 parameter.
- $\lambda_{noobj}$ 객체가 있는 box와 없는 box의 loss 간의 밸런스를 위한 parameter.
- 위의 모든 과정을 다 더해서 손실함수를 만들고 모델을 최적화시킨다.
- 과적합을 막기위해서 dropout(0.5), 원본 이미지의 크기보다 최대 20%까지 랜덤하게 scaling & translation 적용, HSV factor 조절의 data augmentation을 사용했다.
3. Inference
- Test 시에는 PASCAL VOC에서 이미지당 98개의 bounding box가 그려졌고 각 박스에 대해 class를 예측했다.
- 물체가 너무 크거나 물체가 여러 셀의 경계 근처에 있다면 주변의 다른 셀들의 정보를 참고해야 원활한 localized가 수행될 수 있는데 이러한 경우 하나의 객체가 여러 셀에서 발견되는 multiple detection 문제가 생길 수 있다. NMS를 통해 하나의 객체에 하나의 bounding box만 남도록 했고 RCNN이나 DPM처럼 NMS가 매우 중요하진않지만 이를 사용했을 때 mAP가 조금 향상되었다.
4. Limitations of YOLO
- YOLO는 각 그리드 셀마다 두개의 bounding box가 그려지고(각각 다른 종횡비를 가진 bbox를 그리고 NMS기법으로 하나만 남김), 각 셀은 오직 하나의 class로만 예측이되어야하기 떄문에 공간적 제약이 생긴다. 이런 공간적인 제약은 근처에 있는 한 셀에 여러 objects가 있는 경우 모든 objects를 잘 탐지하지 못하고, 특히 크기가 작은 물체를 탐지하는데 어려움을 겪게한다.
- Data로부터 bounding box를 그려내기때문에 가로,세로 비율이나 형태가 익숙하지 않다면 예측에 어려움을 겪는다. 또한, 여러 층을 거치며 downsampling된 features를 사용하기때문에 bbox를 예측하는 단계에서는 input image의 정보가 많이 선명하진 않을 것이다.
- Detection performance를 위해 loss function을 정의했지만, 제곱근을 사용했더라도 작은 bbox나 큰 bbox의 loss function을 결국 유사하게 가져갔고 그 결과 큰 bbox보다 작은 bbox가 IOU에 많은 영향을 미쳤다.
- Error의 가장 큰 문제는 localization이다.
Comparison to Other Detection System
YOLO detection system이 다른 systems들과 어떤 공통점 혹은 차별점을 가지는지 살펴보자.
Deformable parts models
- DPM은 sliding window 개념을 활용하여 객체를 탐지한다.
- DPM은 각각 분리된 형태로 파이프라인을 구축해 features를 추출하고, regions에서 분류를 수행하고, 가장 점수가 높은 regions에서 bbox를 예측한다.
- YOLO는 이 모든 과정을 하나로 통합했다는 점에서 DPM과 차이가 있다. 통합된 아키텍처로 DPM보다 더 빠르고 정확한 모델을 만들었다.
R-CNN
- R-CNN은 sliding window 대신 selective search 알고리즘으로 다양한 region proposals을 생성하고, CNN으로 features를 추출, SVM으로 분류를 수행, linear model로 bbox 예측, 중복된 bbox는 NMS로 제거하는 기법들을 사용했다.
- 굉장히 복잡한 파이프라인이고 각각의 tasks가 모두 독립적으로 이루어져 results를 출력하는 속도가 매우 느리다는 단점을 가지고있다.
- YOLO는 potential bbox를 제안받고, CNN을 통해 features를 추출한다는 점이 동일하지만 selective search가 아닌 공간적 제약을 가진 grid cell로 region proposals을 수행한다는 점에서 차이가 존재한다.
- 또한, 2000개의 bbox가 생성되는 R-CNN에 비해 98개의 bbox만 생성되고, 각각의 components를 하나로 통합했다는 점에서 R-CNN과 차이가 있다.
Other Fast Detectors
- DPM과 R-CNN 모두 각각의 components를 개선시켜 속도와 성능을 높였지만, YOLO는 여전히 애초에 하나의 pipeline에서 속도가 빠른 네트워크를 구축했다는 점에서 차별점이 존재한다.
- YOLO는 general purpose detector로 다양한 객체들을 동시에 예측할 수 있다.
- 이외에도 YOLO는 다른 여러가지 detector system보다 빠르고 한 이미지 내에서 single 뿐만 아니라 multiple objects도 잘 탐지할 수 있는 모델이다.
Experiments
YOLO와 다른 real-time detection systems를 비교하기 위해 VOC 2007 데이터셋을 Fast R-CNN과 비교했다. Error를 계산하는 방법이 다르기때문에 Fast R-CNN을 rescore했고, background false positives 에러를 감소시켜 기존보다 Fast R-CNN의 성능을 높게 조정했다. 또한, VOC 2012의 SOTA mAP와 비교했고, 최종적으로 YOLO의 일반화능력을 평가하기 위해 새로운 도메인인 예술작품에서 다른 시스템과 비교해보았다.
1. Comparison to Other Real-Time Systems
- 많은 연구들이 networks를 빠르게 만드는데 초점을 맞추고있다. 하지만 실제로 초당 30프레임 이상으로 실행되는 시스템은 DPM 모델밖에 없다. 30Hz or 100Hz로 실행되는 DPM의 GPU와 YOLO를 비교했다.
- Fast YOLO는 현존하는 가장 빠른 detection model이고 52.7%의 mAP로 이전보다 2배 이상 정확한 모델이다. YOLO 또한, mAP 63.4%까지 성능을 끌어올렸다.
- VGG-16을 사용해서 YOLO를 훈련시키면 모델이 정확하지만 속도가 기존보다 떨어졌다. 이는 VGG를 기반으로 한 다른 모델들과 비교하기에는 유용하지만 YOLO보다 많이 느려서 VGG로 훈련시키지 않은 원래 YOLO로 비교를 진행했다.
- DPM은 mAP의 큰 희생없이 속도를 효과적으로 높였으나 여전히 성능은 2배 이상 떨어진다. 특히, 딥러닝을 통한 접근에도 성능이 높지 않다는 점에서 한계를 가지고있다.
- R-CNN에서 R을 빼면 selective search가 static bounding box proposals로 바껴서 R-CNN보다 훨씬 빠른 모델이 만들어진다. R-CNN은 여전히 region proposals에 의존하며 좋은 proposals이 없다면 높은 정확도를 낼 수 없다.
- Fast R-CNN은 R-CNN보다 속도가 빨라졌지만 여전히 selective search에 의존하며 이미지당 2초의 proposals 시간이 소요된다. 따라서, mAP는 높지만 0.5fps라서 실시간으로 취급하긴 어렵다.
- 가장 최근 모델인 Faster R-CNN은 selective search를 neural network로 대체해서 성능과 속도에 큰 향상을 이루어냈다. 테스트 결과, 가장 정확도가 높은 모델은 7fps, 더 작은 대신 덜 정확한 모델은 18fps로 실행됐다. Faster R-CNN에 다른 여러가지 모델을 훈련시키면 정확도 향상에 비해 속도가 YOLO보다 크게 느려졌다.
2. VOC 2007 Analysis
- 현존하는 모델 중 PASCAL에서 가장 좋은 성능을 가지고 있는 Fast R-CNN 모델과 성능을 비교했다. Hoiem의 방법론을 사용했고 각 class에 대해 top N predictions(가장 잘맞춘 N개의 class 예측을 평균내어 에러를 비교)을 확인했다. 각각의 prediction은 correct or 다음과 같은 type of error로 구분된다.
- Correct: correct class and IOU > 0.5
- Localization: correct class, 0.1 < IOU < 0.5
- Similar: class is similar, IOU > 0.1
- Other: class is wrong, IOU > 0.1
- Backgruond: IOU < 0.1 for any object
- 아래 Figure .4를 통해 YOLO는 localization error가 다른 에러를 합친 것보다 더 클만큼 localization에 어려움을 겪고있다는 것을 알 수 있다.
- Fast R-CNN은 localization error는 작지만 backgound에 대해 예측을 잘 수행하지 못하고있다. Background error는 모델이 object라고 예측했는데 실제로는 배경이었던 false positives error이다. 즉, 배경을 제대로 맞추지 못하는 것이다.
3. Combining Fast R-CNN and YOLO
- YOLO가 Fast R-CNN보다 배경을 잘 예측하기때문에 이 둘을 조합하여 성능을 향상시켰다. R-CNN이 예측한 bbox에 대해 YOLO가 유사한 박스를 예측하면 YOLO가 미리 지정한 확률과 두 박스 간의 겹침의 정도에 따라 해당 예측에 boost를 준다.
- 그 결과, YOLO와 결합했을 때 Fast R-CNN은 기존보다 3.2% 증가한 75%의 mAP를 달성했다.
- 이 결합 모델은 각 모델을 개별적으로 실행한 다음 결과를 합치기때문에 YOLO의 속도 이점을 누릴 순 없지만, Fast R-CNN의 원래 속도에 비해 계산 시간이 크게 추가되진 않았다.
4. VOC 2012 Results
- VOC 2012 test sets에서 YOLO는 57.9% mAP를 기록했다. 이는 VGG-16을 사용한 original R-CNN 모델과 비슷한 성능이다.
- 병, 모니터 등과 같은 작은 물체를 잘 예측하지 못했지만, 고양이와 기차 등 다른 카테고리에서는 YOLO가 더 높은 성능을 기록했다. Fast R-CNN과 YOLO 결합 모델은 70.7% mAP를 기록하며 최종 5위를 기록했다.
5. Generalizability: Person Detection in Artwork
- 현실에서는 모델이 접해보지 못한 수많은 데이터가 존재한다. 따라서, YOLO의 일반화 능력을 평가하기 위해 Picasso Dataset과 People-Art Dataset을 사용하여 예술 작품 속에서 사람을 탐지하는 test를 진행했다.
- 성능은 사람만 탐지할 것이기때문에 people class에 대한 average precision을 지표로 사용했다. 모든 모델은 VOC 2007의 people 데이터로 학습했고 Picasso model은 VOC 2012로, People-Art 모델은 VOC 2010으로 학습했다.
- R-CNN은 VOC 2007에서 AP가 높았지만, artwork에서는 AP가 크게 떨어졌다. 분류기 단계에서 작은 regions을 보고 좋은 proposals을 수행해야하기때문에 어려울 것이다.
- DPM은 artwork에서도 비슷한 AP를 가졌다. Object에 대해 공간 정보를 잘 간직한 모델이기때문에 R-CNN보다 감소량이 적을 것이다. 하지만, 애초에 AP 자체가 높지않다는 것이 문제다.
- YOLO는 VOC 2007에서 성능이 우수할 뿐만 아니라 artwork에서도 성능 저하가 작다. YOLO도 DPM처럼 객체의 모양이나 공간 정보를 잘 간직하고, artwork는 일반 이미지와 다르지만 객체의 크기와 모양이 비슷했기때문에 좋은 성능을 냈을 것으로 추측한다.
Real-Time Detection In The Wild
- 웹캠을 통해 야생에서 객체 탐지를 수행했을 때 YOLO는 객체가 움직이거나 모양이 변할 때 이를 감지했다.
Conclusion
- YOLO는 전체 이미지를 직접적으로 학습하고 간단하게 구축할 수 있는 모델이다.
- Classifier-based approaches와 달리 detection performance에 알맞은 손실함수를 정의했고 모델에 이를 사용했다.
- Fast YOLO는 가장 빠른 general-purpose object detector이고, YOLO는 real-time 객체 탐지에서 SOTA를 기록했다.
- 새로운 도메인에 대해서도 빠르면서 일반화가 잘된 detection이 가능하다.
개인적인 생각
- YOLO는 class classification과 bounding box regression을 단일 네트워크로 구축했다는 점에서 굉장히 의미있는 연구였다.
- 다른 논문들과는 달리 기존과는 전혀 다른 도메인인 artwork를 사용해서 일반화 능력을 평가했기때문에 YOLO의 일반화 능력에 더 신뢰가 갔다.
- 여러가지 detection model들과 성능을 비교하고 심지어 결합하기까지 해본 실험을 통해 YOLO가 빠르고 성능도 준수한 모델 개발을 위해 얼마나 많은 시간을 들였는지 느껴졌다.
- 2023년 기준, YOLOv8이 나와 객체 탐지에서 매우 좋은 성능을 내고있는데 향후 논문들에서는 기존 YOLO가 가지고 있던 localization과 작은 객체 탐지에 대한 문제를 어떻게 해결할지 기대가 된다.
- YOLO에서는 한 그리드셀에서 하나의 classification만 수행이 되도록했는데 성능을 높이기위해서는 그리드셀 내에서도 multiple detection이 가능해야한다. 이와같은 문제를 어떻게 해결할지 궁금하다.
구현
Pytorch로 YOLOv1 모델을 구현해보자(참고).
1. Model
# 튜플이면 해당 layer가 하나인 것, list이면 해당 layer가 여러개인 것
architecture_config = [
(7, 64, 2, 3), # backbone에서 224 이미지를 2배 키워서 448
"M", # max pooing strdie 2, kernel 2 -> 112
(3, 192, 1, 1), # 112
"M", # 56
(1, 128, 1, 0), # 56
(3, 256, 1, 1), # 56
(1, 256, 1, 0), # 56
(3, 512, 1, 1), # 56
"M", # 28
[(1, 256, 1, 0), (3, 512, 1, 1), 4], # 28
(1, 512, 1, 0), # 28
(3, 1024, 1, 1), # 28
"M", # 14
[(1, 512, 1, 0), (3, 1024, 1, 1), 2], # 14
(3, 1024, 1, 1), # 14
(3, 1024, 2, 1), # 7
(3, 1024, 1, 1), # 7
(3, 1024, 1, 1), # 7
]
class CONVBLOCK(nn.Module):
def __init__(self, in_channels, out_channels, **kwargs):
super(CONVBLOCK, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
self.batchnorm = nn.BatchNorm2d(out_channels)
self.leakyrelu = nn.LeakyReLU(0.1)
def forward(self, x):
return self.leakyrelu(self.batchnorm(self.conv(x)))
class YOLOv1(nn.Module):
def __init__(self, in_channels=3, **kwargs):
super(YOLOv1, self).__init__()
self.architecture = architecture_config
self.in_channels = in_channels
self.darknet = self._create_conv_layers(self.architecture)
self.fc = self._create_fc(**kwargs)
def forward(self,x):
x = self.darknet(x)
x = torch.flatten(x, start_dim= 1)
return self.fc(x)
def _create_conv_layers(self, architecture):
layers = []
in_channles = self.in_channels
for x in architecture:
if type(x) == tuple:
layers += [
CONVBLOCK(in_channles, out_channels=x[1], kernel_size = x[0], stride = x[2], padding = x[3]
)
]
in_channles = x[1]
elif type(x) == str:
layers += [nn.MaxPool2d(kernel_size=(2,2), stride = 2)]
elif type(x) == list:
conv1 = x[0]
conv2 = x[1]
num_repeats = x[2]
for _ in range(num_repeats):
layers += [
CONVBLOCK(in_channles, out_channels = conv1[1], kernel_size = conv1[0], stride = conv1[2],
padding = conv1[3]
)
]
layers += [ # conv1의 output이 conv2의 in
CONVBLOCK(in_channels = conv1[1], out_channels = conv2[1], kernel_size = conv1[0], stride = conv1[2],
padding = conv1[3]
)
]
in_channles = conv2[1]
return nn.Sequential(*layers)
def _create_fc(self, num_split_cell, num_boxes, num_classes):
S, B, C = num_split_cell, num_boxes, num_classes
fc_layer = nn.Sequential(
nn.Flatten(),
nn.Linear(1024 * S * S, 4096),
nn.Dropout(0.0),
nn.LeakyReLU(0.1),
nn.Linear(4096, S * S * (5 * B + C)) # 각 박스마다 4개좌표 + confidence, 각 class에 속할 확률 값
)
return fc_layer
2. Dataset
import torch
from pycocotools.coco import COCO
import os
import cv2
import numpy as np
import transforms as T
from PIL import Image
import matplotlib.pyplot as plt
class SoccerDataset(torch.utils.data.Dataset):
def __init__(self, data_path, label_path, S = 7, B = 2, C = 3, transforms = None):
self.data_path = data_path
self.label_path = label_path
self.transforms = transforms
self.S = S
self.B = B
self.C = C
self.imgs = [x for x in sorted(os.listdir(data_path)) if '.jpg' in x]
self.labs = [x for x in sorted(os.listdir(label_path)) if '.txt' in x]
def __getitem__(self, idx):
lab = self.labs[idx]
boxes = []
with open(self.label_path + lab) as f:
for label in f.readlines():
class_label, x, y, width, height = [
float(x) if float(x) != int(float(x)) else int(x)
for x in label.replace("\n", "").split()
]
boxes.append([class_label, x, y, width, height])
boxes = torch.tensor(boxes)
img_path = self.data_path + lab.split('.txt')[0] + '.jpg'
# image = cv2.imread(img_path)
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = Image.open(img_path)
if self.transforms is not None:
image, boxes = self.transforms(image, boxes)
# 해당 box가 grid cell에서 몇번째 cell에 속하는지
label_matrix = torch.zeros((self.S, self.S, self.C + (5 * self.B)))
for box in boxes:
class_label, x, y, w, h = box.tolist()
class_label = int(class_label)
# i,j는 셀의 행과 열을 의미함
# y가 세로로 이동하니까 행
i,j = int(self.S * y), int(self.S * x) # box가 속해있는 셀의 위치
x_cell, y_cell = self.S * x - j, self.S * y - i # box가 해당 cell에서 어느 위치에 있는지 파악
w_cell, h_cell = (w * self.S, h * self.S)
# 한 그리드 당 하나의 object는 무조건 존재해야하기때문에
if label_matrix[i,j,3] == 0: # 내 클래스가 3개니까 0,1,2 각각 클래스에 대해 존재하는지 여부 다음 3번째 인덱스에서 confidence score.
label_matrix[i,j,3] = 1 # object가 존재한다고 1로 표현
box_coord = torch.tensor([x_cell, y_cell, w_cell, h_cell])
label_matrix[i,j,4:8] = box_coord # 좌표 정보도 추가
label_matrix[i,j,class_label] = 1 # 해당 class 라벨이 존재하면 1
return image, label_matrix
def __len__(self):
return len(self.labs) # label이 있는 이미지만
import torchvision.transforms as transforms # torchvision 내에 있는 transforms은 PIL 이미지를 받아서 전처리 수행
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, img, bboxes):
for t in self.transforms: # bboxes는 448로 resize and 정규화
img, bboxes = t(img), bboxes
return img, bboxes
def get_transform():
transform = Compose([transforms.Resize((448, 448)), transforms.ToTensor(),])
return transform
3. Utils
import torch
import torch.nn as nn
import torch
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from collections import Counter
def intersection_over_union(boxes_preds, boxes_labels, box_format="midpoint"):
if box_format == "midpoint": # x_center, y_center, w, h의 포멧일 때
box1_x1 = boxes_preds[..., 0:1] - boxes_preds[..., 2:3] / 2
box1_y1 = boxes_preds[..., 1:2] - boxes_preds[..., 3:4] / 2
box1_x2 = boxes_preds[..., 0:1] + boxes_preds[..., 2:3] / 2
box1_y2 = boxes_preds[..., 1:2] + boxes_preds[..., 3:4] / 2
box2_x1 = boxes_labels[..., 0:1] - boxes_labels[..., 2:3] / 2
box2_y1 = boxes_labels[..., 1:2] - boxes_labels[..., 3:4] / 2
box2_x2 = boxes_labels[..., 0:1] + boxes_labels[..., 2:3] / 2
box2_y2 = boxes_labels[..., 1:2] + boxes_labels[..., 3:4] / 2
if box_format == "corners": # x1,y1,x2,y2의 포멧일 때
box1_x1 = boxes_preds[..., 0:1]
box1_y1 = boxes_preds[..., 1:2]
box1_x2 = boxes_preds[..., 2:3]
box1_y2 = boxes_preds[..., 3:4] # (N, 1)
box2_x1 = boxes_labels[..., 0:1]
box2_y1 = boxes_labels[..., 1:2]
box2_x2 = boxes_labels[..., 2:3]
box2_y2 = boxes_labels[..., 3:4]
x1 = torch.max(box1_x1, box2_x1)
y1 = torch.max(box1_y1, box2_y1)
x2 = torch.min(box1_x2, box2_x2)
y2 = torch.min(box1_y2, box2_y2)
# clamp는 겹치는 부분이 없으면 0으로 취급
intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0)
box1_area = abs((box1_x2 - box1_x1) * (box1_y2 - box1_y1))
box2_area = abs((box2_x2 - box2_x1) * (box2_y2 - box2_y1))
return intersection / (box1_area + box2_area - intersection + 1e-6) # 분모가 0이 되는 것을 방지 하기 위해 1e-6
def non_max_suppression(bboxes, iou_threshold, threshold, box_format="corners"):
# bboxes: [label, prob_score, x1~y2]
# threshold: confidence score
assert type(bboxes) == list
bboxes = [box for box in bboxes if box[1] > threshold] # conf score로 걸러내기
bboxes = sorted(bboxes, key=lambda x: x[1], reverse=True)
bboxes_after_nms = []
while bboxes:
chosen_box = bboxes.pop(0)
bboxes = [
box for box in bboxes
if box[0] != chosen_box[0] # 현재 박스와 다른 박스들의 클래스가 다르면 다른 객체를 담고있는 것이기때문에 선택되어야함
or intersection_over_union(
torch.tensor(chosen_box[2:]),
torch.tensor(box[2:]),
box_format=box_format,
) < iou_threshold # 만약 같은 클래스라도 많이 겹치지 않으면 다른 객체를 지정하고 있는 것이기때문에 선택되어야함
]
bboxes_after_nms.append(chosen_box)
return bboxes_after_nms
def mean_average_precision(pred_boxes, true_boxes, iou_threshold = 0.5,
box_format = 'midpoint', num_classes = 3):
# bboxes: [train_idx, label, prob_score, x1, y1, x2, y2]
average_precisions = []
epsilon = 1e-6
for c in range(num_classes):
detections = []
ground_truths = []
for detection in pred_boxes:
if detection[1] == c: # 클래스를 맞췄으면
detections.append(detection)
for true_box in true_boxes:
if true_box[1] == c:
ground_truths.append(true_box)
# 0번 이미지에 객체가 3개, 1번에 5개면 amount_bboxes = {0:3, 1:5}
amount_bboxes = Counter([gt[0] for gt in ground_truths]) # train_idx를 카운트 = 실제로 해당 이미지 내 객체가 몇개있는지 카운트
for key, val in amount_bboxes.items():
# ammount_bboxes = {0:torch.tensor[0,0,0], 1:torch.tensor[0,0,0,0,0]}
amount_bboxes[key] = torch.zeros(val)
detections.sort(key=lambda x: x[2], reverse=True) # conf_score에 따라 정렬
TP = torch.zeros((len(detections)))
FP = torch.zeros((len(detections)))
total_true_bboxes = len(ground_truths)
if total_true_bboxes == 0: # 원래 겍체가 없는 이미지라면 다음 클래스로 넘어가기
continue
for detection_idx, detection in enumerate(detections):
ground_truth_img = [
bbox for bbox in ground_truths if bbox[0] == detection[0]
] # 해당 이미지에서 예측한 객체가 실제로 몇개인지
# 하나의 예측 객체에 대해서 정답이랑 가장 가까운 박스가 best iou가 된다.
best_iou = 0
for idx, gt in enumerate(ground_truth_img):
iou = intersection_over_union(
torch.tensor(detection[3:]),
torch.tensor(gt[3:]),
box_format=box_format
)
if iou > best_iou:
best_iou = iou
best_gt_idx = idx
if best_iou > iou_threshold: # iou가 높으면 맞췄다고 가정
if amount_bboxes[detection[0]][best_gt_idx] == 0:
TP[detection_idx] = 1
# amount_bboxes -> {0:torch.tensor[0,0,0], 1:torch.tensor[0,0,0,0,0]}
amount_bboxes[detection[0]][best_gt_idx] = 1
else: # 이미 해당 객체를 best로 예측했었는데 또 예측한 거면 FP
FP[detection_idx] = 1
else:
FP[detection_idx] = 1
TP_cumsum = torch.cumsum(TP, dim=0)
FP_cumsum = torch.cumsum(FP, dim=0)
recalls = TP_cumsum / (total_true_bboxes + epsilon)
precisions = torch.divide(TP_cumsum, (TP_cumsum + FP_cumsum + epsilon))
precisions = torch.cat((torch.tensor([1]), precisions)) # precitions의 시작점 1. y축
recalls = torch.cat((torch.tensor([0]), recalls))# recall의 시작점 0. x축
average_precisions.append(torch.trapz(precisions, recalls)) # 각 클래스마다 밑면적 구하기. 적분 수행
return sum(average_precisions) / len(average_precisions)
def get_bboxes(loader, model, iou_threshold, threshold, pred_format = 'cells',
box_format='midpoint', device = 'cuda'):
all_pred_boxes = []
all_true_boxes = []
model.eval()
train_idx = 0
for batch_idx, (x,labels) in enumerate(loader):
x = x.to(device)
labels = labels.to(device)
with torch.no_grad():
predictions = model(x)
batch_size = x.shape[0]
true_bboxes = cellboxes_to_boxes(labels)
bboxes = cellboxes_to_boxes(predictions)
for idx in range(batch_size):
nms_boxes = non_max_suppression(
bboxes[idx],
iou_threshold=iou_threshold,
threshold=threshold,
box_format=box_format
)
for nms_box in nms_boxes:
all_pred_boxes.append([train_idx] + nms_box) # nms box 앞에 train_idx 추가
for box in true_bboxes[idx]:
if box[1] > threshold:
all_true_boxes.append([train_idx] + box)
train_idx += 1
model.train()
return all_pred_boxes, all_true_boxes
# 그리드 셀의 박스를 전체 이미지에 대한 비율로 다시 바꿈
def convert_cellboxes(predictions, S=7):
predictions = predictions.to('cpu')
batch_size = predictions.shape[0]
predictions = predictions.reshape(batch_size, S, S, 13)
bboxes1 = predictions[..., 4:8]
bboxes2 = predictions[..., 9:13]
scores = torch.cat(
(predictions[..., 3].unsqueeze(0), predictions[..., 8].unsqueeze(0)), dim = 0
)
best_box = scores.argmax(0).unsqueeze(-1)
best_boxes = bboxes1 * (1-best_box) + best_box * bboxes2
# arnage(7): 0부터 6까지 값을 갖는 1차원 텐서
# repeat(): 배치사이즈만큼 S번 반복해서 3차원 텐서로
# unsueeeze(-1): 차원을 하나 추가해서 4차원으로
# cell_ihdices[0,2,3]은 첫번째 이미지에서 (2,3)그리드의 box 좌표
cell_indices = torch.arange(S).repeat(batch_size, S, 1).unsqueeze(-1)
x = 1 / S * (best_boxes[..., :1] + cell_indices)
y = 1 / S * (best_boxes[..., 1:2] + cell_indices.permute(0,2,1,3)) # y값을 기준으로 거리 계산
w_y = 1 / S * best_boxes[..., 2:4] # 2차원
converted_bboxes = torch.cat((x, y, w_y), dim=-1)
predicted_class = predictions[..., :3].argmax(-1).unsqueeze(-1)
best_confidence = torch.max(predictions[..., 3], predictions[..., 8]).unsqueeze(
-1
)
converted_preds = torch.cat(
(predicted_class, best_confidence, converted_bboxes), dim=-1
)
return converted_preds # (batch, S, S, (label,confidence, x, y, w_y)) w_y가 2차원이므로 (batch, S, S, 6)
# 전체 이미지에 대한 비율로 바뀐 boxes들을 각각 개별 바운딩 박스로 저장
def cellboxes_to_boxes(out, S=7):
converted_pred = convert_cellboxes(out).reshape(out.shape[0], S * S, -1) # 7,7, 박스 좌표
converted_pred[..., 0] = converted_pred[..., 0].long()
all_bboxes = []
for ex_idx in range(out.shape[0]):
bboxes = []
for bbox_idx in range(S * S):
bboxes.append([x.item() for x in converted_pred[ex_idx, bbox_idx, :]])
all_bboxes.append(bboxes)
return all_bboxes
def save_checkpoint(state, filename="my_checkpoint.pth.tar"):
print("=> Saving checkpoint")
torch.save(state, filename)
def load_checkpoint(checkpoint, model, optimizer):
print("=> Loading checkpoint")
model.load_state_dict(checkpoint["state_dict"])
optimizer.load_state_dict(checkpoint["optimizer"])
4. Loss
import torch
import torch.nn as nn
class YoloLoss(nn.Module):
def __init__(self, S=7, B=2, C=3):
super(YoloLoss, self).__init__()
self.mse = nn.MSELoss(reduction="sum") # 기본은 MSE를 따르고
self.S = S
self.B = B
self.C = C
# 논문에서 언급한대로 object가 없을 때 0.5, localization에 대해 5의 람다값 부여
self.lambda_noobj = 0.5
self.lambda_coord = 5
def forward(self, predictions, target):
# predictions은 (BATCH_SIZE, S*S(C+B*5)) 사이즈를 가지도록 reshape
# [..., 0 ~ 2] class 확률, 3 box1의 confidence score, 4~7 box1 좌표, 8, box2의 confidence scores, 9~12 box2 좌표
predictions = predictions.reshape(-1, self.S, self.S, self.C + self.B * 5)
# cell당 2개의 box가 그려지기때문에 2개의 iou
iou_b1 = intersection_over_union(predictions[..., 4:8], target[..., 4:8])
iou_b2 = intersection_over_union(predictions[..., 9:13], target[..., 4:8]) # target은 동일함
ious = torch.cat([iou_b1.unsqueeze(0), iou_b2.unsqueeze(0)], dim=0)
# best
iou_maxes, bestbox = torch.max(ious, dim=0) # 두 박스 중 iou 최대값 리턴, 몇번째 박스의 iou가 최대인지 idx 리턴,
exists_box = target[..., 3].unsqueeze(3) # 해당 박스에 객체가 존재하는지. 존재하면 1, 아니면 0. 논문에서 말한 1obj_i
# 1. localization loss 구하기
# 두 박스 중 bestbox를 고르는 방법
box_predictions = exists_box * (
(
bestbox * predictions[..., 9:13] # best box가 1이라면 두번째 박스를 골라야하니까 9:13을
+ (1 - bestbox) * predictions[..., 4:8] # best box가 0이라면 첫번째 박스를 골라야하니까 4:8을 선택. 굳이 복잡하게 표현
)
)
box_targets = exists_box * target[..., 4:8]
# 제곱근 사용
# w,h가 항상 양수가 되도록 sign,abs함수로 조정 후 제곱근.
box_predictions[..., 2:4] = torch.sign(box_predictions[..., 2:4]) * torch.sqrt(
torch.abs(box_predictions[..., 2:4] + 1e-6)
)
box_targets[..., 2:4] = torch.sqrt(box_targets[..., 2:4])
# 제곱근을 취한 두 박스 간의 Loss 계산
box_loss = self.mse(
torch.flatten(box_predictions, end_dim=-2),
torch.flatten(box_targets, end_dim=-2),
)
# 2. Object loss. confidence score
# 위와 같은 방법으로 best box 선정. 9는 box2의 객체 존재여부, 3은 box 1의 객체 존재여부
# best box를 선정했으니까 해당 box에 대해서만 confidence score 손실 계산
# 두 박스 중 best box가 객체가 있는 박스라고 손실을 계산했으면 나머지 다른 하나는 객체가 없는 박스(no_obj_loss)로 손실을 계산해야함.
# 한 그리드당 하나의 객체만 탐지해야하기때문에.
pred_box = (
bestbox * predictions[..., 9:10] + (1 - bestbox) * predictions[..., 3:4]
)
object_loss = self.mse(
torch.flatten(exists_box * pred_box),
torch.flatten(exists_box * target[..., 3:4]), # 타겟의 idx3은 객체 존재여부
)
# 3. No Object loss
# No object에 대해서는 두 박스 모두 계산.
# start_dim 1은 첫번째 차원을 제외하고 그 이후 차원 flatten
# Object loss와는 별개로 두 박스 모두에 객체가 없을 수 있기때문에 여기서는 두 박스를 모두 사용해 손실 계산.
no_object_loss = self.mse(
torch.flatten((1 - exists_box) * predictions[..., 3:4], start_dim=1),
torch.flatten((1 - exists_box) * target[..., 3:4], start_dim=1),
)
no_object_loss += self.mse(
torch.flatten((1 - exists_box) * predictions[..., 8:9], start_dim=1),
torch.flatten((1 - exists_box) * target[..., 3:4], start_dim=1)
)
# 4. Class loss
# 모든 클래스들의 확률값을 사용해서 손실계산
class_loss = self.mse(
torch.flatten(exists_box * predictions[..., :3], end_dim=-2,),
torch.flatten(exists_box * target[..., :3], end_dim=-2,),
)
loss = (
self.lambda_coord * box_loss # 논문 손실함수의 첫번째 두번째 부분
+ object_loss # 세번째 부분
+ self.lambda_noobj * no_object_loss # 네번째 부분
+ class_loss # 다섯번째 부분
)
return loss
5. Training
import torch
import torchvision.transforms as transforms
import torch.optim as optim
import torchvision.transforms.functional as FT
from tqdm import tqdm
from torch.utils.data import DataLoader
# from model import Yolov1
# from dataset import VOCDataset
# from utils import (
# non_max_suppression,
# mean_average_precision,
# intersection_over_union,
# cellboxes_to_boxes,
# get_bboxes,
# plot_image,
# save_checkpoint,
# load_checkpoint,
# )
# from loss import YoloLoss
seed = 2023
torch.manual_seed(seed)
# Hyperparameters etc.
LEARNING_RATE = 2e-5
DEVICE = "cuda" if torch.cuda.is_available else "cpu"
BATCH_SIZE = 16
WEIGHT_DECAY = 0
EPOCHS = 50
NUM_WORKERS = 2
PIN_MEMORY = True
LOAD_MODEL = False
LOAD_MODEL_FILE = "yolo_test.pth.tar"
def train_fn(train_loader, model, optimizer, loss_fn):
loop = tqdm(train_loader, leave=True)
mean_loss = []
for batch_idx, (x, y) in enumerate(loop):
x, y = x.to(DEVICE), y.to(DEVICE)
out = model(x)
loss = loss_fn(out, y)
mean_loss.append(loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
# update progress bar
loop.set_postfix(loss=loss.item())
print(f"Mean loss was {sum(mean_loss)/len(mean_loss)}")
def main():
model = YOLOv1(num_split_cell=7, num_boxes=2, num_classes=3).to(DEVICE)
optimizer = optim.Adam(
model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY
)
loss_fn = YoloLoss()
if LOAD_MODEL:
load_checkpoint(torch.load(LOAD_MODEL_FILE), model, optimizer)
train_dataset = SoccerDataset(
data_path=TR_DATA_PATH,
label_path=TR_LAB_PATH,
S = 7,
B = 2,
C = 3,
transforms=get_transform()
)
val_dataset = SoccerDataset(
data_path=VAL_DATA_PATH,
label_path=VAL_LAB_PATH,
S = 7,
B = 2,
C = 3
)
train_loader = DataLoader(
dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True,
)
val_loader = DataLoader(
dataset=val_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=True,
drop_last=True,
)
for epoch in range(EPOCHS):
# for x, y in train_loader:
# x = x.to(DEVICE)
# for idx in range(8):
# bboxes = cellboxes_to_boxes(model(x))
# bboxes = non_max_suppression(bboxes[idx], iou_threshold=0.5, threshold=0.4, box_format="midpoint")
# plot_image(x[idx].permute(1,2,0).to("cpu"), bboxes)
# import sys
# sys.exit()
pred_boxes, target_boxes = get_bboxes(
train_loader, model, iou_threshold=0.5, threshold=0.4, device=DEVICE
)
mean_avg_prec = mean_average_precision(
pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint"
)
print(f"Train mAP: {mean_avg_prec}")
#if mean_avg_prec > 0.9:
# checkpoint = {
# "state_dict": model.state_dict(),
# "optimizer": optimizer.state_dict(),
# }
# save_checkpoint(checkpoint, filename=LOAD_MODEL_FILE)
# import time
# time.sleep(10)
train_fn(train_loader, model, optimizer, loss_fn)
torch.save(model.state_dict(),f'{WEIGHTS_PATH}yolov1_{EPOCHS}.pt')
if __name__ == "__main__":
main()
6. Evaluation
test_dataset = SoccerDataset(
data_path=TEST_DATA_PATH,
label_path=TEST_LAB_PATH,
S = 7,
B = 2,
C = 3,
transforms = get_transform()
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS,
pin_memory=PIN_MEMORY,
shuffle=False,
drop_last=True
)
model = YOLOv1(num_split_cell=7, num_boxes=2, num_classes=3).to(DEVICE)
model.load_state_dict(torch.load(f'{WEIGHTS_PATH}yolov1_{EPOCHS}.pt'))
pred_boxes, target_boxes = get_bboxes(
test_loader, model, iou_threshold=0.5, threshold=0.4, device=DEVICE
)
mean_avg_prec = mean_average_precision(
pred_boxes, target_boxes, iou_threshold=0.5, box_format="midpoint"
)
print(f"Test mAP: {mean_avg_prec}") # 0.34의 mAP 기록. 과적합
im, t = test_dataset[0]
im, t = im.unsqueeze(0), t.unsqueeze(0)
model.eval()
with torch.no_grad():
prediction = model(im.to(DEVICE))
prediction = prediction.reshape(-1, 7, 7, 3 + 2 * 5)
target = t.to(DEVICE)
# cell당 2개의 box가 그려지기때문에 2개의 iou
iou_b1 = intersection_over_union(prediction[..., 4:8], target[..., 4:8])
iou_b2 = intersection_over_union(prediction[..., 9:13], target[..., 4:8]) # target은 동일함
ious = torch.cat([iou_b1.unsqueeze(0), iou_b2.unsqueeze(0)], dim=0)
# best
iou_maxes, bestbox = torch.max(ious, dim=0) # 두 박스 중 iou 최대값 리턴, 몇번째 박스의 iou가 최대인지 idx 리턴,
exists_box = target[..., 3].unsqueeze(3) # 해당 박스에 객체가 존재하는지. 존재하면 1, 아니면 0. 논문에서 말한 1obj_i
# 1. localization loss 구하기
# 두 박스 중 bestbox를 고르는 방법
box_predictions = exists_box * (
(
bestbox * prediction[..., 9:13] # best box가 1이라면 두번째 박스를 골라야하니까 9:13을
+ (1 - bestbox) * prediction[..., 4:8] # best box가 0이라면 첫번째 박스를 골라야하니까 4:8을 선택.
)
)
class_predictions = exists_box * prediction[...,:3]
img = np.array(im.squeeze(0).permute((1, 2, 0)) * 255).astype(np.uint8).copy()
for y_idx, y in enumerate(box_predictions[0]):
for x_idx, x in enumerate(y):
if x[2] != 0: # width가 0이 아니면 박스가 그려진 것
x_cell, y_cell, w_cell, h_cell = box_predictions[0,y_idx,x_idx].tolist()
label = torch.argmax(class_predictions[0,y_idx,x_idx]).tolist()
x1, y1 = (x_cell + x_idx) / 7, (y_cell + y_idx) / 7 # x_cell, y_cell 형식을 x_center, y_center 형식으로 변환
w, h = w_cell / 7, h_cell / 7
float_x_center = 448 * x1
float_y_center = 448 * y1
float_width = 448 * w
float_height = 448 * h
x1 = int(float_x_center - float_width / 2)
y1 = int(float_y_center - float_height / 2)
x2 = int(float_width) + x1
y2 = int(float_height) + y1
cv2.rectangle(img, (x1, y1), (x2,y2), color = (0,255,0), thickness = 2)
cv2.putText(img, str(label), (x1,y1+10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color = (255,0,0), thickness= 3)
plt.imshow(img)
최종적으로 학습에서는 0.85의 mAP, test에서는 0.34의 mAP를 기록하며 과적합된 모델이다. 아마 데이터셋의 크기가 작은 것이 과적합에 큰 영향을 미쳤을 것 같다. 위의 이미지를 보면 box prediction 부분에서 아쉬움을 보인다. 하지만, 학습 속도는 한 에포크당 20초가 걸렸고 이는 SSD와 비슷한 속도이면서 Faster R-CNN보다 약 10배이상 빠른 속도이다.
논문에서 언급한대로 Faster R-CNN에 비해 mAP 측면에서는 좀 아쉽긴하지만 속도면에서는 매우 빠른 장점을 가진 YOLOv1 모델이다.