AI Computer Vision System for Robots

Разработка Computer Vision для промышленных роботов

Промышленный манипулятор без зрения — это слепой автомат, работающий строго по заданной программе в идеально стерильных условиях. CV добавляет адаптивность: захват произвольно ориентированных деталей (bin picking), инспекция в процессе сборки, навигация мобильных роботов (AMR/AGV) в динамической среде, совместная работа с человеком (collaborative robots, cobot). Ключевые алгоритмы: 6DoF pose estimation для захвата деталей, depth-guided grasping, semantic map building для навигации.

6DoF Pose Estimation для bin picking

import numpy as np
import cv2
import torch
from dataclasses import dataclass
from typing import Optional

@dataclass
class ObjectPose:
    object_class: str
    position_xyz: tuple[float, float, float]    # мм в системе координат камеры
    rotation_matrix: np.ndarray                  # 3×3
    euler_angles: tuple[float, float, float]     # roll, pitch, yaw в градусах
    confidence: float
    grasp_point: tuple[float, float, float]      # рекомендованная точка захвата
    grasp_approach: np.ndarray                   # вектор подхода захвата

class BinPickingSystem:
    """
    Система bin picking: обнаружение и определение позы деталей в контейнере.
    Методы:
    1. FoundationPose / DenseFusion: на основе RGB-D
    2. GDR-Net: геометрически дискретизированный рендеринг
    3. PVPN: point-wise voting
    Камера: Intel RealSense D435i или Azure Kinect.
    CAD модель детали обязательна для FoundationPose.
    """
    def __init__(self, pose_model_path: str,
                  cad_model_path: str,
                  object_classes: list[str],
                  camera_intrinsics: dict,
                  device: str = 'cuda'):
        self.device = device
        self.object_classes = object_classes
        self.camera_intrinsics = camera_intrinsics  # fx, fy, cx, cy

        # Загрузка pose estimation модели
        self.pose_model = torch.load(pose_model_path,
                                      map_location=device).eval()

        # CAD модель для рендеринга (используется FoundationPose)
        self.cad_model = self._load_cad_model(cad_model_path)

        # YOLO для первичного обнаружения объектов
        from ultralytics import YOLO
        self.detector = YOLO(pose_model_path.replace('pose', 'det'))

    def _load_cad_model(self, cad_path: str):
        """Загрузка .ply или .obj CAD модели"""
        try:
            import open3d as o3d
            return o3d.io.read_triangle_mesh(cad_path)
        except ImportError:
            return None

    def estimate_poses(self, rgb: np.ndarray,
                        depth: np.ndarray) -> list[ObjectPose]:
        """
        Оценка позы объектов.
        rgb: (H, W, 3) uint8
        depth: (H, W) float32 в миллиметрах
        """
        # 1. Детекция объектов для ROI
        detections = self.detector(rgb, conf=0.4, verbose=False)

        poses = []
        for box in detections[0].boxes:
            cls_id = int(box.cls.item())
            if cls_id >= len(self.object_classes):
                continue

            x1, y1, x2, y2 = map(int, box.xyxy[0])

            # Вырезать RGB и depth патч
            rgb_crop = rgb[y1:y2, x1:x2]
            depth_crop = depth[y1:y2, x1:x2]

            if rgb_crop.size == 0:
                continue

            # 2. Pose estimation на патче
            pose = self._estimate_single_pose(
                rgb_crop, depth_crop, cls_id, (x1, y1)
            )
            if pose:
                poses.append(pose)

        # Сортировка по высоте Z (верхние детали первыми)
        poses.sort(key=lambda p: p.position_xyz[2])
        return poses

    @torch.no_grad()
    def _estimate_single_pose(self, rgb_crop: np.ndarray,
                               depth_crop: np.ndarray,
                               cls_id: int,
                               offset: tuple) -> Optional[ObjectPose]:
        """Pose estimation для одного объекта"""
        from torchvision import transforms
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                  [0.229, 0.224, 0.225])
        ])
        from PIL import Image
        pil = Image.fromarray(cv2.cvtColor(rgb_crop, cv2.COLOR_BGR2RGB))
        rgb_tensor = transform(pil).unsqueeze(0).to(self.device)
        depth_tensor = torch.from_numpy(depth_crop).unsqueeze(0).unsqueeze(0).float().to(self.device)

        # Конкатенация RGB + depth
        depth_norm = depth_tensor / 1000.0  # мм → метры
        # Упрощённая модель принимает 4-channel input
        input_tensor = torch.cat([
            rgb_tensor,
            torch.nn.functional.interpolate(
                depth_norm, size=rgb_tensor.shape[-2:], mode='bilinear'
            )
        ], dim=1)

        output = self.pose_model(input_tensor)
        # output: (1, 6) — translation(3) + rotation_euler(3)
        if output is None or output.shape[-1] < 6:
            return None

        out_np = output.squeeze().cpu().numpy()
        tx, ty, tz = out_np[:3] * 1000  # метры → мм
        rx, ry, rz = np.degrees(out_np[3:6])

        # Матрица вращения из эйлеровых углов
        R, _ = cv2.Rodrigues(np.array([np.radians(rx),
                                        np.radians(ry),
                                        np.radians(rz)]))

        # Точка захвата: центр объекта + смещение вверх по нормали
        grasp_z = tz - 30  # 30мм выше центра
        grasp_point = (tx, ty, grasp_z)
        approach_vec = R @ np.array([0, 0, -1])  # направление подхода

        conf = float(torch.sigmoid(
            self.pose_model.confidence_head(output) if hasattr(
                self.pose_model, 'confidence_head') else torch.tensor(0.0)
        ).item()) if hasattr(self.pose_model, 'confidence_head') else 0.8

        return ObjectPose(
            object_class=self.object_classes[cls_id],
            position_xyz=(round(tx, 1), round(ty, 1), round(tz, 1)),
            rotation_matrix=R,
            euler_angles=(round(rx, 1), round(ry, 1), round(rz, 1)),
            confidence=round(conf, 3),
            grasp_point=grasp_point,
            grasp_approach=approach_vec
        )


class AMRNavigationVision:
    """
    Computer vision для автономных мобильных роботов (AMR).
    Задачи: obstacle detection, semantic mapping, человек-распознавание
    для cobot safety (ISO/TS 15066 protected/restricted speed zones).
    """
    def __init__(self, obstacle_model_path: str,
                  device: str = 'cuda'):
        from ultralytics import YOLO
        self.obstacle_model = YOLO(obstacle_model_path)
        self.device = device
        # Semantic map: {cell_id: label}
        self.semantic_map: dict = {}

    def process_navigation_frame(self, rgb: np.ndarray,
                                   depth: np.ndarray) -> dict:
        """
        Анализ кадра для навигации AMR.
        Возвращает: obstacles, nearest_human_dist_m, clear_path_sectors.
        """
        results = self.obstacle_model(rgb, conf=0.4, verbose=False)
        obstacles = []
        nearest_human_dist = float('inf')

        h, w = depth.shape[:2]
        sector_width = w // 5  # 5 секторов: LL/L/C/R/RR

        for box in results[0].boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cls_name = results[0].names[int(box.cls.item())]
            cx = (x1 + x2) // 2
            cy = (y1 + y2) // 2

            # Медианная глубина в bbox
            depth_crop = depth[y1:y2, x1:x2]
            valid_depths = depth_crop[depth_crop > 0]
            dist_m = float(np.median(valid_depths)) / 1000.0 if len(valid_depths) > 0 else 0

            obstacles.append({
                'class': cls_name,
                'bbox': [x1, y1, x2, y2],
                'distance_m': round(dist_m, 2),
                'sector': min(cx // sector_width, 4)
            })

            if cls_name == 'person' and dist_m < nearest_human_dist:
                nearest_human_dist = dist_m

        # Определить свободные секторы
        blocked_sectors = {o['sector'] for o in obstacles if o['distance_m'] < 1.5}
        clear_sectors = [s for s in range(5) if s not in blocked_sectors]

        # ISO/TS 15066: если человек < 0.5м → STOP; 0.5–1.5м → reduced speed
        safety_mode = ('STOP' if nearest_human_dist < 0.5
                       else 'REDUCED_SPEED' if nearest_human_dist < 1.5
                       else 'NORMAL')

        return {
            'obstacles': obstacles,
            'nearest_human_m': round(nearest_human_dist, 2),
            'clear_sectors': clear_sectors,
            'safety_mode': safety_mode
        }

Важно для production: синхронизация с PLC обязательна — CV система должна отправлять результаты в контроллер робота через ROS2 topic или OPC-UA с задержкой менее 50мс. Для bin picking добавляется проверка collision-free trajectory перед командой захвата.