首先,我们先加载一下模型:
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
注意的pretrained值true和false的区别,一般我们设置成True就可以,他会在我们的环境里面自动搜索相应的模型参数,如果没有就自动下载。
但是因为下载速度太慢,我们就直接使用迅雷下载:
下载完了之后我们就直接使用model就可以了,毕竟不用咱训练,简单点。。
注意里面的数据增强和归一化,可以直接对opencv对象使用,使其同时成为一个tensor。
model.eval() cap = cv2.VideoCapture(0) transform = transforms.Compose( [ transforms.ToTensor() ] ) while True: ret, frame = cap.read() image = frame frame = transform(frame) pred = model([frame]) # 检测出目标的类别和得分 pred_class = [COCO_INSTANCE_CATEGORY_NAMES[ii] for ii in list(pred[0]['labels'].numpy())] pred_score = list(pred[0]['scores'].detach().numpy()) # 检测出目标的边界框 pred_boxes = [[ii[0], ii[1], ii[2], ii[3]] for ii in list(pred[0]['boxes'].detach().numpy())] # 只保留识别的概率大约 0.5 的结果。 pred_index = [pred_score.index(x) for x in pred_score if x > 0.5] for index in pred_index: box = pred_boxes[index] cv2.rectangle(img=image, pt1=[int(box[0]), int(box[1])], pt2=[int(box[2]), int(box[3])], color=(0, 0, 225), thickness=3) texts = pred_class[index] + ":" + str(np.round(pred_score[index], 2)) font = cv2.FONT_HERSHEY_SIMPLEX cv2.putText(image, texts, (int(box[0]), int(box[1])), font, 1, (200, 255, 155), 2, cv2.LINE_AA) cv2.imshow('摄像头', image) cv2.waitKey(10)
这里想实现实时检测,但是电脑跑不动。。一卡一卡的。。
注意文字的位置,有可能他会加载在图片的外面。。。
【他能识别出那个床是我没有想到的。。。】
欢迎分享,转载请注明来源:内存溢出
评论列表(0条)