优化上料识别与药材识别

bsw215583320

2025-01-16 8446139d79c366fc2c44d72c2ddb963817a8ce7d

优化上料识别与药材识别

 cam_detect.py

@@ -1,6 +1,6 @@
# æ‘„像头画面测试程序

import os

import cv2
import time
import numpy as np
@@ -8,14 +8,14 @@
import win32com.client

if __name__ == '__main__':
# æ‘„像头索引号，通常为0表示第一个摄像头
    camera_index = 1
    # æ‘„像头索引号，通常为0表示第一个摄像头
    camera_index = 3

    # æ‰“开摄像头
    cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
    # è®¾ç½®åˆ†è¾¨çŽ‡
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 3800)  # å®½åº¦
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 2100)  # é«˜åº¦
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 2048)  # å®½åº¦
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1540)  # é«˜åº¦
    # æ£€æŸ¥æ‘„像头是否成功打开
    if not cap.isOpened():
        print("无法打开摄像头")
@@ -62,8 +62,9 @@
        # å°†FPS绘制在图像上
        cv2.putText(resized_frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
                    cv2.LINE_AA)
        resizeframe = cv2.resize(frame, (target_width, target_height))
        # æ˜¾ç¤ºç”»é¢
        cv2.imshow("Camera", resized_frame)
        cv2.imshow("Camera", resizeframe)

        # æ£€æµ‹æŒ‰é”®ï¼Œå¦‚果按下q键则退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
@@ -73,4 +74,4 @@
    cap.release()

    # å…³é—­æ‰€æœ‰çª—口
    cv2.destroyAllWindows()
    cv2.destroyAllWindows()

 camera_onnx.py

@@ -1,5 +1,7 @@
import os
# ä¸Šæ–™è¯†åˆ«ï¼ŒèŽ·å–上料图片，保存到data/images文件夹下


import os
import cv2
import time
import numpy as np
@@ -17,26 +19,26 @@
    for device in webcams:
        name = getattr(device, 'Name', None)
        pnp_class = getattr(device, 'PNPClass', None)
        if name is not None and 'PC Camera' in name:
        if name is not None and cam1 in name:
            # å°†è®¾å¤‡åå­—和索引添加到字典中
            webcam_dict['PC Camera'] = index
            webcam_dict[cam1] = index
            index += 1
        elif name is not None and "USB Camera" in name:
        elif name is not None and cam2 in name:
            # å°†è®¾å¤‡åå­—和索引添加到字典中
            webcam_dict['USB Camera'] = index
            webcam_dict[cam2] = index
            index += 1
    return webcam_dict


# è°ƒç”¨å¦ä¸€ä¸ªé•¿ç„¦é•œå¤´ï¼Œæ‹æ‘„清晰的局部药材图片
def get_image():
    camera2_index = webcams.get('PC Camera')
    camera2_index = webcams.get(cam2)
    print("第二个摄像头索引：" + str(camera2_index))
    # æ‰“开摄像头
    capture = cv2.VideoCapture(camera2_index, cv2.CAP_DSHOW)
    # è®¾ç½®åˆ†è¾¨çŽ‡
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, 3800)  # å®½åº¦
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 2150)  # é«˜åº¦
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)  # å®½åº¦
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)  # é«˜åº¦
    # æ£€æŸ¥æ‘„像头是否成功打开
    if not capture.isOpened():
        print("无法打开摄像头2")
@@ -46,45 +48,49 @@
    print("摄像头2分辨率:", width2, "x", height2)
    # å¾ªçŽ¯è¯»å–摄像头画面
    # Shadows name 'width' from outer scope
    count = 0
    while True:
        ret2, frame2 = capture.read()
        if not ret2:
            print("无法读取摄像头画面")
            break

        # æ˜¾ç¤ºç”»é¢
        # cv2.imshow('Output2', frame2)
        count += 1
        # 1920*1080的图像，中心裁剪640*480的区域
        a2 = int(height2 / 2 - target_height / 2)
        b2 = int(height2 / 2 + target_height / 2)
        c2 = int(width2 / 2 - target_width / 2)
        d2 = int(width2 / 2 + target_width / 2)
        cropped_frame2 = frame2[a2:b2, c2:d2]
        # è°ƒæ•´å›¾åƒå°ºå¯¸
        resized_frame2 = cv2.resize(cropped_frame2, (target_width, target_height))
        # ç”Ÿæˆä¿å­˜æ–‡ä»¶åï¼Œä»¥å½“前时间命名
        save_name = time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
        # ä¿å­˜è°ƒæ•´å°ºå¯¸åŽçš„图片
        cv2.imwrite(save_path + "2/" + save_name, resized_frame2)
        # cv2.imshow("Camera", resized_frame2)
        print("保存图片:", save_name)
        break
        # a2 = int(height2 / 2 - target_height / 2)
        # b2 = int(height2 / 2 + target_height / 2)
        # c2 = int(width2 / 2 - target_width / 2)
        # d2 = int(width2 / 2 + target_width / 2)
        # cropped_frame2 = frame2[a2:b2, c2:d2]
        if count == 2:
            # è°ƒæ•´å›¾åƒå°ºå¯¸
            resized_frame2 = cv2.resize(frame2, (target_width, target_height))
            # æ˜¾ç¤ºç”»é¢
            # cv2.imshow('Output2', resized_frame2)
            # ç”Ÿæˆä¿å­˜æ–‡ä»¶åï¼Œä»¥å½“前时间命名
            save_name2 = time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
            # ä¿å­˜è°ƒæ•´å°ºå¯¸åŽçš„图片
            cv2.imwrite(save_path + "2/" + save_name2, resized_frame2)
            # cv2.imshow("Camera", resized_frame2)
            print("保存图片:", save_name2)
            break
    # ç»“束线程
    capture.release()
    return frame2


if __name__ == '__main__':

    cam1 = "USB Camera"
    cam2 = "PC Camera"
    webcams = list_webcams()
    print(webcams)
    target_width = 1024
    target_height = 768
    save_path = "data/images/"
    # åŠ è½½ONNX模型
    session = onnxruntime.InferenceSession("model/best.onnx")
    session = onnxruntime.InferenceSession("model/loading.onnx")
    # æ‘„像头索引号，通常为0表示第一个摄像头
    camera_index = webcams.get('USB Camera')
    camera_index = webcams.get(cam1)
    print("第一个摄像头索引：" + str(camera_index))
    modelmeta = session.get_modelmeta()
    metadata_map = modelmeta.custom_metadata_map

 herb_identify.py

¶Ô±ÈÐÂÎļþ
@@ -0,0 +1,269 @@
# ä¸Šæ–™è¯†åˆ«ï¼ŒèŽ·å–上料图片，调用药材识别模型，输出识别结果


import os
import cv2
import time
import numpy as np
import onnxruntime
import win32com.client



# è®¡ç®—拉普拉斯响应的方差,判断图像的清晰度
def variance_of_laplacian(image):
    # è®¡ç®—输入图像的拉普拉斯响应的方差
    return cv2.Laplacian(image, cv2.CV_64F).var()
# èŽ·å–摄像头列表
def list_webcams():
    # åˆ›å»ºä¸€ä¸ªWMI客户端实例
    wmi = win32com.client.GetObject("winmgmts:")
    webcams = wmi.InstancesOf("Win32_PnPEntity")
    # åˆ›å»ºæ‘„像头和索引字典
    webcam_dict = {}
    index = 0
    for device in webcams:
        name = getattr(device, 'Name', None)
        pnp_class = getattr(device, 'PNPClass', None)
        if name is not None and cam1 in name:
            # å°†è®¾å¤‡åå­—和索引添加到字典中
            webcam_dict[cam1] = index
            index += 1
        elif name is not None and cam2 in name:
            # å°†è®¾å¤‡åå­—和索引添加到字典中
            webcam_dict[cam2] = index
            index += 1
    return webcam_dict


# è°ƒç”¨å¦ä¸€ä¸ªé•¿ç„¦é•œå¤´ï¼Œæ‹æ‘„清晰的局部药材图片
def get_image():
    camera2_index = webcams.get(cam2)
    print("第二个摄像头索引：" + str(camera2_index))
    # æ‰“开摄像头
    capture = cv2.VideoCapture(camera2_index, cv2.CAP_DSHOW)
    # è®¾ç½®åˆ†è¾¨çŽ‡
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, 2048)  # å®½åº¦
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 1540)  # é«˜åº¦
    # æ£€æŸ¥æ‘„像头是否成功打开
    if not capture.isOpened():
        print("无法打开摄像头2")
        exit()
    width2 = capture.get(cv2.CAP_PROP_FRAME_WIDTH)
    height2 = capture.get(cv2.CAP_PROP_FRAME_HEIGHT)
    print("摄像头2分辨率:", width2, "x", height2)
    # å¾ªçŽ¯è¯»å–摄像头画面
    # Shadows name 'width' from outer scope
    count = 0
    while True:
        ret2, frame2 = capture.read()
        if not ret2:
            print("无法读取摄像头画面")
            break

        count += 1
        # 1920*1080的图像，中心裁剪640*480的区域
        # a2 = int(height2 / 2 - target_height / 2)
        # b2 = int(height2 / 2 + target_height / 2)
        # c2 = int(width2 / 2 - target_width / 2)
        # d2 = int(width2 / 2 + target_width / 2)
        # cropped_frame2 = frame2[a2:b2, c2:d2]
        if count == 2:
            # è°ƒæ•´å›¾åƒå°ºå¯¸
            resized_frame2 = cv2.resize(frame2, (target_width, target_height))

            # é¢„处理
            herb_blob = cv2.dnn.blobFromImage(resized_frame2, 1 / 255.0, (640, 640), swapRB=True, crop=False)

            # æ¨¡åž‹æŽ¨ç†
            herb_outputs = identify_session.run(None, {identify_session.get_inputs()[0].name: herb_blob})

            herb_probabilities = herb_outputs[0]

            top_five_classes = np.argsort(herb_probabilities, axis=1)[0][-5:][::-1]
            name = ""
            for i, class_id in enumerate(top_five_classes):
                # ä¿ç•™ä¸¤ä½å°æ•°
                probability = round(herb_probabilities[0][class_id], 2)

                herb_class = herbs[class_id]
                name = name + herb_class + "=" + str(probability)
            # æ˜¾ç¤ºç”»é¢
            # cv2.imshow('Output2', resized_frame2)

            # è®¡ç®—拉普拉斯响应的方差
            laplacian = variance_of_laplacian(resized_frame2)
            # ç”Ÿæˆä¿å­˜æ–‡ä»¶åï¼Œä»¥å½“前时间命名
            save_name2 = name +"_["+ str(round(laplacian, 2)) +"]_"+ time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
            # åˆ¤æ–­å›¾åƒçš„清晰度
            # ä¿å­˜è°ƒæ•´å°ºå¯¸åŽçš„图片
            if laplacian > 1500:
                c_ = save_path + "2/c/"
                # åˆ¤æ–­æ–‡ä»¶æ˜¯å¦å­˜åœ¨ï¼Œä¸å­˜åœ¨åˆ™åˆ›å»º
                if not os.path.exists(c_):
                    os.makedirs(c_)
                cv2.imwrite(c_ + save_name2, resized_frame2)
            else:
                n_ = save_path + "2/n/"
                # åˆ¤æ–­æ–‡ä»¶æ˜¯å¦å­˜åœ¨ï¼Œä¸å­˜åœ¨åˆ™åˆ›å»º
                if not os.path.exists(n_):
                    os.makedirs(n_)
                cv2.imwrite(n_ + save_name2, resized_frame2)
            # cv2.imshow("Camera", resized_frame2)
            print("保存图片:", save_name2)
            break
    # ç»“束线程
    capture.release()


if __name__ == '__main__':

    cam1 = "USB Camera"
    cam2 = "PC Camera"
    webcams = list_webcams()
    print(webcams)
    target_width = 1024
    target_height = 768
    save_path = "data/images/"
    # åŠ è½½ONNX模型
    session = onnxruntime.InferenceSession("model/loading.onnx")
    identify_session = onnxruntime.InferenceSession("model/herb_identify.onnx")
    herbs = eval(identify_session.get_modelmeta().custom_metadata_map['names'])
    # æ‘„像头索引号，通常为0表示第一个摄像头
    camera_index = webcams.get(cam1)
    print("第一个摄像头索引：" + str(camera_index))
    modelmeta = session.get_modelmeta()
    metadata_map = modelmeta.custom_metadata_map
    classes = eval(metadata_map['names'])
    # æ‰“开摄像头
    cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
    # è®¾ç½®åˆ†è¾¨çŽ‡
    # cap.set(cv2.CAP_PROP_FRAME_WIDTH, 3840)  # å®½åº¦
    # cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 2160)  # é«˜åº¦
    # æ£€æŸ¥æ‘„像头是否成功打开
    if not cap.isOpened():
        print("无法打开摄像头")
        exit()

    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    print("摄像头分辨率:", width, "x", height)

    # ç›®æ ‡å›¾åƒå°ºå¯¸

    # è®¡æ—¶å™¨
    frame_count = 0
    start_time = time.time()
    stime = time.time()

    if not os.path.exists(save_path):
        os.makedirs(save_path)
    class_old = "1"
    count = 0
    status = "没有上料"
    # å¾ªçŽ¯è¯»å–摄像头画面
    while True:
        # ç¡çœ 100毫秒
        time.sleep(0.1)

        ret, frame = cap.read()

        if not ret:
            print("无法读取摄像头画面")
            break

        resized_frame = frame
        if height > target_height and width > target_width:
            # 1920*1080的图像，中心裁剪640*480的区域
            a = int(height / 2 - target_height / 2)
            b = int(height / 2 + target_height / 2)
            c = int(width / 2 - target_width / 2)
            d = int(width / 2 + target_width / 2)
            print(a, b, c, d)
            cropped_frame = frame[a:b, c:d]
            # è°ƒæ•´å›¾åƒå°ºå¯¸
            resized_frame = cv2.resize(cropped_frame, (target_width, target_height))

        # èŽ·å–当前时间
        current_time = time.time()

        # å¦‚果距离上一次保存已经过去1秒，则保存当前画面
        # if current_time - start_time >= 3.0:
        #     # ç”Ÿæˆä¿å­˜æ–‡ä»¶åï¼Œä»¥å½“前时间命名
        #     save_name = time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
        #     # ä¿å­˜è°ƒæ•´å°ºå¯¸åŽçš„图片
        #     cv2.imwrite(save_path + save_name, frame)
        #     print("保存图片:", save_name)
        #     # é‡ç½®è®¡æ—¶å™¨
        #     start_time = time.time()

        # é¢„处理
        blob = cv2.dnn.blobFromImage(resized_frame, 1 / 255.0, (640, 640), swapRB=True, crop=False)

        # æ¨¡åž‹æŽ¨ç†
        outputs = session.run(None, {session.get_inputs()[0].name: blob})
        #
        # print(outputs)
        # åº”用softmax函数
        probabilities = outputs[0]

        # æ‰¾åˆ°æœ€å¤§æ¦‚率的类别
        predicted_class = np.argmax(probabilities, axis=1)[0]
        max_probability = np.max(probabilities, axis=1)[0]
        class_ = classes[predicted_class]

        # è®¡ç®—类型重复的次数，类别更换之后重新计数
        if class_ != class_old:
            count = 0
        else:
            count += 1
        class_old = class_
        print(f"{class_}:{count}: {max_probability}")
        # åˆ¤æ–­æ˜¯å¦ä¸Šæ–™å¹¶ä¸”上料次数大于10次
        if class_ == "shangliao" and count > 10:
            status = "正在上料"
            # æ¯éš”3秒取一帧图像
            # å¦‚果距离上一次保存已经过去1秒，则保存当前画面
            if current_time - stime >= 10.0:
                save_name = time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
                # ä¿å­˜è°ƒæ•´å°ºå¯¸åŽçš„图片
                path_ = save_path + "1/"
                if not os.path.exists(path_):
                    os.makedirs(path_)
                cv2.imwrite(path_ + save_name, resized_frame)
                # é‡ç½®è®¡æ—¶å™¨
                stime = time.time()
                get_image()
        else:
            status = "没有上料"

        print(status)
        # # æ‰¾åˆ°æ¦‚率较高的前十个类别
        # top_ten_classes = np.argsort(probabilities, axis=1)[0][-1:]
        #
        # # è¾“出前十个类别
        # print("Top 5 Classes:")
        # for i in top_ten_classes:
        #     print(f"{classes[i]}: {probabilities[0][i]}")

        # è®¡ç®—帧速率
        frame_count += 1
        end_time = time.time()
        elapsed_time = end_time - start_time
        fps = frame_count / elapsed_time
        # print(f"FPS: {fps:.2f}")
        # å°†FPS绘制在图像上
        cv2.putText(resized_frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
                    cv2.LINE_AA)
        # æ˜¾ç¤ºç”»é¢
        cv2.imshow("Camera", resized_frame)

        # æ£€æµ‹æŒ‰é”®ï¼Œå¦‚果按下q键则退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # å…³é—­æ‘„像头
    cap.release()

    # å…³é—­æ‰€æœ‰çª—口
    cv2.destroyAllWindows()

 laplacian.py

¶Ô±ÈÐÂÎļþ
@@ -0,0 +1,124 @@
import os

import cv2
from scipy.fftpack import fftshift, fft2
import numpy as np


def variance_of_laplacian(image):
    # è®¡ç®—输入图像的拉普拉斯响应的方差
    return cv2.Laplacian(image, cv2.CV_64F).var()

def frequency_domain_analysis(gray):

    # è®¡ç®—傅里叶变换并移动到中心
    f_transform = fftshift(fft2(gray))
    # è®¡ç®—频谱幅度
    magnitude_spectrum = 20 * np.log(np.abs(f_transform))
    # è®¡ç®—高频成分的比例
    high_frequency_magnitude = np.sum(magnitude_spectrum > np.mean(magnitude_spectrum))
    total_magnitude = magnitude_spectrum.size
    ratio = high_frequency_magnitude / total_magnitude
    return ratio
def edge_gradient(gray):

    sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
    magnitude = np.sqrt(sobelx**2 + sobely**2)
    return np.mean(magnitude)

def testOne():
    # åŠ è½½å›¾åƒå¹¶è½¬æ¢ä¸ºç°åº¦å›¾
    image = cv2.imread('H:/images/1/20250114193113.jpg')
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # è®¡ç®—图像的清晰度分数
    score = variance_of_laplacian(gray)
    ratio = frequency_domain_analysis(gray)
    edge_gradient_score = edge_gradient(gray)

    print("Image clarity score: ", score)
    print("High frequency ratio: ", ratio)
    print("Edge gradient score: ", edge_gradient_score)

    # æ ¹æ®ç»éªŒè®¾ç½®é˜ˆå€¼æ¥åˆ¤æ–­å›¾åƒæ˜¯å¦æ¨¡ç³Š
    threshold = 100  # è¿™ä¸ªé˜ˆå€¼å¯ä»¥æ ¹æ®å®žé™…情况调整
    if score < threshold:
        print("The image is considered blurry.")
    else:
        print("The image is considered clear.")


def batch_score():
    global name, join
    # è¯»å–文件夹下的所有图像分别计算分数
    folder_path = 'H:/images/2'
    for filename in os.listdir(folder_path):
        if filename.endswith('.jpg') or filename.endswith('.png'):
            image_path = os.path.join(folder_path, filename)
            image = cv2.imread(image_path)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            score = variance_of_laplacian(gray)
            ratio = frequency_domain_analysis(gray)
            edge_gradient_score = edge_gradient(gray)
            # æ‰€æœ‰ç»“果保留两位小数
            score = round(score, 2)
            ratio = round(ratio, 2)
            edge_gradient_score = round(edge_gradient_score, 2)
            name = "score=" + str(score) + "ratio=" + str(ratio) + "edge_gradient_score=" + str(
                edge_gradient_score) + ".jpg"
            path_ = folder_path + "/22/"
            # åˆ¤æ–­æ–‡ä»¶å¤¹æ˜¯å¦å­˜åœ¨ï¼Œä¸å­˜åœ¨åˆ™æ–°å»º
            if not os.path.exists(path_):
                os.makedirs(path_)
            join = path_ + name
            print(join)
            cv2.imwrite(join, image)


def cam_score():
    # è°ƒç”¨æ‘„像头计算采集到图像的得分
    camera_index = 2
    print("第二个摄像头索引：" + str(camera_index))
    # æ‰“开摄像头
    capture = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
    # è®¾ç½®åˆ†è¾¨çŽ‡
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1921)
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 1081)
    # æ£€æŸ¥æ‘„像头是否成功打开
    if not capture.isOpened():
        print("无法打开摄像头2")
        exit()
    # å¾ªçŽ¯è¯»å–画面
    while True:
        ret, frame = capture.read()
        if not ret:
            print("无法读取摄像头画面")
            break
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        score = variance_of_laplacian(gray)
        # ratio = frequency_domain_analysis(gray)
        # edge_gradient_score = edge_gradient(gray)
        # å°†FPS绘制在图像上
        cv2.putText(frame, f"score={str(score)}"
                    , (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
        # cv2.putText(frame, f"ratio={str(ratio)}"
        #             , (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
        # cv2.putText(frame, f"edge_gradient_score={str(edge_gradient_score)}"
        #             , (10, 110), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
        # æ˜¾ç¤ºç”»é¢
        cv2.imshow("Camera", frame)
        # æ£€æµ‹æŒ‰é”®ï¼Œå¦‚果按下q键则退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    # å…³é—­æ‘„像头
    capture.release()
    # å…³é—­æ‰€æœ‰çª—口
    cv2.destroyAllWindows()


if __name__ == '__main__':
    testOne()
    # batch_score()
    # cam_score()

 model/herb_identify.onnx

Binary files differ

 model/loading.onnx

Binary files differ

 test_herb_identify.py

¶Ô±ÈÐÂÎļþ
@@ -0,0 +1,32 @@

import os
import cv2
import numpy as np
import onnxruntime

if __name__ == '__main__':
    img = cv2.imread("data/images/2/20250114200010.jpg")
    identify_session = onnxruntime.InferenceSession("model/herb_identify.onnx")
    herbs = eval(identify_session.get_modelmeta().custom_metadata_map['names'])
    # é¢„处理
    herb_blob = cv2.dnn.blobFromImage(img, 1 / 255.0, (640, 640), swapRB=True, crop=False)
    # æ¨¡åž‹æŽ¨ç†
    herb_outputs = identify_session.run(None, {identify_session.get_inputs()[0].name: herb_blob})
    herb_probabilities = herb_outputs[0]

    top_five_classes = np.argsort(herb_probabilities, axis=1)[0][-5:][::-1]
    name = ""
    for i, class_id in enumerate(top_five_classes):
        # ä¿ç•™ä¸¤ä½å°æ•°
        probability = round(herb_probabilities[0][class_id], 2)

        herb_class = herbs[class_id]
        name = name + " " + herb_class + ":" + str(probability)
        # å°†è¯†åˆ«ç»“果打印到图片上
        cv2.putText(img, herb_class + ":" + str(probability), (10, (i+1)*40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)

    # æ˜¾ç¤ºå›¾ç‰‡
    cv2.imshow("Output", img)
    print(name)
    # ç­‰å¾…
    cv2.waitKey(0)