Browse Source

update spray algorithm

haminhtien99 5 months ago
parent
commit
1d27b00afc
3 changed files with 159 additions and 0 deletions
  1. 15 0
      README.md
  2. 11 0
      default_config.py
  3. 133 0
      distance_estimation/spray_pesticide.py

+ 15 - 0
README.md

@@ -18,6 +18,16 @@ Di chuyển tới file [`default_config.py`](default_config.py), thay đổi cá
 MARKER_SIZE=5 # kích thước các mã QR-code
 set_resolution='original' # điều chỉnh độ phân giải cho camera, đặt original nếu ko thay đổi
 SQUARE_SIZE = 2.1  # kích thước ô của bàn cờ
+
+# Thiết lập các giá trị cho việc phun thuốc
+ALLOWED_ANGLE_DEG = 20  # Góc lệch cho phép giữa aruco và camera
+ALLOWED_XYZ_RANGE = {   # Phạm vị lệch cho phép x, y, z của aruco tag trong hệ tọa độ gắn với camera
+    "x": (-10, 10),
+    "y": (-10, 10),
+    "z": (20, 50),
+}
+SPRAY_DURATION = 5      # Thời gian phun thuốc
+
 ```
 
 
@@ -79,3 +89,8 @@ Dữ liệu hiệu chuẩn camera lưu tại tệp .npz. Dữ liệu này bao g
 
 `python distance_estimation/camera_position_and_log.py`
 Tùy chọn ước lượng vị trí camera cùng với lưu lịch sử
+
+### Tùy chọn phun thuốc trừ sâu
+`python distance_estimation/spray_pesticide.py`
+- Thiết lập thuộc tính phù hợp với phạm vi lệch về góc, tọa độ xyz của marker so với camera, đảm bảo marker ở trung tâm, thời gian phun
+- In ra trạng thái mỗi marker

+ 11 - 0
default_config.py

@@ -35,3 +35,14 @@ default_positions = np.array([
     [-26.0, 15.5, 0.],   # marker 3
     [23, 13, 0.]   # marker 4
 ])
+
+
+# Thiết lập các giá trị cho việc phun thuốc
+
+ALLOWED_ANGLE_DEG = 20  # Góc lệch cho phép giữa aruco và camera
+ALLOWED_XYZ_RANGE = {   # Phạm vị lệch cho phép x, y, z của aruco tag trong hệ tọa độ gắn với camera
+    "x": (-10, 10),
+    "y": (-10, 10),
+    "z": (20, 50),
+}
+SPRAY_DURATION = 5      # Thời gian phun thuốc

+ 133 - 0
distance_estimation/spray_pesticide.py

@@ -0,0 +1,133 @@
+import cv2
+from cv2 import aruco
+import numpy as np
+import time
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from default_config import *
+visualization = True
+def is_in_center(x_px, y_px, frame_w, frame_h, margin_ratio=0.25):
+    margin_x = frame_w * margin_ratio
+    margin_y = frame_h * margin_ratio
+    return (margin_x < x_px < frame_w - margin_x) and (margin_y < y_px < frame_h - margin_y)
+
+def get_euler_angles_from_rvec(rvec):
+    rot_matrix, _ = cv2.Rodrigues(rvec)
+    sy = np.sqrt(rot_matrix[0, 0]**2 + rot_matrix[1, 0]**2)
+    singular = sy < 1e-6
+    if not singular:
+        roll = np.arctan2(rot_matrix[2, 1], rot_matrix[2, 2])
+        pitch = np.arctan2(-rot_matrix[2, 0], sy)
+        yaw = np.arctan2(rot_matrix[1, 0], rot_matrix[0, 0])
+    else:
+        roll = np.arctan2(-rot_matrix[1, 2], rot_matrix[1, 1])
+        pitch = np.arctan2(-rot_matrix[2, 0], sy)
+        yaw = 0
+    return np.degrees(roll), np.degrees(pitch), np.degrees(yaw)
+
+def is_roll_valid(roll):
+    return (abs(roll) < ALLOWED_ANGLE_DEG) or (160 <= abs(roll) <= 180)
+
+def spray_pesticide(marker_id):
+    print(f"✅ SPRAYING! Marker ID: {marker_id}")
+
+calib_data_path = f"camera_calibration/calib_data/{set_resolution}/MultiMatrix.npz"
+calib_data = np.load(calib_data_path)
+print('resolution:', set_resolution)
+
+cam_mat = calib_data["camMatrix"]
+dist_coef = calib_data["distCoef"]
+
+
+marker_dict = aruco.getPredefinedDictionary(aruco.DICT_7X7_50)
+param_markers = aruco.DetectorParameters()
+
+cap = cv2.VideoCapture(0)
+
+# Thay đổi độ phân giải phù hợp với yêu cầu ban đầu
+new_height = RESOLUTION_STANDARDS[set_resolution]
+new_width, new_height = change_resolution(cap,new_height)    
+print(f'Resolution camera {new_width}x{new_height}')
+
+fps = cap.get(cv2.CAP_PROP_FPS)
+print(f"Frames per second: {fps}")
+
+
+
+def start_spraying(marker_id):
+    print(f"🟡 SPRAYING marker {marker_id}...")
+
+def stop_spraying(marker_id):
+    print(f"✅ SPRAYED marker {marker_id}")
+
+# --- Track state ---
+marker_states = {}  # marker_id -> {"state": "NOT_SPRAYED"/"SPRAYING"/"SPRAYED", "start_time": float}
+
+while True:
+    ret, frame = cap.read()
+    if not ret:
+        break
+
+    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    detector = cv2.aruco.ArucoDetector(marker_dict, param_markers)
+    marker_corners, marker_IDs, _ = detector.detectMarkers(gray_frame)
+
+    current_time = time.time()
+
+    if marker_corners:
+        rVec, tVec, _ = aruco.estimatePoseSingleMarkers(marker_corners, MARKER_SIZE, cam_mat, dist_coef)
+        for i, (ids, corners) in enumerate(zip(marker_IDs, marker_corners)):
+            marker_id = int(ids[0])
+            corners = corners.reshape(4, 2).astype(int)
+            cx = int(np.mean(corners[:, 0]))
+            cy = int(np.mean(corners[:, 1]))
+            x, y, z = tVec[i][0]
+            roll, pitch, yaw = get_euler_angles_from_rvec(rVec[i])
+
+            # Kiểm tra đủ điều kiện
+            in_center = is_in_center(cx, cy, new_width, new_height)
+            within_angle = is_roll_valid(roll) and abs(pitch) < ALLOWED_ANGLE_DEG and abs(yaw) < ALLOWED_ANGLE_DEG
+            within_xyz = (
+                ALLOWED_XYZ_RANGE["x"][0] <= x <= ALLOWED_XYZ_RANGE["x"][1] and
+                ALLOWED_XYZ_RANGE["y"][0] <= y <= ALLOWED_XYZ_RANGE["y"][1] and
+                ALLOWED_XYZ_RANGE["z"][0] <= z <= ALLOWED_XYZ_RANGE["z"][1]
+            )
+
+            if marker_id not in marker_states:
+                marker_states[marker_id] = {"state": "NOT_SPRAYED", "start_time": 0}
+
+            state_info = marker_states[marker_id]
+
+            if in_center and within_angle and within_xyz:
+                if state_info["state"] == "NOT_SPRAYED":
+                    start_spraying(marker_id)
+                    marker_states[marker_id] = {
+                        "state": "SPRAYING",
+                        "start_time": current_time
+                    }
+                elif state_info["state"] == "SPRAYING":
+                    elapsed = current_time - state_info["start_time"]
+                    if elapsed >= SPRAY_DURATION:
+                        stop_spraying(marker_id)
+                        marker_states[marker_id]["state"] = "SPRAYED"
+
+            # --- Visualization ---
+            status = marker_states[marker_id]["state"]
+            color = {"NOT_SPRAYED": (0, 0, 255), "SPRAYING": (0, 255, 255), "SPRAYED": (0, 255, 0)}[status]
+            
+            if visualization:
+                cv2.polylines(frame, [corners], True, color, 2)
+                cv2.drawFrameAxes(frame, cam_mat, dist_coef, rVec[i], tVec[i], 4, 4)
+                cv2.putText(frame, f"ID: {marker_id}", (cx + 10, cy), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+                cv2.putText(frame, f"{status}", (cx + 10, cy + 25), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+                cv2.putText(frame, f"x:{x:.1f} y:{y:.1f} z:{z:.1f}", (cx + 10, cy + 45), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)
+                cv2.putText(frame, f"roll:{roll:.1f} pitch:{pitch:.1f} yaw:{yaw:.1f}", (cx + 10, cy + 65), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 255), 1)
+    if visualization:
+        cv2.imshow("frame", frame)
+        key = cv2.waitKey(1)
+        if key == ord("q"):
+            break
+
+cap.release()
+cv2.destroyAllWindows()