Camera Aruco Tags Dectection
This commit is contained in:
308
src/control/search.py
Normal file
308
src/control/search.py
Normal file
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import math
|
||||
import time
|
||||
import numpy as np
|
||||
from scipy.stats import levy, uniform
|
||||
from time import sleep
|
||||
from enum import Enum
|
||||
|
||||
from control.uav_controller import Controller
|
||||
from vision.object_detector import ObjectDetector
|
||||
from vision.visual_servoing import VisualServoing
|
||||
from utils.helpers import distance_2d
|
||||
|
||||
|
||||
class SearchMode(Enum):
|
||||
SPIRAL = "spiral"
|
||||
LAWNMOWER = "lawnmower"
|
||||
LEVY = "levy"
|
||||
|
||||
|
||||
class Search:
|
||||
|
||||
POSITION_TOLERANCE = 1.0
|
||||
CHECK_INTERVAL = 0.5
|
||||
MAX_TRAVEL_TIME = 30.0
|
||||
CAM_FOV_METERS = 4.0
|
||||
|
||||
def __init__(self, ctrl: Controller, detector: ObjectDetector,
|
||||
camera=None, mode: str = "spiral", altitude: float = 5.0,
|
||||
actions: dict = None):
|
||||
self.ctrl = ctrl
|
||||
self.detector = detector
|
||||
self.camera = camera
|
||||
self.mode = SearchMode(mode.lower())
|
||||
self.altitude = altitude
|
||||
self.found_markers = {}
|
||||
self.running = True
|
||||
self.landed = False
|
||||
|
||||
self.actions = actions or {}
|
||||
self.land_ids = set(self.actions.get("land", []))
|
||||
|
||||
self.servoing = None
|
||||
if self.land_ids:
|
||||
target_id = list(self.land_ids)[0]
|
||||
self.servoing = VisualServoing(ctrl, target_marker_id=target_id)
|
||||
|
||||
self.spiral_max_legs = 12
|
||||
self.spiral_initial_leg = self.CAM_FOV_METERS
|
||||
self.spiral_leg_increment = self.CAM_FOV_METERS * 0.8
|
||||
|
||||
self.lawn_width = 30.0
|
||||
self.lawn_height = 30.0
|
||||
self.lawn_lane_spacing = self.CAM_FOV_METERS * 0.8
|
||||
self.lawn_lanes = 2
|
||||
|
||||
self.levy_max_steps = 20
|
||||
self.levy_field_size = 50.0
|
||||
self.levy_angle_dist = uniform(loc=-180, scale=360)
|
||||
|
||||
def configure(self, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
|
||||
def run(self):
|
||||
print(f"\n{'=' * 50}")
|
||||
print(f" SEARCH: {self.mode.value.upper()} at {self.altitude}m")
|
||||
if self.land_ids:
|
||||
print(f" Landing targets: {self.land_ids}")
|
||||
print(f"{'=' * 50}\n")
|
||||
|
||||
if self.mode == SearchMode.SPIRAL:
|
||||
return self.spiral()
|
||||
elif self.mode == SearchMode.LAWNMOWER:
|
||||
return self.lawnmower()
|
||||
elif self.mode == SearchMode.LEVY:
|
||||
return self.levy_walk()
|
||||
|
||||
def get_camera_frame(self):
|
||||
if self.camera is None:
|
||||
return None
|
||||
frame = self.camera.frames.get("downward")
|
||||
if frame is None:
|
||||
frame = self.camera.frames.get("gimbal")
|
||||
return frame
|
||||
|
||||
def check_for_markers(self):
|
||||
frame = self.get_camera_frame()
|
||||
if frame is None:
|
||||
return []
|
||||
|
||||
detections = self.detector.detect(frame)
|
||||
new_markers = []
|
||||
for d in detections:
|
||||
marker_id = d.get("id")
|
||||
if marker_id is None:
|
||||
continue
|
||||
|
||||
if marker_id not in self.found_markers:
|
||||
self.ctrl.update_state()
|
||||
pos = self.ctrl.get_local_position()
|
||||
self.found_markers[marker_id] = {
|
||||
"id": marker_id,
|
||||
"uav_position": pos.copy(),
|
||||
"distance": d.get("distance", 0),
|
||||
"timestamp": time.time(),
|
||||
}
|
||||
new_markers.append(d)
|
||||
print(f"\n[SEARCH] ArUco ID:{marker_id} detected! "
|
||||
f"distance:{d['distance']:.2f}m "
|
||||
f"UAV at ({pos['x']:.1f}, {pos['y']:.1f})")
|
||||
|
||||
if marker_id in self.land_ids:
|
||||
print(f"\n[SEARCH] Landing target ID:{marker_id} found! Starting approach.")
|
||||
self.execute_landing(detections)
|
||||
return new_markers
|
||||
|
||||
return new_markers
|
||||
|
||||
def execute_landing(self, initial_detections):
|
||||
if self.servoing is None:
|
||||
self.ctrl.land()
|
||||
self.landed = True
|
||||
self.running = False
|
||||
return
|
||||
|
||||
self.servoing.enable()
|
||||
self.servoing.process_detections(initial_detections)
|
||||
|
||||
t0 = time.time()
|
||||
timeout = 60.0
|
||||
while time.time() - t0 < timeout and self.running:
|
||||
frame = self.get_camera_frame()
|
||||
if frame is None:
|
||||
sleep(self.CHECK_INTERVAL)
|
||||
continue
|
||||
|
||||
detections = self.detector.detect(frame)
|
||||
landed = self.servoing.process_detections(detections)
|
||||
|
||||
if landed:
|
||||
print(f"\n[SEARCH] Landed on target!")
|
||||
self.landed = True
|
||||
self.running = False
|
||||
return
|
||||
|
||||
self.ctrl.update_state()
|
||||
if self.ctrl.altitude < 0.3:
|
||||
print(f"\n[SEARCH] Touchdown detected!")
|
||||
self.landed = True
|
||||
self.running = False
|
||||
return
|
||||
|
||||
sleep(0.1)
|
||||
|
||||
if not self.landed:
|
||||
print(f"\n[SEARCH] Landing approach timed out, landing at current position")
|
||||
self.ctrl.land()
|
||||
self.landed = True
|
||||
self.running = False
|
||||
|
||||
def wait_for_position(self, target_x, target_y, timeout=None):
|
||||
if timeout is None:
|
||||
timeout = self.MAX_TRAVEL_TIME
|
||||
t0 = time.time()
|
||||
while time.time() - t0 < timeout and self.running:
|
||||
self.ctrl.update_state()
|
||||
self.check_for_markers()
|
||||
if not self.running:
|
||||
return False
|
||||
pos = self.ctrl.get_local_position()
|
||||
dist = distance_2d(
|
||||
(pos['x'], pos['y']),
|
||||
(target_x, target_y)
|
||||
)
|
||||
elapsed = int(time.time() - t0)
|
||||
print(f"\r[SEARCH] Moving: {dist:.1f}m to target ({elapsed}s) "
|
||||
f"markers found: {len(self.found_markers)} ",
|
||||
end='', flush=True)
|
||||
if dist < self.POSITION_TOLERANCE:
|
||||
print()
|
||||
return True
|
||||
sleep(self.CHECK_INTERVAL)
|
||||
print()
|
||||
return False
|
||||
|
||||
def move_to_local(self, x, y):
|
||||
z = -self.altitude
|
||||
self.ctrl.move_local_ned(x, y, z)
|
||||
return self.wait_for_position(x, y)
|
||||
|
||||
def move_relative(self, dx, dy):
|
||||
self.ctrl.update_state()
|
||||
pos = self.ctrl.get_local_position()
|
||||
target_x = pos['x'] + dx
|
||||
target_y = pos['y'] + dy
|
||||
self.ctrl.move_pos_rel(dx, dy, 0)
|
||||
return self.wait_for_position(target_x, target_y)
|
||||
|
||||
def spiral(self):
|
||||
distance = self.spiral_initial_leg
|
||||
increment = self.spiral_leg_increment
|
||||
travel_x = True
|
||||
direction = 1
|
||||
|
||||
for leg in range(self.spiral_max_legs):
|
||||
if not self.running:
|
||||
break
|
||||
|
||||
self.ctrl.update_state()
|
||||
pos = self.ctrl.get_local_position()
|
||||
print(f"[SEARCH] Spiral leg {leg + 1}/{self.spiral_max_legs} "
|
||||
f"pos:({pos['x']:.1f}, {pos['y']:.1f}) "
|
||||
f"step:{distance:.1f}m markers:{len(self.found_markers)}")
|
||||
|
||||
if travel_x:
|
||||
dx = distance * direction
|
||||
self.move_relative(dx, 0)
|
||||
else:
|
||||
dy = distance * direction
|
||||
self.move_relative(0, dy)
|
||||
direction *= -1
|
||||
distance += increment
|
||||
|
||||
travel_x = not travel_x
|
||||
|
||||
print(f"[SEARCH] Spiral complete. Found {len(self.found_markers)} markers.")
|
||||
return self.found_markers
|
||||
|
||||
def lawnmower(self):
|
||||
lane_spacing = self.lawn_lane_spacing
|
||||
height = self.lawn_height
|
||||
num_lanes = max(1, int(self.lawn_width / lane_spacing))
|
||||
|
||||
self.ctrl.update_state()
|
||||
start_pos = self.ctrl.get_local_position()
|
||||
start_x = start_pos['x']
|
||||
start_y = start_pos['y']
|
||||
|
||||
print(f"[SEARCH] Lawnmower: {num_lanes} lanes, "
|
||||
f"{lane_spacing:.1f}m spacing, {height:.1f}m height")
|
||||
|
||||
for lane in range(num_lanes):
|
||||
if not self.running:
|
||||
break
|
||||
|
||||
lane_x = start_x + lane * lane_spacing
|
||||
|
||||
if lane % 2 == 0:
|
||||
target_y = start_y + height
|
||||
else:
|
||||
target_y = start_y
|
||||
|
||||
print(f"[SEARCH] Lane {lane + 1}/{num_lanes} "
|
||||
f"x:{lane_x:.1f} markers:{len(self.found_markers)}")
|
||||
|
||||
self.move_to_local(lane_x, target_y)
|
||||
|
||||
if not self.running:
|
||||
break
|
||||
|
||||
if lane < num_lanes - 1:
|
||||
next_x = start_x + (lane + 1) * lane_spacing
|
||||
self.move_to_local(next_x, target_y)
|
||||
|
||||
print(f"[SEARCH] Lawnmower complete. Found {len(self.found_markers)} markers.")
|
||||
return self.found_markers
|
||||
|
||||
def levy_walk(self):
|
||||
field_size = self.levy_field_size
|
||||
|
||||
for step in range(self.levy_max_steps):
|
||||
if not self.running:
|
||||
break
|
||||
|
||||
angle_deg = self.levy_angle_dist.rvs()
|
||||
angle_rad = math.radians(angle_deg)
|
||||
|
||||
raw_distance = levy.rvs(loc=1, scale=1)
|
||||
distance = min(max(raw_distance, 1.0), field_size)
|
||||
|
||||
dx = distance * math.cos(angle_rad)
|
||||
dy = distance * math.sin(angle_rad)
|
||||
|
||||
self.ctrl.update_state()
|
||||
pos = self.ctrl.get_local_position()
|
||||
print(f"[SEARCH] Lévy step {step + 1}/{self.levy_max_steps} "
|
||||
f"angle:{angle_deg:.0f}° dist:{distance:.1f}m "
|
||||
f"pos:({pos['x']:.1f}, {pos['y']:.1f}) "
|
||||
f"markers:{len(self.found_markers)}")
|
||||
|
||||
self.move_relative(dx, dy)
|
||||
|
||||
print(f"[SEARCH] Lévy walk complete. Found {len(self.found_markers)} markers.")
|
||||
return self.found_markers
|
||||
|
||||
def stop(self):
|
||||
self.running = False
|
||||
|
||||
def get_results(self):
|
||||
return {
|
||||
"mode": self.mode.value,
|
||||
"markers_found": len(self.found_markers),
|
||||
"markers": self.found_markers,
|
||||
"landed": self.landed,
|
||||
}
|
||||
Reference in New Issue
Block a user