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)