datachain 0.6.8__py3-none-any.whl → 0.6.10__py3-none-any.whl

datachain/lib/models/ultralytics/bbox.py

@@ -0,0 +1,189 @@
+"""
+This module contains the YOLO models.
+
+YOLO stands for "You Only Look Once", a family of object detection models that
+are designed to be fast and accurate. The models are trained to detect objects
+in images by dividing the image into a grid and predicting the bounding boxes
+and class probabilities for each grid cell.
+
+More information about YOLO can be found here:
+- https://pjreddie.com/darknet/yolo/
+- https://docs.ultralytics.com/
+"""
+
+from io import BytesIO
+from typing import TYPE_CHECKING
+
+from PIL import Image
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+from datachain.lib.models.bbox import BBox, OBBox
+
+if TYPE_CHECKING:
+    from ultralytics.engine.results import Results
+    from ultralytics.models import YOLO
+
+    from datachain.lib.file import File
+
+
+class YoloBBox(DataModel):
+    """
+    A class representing a bounding box detected by a YOLO model.
+
+    Attributes:
+        cls: The class of the detected object.
+        name: The name of the detected object.
+        confidence: The confidence score of the detection.
+        box: The bounding box of the detected object
+    """
+
+    cls: int = Field(default=-1)
+    name: str = Field(default="")
+    confidence: float = Field(default=0)
+    box: BBox = Field(default=None)
+
+    @staticmethod
+    def from_file(yolo: "YOLO", file: "File") -> "YoloBBox":
+        results = yolo(Image.open(BytesIO(file.read())))
+        if len(results) == 0:
+            return YoloBBox()
+        return YoloBBox.from_result(results[0])
+
+    @staticmethod
+    def from_result(result: "Results") -> "YoloBBox":
+        summary = result.summary()
+        if not summary:
+            return YoloBBox()
+        name = summary[0].get("name", "")
+        box = (
+            BBox.from_dict(summary[0]["box"], title=name)
+            if "box" in summary[0]
+            else BBox()
+        )
+        return YoloBBox(
+            cls=summary[0]["class"],
+            name=name,
+            confidence=summary[0]["confidence"],
+            box=box,
+        )
+
+
+class YoloBBoxes(DataModel):
+    """
+    A class representing a list of bounding boxes detected by a YOLO model.
+
+    Attributes:
+        cls: A list of classes of the detected objects.
+        name: A list of names of the detected objects.
+        confidence: A list of confidence scores of the detections.
+        box: A list of bounding boxes of the detected objects
+    """
+
+    cls: list[int]
+    name: list[str]
+    confidence: list[float]
+    box: list[BBox]
+
+    @staticmethod
+    def from_file(yolo: "YOLO", file: "File") -> "YoloBBoxes":
+        results = yolo(Image.open(BytesIO(file.read())))
+        return YoloBBoxes.from_results(results)
+
+    @staticmethod
+    def from_results(results: list["Results"]) -> "YoloBBoxes":
+        cls, names, confidence, box = [], [], [], []
+        for r in results:
+            for s in r.summary():
+                name = s.get("name", "")
+                cls.append(s["class"])
+                names.append(name)
+                confidence.append(s["confidence"])
+                box.append(BBox.from_dict(s.get("box", {}), title=name))
+        return YoloBBoxes(
+            cls=cls,
+            name=names,
+            confidence=confidence,
+            box=box,
+        )
+
+
+class YoloOBBox(DataModel):
+    """
+    A class representing an oriented bounding box detected by a YOLO model.
+
+    Attributes:
+        cls: The class of the detected object.
+        name: The name of the detected object.
+        confidence: The confidence score of the detection.
+        box: The oriented bounding box of the detected object.
+    """
+
+    cls: int = Field(default=-1)
+    name: str = Field(default="")
+    confidence: float = Field(default=0)
+    box: OBBox = Field(default=None)
+
+    @staticmethod
+    def from_file(yolo: "YOLO", file: "File") -> "YoloOBBox":
+        results = yolo(Image.open(BytesIO(file.read())))
+        if len(results) == 0:
+            return YoloOBBox()
+        return YoloOBBox.from_result(results[0])
+
+    @staticmethod
+    def from_result(result: "Results") -> "YoloOBBox":
+        summary = result.summary()
+        if not summary:
+            return YoloOBBox()
+        name = summary[0].get("name", "")
+        box = (
+            OBBox.from_dict(summary[0]["box"], title=name)
+            if "box" in summary[0]
+            else OBBox()
+        )
+        return YoloOBBox(
+            cls=summary[0]["class"],
+            name=name,
+            confidence=summary[0]["confidence"],
+            box=box,
+        )
+
+
+class YoloOBBoxes(DataModel):
+    """
+    A class representing a list of oriented bounding boxes detected by a YOLO model.
+
+    Attributes:
+        cls: A list of classes of the detected objects.
+        name: A list of names of the detected objects.
+        confidence: A list of confidence scores of the detections.
+        box: A list of oriented bounding boxes of the detected objects.
+    """
+
+    cls: list[int]
+    name: list[str]
+    confidence: list[float]
+    box: list[OBBox]
+
+    @staticmethod
+    def from_file(yolo: "YOLO", file: "File") -> "YoloOBBoxes":
+        results = yolo(Image.open(BytesIO(file.read())))
+        return YoloOBBoxes.from_results(results)
+
+    @staticmethod
+    def from_results(results: list["Results"]) -> "YoloOBBoxes":
+        cls, names, confidence, box = [], [], [], []
+        for r in results:
+            for s in r.summary():
+                name = s.get("name", "")
+                cls.append(s["class"])
+                names.append(name)
+                confidence.append(s["confidence"])
+                box.append(OBBox.from_dict(s.get("box", {}), title=name))
+        return YoloOBBoxes(
+            cls=cls,
+            name=names,
+            confidence=confidence,
+            box=box,
+        )

datachain/lib/models/ultralytics/pose.py

@@ -0,0 +1,126 @@
+"""
+This module contains the YOLO models.
+
+YOLO stands for "You Only Look Once", a family of object detection models that
+are designed to be fast and accurate. The models are trained to detect objects
+in images by dividing the image into a grid and predicting the bounding boxes
+and class probabilities for each grid cell.
+
+More information about YOLO can be found here:
+- https://pjreddie.com/darknet/yolo/
+- https://docs.ultralytics.com/
+"""
+
+from typing import TYPE_CHECKING
+
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+from datachain.lib.models.bbox import BBox
+from datachain.lib.models.pose import Pose3D
+
+if TYPE_CHECKING:
+    from ultralytics.engine.results import Results
+
+
+class YoloPoseBodyPart:
+    """An enumeration of body parts for YOLO pose keypoints."""
+
+    nose = 0
+    left_eye = 1
+    right_eye = 2
+    left_ear = 3
+    right_ear = 4
+    left_shoulder = 5
+    right_shoulder = 6
+    left_elbow = 7
+    right_elbow = 8
+    left_wrist = 9
+    right_wrist = 10
+    left_hip = 11
+    right_hip = 12
+    left_knee = 13
+    right_knee = 14
+    left_ankle = 15
+    right_ankle = 16
+
+
+class YoloPose(DataModel):
+    """
+    A data model for YOLO pose keypoints.
+
+    Attributes:
+        cls: The class of the pose.
+        name: The name of the pose.
+        confidence: The confidence score of the pose.
+        box: The bounding box of the pose.
+        keypoints: The 3D pose keypoints.
+    """
+
+    cls: int = Field(default=-1)
+    name: str = Field(default="")
+    confidence: float = Field(default=0)
+    box: BBox = Field(default=None)
+    keypoints: Pose3D = Field(default=None)
+
+    @staticmethod
+    def from_result(result: "Results") -> "YoloPose":
+        summary = result.summary()
+        if not summary:
+            return YoloPose()
+        name = summary[0].get("name", "")
+        box = (
+            BBox.from_dict(summary[0]["box"], title=name)
+            if "box" in summary[0]
+            else BBox()
+        )
+        keypoints = (
+            Pose3D.from_dict(summary[0]["keypoints"])
+            if "keypoints" in summary[0]
+            else Pose3D()
+        )
+        return YoloPose(
+            cls=summary[0]["class"],
+            name=name,
+            confidence=summary[0]["confidence"],
+            box=box,
+            keypoints=keypoints,
+        )
+
+
+class YoloPoses(DataModel):
+    """
+    A data model for a list of YOLO pose keypoints.
+
+    Attributes:
+        cls: The classes of the poses.
+        name: The names of the poses.
+        confidence: The confidence scores of the poses.
+        box: The bounding boxes of the poses.
+        keypoints: The 3D pose keypoints of the poses.
+    """
+
+    cls: list[int]
+    name: list[str]
+    confidence: list[float]
+    box: list[BBox]
+    keypoints: list[Pose3D]
+
+    @staticmethod
+    def from_results(results: list["Results"]) -> "YoloPoses":
+        cls, names, confidence, box, keypoints = [], [], [], [], []
+        for r in results:
+            for s in r.summary():
+                name = s.get("name", "")
+                cls.append(s["class"])
+                names.append(name)
+                confidence.append(s["confidence"])
+                box.append(BBox.from_dict(s.get("box", {}), title=name))
+                keypoints.append(Pose3D.from_dict(s.get("keypoints", {})))
+        return YoloPoses(
+            cls=cls,
+            name=names,
+            confidence=confidence,
+            box=box,
+            keypoints=keypoints,
+        )

datachain/lib/models/ultralytics/segment.py

@@ -0,0 +1,121 @@
+"""
+This module contains the YOLO models.
+
+YOLO stands for "You Only Look Once", a family of object detection models that
+are designed to be fast and accurate. The models are trained to detect objects
+in images by dividing the image into a grid and predicting the bounding boxes
+and class probabilities for each grid cell.
+
+More information about YOLO can be found here:
+- https://pjreddie.com/darknet/yolo/
+- https://docs.ultralytics.com/
+"""
+
+from io import BytesIO
+from typing import TYPE_CHECKING
+
+from PIL import Image
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+from datachain.lib.models.bbox import BBox
+from datachain.lib.models.segment import Segments
+
+if TYPE_CHECKING:
+    from ultralytics.engine.results import Results
+    from ultralytics.models import YOLO
+
+    from datachain.lib.file import File
+
+
+class YoloSegment(DataModel):
+    """
+    A data model for a single YOLO segment.
+
+    Attributes:
+        cls (int): The class of the segment.
+        name (str): The name of the segment.
+        confidence (float): The confidence of the segment.
+        box (BBox): The bounding box of the segment.
+        segments (Segments): The segments of the segment.
+    """
+
+    cls: int = Field(default=-1)
+    name: str = Field(default="")
+    confidence: float = Field(default=0)
+    box: BBox = Field(default=None)
+    segments: Segments = Field(default=None)
+
+    @staticmethod
+    def from_file(yolo: "YOLO", file: "File") -> "YoloSegment":
+        results = yolo(Image.open(BytesIO(file.read())))
+        if len(results) == 0:
+            return YoloSegment()
+        return YoloSegment.from_result(results[0])
+
+    @staticmethod
+    def from_result(result: "Results") -> "YoloSegment":
+        summary = result.summary()
+        if not summary:
+            return YoloSegment()
+        name = summary[0].get("name", "")
+        box = (
+            BBox.from_dict(summary[0]["box"], title=name)
+            if "box" in summary[0]
+            else BBox()
+        )
+        segments = (
+            Segments.from_dict(summary[0]["segments"], title=name)
+            if "segments" in summary[0]
+            else Segments()
+        )
+        return YoloSegment(
+            cls=summary[0]["class"],
+            name=summary[0]["name"],
+            confidence=summary[0]["confidence"],
+            box=box,
+            segments=segments,
+        )
+
+
+class YoloSegments(DataModel):
+    """
+    A data model for a list of YOLO segments.
+
+    Attributes:
+        cls (list[int]): The classes of the segments.
+        name (list[str]): The names of the segments.
+        confidence (list[float]): The confidences of the segments.
+        box (list[BBox]): The bounding boxes of the segments.
+        segments (list[Segments]): The segments of the segments.
+    """
+
+    cls: list[int]
+    name: list[str]
+    confidence: list[float]
+    box: list[BBox]
+    segments: list[Segments]
+
+    @staticmethod
+    def from_file(yolo: "YOLO", file: "File") -> "YoloSegments":
+        results = yolo(Image.open(BytesIO(file.read())))
+        return YoloSegments.from_results(results)
+
+    @staticmethod
+    def from_results(results: list["Results"]) -> "YoloSegments":
+        cls, names, confidence, box, segments = [], [], [], [], []
+        for r in results:
+            for s in r.summary():
+                name = s.get("name", "")
+                cls.append(s["class"])
+                names.append(name)
+                confidence.append(s["confidence"])
+                box.append(BBox.from_dict(s.get("box", {}), title=name))
+                segments.append(Segments.from_dict(s.get("segments", {}), title=name))
+        return YoloSegments(
+            cls=cls,
+            name=names,
+            confidence=confidence,
+            box=box,
+            segments=segments,
+        )

datachain/lib/signal_schema.py

@@ -378,7 +378,7 @@ class SignalSchema:
 
     def row_to_features(
         self, row: Sequence, catalog: "Catalog", cache: bool = False
-    ) -> list[DataType]:
+    ) -> list[DataValue]:
         res = []
         pos = 0
         for fr_cls in self.values.values():

datachain/listing.py

@@ -1,6 +1,7 @@
 import glob
 import os
 from collections.abc import Iterable, Iterator
+from functools import cached_property
 from itertools import zip_longest
 from typing import TYPE_CHECKING, Optional
 
@@ -15,28 +16,34 @@ from datachain.utils import suffix_to_number
 if TYPE_CHECKING:
     from datachain.catalog.datasource import DataSource
     from datachain.client import Client
-    from datachain.data_storage import AbstractWarehouse
+    from datachain.data_storage import AbstractMetastore, AbstractWarehouse
     from datachain.dataset import DatasetRecord
 
 
 class Listing:
     def __init__(
         self,
+        metastore: "AbstractMetastore",
         warehouse: "AbstractWarehouse",
         client: "Client",
-        dataset: Optional["DatasetRecord"],
+        dataset_name: Optional["str"] = None,
+        dataset_version: Optional[int] = None,
         object_name: str = "file",
     ):
+        self.metastore = metastore
         self.warehouse = warehouse
         self.client = client
-        self.dataset = dataset  # dataset representing bucket listing
+        self.dataset_name = dataset_name  # dataset representing bucket listing
+        self.dataset_version = dataset_version  # dataset representing bucket listing
         self.object_name = object_name
 
     def clone(self) -> "Listing":
         return self.__class__(
+            self.metastore.clone(),
             self.warehouse.clone(),
             self.client,
-            self.dataset,
+            self.dataset_name,
+            self.dataset_version,
             self.object_name,
         )
 
@@ -53,12 +60,22 @@ class Listing:
     def uri(self):
         from datachain.lib.listing import listing_uri_from_name
 
-        return listing_uri_from_name(self.dataset.name)
+        assert self.dataset_name
 
-    @property
+        return listing_uri_from_name(self.dataset_name)
+
+    @cached_property
+    def dataset(self) -> "DatasetRecord":
+        assert self.dataset_name
+        return self.metastore.get_dataset(self.dataset_name)
+
+    @cached_property
     def dataset_rows(self):
+        dataset = self.dataset
         return self.warehouse.dataset_rows(
-            self.dataset, self.dataset.latest_version, object_name=self.object_name
+            dataset,
+            self.dataset_version or dataset.latest_version,
+            object_name=self.object_name,
         )
 
     def expand_path(self, path, use_glob=True) -> list[Node]:

datachain/toolkit/__init__.py

@@ -0,0 +1,3 @@
+from .split import train_test_split
+
+__all__ = ["train_test_split"]

datachain/toolkit/split.py

@@ -0,0 +1,67 @@
+from datachain import C, DataChain
+
+
+def train_test_split(dc: DataChain, weights: list[float]) -> list[DataChain]:
+    """
+    Splits a DataChain into multiple subsets based on the provided weights.
+
+    This function partitions the rows or items of a DataChain into disjoint subsets,
+    ensuring that the relative sizes of the subsets correspond to the given weights.
+    It is particularly useful for creating training, validation, and test datasets.
+
+    Args:
+        dc (DataChain):
+            The DataChain instance to split.
+        weights (list[float]):
+            A list of weights indicating the relative proportions of the splits.
+            The weights do not need to sum to 1; they will be normalized internally.
+            For example:
+            - `[0.7, 0.3]` corresponds to a 70/30 split;
+            - `[2, 1, 1]` corresponds to a 50/25/25 split.
+
+    Returns:
+        list[DataChain]:
+            A list of DataChain instances, one for each weight in the weights list.
+
+    Examples:
+        Train-test split:
+        ```python
+        from datachain import DataChain
+        from datachain.toolkit import train_test_split
+
+        # Load a DataChain from a storage source (e.g., S3 bucket)
+        dc = DataChain.from_storage("s3://bucket/dir/")
+
+        # Perform a 70/30 train-test split
+        train, test = train_test_split(dc, [0.7, 0.3])
+
+        # Save the resulting splits
+        train.save("dataset_train")
+        test.save("dataset_test")
+        ```
+
+        Train-test-validation split:
+        ```python
+        train, test, val = train_test_split(dc, [0.7, 0.2, 0.1])
+        train.save("dataset_train")
+        test.save("dataset_test")
+        val.save("dataset_val")
+        ```
+
+    Note:
+        The splits are random but deterministic, based on Dataset `sys__rand` field.
+    """
+    if len(weights) < 2:
+        raise ValueError("Weights should have at least two elements")
+    if any(weight < 0 for weight in weights):
+        raise ValueError("Weights should be non-negative")
+
+    weights_normalized = [weight / sum(weights) for weight in weights]
+
+    return [
+        dc.filter(
+            C("sys__rand") % 1000 >= round(sum(weights_normalized[:index]) * 1000),
+            C("sys__rand") % 1000 < round(sum(weights_normalized[: index + 1]) * 1000),
+        )
+        for index, _ in enumerate(weights_normalized)
+    ]