model을 자세히 보면
어떤 건
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
어떤 건
model = torch.hub.load('ultralytics/yolov5', 'yolov5l')
이런 식으로 되어 있는데 차이는 맨 뒤 알파벳에 있습니다.
s는 smal l은 large를 의미하는데
말 그대로 모델의 사이즈가 스몰이냐 라지이냐의 차이입니다.
사이즈가 작으면 레이어가 적고 레이어가 적으니 가볍고 빠릅니다. 그러나 성능은 조금 떨어집니다.
사이즈가 크면 레이어가 많고 레이어가 많으니 무겁고 느립니다. 그러나 성능은 조금 올라갑니다.
yolov5에서 nano, smal, mid, large, xlarge가 있는 것으로 알고 있습니다.
위에서 만든 이미지 네임 리스트에서 개별 이미지 네임들이 실제 이미지에 해당하는 확장자로 구성된 파일인지 확인합니다.
이미지를 열고 사이즈는 640으로 하여 위의 모델에 전달합니다.
자동차 클래스는 욜로 모델에서 '2'로 되어 있습니다.
'2'에 해당하는 객체가 있다면 그 객체를 표시하는 네모 칸의 박스를 추출합니다.
빨간 줄로 네모 칸의 박스를 해당 객체를 둘러싸게 이미지에 그립니다.
그렇게 결과 이미지를 저장하면 바운딩 박스만을 찾아 잘라 내면 크롭 이미지가 추출되는 것입니다.
추출된 크롭 이미지를 데이터셋으로 하여 정제한 다음 다시 그 데이터로 모델 학습을 진행하면 커스터마이징이 됩니다.
크롭된 이미지를 저장하는 코드입니다.
여기서는 yolov5m을 사용했습니다.
계속 동일한 모델을 사용해야 일관성이 있겠지만 저는 실제로 어딘가에 쓰려다기보다는
이것저것 해 보는 것에 중점을 두고 있기 때문에 각 모델의 차이점을 실제 확인해 보려고 여러 가지 모델을 사용하고 있습니다.
쭉 다시 가 보면 이전 코드와 유사한데 마지막 부분이 조금 다릅니다.
이미지에 바운딩 박스를 그립니다.
im.crop()을 사용해서 x 좌표와 y 좌표를 찍고 잘라 냅니다.
잘라 낸 이미지를 받아서 원하는 폴더 경로에 저장하도록 save()를 사용하면 완성입니다.
크롭된 이미지를 욜로 모델이 참조하고 있는 데이터셋 경로에 저장한 다음 train.py를 돌려서 학습합니다.
그리고 val.py나 predict.py 등을 사용해서 결과를 확인해 보고 수정을 진행하면 됩니다.