- 资产集市
- 教学
- 实践
- AI说
- 案例库
- 生态合作
- 专区
3用户已评价
中国的各施工工地每年都要使用大量的钢筋，一车钢筋运到工地现场需要工作人员进行盘点，通常的做法是靠人工一根根数的方式，非常耗时费力。
为了提高钢筋盘点效率，业界提出了对钢筋图片进行拍照，然后使用AI算法检测图片中的钢筋条数，实践证明，该方案不仅准确率高，而且可以极大提高效率。
本案例基于目标检测的方法，使用250张已经人工标注好的图片进行AI模型的训练，训练25分钟，即可检测出图片中钢筋的横截面，从而统计出钢筋的条数。
本案例推荐使用AI框架：Pytorch-1.0.0；
本案例需使用 GPU 运行，请查看《ModelArts JupyterLab 硬件规格使用指南》了解切换硬件规格的方法；
如果您是第一次使用 JupyterLab，请查看《ModelArts JupyterLab使用指导》了解使用方法；
如果您在使用 JupyterLab 过程中碰到报错，请参考《ModelArts JupyterLab常见问题解决办法》尝试解决问题。
复制链接
微信扫一扫
import os
if not os.path.exists('./rebar_count'):
print('Downloading code and datasets...')
os.system("wget -N https://modelarts-labs-bj4-v2.obs.cn-north-4.myhuaweicloud.com/notebook/DL_rebar_count/rebar_count_code.zip")
os.system("wget -N https://cnnorth4-modelhub-datasets-obsfs-sfnua.obs.cn-north-4.myhuaweicloud.com/content/c2c1853f-d6a6-4c9d-ac0e-203d4c304c88/NkxX5K/dataset/rebar_count_datasets.zip")
os.system("unzip rebar_count_code.zip; rm rebar_count_code.zip")
os.system("unzip -q rebar_count_datasets.zip; rm rebar_count_datasets.zip")
os.system("mv rebar_count_code rebar_count; mv rebar_count_datasets rebar_count/datasets")
if os.path.exists('./rebar_count'):
print('Download code and datasets success')
else:
print('Download code and datasets failed, please check the download url is valid or not.')
else:
print('./rebar_count already exists')./rebar_count already exists
import os
import sys
sys.path.insert(0, './rebar_count/src')
import cv2
import time
import random
import torch
import numpy as np
from PIL import Image, ImageDraw
import xml.etree.ElementTree as ET
from datetime import datetime
from collections import OrderedDict
import torch.optim as optim
import torch.utils.data as data
import torch.backends.cudnn as cudnn
from data import VOCroot, VOC_Config, AnnotationTransform, VOCDetection, detection_collate, BaseTransform, preproc
from models.RFB_Net_vgg import build_net
from layers.modules import MultiBoxLoss
from layers.functions import Detect, PriorBox
from utils.visualize import *
from utils.nms_wrapper import nms
from utils.timer import Timer
import matplotlib.pyplot as plt
%matplotlib inline
ROOT_DIR = os.getcwd()
seed = 0
cudnn.benchmark = False
cudnn.deterministic = True
torch.manual_seed(seed) # 为CPU设置随机种子
torch.cuda.manual_seed_all(seed) # 为所有GPU设置随机种子
random.seed(seed)
np.random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed) # 设置hash随机种子def read_xml(xml_path):
'''读取xml标签'''
tree = ET.parse(xml_path)
root = tree.getroot()
boxes = []
labels = []
for element in root.findall('object'):
label = element.find('name').text
if label == 'steel':
bndbox = element.find('bndbox')
xmin = bndbox.find('xmin').text
ymin = bndbox.find('ymin').text
xmax = bndbox.find('xmax').text
ymax = bndbox.find('ymax').text
boxes.append([xmin, ymin, xmax, ymax])
labels.append(label)
return np.array(boxes, dtype=np.float64), labelstrain_img_dir = './rebar_count/datasets/VOC2007/JPEGImages'
train_xml_dir = './rebar_count/datasets/VOC2007/Annotations'
files = os.listdir(train_img_dir)
files.sort()
for index, file_name in enumerate(files[:2]):
img_path = os.path.join(train_img_dir, file_name)
xml_path = os.path.join(train_xml_dir, file_name.split('.jpg')[0] + '.xml')
boxes, labels = read_xml(xml_path)
img = Image.open(img_path)
resize_scale = 2048.0 / max(img.size)
img = img.resize((int(img.size[0] * resize_scale), int(img.size[1] * resize_scale)))
boxes *= resize_scale
plt.figure(figsize=(img.size[0]/100.0, img.size[1]/100.0))
plt.subplot(2,1,1)
plt.imshow(img)
img = img.convert('RGB')
img = np.array(img)
img = img.copy()
for box in boxes:
xmin, ymin, xmax, ymax = box.astype(np.int)
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), thickness=3)
plt.subplot(2,1,2)
plt.imshow(img)
plt.show()
# 定义训练超参
num_classes = 2 # 数据集中只有 steel 一个标签，加上背景，所以总共有2个类
max_epoch = 25 # 默认值为1，调整为大于20的值，训练效果更佳
batch_size = 4
ngpu = 1
initial_lr = 0.01
img_dim = 416 # 模型输入图片大小
train_sets = [('2007', 'trainval')] # 指定训练集
cfg = VOC_Config
rgb_means = (104, 117, 123) # ImageNet数据集的RGB均值
save_folder = './rebar_count/model_snapshots' # 指定训练模型保存路径
if not os.path.exists(save_folder):
os.mkdir(save_folder)
log_path = os.path.join('./rebar_count/logs', datetime.now().isoformat()) # 指定日志保存路径
if not os.path.exists(log_path):
os.makedirs(log_path)net = build_net('train', img_dim, num_classes=num_classes)
if ngpu > 1:
net = torch.nn.DataParallel(net)
net.cuda() # 本案例代码只能在GPU上训练
cudnn.benchmark = True
optimizer = optim.SGD(net.parameters(), lr=initial_lr,
momentum=0.9, weight_decay=0) # 定义优化器
criterion = MultiBoxLoss(num_classes,
overlap_thresh=0.4,
prior_for_matching=True,
bkg_label=0,
neg_mining=True,
neg_pos=3,
neg_overlap=0.3,
encode_target=False) # 定义损失函数
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
priors = priors.cuda()def adjust_learning_rate(optimizer, gamma, epoch, step_index, iteration, epoch_size):
"""
自适应学习率
"""
if epoch < 11:
lr = 1e-8 + (initial_lr-1e-8) * iteration / (epoch_size * 10)
else:
lr = initial_lr * (gamma ** (step_index))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
return lrdef train():
"""
模型训练函数，每10次迭代打印一次日志，20个epoch之后，每个epoch保存一次模型
"""
net.train()
loc_loss = 0
conf_loss = 0
epoch = 0
print('Loading dataset...')
dataset = VOCDetection(VOCroot, train_sets, preproc(img_dim, rgb_means, p=0.0), AnnotationTransform())
epoch_size = len(dataset) // batch_size
max_iter = max_epoch * epoch_size
stepvalues = (25 * epoch_size, 35 * epoch_size)
step_index = 0
start_iter = 0
lr = initial_lr
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
if epoch > 20:
torch.save(net.state_dict(), os.path.join(save_folder, 'epoch_' +
repr(epoch).zfill(3) + '_loss_'+ '%.4f' % loss.item() + '.pth'))
batch_iterator = iter(data.DataLoader(dataset, batch_size,
shuffle=True, num_workers=1, collate_fn=detection_collate))
loc_loss = 0
conf_loss = 0
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, 0.2, epoch, step_index, iteration, epoch_size)
images, targets = next(batch_iterator)
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
load_t1 = time.time()
if iteration % 10 == 0:
print('Epoch:' + repr(epoch) + ' || epochiter: ' + repr(iteration % epoch_size) + '/' + repr(epoch_size)
+ '|| Totel iter ' +
repr(iteration) + ' || L: %.4f C: %.4f||' % (
loss_l.item(),loss_c.item()) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) + 'LR: %.8f' % (lr))
torch.save(net.state_dict(), os.path.join(save_folder, 'epoch_' +
repr(epoch).zfill(3) + '_loss_'+ '%.4f' % loss.item() + '.pth'))t1 = time.time()
print('开始训练，本次训练总共需%d个epoch，每个epoch训练耗时约60秒' % max_epoch)
train()
print('training cost %.2f s' % (time.time() - t1))开始训练，本次训练总共需25个epoch，每个epoch训练耗时约60秒 Loading dataset... Epoch:1 || epochiter: 0/50|| Totel iter 0 || L: 3.7043 C: 3.7730||Batch time: 2.6931 sec. ||LR: 0.00000001 Epoch:1 || epochiter: 10/50|| Totel iter 10 || L: 3.1277 C: 3.1485||Batch time: 1.3692 sec. ||LR: 0.00020001 Epoch:1 || epochiter: 20/50|| Totel iter 20 || L: 3.3249 C: 2.4864||Batch time: 0.7837 sec. ||LR: 0.00040001 Epoch:1 || epochiter: 30/50|| Totel iter 30 || L: 2.8867 C: 2.4690||Batch time: 1.5246 sec. ||LR: 0.00060001 Epoch:1 || epochiter: 40/50|| Totel iter 40 || L: 2.6481 C: 2.1631||Batch time: 1.4777 sec. ||LR: 0.00080001 Epoch:2 || epochiter: 0/50|| Totel iter 50 || L: 3.0177 C: 2.1672||Batch time: 1.5618 sec. ||LR: 0.00100001 Epoch:2 || epochiter: 10/50|| Totel iter 60 || L: 1.9024 C: 1.8743||Batch time: 1.2920 sec. ||LR: 0.00120001 Epoch:2 || epochiter: 20/50|| Totel iter 70 || L: 1.5299 C: 1.7229||Batch time: 1.3726 sec. ||LR: 0.00140001 Epoch:2 || epochiter: 30/50|| Totel iter 80 || L: 1.7592 C: 1.8066||Batch time: 0.9840 sec. ||LR: 0.00160001 Epoch:2 || epochiter: 40/50|| Totel iter 90 || L: 1.4430 C: 1.7445||Batch time: 1.5012 sec. ||LR: 0.00180001 Epoch:3 || epochiter: 0/50|| Totel iter 100 || L: 1.3402 C: 1.5614||Batch time: 1.3830 sec. ||LR: 0.00200001 Epoch:3 || epochiter: 10/50|| Totel iter 110 || L: 1.2771 C: 1.7149||Batch time: 1.4420 sec. ||LR: 0.00220001 Epoch:3 || epochiter: 20/50|| Totel iter 120 || L: 2.1052 C: 2.3860||Batch time: 1.0122 sec. ||LR: 0.00240001 Epoch:3 || epochiter: 30/50|| Totel iter 130 || L: 1.3969 C: 2.0087||Batch time: 1.2500 sec. ||LR: 0.00260001 Epoch:3 || epochiter: 40/50|| Totel iter 140 || L: 1.1426 C: 1.3518||Batch time: 1.3625 sec. ||LR: 0.00280001 Epoch:4 || epochiter: 0/50|| Totel iter 150 || L: 1.3851 C: 1.3837||Batch time: 1.3933 sec. ||LR: 0.00300001 Epoch:4 || epochiter: 10/50|| Totel iter 160 || L: 0.8790 C: 1.0304||Batch time: 1.0430 sec. ||LR: 0.00320001 Epoch:4 || epochiter: 20/50|| Totel iter 170 || L: 1.1230 C: 1.2439||Batch time: 1.0029 sec. ||LR: 0.00340001 Epoch:4 || epochiter: 30/50|| Totel iter 180 || L: 1.0097 C: 1.1061||Batch time: 1.5267 sec. ||LR: 0.00360001 Epoch:4 || epochiter: 40/50|| Totel iter 190 || L: 0.8008 C: 1.0768||Batch time: 1.1727 sec. ||LR: 0.00380001 Epoch:5 || epochiter: 0/50|| Totel iter 200 || L: 1.0015 C: 1.1481||Batch time: 1.3881 sec. ||LR: 0.00400001 Epoch:5 || epochiter: 10/50|| Totel iter 210 || L: 0.9171 C: 1.1305||Batch time: 1.2255 sec. ||LR: 0.00420001 Epoch:5 || epochiter: 20/50|| Totel iter 220 || L: 0.9460 C: 1.0200||Batch time: 1.0095 sec. ||LR: 0.00440001 Epoch:5 || epochiter: 30/50|| Totel iter 230 || L: 0.8780 C: 1.1776||Batch time: 1.3224 sec. ||LR: 0.00460001 Epoch:5 || epochiter: 40/50|| Totel iter 240 || L: 0.8082 C: 0.8878||Batch time: 1.0734 sec. ||LR: 0.00480001 Epoch:6 || epochiter: 0/50|| Totel iter 250 || L: 0.7907 C: 0.9508||Batch time: 1.2835 sec. ||LR: 0.00500001 Epoch:6 || epochiter: 10/50|| Totel iter 260 || L: 0.6690 C: 0.8685||Batch time: 1.4887 sec. ||LR: 0.00520000 Epoch:6 || epochiter: 20/50|| Totel iter 270 || L: 1.1006 C: 0.9525||Batch time: 1.3324 sec. ||LR: 0.00540000 Epoch:6 || epochiter: 30/50|| Totel iter 280 || L: 0.9483 C: 1.0393||Batch time: 1.3198 sec. ||LR: 0.00560000 Epoch:6 || epochiter: 40/50|| Totel iter 290 || L: 0.8986 C: 1.0833||Batch time: 1.3434 sec. ||LR: 0.00580000 Epoch:7 || epochiter: 0/50|| Totel iter 300 || L: 0.8187 C: 0.9676||Batch time: 1.4531 sec. ||LR: 0.00600000 Epoch:7 || epochiter: 10/50|| Totel iter 310 || L: 0.6827 C: 0.9837||Batch time: 0.9223 sec. ||LR: 0.00620000 Epoch:7 || epochiter: 20/50|| Totel iter 320 || L: 0.7325 C: 0.8995||Batch time: 0.9585 sec. ||LR: 0.00640000 Epoch:7 || epochiter: 30/50|| Totel iter 330 || L: 0.9895 C: 1.0482||Batch time: 1.2272 sec. ||LR: 0.00660000 Epoch:7 || epochiter: 40/50|| Totel iter 340 || L: 0.5824 C: 0.8616||Batch time: 1.1445 sec. ||LR: 0.00680000 Epoch:8 || epochiter: 0/50|| Totel iter 350 || L: 1.1853 C: 1.2745||Batch time: 1.5200 sec. ||LR: 0.00700000 Epoch:8 || epochiter: 10/50|| Totel iter 360 || L: 0.7265 C: 1.1777||Batch time: 0.7649 sec. ||LR: 0.00720000 Epoch:8 || epochiter: 20/50|| Totel iter 370 || L: 0.7457 C: 0.8613||Batch time: 1.5218 sec. ||LR: 0.00740000 Epoch:8 || epochiter: 30/50|| Totel iter 380 || L: 0.5295 C: 0.9103||Batch time: 1.2653 sec. ||LR: 0.00760000 Epoch:8 || epochiter: 40/50|| Totel iter 390 || L: 0.7083 C: 1.0060||Batch time: 1.1069 sec. ||LR: 0.00780000 Epoch:9 || epochiter: 0/50|| Totel iter 400 || L: 0.6398 C: 0.9866||Batch time: 1.5802 sec. ||LR: 0.00800000 Epoch:9 || epochiter: 10/50|| Totel iter 410 || L: 0.5987 C: 0.8167||Batch time: 1.0675 sec. ||LR: 0.00820000 Epoch:9 || epochiter: 20/50|| Totel iter 420 || L: 0.5751 C: 0.7944||Batch time: 0.7669 sec. ||LR: 0.00840000 Epoch:9 || epochiter: 30/50|| Totel iter 430 || L: 0.7229 C: 1.0396||Batch time: 1.3895 sec. ||LR: 0.00860000 Epoch:9 || epochiter: 40/50|| Totel iter 440 || L: 0.5569 C: 0.9122||Batch time: 0.8300 sec. ||LR: 0.00880000 Epoch:10 || epochiter: 0/50|| Totel iter 450 || L: 0.6908 C: 0.9928||Batch time: 1.4029 sec. ||LR: 0.00900000 Epoch:10 || epochiter: 10/50|| Totel iter 460 || L: 0.6851 C: 0.8068||Batch time: 1.2804 sec. ||LR: 0.00920000 Epoch:10 || epochiter: 20/50|| Totel iter 470 || L: 0.6783 C: 0.8511||Batch time: 1.7469 sec. ||LR: 0.00940000 Epoch:10 || epochiter: 30/50|| Totel iter 480 || L: 0.7962 C: 0.8040||Batch time: 1.6116 sec. ||LR: 0.00960000 Epoch:10 || epochiter: 40/50|| Totel iter 490 || L: 0.7782 C: 0.9469||Batch time: 1.1979 sec. ||LR: 0.00980000 Epoch:11 || epochiter: 0/50|| Totel iter 500 || L: 0.8902 C: 0.8956||Batch time: 1.8625 sec. ||LR: 0.01000000 Epoch:11 || epochiter: 10/50|| Totel iter 510 || L: 0.8532 C: 0.9259||Batch time: 1.2692 sec. ||LR: 0.01000000 Epoch:11 || epochiter: 20/50|| Totel iter 520 || L: 0.7917 C: 0.7990||Batch time: 1.7494 sec. ||LR: 0.01000000 Epoch:11 || epochiter: 30/50|| Totel iter 530 || L: 0.9688 C: 1.2376||Batch time: 1.1547 sec. ||LR: 0.01000000 Epoch:11 || epochiter: 40/50|| Totel iter 540 || L: 0.7030 C: 0.8440||Batch time: 1.1588 sec. ||LR: 0.01000000 Epoch:12 || epochiter: 0/50|| Totel iter 550 || L: 0.6580 C: 0.8380||Batch time: 1.2196 sec. ||LR: 0.01000000 Epoch:12 || epochiter: 10/50|| Totel iter 560 || L: 0.7978 C: 0.8454||Batch time: 1.1011 sec. ||LR: 0.01000000 Epoch:12 || epochiter: 20/50|| Totel iter 570 || L: 0.6071 C: 0.8394||Batch time: 0.7146 sec. ||LR: 0.01000000 Epoch:12 || epochiter: 30/50|| Totel iter 580 || L: 0.4787 C: 0.6888||Batch time: 1.2482 sec. ||LR: 0.01000000 Epoch:12 || epochiter: 40/50|| Totel iter 590 || L: 0.6505 C: 0.8412||Batch time: 1.1304 sec. ||LR: 0.01000000 Epoch:13 || epochiter: 0/50|| Totel iter 600 || L: 0.6316 C: 0.8319||Batch time: 1.4268 sec. ||LR: 0.01000000 Epoch:13 || epochiter: 10/50|| Totel iter 610 || L: 0.6693 C: 0.7822||Batch time: 1.2204 sec. ||LR: 0.01000000 Epoch:13 || epochiter: 20/50|| Totel iter 620 || L: 0.6773 C: 0.9631||Batch time: 1.2477 sec. ||LR: 0.01000000 Epoch:13 || epochiter: 30/50|| Totel iter 630 || L: 0.4851 C: 0.8346||Batch time: 1.2228 sec. ||LR: 0.01000000 Epoch:13 || epochiter: 40/50|| Totel iter 640 || L: 0.7247 C: 0.9392||Batch time: 1.2318 sec. ||LR: 0.01000000 Epoch:14 || epochiter: 0/50|| Totel iter 650 || L: 0.5716 C: 0.7683||Batch time: 1.8367 sec. ||LR: 0.01000000 Epoch:14 || epochiter: 10/50|| Totel iter 660 || L: 0.7804 C: 1.0285||Batch time: 1.0683 sec. ||LR: 0.01000000 Epoch:14 || epochiter: 20/50|| Totel iter 670 || L: 0.4620 C: 0.8179||Batch time: 1.3811 sec. ||LR: 0.01000000 Epoch:14 || epochiter: 30/50|| Totel iter 680 || L: 0.5459 C: 0.7611||Batch time: 1.4473 sec. ||LR: 0.01000000 Epoch:14 || epochiter: 40/50|| Totel iter 690 || L: 0.4946 C: 0.7604||Batch time: 1.2968 sec. ||LR: 0.01000000 Epoch:15 || epochiter: 0/50|| Totel iter 700 || L: 0.6467 C: 0.6637||Batch time: 1.4271 sec. ||LR: 0.01000000 Epoch:15 || epochiter: 10/50|| Totel iter 710 || L: 0.4383 C: 0.6140||Batch time: 1.1232 sec. ||LR: 0.01000000 Epoch:15 || epochiter: 20/50|| Totel iter 720 || L: 0.5551 C: 0.9027||Batch time: 1.2992 sec. ||LR: 0.01000000 Epoch:15 || epochiter: 30/50|| Totel iter 730 || L: 0.4488 C: 0.7574||Batch time: 0.9148 sec. ||LR: 0.01000000 Epoch:15 || epochiter: 40/50|| Totel iter 740 || L: 0.5179 C: 0.6202||Batch time: 1.5350 sec. ||LR: 0.01000000 Epoch:16 || epochiter: 0/50|| Totel iter 750 || L: 0.4956 C: 0.6740||Batch time: 1.6760 sec. ||LR: 0.01000000 Epoch:16 || epochiter: 10/50|| Totel iter 760 || L: 0.5780 C: 0.8834||Batch time: 1.3318 sec. ||LR: 0.01000000 Epoch:16 || epochiter: 20/50|| Totel iter 770 || L: 0.5829 C: 0.7340||Batch time: 1.0279 sec. ||LR: 0.01000000 Epoch:16 || epochiter: 30/50|| Totel iter 780 || L: 0.4798 C: 0.7019||Batch time: 1.4545 sec. ||LR: 0.01000000 Epoch:16 || epochiter: 40/50|| Totel iter 790 || L: 0.6511 C: 0.7712||Batch time: 1.7330 sec. ||LR: 0.01000000 Epoch:17 || epochiter: 0/50|| Totel iter 800 || L: 0.4281 C: 0.6578||Batch time: 1.6699 sec. ||LR: 0.01000000 Epoch:17 || epochiter: 10/50|| Totel iter 810 || L: 0.5440 C: 0.7102||Batch time: 1.4820 sec. ||LR: 0.01000000 Epoch:17 || epochiter: 20/50|| Totel iter 820 || L: 0.4770 C: 0.7014||Batch time: 1.4020 sec. ||LR: 0.01000000 Epoch:17 || epochiter: 30/50|| Totel iter 830 || L: 0.3601 C: 0.5890||Batch time: 1.0758 sec. ||LR: 0.01000000 Epoch:17 || epochiter: 40/50|| Totel iter 840 || L: 0.4817 C: 0.7329||Batch time: 1.3797 sec. ||LR: 0.01000000 Epoch:18 || epochiter: 0/50|| Totel iter 850 || L: 0.4860 C: 0.7499||Batch time: 1.3214 sec. ||LR: 0.01000000 Epoch:18 || epochiter: 10/50|| Totel iter 860 || L: 0.6856 C: 0.7154||Batch time: 1.4014 sec. ||LR: 0.01000000 Epoch:18 || epochiter: 20/50|| Totel iter 870 || L: 0.6231 C: 0.7692||Batch time: 0.9905 sec. ||LR: 0.01000000 Epoch:18 || epochiter: 30/50|| Totel iter 880 || L: 0.6680 C: 0.8625||Batch time: 1.1373 sec. ||LR: 0.01000000 Epoch:18 || epochiter: 40/50|| Totel iter 890 || L: 0.5535 C: 0.7393||Batch time: 1.1122 sec. ||LR: 0.01000000 Epoch:19 || epochiter: 0/50|| Totel iter 900 || L: 0.4691 C: 0.7235||Batch time: 1.3488 sec. ||LR: 0.01000000 Epoch:19 || epochiter: 10/50|| Totel iter 910 || L: 0.6145 C: 0.7811||Batch time: 1.1163 sec. ||LR: 0.01000000 Epoch:19 || epochiter: 20/50|| Totel iter 920 || L: 0.4698 C: 0.7225||Batch time: 1.6120 sec. ||LR: 0.01000000 Epoch:19 || epochiter: 30/50|| Totel iter 930 || L: 0.5623 C: 0.7341||Batch time: 1.3949 sec. ||LR: 0.01000000 Epoch:19 || epochiter: 40/50|| Totel iter 940 || L: 0.4859 C: 0.5786||Batch time: 0.8949 sec. ||LR: 0.01000000 Epoch:20 || epochiter: 0/50|| Totel iter 950 || L: 0.4193 C: 0.6898||Batch time: 1.4702 sec. ||LR: 0.01000000 Epoch:20 || epochiter: 10/50|| Totel iter 960 || L: 0.4434 C: 0.6261||Batch time: 1.0974 sec. ||LR: 0.01000000 Epoch:20 || epochiter: 20/50|| Totel iter 970 || L: 0.5948 C: 0.8787||Batch time: 1.1951 sec. ||LR: 0.01000000 Epoch:20 || epochiter: 30/50|| Totel iter 980 || L: 0.5842 C: 0.6120||Batch time: 0.9863 sec. ||LR: 0.01000000 Epoch:20 || epochiter: 40/50|| Totel iter 990 || L: 0.4010 C: 0.7356||Batch time: 1.5981 sec. ||LR: 0.01000000 Epoch:21 || epochiter: 0/50|| Totel iter 1000 || L: 0.4719 C: 0.6351||Batch time: 1.1228 sec. ||LR: 0.01000000 Epoch:21 || epochiter: 10/50|| Totel iter 1010 || L: 0.5856 C: 0.7444||Batch time: 1.3812 sec. ||LR: 0.01000000 Epoch:21 || epochiter: 20/50|| Totel iter 1020 || L: 0.5810 C: 0.8371||Batch time: 1.2560 sec. ||LR: 0.01000000 Epoch:21 || epochiter: 30/50|| Totel iter 1030 || L: 0.4583 C: 0.9570||Batch time: 1.1499 sec. ||LR: 0.01000000 Epoch:21 || epochiter: 40/50|| Totel iter 1040 || L: 0.5411 C: 0.5317||Batch time: 1.4007 sec. ||LR: 0.01000000 Epoch:22 || epochiter: 0/50|| Totel iter 1050 || L: 0.3508 C: 0.5599||Batch time: 1.1371 sec. ||LR: 0.01000000 Epoch:22 || epochiter: 10/50|| Totel iter 1060 || L: 0.4045 C: 0.6965||Batch time: 1.1030 sec. ||LR: 0.01000000 Epoch:22 || epochiter: 20/50|| Totel iter 1070 || L: 0.3949 C: 0.6019||Batch time: 1.4505 sec. ||LR: 0.01000000 Epoch:22 || epochiter: 30/50|| Totel iter 1080 || L: 0.3467 C: 0.5563||Batch time: 1.1956 sec. ||LR: 0.01000000 Epoch:22 || epochiter: 40/50|| Totel iter 1090 || L: 0.5757 C: 0.5643||Batch time: 0.8669 sec. ||LR: 0.01000000 Epoch:23 || epochiter: 0/50|| Totel iter 1100 || L: 0.3946 C: 0.6081||Batch time: 1.7117 sec. ||LR: 0.01000000 Epoch:23 || epochiter: 10/50|| Totel iter 1110 || L: 0.3655 C: 0.5579||Batch time: 0.9830 sec. ||LR: 0.01000000 Epoch:23 || epochiter: 20/50|| Totel iter 1120 || L: 0.3912 C: 0.6437||Batch time: 1.2725 sec. ||LR: 0.01000000 Epoch:23 || epochiter: 30/50|| Totel iter 1130 || L: 0.4237 C: 0.6337||Batch time: 1.3346 sec. ||LR: 0.01000000 Epoch:23 || epochiter: 40/50|| Totel iter 1140 || L: 0.3474 C: 0.5517||Batch time: 1.1646 sec. ||LR: 0.01000000 Epoch:24 || epochiter: 0/50|| Totel iter 1150 || L: 0.5573 C: 0.7426||Batch time: 1.5291 sec. ||LR: 0.01000000 Epoch:24 || epochiter: 10/50|| Totel iter 1160 || L: 0.6122 C: 0.6805||Batch time: 1.1861 sec. ||LR: 0.01000000 Epoch:24 || epochiter: 20/50|| Totel iter 1170 || L: 0.3846 C: 0.6484||Batch time: 1.2575 sec. ||LR: 0.01000000 Epoch:24 || epochiter: 30/50|| Totel iter 1180 || L: 0.4183 C: 0.6982||Batch time: 1.1318 sec. ||LR: 0.01000000 Epoch:24 || epochiter: 40/50|| Totel iter 1190 || L: 0.5259 C: 0.7322||Batch time: 1.0091 sec. ||LR: 0.01000000 Epoch:25 || epochiter: 0/50|| Totel iter 1200 || L: 0.4047 C: 0.5544||Batch time: 1.4809 sec. ||LR: 0.01000000 Epoch:25 || epochiter: 10/50|| Totel iter 1210 || L: 0.4519 C: 0.5351||Batch time: 1.2974 sec. ||LR: 0.01000000 Epoch:25 || epochiter: 20/50|| Totel iter 1220 || L: 0.4390 C: 0.6232||Batch time: 1.0032 sec. ||LR: 0.01000000 Epoch:25 || epochiter: 30/50|| Totel iter 1230 || L: 0.4840 C: 0.7323||Batch time: 1.0048 sec. ||LR: 0.01000000 Epoch:25 || epochiter: 40/50|| Totel iter 1240 || L: 0.6699 C: 0.8887||Batch time: 1.7034 sec. ||LR: 0.01000000 training cost 1572.48 s
cfg = VOC_Config
img_dim = 416
rgb_means = (104, 117, 123)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
if torch.cuda.is_available():
priors = priors.cuda()
class ObjectDetector:
"""
定义目标检测类
"""
def __init__(self, net, detection, transform, num_classes=num_classes, thresh=0.01, cuda=True):
self.net = net
self.detection = detection
self.transform = transform
self.num_classes = num_classes
self.thresh = thresh
self.cuda = torch.cuda.is_available()
def predict(self, img):
_t = {'im_detect': Timer(), 'misc': Timer()}
scale = torch.Tensor([img.shape[1], img.shape[0],
img.shape[1], img.shape[0]])
with torch.no_grad():
x = self.transform(img).unsqueeze(0)
if self.cuda:
x = x.cuda()
scale = scale.cuda()
_t['im_detect'].tic()
out = net(x) # forward pass
boxes, scores = self.detection.forward(out, priors)
detect_time = _t['im_detect'].toc()
boxes = boxes[0]
scores = scores[0]
# scale each detection back up to the image
boxes *= scale
boxes = boxes.cpu().numpy()
scores = scores.cpu().numpy()
_t['misc'].tic()
all_boxes = [[] for _ in range(num_classes)]
for j in range(1, num_classes):
inds = np.where(scores[:, j] > self.thresh)[0]
if len(inds) == 0:
all_boxes[j] = np.zeros([0, 5], dtype=np.float32)
continue
c_bboxes = boxes[inds]
c_scores = scores[inds, j]
c_dets = np.hstack((c_bboxes, c_scores[:, np.newaxis])).astype(
np.float32, copy=False)
keep = nms(c_dets, 0.2, force_cpu=False)
c_dets = c_dets[keep, :]
all_boxes[j] = c_dets
nms_time = _t['misc'].toc()
total_time = detect_time + nms_time
return all_boxes, total_timetrained_models = os.listdir(os.path.join(ROOT_DIR, './rebar_count/model_snapshots')) # 模型文件所在目录
lowest_loss = 9999
best_model_name = ''
for model_name in trained_models:
if not model_name.endswith('pth'):
continue
loss = float(model_name.split('_loss_')[1].split('.pth')[0])
if loss < lowest_loss:
lowest_loss = loss
best_model_name = model_name
best_model_path = os.path.join(ROOT_DIR, './rebar_count/model_snapshots', best_model_name)
print('loading model from', best_model_path)
net = build_net('test', img_dim, num_classes) # 加载模型
state_dict = torch.load(best_model_path)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:]
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
net.eval()
print('Finish load model!')
if torch.cuda.is_available():
net = net.cuda()
cudnn.benchmark = True
else:
net = net.cpu()
detector = Detect(num_classes, 0, cfg)
transform = BaseTransform(img_dim, rgb_means, (2, 0, 1))
object_detector = ObjectDetector(net, detector, transform)loading model from /home/ma-user/work/./rebar_count/model_snapshots/epoch_023_loss_1.0207.pth Finish load model!
test_img_dir = r'./rebar_count/datasets/test_dataset' # 待预测的图片目录
files = os.listdir(test_img_dir)
files.sort()
for i, file_name in enumerate(files[:2]):
image_src = cv2.imread(os.path.join(test_img_dir, file_name))
detect_bboxes, tim = object_detector.predict(image_src)
image_draw = image_src.copy()
rebar_count = 0
for class_id, class_collection in enumerate(detect_bboxes):
if len(class_collection) > 0:
for i in range(class_collection.shape[0]):
if class_collection[i, -1] > 0.6:
pt = class_collection[i]
cv2.circle(image_draw, (int((pt[0] + pt[2]) * 0.5), int((pt[1] + pt[3]) * 0.5)), int((pt[2] - pt[0]) * 0.5 * 0.6), (255, 0, 0), -1)
rebar_count += 1
cv2.putText(image_draw, 'rebar_count: %d' % rebar_count, (25, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 3)
plt.figure(i, figsize=(30, 20))
plt.imshow(image_draw)
plt.show()
至此，本案例结束。
250张训练+200张测试_VOC格式_可下载zip包
42个月以前
名称 |
|---|
 epoch_024_loss_1.0182.pth  该文件不支持代码预览 |
评论2
5
3人已评价
请登录后评论