微信公众号:OpenCV学堂 关注获取更多计算机视觉与深度学习知识
纯Python的快速对象检测训练与测试平台,基于tensorflow+opencv构建,支持
在COCO数据集上预训练模型的调用,同时支持自定义对象训练与导出。支持
ImageAI的后台依赖tensorflow框架与keras,所以需要首先安装tensoflow,当前还不支持tensorflow2.0版本
安装imageai,只需要执行如下命令行即可
pip install imageai
1. 图像分类
from imageai.Prediction import ImagePrediction
import os
execution_path = os.getcwd()
prediction = ImagePrediction()
prediction.setModelTypeAsResNet()
prediction.setModelPath(os.path.join(execution_path, "resnet50_weights_tf_dim_ordering_tf_kernels.h5"))
prediction.loadModel()
predictions, probabilities = prediction.predictImage(os.path.join(execution_path, "1.jpg"), result_count=5 )
for eachPrediction, eachProbability in zip(predictions, probabilities):
print(eachPrediction , " : " , eachProbability)
图像
运行输出:
convertible : 52.459555864334106 sports_car : 37.61284649372101 pickup : 3.1751200556755066 car_wheel : 1.817505806684494 minivan : 1.7487050965428352
2. 对象检测
from imageai.Detection import ObjectDetection
import os
execution_path = os.getcwd()
detector = ObjectDetection()
detector.setModelTypeAsYOLOv3()
detector.setModelPath( os.path.join(execution_path , "yolo.h5"))
detector.loadModel()
detections = detector.detectObjectsFromImage(input_image=os.path.join(execution_path , "image2.jpg"), output_image_path=os.path.join(execution_path , "image2new.jpg"), minimum_percentage_probability=30)
for eachObject in detections:
print(eachObject["name"] , " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"] )
print("--------------------------------")
运行输出:
github地址
https://github.com/OlafenwaMoses/ImageAI/