pytorch / vision

Datasets, Transforms and Models specific to Computer Vision
https://pytorch.org/vision
BSD 3-Clause "New" or "Revised" License
15.76k stars 6.89k forks source link

Improved functionality for Oxford IIIT Pet data loader #8364

Open matlabninja opened 3 months ago

matlabninja commented 3 months ago

🚀 The feature

Add the following functionality to the Oxford IIIT Pet data loader

  1. Support binary classification of cat vs dog
  2. With the segmentation target type, produce trimaps with class/background/don't care regions instead of target/background/don't care when the output is a tensor
  3. Support detection as a target type

Motivation, pitch

The Oxford IIIT Pet dataset is a fun dataset for trying out new things and for new practitioners to use to learn. These capabilities allow users to more easily use this dataset with detection and segmentation target types and to use the existing annotation for animal species (rather than breed) as a simpler problem to get started. I have created these capabilities on my local copy of torchvision, and I'm up for creating a PR if the community likes the enhancements. The individual proposed enhancements can be found in the links: Binary cat v dog Class labeled segmentation Detection target type All 3 enhancements

Alternatives

I thought a lot about the ability to write transforms to use with a dataset loader to accomplish this, but it was unclear to me how I could access some of the class members of the dataset loaders that were necessary.

Additional context

Demonstration of the new features can be found in the following notebooks: Class-labeled segmentation maps and binary species classification training Deeplab V3 Detection target type training resnet 50 faster RCNN

NicolasHug commented 2 months ago

Thanks for the feature request @matlabninja .

Supporting 1. and 3. SGTM, we could add target_types="binary-category" and target_types="detection".

I'm not sure I completely follow what you mean by 2. though - could you please share more details about this?

With the segmentation target type, produce trimaps with class/background/don't care regions instead of target/background/don't care when the output is a tensor

In the meantime, please feel free to submit a PR for 1. and 3.

matlabninja commented 2 months ago

Thanks for the feedback @NicolasHug, I've put in a PR for 1. I made some changes to my original code to support "binary-category" as a target type, as that seemed to make more sense from a useability perspective than my previous setup with a separate input for "binary". As a result, I'll probably wait for 1 to clear before doing 3 as it opens up the door for binary detection as well.

I've put it in as draft as I am having some trouble running the unit tests locally (partially initialized module likely due to circular import). It's happening on all unit tests not just mine, so something is whacky with my config maybe. I was going to see how it went with running the unit tests in github instead.

Regarding the segmentation modification, all of the trimap png files in this dataset contain only the values 1, 2, and 3. 1 is used for pixels on the the pet, 2 for background, and 3 for a boundary region that can be used with the "ignore_index" of a loss function to not use that region for training. These png files do not have any class specific information in them, and the existing data loader, these files are loaded up directly. This doesn't let us use the data loader as-is to train a segmentation model supporting class differentiation.

The modification that I made changes that, but because modifying PIL images based on masks is a pain, I had this modification run after the transforms convert it to tensor. I certainly 'can' have it run earlier so that it is not transform dependent. See below for a example outputs with the current data loader. The change results in the background (2) becoming 37, the boundary region (3) becoming 38, and the "foreground region" (1) becoming something on [0,36] corresponding to the class label. image