SegmentationDataset [feature proposal and demo]

[Noiredd]

In addition to #1315, I propose a new dataset subclass for loading samples compatible with the processing pipeline. In essence, if we have a root folder containing images and labels subfolders, this class loads and returns samples in form of tuples (image, label). We can pass a SegmentationCompose directly to it, immediately allowing data preprocessing/augmentations, just like with ImageFolder and Compose.

This is just a proposal/prototype, so some desired features are missing (e.g. checking for file validity, size matching assertions, etc.). Yes, I know that the normal flow for feature proposals is to start with an issue, but since I've already written it (as I need something like this in my own work), I decided to open a PR. I hope you don't mind and we can discuss it here - besides, that might be easier with a code example ;)

Example usage:

import torch
import torchvision as tv

src = tv.datasets.folder.SegmentationDataset(root='path/to/my_dataset')
image, label = src[0]
print(image)
print(label)

But with #1315 we could take this a step further:

import torch
import torchvision as tv

tfs = tv.transforms.SegmentationCompose([
    tv.transforms.RandomCrop(400),
    tv.transforms.ToTensor(),
])
src = tv.datasets.folder.SegmentationDataset(
    root='path/to/my_dataset',
    transforms=tfs
)
ldr = torch.utils.data.DataLoader(src, batch_size=1)
for i, b in enumerate(ldr):
    print(b[0].shape, b[0].type(), b[0].numpy().max())
    print(b[1].shape, b[1].type(), b[1].numpy().max())

[codecov-io]

Codecov Report

Merging #1330 into master will decrease coverage by 0.18%.
The diff coverage is 17.24%.

@@            Coverage Diff            @@
##           master   #1330      +/-   ##
=========================================
- Coverage   65.78%   65.6%   -0.19%     
=========================================
  Files          79      75       -4     
  Lines        5834    5814      -20     
  Branches      887     892       +5     
=========================================
- Hits         3838    3814      -24     
- Misses       1726    1734       +8     
+ Partials      270     266       -4

Impacted Files	Coverage Δ
torchvision/datasets/folder.py	64.48% <17.24%> (-17.57%)	⬇️
torchvision/datasets/kinetics.py	34.78% <0%> (-5.22%)	⬇️
torchvision/datasets/hmdb51.py	27.65% <0%> (-3.3%)	⬇️
torchvision/datasets/ucf101.py	25% <0%> (-3.21%)	⬇️
torchvision/transforms/functional.py	71.26% <0%> (-0.13%)	⬇️
torchvision/io/video.py	71.31% <0%> (-0.12%)	⬇️
torchvision/models/video/__init__.py	100% <0%> (ø)	⬆️
torchvision/ops/boxes.py	94.73% <0%> (ø)	⬆️
torchvision/datasets/mnist.py	51.7% <0%> (ø)	⬆️
torchvision/models/video/video_stems.py
... and 15 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2287c8f...9418bd3. Read the comment docs.

[SebastienEske]

Hello there, I made this pull request which I believe could be simple and generic enough. Looks like there are several pull requests on the same topic... Sorry for spamming.

[fmassa]

Thanks for the PR @Noiredd and sorry for the delay in replying!

This is a very important discussion to have, I'm glad that you've opened this PR.

Following my comment in #1345, I think that we should provide some examples in torchvision on how to use the current abstractions that we have in order to easily build those datasets.

You example in #1345 (comment) is how I would start, with some slight changes:

class BundleDataset(object):
    def __init__(self, *datasets):
        self.datasets = datasets
        # here go sanity checks like asserting that all datasets are equal length

    def __getitem__(self, index):
        output = tuple(dataset[index] for dataset in self.datasets)
        if self.transforms is not None:
                output = self.transforms(*output)
        return output

    def __len__(self):
        return len(self.datasets[0])

Those kind of generic datasets could potentially live somewhere (maybe in torchvision, maybe in PyTorch, maybe in torchdata?)

As such, I think we are not going to be moving forward with this approach, but we should definitely find a common ground solution which is generic, modular and simple, similar to ConcatDataset and Subset.

Let me know your thoughts.

[Noiredd]

Do we want maximum abstraction and generality, or a solution to a specific problem?

In the above example you merge multiple datasets and iterate over them. But as you noticed in this #1345 comment, we need to ensure the mapping between files in constituent folders. Where should that code live? Should BundleDataset inspect and modify its member datasets? Or should the user have to supply their own validation function? This is not as trivial as ConcatDataset or Subset, as it features a problem that I don't think can be easily abstracted away.

Personally, I don't care much if I have to write:

src = tv.datasets.folder.SegmentationDataset(
    root='path/to/my_dataset',
    transforms=tfs
)

or

src_img = tv.datasets.folder.ImageFolder(root='path/to/my_dataset/images')
src_lab = tv.datasets.folder.ImageFolder(root='path/to/my_dataset/labels', loader=label_loader)
src = whatever.BundleDataset((src_img, src_lab), transforms=tfs)

as long as I don't have to copy&paste lines 255-275 (from this PR) every single time I want to run a segmentation task.
But I do a lot of segmentation (and generally image to image tasks) so I'd like that infrastructure to be simple. I am biased ;)

[SebastienEske]

The assumption in #1345 was that the user is responsible for making sure that when python sorts the filenames (this line and the next one ) they are in the proper order for all folders. This allows folders with names such as img1, img2... and label1, label2... to be mapped together in a simple way and we only need to check for the number of files.
Obviously if you want a generic bundle dataset, that's going to be another story... Have fun bundling datasets from torchvision and torchtext.

[yassineAlouini]

Hello @Noiredd and sorry for the very long delay in getting a new update.

It looks like this PR and the #1345 one have been useful to discuss some of the requested segmentation and object detection features. Not yet sure how this could be implemented best or if it is enough to add some documentation using existing tools (for example ImageFolder).

@pmeier What do you think is best to do in this situation? There are some visualization tools here for example but nothing directly linked to segmentation datasets or maybe I haven't looked close enough.

@Noiredd Is this something you are still thinking about or maybe you found an alternative? Please let us know if things have evolved. Many thanks in advance for your feedback and sorry for taking that long to get news.

[pmeier]

IMO the main issue is that there is no "agreed" file structure for segmentation datasets. Thus, we would need to enforce one, which in turn might clash with some datasets or other libraries that also provide dataset wrappers. In such a case I prefer to not provide a builtin solution.

@Noiredd We are refactoring our datasets to use torchdata as backend. With that, loading data from an image folder architecture can be achieved with very few building blocks:

vision/torchvision/prototype/datasets/_folder.py

Lines 44 to 49 in c67a583

	dp = FileLister(str(root), recursive=recursive, masks=masks)
	dp: IterDataPipe = Filter(dp, functools.partial(_is_not_top_level_file, root=root))
	dp = hint_sharding(dp)
	dp = hint_shuffling(dp)
	dp = FileOpener(dp, mode="rb")
	return Mapper(dp, functools.partial(_prepare_sample, root=root, categories=categories)), categories

Expanding that to whatever layout you have locally is not hard. Let me know if you disagree.

@yassineAlouini Thanks for the follow-up.