在Colaboratory中使用ImageAI训练自己的数据集

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)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章