vis_coco_annotations.py

#!/usr/bin/env python

import argparse
import os.path as osp
import sys

import imgviz

try:
    import pycocotools.coco
except ImportError:
    print("Please install pycocotools:\n\n    pip install pycocotools\n")
    sys.exit(1)


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("json_file", help="json file")
    args = parser.parse_args()

    root_dir = osp.dirname(args.json_file)

    coco = pycocotools.coco.COCO(annotation_file=args.json_file)

    class_names = [None] * (max(coco.getCatIds()) + 1)
    for cat_id, cat in coco.cats.items():
        class_names[cat_id] = cat["name"]

    for img_id in coco.getImgIds():
        img_file = coco.loadImgs(ids=[img_id])[0]["file_name"]
        img = imgviz.io.imread(osp.join(root_dir, img_file))

        ann_ids = coco.getAnnIds(imgIds=[img_id])
        anns = coco.loadAnns(ids=ann_ids)

        class_ids = []
        masks = []
        for ann in anns:
            masks.append(coco.annToMask(ann).astype(bool))
            class_ids.append(ann["category_id"])

        captions = [class_names[c] for c in class_ids]
        viz = imgviz.instances2rgb(
            img,
            labels=class_ids,
            masks=masks,
            captions=captions,
            font_size=10,
            line_width=1,
        )
        viz_file = osp.join(root_dir, "Visualization", osp.basename(img_file))
        imgviz.io.imsave(viz_file, viz)
        print("==> Saved to {}".format(viz_file))


if __name__ == "__main__":
    main()