YOLOV8自訂COCO模型

      在〈YOLOV8自訂COCO模型〉中尚無留言

YOLO 預訓練權重就是使用 coco 資料集訓練而成的。即然都訓練好了,為什麼還要重作一次呢。這是因為 coco 裏只能辨識 80 種物件,想要新增其它物件時,就要加入 coco 中,重新訓練一次。。

安裝套件

請記得先安裝 Pytorch for cuda,膃後才安裝 ultralytics。

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

pip install pip install ultralytics labelimg

圖片下載

YOLO 使用約 16 萬 4000 張圖片進行訓練,共有 80 種物件。這些圖片可以由如下網址下載

https://github.com/wongkinyiu/yolov7 網址往下拉,可以看到 train, val, test, labels 的下載網址,分別如下。
train    :  http://images.cocodataset.org/zips/train2017.zip
val       :  http://images.cocodataset.org/zips/val2017.zip
test     :  http://images.cocodataset.org/zips/test2017.zip
labels :  https://github.com/WongKinYiu/yolov7/releases/download/v0.1/coco2017labels-segments.zip

請將上面的四個檔案下載到專案的根目錄下。

先把 coco2017labels-segments.zip 解壓縮至此,產生 coco 目錄。

將 train2017.zip 解壓縮至此,然後將 train2017 目錄改成 images,再將 images 移到 coco/train 目錄下。
將 val2017.zip 解壓縮至此,然後將 val2017 目錄改成 images,再將 images 移到 coco/val 目錄下。
將 test2017.zip 解壓縮至此,然後將 test2017 目錄改成 images,再將 images 移到 coco/test 目錄下。

將 coco/labels/train2017 移到 coco/train 之下,並改成 labels。
將 coco/labels/val2017 移到 coco/valid 之下,並改成 labels。

設定 data.yaml

在 coco 下新增 data.yaml,內容如下。請注意一定要寫絕對路徑。請注意,本專案根目錄為 e:\python_ai\yolov8_coco。

train: e:/python_ai/yolov8_coco/coco/train/images
val: e:/python_ai/yolov8_coco/coco/valid/images
test: e:/python_ai/yolov8_coco/coco/test/images
# number of classes
nc: 80

# class names
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
         'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
         'hair drier', 'toothbrush' ]

下載預訓練模型

到 https://github.com/ultralytics/ultralytics 網站,往下拉下載 YOLOv8n ,儲存在專案根目錄下。

yolo.exe 訓練權重

在 Terminal 執行如下指令,epochs 設定為 200 次。

yolo task=detect mode=train model=./yolov8n.pt data=./coco/data.yaml epochs=200 imgsz=640

使用 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="./coco/data.yaml", epochs=200, imgsz=640)
    t2=time.time()
    print(f'訓練花費時間 : {t2-t1}秒')
    path=model.export()
    print(f'模型匯出路徑 : {path}')

使用 RTX 3080Ti 訓練時間約 48 小時

無法訓練的人,請由本站下載 yolov8_coco.zip ,解開後置於 ./runs/detect/train/weights 之下。

中斷接續訓練

上面訓練要花費二、三天的時間,中途若中斷,可以由中斷點再接續訓練,指令如下

from ultralytics import YOLO
if __name__=='__main__':
    model = YOLO('path/to/last.pt')  # load a partially trained model
    # Resume training
    results = model.train(resume=True)

偵測圖片

請使用如下代碼偵測圖片

import os
import platform
import pylab as plt
import cv2
import numpy as np
from PIL import Image, ImageFont, ImageDraw
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.asarray(pil)

model=YOLO('./runs/detect/train/weights/best.pt')
path="./coco/test/images"

plt.figure(figsize=(12,9))
for i,file in enumerate(os.listdir(path)[:6]):
    full=os.path.join(path, file)
    img=cv2.imdecode(np.fromfile(full, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
    img=img[:,:,::-1].copy()

    results=model.predict(img, save=False)
    boxes=results[0].boxes.xyxy
    names = [results[0].names[int(i.cpu().numpy())] for i in results[0].boxes.cls]
    for box, name in zip(boxes, names):
        print(name)
        box=box.cpu().numpy()
        x1 = int(box[0])
        y1 = int(box[1])
        x2 = int(box[2])
        y2 = int(box[3])
        print(img.shape, x1, y1, x2, y2)
        cv2.rectangle(img,(x1, y1), (x2, y2), (0,255,0) , 2)
        img=text(img, name, (x1, y1), (0,255,0),25)
    plt.subplot(2,3,i+1)
    plt.axis("off")
    plt.imshow(img)
plt.show()

發佈留言

發佈留言必須填寫的電子郵件地址不會公開。 必填欄位標示為 *