Annotation tool 4 0

docs/notebooks/annotation_tool.ipynb

 {
  "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Image annotation tool\n",
-    "This notebook shows how the annotation interface can be used to create masks for images"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -15,113 +7,19 @@
    "outputs": [],
    "source": [
     "import qim3d\n",
+    "from scipy import ndimage\n",
     "import matplotlib.pyplot as plt\n",
-    "import matplotlib as mpl\n",
-    "import numpy as np\n",
-    "%matplotlib inline"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Load 2D example image\n",
-    "img = qim3d.examples.blobs_256x256\n",
+    "qim3d.io.logger.level(\"info\")\n",
+    "\n",
+    "# Load example image\n",
+    "img = qim3d.examples.bone_128x128x128\n",
     "\n",
-    "# Display image\n",
-    "plt.imshow(img)\n",
-    "plt.show()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
     "# Start annotation tool\n",
     "interface = qim3d.gui.annotation_tool.Interface()\n",
-    "interface.max_masks = 4\n",
     "\n",
     "# We can directly pass the image we loaded to the interface\n",
-    "interface.launch(img=img)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# When 'prepare mask for download' is pressed once, the mask can be retrieved with the get_result() method\n",
-    "mask = interface.get_result()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Check the obtained mask"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "print (f\"Original image shape..: {img.shape}\")\n",
-    "print (f\"Mask image shape......: {mask.shape}\")\n",
-    "print (f\"\\nNumber of masks: {np.max(mask)}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Show the masked regions"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%matplotlib inline\n",
-    "\n",
-    "nmasks = np.max(mask)\n",
-    "fig, axs = plt.subplots(nrows=1, ncols=nmasks+2, figsize=(12,3))\n",
-    "\n",
-    "# Show original image\n",
-    "axs[0].imshow(img)\n",
-    "axs[0].set_title(\"Original\")\n",
-    "axs[0].axis('off')\n",
-    "\n",
-    "\n",
-    "# Show masks\n",
-    "cmap = mpl.colormaps[\"rainbow\"].copy()\n",
-    "cmap.set_under(color='black') # Sets the background to black\n",
-    "axs[1].imshow(mask, interpolation='none', cmap=cmap, vmin=1, vmax=nmasks+1)\n",
-    "axs[1].set_title(\"Masks\")\n",
-    "axs[1].axis('off')\n",
-    "\n",
-    "# Show masked regions\n",
-    "for idx in np.arange(2, nmasks+2):\n",
-    "    mask_id = idx-1\n",
-    "    submask = mask.copy()\n",
-    "    submask[submask != mask_id] = 0\n",
-    "    \n",
-    "    masked_img = img.copy()\n",
-    "    masked_img[submask==0] = 0\n",
-    "    axs[idx].imshow(masked_img)\n",
-    "    axs[idx].set_title(f\"Mask {mask_id}\")\n",
-    "\n",
-    "    axs[idx].axis('off')\n",
-    "\n",
-    "plt.show()"
+    "#interface.launch()\n",
+    "interface.launch(ndimage.zoom(img[0], 3, order=0))"
    ]
   }
  ],
@@ -141,9 +39,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

docs/utils.md

@@ -6,6 +6,11 @@ icon: fontawesome/solid/screwdriver-wrench
 
 A set of tools to ease managment of the system, with the common needs for large data in mind.
 
+::: qim3d.utils.img
+    options:
+      members:
+        - overlay_rgb_images
+
 ::: qim3d.utils.system
     options:
       members:

qim3d/css/gradio.css

 /* Override for dark mode */
-.dark{
+.dark {
   --name: default;
   --primary-50: #fff7ed;
   --primary-100: #ffedd5;
@@ -73,9 +73,9 @@
   --link-text-color-hover: var(--secondary-700);
   --link-text-color-visited: var(--secondary-500);
   --body-text-color-subdued: var(--neutral-400);
-  --shadow-drop: rgba(0,0,0,0.05) 0px 1px 2px 0px;
+  --shadow-drop: rgba(0, 0, 0, 0.05) 0px 1px 2px 0px;
   --shadow-drop-lg: 0 1px 3px 0 rgb(0 0 0 / 0.1), 0 1px 2px -1px rgb(0 0 0 / 0.1);
-  --shadow-inset: rgba(0,0,0,0.05) 0px 2px 4px 0px inset;
+  --shadow-inset: rgba(0, 0, 0, 0.05) 0px 2px 4px 0px inset;
   --shadow-spread: 3px;
   --block-background-fill: var(--background-fill-primary);
   --block-border-color: var(--border-color-primary);
@@ -203,6 +203,7 @@
   --button-small-text-weight: 400;
   --button-transition: none;
 }
+
 /* Up here is the override for dark mode */
 /* Now comes our custom CSS */
 
@@ -228,9 +229,10 @@ h2 {
   white-space: pre !important;
 }
 
-.no-border{
+.no-border {
   border-width: 0px !important;
 }
+
 /* Hides Gradio footer */
 footer {
   visibility: hidden;
@@ -341,9 +343,9 @@ input[type="range"]::-webkit-slider-thumb {
 }
 
 /* Buttons */
-.btn-html{
+.btn-html {
   position: sticky;
-  display:flex;
+  display: flex;
   justify-content: center;
   background: white !important;
   border-color: #6c757d !important;
@@ -372,7 +374,7 @@ input[type="range"]::-webkit-slider-thumb {
   text-align: left !important;
 }
 
-.btn-spinner::after{
+.btn-spinner::after {
   content: "";
   position: absolute;
   width: 16px;
@@ -547,7 +549,7 @@ div.svelte-1frtwj3 {
 
 .w-64 {
   width: 64px !important;
-  min-width: 64px !important; 
+  min-width: 64px !important;
 
 }
 
@@ -558,7 +560,7 @@ div.svelte-1frtwj3 {
 
 .w-36 {
   width: 36px !important;
-  min-width: 36px !important; 
+  min-width: 36px !important;
 
 }
 
@@ -569,7 +571,7 @@ div.svelte-1frtwj3 {
 
 .w-32 {
   width: 32px !important;
-  min-width: 32px !important; 
+  min-width: 32px !important;
 
 }
 
@@ -579,7 +581,7 @@ div.svelte-1frtwj3 {
 
 .w-16 {
   width: 16px !important;
-  min-width: 16px !important; 
+  min-width: 16px !important;
 }
 
 .h-0 {
@@ -590,33 +592,124 @@ div.svelte-1frtwj3 {
   width: 0px !important;
 }
 
-.options{
+.options {
   border: 1px solid #bababa !important;
   box-shadow: 0px 0px 8px rgba(0, 0, 0, 0.15) !important;
 
 }
 
-.mt-32{
+.mt-32 {
   margin-top: 32px;
 }
 
-.mt-48{
+.mt-48 {
   margin-top: 48px;
 }
 
-.mt-64{
+.mt-64 {
   margin-top: 64px;
 }
 
-.mt-128{
+.mt-128 {
   margin-top: 128px;
 }
 
-.hide-overflow{
+.hide-overflow {
   overflow: hidden !important;
 }
 
 
-.wrap.default{
+.wrap.default {
   visibility: hidden !important;
-}
+}
\ No newline at end of file
+
+
+/* Annotation tool CSS */
+
+.annotation-tool .layer-wrap {
+  display: none !important;
+}
+
+.annotation-tool button[title="Image button"] {
+  display: none !important;
+}
+
+.annotation-tool:not(.no-img) button[title="Upload button"] {
+  display: none !important;
+}
+
+.annotation-tool button[title="Clear canvas"] {
+  display: none !important;
+}
+
+.annotation-tool button[title="Color button"] {
+  color: var(--block-label-text-color) !important;
+  margin-inline-end: 32px !important;
+}
+
+.annotation-tool button[title="Color button"]::after {
+  content: "Select mask";
+}
+
+.annotation-tool button[title="Size button"] {
+  color: var(--block-label-text-color) !important;
+}
+
+
+.annotation-tool button[title="Size button"]::after {
+  content: "Brush size";
+}
+
+
+.annotation-tool .row-wrap {
+  justify-content: flex-start !important;
+}
+
+.annotation-tool .controls-wrap .small {
+  width: 20px !important;
+  height: 20px !important;
+}
+
+.annotation-tool .controls-wrap {
+  top: 0px !important;
+  left: 0px !important;
+}
+
+.annotation-tool .controls-wrap .padded {
+  padding: 8px !important;
+  border: 0px !important;
+  background-color: rgba(255, 255, 255, 0.8);
+}
+
+.annotation-tool .stage-wrap {
+  margin-top: 0px !important;
+  margin-bottom: 16px !important;
+}
+
+.annotation-tool .svelte-b3dw9m {
+  display: flex !important;
+  visibility: visible !important;
+  opacity: 1 !important;
+}
+
+.annotation-tool .bottom {
+  bottom: 8px !important;
+  left: 0px !important;
+  position: absolute !important;
+}
+
+.annotation-tool .download {
+  min-width: 4rem !important;
+  width: 10%;
+  white-space: nowrap;
+  text-align: right;
+}
+
+.annotation-tool .stage-wrap canvas {
+  border: 0px !important;
+  max-width: 512px !important;
+  max-height: 512px !important;
+  height: 512px !important;
+  width: fit-content !important;
+
+}

qim3d/gui/annotation_tool.py

 import tifffile
 import os
+import time
+import getpass
 import numpy as np
 import gradio as gr
-from qim3d.io import load  # load or DataLoader?
-from qim3d.utils import internal_tools
+import qim3d.utils
+from qim3d.io import load, save
+from qim3d.io.logger import log
+
+
+class Session:
+    def __init__(self):
+        self.n_masks = 3
+        self.img_editor = None
+        self.masks_rgb = None
+        self.mask_names = {0: "red", 1: "green", 2: "blue"}
+        self.temp_files = []
+        self.gradio_temp = None
+        self.username = getpass.getuser()
 
 
 class Interface:
     def __init__(self):
         self.verbose = False
         self.title = "Annotation tool"
-        # self.plot_height = 768
-        self.height = 1024
-        # self.width = 960
-        self.max_masks = 3
-        self.mask_opacity = 0.5
-        self.cmy_hex = ["#00ffff", "#ff00ff", "#ffff00"]  # Colors for max_masks>3?
+        self.height = 768
+        self.interface = None
+        self.username = getpass.getuser()
 
         # CSS path
         current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -23,7 +34,9 @@ class Interface:
 
     def launch(self, img=None, **kwargs):
         # Create gradio interfaces
-        self.interface = self.create_interface(img=img)
+
+        self.interface = self.create_interface(img)
+        self.gradio_temp = self.interface.GRADIO_CACHE
 
         # Set gradio verbose level
         if self.verbose:
@@ -35,7 +48,6 @@ class Interface:
             quiet=quiet,
             height=self.height,
             # width=self.width,
-            show_tips=False,
             **kwargs,
         )
 
@@ -43,322 +55,203 @@ class Interface:
 
     def get_result(self):
         # Get the temporary files from gradio
-        temp_sets = self.interface.temp_file_sets
-        for temp_set in temp_sets:
-            if "mask" in str(temp_set):
+        base = os.path.join(self.gradio_temp, "qim3d", self.username)
+        temp_path_list = []
+        for filename in os.listdir(base):
+            if "mask" in str(filename):
                 # Get the list of the temporary files
-                temp_path_list = list(temp_set)
-
-        # Files are not in creation order,
-        # so we need to get find the latest
-        creation_time_list = []
-        for path in temp_path_list:
-            creation_time_list.append(os.path.getctime(path))
+                temp_path_list.append(os.path.join(base, filename))
 
-        # Get index for the latest file
-        file_idx = np.argmax(creation_time_list)
+        # Make dictionary of maks
+        masks = {}
+        for temp_file in temp_path_list:
+            mask_file = os.path.basename(temp_file)
+            mask_name = os.path.splitext(mask_file)[0]
+            masks[mask_name] = load(temp_file)
 
-        # Load the temporary file
-        mask = load(temp_path_list[file_idx])
+        return masks
 
-        return mask
+    def set_visible(self):
+        return gr.update(visible=True)
 
     def create_interface(self, img=None):
+
+        if img is not None:
+            custom_css = "annotation-tool"
+        else:
+            custom_css = "annotation-tool no-img"
+
         with gr.Blocks(css=self.css_path) as gradio_interface:
-            masks_state = gr.State(value={})
-            counts = gr.Number(value=1, visible=False)
 
+            brush = gr.Brush(
+                colors=[
+                    "rgb(255,50,100)",
+                    "rgb(50,250,100)",
+                    "rgb(50,100,255)",
+                ],
+                color_mode="fixed",
+                default_size=10,
+            )
             with gr.Row():
-                with gr.Column(scale=1, min_width=320):
-                    upload_img_btn = gr.UploadButton(
-                        label="Upload image",
-                        file_types=["image"],
-                        interactive=True if img is None else False,
-                    )
-                    clear_img_btn = gr.Button(
-                        value="Clear image", interactive=False if img is None else True
+
+                with gr.Column(scale=6):
+                    img_editor = gr.ImageEditor(
+                        value=img,
+                        type="numpy",
+                        image_mode="RGB",
+                        brush=brush,
+                        sources="upload",
+                        interactive=True,
+                        show_download_button=True,
+                        container=False,
+                        transforms=[""],
+                        elem_classes=custom_css,
                     )
 
+                with gr.Column(scale=1, min_width=256):
+
                     with gr.Row():
-                        with gr.Column(scale=2, min_width=32):
-                            selected_mask = gr.Radio(
-                                choices=["Mask 1"],
-                                value="Mask 1",
-                                label="Choose which mask to draw",
-                                scale=1,
-                            )
-                        with gr.Column(scale=1, min_width=64):
-                            add_mask_btn = gr.Button(
-                                value="Add mask",
-                                scale=2,
-                            )
+                        btn_update = gr.Button(
+                            value="Update", elem_classes="btn btn-html btn-run"
+                        )
+
                     with gr.Row():
-                        prep_dl_btn = gr.Button(
-                            value="Prepare mask for download",
-                            visible=False if img is None else True,
+                        overlay_img = gr.Image(
+                            show_download_button=False, show_label=False, visible=False
                         )
                     with gr.Row():
-                        save_output = gr.File(
-                            show_label=True,
-                            label="Output file",
+                        masks_download = gr.File(
+                            label="Download masks",
                             visible=False,
+                            elem_classes=custom_css,
                         )
 
-                with gr.Column(scale=4):
-                    with gr.Row():
-                        input_img = gr.Image(
-                            label="Input",
-                            tool="sketch",
-                            value=img,
-                            height=600,
-                            width=600,
-                            brush_color="#00ffff",
-                            mask_opacity=self.mask_opacity,
-                            interactive=False if img is None else True,
-                        )
+            temp_path = gr.Textbox(value=gradio_interface.GRADIO_CACHE, visible=False)
+            session = gr.State([])
+            inputs = [img_editor]
+            operations = Operations()
+            # fmt: off
+            btn_update.click(
+                fn=operations.start_session, inputs=[img_editor,temp_path] , outputs=session).then(
+                fn=operations.preview, inputs=session, outputs=overlay_img).then(
+                fn=self.set_visible, inputs=None, outputs=overlay_img).then(
+                fn=operations.separate_masks, inputs=session, outputs=[session, masks_download]).then(
+                fn=self.set_visible, inputs=None, outputs=masks_download)
+
+            # fmt: on
+        return gradio_interface
 
-                    output_masks = []
-                    for mask_idx in range(self.max_masks):
-                        with gr.Row():  # make a new row for every mask
-                            output_mask = gr.Image(
-                                label=f"Mask {mask_idx+1}",
-                                visible=True if mask_idx == 0 else False,
-                                image_mode="L",
-                                height=600,
-                                width=600,
-                                interactive=False
-                                if img is None
-                                else True,  # If statement added bc of bug after Gradio 3.44.x
-                                show_download_button=False,
-                            )
-                            output_masks.append(output_mask)
-
-            # Operations
-            operations = Operations(max_masks=self.max_masks, cmy_hex=self.cmy_hex)
-
-            # Update component configuration when image is uploaded
-            upload_img_btn.upload(
-                fn=operations.upload_img_update,
-                inputs=upload_img_btn,
-                outputs=[input_img, clear_img_btn, upload_img_btn, prep_dl_btn]
-                + output_masks,
-            )
 
-            # Add mask below when 'add mask' button is clicked
-            add_mask_btn.click(
-                fn=operations.increment_mask,
-                inputs=counts,
-                outputs=[counts, selected_mask] + output_masks,
-            )
+class Operations:
 
-            # Draw mask when input image is edited
-            input_img.edit(
-                fn=operations.update_masks,
-                inputs=[input_img, selected_mask, masks_state, upload_img_btn],
-                outputs=output_masks,
-            )
+    def start_session(self, *args):
+        session = Session()
+        session.img_editor = args[0]
+        session.gradio_temp = args[1]
 
-            # Update brush color according to radio setting
-            selected_mask.change(
-                fn=operations.update_brush_color,
-                inputs=selected_mask,
-                outputs=input_img,
-            )
+        # Clean temp files
+        base = os.path.join(session.gradio_temp, "qim3d", session.username)
 
-            # Make file download visible
-            prep_dl_btn.click(
-                fn=operations.save_mask,
-                inputs=output_masks,
-                outputs=[save_output, save_output],
-            ).success(
-                fn=lambda: os.remove('mask.tif')
-                ) # Remove mask file from working directory immediately after sending it to /tmp/gradio
-
-            # Update 'Add mask' button interactivit according to the current count
-            counts.change(
-                fn=operations.set_add_mask_btn_interactivity,
-                inputs=counts,
-                outputs=add_mask_btn,
-            )
+        try:
+            files = os.listdir(base)
+            for filename in files:
+                # Check if "mask" is in the filename
+                if "mask" in filename:
+                    file_path = os.path.join(base, filename)
+                    os.remove(file_path)
 
-            # Reset component configuration when image is cleared
-            clear_img_btn.click(
-                fn=operations.clear_img_update,
-                inputs=None,
-                outputs=[
-                    selected_mask,
-                    prep_dl_btn,
-                    save_output,
-                    counts,
-                    input_img,
-                    upload_img_btn,
-                    clear_img_btn,
-                ]
-                + output_masks,
-            )
+        except FileNotFoundError:
+            files = None
 
-        return gradio_interface
+        return session
 
+    def overlay_images(self, background, masks, alpha=0.5):
+        """Overlay multiple RGB masks onto an RGB background image using alpha blending.
 
-class Operations:
-    def __init__(self, max_masks, cmy_hex):
-        self.max_masks = max_masks
-        self.cmy_hex = cmy_hex
+        Args:
+            background (numpy.ndarray): The background RGB image with shape (height, width, 3).
+            masks (numpy.ndarray): The RGB mask images with shape (num_masks, height, width, 3).
+            alpha (float, optional): The alpha value for blending. Defaults to 0.5.
 
-    def update_masks(self, input_img, selected_mask, masks_state, file):
-        # Binarize mask (it is not per default due to anti-aliasing)
-        input_mask = input_img["mask"]
-        input_mask[input_mask > 0] = 255
+        Returns:
+            numpy.ndarray: The composite image with overlaid masks.
 
-        try:
-            file_name = file.name
-        except AttributeError:
-            file_name = "nb_img"
-
-        # Add new file to state dictionary when this function sees it first time
-        if file_name not in masks_state.keys():
-            masks_state[file_name] = [[] for _ in range(self.max_masks)]
-
-        # Get index of currently selected and non-selected masks
-        sel_mask_idx = int(selected_mask[-1]) - 1
-        nonsel_mask_idxs = [
-            mask_idx
-            for mask_idx in list(range(self.max_masks))
-            if mask_idx != sel_mask_idx
-        ]
-
-        # Add background to state first time function is invoked in current session
-        if len(masks_state[file_name][0]) == 0:
-            for i in range(len(masks_state[file_name])):
-                masks_state[file_name][i].append(input_mask)
-
-        # Check for discrepancy between what is drawn and what is shown as output masks
-        masks_state_combined = 0
-        for i in range(len(masks_state[file_name])):
-            masks_state_combined += masks_state[file_name][i][-1]
-        discrepancy = masks_state_combined != input_mask
-        if np.any(discrepancy):  # Correct discrepancy in output masks
-            for i in range(self.max_masks):
-                masks_state[file_name][i][-1][discrepancy] = 0
-
-        # Add most recent change in input to currently selected mask
-        mask2append = input_mask
-        for mask_idx in nonsel_mask_idxs:
-            mask2append -= masks_state[file_name][mask_idx][-1]
-        masks_state[file_name][sel_mask_idx].append(mask2append)
-
-        return [masks_state[file_name][i][-1] for i in range(self.max_masks)]
-
-    def save_mask(self, *masks):
-        # Go from multi-channel to single-channel mask
-        stacked_masks = np.stack(masks, axis=-1)
-        final_mask = np.zeros_like(masks[0])
-        final_mask[np.where(stacked_masks == 255)[:2]] = (
-            np.where(stacked_masks == 255)[-1] + 1
-        )
+        Raises:
+            ValueError: If input images have different shapes.
 
-        # Save output image in a temp space (and to current directory which is a bug)
-        filename = "mask.tif"
-        tifffile.imwrite(filename, final_mask)
+        Note:
+            - The function performs alpha blending to overlay the masks onto the background.
+            - It ensures that the background and masks have the same shape before blending.
+            - It calculates the maximum projection of the masks and blends them onto the background.
+            - Brightness outside the masks is adjusted to maintain consistency with the background.
+        """
 
-        save_output_update = gr.File(visible=True)
+        # Igonore alpha in case its there
+        background = background[..., :3]
+        masks = masks[..., :3]
 
-        return save_output_update, filename
+        # Ensure both images have the same shape
+        if background.shape != masks.shape:
+            raise ValueError("Input images must have the same shape")
 
-    def increment_mask(self, counts):
-        # increment count by 1
-        counts += 1
-        counts = int(counts)
+        # Perform alpha blending
+        masks_max_projection = np.amax(masks, axis=2)
+        masks_max_projection = np.stack((masks_max_projection,) * 3, axis=-1)
 
-        counts_update = gr.Number(value=counts)
-        selected_mask_update = gr.Radio(
-            value=f"Mask {counts}", choices=[f"Mask {i+1}" for i in range(counts)]
-        )
-        output_masks_update = [gr.Image(visible=True)] * counts + [
-            gr.Image(visible=False)
-        ] * (self.max_masks - counts)
+        # Normalize if we have something
+        if np.max(masks_max_projection) > 0:
+            masks_max_projection = masks_max_projection / np.max(masks_max_projection)
 
-        return [counts_update, selected_mask_update] + output_masks_update
+        composite = background * (1 - alpha) + masks * alpha
+        composite = np.clip(composite, 0, 255).astype("uint8")
 
-    def update_brush_color(self, selected_mask):
-        sel_mask_idx = int(selected_mask[-1]) - 1
-        if sel_mask_idx < len(self.cmy_hex):
-            input_img_update = gr.Image(brush_color=self.cmy_hex[sel_mask_idx])
-        else:
-            input_img_update = gr.Image(brush_color="#000000")  # Return black brush
+        # Adjust brightness outside masks
+        composite = composite + (background * (1 - alpha)) * (1 - masks_max_projection)
 
-        return input_img_update
+        return composite.astype("uint8")
 
-    def set_add_mask_btn_interactivity(self, counts):
-        add_mask_btn_update = (
-            gr.Button(interactive=True)
-            if counts < self.max_masks
-            else gr.Button(interactive=False)
-        )
-        return add_mask_btn_update
-
-    def clear_img_update(self):
-        selected_mask_update = gr.Radio(
-            choices=["Mask 1"], value="Mask 1"
-        )  # Reset radio component to only show 'Mask 1'
-        prep_dl_btn_update = gr.Button(
-            visible=False
-        )  # Make 'Prepare mask for download' button invisible
-        save_output_update = gr.File(visible=False)  # Make File save box invisible
-        counts_update = gr.Number(value=1)  # Reset invisible counter to 1
-        input_img_update = gr.Image(
-            value=None, interactive=False
-        )  # Set input image component to non-interactive (so a new image cannot be uploaded directly in the component)
-        upload_img_btn_update = gr.Button(
-            interactive=True
-        )  # Make 'Upload image' button interactive
-        clear_img_btn_update = gr.Button(
-            interactive=False
-        )  # Make 'Clear image' button non-interactive
-        output_masks_update = [
-            gr.Image(value=None, visible=True if i == 0 else False, interactive=False)
-            for i in range(self.max_masks)
-        ]  # Remove drawn masks and set as invisible except mask 1. 'interactive=False' added bc of bug after Gradio 3.44.x
-
-        return [
-            selected_mask_update,
-            prep_dl_btn_update,
-            save_output_update,
-            counts_update,
-            input_img_update,
-            upload_img_btn_update,
-            clear_img_btn_update,
-        ] + output_masks_update
-
-    def upload_img_update(self, file):
-        input_img_update = gr.Image(
-            value=load(file.name), interactive=True
-        )  # Upload image from button to Image components
-        clear_img_btn_update = gr.Button(
-            interactive=True
-        )  # Make 'Clear image' button interactive
-        upload_img_btn_update = gr.Button(
-            interactive=False
-        )  # Make 'Upload image' button non-interactive
-        prep_dl_btn_update = gr.Button(
-            visible=True
-        )  # Make 'Prepare mask for download' button visible
-        output_masks_update = [
-            gr.Image(interactive=True)
-        ] * self.max_masks  # This line is added bc of bug in Gradio after 3.44.x
-
-        return [
-            input_img_update,
-            clear_img_btn_update,
-            upload_img_btn_update,
-            prep_dl_btn_update,
-        ] + output_masks_update
-
-
-def run_interface(host = "0.0.0.0"):
+    def preview(self, session):
+        background = session.img_editor["background"]
+        masks = session.img_editor["layers"][0]
+        overlay_image = qim3d.utils.img.overlay_rgb_images(background, masks)
+
+        return overlay_image
+
+    def separate_masks(self, session):
+
+        masks_rgb = session.img_editor["layers"][0]
+        mask_threshold = 200  # This value is based
+
+        mask_list = []
+        files_list = []
+
+        # Go through each channel
+        for idx in np.arange(session.n_masks):
+
+            mask_grayscale = masks_rgb[:, :, idx]
+            mask = mask_grayscale > mask_threshold
+
+            # Save only if we have a mask
+            if np.sum(mask) > 0:
+                mask_list.append(mask)
+                filename = f"mask_{session.mask_names[idx]}.tif"
+                base = os.path.join(session.gradio_temp, "qim3d", session.username)
+                if not os.path.exists(base):
+                    os.makedirs(base)
+                filepath = os.path.join(base, filename)
+                files_list.append(filepath)
+
+                save(filepath, mask, replace=True)
+                session.temp_files.append(filepath)
+
+        return session, files_list
+
+
+def run_interface(host="0.0.0.0"):
     gradio_interface = Interface().create_interface()
-    internal_tools.run_gradio_app(gradio_interface,host)
+    qim3d.utils.internal_tools.run_gradio_app(gradio_interface, host)
+
 
 if __name__ == "__main__":
     # Creates interface
-    run_interface()
+    run_interface()
\ No newline at end of file

qim3d/utils/__init__.py

@@ -4,4 +4,5 @@ from .augmentations import Augmentation
 from .data import Dataset, prepare_datasets, prepare_dataloaders
 #from .doi import get_bibtex, get_reference
 from . import doi
-from .system import Memory
+from .system import Memory
\ No newline at end of file
+from .img import overlay_rgb_images
\ No newline at end of file

qim3d/utils/img.py

+import numpy as np
+
+def overlay_rgb_images(background, foreground, alpha=0.5):
+    """Overlay multiple RGB foreground onto an RGB background image using alpha blending.
+
+    Args:
+        background (numpy.ndarray): The background RGB image.
+        foreground (numpy.ndarray): The foreground RGB image (usually masks).
+        alpha (float, optional): The alpha value for blending. Defaults to 0.5.
+
+    Returns:
+        numpy.ndarray: The composite RGB image with overlaid foreground.
+
+    Raises:
+        ValueError: If input images have different shapes.
+
+    Note:
+        - The function performs alpha blending to overlay the foreground onto the background.
+        - It ensures that the background and foreground have the same shape before blending.
+        - It calculates the maximum projection of the foreground and blends them onto the background.
+        - Brightness outside the foreground is adjusted to maintain consistency with the background.
+    """
+
+    # Igonore alpha in case its there
+    background = background[..., :3]
+    foreground = foreground[..., :3]
+
+    # Ensure both images have the same shape
+    if background.shape != foreground.shape:
+        raise ValueError("Input images must have the same shape")
+
+    # Perform alpha blending
+    foreground_max_projection = np.amax(foreground, axis=2)
+    foreground_max_projection = np.stack((foreground_max_projection,) * 3, axis=-1)
+
+    # Normalize if we have something
+    if np.max(foreground_max_projection) > 0:
+        foreground_max_projection = foreground_max_projection / np.max(foreground_max_projection)
+
+    composite = background * (1 - alpha) + foreground * alpha
+    composite = np.clip(composite, 0, 255).astype("uint8")
+
+    # Adjust brightness outside foreground
+    composite = composite + (background * (1 - alpha)) * (1 - foreground_max_projection)
+
+    return composite.astype("uint8")