Add image logging support (#142)

Author: Saba9

examples/fake-training-images.py

@@ -0,0 +1,80 @@
+import math
+import random
+import time
+
+from PIL import Image as PILImage
+from PIL import ImageDraw
+
+import trackio as wandb
+
+EPOCHS = 12
+W = H = 128
+
+
+def lissajous(t, w=W, h=H):
+    x = (w // 2) + int((w // 3) * math.sin(2.0 * t))
+    y = (h // 2) + int((h // 3) * math.sin(3.0 * t + math.pi / 4))
+    return x, y
+
+
+def render_overlay(target_xy, pred_xy):
+    img = PILImage.new("RGB", (W, H), "black")
+    draw = ImageDraw.Draw(img)
+
+    tx, ty = target_xy
+    draw.ellipse((tx - 5, ty - 5, tx + 5, ty + 5), fill=(0, 255, 0))
+
+    px, py = pred_xy
+    draw.ellipse((px - 5, py - 5, px + 5, py + 5), fill=(255, 80, 80))
+
+    draw.line([(tx, ty), (px, py)], fill=(255, 255, 0), width=1)
+    return img
+
+
+def main():
+    project_id = random.randint(10000, 99999)
+    project_name = f"image-logging-demo-{project_id}"
+
+    for run_index in range(1, 3):
+        run_name = f"image-run-{run_index}"
+        wandb.init(project=project_name, name=run_name)
+
+        pred_x, pred_y = random.randint(0, W - 1), random.randint(0, H - 1)
+
+        for epoch in range(EPOCHS):
+            t = epoch / 3.0
+            target = lissajous(t)
+
+            lr = 0.35
+            noise_scale = max(0.0, 5.0 * (1.0 - epoch / (EPOCHS - 1)))
+            pred_x += lr * (target[0] - pred_x) + random.uniform(
+                -noise_scale, noise_scale
+            )
+            pred_y += lr * (target[1] - pred_y) + random.uniform(
+                -noise_scale, noise_scale
+            )
+            pred = (int(round(pred_x)), int(round(pred_y)))
+
+            loss = math.dist(target, pred)
+
+            overlay = render_overlay(target, pred)
+            wandb.log(
+                {
+                    "loss": loss,
+                    "target_x": target[0],
+                    "target_y": target[1],
+                    "pred_x": pred[0],
+                    "pred_y": pred[1],
+                    "overlay": wandb.Image(
+                        overlay, caption=f"step={epoch}, loss={loss:.2f}"
+                    ),
+                },
+                step=epoch,
+            )
+            time.sleep(0.2)
+
+        wandb.finish()
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -14,7 +14,9 @@ requires-python = ">=3.10"
 dependencies = [
     "pandas<3.0.0",
     "huggingface-hub<1.0.0",
-    "gradio>=5.35.0,<6.0.0"
+    "gradio>=5.43.1,<6.0.0",
+    "numpy<3.0.0",
+    "pillow<12.0.0",
 ]
 classifiers = [
     "Programming Language :: Python :: 3",

tests/conftest.py

@@ -5,8 +5,9 @@
 
 
 @pytest.fixture
-def temp_db(monkeypatch):
-    """Fixture that creates a temporary directory for database storage and patches the TRACKIO_DIR."""
+def temp_dir(monkeypatch):
+    """Fixture that creates a temporary TRACKIO_DIR."""
     with tempfile.TemporaryDirectory() as tmpdir:
-        monkeypatch.setattr("trackio.sqlite_storage.TRACKIO_DIR", Path(tmpdir))
+        for name in ("trackio.sqlite_storage", "trackio.media", "trackio.file_storage"):
+            monkeypatch.setattr(f"{name}.TRACKIO_DIR", Path(tmpdir))
         yield tmpdir

tests/e2e/test_basic_logging.py

@@ -2,13 +2,13 @@
 from trackio.sqlite_storage import SQLiteStorage
 
 
-def test_basic_logging(temp_db):
+def test_basic_logging(temp_dir):
     trackio.init(project="test_project", name="test_run")
     trackio.log(metrics={"loss": 0.1})
     trackio.log(metrics={"loss": 0.2, "acc": 0.9})
     trackio.finish()
 
-    results = SQLiteStorage.get_metrics(project="test_project", run="test_run")
+    results = SQLiteStorage.get_logs(project="test_project", run="test_run")
     assert len(results) == 2
     assert results[0]["loss"] == 0.1
     assert results[0]["step"] == 0
@@ -20,13 +20,13 @@ def test_basic_logging(temp_db):
     assert "timestamp" in results[1]
 
 
-def test_basic_logging_with_step(temp_db):
+def test_basic_logging_with_step(temp_dir):
     trackio.init(project="test_project", name="test_run")
     trackio.log(metrics={"loss": 0.1}, step=0)
     trackio.log(metrics={"loss": 0.2, "acc": 0.9}, step=2)
     trackio.finish()
 
-    results = SQLiteStorage.get_metrics(project="test_project", run="test_run")
+    results = SQLiteStorage.get_logs(project="test_project", run="test_run")
     assert len(results) == 2
     assert results[0]["loss"] == 0.1
     assert results[0]["step"] == 0

tests/e2e/test_bulk_logging.py

@@ -35,7 +35,7 @@ def test_rapid_bulk_logging():
     time.sleep(0.6)  # Wait for the client to send the logs
 
     # Verify run1 metrics
-    metrics_run1 = SQLiteStorage.get_metrics(project_name, run1_name)
+    metrics_run1 = SQLiteStorage.get_logs(project_name, run1_name)
     assert len(metrics_run1) == num_logs_run1, (
         f"Expected {num_logs_run1} logs for run1, but found {len(metrics_run1)}"
     )
@@ -52,7 +52,7 @@ def test_rapid_bulk_logging():
         )
 
     # Verify run2 metrics
-    metrics_run2 = SQLiteStorage.get_metrics(project_name, run2_name)
+    metrics_run2 = SQLiteStorage.get_logs(project_name, run2_name)
     assert len(metrics_run2) == num_logs_run2, (
         f"Expected {num_logs_run2} logs for run2, but found {len(metrics_run2)}"
     )

tests/e2e/test_image_logging.py

@@ -0,0 +1,35 @@
+import numpy as np
+
+import trackio
+from trackio.sqlite_storage import SQLiteStorage
+
+PROJECT_NAME = "test_project"
+
+
+def test_image_logging(temp_dir):
+    trackio.init(project=PROJECT_NAME, name="test_run")
+
+    image1 = trackio.Image(
+        np.random.randint(255, size=(100, 100, 3), dtype=np.uint8),
+        caption="test_caption1",
+    )
+    image2 = trackio.Image(
+        np.random.randint(255, size=(100, 100, 3), dtype=np.uint8),
+        caption="test_caption2",
+    )
+    trackio.log(metrics={"loss": 0.1, "img1": image1})
+    trackio.log(metrics={"loss": 0.2, "img1": image1, "img2": image2})
+    trackio.finish()
+
+    metrics = SQLiteStorage.get_logs(project=PROJECT_NAME, run="test_run")
+
+    assert len(metrics) == 2
+
+    assert metrics[0]["loss"] == 0.1
+    assert metrics[0]["step"] == 0
+    assert metrics[0]["img1"] == image1._to_dict()
+
+    assert metrics[1]["loss"] == 0.2
+    assert metrics[1]["step"] == 1
+    assert metrics[1]["img1"] == image1._to_dict()
+    assert metrics[1]["img2"] == image2._to_dict()

tests/e2e/test_import_from_csv.py

@@ -4,14 +4,14 @@
 from trackio.sqlite_storage import SQLiteStorage
 
 
-def test_import_from_csv(temp_db):
+def test_import_from_csv(temp_dir):
     trackio.import_csv(
         csv_path=str(Path(__file__).parent / "logs.csv"),
         project="test_project",
         name="test_run",
     )
 
-    results = SQLiteStorage.get_metrics(project="test_project", run="test_run")
+    results = SQLiteStorage.get_logs(project="test_project", run="test_run")
     assert len(results) == 4
     assert results[0]["train/loss"] == 12.2
     assert results[0]["train/acc"] == 82.2

tests/e2e/test_import_from_tf.py

@@ -6,7 +6,7 @@
 from trackio.sqlite_storage import SQLiteStorage
 
 
-def test_import_from_tf_events(temp_db):
+def test_import_from_tf_events(temp_dir):
     test_run_dir = "tf_test_run"
 
     def create_tfevents_tensorboardx(log_dir: Path):
@@ -28,7 +28,7 @@ def create_tfevents_tensorboardx(log_dir: Path):
         name="test_run",
     )
 
-    results = SQLiteStorage.get_metrics(project="test_tf_project", run="test_run_main")
+    results = SQLiteStorage.get_logs(project="test_tf_project", run="test_run_main")
     # There should be 5 steps × 2 metrics = 10 entries
     assert len(results) == 10
 

tests/test_media.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import numpy as np
+
+from trackio.media import TrackioImage
+
+PROJECT_NAME = "test_project"
+
+
+def test_image_save(temp_dir):
+    image = TrackioImage(np.random.randint(255, size=(100, 100, 3), dtype=np.uint8))
+    image._save(PROJECT_NAME, "test_run", 0, "PNG")
+
+    assert image._file_format == "PNG"
+
+    expected_rel_dir = Path("media") / PROJECT_NAME / "test_run" / "0"
+    assert str(image._get_relative_file_path()).startswith(str(expected_rel_dir))
+    assert str(image._get_absolute_file_path()).endswith(".png")
+    assert image._get_absolute_file_path().is_file()
+
+
+def test_image_serialization(temp_dir):
+    image = TrackioImage(
+        np.random.randint(255, size=(100, 100, 3), dtype=np.uint8),
+        caption="test_caption",
+    )
+    image._save(PROJECT_NAME, "test_run", 0, "PNG")
+    value = image._to_dict()
+
+    assert value is not None
+    assert value.get("_type") == TrackioImage.TYPE
+    assert value.get("file_path") == str(image._get_relative_file_path())
+    assert value.get("file_format") == "PNG"
+    assert value.get("caption") == "test_caption"
+
+
+def test_image_deserialization(temp_dir):
+    image = TrackioImage(
+        np.random.randint(255, size=(100, 100, 3), dtype=np.uint8),
+        caption="test_caption",
+    )
+    image._save(PROJECT_NAME, "test_run", 0, "PNG")
+    value = image._to_dict()
+
+    image2 = TrackioImage._from_dict(value)
+    assert image2._get_relative_file_path() == image._get_relative_file_path()
+    assert image2._get_absolute_file_path() == image._get_absolute_file_path()
+    assert image2._get_absolute_file_path().is_file()
+    assert image2._file_format == "PNG"
+    assert image2.caption == "test_caption"

tests/test_run.py

@@ -26,7 +26,7 @@ def test_run_log_calls_client():
     )
 
 
-def test_init_resume_modes(temp_db):
+def test_init_resume_modes(temp_dir):
     run = init(
         project="test-project",
         name="new-run",