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