refactor: 重构中药识别项目 分为onnx版和openvino版

baoshiwei

2025-04-22 88fc0f9f9b7fd3eb81c958ca41ed822cf3657c47

 openvino/herb_ai.py

copy from herb_ai.py
copy to openvino/herb_ai.py
Îļþ´Ó herb_ai.py ¸´ÖÆ
@@ -11,6 +11,7 @@
import multiprocessing
from safety_detect import SAFETY_DETECT
from identifier import IDENTIFIER

import os
from logger_config import logger
import threading
@@ -46,7 +47,7 @@

# è°ƒç”¨å¦ä¸€ä¸ªé•¿ç„¦é•œå¤´ï¼Œæ‹æ‘„清晰的局部药材图片
def get_image():
    herb_identifier = IDENTIFIER("model/herb_identify.onnx")
    herb_identifier = IDENTIFIER("./model/herb_id")
    logger.info("识别线程启动")
    global is_loaded, class_count, class_count_max, class_sum
    camera2_index = config['cam']['cam2']
@@ -179,9 +180,7 @@
    print("摄像头分辨率:", width, "x", height)
    logger.info(f"摄像头分辨率:, {width}, x, {height}")
    # ç›®æ ‡å›¾åƒå°ºå¯¸
    # è®¡æ—¶å™¨
    frame_count = 0
    start_time = time.time()

    stime = time.time()
    if not os.path.exists(save_path):
        os.makedirs(save_path)
@@ -197,7 +196,7 @@

    # å¾ªçŽ¯è¯»å–摄像头画面
    while True:
        logger.info("循环读取摄像头画面")
        start_time = time.time()
        # ç¡çœ 100毫秒
        time.sleep(config['cam']['sleep'])
        ret, frame = cap.read()
@@ -211,7 +210,7 @@

        # å®‰å…¨æ£€æµ‹
        boxes, scores, class_ids = safety_detect(frame)
        draw_img = safety_detect.draw_detections(frame, boxes, scores, class_ids)
        draw_img = safety_detect.draw_detections(frame, class_ids, scores,boxes )

        det_res = {}
        if class_ids is not None:
@@ -277,7 +276,7 @@
        # print(status)

        # ä¸Šæ–™æœºä½ç½®è¯†åˆ«
        probabilities2 = hoister_position(frame);
        probabilities2 = hoister_position(frame)
        predicted_class2 = np.argmax(probabilities2, axis=1)[0]
        max_probability2 = np.max(probabilities2, axis=1)[0]
        class_2 = hoister_position.class_names[predicted_class2]
@@ -291,10 +290,8 @@
            logger.info("发送上料机位置识别结果："+str(class_feeder))
            l.send_msg(class_feeder)
        # è®¡ç®—帧速率
        frame_count += 1
        end_time = time.time()
        elapsed_time = end_time - start_time
        fps = frame_count / elapsed_time
        fps = (1 / (end_time - start_time))
        # print(f"FPS: {fps:.2f}")
        # å°†FPS绘制在图像上
        cv2.putText(draw_img, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
@@ -303,6 +300,11 @@
        # èŽ·å–当前窗口大小
        width = cv2.getWindowImageRect("AICamera")[2]
        height = cv2.getWindowImageRect("AICamera")[3]
        # print("width", width, "height", height)

        # å¦‚æžœheight小于1则赋值100
        if height < 1:
            height = 100

        # è°ƒæ•´å›¾åƒå¤§å°ä»¥é€‚应窗口
        resized_frame = cv2.resize(draw_img, (width, height))
@@ -335,21 +337,6 @@
        ('cbData', ctypes.wintypes.DWORD),
        ('lpData', ctypes.c_char_p)
    ]



# logging.info("准备加载安全检测模型..")
# print("准备加载安全检测模型..")
# model_safe = SAFETY_DETECT(config['model']['safe'])
#
# logging.info("安全检测模型加载成功。")
# print("安全检测模型加载成功。")
# logging.info("准备加载药材识别模型..")
# print("准备加载药材识别模型..")
# model_cls = HERB_IDENTIFY(config['model']['cls'])
# logging.info("药材识别模型加载成功。")
# print("药材识别模型加载成功。")


class Listener:
    def __init__(self):
@@ -462,6 +449,8 @@
        return 0

if __name__ == '__main__':


    # ç´¯è®¡æ¯ç§è¯æä¸è®ºåæ¬¡å‡ºçŽ°çš„次数
    class_count = {}
    # ç´¯è®¡æ¯ç§è¯æç½®ä¿¡åº¦æœ€é«˜çš„次数
@@ -478,9 +467,10 @@
    is_loaded = False
    # åŠ è½½ONNX模型

    load_identifier = IDENTIFIER("model/loading.onnx")
    hoister_position = IDENTIFIER("model/hl.onnx")
    safety_detect = SAFETY_DETECT("model/safety_det.onnx")
    print("加载模型===============")
    load_identifier = IDENTIFIER("./model/load_id")
    hoister_position = IDENTIFIER("./model/feeder_id")
    safety_detect = SAFETY_DETECT("./model/safe_det")
    config = read_config()
    PCOPYDATASTRUCT = ctypes.POINTER(COPYDATASTRUCT)