前言

使用python，进行对图片中物品信息识别

开始

项目架构

main.py

from detectorYolo import YOLODetector

dYolo = YOLODetector("yolo11n.pt")

if __name__ == "__main__":
    print("===============================================")
    # 识别图片中的物品
    result = dYolo.detect("./images/yolo4.jpg")
    # result = dYolo.detect("./images/yolo5.jpg")
    print(result)

detectorYolo.py

from ultralytics import YOLO


class YOLODetector:

    def __init__(self, model_path="yolo11n.pt", device=None):
        """
        初始化模型
        :param model_path: 模型路径
        :param device: cuda / cpu
        """

        self.model = YOLO(model_path)

        if device:
            self.model.to(device)


    def detect(self, image):
        """
        图片检测
        :param image:
            图片路径 / numpy / PIL图片
        :return:
            检测结果
        """

        results = self.model(image)

        objects = []

        for result in results:
            boxes = result.boxes
            for box in boxes:
                cls_id = int(box.cls[0])
                confidence = float(box.conf[0])
                xyxy = box.xyxy[0]
                objects.append({
                    "name": self.model.names[cls_id],
                    "confidence": round(confidence,3), "box":[
                        int(xyxy[0]),
                        int(xyxy[1]),
                        int(xyxy[2]),
                        int(xyxy[3])

                    ]
                })
        return objects

    def save_result(self, image, path):
        """
        保存带框图片
        """
        results = self.model(image)
        for result in results:
            result.save(filename=path)

总结

主要 环境安装比较麻烦~，目前这个只是一个小用例，因为还要涉及图片底色如何进行训练