simulation-generation/add_devices.py

310 lines
8.6 KiB
Python

import csv
import math
from pathlib import Path
import re
# -----------------------
# CONFIG
# -----------------------
SCALE = 0.0254
FIXED_Y = 2.4
CONVEYOR_WIDTH = 1.524
EDGE_CLEARANCE = 0.45
BEAM_RANGE_ADJUSTMENT = 0.25
SCRIPT_DIR = Path(__file__).resolve().parent
csv_candidates = list(SCRIPT_DIR.glob("*.csv"))
if not csv_candidates:
raise RuntimeError("No CSV found in script directory.")
if len(csv_candidates) > 1:
raise RuntimeError(
"Multiple CSV files found. Please keep only ONE CSV in the folder:\n"
+ "\n".join(c.name for c in csv_candidates)
)
CSV_PATH = csv_candidates[0]
SCENE_DIR = SCRIPT_DIR.parent
OUTPUT_DIR = SCRIPT_DIR / "with_devices"
OUTPUT_DIR.mkdir(exist_ok=True)
SENSOR_SCENE_PATH = "res://parts/DiffuseSensor.tscn"
SENSOR_RES_ID = "auto_sensor"
# -----------------------
# LOAD CSV DEVICES
# -----------------------
devices = []
with open(CSV_PATH, newline="", encoding="utf-8") as f:
reader = csv.DictReader(f)
for row in reader:
if row.get("record_type") == "TPE":
devices.append(row)
if not devices:
raise RuntimeError("No TPE records found in CSV.")
devices_by_conveyor = {}
for d in devices:
key = (d.get("conveyor_key") or "").strip()
if key:
devices_by_conveyor.setdefault(key, []).append(d)
# -----------------------
# SELECT SCENE
# -----------------------
scene_files = list(SCENE_DIR.glob("*.tscn"))
if not scene_files:
raise RuntimeError("No .tscn files found in parent directory.")
print("\nAvailable scenes:")
for i, s in enumerate(scene_files, 1):
print(f"[{i}] {s.name}")
choice = int(input("\nSelect scene number: ")) - 1
if choice < 0 or choice >= len(scene_files):
raise RuntimeError("Invalid selection.")
SCENE_PATH = scene_files[choice]
# -----------------------
# READ SCENE
# -----------------------
scene_text = SCENE_PATH.read_text(encoding="utf-8")
lines = scene_text.splitlines()
# -----------------------
# PARSE CONVEYORS
# -----------------------
conveyors = {}
node_name = None
has_size = False
pending_basis = None
pending_position = None
conveyor_width = CONVEYOR_WIDTH
conveyor_size = None
for line in lines:
if line.startswith("[node name="):
if node_name and pending_basis and pending_position:
conveyors[node_name] = {
"basis": pending_basis,
"has_size": has_size,
"width": conveyor_width,
"position": pending_position,
"size": conveyor_size
}
m = re.search(r'name="([^"]+)"', line)
node_name = m.group(1) if m else None
has_size = False
pending_basis = None
pending_position = None
conveyor_width = CONVEYOR_WIDTH
conveyor_size = None
continue
if node_name and line.strip().startswith("transform = Transform3D"):
vals = [float(v) for v in re.search(r"\(([^)]+)\)", line).group(1).split(",")]
pending_basis = (vals[0], vals[2], vals[6], vals[8])
pending_position = (vals[9], vals[10], vals[11])
if node_name and line.strip().startswith("size = Vector3"):
has_size = True
conveyor_size = tuple(
float(v) for v in re.search(r"\(([^)]+)\)", line).group(1).split(",")
)
if node_name and line.strip().startswith("conveyor_width ="):
conveyor_width = float(line.split("=")[1])
if node_name and pending_basis and pending_position:
conveyors[node_name] = {
"basis": pending_basis,
"has_size": has_size,
"width": conveyor_width,
"position": pending_position,
"size": conveyor_size
}
# -----------------------
# ENSURE SENSOR RESOURCE
# -----------------------
if SENSOR_SCENE_PATH not in scene_text:
idx = max(i for i, l in enumerate(lines) if l.startswith("[ext_resource"))
lines.insert(idx + 1,
f'[ext_resource type="PackedScene" path="{SENSOR_SCENE_PATH}" id="{SENSOR_RES_ID}"]'
)
scene_text = "\n".join(lines) + "\n"
# -----------------------
# HELPERS
# -----------------------
def yaw_from_x_axis(xx, xz):
return math.atan2(-xz, xx)
def yaw_from_z_axis(zx, zz):
return math.atan2(zx, zz)
def compute_sensor_position(info, yaw, idx, total):
px, _, pz = info["position"]
if not info["size"]:
return None, None
length = info["size"][0]
if total == 1:
off = length / 2
elif total == 2:
off = (-length / 2) if idx == 0 else (length / 2)
else:
return None, None
return (
px + math.cos(yaw) * off,
pz + math.sin(yaw) * off
)
def godot_yaw_from_cad(cad_deg: float) -> float:
"""
AutoCAD → Godot rotation conversion
AutoCAD:
0° = up
90° = right
180° = down
270° = left
clockwise
Godot:
0° = +X
90° = -Z
CCW
With Y→-Z flip already applied in coordinates.
"""
return math.radians((270 - cad_deg) % 360)
def compute_edge_offset(sensor_yaw, width):
"""
Move sensor to the edge in the direction perpendicular to where it's pointing.
The sensor beam points in the sensor_yaw direction, so we offset perpendicular to that.
"""
# Perpendicular to sensor's facing direction (90° to the right of where it points)
side = sensor_yaw - math.pi / 2
d = (width / 2) + EDGE_CLEARANCE
return math.cos(side) * d, math.sin(side) * d
def sensor_yaw_from_csv(info, conveyor_yaw, csv_x, csv_z):
cx, _, cz = info["position"]
# right perpendicular to conveyor forward
rx = math.cos(conveyor_yaw - math.pi / 2)
rz = math.sin(conveyor_yaw - math.pi / 2)
vx = csv_x - cx
vz = csv_z - cz
side = vx * rx + vz * rz
if side >= 0:
return conveyor_yaw + math.pi / 2
return conveyor_yaw - math.pi / 2
# -----------------------
# APPEND DEVICES
# -----------------------
node_blocks = []
for key in sorted(devices_by_conveyor.keys()):
devs = devices_by_conveyor[key]
if key not in conveyors:
print(f"⚠ Conveyor not found: {key}")
continue
info = conveyors[key]
xx, xz, zx, zz = info["basis"]
width = info["width"]
yaw = yaw_from_x_axis(xx, xz) if info["has_size"] else yaw_from_z_axis(zx, zz)
for i, d in enumerate(devs):
name = f"{key}_TPE{i+1}"
tag = f"{name}_OIP"
max_range = width + BEAM_RANGE_ADJUSTMENT
x, z = compute_sensor_position(info, yaw, i, len(devs))
LONGITUDINAL_OFFSET = 0.15
if x is not None:
if len(devs) == 1:
x -= math.cos(yaw) * LONGITUDINAL_OFFSET
z -= math.sin(yaw) * LONGITUDINAL_OFFSET
elif len(devs) == 2:
if i == 0:
x += math.cos(yaw) * LONGITUDINAL_OFFSET
z += math.sin(yaw) * LONGITUDINAL_OFFSET
else:
x -= math.cos(yaw) * LONGITUDINAL_OFFSET
z -= math.sin(yaw) * LONGITUDINAL_OFFSET
else:
x = float(d["tpe_x"]) * SCALE
z = -float(d["tpe_y"]) * SCALE
csv_x = float(d["tpe_x"]) * SCALE
csv_z = -float(d["tpe_y"]) * SCALE
if "tpe_rotation" in d and d["tpe_rotation"] not in ("", None):
tpe_rot = float(d["tpe_rotation"])
else:
tpe_rot = 0.0
if "tpe_block_rotation" in d and d["tpe_block_rotation"] not in ("", None):
blk_rot = float(d["tpe_block_rotation"])
else:
blk_rot = 0.0
# If we have ANY rotation data (block or TPE), use it
if tpe_rot != 0.0 or blk_rot != 0.0:
# Combine both rotations in CAD space
cad_world_rot = (blk_rot + tpe_rot) % 360.0
sensor_yaw = godot_yaw_from_cad(cad_world_rot)
else:
# fallback ONLY if both rotations are missing/zero
sensor_yaw = yaw + math.pi
ox, oz = compute_edge_offset(sensor_yaw, width)
x += ox
z += oz
c, s = math.cos(sensor_yaw), math.sin(sensor_yaw)
transform = (
f"Transform3D({c:.6f},0,{-s:.6f},0,1,0,{s:.6f},0,{c:.6f},"
f"{x:.6f},{FIXED_Y:.6f},{z:.6f})"
)
node_blocks.append(
f'\n[node name="{name}" parent="." instance=ExtResource("{SENSOR_RES_ID}")]\n'
f"transform = {transform}\n"
f"max_range = {max_range:.3f}\n"
f"enable_comms = true\n"
f'tag_name = "{tag}"\n'
)
scene_text += "".join(node_blocks)
# -----------------------
# WRITE OUTPUT
# -----------------------
out_path = OUTPUT_DIR / f"{SCENE_PATH.stem}_devices{SCENE_PATH.suffix}"
out_path.write_text(scene_text, encoding="utf-8")
print(f"\n{len(node_blocks)} sensors added → {out_path}")