improve COCO prototype (#4650)

* improve COCO prototype

* test 2017 annotations

* add option to include captions

* fix categories and add tests

* cleanup

* add correct image size to bounding boxes

* fix annotation collation

* appease mypy

* add benchmark

* always use image as reference

* another refactor

* add support for segmentations

* add support for segmentations

* fix CI dependencies

Author: pmeier

.circleci/config.yml

@@ -351,7 +351,7 @@ jobs:
       - install_torchvision
       - install_prototype_dependencies
       - pip_install:
-          args: scipy
+          args: scipy pycocotools
           descr: Install optional dependencies
       - run:
           name: Enable prototype tests

.circleci/config.yml.in

@@ -1,5 +1,6 @@
 import functools
 import gzip
+import json
 import lzma
 import pathlib
 import pickle
@@ -8,6 +9,7 @@
 from typing import Any, Dict, Tuple
 
 import numpy as np
+import PIL.Image
 import pytest
 import torch
 from datasets_utils import create_image_folder, make_tar, make_zip
@@ -18,7 +20,9 @@
 from torchvision.prototype.datasets._api import find
 from torchvision.prototype.utils._internal import add_suggestion
 
+
 make_tensor = functools.partial(_make_tensor, device="cpu")
+make_scalar = functools.partial(make_tensor, ())
 
 __all__ = ["load"]
 
@@ -490,3 +494,113 @@ def imagenet(info, root, config):
     make_tar(root, f"{devkit_root}.tar.gz", devkit_root, compression="gz")
 
     return num_samples
+
+
+class CocoMockData:
+    @classmethod
+    def _make_images_archive(cls, root, name, *, num_samples):
+        image_paths = create_image_folder(
+            root, name, file_name_fn=lambda idx: f"{idx:012d}.jpg", num_examples=num_samples
+        )
+
+        images_meta = []
+        for path in image_paths:
+            with PIL.Image.open(path) as image:
+                width, height = image.size
+            images_meta.append(dict(file_name=path.name, id=int(path.stem), width=width, height=height))
+
+        make_zip(root, f"{name}.zip")
+
+        return images_meta
+
+    @classmethod
+    def _make_annotations_json(
+        cls,
+        root,
+        name,
+        *,
+        images_meta,
+        fn,
+    ):
+        num_anns_per_image = torch.randint(1, 5, (len(images_meta),))
+        num_anns_total = int(num_anns_per_image.sum())
+        ann_ids_iter = iter(torch.arange(num_anns_total)[torch.randperm(num_anns_total)])
+
+        anns_meta = []
+        for image_meta, num_anns in zip(images_meta, num_anns_per_image):
+            for _ in range(num_anns):
+                ann_id = int(next(ann_ids_iter))
+                anns_meta.append(dict(fn(ann_id, image_meta), id=ann_id, image_id=image_meta["id"]))
+        anns_meta.sort(key=lambda ann: ann["id"])
+
+        with open(root / name, "w") as file:
+            json.dump(dict(images=images_meta, annotations=anns_meta), file)
+
+        return num_anns_per_image
+
+    @staticmethod
+    def _make_instances_data(ann_id, image_meta):
+        def make_rle_segmentation():
+            height, width = image_meta["height"], image_meta["width"]
+            numel = height * width
+            counts = []
+            while sum(counts) <= numel:
+                counts.append(int(torch.randint(5, 8, ())))
+            if sum(counts) > numel:
+                counts[-1] -= sum(counts) - numel
+            return dict(counts=counts, size=[height, width])
+
+        return dict(
+            segmentation=make_rle_segmentation(),
+            bbox=make_tensor((4,), dtype=torch.float32, low=0).tolist(),
+            iscrowd=True,
+            area=float(make_scalar(dtype=torch.float32)),
+            category_id=int(make_scalar(dtype=torch.int64)),
+        )
+
+    @staticmethod
+    def _make_captions_data(ann_id, image_meta):
+        return dict(caption=f"Caption {ann_id} describing image {image_meta['id']}.")
+
+    @classmethod
+    def _make_annotations(cls, root, name, *, images_meta):
+        num_anns_per_image = torch.zeros((len(images_meta),), dtype=torch.int64)
+        for annotations, fn in (
+            ("instances", cls._make_instances_data),
+            ("captions", cls._make_captions_data),
+        ):
+            num_anns_per_image += cls._make_annotations_json(
+                root, f"{annotations}_{name}.json", images_meta=images_meta, fn=fn
+            )
+
+        return int(num_anns_per_image.sum())
+
+    @classmethod
+    def generate(
+        cls,
+        root,
+        *,
+        year,
+        num_samples,
+    ):
+        annotations_dir = root / "annotations"
+        annotations_dir.mkdir()
+
+        for split in ("train", "val"):
+            config_name = f"{split}{year}"
+
+            images_meta = cls._make_images_archive(root, config_name, num_samples=num_samples)
+            cls._make_annotations(
+                annotations_dir,
+                config_name,
+                images_meta=images_meta,
+            )
+
+        make_zip(root, f"annotations_trainval{year}.zip", annotations_dir)
+
+        return num_samples
+
+
+@dataset_mocks.register_mock_data_fn
+def coco(info, root, config):
+    return CocoMockData.generate(root, year=config.year, num_samples=5)

test/builtin_dataset_mocks.py

@@ -866,6 +866,13 @@ def _split_files_or_dirs(root, *files_or_dirs):
 
 def _make_archive(root, name, *files_or_dirs, opener, adder, remove=True):
     archive = pathlib.Path(root) / name
+    if not files_or_dirs:
+        dir = archive.parent / archive.name.replace("".join(archive.suffixes), "")
+        if dir.exists() and dir.is_dir():
+            files_or_dirs = (dir,)
+        else:
+            raise ValueError("No file or dir provided.")
+
     files, dirs = _split_files_or_dirs(root, *files_or_dirs)
 
     with opener(archive) as fh:

test/datasets_utils.py

@@ -13,6 +13,17 @@ def to_bytes(file):
     return file.read()
 
 
+def config_id(name, config):
+    parts = [name]
+    for name, value in config.items():
+        if isinstance(value, bool):
+            part = ("" if value else "no_") + name
+        else:
+            part = str(value)
+        parts.append(part)
+    return "-".join(parts)
+
+
 def dataset_parametrization(*names, decoder=to_bytes):
     if not names:
         # TODO: Replace this with torchvision.prototype.datasets.list() as soon as all builtin datasets are supported
@@ -27,16 +38,17 @@ def dataset_parametrization(*names, decoder=to_bytes):
             "caltech256",
             "caltech101",
             "imagenet",
+            "coco",
         )
 
-    params = []
-    for name in names:
-        for config in datasets.info(name)._configs:
-            id = f"{name}-{'-'.join([str(value) for value in config.values()])}"
-            dataset, mock_info = builtin_dataset_mocks.load(name, decoder=decoder, **config)
-            params.append(pytest.param(dataset, mock_info, id=id))
-
-    return pytest.mark.parametrize(("dataset", "mock_info"), params)
+    return pytest.mark.parametrize(
+        ("dataset", "mock_info"),
+        [
+            pytest.param(*builtin_dataset_mocks.load(name, decoder=decoder, **config), id=config_id(name, config))
+            for name in names
+            for config in datasets.info(name)._configs
+        ],
+    )
 
 
 class TestCommon:

test/test_prototype_builtin_datasets.py

@@ -0,0 +1,91 @@
+__background__,N/A
+person,person
+bicycle,vehicle
+car,vehicle
+motorcycle,vehicle
+airplane,vehicle
+bus,vehicle
+train,vehicle
+truck,vehicle
+boat,vehicle
+traffic light,outdoor
+fire hydrant,outdoor
+N/A,N/A
+stop sign,outdoor
+parking meter,outdoor
+bench,outdoor
+bird,animal
+cat,animal
+dog,animal
+horse,animal
+sheep,animal
+cow,animal
+elephant,animal
+bear,animal
+zebra,animal
+giraffe,animal
+N/A,N/A
+backpack,accessory
+umbrella,accessory
+N/A,N/A
+N/A,N/A
+handbag,accessory
+tie,accessory
+suitcase,accessory
+frisbee,sports
+skis,sports
+snowboard,sports
+sports ball,sports
+kite,sports
+baseball bat,sports
+baseball glove,sports
+skateboard,sports
+surfboard,sports
+tennis racket,sports
+bottle,kitchen
+N/A,N/A
+wine glass,kitchen
+cup,kitchen
+fork,kitchen
+knife,kitchen
+spoon,kitchen
+bowl,kitchen
+banana,food
+apple,food
+sandwich,food
+orange,food
+broccoli,food
+carrot,food
+hot dog,food
+pizza,food
+donut,food
+cake,food
+chair,furniture
+couch,furniture
+potted plant,furniture
+bed,furniture
+N/A,N/A
+dining table,furniture
+N/A,N/A
+N/A,N/A
+toilet,furniture
+N/A,N/A
+tv,electronic
+laptop,electronic
+mouse,electronic
+remote,electronic
+keyboard,electronic
+cell phone,electronic
+microwave,appliance
+oven,appliance
+toaster,appliance
+sink,appliance
+refrigerator,appliance
+N/A,N/A
+book,indoor
+clock,indoor
+vase,indoor
+scissors,indoor
+teddy bear,indoor
+hair drier,indoor
+toothbrush,indoor