YOLO v8 預設模型可辨識 COCO 資料集的 80 種物件。若要辨識其它物件,就需準備自訂資料集自行訓練
安裝套件
請記得安裝 ultralytics 時,自動安裝的 torch 無法啟動 GPU,所以請由如下指令安裝套件
pip install pip install ultralytics labelimg
pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio===2.0.2+cu118 -f https://download.pytorch.org/whl/cu118/torch_stable.html
下載資料集
自訂模型最麻煩的就是資料集準備,要收集大量圖片,然後標示,光這項作業就要耗費相當長的時間。還好 Roboflow 提供許多 公用資料集。請由公用資料集的 Roboflow 100 進入,然後搜尋 “Dollar bill detection”,或由如下網址下載 https://universe.roboflow.com/currency-7mtqe/dollarbilldetection。
按下 Download this Dataset,再選擇 v8 版本,將 Dollar Bill Detection.v24-raw-images.yolov8.zip 存放在專案根目錄下。
目錄安排
將 .zip 檔解壓縮檔案,就會產生 “Dollar Bill Detection.v24-raw-images.yolov8” 目錄,請改名為 “Dollar”。
datasets 設定
datasets 又稱為訓練資料集。也就是包含了 train/valid/test 三個目錄,每個目錄又有 images/labels 二個目錄。專案在安裝 utlralytics 套件時,會在 C:\Users\登入者\AppData\Roaming\Ultralytics\settings.yaml 裏將 datasets_dir 設定為如下
settings_version: 0.0.4
datasets_dir: E:\python_ai\yolov8\datasets
weights_dir: weights
如果新專案(b)又安裝一次的話,那麼 datasets_dir 就會被改成 e:\python_ai\b\datasets,那麼舊專案就找不到 datasets了。所以打開 Dollar/data.yaml , 將前三行的 “../” 改成絕對路徑。
#train: ../train/images #val: ../valid/images #test: ../test/images
train: e:/python_ai/yolov8/Dollar/train/images val: e:/python_ai/yolov8/Dollar/valid/images test: e:/python_ai/yolov8/Dollar/test/images
下載預訓練權重
到 https://github.com/ultralytics/ultralytics 網站,往下拉即有預訓練模型可供下載。預訓練模型有yolov8n、yolov8s、yolov8m、yolov8l、yolov8x 共 5 種,愈往下愈精準但速度愈慢。請將下載的 .pt 檔儲存在專案目錄下。
使用 yolo.exe 訓練模型
進入 Terminal 執行如下指令
yolo task=detect mode=train model=./yolov8n.pt data=./Dollar/data.yaml epochs=200 imgsz=640
以 RTX 3080Ti , 在 epochs 約 193 時就會自動停止。
預訓練模型使用 yolov8n.pt 大約需要 10 分鐘才能訓練完成。
訓練完成,會在 ./runs/detect/train/weights 下產生 best.pt 模型,這個就是日後要偵測的主要東西。
使用 Python 訓練模型
在專案下新增 train.py 檔,由 YOLO 戴入預訓練模型 yolov8n.pt 產生 model 物件,再由 model.train() 即可開始訓練。
請注意這段程式碼一定要寫在 if __name__ 的區塊中,否則會出現需使用 fork 執行子行程的錯誤。
import os
import shutil
import time
from ultralytics import YOLO
#訓練模型時,一定要放在 __name__ 區塊內
#否則會出現需使用 fork來執行子行程的錯誤
if __name__=='__main__':
train_path="./runs/detect/train"
if os.path.exists(train_path):
shutil.rmtree(train_path)
model = YOLO("yolov8n.pt")
print("開始訓練 .........")
t1=time.time()
model.train(data="./Dollar/data.yaml", epochs=200, imgsz=640)
t2=time.time()
print(f'訓練花費時間 : {t2-t1}秒')
path=model.export()
print(f'模型匯出路徑 : {path}')
開始偵測
如下代碼可以開始偵測,並手動繪制方框
import os.path import platform import cv2 from PIL import Image, ImageFont, ImageDraw import numpy as np from ultralytics import YOLO def text(img, text, xy=(0, 0), color=(0, 0, 0), size=12): pil = Image.fromarray(img) s = platform.system() if s == "Linux": font =ImageFont.truetype('/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc', size) elif s == "Darwin": font = ImageFont.truetype('....', size) else: font = ImageFont.truetype('simsun.ttc', size) ImageDraw.Draw(pil).text(xy, text, font=font, fill=color) return np.array(pil) model=YOLO('./best.pt') #model=YOLO('./runs/detect/train/weights/best.pt') path='./dollar/test/images' files=[] for file in os.listdir(path): files.append(os.path.join(path, file)) print(files) import pylab as plt for i in range(10,18): file=files[i] img=cv2.imdecode(np.fromfile(file, dtype=np.uint8), cv2.IMREAD_COLOR)[:,:,::-1].copy() results=model.predict(img, save=False) names=[results[0].names[int(i.cpu().numpy())] for i in results[0].boxes.cls] #沒顯卡的人 #names = [results[0].names[int(i.numpy())] for i in results[0].boxes.cls] boxes=results[0].boxes.xyxy for box, name in zip(boxes, names): box=box.cpu().numpy() #box = box.numpy() x1 = int(box[0]) y1 = int(box[1]) x2 = int(box[2]) y2 = int(box[3]) img=cv2.rectangle(img,(x1, y1), (x2, y2), (0,255,0) , 3) img=text(img, name, (x1, y1), (0,0,255),40) plt.subplot(2,4, i+1-10) plt.axis("off") plt.imshow(img) plt.show()