利用yolo的输出,中心和wh
用的torch框架
cap = cv2.VideoCapture('.//6.mp4')
while True:ret, frame = cap.read()ori_img = frame #存个原图用于显示frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) #opencv读取是BGR的格式img1=np.array(frame)res_img = cv2.resize(img1, (cfg["width"], cfg["height"]), interpolation = cv2.INTER_LINEAR)img = res_img.reshape(1, cfg["height"], cfg["width"], 3)#改成网络需要的大小img = torch.from_numpy(img.transpose(0,3, 1, 2)) #channel first 再加个batchimg = img.to(device).float() / 255.0#特征图后处理output = utils.utils.handel_preds(preds, cfg, device)output_boxes = utils.utils.non_max_suppression(output, conf_thres=0.7, iou_thres = 0.4)h, w, _ = ori_img.shapescale_h, scale_w = h / cfg["height"], w / cfg["width"] #用于改成原图上的相对大小for box in output_boxes[0]:box = box.tolist()obj_score = box[4]category = LABEL_NAMES[int(box[5])]x1, y1 = int(box[0] * scale_w), int(box[1] * scale_h)x2, y2 = int(box[2] * scale_w), int(box[3] * scale_h)cv2.rectangle(ori_img, (x1, y1), (x2, y2), (255, 255, 0), 2)cv2.putText(ori_img, '%.2f' % obj_score, (x1, y1 - 5), 0, 0.7, (0, 255, 0), 2)cv2.putText(ori_img, category, (x1, y1 - 25), 0, 0.7, (0, 255, 0), 2)cv2.imshow("test_result.png", ori_img) #show放在 循环标框的后面if cv2.waitKey(1) & 0xFF==ord('q'): # 按键q后breakcap.release()break
cap.release() # 对应while结束
cv2.destroyWindow()
重点
img1=np.array(frame) #有时候用的 PIL.Image.read() 返回的是个类
res_img = cv2.resize(img1, (cfg["width"], cfg["height"]), interpolation = cv2.INTER_LINEAR)
img = torch.from_numpy(img.transpose(0,3, 1, 2)) #channel first 再加个batchimg = img.to(device).float() / 255.0 #很对时候忘记归一化
还有就是
#这个得有if cv2.waitKey(1) & 0xFF==ord('q'): # 按键q后breakcap.release()break
cap.release() # 对应while结束
cv2.destroyWindow()