add access updates for tiled (#217)

* add access updates for tiled

* Update utils/data_utils.py

Co-authored-by: Wiebke Köpp <wkoepp@lbl.gov>

* remove ggshield

* reduce size of np fixture

---------

Co-authored-by: Wiebke Köpp <wkoepp@lbl.gov>

Author: dylanmcreynolds

.env.example

@@ -52,4 +52,4 @@ MLFLOW_TRACKING_USERNAME=
 MLFLOW_TRACKING_PASSWORD=
 #algorithm registry in mlflow
 MLFLOW_TRACKING_URI_OUTSIDE=http://localhost:5000
-ALGORITHM_JSON_PATH="../assets/models.json"
+ALGORITHM_JSON_PATH="../assets/models.json"

.pre-commit-config.yaml

@@ -12,12 +12,6 @@ repos:
       - id: check-symlinks
       - id: check-yaml
       - id: debug-statements
-  - repo: https://github.com/gitguardian/ggshield
-    rev: v1.25.0
-    hooks:
-      - id: ggshield
-        language_version: python3
-        stages: [commit]
   # Using this mirror lets us use mypyc-compiled black, which is about 2x faster
   - repo: https://github.com/psf/black-pre-commit-mirror
     rev: 24.2.0

scripts/save_mlflow_algorithm.py

@@ -8,7 +8,9 @@
 # Add the project root directory to Python path to fix imports
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-from mlex_utils.mlflow_utils.mlflow_algorithm_client import MlflowAlgorithmClient
+from mlex_utils.mlflow_utils.mlflow_algorithm_client import (  # noqa: E402
+    MlflowAlgorithmClient,
+)
 
 # Load environment variables from .env file
 load_dotenv(dotenv_path="../.env")

tests/unit/test_control_bar.py

@@ -1,10 +1,23 @@
+import sys
+
 import numpy as np
 import pytest
 
 
+class DummyMlflowAlgorithmClient:
+    def __init__(self, *args, **kwargs):
+        self.modelname_list = []
+
+    def load_from_mlflow(self, algorithm_type="segmentation"):
+        return None
+
+    def __getitem__(self, key):
+        raise KeyError(key)
+
+
 @pytest.fixture
 def tiled_data_mock():
-    return {"sample_project": np.zeros((500, 500, 500)), "reconstruction": 0}
+    return {"sample_project": np.zeros((2, 500, 500)), "reconstruction": 0}
 
 
 @pytest.fixture
@@ -17,6 +30,12 @@ def test_reset_filters(mocker, nclicks, tiled_data_mock):
         "tiled.client.from_uri",
         return_value=tiled_data_mock,
     )
+    mocker.patch(
+        "mlex_utils.mlflow_utils.mlflow_algorithm_client.MlflowAlgorithmClient",
+        DummyMlflowAlgorithmClient,
+    )
+    sys.modules.pop("callbacks.control_bar", None)
+    sys.modules.pop("utils.data_utils", None)
     from callbacks.control_bar import reset_filters
 
     assert reset_filters(nclicks) == 100

tests/unit/test_data_utils.py

@@ -0,0 +1,111 @@
+import importlib
+import sys
+from types import SimpleNamespace
+
+
+class DummyMlflowAlgorithmClient:
+    def __init__(self, *args, **kwargs):
+        self.modelname_list = []
+
+    def load_from_mlflow(self, algorithm_type="segmentation"):
+        return None
+
+    def __getitem__(self, key):
+        raise KeyError(key)
+
+
+class FakeContainer:
+    def __init__(self, uri="tiled://root"):
+        self.uri = uri
+        self._children = {}
+        self.metadata = None
+
+    def keys(self):
+        return self._children.keys()
+
+    def create_container(self, key, metadata=None):
+        container = FakeContainer(uri=f"{self.uri}/{key}")
+        container.metadata = metadata
+        self._children[key] = container
+        return container
+
+    def write_array(self, key, array):
+        array_client = SimpleNamespace(uri=f"{self.uri}/{key}", array=array)
+        self._children[key] = array_client
+        return array_client
+
+    def __getitem__(self, key):
+        return self._children[key]
+
+
+def test_save_annotations_data(mocker, monkeypatch):
+    monkeypatch.setenv("DATA_TILED_URI", "http://example.com/api/v1/metadata/data")
+    monkeypatch.setenv("MASK_TILED_URI", "http://example.com/api/v1/metadata/masks")
+    monkeypatch.setenv("SEG_TILED_URI", "http://example.com/api/v1/metadata/seg")
+
+    mocker.patch("tiled.client.from_uri", return_value={})
+    mocker.patch(
+        "mlex_utils.mlflow_utils.mlflow_algorithm_client.MlflowAlgorithmClient",
+        DummyMlflowAlgorithmClient,
+    )
+
+    sys.modules.pop("utils.data_utils", None)
+    data_utils = importlib.import_module("utils.data_utils")
+
+    root_container = FakeContainer()
+    mask_handler = data_utils.TiledMaskHandler.__new__(data_utils.TiledMaskHandler)
+    mask_handler.mask_client = root_container
+
+    mocker.patch.object(
+        data_utils, "from_uri", return_value=SimpleNamespace(access_blob={})
+    )
+    mocker.patch.object(data_utils, "copy_tiled_access_info", return_value=None)
+    mocker.patch.object(
+        data_utils.tiled_datasets,
+        "get_data_uri_by_trimmed_uri",
+        return_value="http://example.com/data/project/sample",
+    )
+    mocker.patch.object(
+        data_utils.tiled_datasets,
+        "get_data_sequence_by_trimmed_uri",
+        return_value=SimpleNamespace(access_blob={}),
+    )
+
+    all_annotations = [
+        {
+            "class_id": "class-1",
+            "label": "Class 1",
+            "color": "#ffffff",
+            "annotations": {
+                "0": [
+                    {
+                        "type": "rect",
+                        "x0": 0,
+                        "y0": 0,
+                        "x1": 1,
+                        "y1": 1,
+                    }
+                ]
+            },
+        }
+    ]
+
+    uri, num_classes, message = mask_handler.save_annotations_data(
+        global_store={"image_shapes": [(4, 4)]},
+        all_annotations=all_annotations,
+        trimmed_uri="project/sample",
+    )
+
+    assert uri is not None
+    assert num_classes == 1
+    assert message == "Annotations saved successfully."
+    assert "project/sample" in uri
+
+    user_container = root_container[data_utils.USER_NAME]
+    project_container = user_container["project"]["sample"]
+    assert len(list(project_container.keys())) == 1
+
+    saved_hash = next(iter(project_container.keys()))
+    saved_container = project_container[saved_hash]
+    assert saved_container.metadata["project_name"] == "project/sample"
+    assert "mask" in saved_container.keys()

utils/data_utils.py

@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 import traceback
 from urllib.parse import urlparse, urlunparse
@@ -15,6 +16,12 @@
 
 load_dotenv()
 
+# Setup basic logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+
 DATA_TILED_URI = os.getenv("DATA_TILED_URI")
 DATA_TILED_API_KEY = os.getenv("DATA_TILED_API_KEY")
 MASK_TILED_URI = os.getenv("MASK_TILED_URI")
@@ -354,6 +361,10 @@ def save_annotations_data(self, global_store, all_annotations, trimmed_uri):
                 key=annotations_hash, metadata=metadata
             )
             mask = last_container.write_array(key="mask", array=mask)
+            image_client = tiled_datasets.get_data_sequence_by_trimmed_uri(trimmed_uri)
+            # match the new container and new array access tags to the original data_client
+            copy_tiled_access_info(image_client, last_container)
+            copy_tiled_access_info(image_client, mask)
         else:
             last_container = last_container[annotations_hash]
         return (
@@ -439,3 +450,20 @@ def assemble_io_parameters_from_uris(data_uri, mask_uri):
         "seg_tiled_uri": SEG_TILED_URI,
     }
     return io_parameters
+
+
+def copy_tiled_access_info(source_client, target_client):
+    """
+    This function copies the access information from a source Tiled URI to a target Tiled client.
+    Input:
+        source_uri: str, The URI of the source Tiled resource.
+        target_client: Tiled client object, The target Tiled client to copy access info to.
+    Output:
+        None
+    """
+    access_blob = source_client.access_blob
+    if access_blob and access_blob.get("tags") is not None:
+        target_client.replace_metadata(access_tags=access_blob["tags"])
+    logging.info(
+        f"Tiled access information copied successfully {source_client} {source_client.access_blob}."
+    )