lightly-ai / lightly-ai/lightly
Supporting Tensor as input to ImageCollateFunction
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.8k
- Forks
- 367
- Avg merge
- 3d 22h
- Merged PRs (30d)
- 5
Description
I'm using ImageCollateFunction to augment images in SimClr and BYOL methods. I've have implemented a torch `Dataset` which is then passed to `LightlyDataset`, like `LightlyDataset.from_torch_dataset(torch_dataset)`. My torch dataset's `__getitem__` method returns image and label pair.
The issue is that [ImageCollateFunction](https://github.com/lightly-ai/lightly/blob/020ff078a4f357c950c4823402f42f2338fda160/lightly/data/collate.py#:~:text=class%20ImageCollateFunction(BaseCollateFunction)%3A) supports only PIL image as input, while my dataset is naturally in Tensor format.
I've added conversion to PIL image on my side, but I think that Tensor input should be supported, as it's a natural format for both Pytorch and many datasets. Also, almost all of the transformation functions support both PIL and Tensor format. It would be faster to avoid torch -> PIL -> torch transformations, and instead just keep data in torch format.
My suggestion is to do the following:
- [GaussianBlur ](https://github.com/lightly-ai/lightly/blob/020ff078a4f357c950c4823402f42f2338fda160/lightly/transforms/gaussian_blur.py#L10) function currently relies on PIL's gaussian blur, which supports only PIL format. This class should be changed to use [torchvision's GaussianBlur](https://pytorch.org/vision/stable/generated/torchvision.transforms.GaussianBlur.html#torchvision.transforms.GaussianBlur) which supports both PIL and Tensor as input (and output)
- Instead of `T.ToTensor()` in [the list of transforms](https://github.com/lightly-ai/lightly/blob/020ff078a4f357c950c4823402f42f2338fda160/lightly/data/collate.py#L167), a custom `ToTensorCustom()' should be called, which supports both Tensor and PIL image as input, and returns Tensor. The logic would be trivial:
```
def ToTensorCustom(img):
if isinstance(img, Tensor):
return img
return T.ToTensor(img)
```
What do you think about this, is there an easier or clearer way to do it? I'd be happy to implement this myself and create a PR.
Kudos to @guarin who helped me with this.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading lightly/data/collate.py and lightly/transforms/gaussian_blur.py, then trace ImageCollateFunction's input and transformation path. Done means tensor inputs work through the SimCLR and BYOL augmentation flow without requiring PIL conversion, while PIL inputs remain supported.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- computer-vision, machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 32/100