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 if args.useCats: cat = spider.dataset.getCatIds()[int(ret_cats[_].argmax())] score = float(ret_cats[_].max()) else: score = float(ret_scores[_]) objn = float(ret_scores[_]) results.append({ 'image_id': image_id, 'category_id': cat, 'segmentation': encode(ret_masks[_]), 'score': score, 'objn': objn }) with open('results/%s.json' % args.model, "wb") as f: f.write(cjson.encode(results))
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 max_result = -999 for _ in range(len(ret_masks)): if ret_scores[_] > args.threshold and reserved[_]:
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()