122 lines
4.6 KiB
Python
Executable File
122 lines
4.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import os
|
|
import yaml
|
|
import sys
|
|
import xml.etree.ElementTree as ET
|
|
|
|
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
PROJECT_DIR = os.path.dirname(SCRIPT_DIR)
|
|
CONFIG_DIR = os.path.join(PROJECT_DIR, "config")
|
|
WORLD_DIR = os.path.join(PROJECT_DIR, "worlds")
|
|
|
|
def load_config(name):
|
|
path = os.path.join(CONFIG_DIR, name)
|
|
if not os.path.exists(path):
|
|
return {}
|
|
with open(path, 'r') as f:
|
|
return yaml.safe_load(f) or {}
|
|
|
|
def generate_world(base_filename="uav_ugv_search_base.sdf", output_filename="uav_ugv_search.sdf"):
|
|
ugv_cfg = load_config("ugv.yaml")
|
|
search_cfg = load_config("search.yaml")
|
|
|
|
ugv_x = ugv_cfg.get("position", {}).get("x", 0.0)
|
|
ugv_y = ugv_cfg.get("position", {}).get("y", 0.0)
|
|
|
|
uav_x = ugv_x
|
|
uav_y = ugv_y
|
|
uav_z = 0.40 # Start on top of UGV (0.18 top + 0.195 legs)
|
|
|
|
marker = search_cfg.get("marker", {})
|
|
target_pos = marker.get("target_position", [8.0, -6.0])
|
|
target_x = target_pos[0]
|
|
target_y = target_pos[1]
|
|
|
|
base_path = os.path.join(WORLD_DIR, base_filename)
|
|
output_path = os.path.join(WORLD_DIR, output_filename)
|
|
|
|
if not os.path.exists(base_path):
|
|
print(f"[ERROR] Base world file not found: {base_path}")
|
|
sys.exit(1)
|
|
|
|
try:
|
|
tree = ET.parse(base_path)
|
|
root = tree.getroot()
|
|
|
|
world = root.find('world')
|
|
if world is None:
|
|
print("[ERROR] No <world> tag found in SDF.")
|
|
sys.exit(1)
|
|
|
|
for include in world.findall('include'):
|
|
uri = include.find('uri')
|
|
if uri is not None and uri.text == 'model://iris_with_gimbal':
|
|
pose = include.find('pose')
|
|
if pose is not None:
|
|
pose.text = f"{uav_x} {uav_y} {uav_z} 0 0 90"
|
|
|
|
name = include.find('name')
|
|
if name is not None and name.text == 'ugv':
|
|
pose = include.find('pose')
|
|
if pose is not None:
|
|
pose.text = f"{ugv_x} {ugv_y} 0 0 0 0"
|
|
|
|
for model in world.findall('model'):
|
|
name_attr = model.get('name')
|
|
if name_attr == 'target_tag_1':
|
|
pose = model.find('pose')
|
|
if pose is not None:
|
|
pose.text = f"{target_x} {target_y} 0.005 0 0 0"
|
|
|
|
geofence_cfg = search_cfg.get("geofence", {})
|
|
if geofence_cfg.get("enabled", False):
|
|
points = geofence_cfg.get("points", [])
|
|
if len(points) >= 3:
|
|
# Remove old geofence visual if it exists
|
|
for old_model in world.findall('model'):
|
|
if old_model.get('name') == 'geofence_visual':
|
|
world.remove(old_model)
|
|
|
|
gf_model = ET.SubElement(world, "model", name="geofence_visual")
|
|
ET.SubElement(gf_model, "static").text = "true"
|
|
link = ET.SubElement(gf_model, "link", name="link")
|
|
|
|
import math
|
|
for i in range(len(points)):
|
|
p1 = points[i]
|
|
p2 = points[(i + 1) % len(points)]
|
|
x1, y1 = float(p1[0]), float(p1[1])
|
|
x2, y2 = float(p2[0]), float(p2[1])
|
|
|
|
dx = x2 - x1
|
|
dy = y2 - y1
|
|
length = math.sqrt(dx*dx + dy*dy)
|
|
cx = x1 + dx / 2.0
|
|
cy = y1 + dy / 2.0
|
|
yaw = math.atan2(dy, dx)
|
|
|
|
visual = ET.SubElement(link, "visual", name=f"edge_{i}")
|
|
ET.SubElement(visual, "pose").text = f"{cx} {cy} 0.01 0 0 {yaw}"
|
|
|
|
geometry = ET.SubElement(visual, "geometry")
|
|
box = ET.SubElement(geometry, "box")
|
|
# size is Length(X), Width(Y), Thickness(Z)
|
|
ET.SubElement(box, "size").text = f"{length} 0.2 0.02"
|
|
|
|
material = ET.SubElement(visual, "material")
|
|
ET.SubElement(material, "ambient").text = "1 0 0 1"
|
|
ET.SubElement(material, "diffuse").text = "1 0 0 1"
|
|
ET.SubElement(material, "emissive").text = "0.8 0 0 0.5"
|
|
|
|
tree.write(output_path, encoding="utf-8", xml_declaration=True)
|
|
print(f"[INFO] Generated world file: {output_path}")
|
|
print(f"[INFO] UGV set to ({ugv_x}, {ugv_y})")
|
|
print(f"[INFO] Target Marker set to ({target_x}, {target_y})")
|
|
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to parse or write XML: {e}")
|
|
sys.exit(1)
|
|
|
|
if __name__ == "__main__":
|
|
generate_world()
|