| |
| |
|
|
| import numpy as np |
| import torch |
|
|
| |
| from detectron2.data.detection_utils import \ |
| annotations_to_instances as base_annotations_to_instances |
| from detectron2.data.detection_utils import \ |
| transform_instance_annotations as base_transform_instance_annotations |
|
|
| |
|
|
|
|
| def annotations_to_instances(annos, image_size, sample_points=0): |
| """ |
| Create an :class:`Instances` object used by the models, |
| from instance annotations in the dataset dict. |
| |
| Args: |
| annos (list[dict]): a list of instance annotations in one image, each |
| element for one instance. |
| image_size (tuple): height, width |
| sample_points (int): subsample points at each iteration |
| |
| Returns: |
| Instances: |
| It will contain fields "gt_boxes", "gt_classes", |
| "gt_point_coords", "gt_point_labels", if they can be obtained from `annos`. |
| This is the format that builtin models with point supervision expect. |
| """ |
| target = base_annotations_to_instances(annos, image_size) |
|
|
| assert ("point_coords" in annos[0]) == ("point_labels" in annos[0]) |
|
|
| if len(annos) and "point_labels" in annos[0]: |
| point_coords = [] |
| point_labels = [] |
| for i, _ in enumerate(annos): |
| |
| point_coords_wrt_image = np.array(annos[i]["point_coords"]) |
| point_labels_wrt_image = np.array(annos[i]["point_labels"]) |
|
|
| if sample_points > 0: |
| random_indices = np.random.choice( |
| point_coords_wrt_image.shape[0], |
| sample_points, |
| replace=point_coords_wrt_image.shape[0] < sample_points, |
| ).astype(int) |
| point_coords_wrt_image = point_coords_wrt_image[random_indices] |
| point_labels_wrt_image = point_labels_wrt_image[random_indices] |
| assert point_coords_wrt_image.shape[0] == point_labels_wrt_image.size |
|
|
| point_coords.append(point_coords_wrt_image) |
| point_labels.append(point_labels_wrt_image) |
|
|
| point_coords = torch.stack([torch.from_numpy(x) for x in point_coords]) |
| point_labels = torch.stack([torch.from_numpy(x) for x in point_labels]) |
| target.gt_point_coords = point_coords |
| target.gt_point_labels = point_labels |
|
|
| return target |
|
|
|
|
| def transform_instance_annotations( |
| annotation, transforms, image_size, *, keypoint_hflip_indices=None |
| ): |
| """ |
| Apply transforms to box, and point annotations of a single instance. |
| It will use `transforms.apply_box` for the box, and |
| `transforms.apply_coords` for points. |
| Args: |
| annotation (dict): dict of instance annotations for a single instance. |
| It will be modified in-place. |
| transforms (TransformList or list[Transform]): |
| image_size (tuple): the height, width of the transformed image |
| keypoint_hflip_indices (ndarray[int]): see `create_keypoint_hflip_indices`. |
| Returns: |
| dict: |
| the same input dict with fields "bbox", "point_coords", "point_labels" |
| transformed according to `transforms`. |
| The "bbox_mode" field will be set to XYXY_ABS. |
| """ |
| annotation = base_transform_instance_annotations( |
| annotation, transforms, image_size, keypoint_hflip_indices |
| ) |
|
|
| assert ("point_coords" in annotation) == ("point_labels" in annotation) |
| if "point_coords" in annotation and "point_labels" in annotation: |
| point_coords = annotation["point_coords"] |
| point_labels = np.array(annotation["point_labels"]).astype(float) |
| point_coords = transforms.apply_coords(point_coords) |
|
|
| |
| inside = (point_coords >= np.array([0, 0])) & (point_coords <= np.array(image_size[::-1])) |
| inside = inside.all(axis=1) |
| point_labels[~inside] = -1 |
|
|
| annotation["point_coords"] = point_coords |
| annotation["point_labels"] = point_labels |
|
|
| return annotation |
|
|
