Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix import & asignement issue #377

Merged
merged 5 commits into from
Sep 18, 2024
Merged

fix import & asignement issue #377

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b496cbd
fix import & asignement issue
Borda Sep 18, 2024
9d6d6d2
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Sep 18, 2024
8009f3f
Apply suggestions from code review
Borda Sep 18, 2024
3fd0fa1
Merge branch 'main' into fix/var-import
bhimrazy Sep 18, 2024
e1aad7f
Merge branch 'main' into fix/var-import
Borda Sep 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions src/litdata/streaming/serializers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@

Borda marked this conversation as resolved.
Show resolved Hide resolved
if TYPE_CHECKING:
from PIL.JpegImagePlugin import JpegImageFile

_PIL_AVAILABLE = RequirementCache("PIL")
_TORCH_VISION_AVAILABLE = RequirementCache("torchvision")
_AV_AVAILABLE = RequirementCache("av")
Expand Down Expand Up @@ -70,6 +69,8 @@ def serialize(self, item: Any) -> Tuple[bytes, Optional[str]]:

@classmethod
def deserialize(cls, data: bytes) -> Any:
if not _PIL_AVAILABLE:
raise ModuleNotFoundError("PIL is required. Run `pip install pillow`")
from PIL import Image

idx = 3 * 4
Expand All @@ -94,6 +95,9 @@ class JPEGSerializer(Serializer):
"""The JPEGSerializer serialize and deserialize JPEG image to and from bytes."""

def serialize(self, item: Any) -> Tuple[bytes, Optional[str]]:
if not _PIL_AVAILABLE:
raise ModuleNotFoundError("PIL is required. Run `pip install pillow`")

from PIL import Image
from PIL.GifImagePlugin import GifImageFile
from PIL.JpegImagePlugin import JpegImageFile
Expand Down Expand Up @@ -122,7 +126,7 @@ def serialize(self, item: Any) -> Tuple[bytes, Optional[str]]:
buff.seek(0)
return buff.read(), None

raise TypeError(f"The provided item should be of type {JpegImageFile}. Found {item}.")
raise TypeError(f"The provided item should be of type `JpegImageFile`. Found {item}.")

def deserialize(self, data: bytes) -> Union["JpegImageFile", torch.Tensor]:
if _TORCH_VISION_AVAILABLE:
Expand Down Expand Up @@ -184,7 +188,7 @@ def deserialize(self, data: bytes) -> torch.Tensor:
shape = []
for shape_idx in range(shape_size):
shape.append(np.frombuffer(data[8 + 4 * shape_idx : 8 + 4 * (shape_idx + 1)], np.uint32).item())
idx_start = 8 + 4 * (shape_idx + 1)
idx_start = 8 + 4 * shape_size
idx_end = len(data)
if idx_end > idx_start:
tensor = torch.frombuffer(data[idx_start:idx_end], dtype=dtype)
Expand Down Expand Up @@ -250,7 +254,7 @@ def deserialize(self, data: bytes) -> np.ndarray:
shape.append(np.frombuffer(data[8 + 4 * shape_idx : 8 + 4 * (shape_idx + 1)], np.uint32).item())

# deserialize the numpy array bytes
tensor = np.frombuffer(data[8 + 4 * (shape_idx + 1) : len(data)], dtype=dtype)
tensor = np.frombuffer(data[8 + 4 * shape_size : len(data)], dtype=dtype)
if tensor.shape == shape:
return tensor
return np.reshape(tensor, shape)
Expand Down
Loading