Support a list of Trackio.Image in a trackio.Table cell (#336)

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>

Author: abidlabs

.changeset/eight-keys-help.md

@@ -0,0 +1,5 @@
+---
+"trackio": patch
+---
+
+feat:Support a list of `Trackio.Image` in a `trackio.Table` cell

examples/deploy-images-on-spaces.py

@@ -0,0 +1,23 @@
+import random
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+import trackio
+
+EPOCHS = 20
+PROJECT_ID = random.randint(100000, 999999)
+
+
+trackio.init(
+    project=f"deploy-images-on-spaces-{PROJECT_ID}",
+    space_id=f"deploy-images-on-spaces-{PROJECT_ID}",
+)
+image = trackio.Image(
+    Image.fromarray(np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8))
+)
+df = pd.DataFrame({"value": [0.1, 0.2, 0.3], "image": [[image, image], image, image]})
+table = trackio.Table(dataframe=df)
+trackio.log({"my_table": table})
+trackio.finish()

examples/deploy-on-spaces.py

@@ -84,4 +84,4 @@ def generate_accuracy_curve(epoch, max_epochs, max_acc=0.95, min_acc=0.1):
 
         time.sleep(0.2)
 
-wandb.finish()
+    wandb.finish()

tests/test_table.py

@@ -73,13 +73,63 @@ def test_table_to_display_format_with_images():
     assert row1["step"] == 1
     assert row1["value"] == 42
     assert row1["text"] == "regular text"
-    assert "![Test Caption](/gradio_api/file=" in row1["image"]
-    assert row1["image"].endswith(
-        "image.png)"
-    )  # The extra ) is due to the Markdown syntax
+    assert '<img src="/gradio_api/file=' in row1["image"]
+    assert 'image.png"' in row1["image"]
+    assert 'alt="Test Caption"' in row1["image"]
 
     row2 = processed_data[1]
     assert row2["step"] == 2
     assert row2["value"] == 84
     assert row2["text"] == "more text"
     assert row2["image"] is None
+
+
+def test_table_to_display_format_with_multiple_images():
+    table_data = [
+        {
+            "step": 1,
+            "images": [
+                {
+                    "_type": "trackio.image",
+                    "file_path": "test/path/image1.png",
+                    "caption": "First Image",
+                },
+                {
+                    "_type": "trackio.image",
+                    "file_path": "test/path/image2.png",
+                    "caption": "Second Image",
+                },
+                {
+                    "_type": "trackio.image",
+                    "file_path": "test/path/image3.png",
+                    "caption": "Third Image",
+                },
+            ],
+            "value": 42,
+        },
+        {
+            "step": 2,
+            "images": [],
+            "value": 84,
+        },
+    ]
+
+    processed_data = Table.to_display_format(table_data)
+
+    assert len(processed_data) == 2
+
+    row1 = processed_data[0]
+    assert row1["step"] == 1
+    assert row1["value"] == 42
+    assert '<img src="/gradio_api/file=' in row1["images"]
+    assert 'alt="First Image"' in row1["images"]
+    assert 'image1.png"' in row1["images"]
+    assert 'alt="Second Image"' in row1["images"]
+    assert 'image2.png"' in row1["images"]
+    assert 'alt="Third Image"' in row1["images"]
+    assert 'image3.png"' in row1["images"]
+
+    row2 = processed_data[1]
+    assert row2["step"] == 2
+    assert row2["value"] == 84
+    assert row2["images"] == []

trackio/run.py

@@ -110,22 +110,61 @@ def _init_client_background(self):
 
         self._batch_sender()
 
+    def _queue_upload(self, file_path, step: int | None):
+        """Queue a media file for upload to space."""
+        upload_entry: UploadEntry = {
+            "project": self.project,
+            "run": self.name,
+            "step": step,
+            "uploaded_file": handle_file(file_path),
+        }
+        with self._client_lock:
+            self._queued_uploads.append(upload_entry)
+
     def _process_media(self, value: TrackioMedia, step: int | None) -> dict:
         """
         Serialize media in metrics and upload to space if needed.
         """
         value._save(self.project, self.name, step)
         if self._space_id:
-            upload_entry: UploadEntry = {
-                "project": self.project,
-                "run": self.name,
-                "step": step,
-                "uploaded_file": handle_file(value._get_absolute_file_path()),
-            }
-            with self._client_lock:
-                self._queued_uploads.append(upload_entry)
+            self._queue_upload(value._get_absolute_file_path(), step)
         return value._to_dict()
 
+    def _scan_and_queue_media_uploads(self, table_dict: dict, step: int | None):
+        """
+        Scan a serialized table for media objects and queue them for upload to space.
+        """
+        if not self._space_id:
+            return
+
+        table_data = table_dict.get("_value", [])
+        for row in table_data:
+            for value in row.values():
+                if isinstance(value, dict) and value.get("_type") in [
+                    "trackio.image",
+                    "trackio.video",
+                    "trackio.audio",
+                ]:
+                    file_path = value.get("file_path")
+                    if file_path:
+                        from trackio.utils import MEDIA_DIR
+
+                        absolute_path = MEDIA_DIR / file_path
+                        self._queue_upload(absolute_path, step)
+                elif isinstance(value, list):
+                    for item in value:
+                        if isinstance(item, dict) and item.get("_type") in [
+                            "trackio.image",
+                            "trackio.video",
+                            "trackio.audio",
+                        ]:
+                            file_path = item.get("file_path")
+                            if file_path:
+                                from trackio.utils import MEDIA_DIR
+
+                                absolute_path = MEDIA_DIR / file_path
+                                self._queue_upload(absolute_path, step)
+
     def log(self, metrics: dict, step: int | None = None):
         renamed_keys = []
         new_metrics = {}
@@ -147,6 +186,7 @@ def log(self, metrics: dict, step: int | None = None):
                 metrics[key] = value._to_dict(
                     project=self.project, run=self.name, step=step
                 )
+                self._scan_and_queue_media_uploads(metrics[key], step)
             elif isinstance(value, Histogram):
                 metrics[key] = value._to_dict()
             elif isinstance(value, TrackioMedia):

trackio/table.py

@@ -13,14 +13,15 @@
 
 class Table:
     """
-    Initializes a Table object. Tables can include image columns using the Image class.
+    Initializes a Table object. Tables can be used to log tabular data including images, numbers, and text.
 
     Args:
         columns (`list[str]`, *optional*):
             Names of the columns in the table. Optional if `data` is provided. Not
             expected if `dataframe` is provided. Currently ignored.
         data (`list[list[Any]]`, *optional*):
-            2D row-oriented array of values.
+            2D row-oriented array of values. Each value can be: a number, a string (treated as Markdown and truncated if too long),
+             or a `Trackio.Image` or list of `Trackio.Image` objects.
         dataframe (`pandas.`DataFrame``, *optional*):
             DataFrame object used to create the table. When set, `data` and `columns`
             arguments are ignored.
@@ -54,10 +55,20 @@ def __init__(
             self.data = dataframe
 
     def _has_media_objects(self, dataframe: DataFrame) -> bool:
-        """Check if dataframe contains any TrackioMedia objects."""
+        """Check if dataframe contains any TrackioMedia objects or lists of TrackioMedia objects."""
         for col in dataframe.columns:
             if dataframe[col].apply(lambda x: isinstance(x, TrackioMedia)).any():
                 return True
+            if (
+                dataframe[col]
+                .apply(
+                    lambda x: isinstance(x, list)
+                    and len(x) > 0
+                    and isinstance(x[0], TrackioMedia)
+                )
+                .any()
+            ):
+                return True
         return False
 
     def _process_data(self, project: str, run: str, step: int = 0):
@@ -73,6 +84,13 @@ def _process_data(self, project: str, run: str, step: int = 0):
                 if isinstance(value, TrackioMedia):
                     value._save(project, run, step)
                     processed_df.at[idx, col] = value._to_dict()
+                if (
+                    isinstance(value, list)
+                    and len(value) > 0
+                    and isinstance(value[0], TrackioMedia)
+                ):
+                    [v._save(project, run, step) for v in value]
+                    processed_df.at[idx, col] = [v._to_dict() for v in value]
 
         return processed_df.to_dict(orient="records")
 
@@ -86,19 +104,33 @@ def to_display_format(table_data: list[dict]) -> list[dict]:
             table_data: List of dictionaries representing table rows (from stored _value)
 
         Returns:
-            Table data with images converted to markdown syntax
+            Table data with images converted to markdown syntax and long text truncated.
         """
         truncate_length = int(os.getenv("TRACKIO_TABLE_TRUNCATE_LENGTH", "250"))
+
+        def convert_image_to_markdown(image_data: dict) -> str:
+            relative_path = image_data.get("file_path", "")
+            caption = image_data.get("caption", "")
+            absolute_path = MEDIA_DIR / relative_path
+            return f'<img src="/gradio_api/file={absolute_path}" alt="{caption}" />'
+
         processed_data = []
         for row in table_data:
             processed_row = {}
             for key, value in row.items():
                 if isinstance(value, dict) and value.get("_type") == "trackio.image":
-                    relative_path = value.get("file_path", "")
-                    caption = value.get("caption", "")
-                    absolute_path = MEDIA_DIR / relative_path
+                    processed_row[key] = convert_image_to_markdown(value)
+                elif (
+                    isinstance(value, list)
+                    and len(value) > 0
+                    and isinstance(value[0], dict)
+                    and value[0].get("_type") == "trackio.image"
+                ):
+                    # This assumes that if the first item is an image, all items are images. Ok for now since we don't support mixed types in a single cell.
                     processed_row[key] = (
-                        f"![{caption}](/gradio_api/file={absolute_path})"
+                        '<div style="display: flex; gap: 10px;">'
+                        + "".join([convert_image_to_markdown(item) for item in value])
+                        + "</div>"
                     )
                 elif isinstance(value, str) and len(value) > truncate_length:
                     truncated = value[:truncate_length]