ESP32-Robotic-Arm/OpenCV code red color tracking.py at main · silentneuron/ESP32-Robotic-Arm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import cv2
import numpy as np
import websocket
import json
import time

ESP32_IP = "192.168.4.1"  # just put ur IP address from the output of arduino


# Servo angle limits
BASE_MIN = 60
BASE_MAX = 120
FRAME_WIDTH = 640
FRAME_HEIGHT = 480

# Gripper positions
GRIPPER_OPEN = 90
GRIPPER_CLOSED = 30

# Convert cx to base angle
def map_base_angle(cx):
    return int(BASE_MIN + (cx / FRAME_WIDTH) * (BASE_MAX - BASE_MIN))

# Convert cy to shoulder and elbow angles
def map_joint_angles(cy):
    shoulder_angle = int(80 + (cy / FRAME_HEIGHT) * 40)  # 80 → 120
    elbow_angle = int(100 - (cy / FRAME_HEIGHT) * 40)    # 100 → 60

    # clamp
    shoulder_angle = max(60, min(120, shoulder_angle))
    elbow_angle = max(60, min(120, elbow_angle))

    return shoulder_angle, elbow_angle

# Send servo command to ESP32
def send_to_esp32(ws, base, shoulder, elbow, gripper):
    msg = {
        "Base": base,
        "Shoulder": shoulder,
        "Elbow": elbow,
        "Gripper": gripper
    }

    print("Sending:", msg)

    for key, value in msg.items():
        ws.send(f"{key},{value}")
        time.sleep(0.02)  # allows ESP32 to process

# HSV color ranges
blue_lower = np.array([100, 150, 70])
blue_upper = np.array([140, 255, 255])

red_lower1 = np.array([0, 120, 70])
red_upper1 = np.array([10, 255, 255])
red_lower2 = np.array([170, 120, 70])
red_upper2 = np.array([180, 255, 255])


def detect_color(frame):
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    blue_mask = cv2.inRange(hsv, blue_lower, blue_upper)

    red_mask1 = cv2.inRange(hsv, red_lower1, red_upper1)
    red_mask2 = cv2.inRange(hsv, red_lower2, red_upper2)
    red_mask = red_mask1 | red_mask2

    return blue_mask, red_mask


def get_largest_contour_center(mask):
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        largest = max(contours, key=cv2.contourArea)
        if cv2.contourArea(largest) > 1000:
            M = cv2.moments(largest)
            if M["m00"] != 0:
                cx = int(M["m10"] / M["m00"])
                cy = int(M["m01"] / M["m00"])
                return cx, cy
    return None


# Connect to ESP32 WebSocket
ws = websocket.WebSocket()
ws.connect(f"ws://{ESP32_IP}/RobotArmInput")

# Start camera
cap = cv2.VideoCapture(0)
cap.set(3, FRAME_WIDTH)
cap.set(4, FRAME_HEIGHT)

while True:
    ret, frame = cap.read()
    if not ret:
        break

    blue_mask, red_mask = detect_color(frame)
    blue_center = get_largest_contour_center(blue_mask)
    red_center = get_largest_contour_center(red_mask)

    if blue_center:
        cx, cy = blue_center
        cv2.circle(frame, (cx, cy), 10, (255, 0, 0), -1)

        base_angle = map_base_angle(cx)
        shoulder, elbow = map_joint_angles(cy)

        send_to_esp32(ws, base_angle, shoulder, elbow, GRIPPER_CLOSED)

    elif red_center:
        cx, cy = red_center
        cv2.circle(frame, (cx, cy), 10, (0, 0, 255), -1)

        base_angle = map_base_angle(cx)
        shoulder, elbow = map_joint_angles(cy)

        send_to_esp32(ws, base_angle, shoulder, elbow, GRIPPER_CLOSED)

    else:
        # idle
        send_to_esp32(ws, 90, 90, 90, GRIPPER_OPEN)

    cv2.imshow("Robot Eye", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
ws.close()