def detect_cv2_camera(weightfile):
import cv2
model = resnet50()
model.load_state_dict(torch.load(weightfile))
print('Loading weights from %s... Done!' % (weightfile))
if use_cuda:
model.cuda()
cap = cv2.VideoCapture(0)
# cap = cv2.VideoCapture("./test.mp4")
cap.set(3, 1280)
cap.set(4, 720)
print("Starting the YOLO loop...")
while True:
ret, img = cap.read() # 获取当前帧
start = time.time()
result_img = plot_boxes_cv2(model, img)
finish = time.time()
print('Predicted in %f seconds.' % (finish - start))
cv2.imshow('Yolov1 demo', result_img)
cv2.waitKey(1)
test_dataset = yoloDataset(root=file_root,
list_file='voc2007test.txt',
train=False,
transform=[transforms.ToTensor()])
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=4)
print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % batch_size)
# ---------------------网络选择---------------------
use_resnet = True
if use_resnet:
net = resnet50()
else:
net = vgg16_bn()
if resume:
print("loading weight from checkpoints/best.pth")
net.load_state_dict(torch.load('checkpoints/best.pth'))
else:
print('loading pre-trined model ......')
if use_resnet:
resnet = models.resnet50(pretrained=True)
new_state_dict = resnet.state_dict()
dd = net.state_dict()
for k in new_state_dict.keys():
print(k)
if k in dd.keys() and not k.startswith('fc'):
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=4)
# val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
print('the train dataset has %d images' % (len(train_dataset)))
# print('the test dataset has %d images' % (len(val_dataset)))
print('the batch_size is %d' % batch_size)
# --------------------- 网络选择 ----------------------
use_resnet = True
if use_resnet:
net = resnet50() # 自己定义的resnet 网络
else:
net = vgg16_bn()
"""
models 模型表示官方代码
自己定义的是resnet50
"""
if resume:
print("loading weight from checkpoints/best.pth")
net.load_state_dict(torch.load('checkpoints/best.pth'))
else:
print('loading pre-trained model ......')
if use_resnet: # 对应自己网络结构的加载官方预训练模型参数
offiNet = officalModel.resnet50(pretrained=True) # 官方模型
offiNew_state_dict = offiNet.state_dict()
dd = net.state_dict()
boxes, cls_indexs, probs = decoder(pred)
for i, box in enumerate(boxes):
x1 = int(box[0]*w)
x2 = int(box[2]*w)
y1 = int(box[1]*h)
y2 = int(box[3]*h)
cls_index = cls_indexs[i]
cls_index = int(cls_index) # convert LongTensor to int
prob = probs[i]
prob = float(prob)
result.append([(x1, y1), (x2, y2), VOC_CLASSES[cls_index], image_name, prob])
return result
if __name__ == '__main__':
model = resnet50()
print('load model...')
num_epochs=33
learning_rate=0.0001
model.load_state_dict(torch.load('val_best.pth'))
#model.load_state_dict(torch.load('best.pth'))
model.eval()
model.cuda()
image_name = 'imgs/person.jpg'
image = cv2.imread(image_name)
print('predicting...')
result = predict_gpu(model, image_name) # 获取预测值
# 绘图
for left_up, right_bottom, class_name, _, prob in result:
print("class_name:%s, prob:%f" % (class_name, prob))
