def get_maps(self, image):
"""
Args:
image (np.ndarray): an image of shape (H, W, C) (in BGR order).
This is the format used by OpenCV.
Returns:
maps (dict): a dictionary of maps.
"""
maps = []
predictions = self.predictor(image)
if "panoptic_seg" in predictions:
panoptic_seg, segments_info = predictions["panoptic_seg"]
panoptic_seg = panoptic_seg.to(self.cpu_device)
self.pred = _PanopticPrediction(panoptic_seg, segments_info)
sem_map, ins_map = self.panoptic_seg2sem_map(panoptic_seg), \
self.panoptic_seg2ins_map(panoptic_seg)
maps = {"sem_map": sem_map, "ins_map": ins_map}
else:
if "sem_seg" in predictions:
raise NotImplementedError()
if "instances" in predictions:
raise NotImplementedError()
return maps
def draw_panoptic_seg_predictions(
self, frame, panoptic_seg, segments_info, area_threshold=None, alpha=0.5
):
frame_visualizer = Visualizer(frame, self.metadata)
pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata)
if self._instance_mode == ColorMode.IMAGE_BW:
frame_visualizer.output.img = frame_visualizer._create_grayscale_image(
pred.non_empty_mask()
)
# draw mask for all semantic segments first i.e. "stuff"
for mask, sinfo in pred.semantic_masks():
category_idx = sinfo["category_id"]
try:
mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]]
except AttributeError:
mask_color = None
frame_visualizer.draw_binary_mask(
mask,
color=mask_color,
text=self.metadata.stuff_classes[category_idx],
alpha=alpha,
area_threshold=area_threshold,
)
all_instances = list(pred.instance_masks())
if len(all_instances) == 0:
return frame_visualizer.output
# draw mask for all instances second
masks, sinfo = list(zip(*all_instances))
num_instances = len(masks)
masks_rles = mask_util.encode(
np.asarray(np.asarray(masks).transpose(1, 2, 0), dtype=np.uint8, order="F")
)
assert len(masks_rles) == num_instances
category_ids = [x["category_id"] for x in sinfo]
detected = [
_DetectedInstance(category_ids[i], bbox=None, mask_rle=masks_rles[i], color=None, ttl=8)
for i in range(num_instances)
]
colors = self._assign_colors(detected)
labels = [self.metadata.thing_classes[k] for k in category_ids]
frame_visualizer.overlay_instances(
boxes=None,
masks=masks,
labels=labels,
keypoints=None,
assigned_colors=colors,
alpha=alpha,
)
return frame_visualizer.output
def decode_panoptic_results(self,
panoptic_seg,
segments_info,
area_threshold=None,
alpha=0.7):
"""
Draw panoptic prediction results on an image.
Args:
panoptic_seg (Tensor): of shape (height, width) where the values are ids for each
segment.
segments_info (list[dict]): Describe each segment in `panoptic_seg`.
Each dict contains keys "id", "category_id", "isthing".
area_threshold (int): stuff segments with less than `area_threshold` are not drawn.
Returns:
output (VisImage): image object with visualizations.
"""
pred = _PanopticPrediction(panoptic_seg.cpu(), segments_info)
fg_class_cnt = len(self.metadata.thing_classes)
bg_class_cnt = len(self.metadata.stuff_classes)
all_mask = np.zeros((panoptic_seg.shape[0], panoptic_seg.shape[1], 3),
dtype=np.uint8)
all_mask_sep = np.zeros((panoptic_seg.shape[0], panoptic_seg.shape[1],
fg_class_cnt + bg_class_cnt + 1),
dtype=np.uint8)
# draw mask for all semantic segments first i.e. "stuff"
for mask, sinfo in pred.semantic_masks():
all_mask[:, :, 2] += mask
category_idx = sinfo["category_id"]
all_mask_sep[:, :, 1 + fg_class_cnt + category_idx] += mask
# draw mask for all instances second
all_instances = list(pred.instance_masks())
if len(all_instances) != 0:
masks, sinfo = list(zip(*all_instances))
category_ids = [x["category_id"] for x in sinfo]
for ind, ins_masks in enumerate(masks):
all_mask[:, :, 1] += ins_masks
all_mask_sep[:, :, 1 + category_ids[ind]] += ins_masks
return all_mask, all_mask_sep
def get_panoptic_seg_predictions(self, panoptic_seg, segments_info):
pred = _PanopticPrediction(panoptic_seg, segments_info)
objects = []
for mask, sinfo in pred.semantic_masks():
obj = {}
category_idx = sinfo["category_id"]
obj["text"] = self.metadata.stuff_classes[category_idx]
obj.update(sinfo)
objects.append(obj)
for mask, sinfo in pred.instance_masks():
obj = {}
instance_idx = sinfo["category_id"]
obj["text"] = self.metadata.thing_classes[instance_idx]
obj.update(sinfo)
objects.append(obj)
return objects
def draw_panoptic_seg_predictions(self, frame_shape, panoptic_seg,
segments_info, save_path):
pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata)
# save semantic mask
semantic_mask = np.zeros(frame_shape)
for mask, sinfo in pred.semantic_masks():
semantic_mask[mask] = sinfo["category_id"]
# print(sinfo)
np.save(f"{save_path}_sem_seg.npy", semantic_mask)
# save instances
all_instances = list(pred.instance_masks())
if len(all_instances) == 0:
return
masks, sinfo = list(zip(*all_instances))
masks_rles = mask_util.encode(
np.asarray(np.asarray(masks).transpose(1, 2, 0),
dtype=np.uint8,
order="F"))
category_ids = [x["category_id"] for x in sinfo]
detected = []
for cat, rle in zip(category_ids, masks_rles):
detected.append(
_DetectedInstance(cat, bbox=None, mask_rle=rle, ttl=8))
idds = self._assign_idd(detected)
pickle.dump(idds, open(f"{save_path}_ids.p", "wb"))
labels = [self.metadata.thing_classes[k] for k in category_ids]
pickle.dump(labels, open(f"{save_path}_labels.p", "wb"))
boxes = [self.mask2bbox(m) for m in masks]
np.save(f"{save_path}_boxes.npy", boxes)
for k, (mask, bb) in enumerate(zip(masks, boxes)):
crop = mask[math.floor(bb[1]):math.ceil(bb[3]),
math.floor(bb[0]):math.ceil(bb[2]), ].copy()
np.save(f"{save_path}_mask{k}.npy", crop)
