YoloV5 快速體驗

feiyangyy94發表於2024-04-12
YOLO

準備工作

1. clone yolov5 的倉庫
git clone https://github.com/ultralytics/yolov5.git
2. 準備資料集

我看網上好多文章都是在講如何標註，自己製作資料集；此處提供一個從外網下載的口罩資料集，已經做好標註

https://www.alipan.com/s/bpR8veU9k48

資料集來源於牆外分享，不得商用

下載解壓後，可以看到obj目錄內有一系列jpg和txt, 每張圖片對應於一個txt， txt的內容描述的是圖片的標註資訊。此資料集內的標註檔案格式都是yolo的格式。

我們需要把這些檔案分割成多個子集：

  • 訓練集: 用於訓練的資料集
  • 驗證集：---驗證-----
  • 測試集：---測試-----

每個資料集有image和labels兩部分，需要分別儲存在不同的目錄，下面的指令碼用於完成這些工作（請酌情修改）:

import torch
# from IPython.display import Image  # for displaying images
import os 
import random
import shutil
from sklearn.model_selection import train_test_split
import xml.etree.ElementTree as ET
from xml.dom import minidom
from tqdm import tqdm
from PIL import Image, ImageDraw
import numpy as np
import matplotlib.pyplot as plt
import glob
import shutil
random.seed(108)

def make_data_set(path:str):
  # get full path
  images = glob.glob(os.path.join(path, "*.jpg"))
  annotations = glob.glob(os.path.join(path, "*.txt"))
  images.sort()
  annotations.sort()
  # 測試集0.2, 訓練集0.8
  train_images, val_images, train_annotations, val_annotations = train_test_split(images, annotations, test_size = 0.2, random_state = 1)
  # 測試0.1 驗證 0.1
  val_images, test_images, val_annotations, test_annotations = train_test_split(val_images, val_annotations, test_size = 0.5, random_state = 1)
  images_ = ["images/train", "images/val", "images/test"]
  annotations_ = ["labels/train", "labels/val", "labels/test"]
  for i in range(len(images_)):
    if not os.path.exists(images_[i]):
      os.makedirs(images_[i], exist_ok=True)
    if not os.path.exists(annotations_[i]):
      os.makedirs(annotations_[i],  exist_ok=True)
  for f in train_images:
    shutil.copy(f, images_[0])
  for f in val_images:
    shutil.copy(f, images_[1])
  for f in test_images:
    shutil.copy(f, images_[2])

  for f in train_annotations:
    shutil.copy(f, annotations_[0])
  for f in val_annotations:
    shutil.copy(f, annotations_[1])
  for f in test_annotations:
    shutil.copy(f, annotations_[2])

這個指令碼從obj目錄提取圖片和標註檔案，按照8:1:1的比例分割訓練、驗證和測試資料集，並且將其分別儲存不同的目錄中

準備訓練所需yaml檔案:
在yolov5/data下建立mask_data.yaml, 內容如下：

# 路徑根據實際情況填寫
train: ../images/train/ 
val:  ../images/val/
test: ../images/test/
# 2個分類
nc: 2

names: ["masked","no_mask"]

開始訓練

進入yolov5目錄，輸入以下命令

python train.py --img 480 --cfg yolov5s.yaml  --batch 32 --epochs 100 --data mask_data.yaml --weights yolov5s.pt  --name mask_data --device 0

上面要注意的是--device 選項，對於無顯示卡的裝置，需要改為--device cpu 進行訓練，使用cpu訓練非常慢

使用雲平臺訓練（推薦）

這裡用的是矩池雲，該平臺為每個新使用者提供5元體驗券，用作本文的實驗性訓練來說足夠了
註冊賬號，公眾號領算力豆，建立例項這些大家按照平臺說明操作就好。

我這裡選的算力產品是：Pytorch 2.1.1

例項建立完畢後，將資料集上傳至伺服器並解壓，按照前述步驟clone yolov5程式碼以及編輯好相關配置，執行上文中的訓練命令。對於本文的資料集，一般30分鐘內即可搞定

使用訓練好的權重資料

訓練完成後，終端會輸出結果儲存在哪個目錄，我這裡是:

runs/train/mask_data7/weights/best.pt

我們可以使用這個權重資料來進行一些簡單的測試

python detect.py --source FF.MP4 --weights best.pt --conf-thres 0.25

此處有兩點需注意：

  1. 如果是在windows 上跑detect, 可能會報出PosixPath問題, 需在detect.py檔案頭部加入:
import pathlib
pathlib.PosixPath = pathlib.WindowsPath
  1. 如果是沒有NV顯示卡的機器，則需要使用cpu 推理，在上述命令末尾加上 --device cpu 即可，使用cpu推理可能會比較慢

推理完成後，終端會輸出結果儲存目錄

相關文章