基于python使用yolo模型的水果识别代码

好的，以下是使用Python和YOLO模型进行水果识别的代码：

# 导入相关的库
import cv2
import numpy as np
import argparse

# 设置命令行参数
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True, help="path to input image")
ap.add_argument("-c", "--config", required=True, help="path to yolo config file")
ap.add_argument("-w", "--weights", required=True, help="path to yolo pre-trained weights")
ap.add_argument("-cl", "--classes", required=True, help="path to text file containing class names")
args = vars(ap.parse_args())

# 加载类别名
classes = None
with open(args["classes"], 'r') as f:
    classes = [line.strip() for line in f.readlines()]

# 加载模型配置和权重
net = cv2.dnn.readNetFromDarknet(args["config"], args["weights"])

# 加载输入图像并进行预处理
image = cv2.imread(args["image"])
blob = cv2.dnn.blobFromImage(image, 1/255.0, (416, 416), swapRB=True, crop=False)

# 设置模型的输入和输出节点
net.setInput(blob)
output_layers_names = net.getUnconnectedOutLayersNames()
layerOutputs = net.forward(output_layers_names)

# 初始化输出结果
boxes = []
confidences = []
classIDs = []

# 循环遍历每个输出层，提取检测结果
for output in layerOutputs:
    for detection in output:
        scores = detection[5:]
        classID = np.argmax(scores)
        confidence = scores[classID]

        # 过滤掉置信度低的检测结果
        if confidence > 0.5:
            box = detection[0:4] * np.array([image.shape[1], image.shape[0], image.shape[1], image.shape[0]])
            (centerX, centerY, width, height) = box.astype("int")

            # 计算边框的左上角坐标
            x = int(centerX - (width / 2))
            y = int(centerY - (height / 2))

            # 更新输出结果
            boxes.append([x, y, int(width), int(height)])
            confidences.append(float(confidence))
            classIDs.append(classID)

# 对输出结果进行NMS处理，去除冗余的检测结果
idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)

# 在图像上绘制检测结果
if len(idxs) > 0:
    for i in idxs.flatten():
        (x, y) = (boxes[i][0], boxes[i][1])
        (w, h) = (boxes[i][2], boxes[i][3])
        color = [int(c) for c in COLORS[classIDs[i]]]
        cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
        text = "{}: {:.4f}".format(classes[classIDs[i]], confidences[i])
        cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

# 显示输出图像
cv2.imshow("Image", image)
cv2.waitKey(0)
cv2.destroyAllWindows()

使用时可以在命令行中运行以下命令：

python fruit_detection.py --image input_image.jpg --config yolo.cfg --weights yolo.weights --classes classes.txt

其中，`input_image.jpg`是要识别的图像，`yolo.cfg`和`yolo.weights`是YOLO模型的配置文件和权重文件，`classes.txt`是包含类别名的文本文件。