SirBlob/simulation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simulation Exit Fixes

Browse Source
This commit is contained in:
SirBlobby
2026-02-13 17:10:41 -05:00
parent 42e4fa28a9
commit 50ef3f0490
9 changed files with 97 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files

111
scripts/generate_tags.py
View File

@@ -1,23 +1,9 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""
Generate ArUco marker tags for use as Gazebo textures.
Usage:
python3 scripts/generate_tags.py # Generate default set
python3 scripts/generate_tags.py --ids 0 5 10 # Generate specific IDs
python3 scripts/generate_tags.py --dict DICT_6X6_250 --ids 0 1 2
python3 scripts/generate_tags.py --all # Generate all 50 tags in DICT_4X4_50
Output:
worlds/tags/aruco_<DICT>_<ID>.png (individual markers)
worlds/tags/land_tag.png (copy of the landing target, ID from config)
"""
import argparse import argparse
import sys import sys
import os import os
from pathlib import Path from pathlib import Path
import cv2 import cv2
import numpy as np import numpy as np
import yaml import yaml
@@ -66,7 +52,6 @@ DICT_SIZES = {
def generate_marker(dict_name, marker_id, pixel_size=600, border_bits=1): def generate_marker(dict_name, marker_id, pixel_size=600, border_bits=1):
"""Generate a single ArUco marker image with a white border."""
dict_id = ARUCO_DICTS[dict_name] dict_id = ARUCO_DICTS[dict_name]
aruco_dict = cv2.aruco.getPredefinedDictionary(dict_id) aruco_dict = cv2.aruco.getPredefinedDictionary(dict_id)
@@ -76,17 +61,12 @@ def generate_marker(dict_name, marker_id, pixel_size=600, border_bits=1):
marker_img = cv2.aruco.drawMarker(aruco_dict, marker_id, pixel_size) marker_img = cv2.aruco.drawMarker(aruco_dict, marker_id, pixel_size)
border_px = pixel_size // 6 border_px = pixel_size // 6
padded = np.ones( padded = np.ones((pixel_size + 2 * border_px, pixel_size + 2 * border_px), dtype=np.uint8) * 255
(pixel_size + 2 * border_px, pixel_size + 2 * border_px),
dtype=np.uint8,
) * 255
padded[border_px:border_px + pixel_size, border_px:border_px + pixel_size] = marker_img padded[border_px:border_px + pixel_size, border_px:border_px + pixel_size] = marker_img
return padded return padded
def get_landing_id(): def get_landing_id():
"""Read the landing target ID from search.yaml config."""
search_cfg = CONFIG_DIR / "search.yaml" search_cfg = CONFIG_DIR / "search.yaml"
if search_cfg.exists(): if search_cfg.exists():
with open(search_cfg) as f: with open(search_cfg) as f:
@@ -98,72 +78,42 @@ def get_landing_id():
def main(): def main():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(description="Generate ArUco marker tags for Gazebo simulation")
description="Generate ArUco marker tags for Gazebo simulation" parser.add_argument("--dict", default="DICT_4X4_50", choices=list(ARUCO_DICTS.keys()))
) parser.add_argument("--ids", type=int, nargs="+", default=None)
parser.add_argument( parser.add_argument("--all", action="store_true")
"--dict", default="DICT_4X4_50", parser.add_argument("--size", type=int, default=600)
choices=list(ARUCO_DICTS.keys()), parser.add_argument("--format", default="png", choices=["png", "jpg"])
help="ArUco dictionary (default: DICT_4X4_50)",
)
parser.add_argument(
"--ids", type=int, nargs="+", default=None,
help="Marker IDs to generate (default: landing target from config)",
)
parser.add_argument(
"--all", action="store_true",
help="Generate all markers in the dictionary",
)
parser.add_argument(
"--size", type=int, default=600,
help="Marker image size in pixels (default: 600)",
)
parser.add_argument(
"--format", default="png", choices=["png", "jpg"],
help="Output image format (default: png)",
)
args = parser.parse_args() args = parser.parse_args()
TAGS_DIR.mkdir(parents=True, exist_ok=True) TAGS_DIR.mkdir(parents=True, exist_ok=True)
dict_name = args.dict
landing_id = get_landing_id() landing_id = get_landing_id()
if args.all: if args.all:
max_id = DICT_SIZES.get(dict_name, 50) ids = list(range(DICT_SIZES.get(args.dict, 50)))
ids = list(range(max_id))
elif args.ids is not None: elif args.ids is not None:
ids = args.ids ids = args.ids
else: else:
ids = [landing_id] ids = [landing_id]
print(f"Generating {len(ids)} ArUco marker(s)") print(f"Generating {len(ids)} ArUco marker(s) in {TAGS_DIR}/")
print(f" Dictionary: {dict_name}")
print(f" Size: {args.size}px")
print(f" Output: {TAGS_DIR}/")
print(f" Landing target ID: {landing_id}")
print()
for marker_id in ids: for marker_id in ids:
marker_img = generate_marker(dict_name, marker_id, args.size) marker_img = generate_marker(args.dict, marker_id, args.size)
filename = f"aruco_{args.dict}_{marker_id}.{args.format}"
filename = f"aruco_{dict_name}_{marker_id}.{args.format}" cv2.imwrite(str(TAGS_DIR / filename), marker_img)
filepath = TAGS_DIR / filename print(f" {filename}")
cv2.imwrite(str(filepath), marker_img)
print(f" {filename} ({marker_img.shape[0]}x{marker_img.shape[1]}px)")
if marker_id == landing_id: if marker_id == landing_id:
land_path = TAGS_DIR / f"land_tag.{args.format}" cv2.imwrite(str(TAGS_DIR / f"land_tag.{args.format}"), marker_img)
cv2.imwrite(str(land_path), marker_img)
print(f" land_tag.{args.format} (copy of ID {landing_id})")
# Verify generated tags are detectable
print("\nVerifying detection...") print("\nVerifying detection...")
aruco_dict = cv2.aruco.getPredefinedDictionary(ARUCO_DICTS[dict_name]) aruco_dict = cv2.aruco.getPredefinedDictionary(ARUCO_DICTS[args.dict])
try: try:
params = cv2.aruco.DetectorParameters() params = cv2.aruco.DetectorParameters()
except AttributeError: except AttributeError:
params = cv2.aruco.DetectorParameters_create() params = cv2.aruco.DetectorParameters_create()
params.minMarkerPerimeterRate = 0.01 params.minMarkerPerimeterRate = 0.01
try: try:
params.cornerRefinementMethod = cv2.aruco.CORNER_REFINE_SUBPIX params.cornerRefinementMethod = cv2.aruco.CORNER_REFINE_SUBPIX
@@ -177,34 +127,13 @@ def main():
_detect = lambda img: cv2.aruco.detectMarkers(img, aruco_dict, parameters=params) _detect = lambda img: cv2.aruco.detectMarkers(img, aruco_dict, parameters=params)
ok = 0 ok = 0
fail = 0
for marker_id in ids: for marker_id in ids:
filename = f"aruco_{dict_name}_{marker_id}.{args.format}" img = cv2.imread(str(TAGS_DIR / f"aruco_{args.dict}_{marker_id}.{args.format}"), cv2.IMREAD_GRAYSCALE)
filepath = TAGS_DIR / filename _, detected_ids, _ = _detect(img)
img = cv2.imread(str(filepath), cv2.IMREAD_GRAYSCALE)
corners, detected_ids, _ = _detect(img)
if detected_ids is not None and marker_id in detected_ids.flatten(): if detected_ids is not None and marker_id in detected_ids.flatten():
ok += 1 ok += 1
else:
print(f" WARNING: {filename} failed detection check!") print(f"Done. {ok}/{len(ids)} markers verified.")
fail += 1
# Also test at simulated camera distances
land_file = TAGS_DIR / f"aruco_{dict_name}_{landing_id}.{args.format}"
if land_file.exists():
land_img = cv2.imread(str(land_file), cv2.IMREAD_GRAYSCALE)
print(f"\nDistance detection test (ID {landing_id}):")
for sim_size in [120, 80, 60, 40]:
small = cv2.resize(land_img, (sim_size, sim_size), interpolation=cv2.INTER_AREA)
bg = np.full((480, 640), 180, dtype=np.uint8)
y, x = (480 - sim_size) // 2, (640 - sim_size) // 2
bg[y:y + sim_size, x:x + sim_size] = small
corners, detected_ids, _ = _detect(bg)
status = "OK" if detected_ids is not None else "FAIL"
print(f" {sim_size}x{sim_size}px in 640x480 frame: {status}")
print(f"\nDone. {ok}/{ok + fail} markers verified.")
if __name__ == "__main__": if __name__ == "__main__":

69
src/control/search.py
View File

@@ -19,7 +19,6 @@ class SearchMode(Enum):
class Search: class Search:
POSITION_TOLERANCE = 1.0 POSITION_TOLERANCE = 1.0
CHECK_INTERVAL = 0.5 CHECK_INTERVAL = 0.5
MAX_TRAVEL_TIME = 30.0 MAX_TRAVEL_TIME = 30.0
@@ -41,7 +40,6 @@ class Search:
self.actions = actions or {} self.actions = actions or {}
self.land_ids = set(self.actions.get("land", [])) self.land_ids = set(self.actions.get("land", []))
# Search pattern parameters
self.spiral_max_legs = 12 self.spiral_max_legs = 12
self.spiral_initial_leg = self.CAM_FOV_METERS self.spiral_initial_leg = self.CAM_FOV_METERS
self.spiral_leg_increment = self.CAM_FOV_METERS * 0.8 self.spiral_leg_increment = self.CAM_FOV_METERS * 0.8
@@ -55,30 +53,16 @@ class Search:
self.levy_field_size = 50.0 self.levy_field_size = 50.0
self.levy_angle_dist = uniform(loc=-180, scale=360) self.levy_angle_dist = uniform(loc=-180, scale=360)
# Pre-flight plan
self.waypoints = [] self.waypoints = []
self.current_wp = 0 self.current_wp = 0
self._on_waypoint_change = None # callback: fn(index) self._on_waypoint_change = None
def configure(self, **kwargs): def configure(self, **kwargs):
for key, value in kwargs.items(): for key, value in kwargs.items():
if hasattr(self, key): if hasattr(self, key):
setattr(self, key, value) setattr(self, key, value)
# ═══════════════════════════════════════════════════════════
# PRE-FLIGHT PLANNING
# ═══════════════════════════════════════════════════════════
def plan(self, start_pos=(0.0, 0.0), geofence_points=None): def plan(self, start_pos=(0.0, 0.0), geofence_points=None):
"""Pre-compute all waypoints for the search pattern.
Args:
start_pos: (x, y) starting position in local NED
geofence_points: list of (x, y) polygon vertices, or None
Returns:
List of (x, y) absolute waypoints
"""
if self.mode == SearchMode.SPIRAL: if self.mode == SearchMode.SPIRAL:
waypoints = self._plan_spiral(start_pos) waypoints = self._plan_spiral(start_pos)
elif self.mode == SearchMode.LAWNMOWER: elif self.mode == SearchMode.LAWNMOWER:
@@ -88,7 +72,6 @@ class Search:
else: else:
waypoints = [start_pos] waypoints = [start_pos]
# Clip to geofence
if geofence_points: if geofence_points:
waypoints = self._clip_to_geofence(waypoints, geofence_points) waypoints = self._clip_to_geofence(waypoints, geofence_points)
@@ -105,7 +88,6 @@ class Search:
return waypoints return waypoints
def _plan_spiral(self, start): def _plan_spiral(self, start):
"""Expanding square spiral from center outward."""
waypoints = [start] waypoints = [start]
x, y = start x, y = start
@@ -135,7 +117,6 @@ class Search:
return waypoints return waypoints
def _plan_lawnmower(self, start): def _plan_lawnmower(self, start):
"""Lawnmower (boustrophedon) pattern."""
waypoints = [start] waypoints = [start]
sx, sy = start sx, sy = start
lane_spacing = self.lawn_lane_spacing lane_spacing = self.lawn_lane_spacing
@@ -158,7 +139,6 @@ class Search:
return waypoints return waypoints
def _plan_levy(self, start): def _plan_levy(self, start):
"""Lévy walk — random but pre-computed."""
waypoints = [start] waypoints = [start]
x, y = start x, y = start
@@ -175,26 +155,18 @@ class Search:
return waypoints return waypoints
def _clip_to_geofence(self, waypoints, polygon): def _clip_to_geofence(self, waypoints, polygon):
"""Clip waypoints to stay within the geofence polygon.
If a waypoint is outside, it is projected to the nearest
point on the polygon boundary. Consecutive out-of-bounds
waypoints are collapsed.
"""
clipped = [] clipped = []
for wx, wy in waypoints: for wx, wy in waypoints:
if self._point_in_polygon(wx, wy, polygon): if self._point_in_polygon(wx, wy, polygon):
clipped.append((wx, wy)) clipped.append((wx, wy))
else: else:
nearest = self._nearest_point_on_polygon(wx, wy, polygon) nearest = self._nearest_point_on_polygon(wx, wy, polygon)
# Avoid duplicate consecutive points
if not clipped or distance_2d(clipped[-1], nearest) > 0.5: if not clipped or distance_2d(clipped[-1], nearest) > 0.5:
clipped.append(nearest) clipped.append(nearest)
return clipped return clipped
@staticmethod @staticmethod
def _point_in_polygon(px, py, polygon): def _point_in_polygon(px, py, polygon):
"""Ray-casting point-in-polygon test."""
n = len(polygon) n = len(polygon)
inside = False inside = False
j = n - 1 j = n - 1
@@ -209,14 +181,12 @@ class Search:
@staticmethod @staticmethod
def _nearest_point_on_polygon(px, py, polygon): def _nearest_point_on_polygon(px, py, polygon):
"""Find the nearest point on the polygon edge to (px, py)."""
min_dist = float('inf') min_dist = float('inf')
nearest = (px, py) nearest = (px, py)
n = len(polygon) n = len(polygon)
for i in range(n): for i in range(n):
x1, y1 = polygon[i] x1, y1 = polygon[i]
x2, y2 = polygon[(i + 1) % n] x2, y2 = polygon[(i + 1) % n]
# Project point onto segment
dx, dy = x2 - x1, y2 - y1 dx, dy = x2 - x1, y2 - y1
length_sq = dx * dx + dy * dy length_sq = dx * dx + dy * dy
if length_sq == 0: if length_sq == 0:
@@ -230,12 +200,7 @@ class Search:
nearest = proj nearest = proj
return nearest return nearest
# ═══════════════════════════════════════════════════════════
# EXECUTION
# ═══════════════════════════════════════════════════════════
def run(self): def run(self):
"""Execute the pre-planned search by following waypoints."""
if not self.waypoints: if not self.waypoints:
self.plan() self.plan()
@@ -282,10 +247,6 @@ class Search:
return self.found_markers return self.found_markers
# ═══════════════════════════════════════════════════════════
# CAMERA & MARKER DETECTION
# ═══════════════════════════════════════════════════════════
def get_camera_frame(self): def get_camera_frame(self):
if self.camera is None: if self.camera is None:
return None return None
@@ -329,13 +290,7 @@ class Search:
return new_markers return new_markers
# ═══════════════════════════════════════════════════════════
# LANDING
# ═══════════════════════════════════════════════════════════
def execute_landing(self, initial_detections): def execute_landing(self, initial_detections):
"""Visual-servoing descent: keep the ArUco tag centered
while descending."""
target_det = None target_det = None
for d in initial_detections: for d in initial_detections:
if d.get("id") in self.land_ids: if d.get("id") in self.land_ids:
@@ -350,12 +305,10 @@ class Search:
print(f"[SEARCH] Target marker ID:{target_det['id']} " print(f"[SEARCH] Target marker ID:{target_det['id']} "
f"distance:{target_det['distance']:.2f}m") f"distance:{target_det['distance']:.2f}m")
# Camera parameters (must match model.sdf)
IMG_W, IMG_H = 640, 480 IMG_W, IMG_H = 640, 480
IMG_CX, IMG_CY = IMG_W / 2, IMG_H / 2 IMG_CX, IMG_CY = IMG_W / 2, IMG_H / 2
HFOV = 1.3962634 # 80° horizontal FOV HFOV = 1.3962634 # 80° horizontal FOV
# Landing parameters
DESCENT_STEP = 1.0 DESCENT_STEP = 1.0
LAND_ALT = 1.5 LAND_ALT = 1.5
CENTER_PX = 50 CENTER_PX = 50
@@ -363,7 +316,6 @@ class Search:
MAX_LOST = 30 MAX_LOST = 30
GAIN = 0.35 GAIN = 0.35
# Start visual servoing from current position
self.ctrl.update_state() self.ctrl.update_state()
current_alt = self.ctrl.altitude current_alt = self.ctrl.altitude
@@ -379,7 +331,6 @@ class Search:
self._landing = False self._landing = False
return return
# Descend one step
current_alt = max(current_alt - DESCENT_STEP, LAND_ALT) current_alt = max(current_alt - DESCENT_STEP, LAND_ALT)
print(f"\n[SEARCH] Descending to {current_alt:.1f}m") print(f"\n[SEARCH] Descending to {current_alt:.1f}m")
self.ctrl.update_state() self.ctrl.update_state()
@@ -387,13 +338,11 @@ class Search:
self.ctrl.move_local_ned(cur['x'], cur['y'], -current_alt) self.ctrl.move_local_ned(cur['x'], cur['y'], -current_alt)
sleep(2.0) sleep(2.0)
# Final precision centering
print(f"[SEARCH] Final centering at {LAND_ALT:.1f}m") print(f"[SEARCH] Final centering at {LAND_ALT:.1f}m")
self._center_over_marker( self._center_over_marker(
LAND_ALT, IMG_W, IMG_H, IMG_CX, IMG_CY, HFOV, LAND_ALT, IMG_W, IMG_H, IMG_CX, IMG_CY, HFOV,
35, MAX_CORRECTIONS, MAX_LOST, 0.25) 35, MAX_CORRECTIONS, MAX_LOST, 0.25)
# Land
print(f"\n[SEARCH] ===== LANDING =====") print(f"\n[SEARCH] ===== LANDING =====")
self.ctrl.land() self.ctrl.land()
self.landed = True self.landed = True
@@ -401,10 +350,6 @@ class Search:
def _center_over_marker(self, target_alt, img_w, img_h, img_cx, img_cy, def _center_over_marker(self, target_alt, img_w, img_h, img_cx, img_cy,
hfov, center_px, max_corrections, max_lost, gain): hfov, center_px, max_corrections, max_lost, gain):
"""Center the UAV over the ArUco marker using pixel error.
Returns True if centered, False if partial, None if lost.
"""
lost_count = 0 lost_count = 0
centered = False centered = False
@@ -465,15 +410,10 @@ class Search:
sleep(0.5) sleep(0.5)
if not centered: if not centered:
print(f"\n[SEARCH] Partially centered at {target_alt:.1f}m, " print(f"\n[SEARCH] Partially centered at {target_alt:.1f}m, continuing")
f"continuing")
return centered return centered
# ═══════════════════════════════════════════════════════════
# MOVEMENT
# ═══════════════════════════════════════════════════════════
def wait_for_position(self, target_x, target_y, timeout=None): def wait_for_position(self, target_x, target_y, timeout=None):
if timeout is None: if timeout is None:
timeout = self.MAX_TRAVEL_TIME timeout = self.MAX_TRAVEL_TIME
@@ -500,7 +440,6 @@ class Search:
return False return False
def _wait_arrival(self, target_x, target_y, timeout=15.0): def _wait_arrival(self, target_x, target_y, timeout=15.0):
"""Wait for position without checking markers (landing)."""
t0 = time.time() t0 = time.time()
while time.time() - t0 < timeout and self.running: while time.time() - t0 < timeout and self.running:
self.ctrl.update_state() self.ctrl.update_state()
@@ -524,10 +463,6 @@ class Search:
self.ctrl.move_local_ned(x, y, z) self.ctrl.move_local_ned(x, y, z)
return self.wait_for_position(x, y) return self.wait_for_position(x, y)
# ═══════════════════════════════════════════════════════════
# RESULTS
# ═══════════════════════════════════════════════════════════
def stop(self): def stop(self):
self.running = False self.running = False

21
src/control/uav_controller.py
View File

@@ -6,7 +6,6 @@ import argparse
import numpy as np import numpy as np
from time import sleep, perf_counter from time import sleep, perf_counter
from typing import Tuple from typing import Tuple
from pymavlink import mavutil from pymavlink import mavutil
import pymavlink.dialects.v20.ardupilotmega as mavlink import pymavlink.dialects.v20.ardupilotmega as mavlink
from pymavlink.dialects.v20.ardupilotmega import ( from pymavlink.dialects.v20.ardupilotmega import (
@@ -25,7 +24,6 @@ LOWEST_ALT = 3.0
GUIDED_MODE = 4 GUIDED_MODE = 4
GUIDED_NOGPS_MODE = 20 GUIDED_NOGPS_MODE = 20
DEFAULT_WPNAV_SPEED = 150 DEFAULT_WPNAV_SPEED = 150
DEFAULT_WPNAV_ACCEL = 100 DEFAULT_WPNAV_ACCEL = 100
DEFAULT_WPNAV_SPEED_UP = 100 DEFAULT_WPNAV_SPEED_UP = 100
@@ -35,7 +33,6 @@ DEFAULT_LOIT_SPEED = 150
class Controller: class Controller:
HOLD_ALT = HOLD_ALT HOLD_ALT = HOLD_ALT
LOWEST_ALT = LOWEST_ALT LOWEST_ALT = LOWEST_ALT
@@ -133,23 +130,14 @@ class Controller:
def arm(self, retries: int = 30): def arm(self, retries: int = 30):
self.set_mode_guided() self.set_mode_guided()
for i in range(retries): for i in range(retries):
print(f"[UAV] Arm attempt {i+1}/{retries}...") print(f"[UAV] Arm attempt {i+1}/{retries}...")
# Force arm: param2=21196 bypasses pre-arm checks
# (equivalent to MAVProxy's "arm throttle force")
self.conn.mav.command_long_send( self.conn.mav.command_long_send(
self.conn.target_system, self.conn.target_system,
self.conn.target_component, self.conn.target_component,
mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM, mavutil.mavlink.MAV_CMD_COMPONENT_ARM_DISARM,
0, 0, 1, 21196, 0, 0, 0, 0, 0)
1, # param1: 1 = arm
21196, # param2: force arm magic number
0, 0, 0, 0, 0)
# Wait up to 4 seconds for armed status instead of
# blocking forever with motors_armed_wait()
armed = False armed = False
for _ in range(20): for _ in range(20):
self._drain_messages() self._drain_messages()
@@ -157,16 +145,12 @@ class Controller:
armed = True armed = True
break break
time.sleep(0.2) time.sleep(0.2)
if armed: if armed:
self.armed = True self.armed = True
print("[UAV] Armed") print("[UAV] Armed")
return True return True
print(f"[UAV] Arm attempt {i+1} failed, retrying...") print(f"[UAV] Arm attempt {i+1} failed, retrying...")
time.sleep(2) time.sleep(2)
print("[UAV] FAILED to arm after all retries")
return False return False
def disarm(self): def disarm(self):
@@ -264,7 +248,6 @@ class Controller:
print(f"\r[UAV] Climbing: {self.altitude:.1f}m / {target_alt:.1f}m ", print(f"\r[UAV] Climbing: {self.altitude:.1f}m / {target_alt:.1f}m ",
end='', flush=True) end='', flush=True)
sleep(0.2) sleep(0.2)
print(f"\n[UAV] Altitude timeout at {self.altitude:.1f}m")
return False return False
def wait_for_gps(self, timeout: float = 120.0): def wait_for_gps(self, timeout: float = 120.0):
@@ -282,7 +265,6 @@ class Controller:
return True return True
elapsed = int(time.time() - t0) elapsed = int(time.time() - t0)
print(f"\r[UAV] Waiting for GPS ... {elapsed}s ", end='', flush=True) print(f"\r[UAV] Waiting for GPS ... {elapsed}s ", end='', flush=True)
print("\n[UAV] GPS timeout")
return False return False
def configure_speed_limits(self, def configure_speed_limits(self,
@@ -292,7 +274,6 @@ class Controller:
wpnav_speed_dn=DEFAULT_WPNAV_SPEED_DN, wpnav_speed_dn=DEFAULT_WPNAV_SPEED_DN,
wpnav_accel_z=DEFAULT_WPNAV_ACCEL_Z, wpnav_accel_z=DEFAULT_WPNAV_ACCEL_Z,
loit_speed=DEFAULT_LOIT_SPEED): loit_speed=DEFAULT_LOIT_SPEED):
params = { params = {
'WPNAV_SPEED': wpnav_speed, 'WPNAV_SPEED': wpnav_speed,
'WPNAV_ACCEL': wpnav_accel, 'WPNAV_ACCEL': wpnav_accel,

26
src/main.py
View File

@@ -65,7 +65,6 @@ def main():
parser.add_argument('--no-record', action='store_true', help='Disable recording') parser.add_argument('--no-record', action='store_true', help='Disable recording')
args = parser.parse_args() args = parser.parse_args()
# ── Start recorder ──────────────────────────────────────────
recorder = None recorder = None
if not args.no_record: if not args.no_record:
recorder = RunRecorder() recorder = RunRecorder()
@@ -78,13 +77,11 @@ def main():
altitude = args.altitude or search_cfg.get('altitude', 5.0) altitude = args.altitude or search_cfg.get('altitude', 5.0)
search_mode = args.search search_mode = args.search
# Marker config (single source of truth)
marker_cfg = search_cfg.get('marker', {}) marker_cfg = search_cfg.get('marker', {})
marker_dict = marker_cfg.get('dictionary', 'DICT_4X4_50') marker_dict = marker_cfg.get('dictionary', 'DICT_4X4_50')
marker_size = marker_cfg.get('size', 0.5) marker_size = marker_cfg.get('size', 0.5)
landing_ids = marker_cfg.get('landing_ids', [0]) landing_ids = marker_cfg.get('landing_ids', [0])
# UGV position
ugv_pos_cfg = ugv_cfg.get('position', {}) ugv_pos_cfg = ugv_cfg.get('position', {})
ugv_position = (ugv_pos_cfg.get('x', 5.0), ugv_pos_cfg.get('y', 5.0)) ugv_position = (ugv_pos_cfg.get('x', 5.0), ugv_pos_cfg.get('y', 5.0))
@@ -128,7 +125,6 @@ def main():
if mode_cfg: if mode_cfg:
search.configure(**{f"{search_mode}_{k}": v for k, v in mode_cfg.items()}) search.configure(**{f"{search_mode}_{k}": v for k, v in mode_cfg.items()})
# ── Pre-flight plan ─────────────────────────────────────────
geofence_cfg = search_cfg.get('geofence', {}) geofence_cfg = search_cfg.get('geofence', {})
geofence_points = None geofence_points = None
if geofence_cfg.get('enabled', False): if geofence_cfg.get('enabled', False):
@@ -143,7 +139,6 @@ def main():
print(f"[MAIN] Search: {search_mode} Altitude: {altitude}m") print(f"[MAIN] Search: {search_mode} Altitude: {altitude}m")
print(f"[MAIN] Marker: {marker_dict} size:{marker_size}m land:{landing_ids}") print(f"[MAIN] Marker: {marker_dict} size:{marker_size}m land:{landing_ids}")
# ── Start flight tracker (show plan before takeoff) ─────────
tracker = FlightTracker( tracker = FlightTracker(
window_size=600, window_size=600,
world_range=30.0, world_range=30.0,
@@ -166,12 +161,10 @@ def main():
ctrl.takeoff(altitude) ctrl.takeoff(altitude)
ctrl.wait_altitude(altitude, tolerance=1.0, timeout=30) ctrl.wait_altitude(altitude, tolerance=1.0, timeout=30)
# ── Start recording tracker + camera ────────────────────────
if recorder: if recorder:
recorder.start_recording(tracker=tracker, camera=camera) recorder.start_recording(tracker=tracker, camera=camera)
recorder.snapshot_camera("pre_search") recorder.snapshot_camera("pre_search")
# Feed position updates to tracker during search
original_check = search.check_for_markers original_check = search.check_for_markers
def tracked_check(): def tracked_check():
result = original_check() result = original_check()
@@ -186,13 +179,11 @@ def main():
return result return result
search.check_for_markers = tracked_check search.check_for_markers = tracked_check
# Wire up waypoint progress to tracker
search._on_waypoint_change = lambda idx: tracker.set_active_waypoint(idx) search._on_waypoint_change = lambda idx: tracker.set_active_waypoint(idx)
results = search.run() results = search.run()
search_results = search.get_results() search_results = search.get_results()
# ── Snapshot on detection ───────────────────────────────────
if recorder and results: if recorder and results:
recorder.snapshot_camera("marker_detected") recorder.snapshot_camera("marker_detected")
@@ -215,7 +206,6 @@ def main():
else: else:
wait_for_landing(ctrl) wait_for_landing(ctrl)
# ── Finalize recording ──────────────────────────────────────
if recorder: if recorder:
recorder.snapshot_camera("post_landing") recorder.snapshot_camera("post_landing")
recorder.save_summary( recorder.save_summary(
@@ -226,8 +216,20 @@ def main():
) )
recorder.stop() recorder.stop()
tracker.stop() print(f"\n{'=' * 50}")
print("[MAIN] Done.") print(f" Simulation finished.")
print(f" Tracker and Gazebo windows are still open.")
print(f" Press Ctrl+C to exit.")
print(f"{'=' * 50}\n")
try:
while True:
sleep(1)
except KeyboardInterrupt:
pass
finally:
tracker.stop()
print("[MAIN] Done.")
if __name__ == '__main__': if __name__ == '__main__':

288
src/navigation/flight_tracker.py
View File

@@ -1,16 +1,4 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""
Flight Tracker - Real-time 2D visualization of UAV and UGV positions.
Opens a separate OpenCV window showing:
- Grid with coordinates
- UAV position (green square) with flight path trail
- UGV position (red square)
- Geofence boundary (yellow dashed polygon with warning zone)
- Position data overlay
Can run standalone or be imported and updated by main.py.
"""
import sys import sys
import os import os
@@ -25,7 +13,6 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
class FlightTracker: class FlightTracker:
# Colors (BGR)
BG_COLOR = (30, 30, 30) BG_COLOR = (30, 30, 30)
GRID_COLOR = (60, 60, 60) GRID_COLOR = (60, 60, 60)
GRID_MAJOR_COLOR = (80, 80, 80) GRID_MAJOR_COLOR = (80, 80, 80)
@@ -36,25 +23,16 @@ class FlightTracker:
TEXT_COLOR = (200, 200, 200) TEXT_COLOR = (200, 200, 200)
TEXT_DIM_COLOR = (120, 120, 120) TEXT_DIM_COLOR = (120, 120, 120)
MARKER_COLOR = (180, 0, 220) MARKER_COLOR = (180, 0, 220)
FENCE_COLOR = (0, 220, 255) # Yellow-cyan FENCE_COLOR = (0, 220, 255)
FENCE_WARN_COLOR = (0, 140, 200) # Darker yellow FENCE_WARN_COLOR = (0, 140, 200)
FENCE_BREACH_COLOR = (0, 0, 255) # Red FENCE_BREACH_COLOR = (0, 0, 255)
PLAN_COLOR = (180, 140, 50) # Muted cyan-blue PLAN_COLOR = (180, 140, 50)
PLAN_ACTIVE_COLOR = (0, 180, 255) # Orange PLAN_ACTIVE_COLOR = (0, 180, 255)
PLAN_VISITED_COLOR = (60, 60, 60) # Dim gray PLAN_VISITED_COLOR = (60, 60, 60)
def __init__(self, window_size=600, world_range=30.0, def __init__(self, window_size=600, world_range=30.0,
ugv_position=(5.0, 5.0), show_ugv=True, ugv_position=(5.0, 5.0), show_ugv=True,
geofence=None): geofence=None):
"""
Args:
window_size: Window dimension in pixels (square)
world_range: World coordinate range (-range to +range) in meters
ugv_position: Known UGV position in local NED (x, y)
show_ugv: Whether to draw the UGV
geofence: Dict with 'points' (list of [x,y]), 'warning_distance',
'min_altitude', 'max_altitude'
"""
self.window_size = window_size self.window_size = window_size
self.world_range = world_range self.world_range = world_range
self.ugv_pos = ugv_position self.ugv_pos = ugv_position
@@ -66,11 +44,9 @@ class FlightTracker:
self.trail = deque() self.trail = deque()
self.markers_found = {} self.markers_found = {}
# Flight plan
self.plan_waypoints = [] self.plan_waypoints = []
self.active_waypoint = 0 self.active_waypoint = 0
# Geofence
self.geofence = geofence self.geofence = geofence
self.fence_points = [] self.fence_points = []
self.fence_warn_dist = 3.0 self.fence_warn_dist = 3.0
@@ -91,14 +67,11 @@ class FlightTracker:
self.total_width = self.window_size + self.panel_width self.total_width = self.window_size + self.panel_width
def world_to_pixel(self, x, y): def world_to_pixel(self, x, y):
"""Convert NED world coords to pixel coords on the map.
NED: x=North(up on screen), y=East(right on screen)"""
px = int(self.window_size / 2 + (y / self.world_range) * (self.window_size / 2)) px = int(self.window_size / 2 + (y / self.world_range) * (self.window_size / 2))
py = int(self.window_size / 2 - (x / self.world_range) * (self.window_size / 2)) py = int(self.window_size / 2 - (x / self.world_range) * (self.window_size / 2))
return (px, py) return (px, py)
def update_uav(self, x, y, altitude=0.0, heading=0.0): def update_uav(self, x, y, altitude=0.0, heading=0.0):
"""Update UAV position (NED coords)."""
with self._lock: with self._lock:
self.uav_pos = (x, y) self.uav_pos = (x, y)
self.uav_alt = altitude self.uav_alt = altitude
@@ -106,28 +79,23 @@ class FlightTracker:
self.trail.append((x, y)) self.trail.append((x, y))
def update_ugv(self, x, y): def update_ugv(self, x, y):
"""Update UGV position."""
with self._lock: with self._lock:
self.ugv_pos = (x, y) self.ugv_pos = (x, y)
def add_marker(self, marker_id, x, y): def add_marker(self, marker_id, x, y):
"""Record a detected marker position."""
with self._lock: with self._lock:
self.markers_found[marker_id] = (x, y) self.markers_found[marker_id] = (x, y)
def set_plan(self, waypoints): def set_plan(self, waypoints):
"""Set pre-computed flight plan waypoints for visualization."""
with self._lock: with self._lock:
self.plan_waypoints = list(waypoints) self.plan_waypoints = list(waypoints)
self.active_waypoint = 0 self.active_waypoint = 0
def set_active_waypoint(self, idx): def set_active_waypoint(self, idx):
"""Update which waypoint the UAV is currently targeting."""
with self._lock: with self._lock:
self.active_waypoint = idx self.active_waypoint = idx
def _point_in_polygon(self, px, py, polygon): def _point_in_polygon(self, px, py, polygon):
"""Simple ray-casting point-in-polygon test."""
n = len(polygon) n = len(polygon)
inside = False inside = False
j = n - 1 j = n - 1
@@ -140,13 +108,11 @@ class FlightTracker:
return inside return inside
def _distance_to_polygon_edge(self, px, py, polygon): def _distance_to_polygon_edge(self, px, py, polygon):
"""Minimum distance from point to polygon edges."""
min_dist = float('inf') min_dist = float('inf')
n = len(polygon) n = len(polygon)
for i in range(n): for i in range(n):
x1, y1 = polygon[i] x1, y1 = polygon[i]
x2, y2 = polygon[(i + 1) % n] x2, y2 = polygon[(i + 1) % n]
# Point-to-segment distance
dx, dy = x2 - x1, y2 - y1 dx, dy = x2 - x1, y2 - y1
length_sq = dx * dx + dy * dy length_sq = dx * dx + dy * dy
if length_sq == 0: if length_sq == 0:
@@ -160,13 +126,11 @@ class FlightTracker:
return min_dist return min_dist
def draw(self): def draw(self):
"""Render the full tracker frame."""
frame = np.full((self.window_size, self.total_width, 3), frame = np.full((self.window_size, self.total_width, 3),
self.BG_COLOR, dtype=np.uint8) self.BG_COLOR, dtype=np.uint8)
self._draw_grid(frame) self._draw_grid(frame)
# Draw geofence before other elements
if self.fence_points: if self.fence_points:
self._draw_geofence(frame) self._draw_geofence(frame)
@@ -180,7 +144,6 @@ class FlightTracker:
plan_copy = list(self.plan_waypoints) plan_copy = list(self.plan_waypoints)
active_wp = self.active_waypoint active_wp = self.active_waypoint
# Draw plan UNDER the trail
if plan_copy: if plan_copy:
self._draw_plan(frame, plan_copy, active_wp) self._draw_plan(frame, plan_copy, active_wp)
@@ -202,57 +165,37 @@ class FlightTracker:
return frame return frame
def _draw_grid(self, frame): def _draw_grid(self, frame):
"""Draw coordinate grid."""
step = 5.0 step = 5.0
r = self.world_range r = self.world_range
n = int(r / step) n = int(r / step)
for i in range(-n, n + 1): for i in range(-n, n + 1):
val = i * step val = i * step
color = self.GRID_MAJOR_COLOR if i % 2 == 0 else self.GRID_COLOR color = self.GRID_MAJOR_COLOR if i % 2 == 0 else self.GRID_COLOR
# Vertical lines (constant y)
px, _ = self.world_to_pixel(0, val) px, _ = self.world_to_pixel(0, val)
cv2.line(frame, (px, 0), (px, self.window_size), color, 1) cv2.line(frame, (px, 0), (px, self.window_size), color, 1)
# Horizontal lines (constant x)
_, py = self.world_to_pixel(val, 0) _, py = self.world_to_pixel(val, 0)
cv2.line(frame, (0, py), (self.window_size, py), color, 1) cv2.line(frame, (0, py), (self.window_size, py), color, 1)
# Axis labels
if i != 0 and i % 2 == 0: if i != 0 and i % 2 == 0:
cv2.putText(frame, f"{val:.0f}", cv2.putText(frame, f"{val:.0f}", (px + 2, self.window_size - 5),
(px + 2, self.window_size - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.3, self.TEXT_DIM_COLOR, 1) cv2.FONT_HERSHEY_SIMPLEX, 0.3, self.TEXT_DIM_COLOR, 1)
cv2.putText(frame, f"{val:.0f}", cv2.putText(frame, f"{val:.0f}", (5, py - 3),
(5, py - 3),
cv2.FONT_HERSHEY_SIMPLEX, 0.3, self.TEXT_DIM_COLOR, 1) cv2.FONT_HERSHEY_SIMPLEX, 0.3, self.TEXT_DIM_COLOR, 1)
# Draw axes (origin cross)
cx, cy = self.world_to_pixel(0, 0) cx, cy = self.world_to_pixel(0, 0)
cv2.line(frame, (cx, 0), (cx, self.window_size), self.AXIS_COLOR, 1) cv2.line(frame, (cx, 0), (cx, self.window_size), self.AXIS_COLOR, 1)
cv2.line(frame, (0, cy), (self.window_size, cy), self.AXIS_COLOR, 1) cv2.line(frame, (0, cy), (self.window_size, cy), self.AXIS_COLOR, 1)
cv2.putText(frame, "N", (cx + 4, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.AXIS_COLOR, 1)
# Axis labels at origin cv2.putText(frame, "E", (self.window_size - 15, cy - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.AXIS_COLOR, 1)
cv2.putText(frame, "N", (cx + 4, 15),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.AXIS_COLOR, 1)
cv2.putText(frame, "E", (self.window_size - 15, cy - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.AXIS_COLOR, 1)
def _draw_geofence(self, frame): def _draw_geofence(self, frame):
"""Draw geofence boundary with warning zone and fill."""
pts = [self.fence_points[i] for i in range(len(self.fence_points))] pts = [self.fence_points[i] for i in range(len(self.fence_points))]
pixel_pts = np.array([self.world_to_pixel(x, y) for x, y in pts], dtype=np.int32) pixel_pts = np.array([self.world_to_pixel(x, y) for x, y in pts], dtype=np.int32)
# Semi-transparent fill inside the geofence
overlay = frame.copy() overlay = frame.copy()
cv2.fillPoly(overlay, [pixel_pts], (40, 50, 40)) cv2.fillPoly(overlay, [pixel_pts], (40, 50, 40))
cv2.addWeighted(overlay, 0.3, frame, 0.7, 0, frame) cv2.addWeighted(overlay, 0.3, frame, 0.7, 0, frame)
# Draw warning zone (inner boundary)
if self.fence_warn_dist > 0: if self.fence_warn_dist > 0:
# Inset the polygon by warning_distance to show warning zone
# Simple approach: draw dashed inner boundary
warn_pts = [] warn_pts = []
cx = sum(p[0] for p in pts) / len(pts) cx = sum(p[0] for p in pts) / len(pts)
cy = sum(p[1] for p in pts) / len(pts) cy = sum(p[1] for p in pts) / len(pts)
@@ -266,25 +209,17 @@ class FlightTracker:
else: else:
warn_pts.append((x, y)) warn_pts.append((x, y))
warn_pixel_pts = np.array( warn_pixel_pts = np.array([self.world_to_pixel(x, y) for x, y in warn_pts], dtype=np.int32)
[self.world_to_pixel(x, y) for x, y in warn_pts], dtype=np.int32
)
self._draw_dashed_polygon(frame, warn_pixel_pts, self.FENCE_WARN_COLOR, 1, 8) self._draw_dashed_polygon(frame, warn_pixel_pts, self.FENCE_WARN_COLOR, 1, 8)
# Draw main fence boundary (solid bright line)
cv2.polylines(frame, [pixel_pts], True, self.FENCE_COLOR, 2) cv2.polylines(frame, [pixel_pts], True, self.FENCE_COLOR, 2)
# Corner markers
for px, py in pixel_pts: for px, py in pixel_pts:
cv2.circle(frame, (int(px), int(py)), 4, self.FENCE_COLOR, -1) cv2.circle(frame, (int(px), int(py)), 4, self.FENCE_COLOR, -1)
# Label
top_px, top_py = pixel_pts[0] top_px, top_py = pixel_pts[0]
cv2.putText(frame, "GEOFENCE", (int(top_px) + 6, int(top_py) - 8), cv2.putText(frame, "GEOFENCE", (int(top_px) + 6, int(top_py) - 8),
cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.FENCE_COLOR, 1) cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.FENCE_COLOR, 1)
def _draw_dashed_polygon(self, frame, pts, color, thickness, dash_len): def _draw_dashed_polygon(self, frame, pts, color, thickness, dash_len):
"""Draw a polygon with dashed lines."""
n = len(pts) n = len(pts)
for i in range(n): for i in range(n):
p1 = pts[i] p1 = pts[i]
@@ -292,7 +227,6 @@ class FlightTracker:
self._draw_dashed_line(frame, tuple(p1), tuple(p2), color, thickness, dash_len) self._draw_dashed_line(frame, tuple(p1), tuple(p2), color, thickness, dash_len)
def _draw_dashed_line(self, frame, pt1, pt2, color, thickness, dash_len): def _draw_dashed_line(self, frame, pt1, pt2, color, thickness, dash_len):
"""Draw a dashed line between two points."""
x1, y1 = pt1 x1, y1 = pt1
x2, y2 = pt2 x2, y2 = pt2
dx = x2 - x1 dx = x2 - x1
@@ -313,268 +247,142 @@ class FlightTracker:
cv2.line(frame, (sx, sy), (ex, ey), color, thickness) cv2.line(frame, (sx, sy), (ex, ey), color, thickness)
def _draw_plan(self, frame, waypoints, active_idx): def _draw_plan(self, frame, waypoints, active_idx):
"""Draw the pre-flight plan: dashed path with numbered waypoints."""
if len(waypoints) < 2: if len(waypoints) < 2:
return return
pts = [self.world_to_pixel(x, y) for x, y in waypoints] pts = [self.world_to_pixel(x, y) for x, y in waypoints]
# Draw connecting lines
for i in range(1, len(pts)): for i in range(1, len(pts)):
if i <= active_idx: color = self.PLAN_VISITED_COLOR if i <= active_idx else self.PLAN_COLOR
color = self.PLAN_VISITED_COLOR
else:
color = self.PLAN_COLOR
self._draw_dashed_line(frame, pts[i - 1], pts[i], color, 1, 8) self._draw_dashed_line(frame, pts[i - 1], pts[i], color, 1, 8)
# Draw waypoint markers
for i, (px, py) in enumerate(pts): for i, (px, py) in enumerate(pts):
if i < active_idx: if i < active_idx:
# Visited — small dim dot
cv2.circle(frame, (px, py), 3, self.PLAN_VISITED_COLOR, -1) cv2.circle(frame, (px, py), 3, self.PLAN_VISITED_COLOR, -1)
elif i == active_idx: elif i == active_idx:
# Active target — bright ring
cv2.circle(frame, (px, py), 8, self.PLAN_ACTIVE_COLOR, 2) cv2.circle(frame, (px, py), 8, self.PLAN_ACTIVE_COLOR, 2)
cv2.circle(frame, (px, py), 3, self.PLAN_ACTIVE_COLOR, -1) cv2.circle(frame, (px, py), 3, self.PLAN_ACTIVE_COLOR, -1)
else: else:
# Future — small cyan dot
cv2.circle(frame, (px, py), 4, self.PLAN_COLOR, 1) cv2.circle(frame, (px, py), 4, self.PLAN_COLOR, 1)
cv2.putText(frame, str(i), (px + 7, py - 5),
# Waypoint number label
label = str(i)
cv2.putText(frame, label, (px + 7, py - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.28, self.PLAN_COLOR, 1) cv2.FONT_HERSHEY_SIMPLEX, 0.28, self.PLAN_COLOR, 1)
def _draw_trail(self, frame, trail): def _draw_trail(self, frame, trail):
"""Draw the flight path trail with persistent visibility."""
if len(trail) < 2: if len(trail) < 2:
return return
# Minimum brightness so old trail never disappears
MIN_ALPHA = 0.3 MIN_ALPHA = 0.3
FADE_POINTS = 200 # only the newest N points fade FADE_POINTS = 200
pts = [self.world_to_pixel(x, y) for x, y in trail] pts = [self.world_to_pixel(x, y) for x, y in trail]
n = len(pts) n = len(pts)
fade_start = max(0, n - FADE_POINTS) fade_start = max(0, n - FADE_POINTS)
for i in range(1, n): for i in range(1, n):
if i <= fade_start: if i <= fade_start:
# Old segments: solid dim color (never invisible)
alpha = MIN_ALPHA alpha = MIN_ALPHA
else: else:
# Recent segments: fade from MIN_ALPHA to 1.0
progress = (i - fade_start) / FADE_POINTS progress = (i - fade_start) / FADE_POINTS
alpha = MIN_ALPHA + (1.0 - MIN_ALPHA) * progress alpha = MIN_ALPHA + (1.0 - MIN_ALPHA) * progress
color = ( color = (int(self.UAV_TRAIL_COLOR[0] * alpha),
int(self.UAV_TRAIL_COLOR[0] * alpha), int(self.UAV_TRAIL_COLOR[1] * alpha),
int(self.UAV_TRAIL_COLOR[1] * alpha), int(self.UAV_TRAIL_COLOR[2] * alpha))
int(self.UAV_TRAIL_COLOR[2] * alpha),
)
cv2.line(frame, pts[i - 1], pts[i], color, 2) cv2.line(frame, pts[i - 1], pts[i], color, 2)
def _draw_uav(self, frame, pos, heading): def _draw_uav(self, frame, pos, heading):
"""Draw UAV as a green square with heading indicator."""
px, py = self.world_to_pixel(pos[0], pos[1]) px, py = self.world_to_pixel(pos[0], pos[1])
size = 8 size = 8
cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), self.UAV_COLOR, -1)
self.UAV_COLOR, -1) cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), (255, 255, 255), 1)
cv2.rectangle(frame, (px - size, py - size), (px + size, py + size),
(255, 255, 255), 1)
# Heading line
hx = int(px + 15 * math.sin(heading)) hx = int(px + 15 * math.sin(heading))
hy = int(py - 15 * math.cos(heading)) hy = int(py - 15 * math.cos(heading))
cv2.line(frame, (px, py), (hx, hy), self.UAV_COLOR, 2) cv2.line(frame, (px, py), (hx, hy), self.UAV_COLOR, 2)
cv2.putText(frame, "UAV", (px + 12, py - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.UAV_COLOR, 1)
# Label
cv2.putText(frame, "UAV", (px + 12, py - 4),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.UAV_COLOR, 1)
def _draw_ugv(self, frame, pos): def _draw_ugv(self, frame, pos):
"""Draw UGV as a red square."""
px, py = self.world_to_pixel(pos[0], pos[1]) px, py = self.world_to_pixel(pos[0], pos[1])
size = 10 size = 10
cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), self.UGV_COLOR, -1)
self.UGV_COLOR, -1) cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), (255, 255, 255), 1)
cv2.rectangle(frame, (px - size, py - size), (px + size, py + size), cv2.putText(frame, "UGV", (px + 14, py - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.UGV_COLOR, 1)
(255, 255, 255), 1)
cv2.putText(frame, "UGV", (px + 14, py - 4),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.UGV_COLOR, 1)
def _draw_panel(self, frame, uav, alt, hdg, ugv, trail): def _draw_panel(self, frame, uav, alt, hdg, ugv, trail):
"""Draw the data panel on the right side."""
x0 = self.window_size + 10 x0 = self.window_size + 10
y = 25 y = 25
cv2.line(frame, (self.window_size, 0), (self.window_size, self.window_size), self.GRID_MAJOR_COLOR, 1)
# Separator line cv2.putText(frame, "FLIGHT TRACKER", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
cv2.line(frame, (self.window_size, 0), (self.window_size, self.window_size),
self.GRID_MAJOR_COLOR, 1)
# Title
cv2.putText(frame, "FLIGHT TRACKER", (x0, y),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
y += 30 y += 30
# Elapsed time
elapsed = time.time() - self.start_time elapsed = time.time() - self.start_time
mins = int(elapsed // 60) mins, secs = int(elapsed // 60), int(elapsed % 60)
secs = int(elapsed % 60) cv2.putText(frame, f"Time: {mins:02d}:{secs:02d}", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.TEXT_DIM_COLOR, 1)
cv2.putText(frame, f"Time: {mins:02d}:{secs:02d}", (x0, y),
cv2.FONT_HERSHEY_SIMPLEX, 0.4, self.TEXT_DIM_COLOR, 1)
y += 30 y += 30
cv2.putText(frame, "UAV", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.UAV_COLOR, 1)
# UAV section
cv2.putText(frame, "UAV", (x0, y),
cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.UAV_COLOR, 1)
y += 22 y += 22
cv2.putText(frame, f"X: {uav[0]:+.1f}m", (x0, y), cv2.putText(frame, f"X: {uav[0]:+.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 18 y += 18
cv2.putText(frame, f"Y: {uav[1]:+.1f}m", (x0, y), cv2.putText(frame, f"Y: {uav[1]:+.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 18 y += 18
cv2.putText(frame, f"Alt: {alt:.1f}m", (x0, y), cv2.putText(frame, f"Alt: {alt:.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 18 y += 18
cv2.putText(frame, f"Hdg: {math.degrees(hdg):.0f} deg", (x0, y), cv2.putText(frame, f"Hdg: {math.degrees(hdg):.0f} deg", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 30 y += 30
# UGV section
if self.show_ugv: if self.show_ugv:
cv2.putText(frame, "UGV", (x0, y), cv2.putText(frame, "UGV", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.UGV_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.UGV_COLOR, 1)
y += 22 y += 22
cv2.putText(frame, f"X: {ugv[0]:+.1f}m", (x0, y), cv2.putText(frame, f"X: {ugv[0]:+.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 18 y += 18
cv2.putText(frame, f"Y: {ugv[1]:+.1f}m", (x0, y), cv2.putText(frame, f"Y: {ugv[1]:+.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 18 y += 18
dist = np.sqrt((uav[0] - ugv[0])**2 + (uav[1] - ugv[1])**2) dist = np.sqrt((uav[0] - ugv[0])**2 + (uav[1] - ugv[1])**2)
cv2.putText(frame, f"Dist: {dist:.1f}m", (x0, y), cv2.putText(frame, f"Dist: {dist:.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 30 y += 30
# Geofence section
if self.fence_points: if self.fence_points:
# Check fence status
inside = self._point_in_polygon(uav[0], uav[1], self.fence_points) inside = self._point_in_polygon(uav[0], uav[1], self.fence_points)
alt_ok = self.fence_alt_min <= alt <= self.fence_alt_max alt_ok = self.fence_alt_min <= alt <= self.fence_alt_max
fence_dist = self._distance_to_polygon_edge(uav[0], uav[1], self.fence_points) fence_dist = self._distance_to_polygon_edge(uav[0], uav[1], self.fence_points)
if not inside or not alt_ok: if not inside or not alt_ok:
status = "BREACHED" status, status_color = "BREACHED", self.FENCE_BREACH_COLOR
status_color = self.FENCE_BREACH_COLOR
elif fence_dist < self.fence_warn_dist: elif fence_dist < self.fence_warn_dist:
status = "WARNING" status, status_color = "WARNING", self.FENCE_WARN_COLOR
status_color = self.FENCE_WARN_COLOR
else: else:
status = "OK" status, status_color = "OK", self.FENCE_COLOR
status_color = self.FENCE_COLOR
cv2.putText(frame, "GEOFENCE", (x0, y), cv2.putText(frame, "GEOFENCE", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.FENCE_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.FENCE_COLOR, 1)
y += 22 y += 22
cv2.putText(frame, f"Status: {status}", (x0, y), cv2.putText(frame, f"Status: {status}", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, status_color, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, status_color, 1)
y += 18 y += 18
cv2.putText(frame, f"Edge: {fence_dist:.1f}m", (x0, y), cv2.putText(frame, f"Edge: {fence_dist:.1f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 18 y += 18
cv2.putText(frame, f"Alt lim: {self.fence_alt_min:.0f}-{self.fence_alt_max:.0f}m", (x0, y), cv2.putText(frame, f"Alt lim: {self.fence_alt_min:.0f}-{self.fence_alt_max:.0f}m", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.38, self.TEXT_COLOR, 1)
y += 30 y += 30
# Markers section cv2.putText(frame, "MARKERS", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.MARKER_COLOR, 1)
cv2.putText(frame, "MARKERS", (x0, y),
cv2.FONT_HERSHEY_SIMPLEX, 0.45, self.MARKER_COLOR, 1)
y += 22 y += 22
if self.markers_found: if self.markers_found:
for mid, (mx, my) in self.markers_found.items(): for mid, (mx, my) in self.markers_found.items():
cv2.putText(frame, f"ID:{mid} ({mx:.1f},{my:.1f})", (x0, y), cv2.putText(frame, f"ID:{mid} ({mx:.1f},{my:.1f})", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.TEXT_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.TEXT_COLOR, 1)
y += 16 y += 16
else: else:
cv2.putText(frame, "None", (x0, y), cv2.putText(frame, "None", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.TEXT_DIM_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.TEXT_DIM_COLOR, 1)
y += 16 y += 16
y += 20 y += 20
cv2.putText(frame, f"Trail: {len(trail)} pts", (x0, y), cv2.putText(frame, f"Trail: {len(trail)} pts", (x0, y), cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.TEXT_DIM_COLOR, 1)
cv2.FONT_HERSHEY_SIMPLEX, 0.35, self.TEXT_DIM_COLOR, 1)
def start(self): def start(self):
"""Start the display loop in a background thread."""
self.running = True self.running = True
self._thread = threading.Thread(target=self._display_loop, daemon=True) self._thread = threading.Thread(target=self._display_loop, daemon=True)
self._thread.start() self._thread.start()
def _display_loop(self): def _display_loop(self):
"""OpenCV display loop (runs in its own thread)."""
print("[TRACKER] Flight tracker window started") print("[TRACKER] Flight tracker window started")
while self.running: while self.running:
frame = self.draw() frame = self.draw()
cv2.imshow("Flight Tracker", frame) cv2.imshow("Flight Tracker", frame)
key = cv2.waitKey(50) & 0xFF if cv2.waitKey(50) & 0xFF == ord('q'):
if key == ord('q'):
break break
cv2.destroyWindow("Flight Tracker") cv2.destroyWindow("Flight Tracker")
print("[TRACKER] Flight tracker stopped") print("[TRACKER] Flight tracker stopped")
def stop(self): def stop(self):
"""Stop the display loop."""
self.running = False self.running = False
if self._thread: if self._thread:
self._thread.join(timeout=2.0) self._thread.join(timeout=2.0)
def main():
"""Standalone test mode - simulates a spiral flight path with geofence."""
geofence = {
'enabled': True,
'warning_distance': 3.0,
'min_altitude': 0.0,
'max_altitude': 15.0,
'points': [[-5, -5], [-5, 25], [25, 25], [25, -5]],
}
tracker = FlightTracker(
window_size=600,
world_range=30.0,
ugv_position=(5.0, 5.0),
geofence=geofence,
)
print("Flight Tracker - Standalone Test (with geofence)")
print("Press 'q' to quit")
x, y = 0.0, 0.0
t = 0.0
while True:
t += 0.05
radius = t * 0.5
x = radius * math.cos(t)
y = radius * math.sin(t)
heading = math.atan2(math.cos(t), -math.sin(t))
tracker.update_uav(x, y, altitude=5.0, heading=heading)
if t > 5 and 0 not in tracker.markers_found:
tracker.add_marker(0, 5.0, 5.0)
frame = tracker.draw()
cv2.imshow("Flight Tracker", frame)
key = cv2.waitKey(30) & 0xFF
if key == ord('q'):
break
cv2.destroyAllWindows()
if __name__ == "__main__":
main()

18
src/utils/helpers.py
View File

@@ -1,5 +1,4 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""Utility helper functions."""
import numpy as np import numpy as np
import yaml import yaml
@@ -77,32 +76,19 @@ class PIDController:
self.ki = ki self.ki = ki
self.kd = kd self.kd = kd
self.output_limits = output_limits self.output_limits = output_limits
self.integral = 0.0 self.integral = 0.0
self.last_error = 0.0 self.last_error = 0.0
self.last_time = None self.last_time = None
def compute(self, error, current_time): def compute(self, error, current_time):
if self.last_time is None: dt = 0.0 if self.last_time is None else current_time - self.last_time
dt = 0.0
else:
dt = current_time - self.last_time
self.integral += error * dt self.integral += error * dt
derivative = (error - self.last_error) / dt if dt > 0 else 0.0
if dt > 0:
derivative = (error - self.last_error) / dt
else:
derivative = 0.0
output = self.kp * error + self.ki * self.integral + self.kd * derivative output = self.kp * error + self.ki * self.integral + self.kd * derivative
if self.output_limits: if self.output_limits:
output = clamp(output, self.output_limits[0], self.output_limits[1]) output = clamp(output, self.output_limits[0], self.output_limits[1])
self.last_error = error self.last_error = error
self.last_time = current_time self.last_time = current_time
return output return output
def reset(self): def reset(self):

41
src/utils/recorder.py
View File

@@ -1,14 +1,4 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""
Run Recorder — Captures logs, flight path, camera frames, and summary for each run.
Creates timestamped folders in /results with:
flight_path.avi — Flight tracker video
camera.avi — Downward camera video
flight_path.png — Final flight tracker snapshot
log.txt — Full console output
summary.json — Run metadata and results
"""
import os import os
import sys import sys
@@ -23,8 +13,6 @@ from datetime import datetime
class RunRecorder: class RunRecorder:
"""Records all outputs from a simulation run."""
def __init__(self, results_dir=None, fps=5): def __init__(self, results_dir=None, fps=5):
if results_dir is None: if results_dir is None:
project_dir = Path(__file__).resolve().parent.parent project_dir = Path(__file__).resolve().parent.parent
@@ -33,7 +21,6 @@ class RunRecorder:
results_dir = Path(results_dir) results_dir = Path(results_dir)
results_dir.mkdir(parents=True, exist_ok=True) results_dir.mkdir(parents=True, exist_ok=True)
# Find next sequential run number
existing = [d.name for d in results_dir.iterdir() existing = [d.name for d in results_dir.iterdir()
if d.is_dir() and d.name.startswith("run_")] if d.is_dir() and d.name.startswith("run_")]
nums = [] nums = []
@@ -51,7 +38,6 @@ class RunRecorder:
self.fps = fps self.fps = fps
self.start_time = time.time() self.start_time = time.time()
# Log capture
self._log_path = self.run_dir / "log.txt" self._log_path = self.run_dir / "log.txt"
self._log_file = open(self._log_path, "w") self._log_file = open(self._log_path, "w")
self._original_stdout = sys.stdout self._original_stdout = sys.stdout
@@ -59,29 +45,24 @@ class RunRecorder:
self._tee_stdout = _TeeWriter(sys.stdout, self._log_file) self._tee_stdout = _TeeWriter(sys.stdout, self._log_file)
self._tee_stderr = _TeeWriter(sys.stderr, self._log_file) self._tee_stderr = _TeeWriter(sys.stderr, self._log_file)
# Video writers (initialized lazily on first frame)
self._tracker_writer = None self._tracker_writer = None
self._camera_writer = None self._camera_writer = None
self._tracker_size = None self._tracker_size = None
self._camera_size = None self._camera_size = None
# Frame counters
self._tracker_frames = 0 self._tracker_frames = 0
self._camera_frames = 0 self._camera_frames = 0
# Snapshot storage
self._last_tracker_frame = None self._last_tracker_frame = None
self._last_camera_frame = None self._last_camera_frame = None
self._camera_snapshots = [] self._camera_snapshots = []
# Recording thread
self._recording = False self._recording = False
self._tracker_ref = None self._tracker_ref = None
self._camera_ref = None self._camera_ref = None
self._record_thread = None self._record_thread = None
self._lock = threading.Lock() self._lock = threading.Lock()
# Metadata
self.metadata = { self.metadata = {
"run": run_num, "run": run_num,
"start_time": datetime.now().isoformat(), "start_time": datetime.now().isoformat(),
@@ -91,17 +72,14 @@ class RunRecorder:
print(f"[REC] Recording to: {self.run_dir}") print(f"[REC] Recording to: {self.run_dir}")
def start_logging(self): def start_logging(self):
"""Redirect stdout/stderr to both console and log file."""
sys.stdout = self._tee_stdout sys.stdout = self._tee_stdout
sys.stderr = self._tee_stderr sys.stderr = self._tee_stderr
def stop_logging(self): def stop_logging(self):
"""Restore original stdout/stderr."""
sys.stdout = self._original_stdout sys.stdout = self._original_stdout
sys.stderr = self._original_stderr sys.stderr = self._original_stderr
def start_recording(self, tracker=None, camera=None): def start_recording(self, tracker=None, camera=None):
"""Start background recording of tracker and camera frames."""
self._tracker_ref = tracker self._tracker_ref = tracker
self._camera_ref = camera self._camera_ref = camera
self._recording = True self._recording = True
@@ -111,12 +89,10 @@ class RunRecorder:
self._record_thread.start() self._record_thread.start()
def _record_loop(self): def _record_loop(self):
"""Periodically capture frames from tracker and camera."""
interval = 1.0 / self.fps interval = 1.0 / self.fps
while self._recording: while self._recording:
t0 = time.time() t0 = time.time()
# Capture tracker frame
if self._tracker_ref is not None: if self._tracker_ref is not None:
try: try:
frame = self._tracker_ref.draw() frame = self._tracker_ref.draw()
@@ -126,7 +102,6 @@ class RunRecorder:
except Exception: except Exception:
pass pass
# Capture camera frame
if self._camera_ref is not None: if self._camera_ref is not None:
try: try:
frame = self._camera_ref.frames.get("downward") frame = self._camera_ref.frames.get("downward")
@@ -141,7 +116,6 @@ class RunRecorder:
time.sleep(sleep_time) time.sleep(sleep_time)
def _write_tracker_frame(self, frame): def _write_tracker_frame(self, frame):
"""Write a frame to the tracker video."""
h, w = frame.shape[:2] h, w = frame.shape[:2]
if self._tracker_writer is None: if self._tracker_writer is None:
self._tracker_size = (w, h) self._tracker_size = (w, h)
@@ -152,7 +126,6 @@ class RunRecorder:
self._tracker_frames += 1 self._tracker_frames += 1
def _write_camera_frame(self, frame): def _write_camera_frame(self, frame):
"""Write a frame to the camera video."""
h, w = frame.shape[:2] h, w = frame.shape[:2]
if self._camera_writer is None: if self._camera_writer is None:
self._camera_size = (w, h) self._camera_size = (w, h)
@@ -163,7 +136,6 @@ class RunRecorder:
self._camera_frames += 1 self._camera_frames += 1
def snapshot_camera(self, label="snapshot"): def snapshot_camera(self, label="snapshot"):
"""Save a named snapshot of the current camera frame."""
if self._camera_ref is None: if self._camera_ref is None:
return return
frame = self._camera_ref.frames.get("downward") frame = self._camera_ref.frames.get("downward")
@@ -178,7 +150,6 @@ class RunRecorder:
def save_summary(self, search_mode="", altitude=0, markers=None, def save_summary(self, search_mode="", altitude=0, markers=None,
landed=False, extra=None): landed=False, extra=None):
"""Write the run summary JSON."""
duration = time.time() - self.start_time duration = time.time() - self.start_time
mins = int(duration // 60) mins = int(duration // 60)
secs = int(duration % 60) secs = int(duration % 60)
@@ -223,29 +194,24 @@ class RunRecorder:
print(f"[REC] Summary saved: {path}") print(f"[REC] Summary saved: {path}")
def stop(self): def stop(self):
"""Stop recording and finalize all outputs."""
self._recording = False self._recording = False
if self._record_thread: if self._record_thread:
self._record_thread.join(timeout=3.0) self._record_thread.join(timeout=3.0)
# Save final flight path image
if self._last_tracker_frame is not None: if self._last_tracker_frame is not None:
path = self.run_dir / "flight_path.png" path = self.run_dir / "flight_path.png"
cv2.imwrite(str(path), self._last_tracker_frame) cv2.imwrite(str(path), self._last_tracker_frame)
print(f"[REC] Flight path saved: {path}") print(f"[REC] Flight path saved: {path}")
# Save final camera frame
if self._last_camera_frame is not None: if self._last_camera_frame is not None:
path = self.run_dir / "camera_final.png" path = self.run_dir / "camera_final.png"
cv2.imwrite(str(path), self._last_camera_frame) cv2.imwrite(str(path), self._last_camera_frame)
# Release video writers
if self._tracker_writer: if self._tracker_writer:
self._tracker_writer.release() self._tracker_writer.release()
if self._camera_writer: if self._camera_writer:
self._camera_writer.release() self._camera_writer.release()
# Stop log capture
self.stop_logging() self.stop_logging()
self._log_file.close() self._log_file.close()
@@ -253,7 +219,6 @@ class RunRecorder:
mins = int(duration // 60) mins = int(duration // 60)
secs = int(duration % 60) secs = int(duration % 60)
# Print to original stdout since we stopped the tee
self._original_stdout.write( self._original_stdout.write(
f"\n[REC] Run recorded: {self.run_dir}\n" f"\n[REC] Run recorded: {self.run_dir}\n"
f"[REC] Duration: {mins}m {secs}s | " f"[REC] Duration: {mins}m {secs}s | "
@@ -263,8 +228,6 @@ class RunRecorder:
class _TeeWriter: class _TeeWriter:
"""Writes to both a stream and a file simultaneously."""
def __init__(self, stream, log_file): def __init__(self, stream, log_file):
self._stream = stream self._stream = stream
self._log = log_file self._log = log_file
@@ -272,9 +235,7 @@ class _TeeWriter:
def write(self, data): def write(self, data):
self._stream.write(data) self._stream.write(data)
try: try:
# Strip ANSI escape codes for the log file self._log.write(data)
clean = data
self._log.write(clean)
self._log.flush() self._log.flush()
except (ValueError, IOError): except (ValueError, IOError):
pass pass

47
src/vision/camera_processor.py
View File

@@ -1,9 +1,4 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
"""
Camera Processor - Handles Gazebo camera feeds for GPS-denied navigation.
Subscribes to gz-transport topics, processes frames, and displays via OpenCV.
Downstream modules (ArUco detector, optical flow) register as callbacks.
"""
import sys import sys
import signal import signal
@@ -13,7 +8,6 @@ import cv2
from gz.transport13 import Node from gz.transport13 import Node
from gz.msgs10.image_pb2 import Image from gz.msgs10.image_pb2 import Image
# PixelFormatType enum values (from gz.msgs10 Image proto)
PF_RGB_INT8 = 3 PF_RGB_INT8 = 3
PF_BGR_INT8 = 8 PF_BGR_INT8 = 8
PF_R_FLOAT32 = 13 PF_R_FLOAT32 = 13
@@ -35,7 +29,6 @@ class CameraProcessor:
} }
self.topics = topics self.topics = topics
for name, topic in self.topics.items(): for name, topic in self.topics.items():
self.callbacks[name] = [] self.callbacks[name] = []
ok = self.node.subscribe(Image, topic, self._make_gz_callback(name)) ok = self.node.subscribe(Image, topic, self._make_gz_callback(name))
@@ -50,19 +43,12 @@ class CameraProcessor:
def cb(msg): def cb(msg):
fmt = msg.pixel_format_type fmt = msg.pixel_format_type
if fmt == PF_RGB_INT8: if fmt == PF_RGB_INT8:
arr = np.frombuffer(msg.data, dtype=np.uint8).reshape( arr = np.frombuffer(msg.data, dtype=np.uint8).reshape(msg.height, msg.width, 3)
msg.height, msg.width, 3
)
bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
elif fmt == PF_BGR_INT8: elif fmt == PF_BGR_INT8:
arr = np.frombuffer(msg.data, dtype=np.uint8).reshape( bgr = np.frombuffer(msg.data, dtype=np.uint8).reshape(msg.height, msg.width, 3)
msg.height, msg.width, 3
)
bgr = arr
elif fmt == PF_R_FLOAT32: elif fmt == PF_R_FLOAT32:
arr = np.frombuffer(msg.data, dtype=np.float32).reshape( arr = np.frombuffer(msg.data, dtype=np.float32).reshape(msg.height, msg.width)
msg.height, msg.width
)
normalized = cv2.normalize(arr, None, 0, 255, cv2.NORM_MINMAX) normalized = cv2.normalize(arr, None, 0, 255, cv2.NORM_MINMAX)
bgr = cv2.cvtColor(normalized.astype(np.uint8), cv2.COLOR_GRAY2BGR) bgr = cv2.cvtColor(normalized.astype(np.uint8), cv2.COLOR_GRAY2BGR)
else: else:
@@ -71,10 +57,8 @@ class CameraProcessor:
processed = self.process_image(bgr) processed = self.process_image(bgr)
self.raw_frames[name] = processed.copy() self.raw_frames[name] = processed.copy()
self.frames[name] = processed self.frames[name] = processed
for fn in self.callbacks.get(name, []): for fn in self.callbacks.get(name, []):
fn(name, processed) fn(name, processed)
return cb return cb
def process_image(self, image): def process_image(self, image):
@@ -82,40 +66,26 @@ class CameraProcessor:
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB) lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
lab[:, :, 0] = clahe.apply(lab[:, :, 0]) lab[:, :, 0] = clahe.apply(lab[:, :, 0])
processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR) return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
return processed
def register_callback(self, camera_name, fn): def register_callback(self, camera_name, fn):
if camera_name in self.callbacks: if camera_name in self.callbacks:
self.callbacks[camera_name].append(fn) self.callbacks[camera_name].append(fn)
print(f"[CAM] Registered callback for {camera_name}: {fn.__name__}")
def stop(self): def stop(self):
self.running = False self.running = False
def spin(self): def spin(self):
print("[CAM] Running. Press 'q' to quit, 's' to screenshot.")
while self.running: while self.running:
for name, frame in list(self.frames.items()): for name, frame in list(self.frames.items()):
cv2.imshow(name, frame) cv2.imshow(name, frame)
if cv2.waitKey(33) & 0xFF == ord("q"):
key = cv2.waitKey(33) & 0xFF
if key == ord("q"):
break break
elif key == ord("s"):
for name, frame in self.frames.items():
fname = f"{name}_{int(time.time())}.png"
cv2.imwrite(fname, frame)
print(f"[CAM] Saved {fname}")
cv2.destroyAllWindows() cv2.destroyAllWindows()
print("[CAM] Stopped.")
def spin_headless(self): def spin_headless(self):
print("[CAM] Running headless (no GUI).")
while self.running: while self.running:
time.sleep(0.1) time.sleep(0.1)
print("[CAM] Stopped.")
def get_frame(self, camera_name): def get_frame(self, camera_name):
return self.frames.get(camera_name) return self.frames.get(camera_name)
@@ -124,7 +94,6 @@ class CameraProcessor:
def main(): def main():
cameras = "both" cameras = "both"
show_gui = True show_gui = True
if len(sys.argv) > 1: if len(sys.argv) > 1:
cameras = sys.argv[1].lower() cameras = sys.argv[1].lower()
if "--headless" in sys.argv: if "--headless" in sys.argv:
@@ -143,7 +112,6 @@ def main():
topics = all_topics topics = all_topics
proc = CameraProcessor(topics=topics, show_gui=show_gui) proc = CameraProcessor(topics=topics, show_gui=show_gui)
try: try:
import os import os
src_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) src_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@@ -155,13 +123,12 @@ def main():
def detection_overlay(camera_name, frame): def detection_overlay(camera_name, frame):
detections = detector.detect(frame) detections = detector.detect(frame)
if detections: if detections:
annotated = detector.annotate_frame(frame, detections) proc.frames[camera_name] = detector.annotate_frame(frame, detections)
proc.frames[camera_name] = annotated
for cam_name in topics: for cam_name in topics:
proc.register_callback(cam_name, detection_overlay) proc.register_callback(cam_name, detection_overlay)
except Exception as e: except Exception as e:
print(f"[CAM] ArUco detection unavailable: {e}") print(f"[CAM] Detection unavailable: {e}")
if show_gui: if show_gui:
proc.spin() proc.spin()

27
src/vision/object_detector.py
View File

@@ -32,13 +32,11 @@ H_RES = 480
class ObjectDetector: class ObjectDetector:
CAM_DEG_FOV = 110 CAM_DEG_FOV = 110
def __init__(self, marker_size=0.15, camera_matrix=None, def __init__(self, marker_size=0.15, camera_matrix=None,
aruco_dict_name="DICT_4X4_50"): aruco_dict_name="DICT_4X4_50"):
self.marker_size = marker_size self.marker_size = marker_size
if camera_matrix is not None: if camera_matrix is not None:
self.camera_matrix = np.array(camera_matrix).reshape(3, 3) self.camera_matrix = np.array(camera_matrix).reshape(3, 3)
else: else:
@@ -46,11 +44,9 @@ class ObjectDetector:
[[500.0, 0.0, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]] [[500.0, 0.0, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]]
) )
self.dist_coeffs = np.zeros(5) self.dist_coeffs = np.zeros(5)
dict_id = ARUCO_DICT.get(aruco_dict_name, cv2.aruco.DICT_4X4_50) dict_id = ARUCO_DICT.get(aruco_dict_name, cv2.aruco.DICT_4X4_50)
self.aruco_dict = cv2.aruco.getPredefinedDictionary(dict_id) self.aruco_dict = cv2.aruco.getPredefinedDictionary(dict_id)
# Support both old (<4.7) and new (>=4.7) OpenCV ArUco API
try: try:
self.aruco_params = cv2.aruco.DetectorParameters() self.aruco_params = cv2.aruco.DetectorParameters()
self._new_api = True self._new_api = True
@@ -80,17 +76,12 @@ class ObjectDetector:
self.found_markers = {} self.found_markers = {}
self.on_detection = None self.on_detection = None
print(f"[DET] ObjectDetector initialized ({aruco_dict_name})") print(f"[DET] ObjectDetector initialized ({aruco_dict_name})")
def detect(self, frame): def detect(self, frame):
if frame is None: if frame is None:
return [] return []
gray = frame if len(frame.shape) == 2 else cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
if len(frame.shape) == 2:
gray = frame
else:
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
if self._new_api: if self._new_api:
corners, ids, rejected = self.aruco_detector.detectMarkers(gray) corners, ids, rejected = self.aruco_detector.detectMarkers(gray)
@@ -110,7 +101,6 @@ class ObjectDetector:
for i, marker_id in enumerate(ids.flatten()): for i, marker_id in enumerate(ids.flatten()):
tvec = tvecs[i][0] tvec = tvecs[i][0]
center = np.mean(corners[i][0], axis=0) center = np.mean(corners[i][0], axis=0)
detection = { detection = {
"id": int(marker_id), "id": int(marker_id),
"corners": corners[i][0].tolist(), "corners": corners[i][0].tolist(),
@@ -129,39 +119,26 @@ class ObjectDetector:
if detections and self.on_detection: if detections and self.on_detection:
self.on_detection(detections) self.on_detection(detections)
return detections return detections
def annotate_frame(self, frame, detections): def annotate_frame(self, frame, detections):
annotated = frame.copy() annotated = frame.copy()
if not detections: if not detections:
return annotated return annotated
overlay = annotated.copy() overlay = annotated.copy()
for d in detections: for d in detections:
corners_arr = np.array(d["corners"], dtype=np.int32) corners_arr = np.array(d["corners"], dtype=np.int32)
# Fill marker area with purple
cv2.fillPoly(overlay, [corners_arr], (180, 0, 220)) cv2.fillPoly(overlay, [corners_arr], (180, 0, 220))
# Draw purple border
cv2.polylines(annotated, [corners_arr], True, (180, 0, 220), 3) cv2.polylines(annotated, [corners_arr], True, (180, 0, 220), 3)
cx, cy = int(d["center_px"][0]), int(d["center_px"][1])
# Center dot
cx = int(d["center_px"][0])
cy = int(d["center_px"][1])
cv2.circle(annotated, (cx, cy), 5, (255, 255, 255), -1) cv2.circle(annotated, (cx, cy), 5, (255, 255, 255), -1)
# Label with ID and distance
label = f"ID:{d['id']} d:{d['distance']:.2f}m" label = f"ID:{d['id']} d:{d['distance']:.2f}m"
cv2.putText(annotated, label, cv2.putText(annotated, label,
(int(corners_arr[0][0]), int(corners_arr[0][1]) - 10), (int(corners_arr[0][0]), int(corners_arr[0][1]) - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (180, 0, 220), 2) cv2.FONT_HERSHEY_SIMPLEX, 0.6, (180, 0, 220), 2)
# Blend overlay at 40% opacity
cv2.addWeighted(overlay, 0.4, annotated, 0.6, 0, annotated) cv2.addWeighted(overlay, 0.4, annotated, 0.6, 0, annotated)
return annotated return annotated
def get_found_markers(self): def get_found_markers(self):