Merge pull request #202 from mlexchange/integrate-tiled-viewer

Replace data section dropdown with `tiled_viewer`

Author: Wiebke

callbacks/control_bar.py

@@ -3,6 +3,7 @@
 import random
 import time
 import uuid
+from urllib.parse import urlparse
 
 import dash_mantine_components as dmc
 import plotly.express as px
@@ -40,6 +41,36 @@
     open(EXPORT_FILE_PATH, "w").close()
 
 
+@callback(
+    Output("image-uri", "value"),
+    Input("tiled-image-selector", "selectedLinks"),
+    prevent_initial_call=True,
+)
+def update_selected_image_uri(selected_links):
+    print(f"DEBUG - Selected image links: {selected_links}")  # Debug print
+
+    if selected_links:
+        # Extract the 'self' key from the dictionary
+        if isinstance(selected_links, dict) and "self" in selected_links:
+            # Extract the full URI from the selected links
+            full_uri = selected_links.get("self", "")
+
+            # Parse the URI and extract the path
+            parsed_uri = urlparse(full_uri)
+            extracted_path = parsed_uri.path
+
+            # Remove the common prefix from the path
+            base_data_path = urlparse(tiled_datasets.data_tiled_uri).path
+            if extracted_path.startswith(base_data_path):
+                extracted_path = extracted_path[len(base_data_path) :]
+
+            # Clean up the extracted path and return
+            return extracted_path.strip("/")
+        # Fall back to string representation if we can't extract 'self'
+        return str(selected_links)
+    return ""
+
+
 @callback(
     Output("current-class-selection", "data", allow_duplicate=True),
     Output("notifications-container", "children", allow_duplicate=True),
@@ -746,10 +777,10 @@ def export_annotation(n_clicks, all_annotations, global_store):
     Input("save-annotations", "n_clicks"),
     State("annotation-store", "data"),
     State({"type": "annotation-class-store", "index": ALL}, "data"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     prevent_initial_call=True,
 )
-def save_data(n_clicks, global_store, all_annotations, image_src):
+def save_data(n_clicks, global_store, all_annotations, image_uri):
     """This callback is responsible for saving the annotation data to the store"""
     if not n_clicks:
         raise PreventUpdate
@@ -758,7 +789,7 @@ def save_data(n_clicks, global_store, all_annotations, image_src):
         # TODO: save store to the server file-user system, this will be changed to DB later
         export_data = {
             "user": USER_NAME,
-            "source": image_src,
+            "source": image_uri,
             "time": time.strftime("%Y-%m-%d-%H:%M:%S"),
             "data": json.dumps(all_annotations),
         }
@@ -787,10 +818,10 @@ def toggle_save_load_modal(n_clicks, opened):
 @callback(
     Output("load-annotations-server-container", "children"),
     Input("open-data-management-modal-button", "n_clicks"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     prevent_initial_call=True,
 )
-def populate_load_annotations_dropdown_menu_options(modal_opened, image_src):
+def populate_load_annotations_dropdown_menu_options(modal_opened, image_uri):
     """
     This callback populates dropdown window with all saved annotation options for the given project name.
     It then creates buttons with info about the save, which when clicked, loads the data from the server.
@@ -799,7 +830,7 @@ def populate_load_annotations_dropdown_menu_options(modal_opened, image_src):
         raise PreventUpdate
 
     data = tiled_masks.DEV_load_exported_json_data(
-        EXPORT_FILE_PATH, USER_NAME, image_src
+        EXPORT_FILE_PATH, USER_NAME, image_uri
     )
     if not data:
         return "No annotations found for the selected data source."
@@ -827,11 +858,11 @@ def populate_load_annotations_dropdown_menu_options(modal_opened, image_src):
     Output("annotation-class-container", "children", allow_duplicate=True),
     Output("data-management-modal", "opened", allow_duplicate=True),
     Input({"type": "load-server-annotations", "index": ALL}, "n_clicks"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     State("image-selection-slider", "value"),
     prevent_initial_call=True,
 )
-def load_and_apply_selected_annotations(selected_annotation, image_src, img_idx):
+def load_and_apply_selected_annotations(selected_annotation, image_uri, img_idx):
     """
     This callback is responsible for loading and applying the selected annotations when user selects them from the modal.
     """
@@ -845,7 +876,7 @@ def load_and_apply_selected_annotations(selected_annotation, image_src, img_idx)
 
     # TODO : when quering from the server, load (data) for user, source, time
     data = tiled_masks.DEV_load_exported_json_data(
-        EXPORT_FILE_PATH, USER_NAME, image_src
+        EXPORT_FILE_PATH, USER_NAME, image_uri
     )
     data = tiled_masks.DEV_filter_json_data_by_timestamp(
         data, str(selected_annotation_timestamp)
@@ -875,16 +906,6 @@ def open_controls_drawer(n_clicks, is_opened):
     return no_update, no_update
 
 
-@callback(Output("project-name-src", "data"), Input("refresh-tiled", "n_clicks"))
-def refresh_data_client(refresh_tiled):
-    if refresh_tiled:
-        tiled_datasets.refresh_data_client()
-    data_options = [
-        item for item in tiled_datasets.get_data_project_names() if "seg" not in item
-    ]
-    return data_options
-
-
 @callback(
     Output("show-result-overlay-toggle", "checked"),
     Output("show-result-overlay-toggle", "disabled"),

callbacks/image_viewer.py

@@ -65,7 +65,7 @@ def hide_show_segmentation_overlay(toggle_seg_result, opacity):
     Input("show-result-overlay-toggle", "checked"),
     Input("annotated-slices-selector", "value"),
     State({"type": "annotation-class-store", "index": ALL}, "data"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     State("annotation-store", "data"),
     State("image-metadata", "data"),
     State("screen-size", "data"),
@@ -81,7 +81,7 @@ def render_image(
     toggle_seg_result,
     slice_selection,
     all_annotation_class_store,
-    project_name,
+    image_uri,
     annotation_store,
     image_metadata,
     screen_size,
@@ -111,7 +111,7 @@ def render_image(
 
     if image_idx:
         image_idx -= 1  # slider starts at 1, so subtract 1 to get the correct index
-        tf = tiled_datasets.get_data_sequence_by_name(project_name)[image_idx]
+        tf = tiled_datasets.get_data_sequence_by_trimmed_uri(image_uri)[image_idx]
         # Auto-scale data
         low = np.percentile(tf.ravel(), 1)
         high = np.percentile(tf.ravel(), 99)
@@ -135,15 +135,15 @@ def render_image(
                     if image_idx in annotation_indices:
                         # Will not return an error since we already checked if image_idx is in the list
                         mapped_index = annotation_indices.index(image_idx)
-                        result = tiled_results.get_data_by_trimmed_uri(
+                        result = tiled_results.get_data_slice_by_trimmed_uri(
                             seg_result["seg_result_trimmed_uri"], slice=mapped_index
                         )
                     else:
                         result = None
                 # if mask_idx is not given in the results,
                 # then the result stems from inference on the full data set
                 else:
-                    result = tiled_results.get_data_by_trimmed_uri(
+                    result = tiled_results.get_data_slice_by_trimmed_uri(
                         seg_result["seg_result_trimmed_uri"], slice=image_idx
                     )
             else:
@@ -219,8 +219,8 @@ def render_image(
 
     # No update is needed for the 'children' of the control components
     # since we just want to trigger the loading overlay with this callback
-    if project_name != image_metadata["name"] or image_metadata["name"] is None:
-        curr_image_metadata = {"size": tf.shape, "name": project_name}
+    if image_uri != image_metadata["name"] or image_metadata["name"] is None:
+        curr_image_metadata = {"size": tf.shape, "name": image_uri}
     else:
         curr_image_metadata = dash.no_update
     return (
@@ -393,7 +393,7 @@ def update_viewfinder(relayout_data, annotation_store):
     }
     """,
     Output("image-viewer-loading", "className", allow_duplicate=True),
-    Input("project-name-src", "value"),
+    Input("image-uri", "value"),
     prevent_initial_call=True,
 )
 clientside_callback(
@@ -478,7 +478,7 @@ def locally_store_annotations(
     Output(
         {"type": "annotation-class-store", "index": ALL}, "data", allow_duplicate=True
     ),
-    Input("project-name-src", "value"),
+    Input("image-uri", "value"),
     State({"type": "annotation-class-store", "index": ALL}, "data"),
     prevent_initial_call=True,
 )
@@ -498,17 +498,17 @@ def clear_annotations_on_dataset_change(change_project, all_annotation_class_sto
     Output("image-selection-slider", "value"),
     Output("image-selection-slider", "disabled"),
     Output("annotation-store", "data"),
-    Input("project-name-src", "value"),
+    Input("image-uri", "value"),
     State("annotation-store", "data"),
 )
-def update_slider_values(project_name, annotation_store):
+def update_slider_values(image_uri, annotation_store):
     """
     When the data source is loaded, this callback will set the slider values and chain call
     "update_selection_and_image" callback which will update image and slider selection component.
     """
-    # Retrieve data shape if project_name is valid and points to a 3d array
+    # Retrieve data shape if image_uri is valid and points to a 3d array
     data_shape = (
-        tiled_datasets.get_data_shape_by_name(project_name) if project_name else None
+        tiled_datasets.get_data_shape_by_trimmed_uri(image_uri) if image_uri else None
     )
     disable_slider = data_shape is None
     if not disable_slider:

callbacks/segmentation.py

@@ -174,7 +174,7 @@
     Input("run-train", "n_clicks"),
     State("annotation-store", "data"),
     State({"type": "annotation-class-store", "index": ALL}, "data"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     State("model-parameters", "children"),
     State("model-list", "value"),
     State("job-name", "value"),
@@ -184,7 +184,7 @@ def run_train(
     n_clicks,
     global_store,
     all_annotations,
-    project_name,
+    image_uri,
     model_parameter_container,
     model_name,
     job_name,
@@ -208,7 +208,7 @@ def run_train(
             )
             return notification, no_update
         mask_uri, num_classes, mask_error_message = tiled_masks.save_annotations_data(
-            global_store, all_annotations, project_name
+            global_store, all_annotations, image_uri
         )
         model_parameters["num_classes"] = num_classes
         model_parameters["network"] = model_name
@@ -227,7 +227,7 @@ def run_train(
             "%Y/%m/%d %H:%M:%S"
         )
         flow_run_name = f"{job_name} {current_time}"
-        data_uri = tiled_datasets.get_data_uri_by_name(project_name)
+        data_uri = tiled_datasets.get_data_uri_by_trimmed_uri(image_uri)
         io_parameters = assemble_io_parameters_from_uris(data_uri, mask_uri)
         io_parameters["uid_retrieve"] = ""
         io_parameters["models_dir"] = RESULTS_DIR
@@ -263,7 +263,7 @@ def run_train(
                     FLOW_NAME,
                     parameters=TRAIN_PARAMS_EXAMPLE,
                     flow_run_name=flow_run_name,
-                    tags=PREFECT_TAGS + ["train", project_name],
+                    tags=PREFECT_TAGS + ["train", image_uri],
                 )
                 job_message = f"Job has been succesfully submitted with uid: {job_uid} and mask uri: {mask_uri}"
                 notification_color = "indigo"
@@ -287,7 +287,7 @@ def run_train(
     Input("run-inference", "n_clicks"),
     State("train-job-selector", "value"),
     State({"type": "annotation-class-store", "index": ALL}, "data"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     State("model-parameters", "children"),
     State("model-list", "value"),
     prevent_initial_call=True,
@@ -296,7 +296,7 @@ def run_inference(
     n_clicks,
     train_job_id,
     all_annotations,
-    project_name,
+    image_uri,
     model_parameter_container,
     model_name,
 ):
@@ -324,7 +324,7 @@ def run_inference(
         model_parameters["network"] = model_name
 
         # Set io_parameters for inference, there will be no mask
-        data_uri = tiled_datasets.get_data_uri_by_name(project_name)
+        data_uri = tiled_datasets.get_data_uri_by_trimmed_uri(image_uri)
         io_parameters = assemble_io_parameters_from_uris(data_uri, "")
         io_parameters["uid_retrieve"] = ""
         io_parameters["models_dir"] = RESULTS_DIR
@@ -380,7 +380,7 @@ def run_inference(
                             FLOW_NAME,
                             parameters=INFERENCE_PARAMS_EXAMPLE,
                             flow_run_name=flow_run_name,
-                            tags=PREFECT_TAGS + ["inference", project_name],
+                            tags=PREFECT_TAGS + ["inference", image_uri],
                         )
                         job_message = (
                             f"Job has been succesfully submitted with uid: {job_uid}"
@@ -445,10 +445,10 @@ def check_train_job(n_intervals):
     Output("infra-state", "data", allow_duplicate=True),
     Input("model-check", "n_intervals"),
     Input("train-job-selector", "value"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     prevent_initial_call=True,
 )
-def check_inference_job(n_intervals, train_job_id, project_name):
+def check_inference_job(n_intervals, train_job_id, image_uri):
     """
     This callback populates the inference job selector dropdown with job names and ids from Prefect.
     The list of jobs is filtered by the selected train job in the train job selector dropdown.
@@ -472,7 +472,7 @@ def check_inference_job(n_intervals, train_job_id, project_name):
             if job_name is not None:
                 data = query_flow_runs(
                     flow_run_name=job_name,
-                    tags=PREFECT_TAGS + ["inference", project_name],
+                    tags=PREFECT_TAGS + ["inference", image_uri],
                 )
             infra_state = no_update
 
@@ -489,7 +489,7 @@ def check_inference_job(n_intervals, train_job_id, project_name):
 
 def populate_segmentation_results(
     job_id,
-    project_name,
+    image_uri,
     job_type="training",
 ):
     """
@@ -498,7 +498,7 @@ def populate_segmentation_results(
     """
     # Nothing has been selected is job_id is None
     if job_id is not None:
-        data_uri = tiled_datasets.get_data_uri_by_name(project_name)
+        data_uri = tiled_datasets.get_data_uri_by_trimmed_uri(image_uri)
         # Only returns the name if the job finished successfully
         job_name = get_flow_run_name(job_id)
         if job_name is not None:
@@ -515,7 +515,7 @@ def populate_segmentation_results(
                 # First refresh the data client,
                 # the root container may not yet have existed on startup
                 tiled_results.refresh_data_client()
-                result_container = tiled_results.get_data_by_trimmed_uri(
+                result_container = tiled_results.get_data_slice_by_trimmed_uri(
                     expected_result_uri
                 )
             except Exception:
@@ -551,16 +551,16 @@ def populate_segmentation_results(
     Output("seg-results-train-store", "data"),
     Output("dvc-training-stats-link", "href"),
     Input("train-job-selector", "value"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     prevent_initial_call=True,
 )
-def populate_segmentation_results_train(train_job_id, project_name):
+def populate_segmentation_results_train(train_job_id, image_uri):
     """
     This callback populates the segmentation results store based on the uids
     if the training job and the inference job.
     """
     notification, result_store, segment_job_id = populate_segmentation_results(
-        train_job_id, project_name, "training"
+        train_job_id, image_uri, "training"
     )
     if segment_job_id is not None:
         results_link = f"/results/{segment_job_id}/dvc_metrics/report.html"
@@ -574,16 +574,16 @@ def populate_segmentation_results_train(train_job_id, project_name):
     Output("notifications-container", "children", allow_duplicate=True),
     Output("seg-results-inference-store", "data"),
     Input("inference-job-selector", "value"),
-    State("project-name-src", "value"),
+    State("image-uri", "value"),
     prevent_initial_call=True,
 )
-def populate_segmentation_results_inference(inference_job_id, project_name):
+def populate_segmentation_results_inference(inference_job_id, image_uri):
     """
     This callback populates the segmentation results store based on the uids
     if the training job and the inference job.
     """
     notification, result_store, _ = populate_segmentation_results(
-        inference_job_id, project_name, "inference"
+        inference_job_id, image_uri, "inference"
     )
     return (
         notification,