YOLOv8로 object detection 해보기

Object Detection with YOLOv8

YOLOv8을 이용하여 object detection을 수행해볼 것이다. Detection할 class는 barcode로 정했다.
YOLOv8를 통해 Classification, Detection, Segmentation 등을 수행할 수 있다. YOLOv8에 대한 자세한 설명은 아래 링크에서 확인할 수 있다.

https://github.com/ultralytics/ultralytics?tab=readme-ov-file

---

Dataset

우선 Dataset을 수집해야 한다. 집에 있는 굴러다니는 제품들 사진과 이전에 갤러리에 있던 바코드의 흔적이 있는 사진들을 종합하여 약 100장의 사진을 모아서 작업 환경으로 옮겨준다.

Preprocessing

하지만 사진별로 size도 다르고 HEIC인 파일도 있다. 우선 HEIC 파일을 png로 바꿔줄 것이다.

import os
import pyheif
from PIL import Image

def convert_heic_to_png(heic_file_path, output_file_path):
    heif_file = pyheif.read(heic_file_path)
    image = Image.frombytes(
        heif_file.mode, 
        heif_file.size, 
        heif_file.data,
        "raw",
        heif_file.mode,
        heif_file.stride,
    )
    image.save(output_file_path, "PNG")
    
    def convert_all_heic_to_png_in_folder(folder_path, output_folder_path):
    if not os.path.exists(output_folder_path):
        os.makedirs(output_folder_path)

    for filename in os.listdir(folder_path):
        if filename.lower().endswith('.heic'):
            heic_file_path = os.path.join(folder_path, filename)
            output_file_name = os.path.splitext(filename)[0] + ".png"
            output_file_path = os.path.join(output_folder_path, output_file_name)
            
            convert_heic_to_png(heic_file_path, output_file_path)
            print(f"{filename} -> {output_file_name} 변환 완료")
            
folder_path = "folder_path"
output_folder_path = "output_folder_path"

convert_all_heic_to_png_in_folder(folder_path, output_folder_path)                        

IMG_1989.HEIC -> IMG_1989.png 변환 완료
IMG_2044.HEIC -> IMG_2044.png 변환 완료
IMG_1990.HEIC -> IMG_1990.png 변환 완료
IMG_1995.HEIC -> IMG_1995.png 변환 완료
IMG_2050.HEIC -> IMG_2050.png 변환 완료
IMG_2055.HEIC -> IMG_2055.png 변환 완료
IMG_2008.HEIC -> IMG_2008.png 변환 완료
IMG_2056.HEIC -> IMG_2056.png 변환 완료
IMG_2026.HEIC -> IMG_2026.png 변환 완료
IMG_1991.HEIC -> IMG_1991.png 변환 완료
IMG_2043.HEIC -> IMG_2043.png 변환 완료
IMG_1992.HEIC -> IMG_1992.png 변환 완료
IMG_2051.HEIC -> IMG_2051.png 변환 완료
IMG_2021.HEIC -> IMG_2021.png 변환 완료
IMG_1997.HEIC -> IMG_1997.png 변환 완료
IMG_2028.HEIC -> IMG_2028.png 변환 완료
IMG_1987.HEIC -> IMG_1987.png 변환 완료
IMG_6766.HEIC -> IMG_6766.png 변환 완료
IMG_2049.HEIC -> IMG_2049.png 변환 완료
IMG_2041.HEIC -> IMG_2041.png 변환 완료
IMG_5662.HEIC -> IMG_5662.png 변환 완료
IMG_2037.HEIC -> IMG_2037.png 변환 완료
IMG_1996.HEIC -> IMG_1996.png 변환 완료
IMG_2014.HEIC -> IMG_2014.png 변환 완료
IMG_2012.HEIC -> IMG_2012.png 변환 완료
...
IMG_2006.HEIC -> IMG_2006.png 변환 완료
IMG_2031.HEIC -> IMG_2031.png 변환 완료
IMG_2027.HEIC -> IMG_2027.png 변환 완료
IMG_2005.HEIC -> IMG_2005.png 변환 완료

그 후 Resize & Crop 해준다.

def resize_and_crop(image, size):
    img_ratio = image.width / image.height
    target_ratio = size[0] / size[1]

    if img_ratio > target_ratio:
        
        new_height = size[1]
        new_width = int(new_height * img_ratio)
    else:
        
        new_width = size[0]
        new_height = int(new_width / img_ratio)

   
    image = image.resize((new_width, new_height), Image.LANCZOS)

    
    left = (new_width - size[0]) / 2
    top = (new_height - size[1]) / 2
    right = (new_width + size[0]) / 2
    bottom = (new_height + size[1]) / 2

    return image.crop((left, top, right, bottom))

def resize_and_crop_all_png_in_folder(folder_path, output_folder_path, size=(512, 512)):
    
    if not os.path.exists(output_folder_path):
        os.makedirs(output_folder_path)
    
    for filename in os.listdir(folder_path):
        if filename.lower().endswith('.png'):
            image_file_path = os.path.join(folder_path, filename)
            output_file_path = os.path.join(output_folder_path, filename)

            
            image = Image.open(image_file_path)

            
            image = resize_and_crop(image, size)

            
            image.save(output_file_path, "PNG")
            
            print(f"{filename} Resize & Crop 완료")
            
            
folder_path = "folder_path" 
output_folder_path = "output_folder_path"
resize_and_crop_all_png_in_folder(folder_path, output_folder_path, size=(512, 512))            

0206071700458023214460.png Resize & Crop 완료
IMG_1043.PNG Resize & Crop 완료
IMG_5689.png Resize & Crop 완료
IMG_2005.png Resize & Crop 완료
IMG_1989.png Resize & Crop 완료
IMG_2029.png Resize & Crop 완료
IMG_6763.png Resize & Crop 완료
IMG_2032.png Resize & Crop 완료
IMG_2006.png Resize & Crop 완료
IMG_2011.png Resize & Crop 완료
IMG_2034.png Resize & Crop 완료
IMG_2002.png Resize & Crop 완료
IMG_2045.png Resize & Crop 완료
IMG_2037.png Resize & Crop 완료
IMG_2026.png Resize & Crop 완료
IMG_1988.png Resize & Crop 완료
IMG_1987.png Resize & Crop 완료
IMG_1985.png Resize & Crop 완료
IMG_1997.png Resize & Crop 완료
IMG_2028.png Resize & Crop 완료
IMG_2044.png Resize & Crop 완료
IMG_5602.png Resize & Crop 완료
IMG_6498.png Resize & Crop 완료
IMG_1979.png Resize & Crop 완료
IMG_2053.png Resize & Crop 완료
...
IMG_2055.png Resize & Crop 완료
IMG_3297.png Resize & Crop 완료
0401081150498043609827.png Resize & Crop 완료
IMG_2041.png Resize & Crop 완료

최종 결과

---

Labeling

이제 전처리된 Data들을 라벨링 해줘야한다. 여러 방법이 있지만 Roboflow에서 라벨링을 진행하기로 결정했다.
라벨링을 위한 프로젝트를 생성하고, Data들을 업로드 해준다.

이제 사진들 하나하나 라벨링 해줘야 한다...
100장 정도라서 오래 걸리진 않을 것이다.
실제 수많은 사진들을 라벨링 해야하는 경우가 있는데, 가장 오래걸리고 중요한 작업이기도 하다.

모두 라벨링을 해줬다.

이제 데이터들을 Train, Valid, Test 데이터로 나눠준다.
Train : Valid : Test = 7 : 2 : 1 로 나눠주었다.

---

Export & Download

본격적으로 데이터셋을 모델에 이용할 수 있도록 하기 전에 몇가지 작업을 더 진행할 수 있다.
Roboflow에서는 최종적으로 데이터셋을 Export 하기 전에 몇가지를 더 Preprocessing 할 수 있는 기능을 지원한다.
아까 Roboflow에 데이터들을 업로드 하기 전에 이미 512X512로 Resize 해줬기 때문에 지금 단계에서 Resize는 생략해줄 것이다.

그 다음에는 데이터 증강을 할 수 있는 단계가 있다.

어떤 방식으로 증강을 할 것인지 선택할 수 있다.
이번에도 데이터 증강은 생략할 것이다.

모든 단계를 끝낸 후 조금 기다리다 보면 이제 최종 데이터셋을 사용할 수 있게 된다!
우측 상단에 있는 Download Dataset을 눌러준다.

YOLOv8 모델을 사용할 것이기 때문에 Format을 YOLOv8로 설정해준다.
그 후에 이 데이터셋을 직접 컴퓨터로 다운로드 할지, 코드를 통해 다운로드 할지 선택할 수 있다. 이 포스팅에서는 코드를 통해 다운로드 할 것이다.

이런 식으로 코드가 제공된는데, 본인이 원하는 코드를 선택한 후 복사해서 입력하면 된다.

코드를 실행 후 다운르도된 데이터셋 폴더(보통 이전에 입력한 프로젝트 이름으로 생성됨)로 들어가면 data.yaml 파일을 찾을 수 있다.
이 파일에는 class가 몇개인지, 어떤 class가 있는지, train, valid, test에 쓰일 이미지가 어떤 경로에 있는지 등에 대한 정보들을 담고있다. 경로를 올바르게 수정해줄 것이다.

이런 식으로 수정해준다.

---

Train

이제 본격적으로 학습을 진행해보겠다.

from ultralytics import YOLO

YOLOv8 모델을 사용할 것이기 때문에 ultralytics에서 YOLO를 import 해준다. 만약 ultralytics가 설치되어있지 않다면, 터미널에 아래 코드를 입력해서 설치하면 된다.

pip install ultralytics

YOLOv8에는 YOLOv8n, YOLOv8s, YOLOv8m, YOLOv8l, YOLOv8x가 있다. YOLOv8n이 가장 단순한 모델이고, YOLOv8x가 가장 복잡한 모델이다.
모델이 복잡해질수록 GPU Memory 사용량, 추론 속도, 성능이 증가한다.

실시간 추론 속도가 중요한 프로젝트에서는 YOLOv8n, YOLOv8s 쪽에서 선택하면 되고, 속도보다는 성능이 더 중요한 프로젝트에서는 YOLOv8l, YOLOv8x 같은 모델을 선택하면 된다. 밸런스가 잡힌 모델은 YOLOv8m이라고 한다. 자신의 목적에 따라 적절한 모델을 선택하면 된다.

이 포스팅에서는 간단한 데이터셋을 통해 학습을 빠르게 진행해볼 것이기 때문에 YOLOv8s 모델을 사용해볼 것이다.

model = YOLO('yolov8s.pt')

아래는 train을 위한 코드이다.

model.train(data='data.yaml/path', epochs = 60, batch = 1, imgsz = 512)

data에는 data.yaml 파일의 경로를 입력해주면 된다.
epochs는 60으로 설정하고, 가벼운 데이터셋을 사용하기 때문에 batch는 1로 설정해준다. imgsz는 데이터셋 이미지 사이즈를 입력해주면 된다.

학습이 모두 완료될 때 까지 기다려준다.

학습이 끝난 후에 생성된 runs/detect/train/weights 경로로 이동하면 best.pt, last.pt 파일이 있는 것을 볼 수 있다.
best.pt는 가장 성능이 좋은 모델 파일이고, last.pt는 학습 과정 중 가장 마지막에 저장된 모델 파일이다.

runs/detect/train에 있는 results.png 파일에선 loss, precision, recall, mAP와 같은 지표들을 확인할 수 있다.

이외 다른 평가 지표들도 해당 경로 내에서 확인할 수 있다.

---

Predict

학습된 모델을 바탕으로 성능이 잘 나오는지 추론해볼 것이다.

model.predict(source='img_path/img.png', save=True)

source에는 추론하고싶은 이미지의 경로를 입력해준다. 잘 탐지하는지 확인하기 위해 save는 True로 설정해준다.

Speed: 95.1ms preprocess, 461.0ms inference, 63.8ms postprocess per image at shape (1, 3, 288, 512)
Results saved to runs/detect/train2

추론이 완료되면 해당 경로로 이동해서 확인하면 된다.

성능이 잘 나오는 것을 확인할 수 있다.

하지만 확인해봐야 할 것이 있다.
이번 포스팅에서는 데이터셋을 100장 정도로 매우 가벼운 데이터셋을 사용하였다. 그렇기 때문에 이후 성능을 잘 낼 수 없는 Underfitting 문제가 발생했을 가능성이 매우 높다.
이를 확인해볼 것이다.

무작위로 선을 그려놓은 그림이다. 절대 바코드일 수가 없는 그림이다.
하지만 이 그림으로 위에서 실행한 코드를 통해 추론을 진행하면,

해당 모델은 이를 0.68의 Confidence Score로 바코드라고 판단해버린다.

이 사진은 사진 전체가 바코드로 채워져 있다. 이 사진으로 추론을 하면,

바코드가 여러개가 있다고 판단해버린다.


이런 결과가 발생한 이유는 가벼운 데이터셋을 사용했기 때문이다. 데이터셋을 더 많이 수집하고 Augmentation을 진행한 후에 이에 맞는 더 복잡한 모델을 이용하면 성능 향상을 기대해 볼 수 있을 것이다.