한남대 졸업작품

AI Camera 제작을 위한 Yolov5n custom train

JUNGCHAN LEE 2024. 11. 13. 22:05

제작한 PCB 보드에서 작동하는 AI 를 만들기 위해 낮은 weight를 가지는 custom training 이 필요했다. 작업은 개인 리눅스 서버에서 진행했으며, COLAB 과 동작이 같지 않으므로 COLAB 에 적용시에는 알맞은 명령어가 필요하다.

~$ mkdir capstone1107
~$ cd capstone1107

일단 capstone 파일을 만들어주었다.

 

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

깃허브에서 yolov5 를 클론해왔다. v5n으로 custome training 해주기 위한 Data는 인터넷 opensource 를 활용했으며, Roboflow에서 annotating을 진행했다.  Roboflow annotation은 insturction 을 잘 따라가면 쉽기 때문에 따로 다루지는 않는다. 다만 annotation dataset은 200장으로 정하고 학습시켰다.

Dataset의 size는 800*480이다. 제작할 PCB에 올라갈 LCD와 Camera(OV5640)의 Capture resolution을 800x480으로 설정할 것이기 때문에 그렇게 정했다. 

 

~$ curl -L "https://app.roboflow.com/ds/{api-key}" > roboflow.zip; unzip roboflow.zip;

리눅스에서 curl으로 dataset을 다운로드 받아왔다. 

 

~$ cd yolov5
~$ ls

yolov5 폴더 안으로 들어가서 파일 리스트를 확인했다. 

 

 

처음에는 yolov5n-seg.pt 파일이 없었다. yolov5n-seg를 wget 해왔다. 깃허브에 업로드 되어있다. 

~$ wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n-seg.pt

 

~$ pip install -r requirements.txt 
~$ python3 train.py --img 800 --epochs 100 --data ../data.yaml --weights yolov5n-seg.pt --cfg models/yolov5n.yaml

위의 코드를 실행시켜서 runs에 내용을 넣어줬다. 학습된 wegiht 파일(pytourch type)과  각종 학습 결과가 나왔다. 

 

STM32CUBE-AI는 pytoruch 를 사용할 수 없다. Onnx나 tflite 파일로 바꿔줘야한다. 

~$ python3 export.py --include onnx --weights {file_path}

이 코드를 실행시켜 onnx 타입으로 전환해주었다.

 

 

검증파일이다. 어느정도 인식률이 있는 것으로 보인다. 

 

 

다음과 같은 Result 파일을 생성했다.