load_config("configs/%s.json" % args.model) else: print("Specified config does not exists, use the default config...") image = load_image(args.input_image) oh, ow = image.shape[:2] im_scale = config.TEST_SCALE * 1.0 / max(oh, ow) input_blob = image - config.RGB_MEAN input_blob = input_blob.transpose((2, 0, 1)) ih, iw = int(oh * im_scale), int(ow * im_scale) ih, iw = ih - ih % 4, iw - iw % 4 input_blob = resize_blob(input_blob, dest_shape=(ih, iw)) input_blob = input_blob[np.newaxis, ...] ret_masks, ret_scores = gen_masks(net, input_blob, config, dest_shape=(oh, ow)) # nms encoded_masks = encode(ret_masks) reserved = np.ones((len(ret_masks))) for i in range(len(reserved)): if ret_scores[i] < args.threshold: reserved[i] = 0 continue if reserved[i]: for j in range(i + 1, len(reserved)): if reserved[j] and iou(encoded_masks[i], encoded_masks[j], [False]) > 0.5: reserved[j] = 0
from spiders.coco_ssm_spider import COCOSSMDemoSpider spider = COCOSSMDemoSpider() spider.dataset.sort(key=lambda item: int(item.image_path[-10:-4])) ds = spider.dataset[:args.end] timer.tic() results = [] for i in range(len(ds)): spider.fetch() img = spider.img_blob image_id = int(ds[i].image_path[-10:-4]) # gen mask ret = gen_masks(net, img, config, dest_shape=(spider.origin_height, spider.origin_width), useCats=args.useCats, vis=args.debug) if args.useCats: ret_masks, ret_scores, ret_cats = ret else: ret_masks, ret_scores = ret printProgress(i, len(ds), prefix='Progress: ', suffix='Complete', barLength=50) for _ in range(len(ret_masks)): cat = 1 if args.useCats: cat = spider.dataset.getCatIds()[int(ret_cats[_].argmax())] score = float(ret_cats[_].max()) else: score = float(ret_scores[_])
spider = COCOSSMDemoSpider() spider.dataset.sort(key=lambda item: int(item.image_path[-10:-4])) ds = spider.dataset[:args.end] timer.tic() results = [] for i in range(len(ds)): spider.fetch() img = spider.img_blob image_id = int(ds[i].image_path[-10:-4]) # gen mask ret = gen_masks(net, img, config, dest_shape=(spider.origin_height, spider.origin_width), image=args.debug) if args.useCats: ret_masks, ret_scores, ret_cats = ret else: ret_masks, ret_scores = ret printProgress(i, len(ds), prefix='Progress: ', suffix='Complete', barLength=50) for _ in range(len(ret_masks)): cat = 1
def handle_fast_mask_segmentation(self, req): if self.net == None: try: self.net = caffe.Net(self.model_path, self.weight_path, caffe.TEST) except: rospy.logerr("Error, cannot load fm net to the GPU") self.net = None self.service_queue = -1 return FastMaskSegmentationResponse() try: image = self.br.imgmsg_to_cv2(req.rgb_img, desired_encoding="bgr8") image = image.astype(np.float64) if self.gpu_id >= 0: caffe.set_mode_gpu() caffe.set_device(self.gpu_id) else: caffe.set_mode_cpu() oh, ow = image.shape[:2] im_scale = config.TEST_SCALE * 1.0 / max(oh, ow) input_blob = image - config.RGB_MEAN input_blob = input_blob.transpose((2, 0, 1)) ih, iw = int(oh * im_scale), int(ow * im_scale) ih, iw = ih - ih % 4, iw - iw % 4 input_blob = resize_blob(input_blob, dest_shape=(ih, iw)) input_blob = input_blob[np.newaxis, ...] ret_masks, ret_scores = gen_masks(self.net, input_blob, config, dest_shape=(oh, ow)) encoded_masks = encode(ret_masks) reserved = np.ones((len(ret_masks))) for i in range(len(reserved)): if ret_scores[i] < self.threshold: reserved[i] = 0 continue if reserved[i]: for j in range(i + 1, len(reserved)): if reserved[j] and iou(encoded_masks[i], encoded_masks[j], [False]) > 0.5: reserved[j] = 0 temp_image = image.copy() fastmask_bbox_arr = FastMaskBB2DArray() for _ in range(len(ret_masks)): if ret_scores[_] > self.threshold and reserved[_]: mask = ret_masks[_].copy() bbox = toBbox(mask) x, y, w, h = bbox fastmask_bbox = FastMaskBB2D() fastmask_bbox.bbox.x = x fastmask_bbox.bbox.y = y fastmask_bbox.bbox.w = w fastmask_bbox.bbox.h = h fastmask_bbox.score = ret_scores[_] fastmask_bbox_arr.fm_bbox_arr.append(fastmask_bbox) if self.save_and_display: bbox = [int(x) for x in bbox] mask[mask == 1] = 0.3 mask[mask == 0] = 1 color = COLORS[_ % len(COLORS)] _color = color & 0xff for k in range(3): image[:, :, k] = image[:, :, k] * mask mask[mask == 1] = 0 mask[mask > 0] = 0.7 for k in range(3): image[:, :, k] += mask * (color & 0xff) color >>= 8 cv2.rectangle(temp_image, (bbox[0], bbox[1]), (bbox[0] + bbox[2], bbox[1] + bbox[3]), (_color, _color, _color), 1) if self.save_and_display: image = image.astype(np.uint8) temp_image = temp_image.astype(np.uint8) cv2.imwrite(fast_mask_root + "/images/mask_result.jpg", image) cv2.imwrite(fast_mask_root + "/images/boundingbox_result.jpg", temp_image) self.net = None return FastMaskSegmentationResponse( segmentation_bbox_arr=fastmask_bbox_arr) except cv_bridge.CvBridgeError as e: rospy.logerr("CvBridge exception %s", e) return FastMaskSegmentationResponse()