通过cv对视频文件的读取，进行肺结节的实时目标检测

import torch import numpy as np import cv2 from torchvision.transforms import transforms from torchvision.models import detection # 预处理 transformer = transforms.Compose([transforms.ToTensor()]) # 运行设备 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' # 类名转换 id_to_class = {'1':'A'} # 加载网络模型 model = detection.fasterrcnn_mobilenet_v3_large_fpn(weight=None, progress=True,num_classes=2) model.load_state_dict(torch.load('./log/pretrained_epoch_16,train_iou_0.978933572769165,test_iou_0.652420699596405.pth')) model.to(DEVICE) # 读取视频 video_path = './video/ezgif-2-e9741894d1.mp4' cap = cv2.VideoCapture(video_path) while cap.isOpened(): ret,fram = cap.read() # 读取一帧数据 img = fram img = np.array(img) img_formodel = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 模型是单通道 # object detection 部分 with torch.no_grad(): model.eval() # 测试模式 img_tensor = transformer(img_formodel) # 预处理变换 img_tensor = img_tensor.to(DEVICE) predict = model([img_tensor]) # 网络预测，输入应该是一个list boxes = predict[0]['boxes'].cpu().detach().numpy().astype(int) # 获取边界框 labels = predict[0]['labels'].cpu().detach().numpy().astype(np.uint8) # 获取类别 scores = predict[0]['scores'].cpu().detach().numpy() # 置信度 for idx in range(boxes.shape[0]): if scores[idx] > 0.5: # 显示高得置信度 x1, y1, x2, y2 = boxes[idx][0], boxes[idx][1], boxes[idx][2], boxes[idx][3] names = id_to_class[str(labels[idx])] score_out = round(scores[idx], 3) cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2) # 绘制矩形边界框 cv2.putText(img, text='{},{}'.format(names, str(score_out)), org=(x1, y1 + 10), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=1, lineType=cv2.LINE_AA, color=(0, 0, 255)) cv2.imshow('video_detection', img) # 退出检测 if cv2.waitKey(1) & 0xFF == ord('q'): cap.release() cv2.destroyAllWindows() break