1import cv2
2import numpy as np
3import pyrealsense2 as rs
4from collections import deque
5import time
6
7# ==========================================
8# CONFIGURATION
9# ==========================================
10
11USE_REALSENSE = True   # Set False for normal USB / Pi camera
12ENABLE_TFLITE = False  # Set True if using TensorFlow Lite (Jetson)
13
14# ==========================================
15# REALSENSE SETUP (Leo Rover Compatible)
16# ==========================================
17
18if USE_REALSENSE:
19    pipeline = rs.pipeline()
20    config = rs.config()
21    config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
22    config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
23    profile = pipeline.start(config)
24    align = rs.align(rs.stream.color)
25    depth_scale = profile.get_device().first_depth_sensor().get_depth_scale()
26else:
27    cap = cv2.VideoCapture(0)
28
29# ==========================================
30# HSV COLOR RANGES (Lighting Robust)
31# ==========================================
32
33COLOR_RANGES = {
34    "red": [(0, 120, 80, 10, 255, 255),
35            (170, 120, 80, 180, 255, 255)],
36    "green": [(35, 80, 80, 85, 255, 255)],
37    "blue": [(90, 80, 80, 130, 255, 255)],
38    "yellow": [(20, 80, 80, 35, 255, 255)],
39    "orange": [(10, 100, 100, 20, 255, 255)]
40}
41
42# ==========================================
43# DEPTH SMOOTHING
44# ==========================================
45
46depth_buffer = deque(maxlen=5)
47
48# ==========================================
49# SHAPE DETECTION
50# ==========================================
51
52def detect_shape(contour):
53    epsilon = 0.02 * cv2.arcLength(contour, True)
54    approx = cv2.approxPolyDP(contour, epsilon, True)
55    vertices = len(approx)
56
57    if vertices == 3:
58        return "triangle"
59    elif vertices == 4:
60        x, y, w, h = cv2.boundingRect(approx)
61        aspect = float(w) / (h + 1e-5)
62        if 0.9 < aspect < 1.1:
63            return "square"
64        return "rectangle"
65    elif vertices > 4:
66        return "circle"
67    return "unknown"
68
69# ==========================================
70# MAIN LOOP
71# ==========================================
72
73print("Leo Rover Dynamic Color Detection Running")
74
75try:
76    while True:
77
78        # ===============================
79        # Capture Frame
80        # ===============================
81        if USE_REALSENSE:
82            frames = pipeline.wait_for_frames()
83            aligned = align.process(frames)
84            depth_frame = aligned.get_depth_frame()
85            color_frame = aligned.get_color_frame()
86
87            if not depth_frame or not color_frame:
88                continue
89
90            frame = np.asanyarray(color_frame.get_data())
91            depth_image = np.asanyarray(depth_frame.get_data())
92        else:
93            ret, frame = cap.read()
94            if not ret:
95                break
96            depth_image = None
97
98        # ===============================
99        # Image Preprocessing
100        # ===============================
101        blurred = cv2.GaussianBlur(frame, (5,5), 0)
102        hsv = cv2.cvtColor(blurred, cv2.COLOR_BGR2HSV)
103
104        detected_objects = []
105
106        # ===============================
107        # Color Segmentation
108        # ===============================
109        for color_name, ranges in COLOR_RANGES.items():
110
111            mask_total = None
112
113            for (h1,s1,v1,h2,s2,v2) in ranges:
114                lower = np.array([h1,s1,v1])
115                upper = np.array([h2,s2,v2])
116                mask = cv2.inRange(hsv, lower, upper)
117                mask_total = mask if mask_total is None else mask_total + mask
118
119            # Noise Filtering
120            kernel = np.ones((5,5), np.uint8)
121            mask_total = cv2.morphologyEx(mask_total, cv2.MORPH_OPEN, kernel)
122            mask_total = cv2.morphologyEx(mask_total, cv2.MORPH_CLOSE, kernel)
123
124            contours, _ = cv2.findContours(mask_total,
125                                           cv2.RETR_EXTERNAL,
126                                           cv2.CHAIN_APPROX_SIMPLE)
127
128            for contour in contours:
129
130                area = cv2.contourArea(contour)
131                if area < 800:
132                    continue
133
134                shape = detect_shape(contour)
135                if shape == "unknown":
136                    continue
137
138                x, y, w, h = cv2.boundingRect(contour)
139
140                M = cv2.moments(contour)
141                if M["m00"] == 0:
142                    continue
143
144                cx = int(M["m10"] / M["m00"])
145                cy = int(M["m01"] / M["m00"])
146
147                depth_m = 0
148                if USE_REALSENSE:
149                    depth_m = depth_frame.get_distance(cx, cy)
150                    if depth_m <= 0:
151                        continue
152
153                detected_objects.append({
154                    "contour": contour,
155                    "color": color_name,
156                    "shape": shape,
157                    "depth": depth_m,
158                    "center": (cx, cy)
159                })
160
161        # ===============================
162        # Select Nearest Object
163        # ===============================
164        if detected_objects:
165
166            if USE_REALSENSE:
167                target = min(detected_objects, key=lambda o: o["depth"])
168                depth_buffer.append(target["depth"])
169                depth_smooth = np.mean(depth_buffer)
170            else:
171                target = detected_objects[0]
172                depth_smooth = 0
173
174            cx, cy = target["center"]
175
176            cv2.drawContours(frame, [target["contour"]], -1, (0,0,255), 3)
177            cv2.circle(frame, (cx,cy), 6, (0,255,255), -1)
178
179            label = f"TARGET: {target['color']} {target['shape']}"
180            depth_text = f"Depth: {depth_smooth:.2f}m"
181
182            cv2.putText(frame, label, (20,40),
183                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,255), 2)
184
185            if USE_REALSENSE:
186                cv2.putText(frame, depth_text, (20,70),
187                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,255), 2)
188
189        # ===============================
190        # Display
191        # ===============================
192        cv2.imshow("Leo Rover Color Detection", frame)
193
194        if cv2.waitKey(1) & 0xFF == ord('q'):
195            break
196
197finally:
198    if USE_REALSENSE:
199        pipeline.stop()
200    else:
201        cap.release()
202    cv2.destroyAllWindows()