Unify different panoptic annotation/prediction formats
| 153 | |
| 154 | |
| 155 | class _PanopticPrediction: |
| 156 | """ |
| 157 | Unify different panoptic annotation/prediction formats |
| 158 | """ |
| 159 | |
| 160 | def __init__(self, panoptic_seg, segments_info, metadata=None): |
| 161 | if segments_info is None: |
| 162 | assert metadata is not None |
| 163 | # If "segments_info" is None, we assume "panoptic_img" is a |
| 164 | # H*W int32 image storing the panoptic_id in the format of |
| 165 | # category_id * label_divisor + instance_id. We reserve -1 for |
| 166 | # VOID label. |
| 167 | label_divisor = metadata.label_divisor |
| 168 | segments_info = [] |
| 169 | for panoptic_label in np.unique(panoptic_seg.numpy()): |
| 170 | if panoptic_label == -1: |
| 171 | # VOID region. |
| 172 | continue |
| 173 | pred_class = panoptic_label // label_divisor |
| 174 | isthing = pred_class in metadata.thing_dataset_id_to_contiguous_id.values() |
| 175 | segments_info.append( |
| 176 | { |
| 177 | "id": int(panoptic_label), |
| 178 | "category_id": int(pred_class), |
| 179 | "isthing": bool(isthing), |
| 180 | } |
| 181 | ) |
| 182 | del metadata |
| 183 | |
| 184 | self._seg = panoptic_seg |
| 185 | |
| 186 | self._sinfo = {s["id"]: s for s in segments_info} # seg id -> seg info |
| 187 | segment_ids, areas = torch.unique(panoptic_seg, sorted=True, return_counts=True) |
| 188 | areas = areas.numpy() |
| 189 | sorted_idxs = np.argsort(-areas) |
| 190 | self._seg_ids, self._seg_areas = segment_ids[sorted_idxs], areas[sorted_idxs] |
| 191 | self._seg_ids = self._seg_ids.tolist() |
| 192 | for sid, area in zip(self._seg_ids, self._seg_areas): |
| 193 | if sid in self._sinfo: |
| 194 | self._sinfo[sid]["area"] = float(area) |
| 195 | |
| 196 | def non_empty_mask(self): |
| 197 | """ |
| 198 | Returns: |
| 199 | (H, W) array, a mask for all pixels that have a prediction |
| 200 | """ |
| 201 | empty_ids = [] |
| 202 | for id in self._seg_ids: |
| 203 | if id not in self._sinfo: |
| 204 | empty_ids.append(id) |
| 205 | if len(empty_ids) == 0: |
| 206 | return np.zeros(self._seg.shape, dtype=np.uint8) |
| 207 | assert ( |
| 208 | len(empty_ids) == 1 |
| 209 | ), ">1 ids corresponds to no labels. This is currently not supported" |
| 210 | return (self._seg != empty_ids[0]).numpy().astype(np.bool) |
| 211 | |
| 212 | def semantic_masks(self): |