Refactor load_image to return torch.Tensor instead of PIL.Image #2366

Ankur-singh · 2025-02-08T09:42:48Z

Context

What is the purpose of this PR? Is it to

add a new feature
fix a bug
update tests and/or documentation
other (refactor)

Please link to any issues this PR addresses. #2303

Changelog

What are the changes made in this PR?

Updated load_image to return torch.Tensor instead of PIL.Image
Updated CLIPImageTransform to support both torch.Tensor and PIL.Image
Added 4x new tests for CLIPImageTransform with torch.Tensor input

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
add unit tests for any new functionality
update docstrings for any new or updated methods or classes
run unit tests via pytest tests
run recipe tests via pytest tests -m integration_test
manually run any new or modified recipes with sufficient proof of correctness
include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

I did not change any public API
I have added an example to docs or docstrings

pytorch-bot · 2025-02-08T09:42:51Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2366

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a069f8a with merge base 7b654ea ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

joecummings

cc @NicolasHug just for a quick glance.

joecummings

Couple of comments, but super happy to see this working!

joecummings · 2025-02-10T18:01:08Z

torchtune/data/_utils.py



-def load_image(image_loc: Union[Path, str]) -> "PIL.Image.Image":
+def load_image(image_loc: Union[Path, str]) -> "torch.Tensor":


You don't need to quote the return type as torch.Tensor is definitely already importable.

joecummings · 2025-02-10T18:02:12Z

torchtune/data/_utils.py

-        image = Image.open(image_loc)
-    except Exception as e:
-        raise ValueError(f"Failed to open image as PIL Image from {image_loc}") from e
+            raise ValueError("Failed to load image as torch.Tensor") from e


Can you be more specific with this error? Something like failed to load image from remote.

joecummings · 2025-02-10T18:02:28Z

torchtune/data/_utils.py

+        try:
+            image = torchvision.io.decode_image(image_loc)
+        except Exception as e:
+            raise ValueError("Failed to open image as torch.Tensor") from e


Same here, lets be specific that this fails to load from a local image.

joecummings · 2025-02-10T18:03:16Z

torchtune/models/clip/inference/_transform.py


-    def __call__(self, *, image: Image.Image, **kwargs) -> Mapping[str, Any]:
+    def __call__(
+        self, *, image: Union[Image.Image, torch.Tensor], **kwargs


cc @pbontrager Do we actually want to allow PIL images through here? Or limit to just tensor?

hi @pbontrager, any comments?

NicolasHug · 2025-02-10T18:23:34Z

cc @NicolasHug just for a quick glance.

This SGTM

NicolasHug · 2025-02-12T13:36:57Z

torchtune/data/_utils.py

-            image_loc = request.urlopen(image_loc)
+            image_loc = request.urlopen(image_loc).read()
+            image = torchvision.io.decode_image(
+                torch.frombuffer(image_loc, dtype=torch.uint8)


Chatting with @pbontrager about this, I think you'll want to ensure that you get RGB images from decode_image by passing mode="RGB". https://pytorch.org/vision/main/generated/torchvision.io.ImageReadMode.html#torchvision.io.ImageReadMode

This is typically needed for images that are encoded as grayscale, or have some transparency channel.

(Here and below)

Ankur-singh · 2025-02-15T05:13:25Z

Sorry guys, I was using python 3.11 and everything seemed to work. It was a typing issue, learnt that in python 3.9 one can't use subscripted generics like Union[. . .] in runtime checks. I was using

if isinstance(x, Union[Image.Image, torch.Tensor]):
    ...

I change it to

if isinstance(x, (Image.Image, torch.Tensor)):
    ...

this seems to work in both python 3.9 and 3.11.

Also, any suggestions on how to manage multiple python versions? What does the workflow look like? Docker, uv, etc.

codecov-commenter · 2025-02-15T20:16:19Z

Codecov Report

Attention: Patch coverage is 21.87500% with 25 lines in your changes missing coverage. Please review.

Project coverage is 23.46%. Comparing base (f67ccda) to head (4150ee0).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
tests/torchtune/data/test_data_utils.py	9.09%	10 Missing ⚠️
torchtune/data/_utils.py	30.00%	7 Missing ⚠️
...torchtune/models/clip/test_clip_image_transform.py	0.00%	4 Missing ⚠️
torchtune/models/clip/_transform.py	0.00%	2 Missing ⚠️
.../torchtune/datasets/multimodal/test_vqa_dataset.py	50.00%	1 Missing ⚠️
torchtune/models/clip/inference/_transform.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2366   +/-   ##
=======================================
  Coverage   23.46%   23.46%           
=======================================
  Files         373      373           
  Lines       22405    22420   +15     
=======================================
+ Hits         5257     5261    +4     
- Misses      17148    17159   +11

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

ebsmothers · 2025-02-24T17:21:44Z

Hi guys, sorry I've been a bit out of the loop on this one. @Ankur-singh what's the status here? Is this ready for another round of review?

Ankur-singh · 2025-02-24T17:33:42Z

@ebsmothers its read for review. Should work 🤞

joecummings

lgtm

…syntax

joecummings · 2025-02-25T16:23:01Z

Thanks for all the hard work @Ankur-singh !!!@

…-pytorch#2366)

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Feb 8, 2025

joecummings reviewed Feb 10, 2025

View reviewed changes

Ankur-singh mentioned this pull request Feb 11, 2025

Move from PIL to torchvision.io.decode_image #2303

Open

NicolasHug reviewed Feb 12, 2025

View reviewed changes

Ankur-singh force-pushed the use-decode-image branch from f77145f to 3f04994 Compare February 14, 2025 07:32

Ankur-singh force-pushed the use-decode-image branch from 4150ee0 to 3cd5709 Compare February 20, 2025 04:22

joecummings approved these changes Feb 24, 2025

View reviewed changes

Ankur-singh added 4 commits February 24, 2025 22:53

Refactor load_image to use torch.Tensor instead of PIL.Image

4dfffac

Improve error messages in load_image function for clarity

b3554cc

Add RGB mode support for image loading in load_image function

ad9a541

Fix type checking for input image in CLIPImageTransform to use tuple …

a069f8a

…syntax

Ankur-singh force-pushed the use-decode-image branch from 3cd5709 to a069f8a Compare February 25, 2025 06:53

joecummings merged commit 0e8f840 into meta-pytorch:main Feb 25, 2025
17 checks passed

joecummings pushed a commit to joecummings/torchtune that referenced this pull request Feb 27, 2025

Refactor load_image to return torch.Tensor instead of PIL.Image (meta…

8a2250a

…-pytorch#2366)

joecummings pushed a commit to joecummings/torchtune that referenced this pull request Feb 27, 2025

Refactor load_image to return torch.Tensor instead of PIL.Image (meta…

2803b92

…-pytorch#2366)

Ankur-singh deleted the use-decode-image branch March 4, 2025 07:25

TJ-Solergibert mentioned this pull request Mar 27, 2025

Adding OBELICS DataLoader pytorch/torchtitan#663

Merged



		def load_image(image_loc: Union[Path, str]) -> "PIL.Image.Image":
		def load_image(image_loc: Union[Path, str]) -> "torch.Tensor":

Refactor load_image to return torch.Tensor instead of PIL.Image #2366

Refactor load_image to return torch.Tensor instead of PIL.Image #2366

Uh oh!

Conversation

Ankur-singh commented Feb 8, 2025

Context

Changelog

Test plan

UX

Uh oh!

pytorch-bot bot commented Feb 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2366

✅ No Failures

Uh oh!

joecummings left a comment

Choose a reason for hiding this comment

Uh oh!

joecummings left a comment

Choose a reason for hiding this comment

Uh oh!

joecummings Feb 10, 2025

Choose a reason for hiding this comment

Uh oh!

joecummings Feb 10, 2025

Choose a reason for hiding this comment

Uh oh!

joecummings Feb 10, 2025

Choose a reason for hiding this comment

Uh oh!

joecummings Feb 10, 2025

Choose a reason for hiding this comment

Uh oh!

Ankur-singh Feb 15, 2025

Choose a reason for hiding this comment

Uh oh!

NicolasHug commented Feb 10, 2025

Uh oh!

NicolasHug Feb 12, 2025

Choose a reason for hiding this comment

Uh oh!

Ankur-singh commented Feb 15, 2025

Uh oh!

codecov-commenter commented Feb 15, 2025

Codecov Report

Uh oh!

ebsmothers commented Feb 24, 2025

Uh oh!

Ankur-singh commented Feb 24, 2025

Uh oh!

joecummings left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

joecummings commented Feb 25, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

pytorch-bot bot commented Feb 8, 2025 •

edited

Loading