datachain 0.11.11__py3-none-any.whl → 0.13.0__py3-none-any.whl

datachain/lib/udf.py

@@ -159,6 +159,7 @@ class UDFBase(AbstractUDF):
         ```
     """
 
+    is_input_batched = False
     is_output_batched = False
     prefetch: int = 0
 
@@ -395,6 +396,7 @@ class Mapper(UDFBase):
 class BatchMapper(UDFBase):
     """Inherit from this class to pass to `DataChain.batch_map()`."""
 
+    is_input_batched = True
     is_output_batched = True
 
     def run(
@@ -481,6 +483,7 @@ class Generator(UDFBase):
 class Aggregator(UDFBase):
     """Inherit from this class to pass to `DataChain.agg()`."""
 
+    is_input_batched = True
     is_output_batched = True
 
     def run(

datachain/lib/udf_signature.py

@@ -1,7 +1,7 @@
 import inspect
 from collections.abc import Generator, Iterator, Sequence
 from dataclasses import dataclass
-from typing import Callable, Union, get_args, get_origin
+from typing import Any, Callable, Union, get_args, get_origin
 
 from datachain.lib.data_model import DataType, DataTypeNames, is_chain_type
 from datachain.lib.signal_schema import SignalSchema
@@ -18,7 +18,7 @@ class UdfSignatureError(DataChainParamsError):
 @dataclass
 class UdfSignature:
     func: Union[Callable, UDFBase]
-    params: Sequence[str]
+    params: dict[str, Union[DataType, Any]]
     output_schema: SignalSchema
 
     DEFAULT_RETURN_TYPE = str
@@ -58,15 +58,23 @@ class UdfSignature:
         if not isinstance(udf_func, UDFBase) and not callable(udf_func):
             raise UdfSignatureError(chain, f"UDF '{udf_func}' is not callable")
 
-        func_params_map_sign, func_outs_sign, is_iterator = (
-            UdfSignature._func_signature(chain, udf_func)
+        func_params_map_sign, func_outs_sign, is_iterator = cls._func_signature(
+            chain, udf_func
         )
+
+        udf_params: dict[str, Union[DataType, Any]] = {}
         if params:
-            udf_params = [params] if isinstance(params, str) else params
-        elif not func_params_map_sign:
-            udf_params = []
-        else:
-            udf_params = list(func_params_map_sign.keys())
+            udf_params = (
+                {params: Any} if isinstance(params, str) else dict.fromkeys(params, Any)
+            )
+        elif func_params_map_sign:
+            udf_params = {
+                param: (
+                    param_type if param_type is not inspect.Parameter.empty else Any
+                )
+                for param, param_type in func_params_map_sign.items()
+            }
+
         if output:
             udf_output_map = UdfSignature._validate_output(
                 chain, signal_name, func, func_outs_sign, output

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 @@ except ImportError as exc:
     ) 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()
@@ -170,7 +172,7 @@ def save_video_frame(
     output_file = posixpath.join(
         output, f"{video.get_file_stem()}_{frame:04d}.{format}"
     )
-    return ImageFile.upload(img, output_file)
+    return ImageFile.upload(img, output_file, catalog=video._catalog)
 
 
 def save_video_fragment(
@@ -218,6 +220,6 @@ def save_video_fragment(
         ).output(output_file_tmp).run(quiet=True)
 
         with open(output_file_tmp, "rb") as f:
-            return VideoFile.upload(f.read(), output_file)
+            return VideoFile.upload(f.read(), output_file, catalog=video._catalog)
     finally:
         shutil.rmtree(temp_dir)

datachain/model/bbox.py

@@ -1,47 +1,216 @@
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Union
+
 from pydantic import Field
 
 from datachain.lib.data_model import DataModel
 
+from .utils import convert_bbox, validate_bbox
+
+if TYPE_CHECKING:
+    from .pose import Pose, Pose3D
+
 
 class BBox(DataModel):
     """
-    A data model for representing bounding box.
+    A data model representing a bounding box.
 
     Attributes:
-        title (str): The title of the bounding box.
-        coords (list[int]): The coordinates of the bounding box.
+        title (str): The title or label associated with the bounding box.
+        coords (list[int]): A list of four bounding box coordinates.
 
-    The bounding box is defined by two points:
-        - (x1, y1): The top-left corner of the box.
-        - (x2, y2): The bottom-right corner of the box.
+    The bounding box follows the PASCAL VOC format, where:
+        - (x1, y1) represents the pixel coordinates of the top-left corner.
+        - (x2, y2) represents the pixel coordinates of the bottom-right corner.
     """
 
     title: str = Field(default="")
     coords: list[int] = Field(default=[])
 
     @staticmethod
-    def from_list(coords: list[float], title: str = "") -> "BBox":
-        assert len(coords) == 4, "Bounding box must be a list of 4 coordinates."
-        assert all(isinstance(value, (int, float)) for value in coords), (
-            "Bounding box coordinates must be floats or integers."
-        )
-        return BBox(
-            title=title,
-            coords=[round(c) for c in coords],
+    def from_albumentations(
+        coords: Sequence[float],
+        img_size: Sequence[int],
+        title: str = "",
+    ) -> "BBox":
+        """
+        Create a bounding box from Albumentations format.
+
+        Albumentations represents bounding boxes as `[x_min, y_min, x_max, y_max]`
+        with normalized coordinates (values between 0 and 1) relative to the image size.
+
+        Args:
+            coords (Sequence[float]): The bounding box coordinates in
+                Albumentations format.
+            img_size (Sequence[int]): The reference image size as `[width, height]`.
+            title (str, optional): The title or label of the bounding box.
+                Defaults to an empty string.
+
+        Returns:
+            BBox: The bounding box data model.
+        """
+        validate_bbox(coords, float)
+        bbox_coords = convert_bbox(coords, img_size, "albumentations", "voc")
+        return BBox(title=title, coords=list(map(round, bbox_coords)))
+
+    def to_albumentations(self, img_size: Sequence[int]) -> list[float]:
+        """
+        Convert the bounding box coordinates to Albumentations format.
+
+        Albumentations represents bounding boxes as `[x_min, y_min, x_max, y_max]`
+        with normalized coordinates (values between 0 and 1) relative to the image size.
+
+        Args:
+            img_size (Sequence[int]): The reference image size as `[width, height]`.
+
+        Returns:
+            list[float]: The bounding box coordinates in Albumentations format.
+        """
+        return convert_bbox(self.coords, img_size, "voc", "albumentations")
+
+    @staticmethod
+    def from_coco(
+        coords: Sequence[float],
+        title: str = "",
+    ) -> "BBox":
+        """
+        Create a bounding box from COCO format.
+
+        COCO format represents bounding boxes as [x_min, y_min, width, height], where:
+        - (x_min, y_min) are the pixel coordinates of the top-left corner.
+        - width and height define the size of the bounding box in pixels.
+
+        Args:
+            coords (Sequence[float]): The bounding box coordinates in COCO format.
+            title (str): The title of the bounding box.
+
+        Returns:
+            BBox: The bounding box data model.
+        """
+        validate_bbox(coords, float, int)
+        bbox_coords = convert_bbox(coords, [], "coco", "voc")
+        return BBox(title=title, coords=list(map(round, bbox_coords)))
+
+    def to_coco(self) -> list[int]:
+        """
+        Return the bounding box coordinates in COCO format.
+
+        COCO format represents bounding boxes as [x_min, y_min, width, height], where:
+        - (x_min, y_min) are the pixel coordinates of the top-left corner.
+        - width and height define the size of the bounding box in pixels.
+
+        Returns:
+            list[int]: The bounding box coordinates in COCO format.
+        """
+        res = convert_bbox(self.coords, [], "voc", "coco")
+        return list(map(round, res))
+
+    @staticmethod
+    def from_voc(
+        coords: Sequence[float],
+        title: str = "",
+    ) -> "BBox":
+        """
+        Create a bounding box from PASCAL VOC format.
+
+        PASCAL VOC format represents bounding boxes as [x_min, y_min, x_max, y_max],
+        where:
+        - (x_min, y_min) are the pixel coordinates of the top-left corner.
+        - (x_max, y_max) are the pixel coordinates of the bottom-right corner.
+
+        Args:
+            coords (Sequence[float]): The bounding box coordinates in VOC format.
+            title (str): The title of the bounding box.
+
+        Returns:
+            BBox: The bounding box data model.
+        """
+        validate_bbox(coords, float, int)
+        return BBox(title=title, coords=list(map(round, coords)))
+
+    def to_voc(self) -> list[int]:
+        """
+        Return the bounding box coordinates in PASCAL VOC format.
+
+        PASCAL VOC format represents bounding boxes as [x_min, y_min, x_max, y_max],
+        where:
+        - (x_min, y_min) are the pixel coordinates of the top-left corner.
+        - (x_max, y_max) are the pixel coordinates of the bottom-right corner.
+
+        Returns:
+            list[int]: The bounding box coordinates in VOC format.
+        """
+        return self.coords
+
+    @staticmethod
+    def from_yolo(
+        coords: Sequence[float],
+        img_size: Sequence[int],
+        title: str = "",
+    ) -> "BBox":
+        """
+        Create a bounding box from YOLO format.
+
+        YOLO format represents bounding boxes as [x_center, y_center, width, height],
+        where:
+        - (x_center, y_center) are the normalized coordinates of the box center.
+        - width and height normalized values define the size of the bounding box.
+
+        Args:
+            coords (Sequence[float]): The bounding box coordinates in YOLO format.
+            img_size (Sequence[int]): The reference image size as `[width, height]`.
+            title (str): The title of the bounding box.
+
+        Returns:
+            BBox: The bounding box data model.
+        """
+        validate_bbox(coords, float)
+        bbox_coords = convert_bbox(coords, img_size, "yolo", "voc")
+        return BBox(title=title, coords=list(map(round, bbox_coords)))
+
+    def to_yolo(self, img_size: Sequence[int]) -> list[float]:
+        """
+        Return the bounding box coordinates in YOLO format.
+
+        YOLO format represents bounding boxes as [x_center, y_center, width, height],
+        where:
+        - (x_center, y_center) are the normalized coordinates of the box center.
+        - width and height normalized values define the size of the bounding box.
+
+        Args:
+            img_size (Sequence[int]): The reference image size as `[width, height]`.
+
+        Returns:
+            list[float]: The bounding box coordinates in YOLO format.
+        """
+        return convert_bbox(self.coords, img_size, "voc", "yolo")
+
+    def point_inside(self, x: int, y: int) -> bool:
+        """
+        Return True if the point is inside the bounding box.
+
+        Assumes that if the point is on the edge of the bounding box,
+        it is considered inside.
+        """
+        x1, y1, x2, y2 = self.coords
+        return x1 <= x <= x2 and y1 <= y <= y2
+
+    def pose_inside(self, pose: Union["Pose", "Pose3D"]) -> bool:
+        """Return True if the pose is inside the bounding box."""
+        return all(
+            self.point_inside(x, y) for x, y in zip(pose.x, pose.y) if x > 0 or y > 0
         )
 
+    @staticmethod
+    def from_list(coords: Sequence[float], title: str = "") -> "BBox":
+        return BBox.from_voc(coords, title=title)
+
     @staticmethod
     def from_dict(coords: dict[str, float], title: str = "") -> "BBox":
-        assert isinstance(coords, dict) and set(coords) == {
-            "x1",
-            "y1",
-            "x2",
-            "y2",
-        }, "Bounding box must be a dictionary with keys 'x1', 'y1', 'x2' and 'y2'."
-        return BBox.from_list(
-            [coords["x1"], coords["y1"], coords["x2"], coords["y2"]],
-            title=title,
-        )
+        keys = ("x1", "y1", "x2", "y2")
+        if not isinstance(coords, dict) or set(coords) != set(keys):
+            raise ValueError("Bounding box must be a dictionary with coordinates.")
+        return BBox.from_voc([coords[k] for k in keys], title=title)
 
 
 class OBBox(DataModel):
@@ -63,40 +232,22 @@ class OBBox(DataModel):
     coords: list[int] = Field(default=[])
 
     @staticmethod
-    def from_list(coords: list[float], title: str = "") -> "OBBox":
-        assert len(coords) == 8, (
-            "Oriented bounding box must be a list of 8 coordinates."
-        )
-        assert all(isinstance(value, (int, float)) for value in coords), (
-            "Oriented bounding box coordinates must be floats or integers."
-        )
-        return OBBox(
-            title=title,
-            coords=[round(c) for c in coords],
-        )
+    def from_list(coords: Sequence[float], title: str = "") -> "OBBox":
+        if not isinstance(coords, (list, tuple)):
+            raise TypeError("Oriented bounding box must be a list of coordinates.")
+        if len(coords) != 8:
+            raise ValueError("Oriented bounding box must have 8 coordinates.")
+        if not all(isinstance(value, (int, float)) for value in coords):
+            raise ValueError(
+                "Oriented bounding box coordinates must be floats or integers."
+            )
+        return OBBox(title=title, coords=list(map(round, coords)))
 
     @staticmethod
     def from_dict(coords: dict[str, float], title: str = "") -> "OBBox":
-        assert isinstance(coords, dict) and set(coords) == {
-            "x1",
-            "y1",
-            "x2",
-            "y2",
-            "x3",
-            "y3",
-            "x4",
-            "y4",
-        }, "Oriented bounding box must be a dictionary with coordinates."
-        return OBBox.from_list(
-            [
-                coords["x1"],
-                coords["y1"],
-                coords["x2"],
-                coords["y2"],
-                coords["x3"],
-                coords["y3"],
-                coords["x4"],
-                coords["y4"],
-            ],
-            title=title,
-        )
+        keys = ("x1", "y1", "x2", "y2", "x3", "y3", "x4", "y4")
+        if not isinstance(coords, dict) or set(coords) != set(keys):
+            raise ValueError(
+                "Oriented bounding box must be a dictionary with coordinates."
+            )
+        return OBBox.from_list([coords[k] for k in keys], title=title)

datachain/model/pose.py

@@ -1,3 +1,5 @@
+from collections.abc import Sequence
+
 from pydantic import Field
 
 from datachain.lib.data_model import DataModel
@@ -19,27 +21,30 @@ class Pose(DataModel):
     y: list[int] = Field(default=[])
 
     @staticmethod
-    def from_list(points: list[list[float]]) -> "Pose":
-        assert len(points) == 2, "Pose must be a list of 2 lists: x and y coordinates."
+    def from_list(points: Sequence[Sequence[float]]) -> "Pose":
+        if not isinstance(points, (list, tuple)):
+            raise TypeError("Pose must be a list of coordinates.")
+        if len(points) != 2:
+            raise ValueError("Pose must be a list of 2 lists: x and y coordinates.")
         points_x, points_y = points
-        assert len(points_x) == len(points_y) == 17, (
-            "Pose x and y coordinates must have the same length of 17."
-        )
-        assert all(
-            isinstance(value, (int, float)) for value in [*points_x, *points_y]
-        ), "Pose coordinates must be floats or integers."
-        return Pose(
-            x=[round(coord) for coord in points_x],
-            y=[round(coord) for coord in points_y],
-        )
+        if not isinstance(points_x, (list, tuple)) or not isinstance(
+            points_y, (list, tuple)
+        ):
+            raise TypeError("Pose x and y coordinates must be lists.")
+        if len(points_x) != len(points_y) != 17:
+            raise ValueError(
+                "Pose x and y coordinates must have the same length of 17."
+            )
+        if not all(isinstance(value, (int, float)) for value in [*points_x, *points_y]):
+            raise ValueError("Pose coordinates must be floats or integers.")
+        return Pose(x=list(map(round, points_x)), y=list(map(round, points_y)))
 
     @staticmethod
-    def from_dict(points: dict[str, list[float]]) -> "Pose":
-        assert isinstance(points, dict) and set(points) == {
-            "x",
-            "y",
-        }, "Pose must be a dict with keys 'x' and 'y'."
-        return Pose.from_list([points["x"], points["y"]])
+    def from_dict(points: dict[str, Sequence[float]]) -> "Pose":
+        keys = ("x", "y")
+        if not isinstance(points, dict) or set(points) != set(keys):
+            raise ValueError("Pose must be a dictionary with coordinates.")
+        return Pose.from_list([points[k] for k in keys])
 
 
 class Pose3D(DataModel):
@@ -60,29 +65,36 @@ class Pose3D(DataModel):
     visible: list[float] = Field(default=[])
 
     @staticmethod
-    def from_list(points: list[list[float]]) -> "Pose3D":
-        assert len(points) == 3, (
-            "Pose3D must be a list of 3 lists: x, y coordinates and visible."
-        )
+    def from_list(points: Sequence[Sequence[float]]) -> "Pose3D":
+        if not isinstance(points, (list, tuple)):
+            raise TypeError("Pose3D must be a list of coordinates.")
+        if len(points) != 3:
+            raise ValueError(
+                "Pose3D must be a list of 3 lists: x, y coordinates and visible."
+            )
         points_x, points_y, points_v = points
-        assert len(points_x) == len(points_y) == len(points_v) == 17, (
-            "Pose3D x, y coordinates and visible must have the same length of 17."
-        )
-        assert all(
+        if (
+            not isinstance(points_x, (list, tuple))
+            or not isinstance(points_y, (list, tuple))
+            or not isinstance(points_v, (list, tuple))
+        ):
+            raise TypeError("Pose3D x, y and visible must be lists.")
+        if len(points_x) != len(points_y) != len(points_v) != 17:
+            raise ValueError("Pose3D x, y and visible must have the same length of 17.")
+        if not all(
             isinstance(value, (int, float))
             for value in [*points_x, *points_y, *points_v]
-        ), "Pose3D coordinates must be floats or integers."
+        ):
+            raise ValueError("Pose3D coordinates must be floats or integers.")
         return Pose3D(
-            x=[round(coord) for coord in points_x],
-            y=[round(coord) for coord in points_y],
-            visible=points_v,
+            x=list(map(round, points_x)),
+            y=list(map(round, points_y)),
+            visible=list(points_v),
         )
 
     @staticmethod
-    def from_dict(points: dict[str, list[float]]) -> "Pose3D":
-        assert isinstance(points, dict) and set(points) == {
-            "x",
-            "y",
-            "visible",
-        }, "Pose3D must be a dict with keys 'x', 'y' and 'visible'."
-        return Pose3D.from_list([points["x"], points["y"], points["visible"]])
+    def from_dict(points: dict[str, Sequence[float]]) -> "Pose3D":
+        keys = ("x", "y", "visible")
+        if not isinstance(points, dict) or set(points) != set(keys):
+            raise ValueError("Pose3D must be a dictionary with coordinates.")
+        return Pose3D.from_list([points[k] for k in keys])

datachain/model/segment.py

@@ -1,3 +1,5 @@
+from collections.abc import Sequence
+
 from pydantic import Field
 
 from datachain.lib.data_model import DataModel
@@ -21,27 +23,29 @@ class Segment(DataModel):
     y: list[int] = Field(default=[])
 
     @staticmethod
-    def from_list(points: list[list[float]], title: str = "") -> "Segment":
-        assert len(points) == 2, (
-            "Segment must be a list of 2 lists: x and y coordinates."
-        )
+    def from_list(points: Sequence[Sequence[float]], title: str = "") -> "Segment":
+        if not isinstance(points, (list, tuple)):
+            raise TypeError("Segment must be a list of coordinates.")
+        if len(points) != 2:
+            raise ValueError("Segment must be a list of 2 lists: x and y coordinates.")
         points_x, points_y = points
-        assert len(points_x) == len(points_y), (
-            "Segment x and y coordinates must have the same length."
-        )
-        assert all(
-            isinstance(value, (int, float)) for value in [*points_x, *points_y]
-        ), "Segment coordinates must be floats or integers."
+        if not isinstance(points_x, (list, tuple)) or not isinstance(
+            points_y, (list, tuple)
+        ):
+            raise TypeError("Segment x and y coordinates must be lists.")
+        if len(points_x) != len(points_y):
+            raise ValueError("Segment x and y coordinates must have the same length.")
+        if not all(isinstance(value, (int, float)) for value in [*points_x, *points_y]):
+            raise ValueError("Segment coordinates must be floats or integers.")
         return Segment(
             title=title,
-            x=[round(coord) for coord in points_x],
-            y=[round(coord) for coord in points_y],
+            x=list(map(round, points_x)),
+            y=list(map(round, points_y)),
         )
 
     @staticmethod
-    def from_dict(points: dict[str, list[float]], title: str = "") -> "Segment":
-        assert isinstance(points, dict) and set(points) == {
-            "x",
-            "y",
-        }, "Segment must be a dict with keys 'x' and 'y'."
-        return Segment.from_list([points["x"], points["y"]], title=title)
+    def from_dict(points: dict[str, Sequence[float]], title: str = "") -> "Segment":
+        keys = ("x", "y")
+        if not isinstance(points, dict) or set(points) != set(keys):
+            raise ValueError("Segment must be a dictionary with coordinates.")
+        return Segment.from_list([points[k] for k in keys], title=title)

datachain/model/ultralytics/bbox.py

@@ -23,7 +23,7 @@ class YoloBBox(DataModel):
     cls: int = Field(default=-1)
     name: str = Field(default="")
     confidence: float = Field(default=0)
-    box: BBox
+    box: BBox = Field(default=BBox())
 
     @staticmethod
     def from_result(result: "Results") -> "YoloBBox":
@@ -55,10 +55,10 @@ class YoloBBoxes(DataModel):
         box: A list of bounding boxes of the detected objects
     """
 
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[BBox]
+    cls: list[int] = Field(default=[])
+    name: list[str] = Field(default=[])
+    confidence: list[float] = Field(default=[])
+    box: list[BBox] = Field(default=[])
 
     @staticmethod
     def from_results(results: list["Results"]) -> "YoloBBoxes":
@@ -124,10 +124,10 @@ class YoloOBBoxes(DataModel):
         box: A list of oriented bounding boxes of the detected objects.
     """
 
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[OBBox]
+    cls: list[int] = Field(default=[])
+    name: list[str] = Field(default=[])
+    confidence: list[float] = Field(default=[])
+    box: list[OBBox] = Field(default=[])
 
     @staticmethod
     def from_results(results: list["Results"]) -> "YoloOBBoxes":

datachain/model/ultralytics/pose.py

@@ -47,8 +47,8 @@ class YoloPose(DataModel):
     cls: int = Field(default=-1)
     name: str = Field(default="")
     confidence: float = Field(default=0)
-    box: BBox
-    pose: Pose3D
+    box: BBox = Field(default=BBox())
+    pose: Pose3D = Field(default=Pose3D())
 
     @staticmethod
     def from_result(result: "Results") -> "YoloPose":
@@ -87,11 +87,11 @@ class YoloPoses(DataModel):
         pose: The 3D pose keypoints of the poses.
     """
 
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[BBox]
-    pose: list[Pose3D]
+    cls: list[int] = Field(default=[])
+    name: list[str] = Field(default=[])
+    confidence: list[float] = Field(default=[])
+    box: list[BBox] = Field(default=[])
+    pose: list[Pose3D] = Field(default=[])
 
     @staticmethod
     def from_results(results: list["Results"]) -> "YoloPoses":