|
@@ -0,0 +1,265 @@
|
|
|
+import os
|
|
|
+import sys
|
|
|
+import onnx
|
|
|
+import onnxruntime as ort
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
+import time # 导入时间模块
|
|
|
+
|
|
|
+# CLASSES = ['person']
|
|
|
+CLASSES = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
|
|
|
+ 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
|
|
|
+ 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
|
|
|
+ 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
|
|
|
+ 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
|
|
|
+ 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
|
|
|
+ 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
|
|
|
+ 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
|
|
|
+ 'hair drier', 'toothbrush'] # coco80类别
|
|
|
+
|
|
|
+
|
|
|
+class Yolov5ONNX(object):
|
|
|
+ def __init__(self, onnx_path):
|
|
|
+ onnx_model = onnx.load(onnx_path)
|
|
|
+ try:
|
|
|
+ onnx.checker.check_model(onnx_model)
|
|
|
+ except Exception:
|
|
|
+ print("Model incorrect")
|
|
|
+ else:
|
|
|
+ print("Model correct")
|
|
|
+
|
|
|
+ # 设置 ONNX Runtime 会尝试优先使用 CUDAExecutionProvider,如果不可用则使用 CPUExecutionProvider
|
|
|
+ self.onnx_session = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
|
|
+
|
|
|
+ # 检查是否使用了 CUDA
|
|
|
+ providers = self.onnx_session.get_providers()
|
|
|
+ if 'CUDAExecutionProvider' in providers:
|
|
|
+ print("Using CUDA for inference.")
|
|
|
+ else:
|
|
|
+ print("CUDA is not available, using CPU for inference.")
|
|
|
+
|
|
|
+ self.input_name = self.get_input_name()
|
|
|
+ self.output_name = self.get_output_name()
|
|
|
+ self.input_size = (640, 640)
|
|
|
+
|
|
|
+ def get_input_name(self):
|
|
|
+ input_name = []
|
|
|
+ for node in self.onnx_session.get_inputs():
|
|
|
+ input_name.append(node.name)
|
|
|
+ return input_name
|
|
|
+
|
|
|
+ def get_output_name(self):
|
|
|
+ output_name = []
|
|
|
+ for node in self.onnx_session.get_outputs():
|
|
|
+ output_name.append(node.name)
|
|
|
+ return output_name
|
|
|
+
|
|
|
+ def get_input_feed(self, image_numpy):
|
|
|
+ input_feed = {}
|
|
|
+ for name in self.input_name:
|
|
|
+ input_feed[name] = image_numpy
|
|
|
+ return input_feed
|
|
|
+
|
|
|
+ def inference(self, img):
|
|
|
+ h, w, _ = img.shape
|
|
|
+ new_w, new_h = self.input_size
|
|
|
+
|
|
|
+ scale = min(new_w / w, new_h / h)
|
|
|
+ new_w = int(w * scale)
|
|
|
+ new_h = int(h * scale)
|
|
|
+
|
|
|
+ img_resized = cv2.resize(img, (new_w, new_h))
|
|
|
+
|
|
|
+ padded_img = np.zeros((self.input_size[1], self.input_size[0], 3), dtype=np.uint8)
|
|
|
+ pad_x = (self.input_size[0] - new_w) // 2
|
|
|
+ pad_y = (self.input_size[1] - new_h) // 2
|
|
|
+ padded_img[pad_y:pad_y + new_h, pad_x:pad_x + new_w] = img_resized
|
|
|
+
|
|
|
+ img_rgb = cv2.cvtColor(padded_img, cv2.COLOR_BGR2RGB).transpose(2, 0, 1)
|
|
|
+ img_rgb = img_rgb.astype(np.float32)
|
|
|
+ img_rgb /= 255.0
|
|
|
+ img_rgb = np.expand_dims(img_rgb, axis=0)
|
|
|
+
|
|
|
+ input_feed = self.get_input_feed(img_rgb)
|
|
|
+
|
|
|
+ start_time = time.time() # 记录开始时间
|
|
|
+ pred = self.onnx_session.run(None, input_feed)[0]
|
|
|
+ end_time = time.time() # 记录结束时间
|
|
|
+
|
|
|
+ inference_time = end_time - start_time # 计算推理时间
|
|
|
+ print(f"Inference time: {inference_time:.4f} seconds") # 输出推理时间
|
|
|
+
|
|
|
+ return pred, padded_img
|
|
|
+
|
|
|
+
|
|
|
+def nms(dets, thresh):
|
|
|
+ x1 = dets[:, 0]
|
|
|
+ y1 = dets[:, 1]
|
|
|
+ x2 = dets[:, 2]
|
|
|
+ y2 = dets[:, 3]
|
|
|
+
|
|
|
+ areas = (y2 - y1 + 1) * (x2 - x1 + 1)
|
|
|
+ scores = dets[:, 4]
|
|
|
+
|
|
|
+ keep = []
|
|
|
+ index = scores.argsort()[::-1]
|
|
|
+
|
|
|
+ while index.size > 0:
|
|
|
+ i = index[0]
|
|
|
+ keep.append(i)
|
|
|
+
|
|
|
+ x11 = np.maximum(x1[i], x1[index[1:]])
|
|
|
+ y11 = np.maximum(y1[i], y1[index[1:]])
|
|
|
+ x22 = np.minimum(x2[i], x2[index[1:]])
|
|
|
+ y22 = np.minimum(y2[i], y2[index[1:]])
|
|
|
+
|
|
|
+ w = np.maximum(0, x22 - x11 + 1)
|
|
|
+ h = np.maximum(0, y22 - y11 + 1)
|
|
|
+
|
|
|
+ overlaps = w * h
|
|
|
+ ious = overlaps / (areas[i] + areas[index[1:]] - overlaps)
|
|
|
+
|
|
|
+ idx = np.where(ious <= thresh)[0]
|
|
|
+ index = index[idx + 1]
|
|
|
+ return keep
|
|
|
+
|
|
|
+
|
|
|
+def xywh2xyxy(x):
|
|
|
+ y = np.copy(x)
|
|
|
+ y[:, 0] = x[:, 0] - x[:, 2] / 2
|
|
|
+ y[:, 1] = x[:, 1] - x[:, 3] / 2
|
|
|
+ y[:, 2] = x[:, 0] + x[:, 2] / 2
|
|
|
+ y[:, 3] = x[:, 1] + x[:, 3] / 2
|
|
|
+
|
|
|
+ return y
|
|
|
+
|
|
|
+
|
|
|
+def filter_box(org_box, conf_thres, iou_thres):
|
|
|
+ org_box = np.squeeze(org_box)
|
|
|
+ conf = org_box[..., 4] > conf_thres
|
|
|
+ box = org_box[conf == True]
|
|
|
+
|
|
|
+ if box.size == 0:
|
|
|
+ return np.array([]) # 确保返回一个空的numpy数组
|
|
|
+
|
|
|
+ cls_cinf = box[..., 5:]
|
|
|
+ cls = [int(np.argmax(cls_cinf[i])) for i in range(len(cls_cinf))]
|
|
|
+
|
|
|
+ # 只保留类别为“person”的框
|
|
|
+ person_boxes = [box[i] for i in range(len(cls)) if cls[i] == 0] # 类别“person”对应的索引是0
|
|
|
+
|
|
|
+ if len(person_boxes) == 0:
|
|
|
+ return np.array([]) # 如果没有“person”框,返回一个空的numpy数组
|
|
|
+
|
|
|
+ person_boxes = np.array(person_boxes)
|
|
|
+ person_boxes = xywh2xyxy(person_boxes)
|
|
|
+
|
|
|
+ person_out_box = nms(person_boxes, iou_thres)
|
|
|
+ output = [person_boxes[k] for k in person_out_box]
|
|
|
+
|
|
|
+ return np.array(output)
|
|
|
+
|
|
|
+
|
|
|
+def draw(image, box_data):
|
|
|
+ if box_data.size == 0:
|
|
|
+ return image
|
|
|
+
|
|
|
+ boxes = box_data[..., :4].astype(np.int32)
|
|
|
+ scores = box_data[..., 4]
|
|
|
+ classes = box_data[..., 5].astype(np.int32)
|
|
|
+ for box, score, cl in zip(boxes, scores, classes):
|
|
|
+ top, left, right, bottom = box
|
|
|
+
|
|
|
+ cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2)
|
|
|
+ cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score),
|
|
|
+ (top, left),
|
|
|
+ cv2.FONT_HERSHEY_SIMPLEX,
|
|
|
+ 0.6, (0, 0, 255), 2)
|
|
|
+ return image
|
|
|
+
|
|
|
+
|
|
|
+# def main():
|
|
|
+# onnx_path = 'D:/Thework/yolov5-5.0/yolov5-5.0/runs/train/exp/weights/yolov5s.onnx'
|
|
|
+# model = Yolov5ONNX(onnx_path)
|
|
|
+#
|
|
|
+# cap = cv2.VideoCapture(0) # 打开摄像头
|
|
|
+# if not cap.isOpened():
|
|
|
+# print("无法打开摄像头")
|
|
|
+# sys.exit(0)
|
|
|
+#
|
|
|
+# while True:
|
|
|
+# ret, frame = cap.read()
|
|
|
+# if not ret:
|
|
|
+# print("无法读取摄像头图像")
|
|
|
+# break
|
|
|
+#
|
|
|
+# start_time = time.time() # 记录开始时间
|
|
|
+# output, org_img = model.inference(frame)
|
|
|
+# outbox = filter_box(output, 0.5, 0.5)
|
|
|
+# org_img = draw(org_img, outbox)
|
|
|
+# end_time = time.time() # 记录结束时间
|
|
|
+#
|
|
|
+# inference_time = end_time - start_time # 计算推理时间
|
|
|
+# print(f"Total processing time: {inference_time:.4f} seconds") # 输出总处理时间
|
|
|
+#
|
|
|
+# cv2.imshow('result', org_img)
|
|
|
+#
|
|
|
+# if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
|
+# break
|
|
|
+#
|
|
|
+# cap.release()
|
|
|
+# cv2.destroyAllWindows()
|
|
|
+#
|
|
|
+#
|
|
|
+# if __name__ == "__main__":
|
|
|
+# main()
|
|
|
+def main():
|
|
|
+ onnx_path = 'yolov5s.onnx'
|
|
|
+ model = Yolov5ONNX(onnx_path)
|
|
|
+
|
|
|
+ # cap = cv2.VideoCapture(0) # 打开摄像头
|
|
|
+ cap = cv2.VideoCapture(7, cv2.CAP_V4L2)
|
|
|
+ cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
|
|
|
+ cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
|
|
|
+ if not cap.isOpened():
|
|
|
+ print("无法打开摄像头")
|
|
|
+ sys.exit(0)
|
|
|
+
|
|
|
+ last_save_time = time.time() # 记录最后一次保存的时间
|
|
|
+ frame_count = 0 # 用于命名保存的图像
|
|
|
+
|
|
|
+ # 创建保存图片的文件夹
|
|
|
+ output_dir = "saved_images"
|
|
|
+ if not os.path.exists(output_dir):
|
|
|
+ os.makedirs(output_dir)
|
|
|
+
|
|
|
+ while True:
|
|
|
+ ret, frame = cap.read()
|
|
|
+ if not ret:
|
|
|
+ print("无法读取摄像头图像")
|
|
|
+ break
|
|
|
+
|
|
|
+ output, org_img = model.inference(frame)
|
|
|
+ outbox = filter_box(output, 0.5, 0.5)
|
|
|
+ org_img = draw(org_img, outbox)
|
|
|
+
|
|
|
+ # 如果距离上次保存超过2秒,则保存图像
|
|
|
+ if time.time() - last_save_time >= 2:
|
|
|
+ frame_count += 1
|
|
|
+ image_path = os.path.join(output_dir, f'result_{frame_count}.jpg')
|
|
|
+ cv2.imwrite(image_path, org_img) # 保存图像
|
|
|
+ print(f"Image saved: {image_path}") # 输出保存路径
|
|
|
+ last_save_time = time.time() # 更新最后一次保存的时间
|
|
|
+
|
|
|
+ # 如果保存了5张图片,退出循环
|
|
|
+ if frame_count >= 5:
|
|
|
+ print("保存了5张图片,程序退出")
|
|
|
+ break
|
|
|
+
|
|
|
+ cap.release()
|
|
|
+ cv2.destroyAllWindows()
|
|
|
+
|
|
|
+if __name__ == "__main__":
|
|
|
+ main()
|
|
|
+
|
|
|
+
|