[ENHANCEMENT] Make exporting records with image consistent (#5454)

# Description
<!-- Please include a summary of the changes and the related issue.
Please also include relevant motivation and context. List any
dependencies that are required for this change. -->

This PR implements a casting step from data uri to pil in the export
method for hf_datasets.

**Type of change**
<!-- Please delete options that are not relevant. Remember to title the
PR according to the type of change -->

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)
- Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- Refactor (change restructuring the codebase without changing
functionality)
- Improvement (change adding some improvement to an existing
functionality)
- Documentation update

**How Has This Been Tested**
<!-- Please add some reference about how your feature has been tested.
-->

**Checklist**
<!-- Please go over the list and make sure you've taken everything into
account -->

- I added relevant documentation
- I followed the style guidelines of this project
- I did a self-review of my code
- I made corresponding changes to the documentation
- I confirm My changes generate no new warnings
- I have added tests that prove my fix is effective or that my feature
works
- I have added relevant notes to the CHANGELOG.md file (See
https://keepachangelog.com/)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 2 people

argilla/src/argilla/records/_dataset_records.py

@@ -121,7 +121,7 @@ def to_json(self, path: Union[Path, str]) -> Path:
         return JsonIO.to_json(records=list(self), path=path)
 
     def to_datasets(self) -> "HFDataset":
-        return HFDatasetsIO.to_datasets(records=list(self))
+        return HFDatasetsIO.to_datasets(records=list(self), dataset=self.__dataset)
 
 
 class DatasetRecords(Iterable[Record], LoggingMixin):

argilla/src/argilla/records/_io/_datasets.py

@@ -19,10 +19,11 @@
 from datasets import IterableDataset, Image
 
 from argilla.records._io._generic import GenericIO
-from argilla._helpers._media import pil_to_data_uri
+from argilla._helpers._media import pil_to_data_uri, uncast_image
 
 if TYPE_CHECKING:
     from argilla.records import Record
+    from argilla.datasets import Dataset
 
 
 class HFDatasetsIO:
@@ -39,16 +40,19 @@ def _is_hf_dataset(dataset: Any) -> bool:
         return isinstance(dataset, HFDataset)
 
     @staticmethod
-    def to_datasets(records: List["Record"]) -> HFDataset:
+    def to_datasets(records: List["Record"], dataset: "Dataset") -> HFDataset:
         """
         Export the records to a Hugging Face dataset.
 
         Returns:
             The dataset containing the records.
         """
         record_dicts = GenericIO.to_dict(records, flatten=True)
-        dataset = HFDataset.from_dict(record_dicts)
-        return dataset
+        hf_dataset = HFDataset.from_dict(record_dicts)
+        image_fields = HFDatasetsIO._get_image_fields(schema=dataset.schema)
+        if image_fields:
+            hf_dataset = HFDatasetsIO._cast_uris_as_images(hf_dataset=hf_dataset, columns=image_fields)
+        return hf_dataset
 
     @staticmethod
     def _record_dicts_from_datasets(dataset: HFDataset) -> List[Dict[str, Union[str, float, int, list]]]:
@@ -72,6 +76,18 @@ def _record_dicts_from_datasets(dataset: HFDataset) -> List[Dict[str, Union[str,
             record_dicts.append(example)
         return record_dicts
 
+    @staticmethod
+    def _get_image_fields(schema: Dict) -> List[str]:
+        """Get the names of the Argilla fields that contain image data.
+
+        Parameters:
+            dataset (Dataset): The dataset to check.
+
+        Returns:
+            List[str]: The names of the Argilla fields that contain image data.
+        """
+        return [field_name for field_name, field in schema.items() if field.type == "image"]
+
     @staticmethod
     def _get_image_features(dataset: "HFDataset") -> List[str]:
         """Check if the Hugging Face dataset contains image features.
@@ -114,3 +130,32 @@ def batch_fn(batch):
             hf_dataset = hf_dataset.rename_column(original_column_name=unique_identifier, new_column_name=column)
 
         return hf_dataset
+
+    @staticmethod
+    def _cast_uris_as_images(hf_dataset: "HFDataset", columns: List[str]) -> "HFDataset":
+        """Cast the image features in the Hugging Face dataset as PIL images.
+
+        Parameters:
+            hf_dataset (HFDataset): The Hugging Face dataset to cast.
+            columns (List[str]): The names of the columns containing the image features.
+
+        Returns:
+            HFDataset: The Hugging Face dataset with image features cast as PIL images.
+        """
+        unique_identifier = uuid4().hex
+
+        def batch_fn(batch):
+            images = [uncast_image(sample) for sample in batch]
+            return {unique_identifier: images}
+
+        for column in columns:
+            hf_dataset = hf_dataset.map(
+                function=batch_fn,
+                with_indices=False,
+                batched=True,
+                input_columns=[column],
+                remove_columns=[column],
+            )
+            hf_dataset = hf_dataset.rename_column(original_column_name=unique_identifier, new_column_name=column)
+
+        return hf_dataset

argilla/tests/integration/test_export_dataset.py

@@ -333,6 +333,7 @@ def test_import_dataset_from_hub_using_wrong_settings(
         with_records_export: bool,
     ):
         repo_id = f"argilla-internal-testing/test_import_dataset_from_hub_using_wrong_settings_with_records_{with_records_export}"
+        mock_unique_name = f"test_import_dataset_from_hub_using_wrong_settings_{uuid.uuid4()}"
         dataset.records.log(records=mock_data)
 
         dataset.to_hub(repo_id=repo_id, with_records=with_records_export, token=token)
@@ -346,6 +347,8 @@ def test_import_dataset_from_hub_using_wrong_settings(
         )
         if with_records_export:
             with pytest.raises(SettingsError):
-                rg.Dataset.from_hub(repo_id=repo_id, client=client, token=token, settings=settings)
+                rg.Dataset.from_hub(
+                    repo_id=repo_id, client=client, token=token, settings=settings, name=mock_unique_name
+                )
         else:
-            rg.Dataset.from_hub(repo_id=repo_id, client=client, token=token, settings=settings)
+            rg.Dataset.from_hub(repo_id=repo_id, client=client, token=token, settings=settings, name=mock_unique_name)

argilla/tests/integration/test_export_records.py

@@ -20,6 +20,7 @@
 from tempfile import TemporaryDirectory
 
 import pytest
+from PIL import Image
 from datasets import Dataset as HFDataset
 
 import argilla as rg
@@ -331,3 +332,47 @@ def test_import_records_from_hf_dataset(dataset: rg.Dataset) -> None:
         assert record.fields["text"] == mock_data[i]["text"]
         assert record.suggestions["label"].value == mock_data[i]["label"]
         assert record.id == str(mock_data[i]["id"])
+
+
+def test_export_records_with_images_to_hf_datasets(client):
+    mock_dataset_name = "".join(random.choices(ascii_lowercase, k=16))
+    settings = rg.Settings(
+        fields=[
+            rg.ImageField(name="image"),
+        ],
+        questions=[
+            rg.TextQuestion(name="label", use_markdown=False),
+        ],
+    )
+    dataset = rg.Dataset(
+        name=mock_dataset_name,
+        settings=settings,
+        client=client,
+    )
+    dataset.create()
+    mock_data = [
+        {
+            "image": Image.new("RGB", (100, 100)),
+            "label": "positive",
+            "id": uuid.uuid4(),
+        },
+        {
+            "image": Image.new("RGB", (100, 100)),
+            "label": "negative",
+            "id": uuid.uuid4(),
+        },
+        {
+            "image": Image.new("RGB", (100, 100)),
+            "label": "positive",
+            "id": uuid.uuid4(),
+        },
+    ]
+    dataset.records.log(records=mock_data)
+    hf_dataset = dataset.records.to_datasets()
+
+    assert isinstance(hf_dataset, HFDataset)
+    assert hf_dataset.num_rows == len(mock_data)
+    assert "image" in hf_dataset.column_names
+    assert "label.suggestion" in hf_dataset.column_names
+    for i, image in enumerate(hf_dataset["image"]):
+        assert isinstance(image, Image.Image)

argilla/tests/unit/test_io/test_hf_datasets.py

@@ -22,6 +22,17 @@
 
 class TestHFDatasetsIO:
     def test_to_datasets_with_partial_values_in_records(self):
+        mock_dataset = rg.Dataset(
+            name="test",
+            settings=rg.Settings(
+                fields=[
+                    rg.TextField(name="field"),
+                ],
+                questions=[
+                    rg.TextQuestion(name="question"),
+                ],
+            ),
+        )
         records = [
             rg.Record(fields={"field": "The field"}, metadata={"a": "a"}),
             rg.Record(fields={"field": "Other field", "other": "Field"}, metadata={"b": "b"}),
@@ -44,7 +55,7 @@ def test_to_datasets_with_partial_values_in_records(self):
             ),
         ]
 
-        ds = HFDatasetsIO.to_datasets(records)
+        ds = HFDatasetsIO.to_datasets(records, dataset=mock_dataset)
         assert ds.features == {
             "status": Value(dtype="string", id=None),
             "_server_id": Value(dtype="null", id=None),