Conception/drake-master/bindings/pydrake/systems/test/perception_test.py

# -*- coding: utf-8 -*-

import unittest
import numpy as np

from pydrake.common.value import AbstractValue
from pydrake.math import RigidTransform, RollPitchYaw, RotationMatrix
from pydrake.perception import BaseField, Fields, PointCloud
from pydrake.systems.analysis import Simulator
from pydrake.systems.framework import DiagramBuilder
from pydrake.systems.perception import (
    PointCloudConcatenation, _ConcatenatePointClouds, _TileColors,
    _TransformPoints)


class TestConcatenatePointClouds(unittest.TestCase):
    def setUp(self):
        self.points_0 = np.array([[1.0], [2.0], [3.0]])
        self.colors_0 = np.array([[0], [128], [255]])

        self.points_1 = np.array([[4.0], [5.0], [6.0]])
        self.colors_1 = np.array([[50], [100], [200]])

        self.points_dict = {"0": self.points_0, "1": self.points_1}
        self.colors_dict = {"0": self.colors_0, "1": self.colors_1}

    def test_concatenation(self):
        scene_points, scene_colors = _ConcatenatePointClouds(
            self.points_dict, self.colors_dict)

        self.assertEqual(scene_points.shape, (3, len(self.points_dict)))
        self.assertEqual(scene_colors.shape, (3, len(self.colors_dict)))
        self.assertEqual(scene_points.shape, scene_colors.shape)

        for i, value in enumerate(self.points_0.flatten()):
            self.assertTrue(value in scene_points[i, :])

        for i, value in enumerate(self.points_1.flatten()):
            self.assertTrue(value in scene_points[i, :])

        for i, value in enumerate(self.colors_0.flatten()):
            self.assertTrue(value in scene_colors[i, :])

        for i, value in enumerate(self.colors_0.flatten()):
            self.assertTrue(value in scene_colors[i, :])


class TestTileColors(unittest.TestCase):
    def setUp(self):
        self.red = [255, 0, 0]
        self.blue = [0, 0, 255]

    def test_one_dim(self):
        tiled = _TileColors(self.red, 1)
        expected_tiled = np.array([[255], [0], [0]])
        self.assertTrue(np.allclose(tiled, expected_tiled))

    def test_three_dims(self):
        tiled = _TileColors(self.blue, 1)
        expected_tiled = np.array([[0, 0, 0], [0, 0, 0], [255, 255, 255]])
        self.assertTrue(np.allclose(tiled, expected_tiled))


class TestTransformPoints(unittest.TestCase):
    def setUp(self):
        self.points = np.array([[1, 1, 0], [2, 1, 0]]).T
        self.translation = RigidTransform(p=[1, 2, 3])
        self.rotation = RigidTransform(
            RotationMatrix(RollPitchYaw(0, 0, np.pi/2)))

    def test_translation(self):
        transformed_points = _TransformPoints(
            self.points, self.translation.GetAsMatrix4())
        expected_translated_points = np.array([[2, 3, 3], [3, 3, 3]]).T

        self.assertTrue(
            np.allclose(transformed_points, expected_translated_points))

    def test_rotation(self):
        transformed_points = _TransformPoints(
            self.points, self.rotation.GetAsMatrix4())
        expected_rotated_points = np.array([[-1, 1, 0], [-1, 2, 0]]).T

        self.assertTrue(
            np.allclose(transformed_points, expected_rotated_points))


class TestPointCloudConcatenation(unittest.TestCase):
    def setUp(self):
        builder = DiagramBuilder()

        X_WP_0 = RigidTransform.Identity()
        X_WP_1 = RigidTransform.Identity()
        X_WP_1.set_translation([1.0, 0, 0])

        id_list = ["0", "1"]

        self.pc_concat = builder.AddSystem(PointCloudConcatenation(id_list))

        self.num_points = 10000
        xyzs = np.random.uniform(-0.1, 0.1, (3, self.num_points))
        # Only go to 254 to distinguish between point clouds with and without
        # color.
        rgbs = np.random.uniform(0., 254.0, (3, self.num_points))

        self.pc = PointCloud(
            self.num_points,
            Fields(BaseField.kXYZs | BaseField.kRGBs))
        self.pc.mutable_xyzs()[:] = xyzs
        self.pc.mutable_rgbs()[:] = rgbs

        self.pc_no_rgbs = PointCloud(
            self.num_points, Fields(BaseField.kXYZs))
        self.pc_no_rgbs.mutable_xyzs()[:] = xyzs

        diagram = builder.Build()

        simulator = Simulator(diagram)

        self.context = diagram.GetMutableSubsystemContext(
            self.pc_concat, simulator.get_mutable_context())

        self.pc_concat.GetInputPort("X_FCi_0").FixValue(self.context, X_WP_0)
        self.pc_concat.GetInputPort("X_FCi_1").FixValue(self.context, X_WP_1)

    def test_no_rgb(self):
        self.pc_concat.GetInputPort("point_cloud_CiSi_0").FixValue(
            self.context, AbstractValue.Make(self.pc_no_rgbs))
        self.pc_concat.GetInputPort("point_cloud_CiSi_1").FixValue(
            self.context, AbstractValue.Make(self.pc_no_rgbs))

        fused_pc = self.pc_concat.GetOutputPort("point_cloud_FS").Eval(
            self.context)

        self.assertEqual(fused_pc.size(), 2 * self.num_points)

        # The first point cloud should be from [-0.1 to 0.1].
        # The second point cloud should be from [0.9 to 1.1].
        self.assertTrue(np.max(fused_pc.xyzs()[0, :]) >= 1.0)
        self.assertTrue(np.min(fused_pc.xyzs()[0, :]) <= 0.0)

        # Even if both input point clouds don't have rgbs, the fused point
        # cloud should contain rgbs of the default color.
        self.assertTrue(fused_pc.has_rgbs())
        self.assertTrue(
            np.all(fused_pc.rgbs()[:, 0] == np.array([255, 255, 255])))
        self.assertTrue(
            np.all(fused_pc.rgbs()[:, -1] == np.array([255, 255, 255])))

    def test_rgb(self):
        self.pc_concat.GetInputPort("point_cloud_CiSi_0").FixValue(
            self.context, self.pc)
        self.pc_concat.GetInputPort("point_cloud_CiSi_1").FixValue(
            self.context, self.pc)

        fused_pc = self.pc_concat.GetOutputPort("point_cloud_FS").Eval(
            self.context)

        self.assertEqual(fused_pc.size(), 2 * self.num_points)

        # The first point cloud should be from [-0.1 to 0.1].
        # The second point cloud should be from [0.9 to 1.1].
        self.assertTrue(np.max(fused_pc.xyzs()[0, :]) >= 1.0)
        self.assertTrue(np.min(fused_pc.xyzs()[0, :]) <= 0.0)

        self.assertTrue(fused_pc.has_rgbs())
        self.assertTrue(
            np.all(fused_pc.rgbs()[:, 0] != np.array([255, 255, 255])))
        self.assertTrue(
            np.all(fused_pc.rgbs()[:, -1] != np.array([255, 255, 255])))

    def test_mix_rgb(self):
        self.pc_concat.GetInputPort("point_cloud_CiSi_0").FixValue(
            self.context, AbstractValue.Make(self.pc))
        self.pc_concat.GetInputPort("point_cloud_CiSi_1").FixValue(
            self.context, AbstractValue.Make(self.pc_no_rgbs))

        fused_pc = self.pc_concat.GetOutputPort("point_cloud_FS").Eval(
            self.context)

        self.assertEqual(fused_pc.size(), 2 * self.num_points)

        # The first point cloud should be from [-0.1 to 0.1].
        # The second point cloud should be from [0.9 to 1.1].
        self.assertTrue(np.max(fused_pc.xyzs()[0, :]) >= 1.0)
        self.assertTrue(np.min(fused_pc.xyzs()[0, :]) <= 0.0)

        self.assertTrue(fused_pc.has_rgbs())

        # We don't know in what order the two point clouds will be combined.
        rgb_first = np.all(fused_pc.rgbs()[:, 0] != np.array([255, 255, 255]))
        rgb_last = np.all(fused_pc.rgbs()[:, -1] != np.array([255, 255, 255]))
        no_rgb_first = np.all(
            fused_pc.rgbs()[:, 0] == np.array([255, 255, 255]))
        no_rgb_last = np.all(
            fused_pc.rgbs()[:, -1] == np.array([255, 255, 255]))

        self.assertTrue(
            (rgb_first and no_rgb_last) or (no_rgb_first and rgb_last))