ImageFolder work-a-like for regression tasks

[asemic-horizon]

🚀 Feature

The requested/proposed feature is a close analog of torchvision.datasets.ImageFolder for regression tasks. The target values could be in a dict or an external JSON file.

Motivation

Thanks to torchvision.models and torchvision.datasets, the workflow for using transfer learning from alexnet, resnet, etc. to custom image classification tasks is very streamlined and polished. Datasets can be handled directly in the file system, which makes it easy to use a number of labeling tools that don't necessarily understand PyTorch, but are able to place images in folders like data/train/class1, data/val/class1, etc.

It's not as easy to do this for regression tasks, despite these being useful variations (e.g. for NSFW scoring rather than binary classification) of popular transfer learning tasks.

Pitch

Ideally, there would be a drop-in replacement for ImageFolder accepting the same arguments (such as transform callables) and an additional target_scores variable accepting a JSON filename or a dict. The method __getitem__ would return a (sample, score) pair.

To facilitate switching from classification to regression tasks, this RegressionImageFolder would ignore the directory structure that's used for target classes in ImageFolder.

Alternatives

It's possible that there's no rationale for having ImageFolder and RegressionImageFolder as separate dataset loaders and the functionality can be folded into what already exists.

cc @pmeier

[pmeier]

Hey @asemic-horizon.

target_scores variable accepting a JSON filename or a dict

What would that dictionary look like? Mapping file names to a score?

[asemic-horizon]

Hey @asemic-horizon.

target_scores variable accepting a JSON filename or a dict

What would that dictionary look like? Mapping file names to a score?

That's the simplest thing that occurs to me.

Ideally (in my view) one would be able to use existing the existing data-loading and transforming pipeline to maximize code reuse between classification and regression workflows; this would include ignoring the existing filesystem structure where subfolders are category names.

But, to think of it: I have the habit of renaming image files to their checksum (as in cksum) to minimize the risk of duplicates. I'm not sure what should be done in cases where there are different images in subfolders with the same filename. So maybe subfolder names can't be ignored.

[pmeier]

If you want to have the same functionality as ImageFolder this could be as easy as

from typing import Dict, Any
from torchvision import datasets


class RegressionImageFolder(datasets.ImageFolder):
    def __init__(
        self, root: str, image_scores: Dict[str, float], **kwargs: Any
    ) -> None:
        super().__init__(root, **kwargs)
        paths, _ = zip(*self.imgs)
        self.targets = [image_scores[path] for path in paths]
        self.samples = self.imgs = list(zip(paths, self.targets))

Did I miss something?

[mehdiborji]

Thank you so much @pmeier for the code, this made my day!

I would suggest *args instead of **kwargs, since the latter requires explicitly keywording parent arguments.