本工程包含AI模型训练，预测的自动机器学习库
先看效果
lol英雄目标检测
Ai框架请下载
基础环境
pip install -r requirements.txt
准备数据集
数据标注工具
- 分类
# 数据集结构如下
# main_directory/
# ...class_a/
# ......a_image_1.jpg
# ......a_image_2.jpg
# ...class_b/
# ......b_image_1.jpg
# ......b_image_2.jpg
- 目标检测
# 数据集结构如下
# main_directory/
# ...images/
# ......a_image_1.jpg
# ......a_image_2.jpg
# ...labels/
# ......b_image_1.txt
# ......b_image_2.txt
# ...classes.txt
# ...notes.json
训练任务
参考上面准备数据集
拉取配置
$ python3 -m utils.automlUtils.base get {dataset_path:例如./main_directory} class # 拉取分类任务的配置信息
$ python3 -m utils.automlUtils.base get {dataset_path:例如./main_directory} detectioon # 拉取目标检测任务的配置信息
config.yaml
请参考配置文件修改或者添加配置信息
训练任务
import cv2
import numpy as np
def 图像分类预测():
from utils.automlUtils import Config, ImageClassifierConfig, ImageClassifierPredictor
conf = ImageClassifierConfig(**(Config.json()))
print(conf)
import cv2
modeler = ImageClassifierPredictor(config=conf)
import time
for url in ["http://imgservice.suning.cn/uimg1/b2c/image/i3sVxZZHr9SbY4TpH1SmvQ.jpg"]:
start_time = time.time()
img = modeler.url2numpy(url=url)
img = cv2.resize(img, (256, 256))
predict_result = modeler.predict(x=np.array([img]))
print("inference time: ", time.time() - start_time)
print("predict_result:", predict_result)
post_result = modeler.post(predict_result)[0]
print("post_result", post_result)
print("-" * 20)
# print(label, post_result, label == post_result)
# inference
# time: 0.7593982219696045
# predict_result: [[0.38690686 0.23805034 0.26550755 0.10953525]]
# post_result
# {'class': 'BabyPants', 'confidence': 0.38690686}
for path in ["./img.png"]:
start_time = time.time()
img = modeler.file2numpy(path=path)
img = cv2.resize(img, (256, 256))
predict_result = modeler.predict(x=np.array([img]))
print("inference time: ", time.time() - start_time)
print("predict_result:", predict_result)
post_result = modeler.post(predict_result)[0]
print("post_result", post_result)
print("-" * 20)
# print(label, post_result, label == post_result)
# inference
# time: 0.0785212516784668
# predict_result: [[0.02243473 0.00355711 0.9482294 0.02577879]]
# post_result
# {'class': 'womencasualshoes', 'confidence': 0.9482294}
def 图片分类训练():
"""
数据集下载地址：https://autogluon.s3.amazonaws.com/datasets/shopee-iet.zip
"""
from utils.automlUtils import Config, ImageClassifierConfig, ImageClassifierPredictor
conf = ImageClassifierConfig(**(Config.json()))
print(conf)
ImageClassifierPredictor(config=conf).train_and_export()
def 图片检测训练():
from utils.automlUtils import Config, ImageDetectorPredictor, ImageDetectorConfig
conf = ImageDetectorConfig(**(Config.json()))
modeler = ImageDetectorPredictor(config=conf)
modeler.train_and_export()
def 图片检测预测():
from utils.automlUtils import Config, ImageDetectorPredictor, ImageDetectorConfig
conf = ImageDetectorConfig(**(Config.json()))
modeler = ImageDetectorPredictor(config=conf)
import time
for url in ["http://n.sinaimg.cn/spider2020824/264/w1080h784/20200824/8508-iyaiihm7246958.jpg"]:
img = modeler.url2Image(url=url)
start_time = time.time()
predict_result = modeler.predict(x=[img], size=640)
print("inference time: ", time.time() - start_time)
print("predict_result:", predict_result)
post_result = modeler.post(predict_result)
post_result[0].update({"name": "行人"})
img = modeler.draw(img, pred_val=post_result, size=40, show=True)
print("post_result", post_result)
print("-" * 20)
def show(path="/Users/rockontrol/Desktop/lol1.mp4"):
from utils.automlUtils import Config, ImageDetectorPredictor, ImageDetectorConfig
conf = ImageDetectorConfig(**(Config.json()))
modeler = ImageDetectorPredictor(config=conf)
video = cv2.VideoCapture(path)
import time
while video.isOpened():
ret, frame = video.read()
if not ret:
break
img = modeler.numpy2Image(img=frame)
start_time = time.time()
predict_result = modeler.predict(x=[img], size=640)
print("inference time: ", time.time() - start_time)
print("predict_result:", predict_result)
post_result = modeler.post(predict_result)
img = modeler.draw(img, pred_val=post_result, size=40, show=False)
frame = np.array(img)
cv2.imshow('frame', frame)
# 若没有按下q键，则每1毫秒显示一帧
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# 所有操作结束后不要忘记释放
video.release()
cv2.destroyWindow("frame")
if __name__ == '__main__':
import fire
fire.Fire()
- 训练
$ python3 main.py 图片分类训练
...
$ python3 main.py 图像分类预测
...
$ python3 main.py 图片检测训练
2 -1 1 18816 yolov5.models.common.C3 [64, 64, 1]
3 -1 1 73984 yolov5.models.common.Conv [64, 128, 3, 2]
4 -1 2 115712 yolov5.models.common.C3 [128, 128, 2]
5 -1 1 295424 yolov5.models.common.Conv [128, 256, 3, 2]
6 -1 3 625152 yolov5.models.common.C3 [256, 256, 3]
7 -1 1 1180672 yolov5.models.common.Conv [256, 512, 3, 2]
8 -1 1 1182720 yolov5.models.common.C3 [512, 512, 1]
9 -1 1 656896 yolov5.models.common.SPPF [512, 512, 5]
10 -1 1 131584 yolov5.models.common.Conv [512, 256, 1, 1]
11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
12 [-1, 6] 1 0 yolov5.models.common.Concat [1]
13 -1 1 361984 yolov5.models.common.C3 [512, 256, 1, False]
14 -1 1 33024 yolov5.models.common.Conv [256, 128, 1, 1]
15 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
16 [-1, 4] 1 0 yolov5.models.common.Concat [1]
17 -1 1 90880 yolov5.models.common.C3 [256, 128, 1, False]
18 -1 1 147712 yolov5.models.common.Conv [128, 128, 3, 2]
19 [-1, 14] 1 0 yolov5.models.common.Concat [1]
20 -1 1 296448 yolov5.models.common.C3 [256, 256, 1, False]
21 -1 1 590336 yolov5.models.common.Conv [256, 256, 3, 2]
22 [-1, 10] 1 0 yolov5.models.common.Concat [1]
23 -1 1 1182720 yolov5.models.common.C3 [512, 512, 1, False]
24 [17, 20, 23] 1 18879 yolov5.models.yolo.Detect [2, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
Model Summary: 270 layers, 7025023 parameters, 7025023 gradients, 15.9 GFLOPs
...
Starting training for 500 epochs...
Epoch gpu_mem box obj cls labels img_size
0/499 2.12G 0.07039 0.01503 0.01063 132 640: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:07<00:00, 1.27s/it]
Class Images Labels P R mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00, 1.65it/s]
all 91 366 0.0024 0.0533 0.00101 0.00015
Epoch gpu_mem box obj cls labels img_size
1/499 2.12G 0.07097 0.01468 0.01075 100 640: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:01<00:00, 3.75it/s]
Class Images Labels P R mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00, 2.38it/s]
all 91 366 0.00211 0.0266 0.000798 0.000117
Epoch gpu_mem box obj cls labels img_size
2/499 2.12G 0.07008 0.01568 0.01056 143 640: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:01<00:00, 3.57it/s]
Class Images Labels P R mAP@.5 mAP@.5:.95:
Epoch gpu_mem box obj cls labels img_size
197/499 2.12G 0.03268 0.01033 0.003047 111 640: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:01<00:00, 4.10it/s]
Class Images Labels P R mAP@.5 mAP@.5:.95: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 4.10it/s]
all 91 366 0.804 0.867 0.829 0.253
Epoch gpu_mem box obj cls labels img_size
198/499 2.12G 0.02963 0.009621 0.003108 141 640: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:01<00:00, 3.65it/s]
Class Images Labels P R mAP@.5 mAP@.5:.95:
$ python3 main.py 图片检测预测
inference time: 0.20399999618530273
post_result: [{'bbox': [385.8732604980469, 107.14449310302734, 447.9578857421875, 170.21142578125], 'confidence': 0.38369208574295044, 'name': '英雄', 'cls': 0}, {'bbox': [537.8407592773438, 118.21878814697266, 611.421142578125, 1813377075195312], 'confidence': 0.38793814182281494, 'name': '小兵', 'cls': 1}, {'bbox': [366.82672119140625, 226.30960083007812, 450.38079833984375, 297.9781188964844], 'confidence': 0.46867531538009644, 'name': '小兵', 'cls': 1}, {'': [383.04864501953125, 105.42671966552734, 450.8124084472656, 167.3153533935547], 'confidence': 0.5011225342750549, 'name': '小兵', 'cls': 1}, {'bbox': [486.0379638671875, 162.92074584960938, 650.8978881835938, 348.2691955566406], onfidence': 0.6519156098365784, 'name': '英雄', 'cls': 0}]
/Users/rockontrol/Desktop/python_code/venv/lib/python3.8/site-packages/torch/autocast_mode.py:141: UserWarning: User provided device_type of 'cuda', but CUDA is not available. Disabling
warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling')
inference time: 0.19359803199768066
post_result: [{'bbox': [386.54046630859375, 105.60713195800781, 448.0963439941406, 171.9403076171875], 'confidence': 0.32672685384750366, 'name': '英雄', 'cls': 0}, {'bbox': [537.49169921875, 118.83441925048828, 611.8263549804688, 1.37303161621094], 'confidence': 0.4074070155620575, 'name': '小兵', 'cls': 1}, {'bbox': [366.115966796875, 226.2453155517578, 449.3259582519531, 295.8804626464844], 'confidence': 0.41025012731552124, 'name': '小兵', 'cls': 1}, {'bbo[383.3480529785156, 105.75922393798828, 450.8968200683594, 167.15966796875], 'confidence': 0.5280213356018066, 'name': '小兵', 'cls': 1}, {'bbox': [486.89239501953125, 163.3772430419922, 648.9190673828125, 344.3421325683594], 'confince': 0.6505584716796875, 'name': '英雄', 'cls': 0}]
...
建议训练时，使用GPU训练
训练时用到的数据集链接: https://pan.baidu.com/s/1VO1CxCiWBRnvNXW1HgNgjg 密码: sea7
1070
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
