From 88fc0f9f9b7fd3eb81c958ca41ed822cf3657c47 Mon Sep 17 00:00:00 2001
From: baoshiwei <baoshiwei@shlanbao.cn>
Date: 星期二, 22 四月 2025 15:50:22 +0800
Subject: [PATCH] refactor: 重构中药识别项目 分为onnx版和openvino版
---
onnx/herb_ai_.py | 257 ++++++++++++++++++++------------------------------
1 files changed, 103 insertions(+), 154 deletions(-)
diff --git a/herb_ai.py b/onnx/herb_ai_.py
similarity index 63%
copy from herb_ai.py
copy to onnx/herb_ai_.py
index 66be1e8..01006a6 100644
--- a/herb_ai.py
+++ b/onnx/herb_ai_.py
@@ -10,6 +10,7 @@
import win32gui
import multiprocessing
from safety_detect import SAFETY_DETECT
+from cam_util import CAM_UTIL
from identifier import IDENTIFIER
import os
from logger_config import logger
@@ -22,34 +23,11 @@
def variance_of_laplacian(image):
# 璁＄畻杈撳叆鍥惧儚鐨勬媺鏅媺鏂搷搴旂殑鏂瑰樊
return cv2.Laplacian(image, cv2.CV_64F).var()
-
-
-def clean_directory(path, days_threshold, max_files):
- """娓呯悊瓒呰繃鏃堕棿鎴栨暟閲忕殑鏂囦欢"""
- now = time.time()
- threshold = now - days_threshold * 24 * 3600
- files = []
- for f in os.listdir(path):
- file_path = os.path.join(path, f)
- if os.path.isfile(file_path):
- files.append((file_path, os.path.getmtime(file_path)))
- # 鎸変慨鏀规椂闂撮檷搴忔帓搴忥紝淇濈暀鏈�鏂版枃浠�
- files.sort(key=lambda x: x[1], reverse=True)
- # 鍒犻櫎杩囨湡鏂囦欢
- for file_info in files:
- if file_info[1] < threshold:
- os.remove(file_info[0])
- # 淇濈暀鏈�鏂版枃浠讹紝鍒犻櫎澶氫綑鏂囦欢
- for file_info in files[max_files:]:
- os.remove(file_info[0])
-
-
# 璋冪敤鍙︿竴涓暱鐒﹂暅澶达紝鎷嶆憚娓呮櫚鐨勫眬閮ㄨ嵂鏉愬浘鐗�
def get_image():
- herb_identifier = IDENTIFIER("model/herb_identify.onnx")
logger.info("璇嗗埆绾跨▼鍚姩")
global is_loaded, class_count, class_count_max, class_sum
- camera2_index = config['cam']['cam2']
+ camera2_index = 1
print("绗簩涓憚鍍忓ご绱㈠紩锛�" + str(camera2_index))
# 鎵撳紑鎽勫儚澶�
capture = cv2.VideoCapture(camera2_index, cv2.CAP_DSHOW)
@@ -76,7 +54,7 @@
break
count += 1
- if count == config['cam']['frames']:
+ if count == 10:
herb_probabilities = herb_identifier(frame2)
top_five_classes = np.argsort(herb_probabilities, axis=1)[0][-5:][::-1]
name = ""
@@ -90,17 +68,16 @@
# 璁＄畻鎷夋櫘鎷夋柉鍝嶅簲鐨勬柟宸�
laplacian = variance_of_laplacian(frame2)
# 鐢熸垚淇濆瓨鏂囦欢鍚嶏紝浠ュ綋鍓嶆椂闂村懡鍚�
- save_name2 = time.strftime("%Y%m%d%H%M%S", time.localtime()) + "_" +name +"_["+ str(round(laplacian, 2)) +"]"+ ".jpg"
+ save_name2 = name +"_["+ str(round(laplacian, 2)) +"]_"+ time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
logger.info(f"璇嗗埆缁撴灉杞崲涓轰繚瀛樺浘鐗囧悕绉�:, {save_name2}")
# 鍒ゆ柇鍥惧儚鐨勬竻鏅板害
# 淇濆瓨璋冩暣灏哄鍚庣殑鍥剧墖
if laplacian > 200:
c_ = save_path + "2/c/"
+ # 鍒ゆ柇鏂囦欢鏄惁瀛樺湪锛屼笉瀛樺湪鍒欏垱寤�
if not os.path.exists(c_):
os.makedirs(c_)
cv2.imwrite(c_ + save_name2, frame2)
- # 鏂板娓呯悊璋冪敤
- clean_directory(c_, config['cam']['days_threshold'], config['cam']['max_files'])
# 娓呮櫚鏉′欢涓嬬疮璁¤瘑鍒粨鏋滀腑鑽潗鍚嶇О鍑虹幇鐨勬鏁�
# 绱姣忕鑽潗涓嶈鍚嶆鍑虹幇鐨勬鏁�,绱姣忕鑽潗缃俊搴︽渶楂樼殑娆℃暟,绱姣忕鑽潗鐨勭疆淇″害鎬诲拰
# class_count = {}
@@ -127,11 +104,10 @@
is_loaded = True
else:
n_ = save_path + "2/n/"
+ # 鍒ゆ柇鏂囦欢鏄惁瀛樺湪锛屼笉瀛樺湪鍒欏垱寤�
if not os.path.exists(n_):
os.makedirs(n_)
cv2.imwrite(n_ + save_name2, frame2)
- # 鏂板娓呯悊璋冪敤
- clean_directory(n_, config['cam']['days_threshold'], config['cam']['max_files'])
# cv2.imshow("Camera", resized_frame2)
print("淇濆瓨鍥剧墖:", save_name2)
break
@@ -139,30 +115,19 @@
capture.release()
def send_result():
global is_loaded,class_count, class_count_max, class_sum
- # 瀵筩lass_count杩涜鎺掑簭锛屾寜鐓у�间粠澶у埌灏忔帓搴�,杩斿洖鍊兼渶澶х殑鍓嶄簲涓�
- sorted_class_count = dict(sorted(class_count.items(), key=lambda x: x[1], reverse=True)[:5])
- # 瀵筩lass_count_max杩涜鎺掑簭锛屾寜鐓у�间粠澶у埌灏忔帓搴�,杩斿洖鍊兼渶澶х殑鍓嶄簲涓�
- sorted_class_count_max = dict(sorted(class_count_max.items(), key=lambda x: x[1], reverse=True)[:5])
- # 瀵� class_sum杩涜鎺掑簭锛屾寜鐓у�间粠澶у埌灏忔帓搴�,杩斿洖鍊兼渶澶х殑鍓嶄簲涓�
- sorted_class_sum = dict(sorted(class_sum.items(), key=lambda x: x[1], reverse=True)[:5])
# 灏嗕笁绉嶇粺璁＄粨鏋滆緭鍑哄埌鏃ュ織涓�
logger.info("class_count:"+str(class_count))
- logger.info("sorted_class_count:"+str(sorted_class_count))
logger.info("class_count_max:"+str(class_count_max))
- logger.info("sorted_class_count_max:"+str(sorted_class_count_max))
logger.info("class_sum:"+str(class_sum))
- logger.info("sorted_class_sum:"+str(sorted_class_sum))
is_loaded = False
- count_msg = "airecognize," + f"{sorted_class_count}"
- logger.info("鍙戦�佽嵂鏉愯瘑鍒粨鏋滐細"+str(count_msg))
- l.send_msg(count_msg)
+ l.send_msg("airecognize," + f"{class_count}")
pass
def load_identify():
- global is_loaded
+ global is_loaded, frame
# 鎽勫儚澶寸储寮曞彿锛岄�氬父涓�0琛ㄧず绗竴涓憚鍍忓ご
- camera_index = config['cam']['cam1']
+ camera_index = 0
print("绗竴涓憚鍍忓ご绱㈠紩锛�" + str(camera_index))
# 鎵撳紑鎽勫儚澶�
cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
@@ -182,24 +147,21 @@
# 璁℃椂鍣�
frame_count = 0
start_time = time.time()
- stime = time.time()
+ sstime = time.time()
+
if not os.path.exists(save_path):
os.makedirs(save_path)
- # 涓婃璇嗗埆缁撴灉
- class_old = "1"
- # 绱娆℃暟
- count = 0
+
+
# 涓婃枡鐘舵��
status = "娌℃湁涓婃枡"
- # 鍒涘缓绐楀彛骞惰缃负鍙皟鏁村ぇ灏�
- cv2.namedWindow("AICamera", cv2.WINDOW_NORMAL)
# 寰幆璇诲彇鎽勫儚澶寸敾闈�
while True:
- logger.info("寰幆璇诲彇鎽勫儚澶寸敾闈�")
+ # logger.info("寰幆璇诲彇鎽勫儚澶寸敾闈�")
# 鐫＄湢100姣
- time.sleep(config['cam']['sleep'])
+ time.sleep(0.1)
ret, frame = cap.read()
if not ret:
print("鏃犳硶璇诲彇鎽勫儚澶寸敾闈�")
@@ -207,107 +169,26 @@
break
# 鑾峰彇褰撳墠鏃堕棿
current_time = time.time()
- # 姣忛殧3绉掑彇涓�甯у浘鍍�
+ # 姣忛殧n绉掑彇涓�甯у浘鍍�
+ if current_time - sstime >= config['cam']['sleep']:
+ sstime = current_time
- # 瀹夊叏妫�娴�
- boxes, scores, class_ids = safety_detect(frame)
- draw_img = safety_detect.draw_detections(frame, boxes, scores, class_ids)
-
- det_res = {}
- if class_ids is not None:
- # 閬嶅巻class_ids 杞崲鎴愮被鍒悕绉�
- for i in range(len(class_ids)):
- class_id = class_ids[i]
- class_name = safety_detect.class_names[class_id]
- # 瀛樺叆鍒癲et_res涓�
- if class_name in det_res:
- det_res[class_name] = det_res[class_name] if det_res[class_name] > scores[i] else scores[i]
- else:
- det_res[class_name] = scores[i]
- print(det_res)
- logger.info(f"瀹夊叏妫�娴嬭瘑鍒粨鏋�, {det_res}")
- # 濡傛灉cass_ids涓寘鍚�0锛屽垯琛ㄧず鏈夊畨鍏ㄦ娴嬪埌浜轰綋
- if 0 in class_ids:
- res_ = "aidetect," + f"{det_res}"
- logger.info("鍙戦�佸畨鍏ㄦ娴嬬粨鏋滐細"+str(res_))
- l.send_msg(res_)
-
- # 涓婃枡璇嗗埆
- probabilities = load_identifier(frame)
- # 鎵惧埌鏈�澶ф鐜囩殑绫诲埆
- predicted_class = np.argmax(probabilities, axis=1)[0]
- max_probability = np.max(probabilities, axis=1)[0]
- class_ = load_identifier.class_names[predicted_class]
- # 璁＄畻绫诲瀷閲嶅鐨勬鏁帮紝绫诲埆鏇存崲涔嬪悗閲嶆柊璁℃暟
- if class_ != class_old:
- count = 0
- else:
- count += 1
- class_old = class_
- print(f"{class_}:{count}: {max_probability}")
- logger.info(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, frame)
- # 鏂板娓呯悊璋冪敤
- clean_directory(path_, config['cam']['days_threshold'], config['cam']['max_files'])
- # 閲嶇疆璁℃椂鍣�
- stime = time.time()
-
- thread = threading.Thread(target=get_image)
- thread.start()
-
- else:
- status = "娌℃湁涓婃枡"
- if class_ == "meishangliao" and count == 3 and is_loaded:
- logger.info("鍋滄涓婃枡鍚庡彂閫佽瘑鍒粨鏋�")
- send_result()
- if class_ == "meishangliao" and count == 1000:
- is_loaded = False
- logger.info("闀挎椂闂存湭涓婃枡锛岄噸缃鍦ㄤ笂鏂欑姸鎬�")
+ thread1 = threading.Thread(target=method_name)
+ thread1.start()
+ # frame = draw_img
# print(status)
- # 涓婃枡鏈轰綅缃瘑鍒�
- 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]
- print(f"-----------{class_2}:{predicted_class2}: {max_probability2}")
- logger.info(f"-----------{class_2}:{predicted_class2}: {max_probability2}")
-
- if predicted_class2 == 0:
- feeder_res = {class_2: max_probability2}
- class_feeder = "aifeeder," + f"{feeder_res}"
- print("send_msg", class_feeder)
- 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
- # print(f"FPS: {fps:.2f}")
- # 灏咶PS缁樺埗鍦ㄥ浘鍍忎笂
- cv2.putText(draw_img, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
- cv2.LINE_AA)
+ # frame_count += 1
+ # end_time = time.time()
+ # elapsed_time = end_time - start_time
+ # fps = frame_count / elapsed_time
+ # # print(f"FPS: {fps:.2f}")
+ # # 灏咶PS缁樺埗鍦ㄥ浘鍍忎笂
+ # cv2.putText(frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
+ # cv2.LINE_AA)
# 鏄剧ず鐢婚潰
- # 鑾峰彇褰撳墠绐楀彛澶у皬
- width = cv2.getWindowImageRect("AICamera")[2]
- height = cv2.getWindowImageRect("AICamera")[3]
-
- # 璋冩暣鍥惧儚澶у皬浠ラ�傚簲绐楀彛
- resized_frame = cv2.resize(draw_img, (width, height))
-
- cv2.imshow("AICamera", resized_frame)
+ cv2.imshow("AICamera", frame)
# 妫�娴嬫寜閿紝濡傛灉鎸変笅q閿垯閫�鍑哄惊鐜�
if cv2.waitKey(1) & 0xFF == ord('q'):
break
@@ -316,6 +197,69 @@
# 鍏抽棴鎵�鏈夌獥鍙�
cv2.destroyAllWindows()
+
+def method_name( ):
+ global is_loaded,count,class_old,stime,frame
+ # 瀹夊叏妫�娴�
+ boxes, scores, class_ids = safety_detect(frame)
+ draw_img = safety_detect.draw_detections(frame, boxes, scores, class_ids)
+ print(boxes, scores, class_ids)
+ det_res = {}
+ if class_ids is not None:
+ # 閬嶅巻class_ids 杞崲鎴愮被鍒悕绉�
+ for i in range(len(class_ids)):
+ class_id = class_ids[i]
+ class_name = safety_detect.class_names[class_id]
+ # 瀛樺叆鍒癲et_res涓�
+ if class_name in det_res:
+ det_res[class_name] = det_res[class_name] if det_res[class_name] > scores[i] else scores[i]
+ else:
+ det_res[class_name] = scores[i]
+ logger.info(f"瀹夊叏妫�娴嬭瘑鍒粨鏋�, {det_res}")
+ # 濡傛灉cass_ids涓寘鍚�0锛屽垯琛ㄧず鏈夊畨鍏ㄦ娴嬪埌浜轰綋
+ if 0 in class_ids:
+ l.send_msg("aidetect," + f"{det_res}")
+ # 涓婃枡璇嗗埆
+ probabilities = load_identifier(frame)
+ # 鎵惧埌鏈�澶ф鐜囩殑绫诲埆
+ predicted_class = np.argmax(probabilities, axis=1)[0]
+ max_probability = np.max(probabilities, axis=1)[0]
+ class_ = load_identifier.class_names[predicted_class]
+ # 璁＄畻绫诲瀷閲嶅鐨勬鏁帮紝绫诲埆鏇存崲涔嬪悗閲嶆柊璁℃暟
+ if class_ != class_old:
+ count = 0
+ else:
+ count += 1
+ class_old = class_
+ print(f"{class_}:{count}: {max_probability}")
+ logger.info(f"{class_}:{count}: {max_probability}")
+ # 鍒ゆ柇鏄惁涓婃枡骞朵笖涓婃枡娆℃暟澶т簬10娆�
+ if class_ == "shangliao" and count > 10:
+ status = "姝ｅ湪涓婃枡"
+ # 姣忛殧3绉掑彇涓�甯у浘鍍�
+ # 濡傛灉璺濈涓婁竴娆′繚瀛樺凡缁忚繃鍘�1绉掞紝鍒欎繚瀛樺綋鍓嶇敾闈�
+ current_time = time.time()
+ 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, frame)
+ # 閲嶇疆璁℃椂鍣�
+ stime = time.time()
+
+ thread = threading.Thread(target=get_image)
+ thread.start()
+
+ else:
+ status = "娌℃湁涓婃枡"
+ if class_ == "meishangliao" and count == 3 and is_loaded:
+ logger.info("鍋滄涓婃枡鍚庡彂閫佽瘑鍒粨鏋�")
+ send_result()
+ if class_ == "meishangliao" and count == 1000:
+ is_loaded = False
+ logger.info("闀挎椂闂存湭涓婃枡锛岄噸缃鍦ㄤ笂鏂欑姸鎬�")
# 璇诲彇閰嶇疆鏂囦欢
@@ -386,7 +330,7 @@
def OnCopyData(self, hwnd, msg, wparam, lparam):
try:
# 璁板綍寮�濮嬫椂闂�
- start_time = time.time()
+ startTime = time.time()
pCDS = ctypes.cast(lparam, PCOPYDATASTRUCT)
s = ctypes.string_at(pCDS.contents.lpData).decode()
strArr = s.split(",")
@@ -426,9 +370,9 @@
logger.info(f"璇嗗埆缁撴灉锛歿res}")
self.send_msg(msg)
# 璁板綍缁撴潫鏃堕棿
- end_time = time.time()
+ endTime = time.time()
# 璁＄畻鎵ц鏃堕棿
- execution_time = end_time - start_time
+ execution_time = endTime - startTime
# 鎵撳嵃鎵ц鏃堕棿
print(f"绋嬪簭鎵ц鏃堕棿涓�:{execution_time}绉�")
logger.info(f"绋嬪簭鎵ц鏃堕棿涓�:{execution_time,}绉�")
@@ -468,6 +412,12 @@
class_count_max = {}
# 绱姣忕鑽潗鐨勭疆淇″害鎬诲拰
class_sum = {}
+ # 涓婃璇嗗埆缁撴灉
+ class_old = "1"
+ # 绱娆℃暟
+ count = 0
+ stime = time.time()
+ frame = None
# cam1 = "USB Camera"
# cam2 = "USB ZOOM Camera"
# camUtil = CAM_UTIL(cam1, cam2)
@@ -477,9 +427,8 @@
# 鏄惁涓婅繃鏂�
is_loaded = False
# 鍔犺浇ONNX妯″瀷
-
+ herb_identifier = IDENTIFIER("model/herb_identify.onnx")
load_identifier = IDENTIFIER("model/loading.onnx")
- hoister_position = IDENTIFIER("model/hl.onnx")
safety_detect = SAFETY_DETECT("model/safety_det.onnx")
config = read_config()
PCOPYDATASTRUCT = ctypes.POINTER(COPYDATASTRUCT)
--
Gitblit v1.9.3