improve COCO prototype

[pmeier]

This adds a few improvements:

1. Add the ability to select which annotations to load, e.g. instances (default) or captions. None is a special value that will only load the images which is much faster if the annotations are not needed. This is for example useful for unsupervised training.

2. Provide categories for each bounding box.

3. Add tests.

4. Add support for segmentation masks. Now you can do something like this:

   from torchvision.prototype import datasets
   from torchvision.utils import draw_segmentation_masks
   
   for sample in datasets.load("coco"):
       draw_segmentation_masks(sample["image"], sample["segmentations"])
       break

   

cc @pmeier @bjuncek

[pmeier]

The new implementation is working, but is a major perf regression. Cold start now takes ~20 minutes on my system. With minmal fake data warm up is minimal, so I'm guessing one or more things I used in the implementation does not scale to more inputs.

[pmeier]

I've added COCO to the benchmarks. Running

python -m torchvision.prototype.datasets.benchmark -n1 --no-start coco

with the patch described in #4650 (comment) gives:

################################################################################
coco (instances)
legacy iteration 185.052 it/s
new iteration 152.523 it/s
################################################################################
coco (captions)
legacy iteration 240.053 it/s
new iteration 198.828 it/s

[fmassa]

LGTM, thanks!

[fmassa]

For a future PR, I think it might be preferable to rename the class to BoundingBoxes as we hold more than one box now.

[fmassa]

Yes, we can even have a custom class that holds the raw polygons, and knows how to convert itself if needed. But can be discussed at a separate stage.

[fmassa]

ONNX-test error seems suspicious, but probably unrelated. I'm merging this but let's keep an eye on it