Files
simulation/scripts/generate_world.py

122 lines
4.3 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
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:
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")
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()