806 lines
32 KiB
Python
806 lines
32 KiB
Python
#!/usr/bin/env python3
|
||
"""Модуль В: распознать деталь, взять ARM95 и вернуть на полку."""
|
||
|
||
import argparse
|
||
import json
|
||
import math
|
||
import os
|
||
import sys
|
||
import threading
|
||
import time
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from tempfile import NamedTemporaryFile
|
||
|
||
|
||
HERE = Path(__file__).resolve().parent
|
||
TARGETS = {"1": "hammer", "2": "wrench", "3": "pliers"}
|
||
|
||
|
||
class EventLog:
|
||
def __init__(self):
|
||
directory = HERE / "logs"
|
||
directory.mkdir(parents=True, exist_ok=True)
|
||
self.path = directory / f"module3_{datetime.now():%Y%m%d_%H%M%S}.log"
|
||
|
||
def write(self, event, **fields):
|
||
record = {"time": datetime.now().isoformat(timespec="milliseconds"), "event": event}
|
||
record.update(fields)
|
||
line = json.dumps(record, ensure_ascii=False)
|
||
print(line, flush=True)
|
||
with self.path.open("a", encoding="utf-8") as stream:
|
||
stream.write(line + "\n")
|
||
|
||
def image_path(self, suffix):
|
||
return self.path.with_name(f"{self.path.stem}_{suffix}.jpg")
|
||
|
||
|
||
def local_path(value):
|
||
path = Path(value).expanduser()
|
||
return path if path.is_absolute() else HERE / path
|
||
|
||
|
||
def target_name(value):
|
||
value = value.strip().lower()
|
||
return TARGETS.get(value, value)
|
||
|
||
|
||
def arguments():
|
||
parser = argparse.ArgumentParser(description=__doc__)
|
||
parser.add_argument("--target", help="1/hammer, 2/wrench или 3/pliers")
|
||
parser.add_argument("--weights", default="models/latest.pt")
|
||
parser.add_argument("--sim", action="store_true", help="Камера/TF и часы Webots")
|
||
parser.add_argument("--topic", help="ROS-топик камеры; по умолчанию выбирается по --sim")
|
||
parser.add_argument("--camera-info-topic")
|
||
parser.add_argument(
|
||
"--compressed",
|
||
action="store_true",
|
||
help="Топик имеет тип sensor_msgs/CompressedImage, а не Image",
|
||
)
|
||
parser.add_argument("--base-frame", default="Base_link")
|
||
parser.add_argument("--fx", type=float, help="Фокус X, если CameraInfo отсутствует")
|
||
parser.add_argument("--fy", type=float, help="Фокус Y, если CameraInfo отсутствует")
|
||
parser.add_argument("--cx", type=float, help="Главная точка X")
|
||
parser.add_argument("--cy", type=float, help="Главная точка Y")
|
||
parser.add_argument("--conf", type=float, default=0.20)
|
||
parser.add_argument("--imgsz", type=int, default=512)
|
||
parser.add_argument("--device", default="cpu")
|
||
parser.add_argument("--timeout", type=float, default=60.0)
|
||
parser.add_argument(
|
||
"--plane-z",
|
||
type=float,
|
||
default=0.128,
|
||
help="Высота верхней грани детали относительно Base_link, м",
|
||
)
|
||
parser.add_argument("--pick-z", type=float, default=0.210)
|
||
parser.add_argument("--approach-z", type=float, default=0.38)
|
||
parser.add_argument("--drop-x", type=float, help="Точка возврата; по умолчанию исходная")
|
||
parser.add_argument("--drop-y", type=float, help="Точка возврата; по умолчанию исходная")
|
||
parser.add_argument("--no-window", action="store_true", help="Не открывать видеопоток")
|
||
parser.add_argument("--auto-start", action="store_true", help="Не ждать клавишу G/Enter")
|
||
parser.add_argument("--dry-run", action="store_true", help="Только распознавание и координаты")
|
||
parser.add_argument("--plan-only", action="store_true", help="Спланировать подход, не двигать руку")
|
||
args = parser.parse_args()
|
||
|
||
if args.topic is None:
|
||
args.topic = (
|
||
"/RMC1/arm95/camera_gripper/image_color"
|
||
if args.sim
|
||
else "/RMC1/arm95/svcam/right/image/compressed"
|
||
)
|
||
|
||
if not 0.0 <= args.conf <= 1.0:
|
||
parser.error("--conf должен быть от 0 до 1")
|
||
if args.timeout <= 0 or args.approach_z <= args.pick_z:
|
||
parser.error("timeout > 0, approach-z должен быть выше pick-z")
|
||
if (args.drop_x is None) != (args.drop_y is None):
|
||
parser.error("--drop-x и --drop-y задаются вместе")
|
||
intrinsics = (args.fx, args.fy, args.cx, args.cy)
|
||
if any(value is not None for value in intrinsics) and not all(
|
||
value is not None and math.isfinite(value) for value in intrinsics
|
||
):
|
||
parser.error("--fx, --fy, --cx и --cy задаются только вместе")
|
||
if args.fx is not None and (args.fx <= 0 or args.fy <= 0):
|
||
parser.error("--fx и --fy должны быть больше нуля")
|
||
return args
|
||
|
||
|
||
def rotate_vector(vector, quaternion):
|
||
"""Поворот вектора кватернионом geometry_msgs без внешней математики."""
|
||
x, y, z = vector
|
||
qx, qy, qz, qw = quaternion
|
||
tx = 2.0 * (qy * z - qz * y)
|
||
ty = 2.0 * (qz * x - qx * z)
|
||
tz = 2.0 * (qx * y - qy * x)
|
||
return (
|
||
x + qw * tx + qy * tz - qz * ty,
|
||
y + qw * ty + qz * tx - qx * tz,
|
||
z + qw * tz + qx * ty - qy * tx,
|
||
)
|
||
|
||
|
||
def normalize_half_turn(angle, reference=math.pi / 2):
|
||
"""Выбирает эквивалентный yaw (схват симметричен через 180 градусов)."""
|
||
choices = (angle - math.pi, angle, angle + math.pi)
|
||
return min(choices, key=lambda item: abs(item - reference))
|
||
|
||
|
||
def card_short_axis(frame, xyxy, cv2, np):
|
||
"""Находит направление короткой стороны светлой карточки внутри YOLO bbox."""
|
||
height, width = frame.shape[:2]
|
||
x1, y1, x2, y2 = (int(round(value)) for value in xyxy)
|
||
x1, y1 = max(0, x1), max(0, y1)
|
||
x2, y2 = min(width, x2), min(height, y2)
|
||
crop = frame[y1:y2, x1:x2]
|
||
if crop.size == 0:
|
||
return None
|
||
|
||
hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV)
|
||
mask = cv2.inRange(hsv, (0, 0, 105), (180, 115, 255))
|
||
kernel = np.ones((5, 5), np.uint8)
|
||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
|
||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if not contours:
|
||
return None
|
||
|
||
crop_area = crop.shape[0] * crop.shape[1]
|
||
candidates = [contour for contour in contours if cv2.contourArea(contour) > crop_area * 0.12]
|
||
if not candidates:
|
||
return None
|
||
rectangle = cv2.minAreaRect(max(candidates, key=cv2.contourArea))
|
||
points = cv2.boxPoints(rectangle)
|
||
edges = [points[(index + 1) % 4] - points[index] for index in range(4)]
|
||
short_edge = min(edges, key=lambda edge: float(np.linalg.norm(edge)))
|
||
length = float(np.linalg.norm(short_edge))
|
||
if length < 2.0:
|
||
return None
|
||
return float(short_edge[0] / length), float(short_edge[1] / length)
|
||
|
||
|
||
class Detection:
|
||
def __init__(self, name, confidence, xyxy, short_axis):
|
||
self.name = name
|
||
self.confidence = confidence
|
||
self.xyxy = xyxy
|
||
self.short_axis = short_axis
|
||
|
||
@property
|
||
def center(self):
|
||
x1, y1, x2, y2 = self.xyxy
|
||
return (x1 + x2) / 2.0, (y1 + y2) / 2.0
|
||
|
||
|
||
def find_targets(result, frame, requested, cv2, np):
|
||
found = []
|
||
if result.boxes is None:
|
||
return found
|
||
for box in result.boxes:
|
||
class_id = int(box.cls.item())
|
||
name = str(result.names[class_id]).lower()
|
||
if name != requested:
|
||
continue
|
||
confidence = float(box.conf.item())
|
||
xyxy = tuple(float(value) for value in box.xyxy[0])
|
||
found.append(Detection(name, confidence, xyxy, card_short_axis(frame, xyxy, cv2, np)))
|
||
return found
|
||
|
||
|
||
def choose_target(node, detections, plane_z):
|
||
"""При одинаковых классах выбирает ближайшую к базе доступную деталь."""
|
||
reachable = []
|
||
for detection in detections:
|
||
try:
|
||
x, y, _ = estimate_pose(node, detection, plane_z)
|
||
except Exception:
|
||
continue
|
||
reachable.append((math.hypot(x, y), -detection.confidence, detection))
|
||
if reachable:
|
||
return min(reachable, key=lambda item: (item[0], item[1]))[2]
|
||
return max(detections, key=lambda item: item.confidence, default=None)
|
||
|
||
|
||
def card_near_center(frame, requested, cv2, np):
|
||
"""На близком кадре находит светлую коробку около центра камеры."""
|
||
height, width = frame.shape[:2]
|
||
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
|
||
mask = cv2.inRange(hsv, (0, 0, 115), (180, 110, 255))
|
||
kernel = np.ones((7, 7), np.uint8)
|
||
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
|
||
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
candidates = []
|
||
for contour in contours:
|
||
rectangle = cv2.minAreaRect(contour)
|
||
rect_width, rect_height = rectangle[1]
|
||
rect_area = rect_width * rect_height
|
||
if rect_area < 2000 or min(rect_width, rect_height) < 20:
|
||
continue
|
||
aspect = max(rect_width, rect_height) / min(rect_width, rect_height)
|
||
if not 1.2 <= aspect <= 4.0:
|
||
continue
|
||
center_x, center_y = rectangle[0]
|
||
distance = math.hypot(center_x - width / 2.0, center_y - height / 2.0)
|
||
candidates.append((distance, rectangle))
|
||
|
||
if not candidates:
|
||
raise RuntimeError("На близком кадре не найдена коробка")
|
||
distance, rectangle = min(candidates, key=lambda item: item[0])
|
||
if distance > min(width, height) * 0.30:
|
||
raise RuntimeError(f"Коробка слишком далеко от центра кадра: {distance:.0f} px")
|
||
|
||
points = cv2.boxPoints(rectangle)
|
||
edges = [points[(index + 1) % 4] - points[index] for index in range(4)]
|
||
short_edge = min(edges, key=lambda edge: float(np.linalg.norm(edge)))
|
||
length = float(np.linalg.norm(short_edge))
|
||
short_axis = float(short_edge[0] / length), float(short_edge[1] / length)
|
||
x1, y1 = np.min(points, axis=0)
|
||
x2, y2 = np.max(points, axis=0)
|
||
return Detection(requested, 1.0, (float(x1), float(y1), float(x2), float(y2)), short_axis)
|
||
|
||
|
||
def estimate_pose(node, detection, plane_z):
|
||
if node.camera_info is None and node.manual_k is None:
|
||
raise RuntimeError(
|
||
f"Нет CameraInfo с {node.info_topic}; задайте калибровку --fx/--fy/--cx/--cy"
|
||
)
|
||
frame = node.camera_frame or (
|
||
node.camera_info.header.frame_id if node.camera_info is not None else ""
|
||
)
|
||
if not frame:
|
||
raise RuntimeError("В сообщении камеры не указан frame_id")
|
||
transform = node.tf_buffer.lookup_transform(
|
||
node.base_frame, frame, node.Time(), timeout=node.Duration(seconds=1.0)
|
||
)
|
||
translation = transform.transform.translation
|
||
rotation = transform.transform.rotation
|
||
quaternion = (rotation.x, rotation.y, rotation.z, rotation.w)
|
||
|
||
u, v = detection.center
|
||
matrix = node.camera_info.k if node.camera_info is not None else node.manual_k
|
||
fx, fy, cx, cy = matrix[0], matrix[4], matrix[2], matrix[5]
|
||
ray_camera = ((u - cx) / fx, (v - cy) / fy, 1.0)
|
||
ray_base = rotate_vector(ray_camera, quaternion)
|
||
if abs(ray_base[2]) < 1e-6:
|
||
raise RuntimeError("Луч камеры параллелен полке")
|
||
distance = (plane_z - translation.z) / ray_base[2]
|
||
if distance <= 0:
|
||
raise RuntimeError("Полка находится позади камеры")
|
||
|
||
x = translation.x + distance * ray_base[0]
|
||
y = translation.y + distance * ray_base[1]
|
||
if math.hypot(x, y) > 0.9:
|
||
raise RuntimeError(f"Координаты вне рабочей зоны: x={x:.3f}, y={y:.3f}")
|
||
|
||
yaw = math.pi / 2
|
||
if detection.short_axis is not None:
|
||
vx, vy = detection.short_axis
|
||
axis_base = rotate_vector((vx, vy, 0.0), quaternion)
|
||
short_angle = math.atan2(axis_base[1], axis_base[0])
|
||
yaw = normalize_half_turn(short_angle - math.pi / 2)
|
||
return x, y, yaw
|
||
|
||
|
||
def build_vision_node(args, imports):
|
||
(
|
||
Node,
|
||
CvBridge,
|
||
Image,
|
||
CompressedImage,
|
||
CameraInfo,
|
||
qos,
|
||
Buffer,
|
||
TransformListener,
|
||
Duration,
|
||
Time,
|
||
TFMessage,
|
||
Parameter,
|
||
QoSProfile,
|
||
ReliabilityPolicy,
|
||
DurabilityPolicy,
|
||
) = imports
|
||
|
||
class VisionNode(Node):
|
||
def __init__(self):
|
||
super().__init__("mvch_module3_vision")
|
||
self.set_parameters(
|
||
[Parameter("use_sim_time", value=args.sim)]
|
||
)
|
||
self.Duration = Duration
|
||
self.Time = Time
|
||
self.base_frame = args.base_frame
|
||
self.bridge = CvBridge()
|
||
self.frame = None
|
||
self.camera_info = None
|
||
self.manual_k = (
|
||
(args.fx, 0.0, args.cx, 0.0, args.fy, args.cy, 0.0, 0.0, 1.0)
|
||
if args.fx is not None
|
||
else None
|
||
)
|
||
self.camera_frame = None
|
||
self.lock = threading.Lock()
|
||
self.tf_buffer = Buffer()
|
||
self.tf_listener = TransformListener(self.tf_buffer, self)
|
||
if not args.sim:
|
||
tf_qos = QoSProfile(
|
||
depth=100,
|
||
reliability=ReliabilityPolicy.BEST_EFFORT,
|
||
durability=DurabilityPolicy.VOLATILE,
|
||
)
|
||
static_tf_qos = QoSProfile(
|
||
depth=100,
|
||
reliability=ReliabilityPolicy.RELIABLE,
|
||
durability=DurabilityPolicy.TRANSIENT_LOCAL,
|
||
)
|
||
self.create_subscription(
|
||
TFMessage, "/RMC1/tf", self.on_tf, tf_qos
|
||
)
|
||
self.create_subscription(
|
||
TFMessage, "/RMC1/tf_static", self.on_tf_static, static_tf_qos
|
||
)
|
||
image_type = CompressedImage if args.compressed else Image
|
||
self.create_subscription(image_type, args.topic, self.on_image, qos)
|
||
info_topic = args.camera_info_topic
|
||
if info_topic is None:
|
||
if args.topic.endswith("/image/compressed"):
|
||
info_topic = args.topic.removesuffix("/image/compressed") + "/camera_info"
|
||
else:
|
||
info_topic = args.topic.rsplit("/", 1)[0] + "/camera_info"
|
||
self.create_subscription(CameraInfo, info_topic, self.on_info, qos)
|
||
self.info_topic = info_topic
|
||
|
||
def on_tf(self, message):
|
||
for transform in message.transforms:
|
||
self.tf_buffer.set_transform(transform, "RMC1 namespaced tf")
|
||
|
||
def on_tf_static(self, message):
|
||
for transform in message.transforms:
|
||
self.tf_buffer.set_transform_static(
|
||
transform, "RMC1 namespaced tf_static"
|
||
)
|
||
|
||
def on_image(self, message):
|
||
try:
|
||
if args.compressed:
|
||
frame = self.bridge.compressed_imgmsg_to_cv2(
|
||
message, desired_encoding="bgr8"
|
||
)
|
||
else:
|
||
frame = self.bridge.imgmsg_to_cv2(
|
||
message, desired_encoding="bgr8"
|
||
)
|
||
except Exception as error:
|
||
self.get_logger().error(f"Ошибка кадра: {error}")
|
||
return
|
||
with self.lock:
|
||
self.frame = frame
|
||
self.camera_frame = message.header.frame_id
|
||
|
||
def on_info(self, message):
|
||
self.camera_info = message
|
||
|
||
def latest_frame(self):
|
||
with self.lock:
|
||
return None if self.frame is None else self.frame.copy()
|
||
|
||
return VisionNode()
|
||
|
||
|
||
class Arm:
|
||
def __init__(
|
||
self,
|
||
moveit,
|
||
params_file,
|
||
PoseStamped,
|
||
quaternion_from_euler,
|
||
event_log,
|
||
plan_only=False,
|
||
base_frame="Base_link",
|
||
):
|
||
self.robot = moveit(
|
||
node_name="mvch_module3_moveit",
|
||
name_space="/RMC1/arm95",
|
||
launch_params_filepaths=[params_file],
|
||
config_dict=None,
|
||
)
|
||
self.arm = self.robot.get_planning_component("arm95_group")
|
||
self.gripper = self.robot.get_planning_component("gripper")
|
||
self.PoseStamped = PoseStamped
|
||
self.quaternion_from_euler = quaternion_from_euler
|
||
self.log = event_log
|
||
self.plan_only = plan_only
|
||
self.base_frame = base_frame
|
||
|
||
def execute(self, component, stage):
|
||
plan = component.plan()
|
||
if not plan:
|
||
raise RuntimeError(f"MoveIt не построил траекторию: {stage}")
|
||
self.log.write("PLAN_READY", stage=stage)
|
||
if not self.plan_only:
|
||
self.robot.execute(plan.trajectory, controllers=[])
|
||
time.sleep(0.4)
|
||
self.log.write("MOTION_DONE", stage=stage)
|
||
|
||
def named(self, component, state, stage):
|
||
component.set_start_state_to_current_state()
|
||
component.set_goal_state(configuration_name=state)
|
||
self.execute(component, stage)
|
||
|
||
def pose(self, x, y, z, yaw, stage):
|
||
goal = self.PoseStamped()
|
||
goal.header.frame_id = self.base_frame
|
||
qx, qy, qz, qw = self.quaternion_from_euler(math.pi, 0.0, yaw)
|
||
goal.pose.orientation.x = qx
|
||
goal.pose.orientation.y = qy
|
||
goal.pose.orientation.z = qz
|
||
goal.pose.orientation.w = qw
|
||
goal.pose.position.x = x
|
||
goal.pose.position.y = y
|
||
goal.pose.position.z = z
|
||
self.arm.set_start_state_to_current_state()
|
||
self.arm.set_goal_state(pose_stamped_msg=goal, pose_link="gripper_base")
|
||
self.execute(self.arm, stage)
|
||
|
||
def open(self):
|
||
self.named(self.gripper, "open", "открыть схват")
|
||
if not self.plan_only:
|
||
self.log.write("GRIPPER", state="open")
|
||
|
||
def close(self):
|
||
self.named(self.gripper, "closed", "закрыть схват")
|
||
if not self.plan_only:
|
||
self.log.write("GRIPPER", state="closed")
|
||
|
||
def initial(self, stage="исходное положение"):
|
||
self.named(self.arm, "table_pos", stage)
|
||
|
||
|
||
def wait_for_trigger(args, node, model, requested, cv2, np, event_log):
|
||
started = time.monotonic()
|
||
last_detection = None
|
||
last_frame_time = 0.0
|
||
last_status_time = 0.0
|
||
frame_count = 0
|
||
best_classes = {}
|
||
terminal_trigger = threading.Event()
|
||
target_visible = threading.Event()
|
||
window = not args.no_window and bool(os.environ.get("DISPLAY"))
|
||
if not args.no_window and not window:
|
||
print("[WARN] DISPLAY не найден, работаю без окна")
|
||
|
||
def read_terminal_command():
|
||
while not terminal_trigger.is_set():
|
||
try:
|
||
command = input().strip().lower()
|
||
except EOFError:
|
||
return
|
||
if command not in ("", "g", "go"):
|
||
print("[WAIT] Команда не распознана. Введите g и нажмите Enter.")
|
||
continue
|
||
if not target_visible.is_set():
|
||
print(f"[WAIT] Класс {requested} пока не найден, команда не принята.")
|
||
continue
|
||
terminal_trigger.set()
|
||
|
||
if not args.auto_start and not args.dry_run and sys.stdin.isatty():
|
||
threading.Thread(target=read_terminal_command, daemon=True).start()
|
||
|
||
if args.dry_run:
|
||
print(f"[READY] Dry-run: жду первое распознавание класса {requested}.")
|
||
else:
|
||
print(
|
||
f"[READY] Цель: {requested}. После фиксации нажмите G в окне "
|
||
"или введите g + Enter в терминале."
|
||
)
|
||
while time.monotonic() - started < args.timeout:
|
||
frame = node.latest_frame()
|
||
if frame is None:
|
||
time.sleep(0.05)
|
||
continue
|
||
frame_count += 1
|
||
result = model.predict(
|
||
source=frame,
|
||
conf=args.conf,
|
||
imgsz=args.imgsz,
|
||
device=args.device,
|
||
verbose=False,
|
||
)[0]
|
||
if result.boxes is not None:
|
||
for box in result.boxes:
|
||
name = str(result.names[int(box.cls.item())]).lower()
|
||
confidence = float(box.conf.item())
|
||
best_classes[name] = max(best_classes.get(name, 0.0), confidence)
|
||
detections = find_targets(result, frame, requested, cv2, np)
|
||
last_detection = choose_target(node, detections, args.plane_z)
|
||
if last_detection is None:
|
||
target_visible.clear()
|
||
else:
|
||
target_visible.set()
|
||
annotated = result.plot(labels=True, conf=True, line_width=2)
|
||
if last_detection is not None:
|
||
u, v = (int(value) for value in last_detection.center)
|
||
cv2.circle(annotated, (u, v), 7, (0, 255, 255), -1)
|
||
if last_detection.short_axis is not None:
|
||
vx, vy = last_detection.short_axis
|
||
cv2.line(
|
||
annotated,
|
||
(u - int(vx * 45), v - int(vy * 45)),
|
||
(u + int(vx * 45), v + int(vy * 45)),
|
||
(255, 255, 0),
|
||
3,
|
||
)
|
||
|
||
now = time.monotonic()
|
||
if last_detection is not None:
|
||
if now - last_frame_time > 0.8:
|
||
event_log.write(
|
||
"DETECTION",
|
||
target=last_detection.name,
|
||
confidence=round(last_detection.confidence, 4),
|
||
center=[round(value, 1) for value in last_detection.center],
|
||
)
|
||
last_frame_time = now
|
||
elif now - last_status_time > 5.0:
|
||
event_log.write(
|
||
"VISION_WAIT",
|
||
target=requested,
|
||
frames=frame_count,
|
||
seen={name: round(value, 3) for name, value in sorted(best_classes.items())},
|
||
)
|
||
last_status_time = now
|
||
|
||
key = -1
|
||
if window:
|
||
cv2.imshow("CHVT module 3 - G: start, ESC: exit", annotated)
|
||
key = cv2.waitKey(1) & 0xFF
|
||
if key == 27:
|
||
raise KeyboardInterrupt
|
||
if last_detection is not None and (
|
||
args.dry_run
|
||
or args.auto_start
|
||
or terminal_trigger.is_set()
|
||
or key in (ord("g"), ord("G"), 13, 32)
|
||
):
|
||
return last_detection
|
||
if frame_count == 0:
|
||
raise RuntimeError(f"Нет кадров с топика {args.topic} за {args.timeout:.0f} с")
|
||
seen = ", ".join(f"{name}={value:.2f}" for name, value in sorted(best_classes.items()))
|
||
raise RuntimeError(
|
||
f"За {args.timeout:.0f} с класс {requested} не найден; "
|
||
f"обработано кадров: {frame_count}; распознано: {seen or 'ничего'}"
|
||
)
|
||
|
||
|
||
def fresh_detection(node, model, requested, args, cv2, np):
|
||
time.sleep(0.7)
|
||
frame = node.latest_frame()
|
||
if frame is None:
|
||
raise RuntimeError("Нет свежего кадра после возврата руки")
|
||
result = model.predict(
|
||
source=frame,
|
||
conf=args.conf,
|
||
imgsz=args.imgsz,
|
||
device=args.device,
|
||
verbose=False,
|
||
)[0]
|
||
detections = find_targets(result, frame, requested, cv2, np)
|
||
detection = choose_target(node, detections, args.plane_z)
|
||
if detection is None:
|
||
raise RuntimeError(f"После установки руки не найден класс {requested}")
|
||
return detection
|
||
|
||
|
||
def save_snapshot(node, event_log, suffix, cv2):
|
||
time.sleep(0.5)
|
||
frame = node.latest_frame()
|
||
if frame is None:
|
||
event_log.write("SNAPSHOT_FAILED", stage=suffix, reason="нет кадра")
|
||
return
|
||
path = event_log.image_path(suffix)
|
||
if cv2.imwrite(str(path), frame):
|
||
event_log.write("SNAPSHOT", stage=suffix, path=str(path))
|
||
else:
|
||
event_log.write("SNAPSHOT_FAILED", stage=suffix, reason="ошибка записи")
|
||
|
||
|
||
def main():
|
||
args = arguments()
|
||
requested = target_name(args.target or input("Цель (1 hammer, 2 wrench, 3 pliers): "))
|
||
if requested not in TARGETS.values():
|
||
raise SystemExit(f"Неизвестная цель: {requested}")
|
||
weights = local_path(args.weights)
|
||
if not weights.is_file():
|
||
raise SystemExit(f"Не найдены веса: {weights}")
|
||
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
if int(np.__version__.split(".", 1)[0]) >= 2:
|
||
raise SystemExit(
|
||
"NumPy 2.x несовместим с cv_bridge из ROS Jazzy. "
|
||
"Выполните: python3 -m pip install --force-reinstall 'numpy<2'"
|
||
)
|
||
import rclpy
|
||
import yaml
|
||
from ament_index_python.packages import get_package_share_directory
|
||
from cv_bridge import CvBridge
|
||
from geometry_msgs.msg import PoseStamped
|
||
from moveit.planning import MoveItPy
|
||
from moveit_configs_utils import MoveItConfigsBuilder
|
||
from rclpy.duration import Duration
|
||
from rclpy.executors import SingleThreadedExecutor
|
||
from rclpy.node import Node
|
||
from rclpy.parameter import Parameter
|
||
from rclpy.qos import (
|
||
DurabilityPolicy,
|
||
QoSProfile,
|
||
ReliabilityPolicy,
|
||
qos_profile_sensor_data,
|
||
)
|
||
from rclpy.time import Time
|
||
from sensor_msgs.msg import CameraInfo, CompressedImage, Image
|
||
from tf2_msgs.msg import TFMessage
|
||
from tf2_ros import Buffer, TransformListener
|
||
from tf_transformations import quaternion_from_euler
|
||
from ultralytics import YOLO
|
||
except ImportError as error:
|
||
if error.name == "transforms3d":
|
||
raise SystemExit(
|
||
"Venv не видит системные зависимости ROS. Включите их: "
|
||
"sed -i 's/include-system-site-packages = false/"
|
||
"include-system-site-packages = true/' .venv/pyvenv.cfg"
|
||
)
|
||
raise SystemExit(
|
||
f"Не удалось импортировать {error.name}: {error}. "
|
||
"Выполните pip install -r requirements.txt после source ROS"
|
||
)
|
||
|
||
event_log = EventLog()
|
||
event_log.write("START", target=requested, log=str(event_log.path))
|
||
rclpy.init()
|
||
imports = (
|
||
Node,
|
||
CvBridge,
|
||
Image,
|
||
CompressedImage,
|
||
CameraInfo,
|
||
qos_profile_sensor_data,
|
||
Buffer,
|
||
TransformListener,
|
||
Duration,
|
||
Time,
|
||
TFMessage,
|
||
Parameter,
|
||
QoSProfile,
|
||
ReliabilityPolicy,
|
||
DurabilityPolicy,
|
||
)
|
||
node = build_vision_node(args, imports)
|
||
executor = SingleThreadedExecutor()
|
||
executor.add_node(node)
|
||
spin_thread = threading.Thread(target=executor.spin, daemon=True)
|
||
spin_thread.start()
|
||
model = YOLO(str(weights))
|
||
moveit_active = False
|
||
exit_code = 0
|
||
|
||
try:
|
||
detection = wait_for_trigger(args, node, model, requested, cv2, np, event_log)
|
||
if args.dry_run:
|
||
x, y, yaw = estimate_pose(node, detection, args.plane_z)
|
||
event_log.write(
|
||
"DRY_RUN",
|
||
x=round(x, 4),
|
||
y=round(y, 4),
|
||
plane_z=args.plane_z,
|
||
yaw=round(yaw, 4),
|
||
)
|
||
return
|
||
|
||
package_dir = Path(get_package_share_directory("ar_webots_fms_ros2"))
|
||
moveit_config = (
|
||
MoveItConfigsBuilder("arm95")
|
||
.robot_description(file_path=str(package_dir / "resource/urdf/arm95_webots.urdf"))
|
||
.trajectory_execution(
|
||
file_path=str(package_dir / "resource/config/moveit_controllers.yaml")
|
||
)
|
||
.moveit_cpp(file_path=str(package_dir / "resource/config/moveit_cpp.yaml"))
|
||
.to_moveit_configs()
|
||
.to_dict()
|
||
)
|
||
moveit_config["use_sim_time"] = args.sim
|
||
# Wildcard нужен из-за namespace /RMC1/arm95: обычный config_dict
|
||
# записывает параметры только для имени узла без namespace.
|
||
with NamedTemporaryFile("w", suffix=".yaml", delete=False) as stream:
|
||
yaml.safe_dump({"/**": {"ros__parameters": moveit_config}}, stream)
|
||
params_file = stream.name
|
||
try:
|
||
arm = Arm(
|
||
MoveItPy,
|
||
params_file,
|
||
PoseStamped,
|
||
quaternion_from_euler,
|
||
event_log,
|
||
args.plan_only,
|
||
args.base_frame,
|
||
)
|
||
moveit_active = True
|
||
finally:
|
||
Path(params_file).unlink(missing_ok=True)
|
||
arm.open()
|
||
arm.initial("установить исходное положение")
|
||
detection = fresh_detection(node, model, requested, args, cv2, np)
|
||
x, y, yaw = estimate_pose(node, detection, args.plane_z)
|
||
drop_x = x if args.drop_x is None else args.drop_x
|
||
drop_y = y if args.drop_y is None else args.drop_y
|
||
event_log.write(
|
||
"TARGET",
|
||
target=requested,
|
||
confidence=round(detection.confidence, 4),
|
||
x=round(x, 4),
|
||
y=round(y, 4),
|
||
yaw=round(yaw, 4),
|
||
)
|
||
|
||
arm.pose(x, y, args.approach_z, yaw, "подход над деталью")
|
||
if args.plan_only:
|
||
print("[PLAN-ONLY] Подход построен, движения не было")
|
||
return
|
||
time.sleep(0.5)
|
||
close_frame = node.latest_frame()
|
||
if close_frame is None:
|
||
raise RuntimeError("Нет близкого кадра для уточнения")
|
||
close_detection = card_near_center(close_frame, requested, cv2, np)
|
||
close_x, close_y, close_yaw = estimate_pose(node, close_detection, args.plane_z)
|
||
correction = math.hypot(close_x - x, close_y - y)
|
||
if correction > 0.12:
|
||
raise RuntimeError(f"Слишком большая поправка координат: {correction:.3f} м")
|
||
x, y, yaw = close_x, close_y, close_yaw
|
||
event_log.write(
|
||
"TARGET_REFINED",
|
||
method="card_near_center",
|
||
correction=round(correction, 4),
|
||
x=round(x, 4),
|
||
y=round(y, 4),
|
||
yaw=round(yaw, 4),
|
||
)
|
||
arm.pose(x, y, args.approach_z, yaw, "точный подход над деталью")
|
||
arm.pose(x, y, args.pick_z, yaw, "опустить схват")
|
||
arm.close()
|
||
save_snapshot(node, event_log, "gripped", cv2)
|
||
arm.pose(x, y, args.approach_z, yaw, "поднять деталь")
|
||
arm.initial("исходное положение с деталью")
|
||
save_snapshot(node, event_log, "with_detail", cv2)
|
||
|
||
arm.pose(drop_x, drop_y, args.approach_z, yaw, "подход к пустому месту")
|
||
arm.pose(drop_x, drop_y, args.pick_z, yaw, "опустить деталь")
|
||
arm.open()
|
||
arm.pose(drop_x, drop_y, args.approach_z, yaw, "отойти от детали")
|
||
arm.initial("исходное положение без детали")
|
||
save_snapshot(node, event_log, "finished", cv2)
|
||
event_log.write("DONE", message="Деталь возвращена, манипулятор в исходном положении")
|
||
except KeyboardInterrupt:
|
||
print("\n[STOP] Остановлено оператором")
|
||
exit_code = 130
|
||
except Exception as error:
|
||
event_log.write("FAILED", reason=str(error))
|
||
exit_code = 1
|
||
finally:
|
||
cv2.destroyAllWindows()
|
||
executor.shutdown()
|
||
node.destroy_node()
|
||
if rclpy.ok():
|
||
rclpy.shutdown()
|
||
if moveit_active:
|
||
# MoveItPy в поставке Jazzy падает в C++-деструкторе после штатной
|
||
# работы. Все движения и логи уже завершены, поэтому не вызываем
|
||
# проблемный teardown при выходе единственного зачетного скрипта.
|
||
sys.stdout.flush()
|
||
sys.stderr.flush()
|
||
os._exit(exit_code)
|
||
if exit_code:
|
||
raise SystemExit(exit_code)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|