working_mask_rcnn_node

#!/usr/bin/env python
import os
import threading
from Queue import Queue
import numpy as np

import cv2
from cv_bridge import CvBridge
import rospy
from sensor_msgs.msg import Image
from sensor_msgs.msg import RegionOfInterest
from sensor_msgs.msg import CameraInfo
from std_msgs.msg import UInt8MultiArray

from mask_rcnn_ros import coco
from mask_rcnn_ros import utils
from mask_rcnn_ros import model as modellib
from mask_rcnn_ros import visualize
from mask_rcnn_ros.msg import Result

# Local path to trained weights file
ROS_HOME = os.environ.get('ROS_HOME', os.path.join(os.environ['HOME'], '.ros'))
COCO_MODEL_PATH = os.path.join(ROS_HOME, 'mask_rcnn_coco.h5')

# COCO Class names
# Index of the class in the list is its ID. For example, to get ID of
# the teddy bear class, use: CLASS_NAMES.index('teddy bear')
CLASS_NAMES = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
               'bus', 'train', 'truck', 'boat', 'traffic light',
               'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
               'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
               'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
               'kite', 'baseball bat', 'baseball glove', 'skateboard',
               'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
               'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
               'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
               'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
               'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
               'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
               'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
               'teddy bear', 'hair drier', 'toothbrush']


class InferenceConfig(coco.CocoConfig):
    # Set batch size to 1 since we'll be running inference on
    # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

class MaskRCNNNode(object):
    def __init__(self):
        self._cv_bridge = CvBridge()

        config = InferenceConfig()
        config.display()

        self._visualization = rospy.get_param('~visualization', True)

        # Create model object in inference mode.
        self._model = modellib.MaskRCNN(mode="inference", model_dir="",
                                        config=config)
        # Load weights trained on MS-COCO
        model_path = rospy.get_param('~model_path', COCO_MODEL_PATH)
        # Download COCO trained weights from Releases if needed
        if model_path == COCO_MODEL_PATH and not os.path.exists(COCO_MODEL_PATH):
            utils.download_trained_weights(COCO_MODEL_PATH)

        self._model.load_weights(model_path, by_name=True)

        self._class_names = rospy.get_param('~class_names', CLASS_NAMES)

        self._last_msg = None
        self._msg_lock = threading.Lock()

        self._last_info = None
        self._info_lock = threading.Lock()

        self._last_depth = None
        self._depth_lock = threading.Lock()

        self._class_colors = visualize.random_colors(len(CLASS_NAMES))

        self._publish_rate = rospy.get_param('~publish_rate', 10)

    def run(self):
        self._result_pub = rospy.Publisher('~result', Result, queue_size=1)
        vis_pub = rospy.Publisher('/kinect2/seg/image_raw', Image, queue_size=1)
        info_pub = rospy.Publisher('kinect2/seg/camera_info', CameraInfo, queue_size=1)
        depth_pub = rospy.Publisher('kinect2/seg/image_depth_rect', Image, queue_size=1)
        sub = rospy.Subscriber('/kinect2/qhd/image_color_rect', Image,
                               self._image_callback, queue_size=1)
        sub_info = rospy.Subscriber('/kinect2/qhd/camera_info', CameraInfo,
                               self._info_callback, queue_size=1)
        sub_depth = rospy.Subscriber('/kinect2/qhd/image_depth_rect', Image,
                               self._depth_callback, queue_size=1)

        rate = rospy.Rate(self._publish_rate)
        while not rospy.is_shutdown():
            if self._msg_lock.acquire(False):
                msg = self._last_msg
                self._last_msg = None
                self._msg_lock.release()
            else:
                rate.sleep()
                continue

            if self._info_lock.acquire(False):
                info = self._last_info
                self._last_info = None
                self._info_lock.release()
            else:
                rate.sleep()
                continue

            if self._depth_lock.acquire(False):
                depth = self._last_depth
                self._last_depth = None
                self._depth_lock.release()
            else:
                rate.sleep()
                continue

            if msg is not None and info is not None and depth is not None:
                np_image = self._cv_bridge.imgmsg_to_cv2(msg, "bgr8")
                np_depth = self._cv_bridge.imgmsg_to_cv2(depth, "passthrough")
                
                
                #np_image = cv2.resize(np_image, (960, 540))
                
                # Run detection
                results = self._model.detect([np_image], verbose=0)
                result = results[0]
                result_msg = self._build_result_msg(msg, result)
                self._result_pub.publish(result_msg)

                # Visualize results
                if self._visualization:
                    
                    #vis_image = self._visualize(result, np_image)
                    np_image = display_instances(np_image, result['rois'], result['masks'], result['class_ids'], CLASS_NAMES, result['scores'])
                    #cv_result = np.zeros(shape=vis_image.shape, dtype=np.uint8)
                    #cv2.convertScaleAbs(vis_image, cv_result)
                    image_msg = self._cv_bridge.cv2_to_imgmsg(np_image, "bgr8")
                    depth_msg = self._cv_bridge.cv2_to_imgmsg(np_depth, "mono16")

                    curr_time = rospy.Time.now()
                    #pack_msg = Image()
                    #pack_msg.format = "jpeg"
                    image_msg.header.stamp = curr_time
                    image_msg.header.frame_id = 'kinect2_link'
                    #pack_msg.data = np.array(cv2.imencode('.jpg', np_image)[1]).tostring()
                    #pack_msg.encoding = "bgr8"

                    vis_pub.publish(image_msg)

                    #info_msg = CameraInfo()
                    info.header.stamp = curr_time
                    #info_msg.data = info
                    info_pub.publish(info)

                    #deppack = Image()
                    depth_msg.header.stamp = curr_time
                    #deppack.data = 
                    depth_pub.publish(depth_msg) 

            rate.sleep()

    def _build_result_msg(self, msg, result):
        result_msg = Result()
        result_msg.header = msg.header
        for i, (y1, x1, y2, x2) in enumerate(result['rois']):
            box = RegionOfInterest()
            box.x_offset = np.asscalar(x1)
            box.y_offset = np.asscalar(y1)
            box.height = np.asscalar(y2 - y1)
            box.width = np.asscalar(x2 - x1)
            result_msg.boxes.append(box)

            class_id = result['class_ids'][i]
            result_msg.class_ids.append(class_id)

            class_name = self._class_names[class_id]
            result_msg.class_names.append(class_name)

            score = result['scores'][i]
            result_msg.scores.append(score)

            mask = Image()
            mask.header = msg.header
            mask.height = result['masks'].shape[0]
            mask.width = result['masks'].shape[1]
            mask.encoding = "mono8"
            mask.is_bigendian = False
            mask.step = mask.width
            mask.data = (result['masks'][:, :, i] * 255).tobytes()
            result_msg.masks.append(mask)
        return result_msg

    def _visualize(self, result, image):
        from matplotlib.backends.backend_agg import FigureCanvasAgg
        from matplotlib.figure import Figure

        fig = Figure()
        canvas = FigureCanvasAgg(fig)
        axes = fig.gca()
        visualize.display_instances(image, result['rois'], result['masks'],
                                    result['class_ids'], CLASS_NAMES,
                                    result['scores'], ax=axes,
                                    class_colors=self._class_colors)
        fig.tight_layout()
        canvas.draw()
        result = np.fromstring(canvas.tostring_rgb(), dtype='uint8')

        _, _, w, h = fig.bbox.bounds
        result = result.reshape((int(h), int(w), 3))
        return result

    def _image_callback(self, msg):
        rospy.logdebug("Get an image")
        if self._msg_lock.acquire(False):
            self._last_msg = msg
            self._msg_lock.release()

    def _info_callback(self, msg):
        rospy.logdebug("Get camera info")
        if self._info_lock.acquire(False):
            self._last_info = msg
            self._info_lock.release()

    def _depth_callback(self, msg):
        rospy.logdebug("Get depth image")
        if self._depth_lock.acquire(False):
            self._last_depth = msg
            self._depth_lock.release()

def main():
    rospy.init_node('mask_rcnn')

    node = MaskRCNNNode()
    node.run()

def display_instances(image, boxes, masks, ids, names, scores):
    """
        take the image and results and apply the mask, box, and Label
    """
    n_instances = boxes.shape[0]
    colors = random_colors(n_instances)
    class_colors = random_colors(len(CLASS_NAMES))
    N = len(colors)

    if not n_instances:
        print('NO INSTANCES TO DISPLAY')
    else:
        assert boxes.shape[0] == masks.shape[-1] == ids.shape[0]

    for i, color in enumerate(colors):
        if not np.any(boxes[i]):
            continue

        y1, x1, y2, x2 = boxes[i]
        label = names[ids[i]]
        score = scores[i] if scores is not None else None
        caption = '{} {:.2f}'.format(label, score) if score else label
        mask = masks[:, :, i]

        if(label == "tv"):
           color = (0, 0, 255)
           image = apply_mask(image, mask, color)
        
        color = class_colors[i]
        #image = apply_mask(image, mask, color)



        #image = cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        #image = cv2.putText(
        #    image, caption, (x1, y1), cv2.FONT_HERSHEY_COMPLEX, 0.7, color, 2
        #)

    return image

def random_colors(N):
    np.random.seed(1)
    colors = [tuple(255 * np.random.rand(3)) for _ in range(N)]
    return colors

def apply_mask(image, mask, color, alpha=0.9):
    """apply mask to image"""
    for n, c in enumerate(color):
        image[:, :, n] = np.where(
            mask == 1,
            image[:, :, n] * (1 - alpha) + alpha * c,
            image[:, :, n]
        )
    return image

if __name__ == '__main__':
    main()