Improve error messages in load_image function for clarity

Author: Ankur-singh

tests/torchtune/data/test_data_utils.py

@@ -122,7 +122,7 @@ def mock_urlopen(url):
     assert image.size() == (3, 403, 580)
 
     # Test that a ValueError is raised when the image path is invalid
-    with pytest.raises(ValueError, match="Failed to open image as torch.Tensor"):
+    with pytest.raises(ValueError, match="Failed to load local image as torch.Tensor"):
         load_image("invalid_path")
 
     # Test a temporary file with invalid image data
@@ -131,16 +131,16 @@ def mock_urlopen(url):
         f.write("Invalid image data")
 
     # Test that a ValueError is raised when the image data is invalid
-    with pytest.raises(ValueError, match="Failed to open image as torch.Tensor"):
+    with pytest.raises(ValueError, match="Failed to load local image as torch.Tensor"):
         load_image(str(image_path))
 
     # Test that a ValueError is raised when there is an HTTP error
     # Mock the urlopen function to raise an exception
     def mock_urlopen(url):
-        raise Exception("Failed to load image")
+        raise Exception("Failed to load remote image as torch.Tensor")
 
     monkeypatch.setattr("urllib.request.urlopen", mock_urlopen)
-    with pytest.raises(ValueError, match="Failed to load image"):
+    with pytest.raises(ValueError, match="Failed to load remote image as torch.Tensor"):
         load_image("http://example.com/test_image.jpg")
 
     # Test that a ValueError is raised when there is an IO error
@@ -149,7 +149,7 @@ def mock_urlopen(url):
     with open(image_path, "w") as f:
         f.write("Test data")
     os.chmod(image_path, 0o000)  # Remove read permissions
-    with pytest.raises(ValueError, match="Failed to open image as torch.Tensor"):
+    with pytest.raises(ValueError, match="Failed to load local image as torch.Tensor"):
         load_image(str(image_path))
     os.chmod(image_path, 0o644)  # Restore read permissions
 
@@ -158,5 +158,5 @@ def mock_urlopen(url):
     image_path = tmp_path / "test_image.jpg"
     with open(image_path, "wb") as f:
         f.write(b"Invalid image data")
-    with pytest.raises(ValueError, match="Failed to open image as torch.Tensor"):
+    with pytest.raises(ValueError, match="Failed to load local image as torch.Tensor"):
         load_image(str(image_path))

torchtune/data/_utils.py

@@ -46,7 +46,7 @@ def truncate(
     return tokens_truncated
 
 
-def load_image(image_loc: Union[Path, str]) -> "torch.Tensor":
+def load_image(image_loc: Union[Path, str]) -> torch.Tensor:
     """
     Convenience method to load an image in torch.Tensor format from a local file path or remote source.
 
@@ -81,14 +81,14 @@ def load_image(image_loc: Union[Path, str]) -> "torch.Tensor":
                 torch.frombuffer(image_loc, dtype=torch.uint8)
             )
         except Exception as e:
-            raise ValueError("Failed to load image as torch.Tensor") from e
+            raise ValueError("Failed to load remote image as torch.Tensor") from e
 
     # Open the local image as a Tensor image
     else:
         try:
             image = torchvision.io.decode_image(image_loc)
         except Exception as e:
-            raise ValueError("Failed to open image as torch.Tensor") from e
+            raise ValueError("Failed to load local image as torch.Tensor") from e
     return image