scripts/survey_mapping/run_survey.py

#!/usr/bin/env python3
"""Move through survey poses, capture one Zivid cloud per pose, and record a bag."""

import argparse
import json
import math
import os
import signal
import subprocess
import sys
import time
from pathlib import Path

import rclpy
import yaml
from beambot_interfaces.action import MoveToAction
from rcl_interfaces.msg import Parameter, ParameterType, ParameterValue
from rcl_interfaces.srv import SetParameters
from rclpy.action import ActionClient
from rclpy.node import Node
from rclpy.qos import DurabilityPolicy, HistoryPolicy, QoSProfile, ReliabilityPolicy
from sensor_msgs.msg import PointCloud2
from std_srvs.srv import Trigger

MOVE_ACTION = "beambot_moveto"
CAPTURE_SERVICE = "/capture"
CLOUD_TOPIC = "/points/xyzrgba"
CAMERA_NODE = "zivid_camera"
SETTINGS_FILE = "manufacturing_specular.yml"
BAG_TOPICS = [CLOUD_TOPIC, "/tf", "/tf_static", "/joint_states"]

ZIVID_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    durability=DurabilityPolicy.VOLATILE,
    depth=1,
)

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_POSES = SCRIPT_DIR / "survey_poses.yaml"
DEFAULT_QOS_OVERRIDE = SCRIPT_DIR / "tf_qos_override.yaml"


def default_settings_file():
    try:
        from ament_index_python.packages import get_package_share_path

        installed = get_package_share_path("beambot") / "config" / "zivid" / SETTINGS_FILE
        if installed.is_file():
            return installed
    except Exception:  # Package may not be built or sourced yet.
        pass
    return SCRIPT_DIR.parents[1] / "src" / "beambot" / "config" / "zivid" / SETTINGS_FILE


def load_poses(path):
    data = yaml.safe_load(Path(path).read_text()) or {}
    if not isinstance(data, dict):
        raise ValueError(f"{path} must contain a YAML mapping")

    poses = {}
    for name, joints in data.items():
        if (
            not isinstance(name, str)
            or not isinstance(joints, list)
            or len(joints) != 6
        ):
            print(f"WARN: skipping invalid pose {name!r}", file=sys.stderr)
            continue
        try:
            values = [float(value) for value in joints]
            if not all(math.isfinite(value) for value in values):
                raise ValueError
            poses[name] = values
        except (TypeError, ValueError):
            print(f"WARN: skipping invalid pose {name!r}", file=sys.stderr)
    return poses


class SurveyRunner(Node):
    def __init__(self, args):
        super().__init__("survey_runner")
        self.args = args
        self.move_client = ActionClient(self, MoveToAction, MOVE_ACTION)
        self.capture_client = self.create_client(Trigger, CAPTURE_SERVICE)
        self.cloud_count = 0
        self.create_subscription(PointCloud2, CLOUD_TOPIC, self._on_cloud, ZIVID_QOS)

    def _on_cloud(self, _message):
        self.cloud_count += 1

    def wait_for_servers(self):
        self.get_logger().info(f"Waiting for {MOVE_ACTION}")
        if not self.move_client.wait_for_server(timeout_sec=15):
            self.get_logger().error(f"{MOVE_ACTION} is unavailable")
            return False
        if not self.args.dry_run and not self.capture_client.wait_for_service(
            timeout_sec=15
        ):
            self.get_logger().error(f"{CAPTURE_SERVICE} is unavailable")
            return False
        return True

    def apply_settings(self):
        path = Path(self.args.settings_file).resolve()
        if not path.is_file():
            self.get_logger().warning(f"Capture settings not found: {path}")
            return

        client = self.create_client(SetParameters, f"/{CAMERA_NODE}/set_parameters")
        if not client.wait_for_service(timeout_sec=10):
            self.get_logger().warning(
                f"Cannot configure {CAMERA_NODE}; keeping current settings"
            )
            return

        request = SetParameters.Request()
        request.parameters = [
            Parameter(
                name="settings_file_path",
                value=ParameterValue(
                    type=ParameterType.PARAMETER_STRING,
                    string_value=str(path),
                ),
            ),
            Parameter(
                name="settings_yaml",
                value=ParameterValue(
                    type=ParameterType.PARAMETER_STRING, string_value=""
                ),
            ),
        ]
        future = client.call_async(request)
        rclpy.spin_until_future_complete(self, future, timeout_sec=10)
        result = future.result() if future.done() else None
        if result and all(item.successful for item in result.results):
            self.get_logger().info(f"Applied capture settings: {path.name}")
        else:
            self.get_logger().warning(
                "Could not apply capture settings; keeping current settings"
            )

    def move_to(self, name, joints):
        goal = MoveToAction.Goal()
        goal.target = name
        goal.planning_type = "joint"
        goal.poses_json = json.dumps({name: joints})

        self.get_logger().info(f"Moving to {name}")
        future = self.move_client.send_goal_async(goal)
        rclpy.spin_until_future_complete(self, future, timeout_sec=15)
        handle = future.result() if future.done() else None
        if handle is None or not handle.accepted:
            self.get_logger().error(f"Move to {name} was rejected")
            return False

        result_future = handle.get_result_async()
        rclpy.spin_until_future_complete(
            self, result_future, timeout_sec=self.args.move_timeout
        )
        if not result_future.done():
            self.get_logger().error(f"Move to {name} timed out; cancelling")
            cancel_future = handle.cancel_goal_async()
            rclpy.spin_until_future_complete(self, cancel_future, timeout_sec=5)
            return False

        wrapped_result = result_future.result()
        result = wrapped_result.result if wrapped_result else None
        if result is None or not result.success:
            message = result.error_message if result else "no result"
            self.get_logger().error(f"Move to {name} failed: {message}")
            return False
        return True

    def capture(self, name):
        count_before = self.cloud_count
        future = self.capture_client.call_async(Trigger.Request())
        rclpy.spin_until_future_complete(
            self, future, timeout_sec=self.args.capture_timeout
        )
        result = future.result() if future.done() else None
        if result is None:
            self.get_logger().error(f"Capture timed out at {name}")
            return False
        if not result.success:
            self.get_logger().warning(
                f"Capture reported failure at {name}: {result.message}"
            )

        deadline = time.monotonic() + self.args.cloud_timeout
        while self.cloud_count == count_before:
            if time.monotonic() >= deadline:
                self.get_logger().error(f"No cloud received at {name}")
                return False
            rclpy.spin_once(self, timeout_sec=0.1)
        return True

    def run(self, poses):
        completed = 0
        total = len(poses)
        for index, (name, joints) in enumerate(poses.items(), start=1):
            self.get_logger().info(f"[{index}/{total}] {name}")
            if not self.move_to(name, joints):
                return False
            if self.args.settle > 0:
                time.sleep(self.args.settle)
            if self.args.dry_run or self.capture(name):
                completed += 1

        self.get_logger().info(f"Survey complete: {completed}/{total}")
        return completed == total


def start_bag(output, qos_override):
    command = ["ros2", "bag", "record", "-o", output]
    if qos_override:
        command += ["--qos-profile-overrides-path", qos_override]
    command += BAG_TOPICS
    print(f"[bag] {' '.join(command)}")
    return subprocess.Popen(command, start_new_session=True)


def stop_bag(process):
    if process is None or process.poll() is not None:
        return
    print("[bag] finalizing")
    try:
        os.killpg(process.pid, signal.SIGINT)
        process.wait(timeout=15)
    except ProcessLookupError:
        pass
    except subprocess.TimeoutExpired:
        process.terminate()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--poses", default=str(DEFAULT_POSES))
    parser.add_argument("--bag-out", default="survey_session")
    parser.add_argument("--no-bag", action="store_true")
    parser.add_argument("--settle", type=float, default=1.5)
    parser.add_argument("--move-timeout", type=float, default=120)
    parser.add_argument("--capture-timeout", type=float, default=25)
    parser.add_argument("--cloud-timeout", type=float, default=20)
    parser.add_argument("--settings-file", default=str(default_settings_file()))
    parser.add_argument("--no-set-settings", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--max-poses", type=int, default=0)
    parser.add_argument("--qos-overrides", default=str(DEFAULT_QOS_OVERRIDE))
    args, _ = parser.parse_known_args()

    if (
        args.settle < 0
        or min(args.move_timeout, args.capture_timeout, args.cloud_timeout) <= 0
        or args.max_poses < 0
    ):
        parser.error(
            "timeouts must be positive; --settle and --max-poses cannot be negative"
        )

    try:
        poses = load_poses(args.poses)
    except (OSError, ValueError, yaml.YAMLError) as exc:
        print(f"ERROR: cannot load poses: {exc}", file=sys.stderr)
        return 1
    if not poses:
        print("ERROR: no valid survey poses", file=sys.stderr)
        return 1
    if args.max_poses > 0:
        poses = dict(list(poses.items())[: args.max_poses])

    if (
        not args.no_bag
        and not args.dry_run
        and args.qos_overrides
        and not Path(args.qos_overrides).is_file()
    ):
        print(f"ERROR: QoS override not found: {args.qos_overrides}", file=sys.stderr)
        return 1
    if not args.no_bag and not args.dry_run and Path(args.bag_out).exists():
        print(f"ERROR: bag output already exists: {args.bag_out}", file=sys.stderr)
        return 1

    rclpy.init()
    node = SurveyRunner(args)
    bag_process = None
    success = False
    try:
        if not node.wait_for_servers():
            return 1
        if not args.dry_run and not args.no_set_settings:
            node.apply_settings()
        if not args.dry_run and not args.no_bag:
            bag_process = start_bag(args.bag_out, args.qos_overrides)
            time.sleep(2)
            if bag_process.poll() is not None:
                print("ERROR: bag recorder failed to start", file=sys.stderr)
                return 1
        success = node.run(poses)
    except KeyboardInterrupt:
        print("\nSurvey interrupted", file=sys.stderr)
    finally:
        stop_bag(bag_process)
        node.destroy_node()
        rclpy.shutdown()

    if success and not args.dry_run and not args.no_bag:
        print(f"Bag written to {args.bag_out}")
    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())