#!/usr/bin/env python """ \author Yubao Liu \date Dec. 2020 """ from __future__ import division from __future__ import print_function import mrcnn.model as modellib from mrcnn import visualize from mrcnn import utils from tensorflow.python.keras.models import load_model from tensorflow.python.keras.backend import set_session from maskrcnn_ros.msg import * import maskrcnn_ros.msg import actionlib import PIL import matplotlib.pyplot as plt import matplotlib import skimage.io import time import logging import threading import copy import cv2 from skimage.transform import resize import message_filters import numpy as np from cv_bridge import CvBridge, CvBridgeError from sensor_msgs.msg import PointCloud2 from sensor_msgs.msg import Image import rospy import math import random import sys import os os.environ["CUDA_VISIBLE_DEVICES"]="1" import tensorflow as tf config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction=0.8 config.gpu_options.allow_growth = True sess = tf.Session(config=config) # Here should be the project name # logging.basicConfig( level=logging.DEBUG, format='(%(threadName)-9s) %(message)s',) logging.basicConfig(filename='maskrcnn_action_server.log', format='%(asctime)s %(levelname)s:%(message)s', level=logging.INFO, datefmt='%m/%d/%Y %I:%M:%S %p') # Maskrcnn # Root directory of the project # TODO: Set ROOT_DIR in launch file ROOT_DIR = os.path.abspath( "/root/catkin_ws/src/MaskRCNN_ROS/include/MaskRCNN/") # ROOT_DIR = os.path.abspath( "/root/catkin_ws/src/MaskRCNN_ROS/include/MaskRCNN/") print("ROOT_DIR: ", ROOT_DIR) # why cannot use the relative path # ROOT_DIR = os.path.abspath("include/MaskRCNN/") # Import Mask RCNN sys.path.append(ROOT_DIR) # To find local version of the library # from mrcnn import utils # Import COCO config # To find local version sys.path.append(os.path.join(ROOT_DIR, "samples/coco/")) import coco # Directory to save logs and trained model MODEL_DIR = os.path.join(ROOT_DIR, "logs") # Local path to trained weights file # COCO_MODEL_PATH = os.path.join(ROOT_DIR, "model/mask_rcnn_coco.h5") # Directory of images to run detection on IMAGE_DIR = os.path.join(ROOT_DIR, "images") # RESULT_DIR = os.path.join(ROOT_DIR, "results") # 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' ] # Totoal time consuming total_time = float(0.0) # Total number of images total_number = int(0) # To sync semantic segmentation thread and action server thread cn_task = threading.Condition() cn_ready = threading.Condition() batch_size = 2 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 = 2 IMAGE_MIN_DIM = 128 IMAGE_MAX_DIM = 128 class MaskRcnn(object): def __init__(self): print('Init MaskRCNN') # Init Mask rcnn # Initialize Maskrcnn config = InferenceConfig() config.display() # Create model object in inference mode. self.model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config) self.coco_model_path_ = rospy.get_param( 'model_path', '/root/cnnmodel/mask_rcnn_coco.h5') # Load weights trained on MS-COCO # self.model.load_weights(COCO_MODEL_PATH, by_name=True) self.model.load_weights(self.coco_model_path_, by_name=True) # if bPublish_result: self.masked_image__pub = rospy.Publisher("masked_image", Image, queue_size=10) print("=============Initialized MaskRCNN==============") def segment(self, image): timer_start = rospy.Time.now() results = self.model.detect(image, verbose=1) segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000 print("predict time: %f ms \n" % segment_time) timer_start = rospy.Time.now() r = results[0] self.masked_image_ = visualize.ros_semantic_result(image[0], r['rois'], r['masks'], r['class_ids'], class_names, r['scores'], show_opencv=False) segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000 print("Visualize time: %f ms \n" % segment_time) return self.masked_image_, results # visualize.display_instances(color_image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) def publish_result(self, image, stamp): if image is None or stamp is None: print("Image invalid") return # h, w, c = image.shape # assert c == 3 msg_img = CvBridge().cv2_to_imgmsg(image, encoding='passthrough') msg_img.header.stamp = stamp self.masked_image__pub.publish(msg_img) # Segmentation cannot putted into ROS action callback def worker(maskrcnn, bPublish_result=False): # Control whether start maskrcnn thread global bStartMaskRCNN global color_image global stamp global masked_image global result maskrcnn = MaskRcnn() bStartMaskRCNN = True color_image = [] stamp = [] while bStartMaskRCNN: with cn_task: cn_task.wait() print("New task comming") timer_start = rospy.Time.now() masked_image, result = maskrcnn.segment(color_image) segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000 print("MaskRCNN segment time for cuurent image: %f ms \n" % segment_time) with cn_ready: cn_ready.notifyAll() if bPublish_result: maskrcnn.publish_result(masked_image, stamp) print("Exit MaskRCNN thread") class SemanticActionServer(object): _feedback = maskrcnn_ros.msg.batchFeedback() _result = maskrcnn_ros.msg.batchResult() def __init__(self): print("Initialize Action Server") # Get action server name self._action_name = rospy.get_param('/semantic/action_name', '/maskrcnn_action_server') print("Action name: ", self._action_name) # Start Action server self._as = actionlib.SimpleActionServer( self._action_name, maskrcnn_ros.msg.batchAction, execute_cb=self.execute_cb, auto_start=False) self._as.start() # Semantic segentation result # self.label_topic = rospy.get_param('/semantic/semantic_label', # '/semantic_label') # self.label_color_topic = rospy.get_param( # '/semantic/semantic_label_color', '/semantic_label_color') # publish result # self.semantic_label_pub = rospy.Publisher(self.label_topic, # Image, # queue_size=10) # self.semantic_label_color_pub = rospy.Publisher(self.label_color_topic, # Image, # queue_size=10) # self.graph = tf.get_default_graph() # image = cv2.imread('/root/catkin_ws/src/MaskRCNN_ROS/include/MaskRCNN/images/tum_rgb.png') # results = self.model.detect([image], verbose=1) # r = results[0] # visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) print("MaskRCNN action server start...") # Notice: Inconvinient to add MaskRCNN detection here, because this callback running in a seperate thread, # The initializaiotn of MaskRCNN will not work here def execute_cb(self, goal): global total_time global total_number global color_image global stamp global masked_image global result # global bSaveResult # clear the old data color_image = [] stamp = [] if not self._as.is_active(): logging.debug("[Error] MaskRCNN action server cannot active") return time_each_loop_start = rospy.Time.now() print("----------------------------------") print("ID: %d" % goal.id) # batch_size = goal.batch_size # print("batch size: %d" % batch_size) # Receive source image from client for i in range(batch_size): try: color_image.append( CvBridge().imgmsg_to_cv2(goal.image[i], 'bgr8')) stamp.append(goal.image[i].header.stamp) # logging.debug('Color: ', color_image.shape) except CvBridgeError as e: print(e) return # Test image communication # (rows, cols, channels) = color_image[0].shape # print("rows: ", rows, " cols: ", cols) # verify the image received # cv2.imshow("Image window", color_image) # cv2.waitKey(1) timer_start = rospy.Time.now() # perform segmenting if np.any(color_image[0]): with cn_task: print("Inform new task") cn_task.notifyAll() with cn_ready: cn_ready.wait() print("semantic result ready") # save_file = RESULT_DIR +'/'+ str(goal.id) + '.png' # cv2.imwrite(save_file, color_image) # skimage_image = skimage.io.imread(save_file) # print("Shape of skimage: ") # Visualize results # r = results[0] # visualize.display_instances(color_image, r['rois'], r['masks'], r['class_ids'], class_names, r['scores']) # # to ROS msg # msg_label_color = CvBridge().cv2_to_imgmsg(decoded, encoding="bgr8") # msg_label = CvBridge().cv2_to_imgmsg(label, encoding="mono8") # calculate time cost segment_time = (rospy.Time.now() - timer_start).to_sec() * 1000 print("MaskRCNN segment time: %f ms \n" % segment_time) # calculate average time consuming total_time = float(total_time) + float(segment_time) total_number = total_number + 1 if int(total_number) > 0: average = total_time / float(total_number) print("Average time: %f ms" % average) # object_num = [] labelMsgs = [] # scoreMsgs = [] for i in range(batch_size): # Original results of MaskRCNN boxes = result[i]['rois'] masks = result[i]['masks'] class_ids = result[i]['class_ids'] # scores = result[i]['scores'] # Total number of objects N = boxes.shape[0] # Result of request label_img = np.zeros(masks.shape[0:2], dtype=np.uint8) # score_img = np.zeros(masks.shape[0:2], dtype=np.float32) logging.info('shap of label image: %s', label_img.shape) # label_img_color = np.zeros((masks.shape[0], masks.shape[1], 3), dtype=np.uint8) # logging.info('shap of label_color image: %s', label_img_color.shape) for i in range(N): # masks is 0/1 two value image, extents to [0-255] mask = masks[:, :, i] * 255 # print('type of mask: ', mask.dtype) # merge mask of each object to one mask image label_img += (masks[:, :, i] * class_ids[i]).astype(np.uint8) # score_img += (masks[:, :, i] * scores[i]).astype(np.float32) # semantic label msg_label = CvBridge().cv2_to_imgmsg(label_img.astype(np.uint8), encoding='mono8') # msg_score = CvBridge().cv2_to_imgmsg( # score_img.astype(np.float32), encoding='passthrough') labelMsgs.append(msg_label) # scoreMsgs.append(msg_score) # object_num.append(N) # return the request result # Total number of objects # self._result.object_num = object_num # Return segment result self._result.id = goal.id self._result.label = labelMsgs # self._result.score = scoreMsgs # feedback self._feedback.complete = True self._as.set_succeeded(self._result) self._as.publish_feedback(self._feedback) # calculate time consuming time_each_loop = (rospy.Time.now() - time_each_loop_start).to_sec() * 1000 print("Time of each request: %f ms \n" % time_each_loop) # Save results after sending back the results # if bSaveResult: # boxes = result['rois'] # masks = result['masks'] # class_ids = result['class_ids'] # scores = result['scores'] # N = boxes.shape[0] # for i in range(N): # save_file = RESULT_DIR + '/mask/' + str(goal.id) + '_' + str(class_ids[i]) + '.png' # m_img = (masks[:, :, i] * 255).astype(np.uint8) # cv2.imwrite(save_file, m_img) # # print('type of m_img: ', m_img.dtype) # # Save semantic label # save_file = RESULT_DIR + '/label/' + str(goal.id) + '.png' # cv2.imwrite(save_file, label_img) # print('type of label_img: ', label_img.dtype) # publish result # masked_image: just for visualization # msg_label_color = CvBridge().cv2_to_imgmsg(masked_image, encoding="bgr8") # if self.semantic_label_pub.get_num_connections() > 0: # msg_label.header.stamp = goal.image.header.stamp # self.semantic_label_pub.publish(msg_label) # # if self.semantic_label_color_pub.get_num_connections() > 0: # msg_label_color.header.stamp = goal.image.header.stamp # self.semantic_label_color_pub.publish(msg_label_color) def main(args): global total_time global total_number # global bSaveResult # save result # bSaveResult = True # bSaveResult = False rospy.init_node('maskrcnn_server', anonymous=False) th_maskrcnn = threading.Thread(name='maskrcnn', target=worker, args=( MaskRcnn, False, )) th_maskrcnn.start() actionserver = SemanticActionServer() print('Setting up MaskRCNN Action Server...') logging.debug('Waiting for worker threads') bStartMaskRCNN = False main_thread = threading.currentThread() try: rospy.spin() except KeyboardInterrupt: print("Shutting down") for t in threading.enumerate(): if t is not main_thread: t.join() if __name__ == '__main__': main(sys.argv)