Aruco Tag Landing and Flight Tracker
This commit is contained in:
211
scripts/generate_tags.py
Normal file
211
scripts/generate_tags.py
Normal file
@@ -0,0 +1,211 @@
|
||||
#!/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 sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
PROJECT_DIR = Path(__file__).resolve().parent.parent
|
||||
TAGS_DIR = PROJECT_DIR / "worlds" / "tags"
|
||||
CONFIG_DIR = PROJECT_DIR / "config"
|
||||
|
||||
ARUCO_DICTS = {
|
||||
"DICT_4X4_50": cv2.aruco.DICT_4X4_50,
|
||||
"DICT_4X4_100": cv2.aruco.DICT_4X4_100,
|
||||
"DICT_4X4_250": cv2.aruco.DICT_4X4_250,
|
||||
"DICT_4X4_1000": cv2.aruco.DICT_4X4_1000,
|
||||
"DICT_5X5_50": cv2.aruco.DICT_5X5_50,
|
||||
"DICT_5X5_100": cv2.aruco.DICT_5X5_100,
|
||||
"DICT_5X5_250": cv2.aruco.DICT_5X5_250,
|
||||
"DICT_5X5_1000": cv2.aruco.DICT_5X5_1000,
|
||||
"DICT_6X6_50": cv2.aruco.DICT_6X6_50,
|
||||
"DICT_6X6_100": cv2.aruco.DICT_6X6_100,
|
||||
"DICT_6X6_250": cv2.aruco.DICT_6X6_250,
|
||||
"DICT_6X6_1000": cv2.aruco.DICT_6X6_1000,
|
||||
"DICT_7X7_50": cv2.aruco.DICT_7X7_50,
|
||||
"DICT_7X7_100": cv2.aruco.DICT_7X7_100,
|
||||
"DICT_7X7_250": cv2.aruco.DICT_7X7_250,
|
||||
"DICT_7X7_1000": cv2.aruco.DICT_7X7_1000,
|
||||
}
|
||||
|
||||
DICT_SIZES = {
|
||||
"DICT_4X4_50": 50,
|
||||
"DICT_4X4_100": 100,
|
||||
"DICT_4X4_250": 250,
|
||||
"DICT_4X4_1000": 1000,
|
||||
"DICT_5X5_50": 50,
|
||||
"DICT_5X5_100": 100,
|
||||
"DICT_5X5_250": 250,
|
||||
"DICT_5X5_1000": 1000,
|
||||
"DICT_6X6_50": 50,
|
||||
"DICT_6X6_100": 100,
|
||||
"DICT_6X6_250": 250,
|
||||
"DICT_6X6_1000": 1000,
|
||||
"DICT_7X7_50": 50,
|
||||
"DICT_7X7_100": 100,
|
||||
"DICT_7X7_250": 250,
|
||||
"DICT_7X7_1000": 1000,
|
||||
}
|
||||
|
||||
|
||||
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]
|
||||
aruco_dict = cv2.aruco.getPredefinedDictionary(dict_id)
|
||||
|
||||
try:
|
||||
marker_img = cv2.aruco.generateImageMarker(aruco_dict, marker_id, pixel_size)
|
||||
except AttributeError:
|
||||
marker_img = cv2.aruco.drawMarker(aruco_dict, marker_id, pixel_size)
|
||||
|
||||
border_px = pixel_size // 6
|
||||
padded = np.ones(
|
||||
(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
|
||||
|
||||
return padded
|
||||
|
||||
|
||||
def get_landing_id():
|
||||
"""Read the landing target ID from search.yaml config."""
|
||||
search_cfg = CONFIG_DIR / "search.yaml"
|
||||
if search_cfg.exists():
|
||||
with open(search_cfg) as f:
|
||||
cfg = yaml.safe_load(f) or {}
|
||||
land_ids = cfg.get("actions", {}).get("land", [])
|
||||
if land_ids:
|
||||
return land_ids[0]
|
||||
return 0
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate ArUco marker tags for Gazebo simulation"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dict", default="DICT_4X4_50",
|
||||
choices=list(ARUCO_DICTS.keys()),
|
||||
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()
|
||||
|
||||
TAGS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
dict_name = args.dict
|
||||
landing_id = get_landing_id()
|
||||
|
||||
if args.all:
|
||||
max_id = DICT_SIZES.get(dict_name, 50)
|
||||
ids = list(range(max_id))
|
||||
elif args.ids is not None:
|
||||
ids = args.ids
|
||||
else:
|
||||
ids = [landing_id]
|
||||
|
||||
print(f"Generating {len(ids)} ArUco marker(s)")
|
||||
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:
|
||||
marker_img = generate_marker(dict_name, marker_id, args.size)
|
||||
|
||||
filename = f"aruco_{dict_name}_{marker_id}.{args.format}"
|
||||
filepath = TAGS_DIR / 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:
|
||||
land_path = TAGS_DIR / f"land_tag.{args.format}"
|
||||
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...")
|
||||
aruco_dict = cv2.aruco.getPredefinedDictionary(ARUCO_DICTS[dict_name])
|
||||
try:
|
||||
params = cv2.aruco.DetectorParameters()
|
||||
except AttributeError:
|
||||
params = cv2.aruco.DetectorParameters_create()
|
||||
params.minMarkerPerimeterRate = 0.01
|
||||
try:
|
||||
params.cornerRefinementMethod = cv2.aruco.CORNER_REFINE_SUBPIX
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
detector = cv2.aruco.ArucoDetector(aruco_dict, params)
|
||||
_detect = lambda img: detector.detectMarkers(img)
|
||||
except AttributeError:
|
||||
_detect = lambda img: cv2.aruco.detectMarkers(img, aruco_dict, parameters=params)
|
||||
|
||||
ok = 0
|
||||
fail = 0
|
||||
for marker_id in ids:
|
||||
filename = f"aruco_{dict_name}_{marker_id}.{args.format}"
|
||||
filepath = TAGS_DIR / filename
|
||||
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():
|
||||
ok += 1
|
||||
else:
|
||||
print(f" WARNING: {filename} failed detection check!")
|
||||
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__":
|
||||
main()
|
||||
Reference in New Issue
Block a user