Improve models and file methods (#966)

* Improve models and file methods

- BBox model now have methods to convert to/from different bbox formats:
  Albumentations, COCO, PASCAL VOC and YOLO
- Add more validations and tests for BBox, Pose and Segment models
- Add tests for Ultralytics models
- Add File conversion methos: to TextFile, ImageFile and VideoFile
  (see Ultralytcs examples for usage)
- Add 'image_info' method to get meta from ImageFile (same as for video
  file)

* Update 'openimage-detect' example to use new BBox.from_albumentations method

* Add more tests, refactor File class 'read' and 'read_bytes' methods

Author: dreadatour

examples/computer_vision/openimage-detect.py

@@ -22,13 +22,9 @@ def openimage_detect(args):
         detections = json.load(stream_json).get("detections", [])
 
     for i, detect in enumerate(detections):
-        bbox = model.BBox.from_list(
-            [
-                detect["XMin"] * img.width,
-                detect["XMax"] * img.width,
-                detect["YMin"] * img.height,
-                detect["YMax"] * img.height,
-            ]
+        bbox = model.BBox.from_albumentations(
+            [detect[k] for k in ("XMin", "YMin", "XMax", "YMax")],
+            img_size=(img.width, img.height),
         )
 
         fstream = File(

examples/computer_vision/ultralytics-bbox.py

@@ -1,19 +1,11 @@
-import os
-
-os.environ["YOLO_VERBOSE"] = "false"
-
-
-from io import BytesIO
-
-from PIL import Image
 from ultralytics import YOLO
 
 from datachain import C, DataChain, File
 from datachain.model.ultralytics import YoloBBoxes
 
 
 def process_bboxes(yolo: YOLO, file: File) -> YoloBBoxes:
-    results = yolo(Image.open(BytesIO(file.read())))
+    results = yolo(file.as_image_file().read(), verbose=False)
     return YoloBBoxes.from_results(results)
 
 

examples/computer_vision/ultralytics-pose.py

@@ -1,19 +1,11 @@
-import os
-
-os.environ["YOLO_VERBOSE"] = "false"
-
-
-from io import BytesIO
-
-from PIL import Image
 from ultralytics import YOLO
 
 from datachain import C, DataChain, File
 from datachain.model.ultralytics import YoloPoses
 
 
 def process_poses(yolo: YOLO, file: File) -> YoloPoses:
-    results = yolo(Image.open(BytesIO(file.read())))
+    results = yolo(file.as_image_file().read(), verbose=False)
     return YoloPoses.from_results(results)
 
 

examples/computer_vision/ultralytics-segment.py

@@ -1,19 +1,11 @@
-import os
-
-os.environ["YOLO_VERBOSE"] = "false"
-
-
-from io import BytesIO
-
-from PIL import Image
 from ultralytics import YOLO
 
 from datachain import C, DataChain, File
 from datachain.model.ultralytics import YoloSegments
 
 
 def process_segments(yolo: YOLO, file: File) -> YoloSegments:
-    results = yolo(Image.open(BytesIO(file.read())))
+    results = yolo(file.as_image_file().read(), verbose=False)
     return YoloSegments.from_results(results)
 
 

pyproject.toml

@@ -102,7 +102,8 @@ tests = [
   "hypothesis",
   "aiotools>=1.7.0",
   "requests-mock",
-  "scipy"
+  "scipy",
+  "ultralytics"
 ]
 dev = [
   "datachain[docs,tests]",
@@ -118,7 +119,7 @@ examples = [
   "defusedxml",
   "accelerate",
   "huggingface_hub[hf_transfer]",
-  "ultralytics==8.3.87",
+  "ultralytics",
   "open_clip_torch"
 ]
 

src/datachain/lib/file.py

@@ -242,6 +242,30 @@ def __init__(self, **kwargs):
         self._catalog = None
         self._caching_enabled: bool = False
 
+    def as_text_file(self) -> "TextFile":
+        """Convert the file to a `TextFile` object."""
+        if isinstance(self, TextFile):
+            return self
+        file = TextFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
+    def as_image_file(self) -> "ImageFile":
+        """Convert the file to a `ImageFile` object."""
+        if isinstance(self, ImageFile):
+            return self
+        file = ImageFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
+    def as_video_file(self) -> "VideoFile":
+        """Convert the file to a `VideoFile` object."""
+        if isinstance(self, VideoFile):
+            return self
+        file = VideoFile(**self.model_dump())
+        file._set_stream(self._catalog, caching_enabled=self._caching_enabled)
+        return file
+
     @classmethod
     def upload(
         cls, data: bytes, path: str, catalog: Optional["Catalog"] = None
@@ -291,20 +315,20 @@ def open(self, mode: Literal["rb", "r"] = "rb") -> Iterator[Any]:
             ) as f:
                 yield io.TextIOWrapper(f) if mode == "r" else f
 
-    def read(self, length: int = -1):
-        """Returns file contents."""
+    def read_bytes(self, length: int = -1):
+        """Returns file contents as bytes."""
         with self.open() as stream:
             return stream.read(length)
 
-    def read_bytes(self):
-        """Returns file contents as bytes."""
-        return self.read()
-
     def read_text(self):
         """Returns file contents as text."""
         with self.open(mode="r") as stream:
             return stream.read()
 
+    def read(self, length: int = -1):
+        """Returns file contents."""
+        return self.read_bytes(length)
+
     def save(self, destination: str, client_config: Optional[dict] = None):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
@@ -547,20 +571,36 @@ def save(self, destination: str, client_config: Optional[dict] = None):
 class ImageFile(File):
     """`DataModel` for reading image files."""
 
+    def get_info(self) -> "Image":
+        """
+        Retrieves metadata and information about the image file.
+
+        Returns:
+            Image: A Model containing image metadata such as width, height and format.
+        """
+        from .image import image_info
+
+        return image_info(self)
+
     def read(self):
         """Returns `PIL.Image.Image` object."""
         from PIL import Image as PilImage
 
         fobj = super().read()
         return PilImage.open(BytesIO(fobj))
 
-    def save(self, destination: str, client_config: Optional[dict] = None):
+    def save(  # type: ignore[override]
+        self,
+        destination: str,
+        format: Optional[str] = None,
+        client_config: Optional[dict] = None,
+    ):
         """Writes it's content to destination"""
         destination = stringify_path(destination)
 
         client: Client = self._catalog.get_client(destination, **(client_config or {}))
         with client.fs.open(destination, mode="wb") as f:
-            self.read().save(f)
+            self.read().save(f, format=format)
 
 
 class Image(DataModel):

src/datachain/lib/image.py

@@ -1,17 +1,41 @@
 from typing import Callable, Optional, Union
 
 import torch
-from PIL import Image
+from PIL import Image as PILImage
+
+from datachain.lib.file import File, FileError, Image, ImageFile
+
+
+def image_info(file: Union[File, ImageFile]) -> Image:
+    """
+    Returns image file information.
+
+    Args:
+        file (ImageFile): Image file object.
+
+    Returns:
+        Image: Image file information.
+    """
+    try:
+        img = file.as_image_file().read()
+    except Exception as exc:
+        raise FileError(file, "unable to open image file") from exc
+
+    return Image(
+        width=img.width,
+        height=img.height,
+        format=img.format or "",
+    )
 
 
 def convert_image(
-    img: Image.Image,
+    img: PILImage.Image,
     mode: str = "RGB",
     size: Optional[tuple[int, int]] = None,
     transform: Optional[Callable] = None,
     encoder: Optional[Callable] = None,
     device: Optional[Union[str, torch.device]] = None,
-) -> Union[Image.Image, torch.Tensor]:
+) -> Union[PILImage.Image, torch.Tensor]:
     """
     Resize, transform, and otherwise convert an image.
 
@@ -47,13 +71,13 @@ def convert_image(
 
 
 def convert_images(
-    images: Union[Image.Image, list[Image.Image]],
+    images: Union[PILImage.Image, list[PILImage.Image]],
     mode: str = "RGB",
     size: Optional[tuple[int, int]] = None,
     transform: Optional[Callable] = None,
     encoder: Optional[Callable] = None,
     device: Optional[Union[str, torch.device]] = None,
-) -> Union[list[Image.Image], torch.Tensor]:
+) -> Union[list[PILImage.Image], torch.Tensor]:
     """
     Resize, transform, and otherwise convert one or more images.
 
@@ -65,7 +89,7 @@ def convert_images(
         encoder (Callable): Encode image using model.
         device (str or torch.device): Device to use.
     """
-    if isinstance(images, Image.Image):
+    if isinstance(images, PILImage.Image):
         images = [images]
 
     converted = [

src/datachain/lib/video.py

@@ -1,11 +1,11 @@
 import posixpath
 import shutil
 import tempfile
-from typing import Optional
+from typing import Optional, Union
 
 from numpy import ndarray
 
-from datachain.lib.file import FileError, ImageFile, Video, VideoFile
+from datachain.lib.file import File, FileError, ImageFile, Video, VideoFile
 
 try:
     import ffmpeg
@@ -18,7 +18,7 @@
     ) from exc
 
 
-def video_info(file: VideoFile) -> Video:
+def video_info(file: Union[File, VideoFile]) -> Video:
     """
     Returns video file information.
 
@@ -28,6 +28,8 @@ def video_info(file: VideoFile) -> Video:
     Returns:
         Video: Video file information.
     """
+    file = file.as_video_file()
+
     if not (file_path := file.get_local_path()):
         file.ensure_cached()
         file_path = file.get_local_path()