Refactor load_image to return torch.Tensor instead of PIL.Image #2366
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2366
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 3 Cancelled JobsAs of commit f77145f with merge base 8c9235e ( NEW FAILURES - The following jobs have failed:
CANCELLED JOBS - The following jobs were cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
torchtune/data/_utils.py
Outdated
| @@ -44,9 +46,9 @@ def truncate( | |||
| return tokens_truncated | |||
|
|
|||
|
|
|||
| def load_image(image_loc: Union[Path, str]) -> "PIL.Image.Image": | |||
| def load_image(image_loc: Union[Path, str]) -> "torch.Tensor": | |||
There was a problem hiding this comment.
You don't need to quote the return type as torch.Tensor is definitely already importable.
torchtune/data/_utils.py
Outdated
| 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 |
There was a problem hiding this comment.
Can you be more specific with this error? Something like failed to load image from remote.
torchtune/data/_utils.py
Outdated
| try: | ||
| image = torchvision.io.decode_image(image_loc) | ||
| except Exception as e: | ||
| raise ValueError("Failed to open image as torch.Tensor") from e |
There was a problem hiding this comment.
Same here, lets be specific that this fails to load from a local image.
| @@ -252,9 +252,13 @@ def __init__( | |||
| antialias=self.antialias, | |||
| ) | |||
|
|
|||
| def __call__(self, *, image: Image.Image, **kwargs) -> Mapping[str, Any]: | |||
| def __call__( | |||
| self, *, image: Union[Image.Image, torch.Tensor], **kwargs | |||
There was a problem hiding this comment.
cc @pbontrager Do we actually want to allow PIL images through here? Or limit to just tensor?
This SGTM |
| 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) |
There was a problem hiding this comment.
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)
Context
What is the purpose of this PR? Is it to
Please link to any issues this PR addresses. #2303
Changelog
What are the changes made in this PR?
load_imageto returntorch.Tensorinstead ofPIL.ImageCLIPImageTransformto support bothtorch.TensorandPIL.ImageCLIPImageTransformwithtorch.TensorinputTest 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.
pre-commit install)pytest testspytest tests -m integration_testUX
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