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)