yolov5로 토마토 열매 인식 모델 만들기

YOLOv5는 실시간 객체 탐지를 위한 딥러닝 모델 중 하나이다. YOLO는 "You Only Look Once"의 약자로, 한 번의 순방향 전파를 통해 이미지 전체를 분석하고 객체를 탐지한다. YOLOv5는 다양한 환경에서 사용되며, 컴퓨터 비전, 자율주행차, 보안 등 여러 분야에서 객체 감지 및 추적에 활용될 수 있다.

 

 

캐글에서 토마토 열매 데이터셋을 다운받는다.(https://www.kaggle.com/datasets/andrewmvd/tomato-detection/data)

 

다운받은 파일은 토마토 사진들과 각 사진들의 라벨 데이터로 구성되어 있다.

 

 

구글 드라이브에서 새 코랩 파일을 만든다.

 

왼쪽에서 드라이브를 마운트한다.

 

하드웨어 가속기를 설정한다.

 

내 드라이브 폴더에 적당한 폴더를 생성하고 그리로 이동한다.

%cd drive/MyDrive/tomato-yolov5

 

 

깃허브에서 yolov5 프로젝트를 클론해온다.(https://github.com/ultralytics/yolov5)

!git clone https://github.com/ultralytics/yolov5.git

 

 

yolov5 폴더로 이동하고 requirements.txt로 패키지들을 설치한다.

%cd yolov5
!pip install -r requirements.txt

 

 

데이터셋 파일들을 코랩에 업로드하기 전에 약간의 처리가 필요하다.

 

데이터셋 파일들은 895개의 사진 파일과 각 사진에 대한 라벨링 파일이 xml로 되어있다.

하지만 여기서 사용하는 yolov5 프로젝트는 xml로 된 라벨링 파일을 사용하지 못한다. 그래서 txt로 변환해야 한다.

 

컴퓨터의 아무곳에나 적당한 폴더를 만들고 깃허브에서 XmlToTxt 프로젝트를 클론해온다.(https://github.com/isabek/XmlToTxt)

git clone https://github.com/isabek/XmlToTxt.git

 

 

 

프로젝트의 xml 폴더에 데이터셋의 라벨링 파일들을 전부 넣는다.

 

 

classes.txt의 내용을 전부 지우고 tomato 클래스를 추가해준다.

 

 

xmltotxt.py 파일을 실행해서 라벨링 파일들을 변환한다.

python xmltotxt.py -xml xml -out out

 

 

변환된 파일들은 out 폴더에 들어가있다.

 

 

이제 데이터셋을 train 셋과 validation 셋으로 나누어야 한다.

 

train 폴더와 valid 폴더를 생성하고 각각의 폴더에 images 폴더와 labels 폴더를 생성한다.

 

 

원본 데이터셋에 있는 895개의 이미지와 라벨링 파일들을 train 폴더와 valid 폴더에 각각 나누어 넣는다. 여기서는 0~799까지 train에, 800~894까지 valid에 넣었다.

 

 

train과 valid 폴더를 압축한다.

 

 

압축한 두 파일들을 코랩의 yolov5 폴더에 업로드한다.

 

 

업로드한 두 파일들의 압축을 푼다.

!unzip -qq './train.zip' -d './train'
!unzip -qq './valid.zip' -d './valid'

 

 

yolov5폴더에 data.yaml 파일을 생성하고 다음 내용을 넣는다.

train: ../train/images
val: ../valid/images

nc: 1
names: ['tomato']

 

 

train.py 파일을 실행하여 학습을 진행한다.

!python train.py --img 512 --batch 16 --epochs 50 --data ./data.yaml --cfg ./models/yolov5s.yaml --weight yolov5s.pt --name tomato_yolov5_results

• --img : 학습할 사진의 크기
• --batch : 전체데이터를 얼만큼 쪼개서 학습 할 건지를 지정
• --epochs : 전체 데이터를 몇 회 반복 학습 할 건지를 지정
• --data : 데이터셋의 경로정보를 가지고 있는 yaml 파일 지정
• --cfg : 학습에 사용할 모델의 설정파일(yaml) 지정
• --weight : 학습에 사용할 모델 파일 지정
• --name : 학습결과를 저장할 폴더명

결과파일이 지정한 폴더에 저장되었다.

 

best.pt 파일이 학습 과정에서 생성된 최적의 가중치를 가진 파일이다. 이 파일을 코랩에서 다운받는다.

 

 

다시 컴퓨터의 적당한 곳에 yolov5 프로젝트를 클론해온다.

git clone https://github.com/ultralytics/yolov5.git

 

 

requirements.txt 파일로 yolov5에 필요한 패키지들을 설치한다.

pip install requirements.txt

 

 

프로젝트 폴더에 best.pt 파일과 실행시킬 이미지를 넣고 주피터 노트북 파일을 생성한다.

 

모델 실행에 사용한 이미

 

 

 

주피터 노트북 파일에 다음 내용을 넣는다.

import torch 

import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath

#torch 모델 파일을 불러와서 함수를 만든다.
model = torch.hub.load('', 'custom', path='best.pt', source='local')


#객체를 탐지할 이미지 경로
img = "tomato.jpg"


#객체 탐지를 실행하여 결과를 저장한다.
results = model(img)


#탐지된 객체들의 좌표와 확률 값이 들어있다.
result_list = results.xyxy[0].tolist()


#토마토일 확률이 80%이상인 것만 남긴다.
fruit_num = len(list(filter(lambda x : x[4]>0.8,results.xyxy[0])))


#토마토 개수를 출력한다.
print(fruit_num)


#실행결과를 보여준다.
results.show()

 

실행결과는 이렇게 나온다.

 

분명 예전에 했을땐 없는 오류였는데 최근에 다시 모델을 만들어서 실행시키니까 이런 오류가 생겼다..

 

https://github.com/ultralytics/yolov5/issues/10240

NotImplementedError: cannot instantiate 'PosixPath' on your system · Issue #10240 · ultralytics/yolov5

 

윈도우에서 실행하려고 하면 생기는 오류 같은데 이 코드를 실행시키면 해결된다.

import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath