Function _create_annotation

axlearn/vision/coco_utils.py:284–325

(*, batch_id: int, image_id: int, box_id: int, class_label: int)

    ]

    def _create_annotation(*, batch_id: int, image_id: int, box_id: int, class_label: int) -> dict:
        ann = {}
        ann["image_id"] = int(groundtruths["source_id"][batch_id][image_id])
        if "is_crowds" in groundtruths:
            ann["iscrowd"] = int(groundtruths["is_crowds"][batch_id][image_id, box_id])
        else:
            ann["iscrowd"] = 0
        ann["category_id"] = class_label
        boxes = groundtruths["boxes"][batch_id]
        ann["bbox"] = [
            float(boxes[image_id, box_id, 1]),
            float(boxes[image_id, box_id, 0]),
            float(boxes[image_id, box_id, 3] - boxes[image_id, box_id, 1]),
            float(boxes[image_id, box_id, 2] - boxes[image_id, box_id, 0]),
        ]
        if "areas" in groundtruths:
            ann["area"] = float(groundtruths["areas"][batch_id][image_id, box_id])
        else:
            ann["area"] = float(
                (boxes[image_id, box_id, 3] - boxes[image_id, box_id, 1])
                * (boxes[image_id, box_id, 2] - boxes[image_id, box_id, 0])
            )
        if "masks" in groundtruths:
            if isinstance(groundtruths["masks"][batch_id][image_id, box_id], tf.Tensor):
                mask = Image.open(
                    six.BytesIO(groundtruths["masks"][batch_id][image_id, box_id].numpy())
                )
                width, height = mask.size
                np_mask = np.array(mask.getdata()).reshape(height, width).astype(np.uint8)
            else:
                mask = Image.open(six.BytesIO(groundtruths["masks"][batch_id][image_id, box_id]))
                width, height = mask.size
                np_mask = np.array(mask.getdata()).reshape(height, width).astype(np.uint8)
            np_mask[np_mask > 0] = 255
            encoded_mask = mask_api.encode(np.asfortranarray(np_mask))
            ann["segmentation"] = encoded_mask
            # Ensure the content of `counts` is JSON serializable string.
            if "counts" in ann["segmentation"]:
                ann["segmentation"]["counts"] = six.ensure_str(ann["segmentation"]["counts"])
            if "areas" not in groundtruths:
                ann["area"] = mask_api.area(encoded_mask)
        return ann

    gt_annotations = []
    num_batches = len(groundtruths["source_id"])