From fba9ef8743b9c91a02cb822f5d441583bc3deba6 Mon Sep 17 00:00:00 2001
From: bsw215583320 <baoshiwei121@163.com>
Date: 星期三, 16 四月 2025 18:51:18 +0800
Subject: [PATCH] 增加上料机位置识别

---
 herb_ai.py |  125 ++++++++++++++++++++++++++++++++---------
 1 files changed, 96 insertions(+), 29 deletions(-)

diff --git a/herb_ai.py b/herb_ai.py
index 5721084..66be1e8 100644
--- a/herb_ai.py
+++ b/herb_ai.py
@@ -10,7 +10,6 @@
 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
@@ -23,11 +22,34 @@
 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 = webcams.get(cam2)
+    camera2_index = config['cam']['cam2']
     print("绗簩涓憚鍍忓ご绱㈠紩锛�" + str(camera2_index))
     # 鎵撳紑鎽勫儚澶�
     capture = cv2.VideoCapture(camera2_index, cv2.CAP_DSHOW)
@@ -54,7 +76,7 @@
             break
         count += 1
 
-        if count == 2:
+        if count == config['cam']['frames']:
             herb_probabilities = herb_identifier(frame2)
             top_five_classes = np.argsort(herb_probabilities, axis=1)[0][-5:][::-1]
             name = ""
@@ -68,16 +90,17 @@
             # 璁$畻鎷夋櫘鎷夋柉鍝嶅簲鐨勬柟宸�
             laplacian = variance_of_laplacian(frame2)
             # 鐢熸垚淇濆瓨鏂囦欢鍚嶏紝浠ュ綋鍓嶆椂闂村懡鍚�
-            save_name2 = name +"_["+ str(round(laplacian, 2)) +"]_"+ time.strftime("%Y%m%d%H%M%S", time.localtime()) + ".jpg"
+            save_name2 = time.strftime("%Y%m%d%H%M%S", time.localtime()) + "_" +name +"_["+ str(round(laplacian, 2)) +"]"+  ".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 = {}
@@ -104,30 +127,42 @@
                 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
     # 缁撴潫绾跨▼
     capture.release()
-def send_result(class_count, class_count_max, class_sum):
-    global is_loaded
+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
-    l.send_msg("airecognize," + f"{class_count}")
+    count_msg = "airecognize," + f"{sorted_class_count}"
+    logger.info("鍙戦�佽嵂鏉愯瘑鍒粨鏋滐細"+str(count_msg))
+    l.send_msg(count_msg)
     pass
 
 
 def load_identify():
     global is_loaded
     # 鎽勫儚澶寸储寮曞彿锛岄�氬父涓�0琛ㄧず绗竴涓憚鍍忓ご
-    camera_index = webcams.get(cam1)
+    camera_index = config['cam']['cam1']
     print("绗竴涓憚鍍忓ご绱㈠紩锛�" + str(camera_index))
     # 鎵撳紑鎽勫儚澶�
     cap = cv2.VideoCapture(camera_index, cv2.CAP_DSHOW)
@@ -157,17 +192,14 @@
     # 涓婃枡鐘舵��
     status = "娌℃湁涓婃枡"
 
-    # 绱姣忕鑽潗涓嶈鍚嶆鍑虹幇鐨勬鏁�
-    class_count = {}
-    # 绱姣忕鑽潗缃俊搴︽渶楂樼殑娆℃暟
-    class_count_max = {}
-    # 绱姣忕鑽潗鐨勭疆淇″害鎬诲拰
-    class_sum = {}
+    # 鍒涘缓绐楀彛骞惰缃负鍙皟鏁村ぇ灏�
+    cv2.namedWindow("AICamera", cv2.WINDOW_NORMAL)
+
     # 寰幆璇诲彇鎽勫儚澶寸敾闈�
     while True:
         logger.info("寰幆璇诲彇鎽勫儚澶寸敾闈�")
         # 鐫$湢100姣
-        time.sleep(0.1)
+        time.sleep(config['cam']['sleep'])
         ret, frame = cap.read()
         if not ret:
             print("鏃犳硶璇诲彇鎽勫儚澶寸敾闈�")
@@ -175,10 +207,12 @@
             break
         # 鑾峰彇褰撳墠鏃堕棿
         current_time = time.time()
+        # 姣忛殧3绉掑彇涓�甯у浘鍍�
+
         # 瀹夊叏妫�娴�
         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 杞崲鎴愮被鍒悕绉�
@@ -190,11 +224,14 @@
                     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:
-            l.send_msg("aidetect," + f"{det_res}")
+            res_ = "aidetect," + f"{det_res}"
+            logger.info("鍙戦�佸畨鍏ㄦ娴嬬粨鏋滐細"+str(res_))
+            l.send_msg(res_)
+
         # 涓婃枡璇嗗埆
         probabilities = load_identifier(frame)
         # 鎵惧埌鏈�澶ф鐜囩殑绫诲埆
@@ -221,6 +258,8 @@
                 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()
 
@@ -231,12 +270,26 @@
             status = "娌℃湁涓婃枡"
             if class_ == "meishangliao" and count == 3 and is_loaded:
                 logger.info("鍋滄涓婃枡鍚庡彂閫佽瘑鍒粨鏋�")
-                send_result(class_count, class_count_max, class_sum)
+                send_result()
             if class_ == "meishangliao" and count == 1000:
                 is_loaded = False
                 logger.info("闀挎椂闂存湭涓婃枡锛岄噸缃鍦ㄤ笂鏂欑姸鎬�")
         # 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()
@@ -244,10 +297,17 @@
         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.putText(draw_img, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
                     cv2.LINE_AA)
         # 鏄剧ず鐢婚潰
-        cv2.imshow("Camera", draw_img)
+        # 鑾峰彇褰撳墠绐楀彛澶у皬
+        width = cv2.getWindowImageRect("AICamera")[2]
+        height = cv2.getWindowImageRect("AICamera")[3]
+
+        # 璋冩暣鍥惧儚澶у皬浠ラ�傚簲绐楀彛
+        resized_frame = cv2.resize(draw_img, (width, height))
+
+        cv2.imshow("AICamera", resized_frame)
         # 妫�娴嬫寜閿紝濡傛灉鎸変笅q閿垯閫�鍑哄惊鐜�
         if cv2.waitKey(1) & 0xFF == ord('q'):
             break
@@ -402,17 +462,24 @@
         return 0
 
 if __name__ == '__main__':
-    cam1 = "USB Camera"
-    cam2 = "PC Camera"
-    camUtil = CAM_UTIL(cam1, cam2)
-    webcams = camUtil.list_webcams()
-    print("鎽勫儚澶�", webcams)
+    # 绱姣忕鑽潗涓嶈鍚嶆鍑虹幇鐨勬鏁�
+    class_count = {}
+    # 绱姣忕鑽潗缃俊搴︽渶楂樼殑娆℃暟
+    class_count_max = {}
+    # 绱姣忕鑽潗鐨勭疆淇″害鎬诲拰
+    class_sum = {}
+    # cam1 = "USB Camera"
+    # cam2 = "USB ZOOM Camera"
+    # camUtil = CAM_UTIL(cam1, cam2)
+    # webcams = camUtil.webcam_list
+    # print("鎽勫儚澶�", webcams)
     save_path = "data/images/"
     # 鏄惁涓婅繃鏂�
     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