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()