1.按《在Google Colaboratory測試imageAI》中介紹的方法,新建筆記本。
2.點擊菜單【代碼執行程序】【更改運行時類型】修改運行時類型爲GPU(否則會報錯:ImportError: libcublas.so.10.0: cannot open shared object file: No such file or directory)
3.下載並解壓數據集。
!wget https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/hololens.zip
!unzip hololens.zip
數據集也可以自己手動上傳,通過!unzip命令解壓。
4.下載pretrained-yolov3.h5。
!wget https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/pretrained-yolov3.h5
4.安裝Tensorflow和Tensorflow-GPU(經測試只有1.13.1版可用,其他版本都報錯),運行完點擊【RESTART RUNTIME】
!pip3 install tensorflow==1.13.1
!pip3 install tensorflow-gpu==1.13.1
5.安裝imageai
!pip3 install imageai
6.編寫訓練代碼,開始訓練。接下來就是等待,可能待上幾個小時,也可能等待幾十個小時。
from imageai.Detection.Custom import DetectionModelTrainer
trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="hololens")
trainer.setTrainConfig(object_names_array=["hololens"], batch_size=4, num_experiments=100, train_from_pretrained_model="pretrained-yolov3.h5")
trainer.trainModel()
7.訓練結束,會成hololens\models下生成一系到*.h5文件。下載h損失值(文件名中loss後面那個數據)最小的文件(例detection_model-ex-090--loss-0001.687.h5)、detection_config.json。
8.測試訓練模型
將detection_model-ex-090--loss-0001.687.h5)、detection_config.json複製到temp.py文件同一目錄下,進行測試。
from imageai.Detection.Custom import CustomObjectDetection
import os
import cv2
execution_path = os.getcwd()
detector = CustomObjectDetection()
detector.setModelTypeAsYOLOv3()
detector.setModelPath("detection_model-ex-090--loss-0001.687.h5")
detector.setJsonPath("detection_config.json")
detector.loadModel()
cap = cv2.VideoCapture(0)
while True:
ret, frame = cap.read()
x, y = frame.shape[0:2]
returned_image, detections = detector.detectObjectsFromImage(input_image=frame,input_type='array', output_type='array', minimum_percentage_probability=30)
cv2.imshow("video", returned_image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
下圖爲花朵訓練結果,僅爲某次訓練的截圖
其他操作記錄:刪除文件夾
import shutil
shutil.rmtree("validation", ignore_errors=True)