init
@@ -0,0 +1,172 @@
|
||||
# Модуль В: датасет и YOLO
|
||||
|
||||
Камера: `/RMC1/arm95/camera_gripper/image_color` (`1600x1200`).
|
||||
|
||||
Основной зачетный скрипт: `module3.py`. Он показывает обработанный видеопоток,
|
||||
определяет координаты выбранной детали относительно `Base_link`, захватывает ее,
|
||||
возвращает ARM95 в исходное положение с деталью, кладет деталь обратно на
|
||||
освободившееся место и снова возвращает руку в исходное положение.
|
||||
После первого безопасного подхода координаты повторно уточняются по близкому
|
||||
кадру, поэтому рука не едет к детали только по одному дальнему измерению.
|
||||
Если в кадре несколько деталей одного класса, скрипт выбирает ближайшую к базе.
|
||||
|
||||
## Запуск задания
|
||||
|
||||
В VM один раз создайте окружение с доступом к системным пакетам ROS:
|
||||
|
||||
```bash
|
||||
cd ~/scripts/mvch/module3
|
||||
sudo apt install python3-venv
|
||||
python3 -m venv --system-site-packages .venv
|
||||
.venv/bin/pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
.venv/bin/pip install -r requirements.txt tqdm ultralytics-thop
|
||||
```
|
||||
|
||||
Если окружение уже обновило NumPy до 2.x и `cv_bridge` ругается на ABI:
|
||||
|
||||
```bash
|
||||
.venv/bin/pip install --force-reinstall "numpy<2" "opencv-python<4.12" lark
|
||||
```
|
||||
|
||||
`lark` — зависимость ROS 2 Launch/MoveIt, а не часть алгоритма распознавания.
|
||||
Ограничение `numpy<2` нужно потому, что бинарный `cv_bridge` в этом образе ROS
|
||||
собран с NumPy 1.x.
|
||||
|
||||
Если venv уже был создан без `--system-site-packages`, пересоздавать его не
|
||||
нужно. Включите системные зависимости ROS и заново активируйте окружение:
|
||||
|
||||
```bash
|
||||
sed -i 's/include-system-site-packages = false/include-system-site-packages = true/' .venv/pyvenv.cfg
|
||||
deactivate
|
||||
source .venv/bin/activate
|
||||
```
|
||||
|
||||
Запустите симулятор модуля В, затем в новом терминале:
|
||||
|
||||
```bash
|
||||
source /opt/ros/jazzy/setup.bash
|
||||
source ~/ros2_ws/install/setup.bash
|
||||
cd ~/scripts/mvch
|
||||
module3/.venv/bin/python module3/module3.py --target 1
|
||||
```
|
||||
|
||||
Цели: `1/hammer`, `2/wrench`, `3/pliers`. После того как эксперт зафиксировал
|
||||
видеопоток, нажмите `G` или `Space` в окне либо введите `g` и
|
||||
нажмите Enter в терминале. `Esc` останавливает скрипт до начала движения.
|
||||
|
||||
Сначала обязательно проверьте только координаты:
|
||||
|
||||
```bash
|
||||
module3/.venv/bin/python module3/module3.py --target hammer --dry-run
|
||||
```
|
||||
|
||||
В режиме `--dry-run` подтверждать запуск клавишей не нужно: скрипт завершится
|
||||
после первой валидной детекции и вывода рассчитанных координат.
|
||||
|
||||
Для проверки через SSH без окна:
|
||||
|
||||
```bash
|
||||
module3/.venv/bin/python module3/module3.py --target hammer --dry-run --no-window
|
||||
```
|
||||
|
||||
Геометрию реального стенда нельзя брать из симулятора вслепую. Перед зачетом
|
||||
измерьте и подстройте `--plane-z`, `--pick-z` и `--approach-z`. По умолчанию
|
||||
деталь возвращается туда, откуда была взята; другую свободную точку можно задать
|
||||
через `--drop-x X --drop-y Y`.
|
||||
|
||||
События распознавания, координаты, построение движений и состояния схвата
|
||||
одновременно выводятся в терминал и сохраняются в `module3/logs/*.log`.
|
||||
Туда же пишутся контрольные кадры `*_gripped.jpg`, `*_with_detail.jpg` и
|
||||
`*_finished.jpg`.
|
||||
|
||||
## Что размечать
|
||||
|
||||
Классы строго в таком написании:
|
||||
|
||||
```text
|
||||
hammer
|
||||
pliers
|
||||
wrench
|
||||
```
|
||||
|
||||
Рамкой выделяйте изображение инструмента на карточке, а класс задавайте по
|
||||
картинке. Используйте одинаковый стиль разметки в симуляторе и на реальном поле.
|
||||
Размечайте каждый видимый экземпляр: в штатной сцене две карточки `pliers`, одна
|
||||
`hammer` и одна `wrench`. Классы `box` и `background` не создавайте.
|
||||
|
||||
В Roboflow создайте `Object Detection` project, загрузите `dataset/raw/*.jpg`,
|
||||
разметьте кадры, сделайте split 70/20/10 и экспортируйте в формате YOLO11.
|
||||
Распакуйте экспорт в `module3/training`: там должны появиться `data.yaml` и
|
||||
каталоги `train`, `valid`, `test`. Не переставляйте номера классов в `data.yaml`:
|
||||
они назначаются Roboflow и должны совпадать с ID внутри label-файлов.
|
||||
|
||||
Локально в `dataset/raw` могут лежать собранные кадры, но Git их намеренно
|
||||
игнорирует: после клонирования датасет нужно перенести отдельно или собрать
|
||||
заново. Расположение объектов менять не требуется: дополнительные повороты и
|
||||
сдвиги выполняются аугментациями во время обучения.
|
||||
|
||||
## Сбор кадров из симулятора
|
||||
|
||||
```bash
|
||||
source /opt/ros/jazzy/setup.bash
|
||||
source ~/ros2_ws/install/setup.bash
|
||||
python3 module3/capture_dataset.py --count 30
|
||||
```
|
||||
|
||||
Во время сбора плавно меняйте положение RMC1/манипулятора. Скрипт пропускает почти
|
||||
одинаковые кадры. Для одиночного снимка используйте `--count 1 --min-change 0`.
|
||||
|
||||
## Установка для обучения
|
||||
|
||||
```bash
|
||||
cd module3
|
||||
python3 -m venv .venv
|
||||
source .venv/bin/activate
|
||||
python -m pip install -U pip
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
Ultralytics ставит PyTorch как зависимость. Для конкретной NVIDIA CUDA или AMD
|
||||
ROCm лучше сначала установить подходящую сборку PyTorch, затем `requirements.txt`.
|
||||
|
||||
## Обучение
|
||||
|
||||
Если часть объектов была размечена полигонами, один раз приведите экспорт к
|
||||
обычным bounding boxes:
|
||||
|
||||
```bash
|
||||
python prepare_dataset.py
|
||||
```
|
||||
|
||||
Для теста используются максимум 30 эпох и early stopping после 7 эпох без
|
||||
улучшения validation metrics.
|
||||
|
||||
Автовыбор: NVIDIA CUDA / AMD ROCm, а если GPU недоступен — CPU:
|
||||
|
||||
```bash
|
||||
python train.py --device auto
|
||||
```
|
||||
|
||||
Только CPU (подойдёт на AMD без настроенного ROCm):
|
||||
|
||||
```bash
|
||||
python train.py --device cpu --batch 4 --workers 2
|
||||
```
|
||||
|
||||
Принудительно NVIDIA CUDA или настроенный AMD ROCm:
|
||||
|
||||
```bash
|
||||
python train.py --device cuda
|
||||
```
|
||||
|
||||
Лучшие веса сохраняются в `training/runs/tools_yolo11n/weights/best.pt` и
|
||||
автоматически копируются в `models/latest.pt` для простого запуска проверки.
|
||||
|
||||
## Проверка на фотографии
|
||||
|
||||
```bash
|
||||
python test_model.py training/test/images/example.jpg --device cpu
|
||||
```
|
||||
|
||||
С `--show` откроется окно. Независимо от окна размеченная фотография сохраняется
|
||||
в `predictions/`: на ней видны bbox, названия классов и confidence.
|
||||
@@ -0,0 +1,115 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Сохраняет непохожие кадры с камеры захвата RMC1 для разметки."""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def arguments():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--topic",
|
||||
default="/RMC1/arm95/camera_gripper/image_color",
|
||||
help="ROS-топик sensor_msgs/Image",
|
||||
)
|
||||
parser.add_argument("--output", default="dataset/raw", help="Каталог для JPG")
|
||||
parser.add_argument("--count", type=int, default=30, help="Сколько кадров сохранить")
|
||||
parser.add_argument("--interval", type=float, default=0.5, help="Минимум секунд между JPG")
|
||||
parser.add_argument(
|
||||
"--min-change",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Минимальное среднее изменение пикселя; 0 сохраняет одинаковые кадры",
|
||||
)
|
||||
parser.add_argument("--timeout", type=float, default=180.0, help="Общий таймаут")
|
||||
parser.add_argument("--jpeg-quality", type=int, default=95)
|
||||
args = parser.parse_args()
|
||||
if args.count < 1 or args.interval < 0 or args.timeout <= 0:
|
||||
parser.error("count >= 1, interval >= 0, timeout > 0")
|
||||
if not 1 <= args.jpeg_quality <= 100 or args.min_change < 0:
|
||||
parser.error("jpeg-quality должен быть 1..100, min-change >= 0")
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = arguments()
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
import rclpy
|
||||
from cv_bridge import CvBridge
|
||||
from rclpy.node import Node
|
||||
from rclpy.qos import qos_profile_sensor_data
|
||||
from sensor_msgs.msg import Image
|
||||
except ImportError as error:
|
||||
raise SystemExit(
|
||||
f"Нет ROS/OpenCV зависимости: {error}. Запускайте после source ROS 2 в VM."
|
||||
)
|
||||
|
||||
output = Path(args.output).expanduser()
|
||||
if not output.is_absolute():
|
||||
output = HERE / output
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
class Collector(Node):
|
||||
def __init__(self):
|
||||
super().__init__("mvch_dataset_collector")
|
||||
self.bridge = CvBridge()
|
||||
self.saved = 0
|
||||
self.previous = None
|
||||
self.last_save = 0.0
|
||||
self.create_subscription(Image, args.topic, self.on_image, qos_profile_sensor_data)
|
||||
|
||||
def on_image(self, message):
|
||||
now = time.monotonic()
|
||||
if now - self.last_save < args.interval:
|
||||
return
|
||||
try:
|
||||
frame = self.bridge.imgmsg_to_cv2(message, desired_encoding="bgr8")
|
||||
except Exception as error:
|
||||
self.get_logger().error(f"Не удалось преобразовать кадр: {error}")
|
||||
return
|
||||
preview = cv2.resize(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY), (64, 48))
|
||||
if self.previous is not None:
|
||||
change = float(np.mean(cv2.absdiff(preview, self.previous)))
|
||||
if change < args.min_change:
|
||||
return
|
||||
filename = output / f"frame_{datetime.now():%Y%m%d_%H%M%S_%f}_{self.saved:03d}.jpg"
|
||||
if not cv2.imwrite(
|
||||
str(filename), frame, [cv2.IMWRITE_JPEG_QUALITY, args.jpeg_quality]
|
||||
):
|
||||
self.get_logger().error(f"Не удалось записать {filename}")
|
||||
return
|
||||
self.previous = preview
|
||||
self.last_save = now
|
||||
self.saved += 1
|
||||
print(f"[{self.saved}/{args.count}] {filename}", flush=True)
|
||||
|
||||
rclpy.init()
|
||||
node = Collector()
|
||||
started = time.monotonic()
|
||||
print(f"Камера: {args.topic}")
|
||||
print("Плавно двигайте RMC1 или манипулятор, чтобы ракурсы менялись.")
|
||||
try:
|
||||
while rclpy.ok() and node.saved < args.count:
|
||||
if time.monotonic() - started >= args.timeout:
|
||||
break
|
||||
rclpy.spin_once(node, timeout_sec=0.2)
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
finally:
|
||||
saved = node.saved
|
||||
node.destroy_node()
|
||||
rclpy.shutdown()
|
||||
print(f"Готово: сохранено {saved} кадров в {output}")
|
||||
if saved == 0:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1 @@
|
||||
raw
|
||||
@@ -0,0 +1,3 @@
|
||||
hammer
|
||||
pliers
|
||||
wrench
|
||||
@@ -0,0 +1,9 @@
|
||||
path: .
|
||||
train: train/images
|
||||
val: valid/images
|
||||
test: test/images
|
||||
|
||||
names:
|
||||
0: hammer
|
||||
1: pliers
|
||||
2: wrench
|
||||
@@ -0,0 +1,2 @@
|
||||
*
|
||||
!.gitignore
|
||||
@@ -0,0 +1,691 @@
|
||||
#!/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("--topic", default="/RMC1/arm95/camera_gripper/image_color")
|
||||
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 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 задаются вместе")
|
||||
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:
|
||||
raise RuntimeError("Нет CameraInfo")
|
||||
frame = node.camera_frame or node.camera_info.header.frame_id
|
||||
transform = node.tf_buffer.lookup_transform(
|
||||
"Base_link", 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
|
||||
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, CameraInfo, qos, Buffer, TransformListener, Duration, Time = imports
|
||||
|
||||
class VisionNode(Node):
|
||||
def __init__(self):
|
||||
super().__init__("mvch_module3_vision")
|
||||
self.Duration = Duration
|
||||
self.Time = Time
|
||||
self.bridge = CvBridge()
|
||||
self.frame = None
|
||||
self.camera_info = None
|
||||
self.camera_frame = None
|
||||
self.lock = threading.Lock()
|
||||
self.tf_buffer = Buffer()
|
||||
self.tf_listener = TransformListener(self.tf_buffer, self)
|
||||
self.create_subscription(Image, args.topic, self.on_image, qos)
|
||||
info_topic = args.topic.rsplit("/", 1)[0] + "/camera_info"
|
||||
self.create_subscription(CameraInfo, info_topic, self.on_info, qos)
|
||||
|
||||
def on_image(self, message):
|
||||
try:
|
||||
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,
|
||||
):
|
||||
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
|
||||
|
||||
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 = "Base_link"
|
||||
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.qos import qos_profile_sensor_data
|
||||
from rclpy.time import Time
|
||||
from sensor_msgs.msg import CameraInfo, Image
|
||||
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,
|
||||
CameraInfo,
|
||||
qos_profile_sensor_data,
|
||||
Buffer,
|
||||
TransformListener,
|
||||
Duration,
|
||||
Time,
|
||||
)
|
||||
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"] = True
|
||||
# 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,
|
||||
)
|
||||
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()
|
||||
@@ -0,0 +1,61 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Преобразует polygon-разметку Roboflow в обычные YOLO bounding boxes."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def arguments():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--dataset", default="training", help="Каталог Roboflow export")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def convert_line(line, filename):
|
||||
values = line.split()
|
||||
if len(values) == 5:
|
||||
return line.strip(), False
|
||||
if len(values) < 7 or (len(values) - 1) % 2:
|
||||
raise ValueError(f"Некорректная строка в {filename}: {line.strip()}")
|
||||
|
||||
class_id = int(values[0])
|
||||
coordinates = [float(value) for value in values[1:]]
|
||||
xs = coordinates[0::2]
|
||||
ys = coordinates[1::2]
|
||||
x1, x2 = min(xs), max(xs)
|
||||
y1, y2 = min(ys), max(ys)
|
||||
box = (class_id, (x1 + x2) / 2, (y1 + y2) / 2, x2 - x1, y2 - y1)
|
||||
return f"{box[0]} {box[1]:.6f} {box[2]:.6f} {box[3]:.6f} {box[4]:.6f}", True
|
||||
|
||||
|
||||
def main():
|
||||
args = arguments()
|
||||
dataset = Path(args.dataset).expanduser()
|
||||
if not dataset.is_absolute():
|
||||
dataset = HERE / dataset
|
||||
|
||||
files = sorted(dataset.glob("*/labels/*.txt"))
|
||||
if not files:
|
||||
raise SystemExit(f"Не найдены labels в {dataset}")
|
||||
|
||||
converted = 0
|
||||
objects = 0
|
||||
for filename in files:
|
||||
output = []
|
||||
for line in filename.read_text().splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
result, changed = convert_line(line, filename)
|
||||
output.append(result)
|
||||
converted += int(changed)
|
||||
objects += 1
|
||||
filename.write_text("\n".join(output) + "\n")
|
||||
|
||||
print(f"Готово: файлов={len(files)}, объектов={objects}, polygon->bbox={converted}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
numpy>=1.23.5,<2
|
||||
opencv-python>=4.6,<4.12
|
||||
lark>=1.1
|
||||
transforms3d>=0.4
|
||||
ultralytics==8.3.0
|
||||
@@ -0,0 +1,92 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Рисует bbox, название класса и confidence модели на фотографии."""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def local_path(value):
|
||||
path = Path(value).expanduser()
|
||||
return path if path.is_absolute() else HERE / path
|
||||
|
||||
|
||||
def select_device(requested, torch):
|
||||
if requested == "auto":
|
||||
if torch.cuda.is_available():
|
||||
return 0
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
if requested in ("cuda", "gpu"):
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA/ROCm не найден. Используйте --device cpu.")
|
||||
return 0
|
||||
return requested
|
||||
|
||||
|
||||
def arguments():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("source", help="Фотография для проверки")
|
||||
parser.add_argument("--weights", default="models/latest.pt")
|
||||
parser.add_argument("--output", default="predictions", help="Каталог результата")
|
||||
parser.add_argument("--conf", type=float, default=0.35)
|
||||
parser.add_argument("--imgsz", type=int, default=512)
|
||||
parser.add_argument("--device", default="auto", help="auto, cpu, cuda или номер GPU")
|
||||
parser.add_argument("--show", action="store_true", help="Открыть окно с результатом")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = arguments()
|
||||
try:
|
||||
import cv2
|
||||
import torch
|
||||
from ultralytics import YOLO
|
||||
except ImportError as error:
|
||||
raise SystemExit(f"Нет зависимости {error.name}. Выполните: pip install -r requirements.txt")
|
||||
|
||||
source = local_path(args.source)
|
||||
weights = local_path(args.weights)
|
||||
output = local_path(args.output)
|
||||
if not source.is_file():
|
||||
raise SystemExit(f"Не найдена фотография: {source}")
|
||||
if not weights.is_file():
|
||||
raise SystemExit(f"Не найдены веса: {weights}")
|
||||
if not 0 <= args.conf <= 1:
|
||||
raise SystemExit("--conf должен быть от 0 до 1")
|
||||
|
||||
model = YOLO(str(weights))
|
||||
result = model.predict(
|
||||
source=str(source),
|
||||
conf=args.conf,
|
||||
imgsz=args.imgsz,
|
||||
device=select_device(args.device, torch),
|
||||
verbose=False,
|
||||
)[0]
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
destination = output / f"{source.stem}_detected.jpg"
|
||||
annotated = result.plot(labels=True, conf=True, line_width=2)
|
||||
if not cv2.imwrite(str(destination), annotated):
|
||||
raise SystemExit(f"Не удалось сохранить {destination}")
|
||||
|
||||
if result.boxes is None or len(result.boxes) == 0:
|
||||
print("Объекты не найдены")
|
||||
else:
|
||||
for box in result.boxes:
|
||||
class_id = int(box.cls.item())
|
||||
confidence = float(box.conf.item())
|
||||
x1, y1, x2, y2 = (round(float(v), 1) for v in box.xyxy[0])
|
||||
print(f"{result.names[class_id]}: {confidence:.3f}, bbox=({x1}, {y1}, {x2}, {y2})")
|
||||
print(f"Изображение с классами: {destination}")
|
||||
|
||||
if args.show:
|
||||
cv2.imshow("Module 3 detection - press any key", annotated)
|
||||
cv2.waitKey(0)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,108 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Обучение YOLO11 на размеченном датасете инструментов."""
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
TRAINING = HERE / "training"
|
||||
|
||||
|
||||
def local_path(value):
|
||||
path = Path(value).expanduser()
|
||||
return path if path.is_absolute() else HERE / path
|
||||
|
||||
|
||||
def select_device(requested, torch):
|
||||
if requested == "cpu":
|
||||
return "cpu"
|
||||
if requested in ("cuda", "gpu"):
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA/ROCm не найден. Используйте --device cpu или установите PyTorch для GPU.")
|
||||
return 0
|
||||
if requested == "auto":
|
||||
if torch.cuda.is_available():
|
||||
return 0 # NVIDIA CUDA и AMD ROCm используют этот интерфейс PyTorch.
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
return "cpu"
|
||||
return requested
|
||||
|
||||
|
||||
def arguments():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--data", default="training/data.yaml", help="data.yaml из Roboflow")
|
||||
parser.add_argument("--model", default="yolo11n.pt", help="Стартовые веса Ultralytics")
|
||||
parser.add_argument("--epochs", type=int, default=30)
|
||||
parser.add_argument(
|
||||
"--patience",
|
||||
type=int,
|
||||
default=7,
|
||||
help="Остановиться после N эпох без улучшения validation metrics",
|
||||
)
|
||||
parser.add_argument("--imgsz", type=int, default=512)
|
||||
parser.add_argument("--batch", type=int, default=8)
|
||||
parser.add_argument("--workers", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--device",
|
||||
default="auto",
|
||||
help="auto, cpu, cuda или номер GPU (например 0)",
|
||||
)
|
||||
parser.add_argument("--name", default="tools_yolo11n")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = arguments()
|
||||
try:
|
||||
import torch
|
||||
from ultralytics import YOLO
|
||||
except ImportError as error:
|
||||
raise SystemExit(f"Нет зависимости {error.name}. Выполните: pip install -r requirements.txt")
|
||||
|
||||
data = local_path(args.data)
|
||||
if not data.is_file():
|
||||
raise SystemExit(f"Не найден датасет: {data}\nРаспакуйте YOLO11 export из Roboflow в module3/dataset.")
|
||||
device = select_device(args.device, torch)
|
||||
if device == 0:
|
||||
backend = "ROCm" if getattr(torch.version, "hip", None) else "CUDA"
|
||||
print(f"Устройство: {backend} GPU — {torch.cuda.get_device_name(0)}")
|
||||
else:
|
||||
print(f"Устройство: {device}")
|
||||
|
||||
model = YOLO(args.model)
|
||||
result = model.train(
|
||||
data=str(data),
|
||||
epochs=args.epochs,
|
||||
patience=args.patience,
|
||||
imgsz=args.imgsz,
|
||||
batch=args.batch,
|
||||
workers=args.workers,
|
||||
device=device,
|
||||
project=str(TRAINING / "runs"),
|
||||
name=args.name,
|
||||
seed=42,
|
||||
deterministic=True,
|
||||
hsv_h=0.0,
|
||||
hsv_s=0.0,
|
||||
hsv_v=0.0,
|
||||
degrees=12.0,
|
||||
translate=0.05,
|
||||
scale=0.15,
|
||||
shear=2.0,
|
||||
fliplr=0.5,
|
||||
flipud=0.5,
|
||||
)
|
||||
save_dir = Path(result.save_dir)
|
||||
best_weights = save_dir / "weights" / "best.pt"
|
||||
model_copy = HERE / "models" / "latest.pt"
|
||||
model_copy.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(best_weights, model_copy)
|
||||
print(f"Лучшие веса: {best_weights}")
|
||||
print(f"Модель для test_model.py: {model_copy}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,6 @@
|
||||
*.zip
|
||||
.venv
|
||||
runs
|
||||
*.cache
|
||||
test
|
||||
valid
|
||||
@@ -0,0 +1,5 @@
|
||||
# chvt-2026 > v0.1.0
|
||||
https://universe.roboflow.com/jons-workspace-8viwr/chvt-2026
|
||||
|
||||
Provided by a Roboflow user
|
||||
License: CC BY 4.0
|
||||
@@ -0,0 +1,27 @@
|
||||
|
||||
chvt-2026 - v1 v0.1.0
|
||||
==============================
|
||||
|
||||
This dataset was exported via roboflow.com on September 5, 2026 at 7:40 PM GMT
|
||||
|
||||
Roboflow is an end-to-end computer vision platform that helps you
|
||||
* collaborate with your team on computer vision projects
|
||||
* collect & organize images
|
||||
* understand and search unstructured image data
|
||||
* annotate, and create datasets
|
||||
* export, train, and deploy computer vision models
|
||||
* use active learning to improve your dataset over time
|
||||
|
||||
For state of the art Computer Vision training notebooks you can use with this dataset,
|
||||
visit https://github.com/roboflow/notebooks
|
||||
|
||||
To find over 100k other datasets and pre-trained models, visit https://universe.roboflow.com
|
||||
|
||||
The dataset includes 40 images.
|
||||
Chvt-2026 are annotated in YOLOv11 format.
|
||||
|
||||
The following pre-processing was applied to each image:
|
||||
* Auto-orientation of pixel data (with EXIF-orientation stripping)
|
||||
* Resize to 512x512 (Stretch)
|
||||
|
||||
No image augmentation techniques were applied.
|
||||
@@ -0,0 +1,18 @@
|
||||
train: train/images
|
||||
val: valid/images
|
||||
test: test/images
|
||||
|
||||
names:
|
||||
0: hammer
|
||||
1: pliers
|
||||
2: wrench
|
||||
|
||||
# nc: 3
|
||||
# names: ['hammer', 'pliers', 'wrench']
|
||||
|
||||
# roboflow:
|
||||
# workspace: jons-workspace-8viwr
|
||||
# project: chvt-2026
|
||||
# version: 1
|
||||
# license: CC BY 4.0
|
||||
# url: https://universe.roboflow.com/jons-workspace-8viwr/chvt-2026/dataset/1
|
||||
@@ -0,0 +1,9 @@
|
||||
# пайторч без куда для амд
|
||||
```
|
||||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
```
|
||||
|
||||
# база
|
||||
```
|
||||
pip install -r ../requirements.txt
|
||||
```
|
||||
|
After Width: | Height: | Size: 20 KiB |
|
After Width: | Height: | Size: 20 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 18 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 20 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 22 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 21 KiB |
|
After Width: | Height: | Size: 21 KiB |
@@ -0,0 +1,4 @@
|
||||
0 0.663086 0.459961 0.060547 0.111328
|
||||
1 0.326172 0.395508 0.046875 0.115234
|
||||
2 0.494141 0.460938 0.074219 0.070312
|
||||
1 0.609375 0.326172 0.070312 0.093750
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.326172 0.395508 0.046875 0.115234
|
||||
1 0.608398 0.327148 0.068359 0.095703
|
||||
0 0.663086 0.459961 0.064453 0.111328
|
||||
2 0.494141 0.460938 0.074219 0.070312
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.330078 0.379883 0.046875 0.111328
|
||||
1 0.619141 0.338867 0.066406 0.099609
|
||||
2 0.495117 0.460938 0.076172 0.070312
|
||||
0 0.663086 0.475586 0.068359 0.107422
|
||||
@@ -0,0 +1,4 @@
|
||||
2 0.494141 0.458008 0.078125 0.064453
|
||||
1 0.3359375 0.3671875 0.0546875 0.11328125
|
||||
1 0.625000 0.348633 0.062500 0.103516
|
||||
0 0.663086 0.488281 0.072266 0.105469
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.336914 0.359375 0.056641 0.109375
|
||||
2 0.495117 0.458008 0.080078 0.060547
|
||||
1 0.628906 0.353516 0.062500 0.105469
|
||||
0 0.662109 0.496094 0.074219 0.101562
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.340820 0.351562 0.056641 0.109375
|
||||
2 0.495117 0.457031 0.080078 0.058594
|
||||
1 0.632812 0.360352 0.062500 0.107422
|
||||
0 0.662109 0.503906 0.074219 0.101562
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.344727 0.343750 0.060547 0.109375
|
||||
2 0.496094 0.457031 0.082031 0.054688
|
||||
1 0.635742 0.365234 0.056641 0.109375
|
||||
0 0.660156 0.510742 0.074219 0.099609
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.349609 0.334961 0.066406 0.107422
|
||||
1 0.638672 0.373047 0.054688 0.109375
|
||||
2 0.496094 0.457031 0.082031 0.050781
|
||||
0 0.660156 0.518555 0.074219 0.091797
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.348633 0.334961 0.064453 0.107422
|
||||
1 0.639648 0.373047 0.056641 0.109375
|
||||
2 0.496094 0.457031 0.082031 0.050781
|
||||
0 0.658203 0.520508 0.078125 0.095703
|
||||
@@ -0,0 +1,4 @@
|
||||
2 0.496094 0.457031 0.082031 0.050781
|
||||
0 0.658203 0.520508 0.078125 0.095703
|
||||
1 0.639648 0.373047 0.056641 0.109375
|
||||
1 0.348633 0.334961 0.064453 0.107422
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.349609 0.335938 0.062500 0.105469
|
||||
2 0.496094 0.457031 0.082031 0.050781
|
||||
1 0.639648 0.371094 0.056641 0.109375
|
||||
0 0.660156 0.517578 0.074219 0.093750
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.345703 0.342773 0.058594 0.107422
|
||||
1 0.636719 0.367188 0.054688 0.109375
|
||||
2 0.496094 0.456055 0.082031 0.052734
|
||||
0 0.661133 0.511719 0.076172 0.097656
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.340820 0.350586 0.056641 0.107422
|
||||
1 0.632812 0.360352 0.062500 0.107422
|
||||
2 0.495117 0.458008 0.080078 0.060547
|
||||
0 0.662109 0.505859 0.074219 0.097656
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.329102 0.379883 0.048828 0.111328
|
||||
1 0.616211 0.334961 0.068359 0.099609
|
||||
2 0.495117 0.458984 0.076172 0.070312
|
||||
0 0.664062 0.473633 0.070312 0.107422
|
||||
@@ -0,0 +1,4 @@
|
||||
2 0.494141 0.459961 0.074219 0.072266
|
||||
1 0.328125 0.388672 0.046875 0.113281
|
||||
1 0.614258 0.332031 0.072266 0.097656
|
||||
0 0.663086 0.467773 0.068359 0.107422
|
||||
@@ -0,0 +1,4 @@
|
||||
2 0.494141 0.460938 0.074219 0.074219
|
||||
1 0.608398 0.325195 0.072266 0.095703
|
||||
1 0.325195 0.398438 0.044922 0.113281
|
||||
0 0.663086 0.458984 0.060547 0.109375
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.32421875 0.40625 0.046875 0.1171875
|
||||
1 0.602539 0.319336 0.072266 0.091797
|
||||
2 0.493164 0.460938 0.072266 0.074219
|
||||
0 0.664062 0.449219 0.058594 0.109375
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.322266 0.416016 0.042969 0.113281
|
||||
1 0.597656 0.315430 0.074219 0.087891
|
||||
2 0.494141 0.461914 0.070312 0.080078
|
||||
0 0.663086 0.441406 0.056641 0.113281
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.319336 0.458984 0.041016 0.113281
|
||||
2 0.492188 0.464844 0.062500 0.089844
|
||||
1 0.573242 0.296875 0.083984 0.078125
|
||||
0 0.656250 0.406250 0.046875 0.117188
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.319336 0.462891 0.041016 0.113281
|
||||
1 0.569336 0.295898 0.080078 0.076172
|
||||
2 0.491211 0.464844 0.060547 0.089844
|
||||
0 0.657227 0.401367 0.044922 0.115234
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.320312 0.463867 0.042969 0.115234
|
||||
1 0.569336 0.293945 0.080078 0.076172
|
||||
2 0.491211 0.464844 0.060547 0.089844
|
||||
0 0.656250 0.400391 0.042969 0.117188
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.321289 0.454102 0.041016 0.115234
|
||||
1 0.576172 0.296875 0.082031 0.078125
|
||||
2 0.494141 0.463867 0.062500 0.091797
|
||||
0 0.660156 0.407227 0.046875 0.115234
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.321289 0.442383 0.041016 0.115234
|
||||
1 0.584961 0.301758 0.080078 0.083984
|
||||
0 0.664062 0.416992 0.050781 0.115234
|
||||
2 0.495117 0.461914 0.064453 0.087891
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.326172 0.419922 0.042969 0.117188
|
||||
1 0.603516 0.311523 0.078125 0.091797
|
||||
2 0.500977 0.462891 0.068359 0.078125
|
||||
0 0.669922 0.438477 0.058594 0.115234
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.330078 0.407227 0.046875 0.119141
|
||||
1 0.615234 0.318359 0.074219 0.097656
|
||||
2 0.502930 0.461914 0.072266 0.080078
|
||||
0 0.674805 0.450195 0.060547 0.111328
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.350586 0.360352 0.056641 0.115234
|
||||
2 0.512695 0.465820 0.080078 0.064453
|
||||
1 0.653320 0.357422 0.064453 0.113281
|
||||
0 0.680664 0.506836 0.076172 0.103516
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.353516 0.355469 0.058594 0.117188
|
||||
1 0.655273 0.363281 0.060547 0.113281
|
||||
2 0.514648 0.465820 0.080078 0.060547
|
||||
0 0.681641 0.514648 0.078125 0.103516
|
||||
@@ -0,0 +1,4 @@
|
||||
1 0.353516 0.355469 0.058594 0.117188
|
||||
1 0.656250 0.363281 0.062500 0.113281
|
||||
2 0.513672 0.465820 0.082031 0.060547
|
||||
0 0.681641 0.517578 0.078125 0.105469
|
||||