color_primitive

color_detection.py

import cv2
import time

from pypot.primitive import Primitive, LoopPrimitive


rgb_ranges = {
    'green': ((0, 80, 0), (120, 255, 120)),
    'red': ((0, 0, 90), (100, 100, 255)),
}

hsv_min, hsv_max = ((0, 0, 0), (255, 255, 255))


class ColorDetection(LoopPrimitive):


    def crop(self, image):
        return image  # [100:300, 250:450]

    def mean_color(self, image):
        return image.mean(axis=(0, 1))

    def get_color(self):
        img = self.robot.camera.frame
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        hsv_mask = cv2.inRange(hsv, hsv_min, hsv_max)

        score = {}

        for color, (lower, upper) in rgb_ranges.items():
            mask = cv2.inRange(img, lower, upper)
            mask[hsv_mask < 255] = 0
            mask = cv2.erode(mask, None, iterations=4)

            score[color] = mask.sum()

        return max(score, key=score.get)

    def update(self):
        self.__centroid = self._get_centroid()
        
    @property
    def centroid(self):
        return self.__centroid

    def _get_centroid(self, color=None, threshold=4000):
        img = self.robot.camera.frame
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        hsv_mask = cv2.inRange(hsv, hsv_min, hsv_max)

        score = {}
        masks = {}
        if color is not None:
            mask = cv2.inRange(img, rgb_ranges[color][0], rgb_ranges[color][1])
            mask[hsv_mask < 255] = 0
            mask = cv2.erode(mask, None, iterations=4)

        else:
            for _color, (lower, upper) in rgb_ranges.items():
                rgb_mask = cv2.inRange(img, lower, upper)
                tmp_mask = rgb_mask.copy()
                tmp_mask[hsv_mask < 255] = 0
                masks[_color] = tmp_mask
                masks[_color] = cv2.erode(masks[_color], None, iterations=4)
                score[_color] = masks[_color].sum()
            color = max(score, key=score.get)
            mask = masks[color]
        if score[color] < threshold:
            return "0;0"
        else:
            moments = cv2.moments(mask, True)
            cx = moments['m10'] / moments['m00']
            cy = moments['m01'] / moments['m00']
            return "%s;%s;%s;%s" % (cx * 100 / img.shape[1], cy * 100 / img.shape[0], color, score[color])

poppy-ergo-jr.py

from numpy import sum
from functools import partial

from poppy.creatures import AbstractPoppyCreature
from poppy.creatures.ik import IKChain

from .primitives.dance import Dance
from .primitives.face_tracking import FaceTracking
from .primitives.tracking_feedback import TrackingFeedback
from .primitives.postures import (BasePosture, RestPosture,
                                  CuriousPosture, TetrisPosture,
                                  SafePowerUp)

from .primitives.color_detection import ColorDetection


class PoppyErgoJr(AbstractPoppyCreature):
    @classmethod
    def setup(cls, robot):
        robot._primitive_manager._filter = partial(sum, axis=0)

        c = IKChain.from_poppy_creature(robot,
                                        motors=robot.motors,
                                        passiv=[],
                                        tip=[0, 0, -0.07])

        robot.chain = c

        robot.attach_primitive(SafePowerUp(robot), 'safe_power_up')

        robot.attach_primitive(Dance(robot), 'dance')

        robot.attach_primitive(BasePosture(robot, 2.), 'base_posture')
        robot.attach_primitive(RestPosture(robot, 2.), 'rest_posture')
        robot.attach_primitive(CuriousPosture(robot, 2.), 'curious_posture')
        robot.attach_primitive(TetrisPosture(robot, 2.), 'tetris_posture')

        if not robot.simulated and hasattr(robot, 'marker_detector'):
            robot.attach_primitive(TrackingFeedback(robot, 25.),
                                   'tracking_feedback')
            
        if not robot.simulated:
            robot.attach_primitive(ColorDetection(robot, 2), 'color_detection')
            

        for m in robot.motors:
            m.pid = (4, 2, 0)
            m.torque_limit = 70.
            m.led = 'off'

        if not robot.simulated and hasattr(robot, 'face_tracking'):
            robot.attach_primitive(FaceTracking(robot, 10,
                                                robot.face_detector),
                                   'face_tracking')

Setup the filtering ranges for color detection.ipynb