JPEGSerializer: Fix serializer io.bytes image (#19369)

Author: tchaton

src/lightning/data/streaming/data_processor.py

@@ -151,6 +151,7 @@ def _download_data_target(input_dir: Dir, cache_dir: str, queue_in: Queue, queue
                         s3.client.download_fileobj(obj.netloc, obj.path.lstrip("/"), f)
 
                 elif os.path.isfile(path):
+                    os.makedirs(os.path.dirname(local_path), exist_ok=True)
                     shutil.copyfile(path, local_path)
                 else:
                     raise ValueError(f"The provided {input_dir.url} isn't supported.")

src/lightning/data/streaming/serializers.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import io
 import os
 import pickle
 import tempfile
@@ -109,8 +110,16 @@ def serialize(self, item: Image) -> Tuple[bytes, Optional[str]]:
                 raise ValueError(
                     "The JPEG Image's filename isn't defined. HINT: Open the image in your Dataset __getitem__ method."
                 )
-            with open(item.filename, "rb") as f:
-                return f.read(), None
+            if item.filename and os.path.exists(item.filename):
+                # read the content of the file directly
+                with open(item.filename, "rb") as f:
+                    return f.read(), None
+            else:
+                item_bytes = io.BytesIO()
+                item.save(item_bytes, format="JPEG")
+                item_bytes = item_bytes.getvalue()
+                return item_bytes, None
+
         raise TypeError(f"The provided itemect should be of type {JpegImageFile}. Found {item}.")
 
     def deserialize(self, data: bytes) -> Union[JpegImageFile, torch.Tensor]:

tests/tests_data/streaming/test_serializer.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import io
 import os
 import sys
 from time import time
@@ -26,6 +27,7 @@
     _TORCH_DTYPES_MAPPING,
     _TORCH_VISION_AVAILABLE,
     IntSerializer,
+    JPEGSerializer,
     NoHeaderNumpySerializer,
     NoHeaderTensorSerializer,
     NumpySerializer,
@@ -86,6 +88,27 @@ def test_pil_serializer(mode):
     assert np.array_equal(np_data, np_dec_data)
 
 
+@pytest.mark.skipif(condition=not _PIL_AVAILABLE, reason="Requires: ['pil']")
+def test_jpeg_serializer():
+    serializer = JPEGSerializer()
+
+    from PIL import Image
+
+    array = np.random.randint(255, size=(28, 28, 3), dtype=np.uint8)
+    img = Image.fromarray(array)
+    img_bytes = io.BytesIO()
+    img.save(img_bytes, format="JPEG")
+    img_bytes = img_bytes.getvalue()
+
+    img = Image.open(io.BytesIO(img_bytes))
+
+    data, _ = serializer.serialize(img)
+    assert isinstance(data, bytes)
+
+    deserialized_img = serializer.deserialize(data)
+    assert deserialized_img.shape == torch.Size([3, 28, 28])
+
+
 @pytest.mark.flaky(reruns=3)
 @pytest.mark.skipif(sys.platform == "win32", reason="Not supported on windows")
 def test_tensor_serializer():