Thu Jul 17 18:47:00 CST 2025 inscode

main.py

-print('欢迎来到 InsCode')
+import cv2
\ No newline at end of file
+import numpy as np
+import onnxruntime as ort
+
+class YOLOv5_Lite:
+    def __init__(self, onnx_model_path, class_names_path, conf_threshold=0.5, iou_threshold=0.5):
+        """
+        初始化 YOLOv5_Lite 类
+        :param onnx_model_path: ONNX 模型文件路径
+        :param class_names_path: 类别名称文件路径
+        :param conf_threshold: 置信度阈值
+        :param iou_threshold: IOU 阈值
+        """
+        # 创建 ONNX Runtime 会话
+        so = ort.SessionOptions()
+        so.log_severity_level = 3
+        self.session = ort.InferenceSession(onnx_model_path, so)
+        # 获取输入输出节点名称
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_name = self.session.get_outputs()[0].name
+        # 加载类别名称
+        with open(class_names_path, 'r') as f:
+            self.class_names = [line.strip() for line in f.readlines()]
+        # 设置阈值
+        self.conf_threshold = conf_threshold
+        self.iou_threshold = iou_threshold
+
+    def preprocess(self, image):
+        """
+        图像预处理
+        :param image: 输入图像
+        :return: 预处理后的图像数据
+        """
+        # 转为 RGB 格式
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        # 缩放至模型输入尺寸
+        img = cv2.resize(image, (640, 640))
+        # 归一化
+        img = img.astype(np.float32) / 255.0
+        # 调整维度
+        img = np.expand_dims(img.transpose(2, 0, 1), axis=0).astype(np.float32)
+        return img
+
+    def postprocess(self, outputs, img_shape, origin_shape):
+        """
+        后处理函数
+        :param outputs: 模型输出
+        :param img_shape: 输入图像尺寸
+        :param origin_shape: 原始图像尺寸
+        :return: 检测框坐标、类别名称、置信度
+        """
+        # 获取模型输出
+        outputs = outputs[0]
+        # 过滤低置信度的检测结果
+        outputs = outputs[outputs[:, 4] >= self.conf_threshold]
+        # 初始化结果列表
+        boxes = []
+        scores = []
+        classes = []
+        # 遍历每个检测结果
+        for output in outputs:
+            # 获取类别索引和置信度
+            class_id = np.argmax(output[5:])
+            confidence = output[5 + class_id]
+            # 过滤低置信度的检测结果
+            if confidence < self.conf_threshold:
+                continue
+            # 获取边界框坐标
+            x, y, w, h = output[:4]
+            # 反归一化
+            x *= img_shape[1]
+            y *= img_shape[0]
+            w *= img_shape[1]
+            h *= img_shape[0]
+            # 转换为左上角和右下角坐标
+            x1 = int(x - w / 2)
+            y1 = int(y - h / 2)
+            x2 = int(x + w / 2)
+            y2 = int(y + h / 2)
+            # 调整边界框坐标以适配原始图像尺寸
+            x1 = max(0, min(x1, origin_shape[1]))
+            y1 = max(0, min(y1, origin_shape[0]))
+            x2 = max(0, min(x2, origin_shape[1]))
+            y2 = max(0, min(y2, origin_shape[0]))
+            # 添加到结果列表
+            boxes.append([x1, y1, x2, y2])
+            scores.append(float(confidence))
+            classes.append(self.class_names[class_id])
+        # 非极大值抑制
+        indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_threshold, self.iou_threshold)
+        # 过滤后的检测结果
+        boxes = [boxes[i] for i in indices]
+        classes = [classes[i] for i in indices]
+        scores = [scores[i] for i in indices]
+        return boxes, classes, scores
+
+    def detect(self, image):
+        """
+        物体检测函数
+        :param image: 输入图像
+        :return: 检测框坐标、类别名称、置信度
+        """
+        # 图像预处理
+        img = self.preprocess(image)
+        # 模型推理
+        outputs = self.session.run([self.output_name], {self.input_name: img})
+        # 后处理
+        boxes, classes, scores = self.postprocess(outputs, img.shape[2:], image.shape[:2])
+        return boxes, classes, scores
+
+if __name__ == "__main__":
+    # 模型路径
+    onnx_model_path = "best (3).onnx"
+    # 类别文件路径
+    class_names_path = "rubbish.names"
+    # 初始化 YOLOv5_Lite 类
+    detector = YOLOv5_Lite(onnx_model_path, class_names_path)
+    # 打开摄像头
+    cap = cv2.VideoCapture(0)
+    while True:
+        # 读取一帧图像
+        ret, frame = cap.read()
+        if not ret:
+            break
+        # 检测物体
+        boxes, classes, scores = detector.detect(frame)
+        # 在图像上绘制检测结果
+        for box, cls, score in zip(boxes, classes, scores):
+            x1, y1, x2, y2 = box
+            # 绘制边界框
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            # 显示类别名称和置信度
+            label = f"{cls}: {score:.2f}"
+            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+        # 显示图像
+        cv2.imshow("Object Detection", frame)
+        # 按下 'q' 键退出循环
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
+    # 释放摄像头和关闭窗口
+    cap.release()
+    cv2.destroyAllWindows()
\ No newline at end of file

requirements.txt

+onnxruntime
\ No newline at end of file

rubbish.names

+person
+bicycle
+car
+motorcycle
\ No newline at end of file