227 lines
8.1 KiB
Python
Executable File
227 lines
8.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import os
|
|
import math
|
|
import random
|
|
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")
|
|
|
|
FEET_TO_METERS = 0.3048
|
|
|
|
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_obstacle_positions(obstacle_cfg, geofence_points, exclusion_zones):
|
|
if not obstacle_cfg.get("enabled", False):
|
|
return []
|
|
|
|
min_spacing = obstacle_cfg.get("min_spacing_ft", 5.0) * FEET_TO_METERS
|
|
obstacle_types = obstacle_cfg.get("types", [])
|
|
|
|
if not obstacle_types:
|
|
return []
|
|
|
|
if geofence_points and len(geofence_points) >= 3:
|
|
xs = [p[0] for p in geofence_points]
|
|
ys = [p[1] for p in geofence_points]
|
|
margin = 2.0
|
|
x_min, x_max = min(xs) + margin, max(xs) - margin
|
|
y_min, y_max = min(ys) + margin, max(ys) - margin
|
|
else:
|
|
x_min, x_max = -12.0, 12.0
|
|
y_min, y_max = -12.0, 12.0
|
|
|
|
placements = []
|
|
|
|
for obs_type in obstacle_types:
|
|
name = obs_type.get("name", "box_small")
|
|
height = obs_type.get("height", 0.3)
|
|
count_min = obs_type.get("count_min", 1)
|
|
count_max = obs_type.get("count_max", 3)
|
|
count = random.randint(count_min, count_max)
|
|
|
|
for _ in range(count):
|
|
placed = False
|
|
for _attempt in range(200):
|
|
x = random.uniform(x_min, x_max)
|
|
y = random.uniform(y_min, y_max)
|
|
|
|
too_close = False
|
|
for ex, ey, er in exclusion_zones:
|
|
if math.sqrt((x - ex)**2 + (y - ey)**2) < er:
|
|
too_close = True
|
|
break
|
|
|
|
if too_close:
|
|
continue
|
|
|
|
for px, py, _, _ in placements:
|
|
if math.sqrt((x - px)**2 + (y - py)**2) < min_spacing:
|
|
too_close = True
|
|
break
|
|
|
|
if too_close:
|
|
continue
|
|
|
|
yaw = random.uniform(0, 2 * math.pi)
|
|
placements.append((x, y, name, height))
|
|
placed = True
|
|
break
|
|
|
|
if not placed:
|
|
print(f"[WARN] Could not place {name} with minimum spacing")
|
|
|
|
return placements
|
|
|
|
|
|
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_ids = marker.get("target_ids", [1])
|
|
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")
|
|
|
|
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"
|
|
|
|
for old_model in list(world.findall('model')):
|
|
name_attr = old_model.get('name', '')
|
|
if name_attr.startswith('obstacle_'):
|
|
world.remove(old_model)
|
|
for old_include in list(world.findall('include')):
|
|
name_el = old_include.find('name')
|
|
if name_el is not None and name_el.text and name_el.text.startswith('obstacle_'):
|
|
world.remove(old_include)
|
|
|
|
obstacle_cfg = search_cfg.get("obstacles", {})
|
|
geofence_points = []
|
|
if geofence_cfg.get("enabled", False):
|
|
geofence_points = [(float(p[0]), float(p[1]))
|
|
for p in geofence_cfg.get("points", [])]
|
|
|
|
exclusion_zones = [
|
|
(ugv_x, ugv_y, 2.0),
|
|
(target_x, target_y, 2.0),
|
|
]
|
|
|
|
if obstacle_cfg.get("enabled", False):
|
|
placements = generate_obstacle_positions(
|
|
obstacle_cfg, geofence_points, exclusion_zones)
|
|
|
|
for idx, (ox, oy, obs_name, obs_height) in enumerate(placements):
|
|
oz = obs_height / 2.0
|
|
yaw = random.uniform(0, 2 * math.pi)
|
|
|
|
include_el = ET.SubElement(world, "include")
|
|
ET.SubElement(include_el, "uri").text = f"model://{obs_name}"
|
|
ET.SubElement(include_el, "name").text = f"obstacle_{obs_name}_{idx}"
|
|
ET.SubElement(include_el, "pose").text = f"{ox:.3f} {oy:.3f} {oz:.3f} 0 0 {yaw:.3f}"
|
|
|
|
print(f"[INFO] Placed {len(placements)} obstacles "
|
|
f"(min spacing: {obstacle_cfg.get('min_spacing_ft', 5.0)} ft)")
|
|
for ox, oy, obs_name, obs_height in placements:
|
|
print(f"[INFO] {obs_name} at ({ox:.1f}, {oy:.1f})")
|
|
|
|
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()
|