Added mask interface accessible from the UI, with it you can create slices from your masks or your images of the samples, included masking in the main processing, on the spectra creation, and on the 3d surface plot, started with interface for pre processing spectral data

app.jl

@@ -1,3 +1,5 @@
+# app.jl
+
 module App
 # ==Packages ==
 using GenieFramework # Set up Genie development environment.
@@ -6,7 +8,6 @@ using Libz
 using PlotlyBase
 using CairoMakie
 using Colors
-# using julia_mzML_imzML
 using MSI_src # Import the new MSIData library
 using Statistics
 using NaturalSort
@@ -20,9 +21,50 @@ using JSON
 using Dates
 
 # Bring MSIData into App module's scope
-using .MSI_src: MSIData, OpenMSIData, #=GetSpectrum,=# process_spectrum, IterateSpectra, ImzMLSource, _iterate_spectra_fast, MzMLSource, find_mass, ViridisPalette, get_mz_slice, get_multiple_mz_slices, quantize_intensity, save_bitmap, median_filter, save_bitmap, downsample_spectrum, TrIQ, precompute_analytics, ImportMzmlFile
+using .MSI_src: MSIData, OpenMSIData, process_spectrum, IterateSpectra, ImzMLSource, _iterate_spectra_fast, MzMLSource, find_mass, ViridisPalette, get_mz_slice, get_multiple_mz_slices, quantize_intensity, save_bitmap, median_filter, save_bitmap, downsample_spectrum, TrIQ, precompute_analytics, ImportMzmlFile, generate_colorbar_image, load_and_prepare_mask
 
-include("./julia_imzML_visual.jl")
+if !@isdefined(increment_image)
+    include("./julia_imzML_visual.jl")
+end
+
+# --- Memory Validation Logging ---
+if get(ENV, "GENIE_ENV", "dev") != "prod"
+    function get_rss_mb()
+        if !Sys.islinux()
+            return 0.0
+        end
+        try
+            pid = getpid()
+            cmd = `ps -p $pid -o rss=`
+            rss_kb_str = read(cmd, String)
+            rss_kb = parse(Int, strip(rss_kb_str))
+            return round(rss_kb / 1024, digits=2)
+        catch e
+            @warn "Could not get RSS via `ps` command. Error: $e"
+            return 0.0
+        end
+    end
+
+    function log_memory_usage(context::String, msi_data_val)
+        rss_mb = get_rss_mb()
+        
+        msi_data_size_mb = 0
+        if msi_data_val !== nothing
+            msi_data_size_mb = round(Base.summarysize(msi_data_val) / (1024^2), digits=2)
+        end
+
+        gc_time_s = round(GC.time(), digits=3)
+        
+        println("--- MEMORY LOG [$(context)] ---")
+        println("  Timestamp: $(now())")
+        println("  Process RSS: $(rss_mb) MB")
+        println("  msi_data size: $(msi_data_size_mb) MB")
+        println("  Cumulative GC time: $(gc_time_s) s")
+        println("--------------------------")
+    end
+else
+    log_memory_usage(context::String, msi_data_val) = nothing # No-op for production
+end
 
 @genietools
 
@@ -47,6 +89,7 @@ include("./julia_imzML_visual.jl")
     @in triqEnabled=false
     @in SpectraEnabled=false
     @in MFilterEnabled=false
+    @in maskEnabled=false
     # Dialogs
     @in warning_msg=false
     @in CompareDialog=false
@@ -160,6 +203,9 @@ include("./julia_imzML_visual.jl")
     @out msg_conversion = ""
     @out btnConvertDisable = true
 
+    # == Pre Processing Variables ==
+    @in pre_tab = "stabilization"
+
     # == Batch Summary Dialog ==
     @in showBatchSummary = false
     @out batch_summary = ""
@@ -271,7 +317,7 @@ include("./julia_imzML_visual.jl")
         margin=attr(l=0,r=0,t=120,b=0,pad=0)
     )
     # Dummy 2D scatter plot
-    traceSpectra=PlotlyBase.scatter(x=Vector{Float64}(), y=Vector{Float64}(),marker=attr(size=1, color="blue", opacity=0.1))
+    traceSpectra=PlotlyBase.stem(x=Vector{Float64}(), y=Vector{Float64}(),marker=attr(size=1, color="blue", opacity=0.1))
     # Create conection to frontend
     @out plotdata=[traceSpectra]
     @out plotlayout=layoutSpectra
@@ -384,6 +430,7 @@ include("./julia_imzML_visual.jl")
                     imgWidth, imgHeight = dims[1], dims[2]
 
                     msi_data = nothing # Ensure data is not held in memory
+                    log_memory_usage("Fast Load (msi_data cleared)", msi_data)
                     btnMetadataDisable = false
                     btnStartDisable = false
                     btnPlotDisable = false
@@ -443,6 +490,7 @@ include("./julia_imzML_visual.jl")
 
                 selected_folder_main = dataset_name
                 msi_data = loaded_data
+                log_memory_usage("Full Load", msi_data)
 
                 eTime = round(time() - sTime, digits=3)
                 msg = "Active file loaded in $(eTime) seconds. Dataset '$(dataset_name)' is ready for analysis."
@@ -612,7 +660,7 @@ include("./julia_imzML_visual.jl")
         overall_progress = 0.0
         progress_message = "Preparing batch process..."
 
-        # --- CAPTURE CURRENT VALUES HERE (NO []) ---
+        # --- CAPTURE CURRENT VALUES HERE ---
         current_selected_files = selected_files
         current_nmass = Nmass
         current_tol = Tol
@@ -620,6 +668,7 @@ include("./julia_imzML_visual.jl")
         current_triq_enabled = triqEnabled
         current_triq_prob = triqProb
         current_mfilter_enabled = MFilterEnabled
+        current_mask_enabled = maskEnabled
         current_registry_path = registry_path
 
         println("starting main process with $(length(current_selected_files)) files")
@@ -635,62 +684,33 @@ include("./julia_imzML_visual.jl")
                     return
                 end
 
-                println("Nmass value: '$current_nmass'")
-                println("Type of current_nmass: $(typeof(current_nmass))")
-
-                masses_str = split(current_nmass, ',', keepempty=false)
-                println("Parsed masses strings: $masses_str")
-
                 masses = Float64[]
                 try
-                    masses = [parse(Float64, strip(m)) for m in masses_str]
+                    masses = [parse(Float64, strip(m)) for m in split(current_nmass, ',', keepempty=false)]
-                    println("Parsed masses: $masses")
                 catch e
                     progress_message = "Invalid m/z value(s). Please provide a comma-separated list of numbers. Error: $e"
                     warning_msg = true
-                    println(progress_message)
                     return
                 end
 
-                println("Masses array: $masses, type: $(typeof(masses))")
+                if isempty(masses)
-
-                if !(masses isa AbstractArray) || isempty(masses)
                     progress_message = "No valid m/z values found. Please provide comma-separated positive numbers."
                     warning_msg = true
-                    println(progress_message)
                     return
                 end
 
-                # Check other parameters
-                if !(0 < current_tol <= 1)
-                    progress_message = "Tolerance must be between 0 and 1."
-                    warning_msg = true
-                    println(progress_message)
-                    return
-                end
-
-                if !(1 < current_color_level < 257)
-                    progress_message = "Color levels must be between 2 and 256."
-                    warning_msg = true
-                    println(progress_message)
-                    return
-                end
-
-                println("entering batch processing loop with masses: $masses")
-                
                 # --- 2. Batch Processing Loop ---
                 num_files = length(current_selected_files)
                 total_steps = num_files
                 current_step = 0
                 errors = Dict("load_errors" => String[], "slice_errors" => String[], "io_errors" => String[])
                 newly_created_folders = String[]
+                files_without_mask = 0
 
                 for (file_idx, file_path) in enumerate(current_selected_files)
-                    # Update progress for current file
                     progress_message = "Processing file $file_idx/$num_files: $(basename(file_path))"
                     overall_progress = current_step / total_steps
                     
-                    # Create parameters for this specific file
                     all_params = (
                         tolerance = current_tol,
                         colorL = current_color_level,
@@ -702,14 +722,12 @@ include("./julia_imzML_visual.jl")
                         nFiles = num_files
                     )
 
-                    # Process the file
+                    success, error_msg = process_file_safely(file_path, masses, all_params, progress_message, overall_progress, use_mask=current_mask_enabled)
-                    success, error_msg = process_file_safely(file_path, masses, all_params, progress_message, overall_progress)
 
                     if !success
                         push!(errors["load_errors"], error_msg)
                     else
-                        dataset_name = replace(basename(file_path), r"\.imzML$"i => "")
+                        push!(newly_created_folders, replace(basename(file_path), r"\.imzML$"i => ""))
-                        push!(newly_created_folders, dataset_name)
                     end
                     current_step += 1
                 end
@@ -717,7 +735,6 @@ include("./julia_imzML_visual.jl")
                 # --- 3. Final Report ---
                 total_time_end = round(time() - total_time_start, digits=3)
                 
-                # Update folder lists in UI
                 registry = load_registry(current_registry_path)
                 all_folders = sort(collect(keys(registry)), lt=natural)
                 img_folders = filter(folder -> get(get(registry, folder, Dict()), "is_imzML", false), all_folders)
@@ -738,8 +755,11 @@ include("./julia_imzML_visual.jl")
                     warning_msg = true
                 end
 
+                mask_summary = current_mask_enabled ? "\nFiles processed without a mask: $(files_without_mask)" : ""
+
                 batch_summary = """
                 Processed $(successful_files)/$(num_files) files successfully.
+                $(mask_summary)
 
                 Errors by category:
                 • Load failures: $(length(errors["load_errors"]))
@@ -751,6 +771,56 @@ include("./julia_imzML_visual.jl")
                 """
                 showBatchSummary = true
 
+                # Update UI to display the last generated image
+                if !isempty(newly_created_folders)
+                    timestamp = string(time_ns())
+                    folder_path = joinpath("public", selected_folder_main)
+
+                    if current_triq_enabled
+                        triq_files = filter(filename -> startswith(filename, "TrIQ_") && endswith(filename, ".bmp"), readdir(folder_path))
+                        col_triq_files = filter(filename -> startswith(filename, "colorbar_TrIQ_") && endswith(filename, ".png"), readdir(folder_path))
+
+                        if !isempty(triq_files)
+                            latest_triq = triq_files[argmax([mtime(joinpath(folder_path, f)) for f in triq_files])]
+                            current_triq = latest_triq
+                            imgIntT = "/$(selected_folder_main)/$(current_triq)?t=$(timestamp)"
+                            plotdataImgT, plotlayoutImgT, _, _ = loadImgPlot(imgIntT)
+                            text_nmass = replace(current_triq, r"TrIQ_|.bmp" => "")
+                            msgtriq = "TrIQ <i>m/z</i>: $(replace(text_nmass, "_" => "."))"
+                            
+                            if !isempty(col_triq_files)
+                                latest_col_triq = col_triq_files[argmax([mtime(joinpath(folder_path, f)) for f in col_triq_files])]
+                                current_col_triq = latest_col_triq
+                                colorbarT = "/$(selected_folder_main)/$(current_col_triq)?t=$(timestamp)"
+                            else
+                                colorbarT = ""
+                            end
+                            selectedTab = "tab1"
+                        end
+                    else # Not TrIQ enabled, display regular MSI image
+                        msi_files = filter(filename -> startswith(filename, "MSI_") && endswith(filename, ".bmp"), readdir(folder_path))
+                        col_msi_files = filter(filename -> startswith(filename, "colorbar_MSI_") && endswith(filename, ".png"), readdir(folder_path))
+
+                        if !isempty(msi_files)
+                            latest_msi = msi_files[argmax([mtime(joinpath(folder_path, f)) for f in msi_files])]
+                            current_msi = latest_msi
+                            imgInt = "/$(selected_folder_main)/$(current_msi)?t=$(timestamp)"
+                            plotdataImg, plotlayoutImg, _, _ = loadImgPlot(imgInt)
+                            text_nmass = replace(current_msi, r"MSI_|.bmp" => "")
+                            msgimg = "<i>m/z</i>: $(replace(text_nmass, "_" => "."))"
+
+                            if !isempty(col_msi_files)
+                                latest_col_msi = col_msi_files[argmax([mtime(joinpath(folder_path, f)) for f in col_msi_files])]
+                                current_col_msi = latest_col_msi
+                                colorbar = "/$(selected_folder_main)/$(current_col_msi)?t=$(timestamp)"
+                            else
+                                colorbar = ""
+                            end
+                            selectedTab = "tab0"
+                        end
+                    end
+                end
+
             catch e
                 println("Error in main process: $e")
                 msg = "Batch processing failed: $e"
@@ -766,6 +836,10 @@ include("./julia_imzML_visual.jl")
                 SpectraEnabled = true
                 overall_progress = 0.0
                 println("Done")
+                GC.gc()
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
+                end
             end
         end
     end
@@ -786,7 +860,9 @@ include("./julia_imzML_visual.jl")
             try
                 sTime = time()
                 registry = load_registry(registry_path)
-                target_path = registry[selected_folder_main]["source_path"]
+                entry = registry[selected_folder_main]
+                target_path = entry["source_path"]
+
                 if target_path == "unknown (manually added)"
                     msg = "Dataset selected contained no route."
                     warning_msg = true
@@ -797,22 +873,29 @@ include("./julia_imzML_visual.jl")
                     msg = "Reloading $(basename(target_path)) for analysis..."
                     full_route = target_path
                     msi_data = OpenMSIData(target_path)
-
+                    if haskey(get(entry, "metadata", Dict()), "global_min_mz") && entry["metadata"]["global_min_mz"] !== nothing
-                    existing_entry = get(registry, selected_folder_main, nothing)
+                        msi_data.global_min_mz = entry["metadata"]["global_min_mz"]
-                    if existing_entry !== nothing && haskey(get(existing_entry, "metadata", Dict()), "global_min_mz") && existing_entry["metadata"]["global_min_mz"] !== nothing
+                        msi_data.global_max_mz = entry["metadata"]["global_max_mz"]
-                        println("Injecting cached m/z range to skip Pass 1...")
-                        msi_data.global_min_mz = existing_entry["metadata"]["global_min_mz"]
-                        msi_data.global_max_mz = existing_entry["metadata"]["global_max_mz"]
                     else
                         precompute_analytics(msi_data)
                     end
                 end
 
-                plotdata, plotlayout, xSpectraMz, ySpectraMz = meanSpectrumPlot(msi_data, selected_folder_main)
+                local mask_path_for_plot::Union{String, Nothing} = nothing
+                if maskEnabled && get(entry, "has_mask", false)
+                    mask_path_for_plot = get(entry, "mask_path", "")
+                    if !isfile(mask_path_for_plot)
+                        @warn "Mask not found for plotting: $(mask_path_for_plot). Plotting without mask."
+                        mask_path_for_plot = nothing
+                    end
+                end
+
+                plotdata, plotlayout, xSpectraMz, ySpectraMz = meanSpectrumPlot(msi_data, selected_folder_main, mask_path=mask_path_for_plot)
                 selectedTab = "tab2"
                 fTime = time()
                 eTime = round(fTime - sTime, digits=3)
                 msg = "Plot loaded in $(eTime) seconds"
+                log_memory_usage("Mean Plot Generated", msi_data)
             catch e
                 msg = "Could not generate mean spectrum plot: $e"
                 warning_msg = true
@@ -822,6 +905,10 @@ include("./julia_imzML_visual.jl")
                 btnPlotDisable = false
                 btnSpectraDisable = false
                 btnStartDisable = false
+                GC.gc()
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
+                end
             end
         end
     end
@@ -842,7 +929,9 @@ include("./julia_imzML_visual.jl")
             try
                 sTime = time()
                 registry = load_registry(registry_path)
-                target_path = registry[selected_folder_main]["source_path"]
+                entry = registry[selected_folder_main]
+                target_path = entry["source_path"]
+
                 if target_path == "unknown (manually added)"
                     msg = "Dataset selected contained no route."
                     warning_msg = true
@@ -853,22 +942,29 @@ include("./julia_imzML_visual.jl")
                     msg = "Reloading $(basename(target_path)) for analysis..."
                     full_route = target_path
                     msi_data = OpenMSIData(target_path)
-
+                    if haskey(get(entry, "metadata", Dict()), "global_min_mz") && entry["metadata"]["global_min_mz"] !== nothing
-                    existing_entry = get(registry, selected_folder_main, nothing)
+                        msi_data.global_min_mz = entry["metadata"]["global_min_mz"]
-                    if existing_entry !== nothing && haskey(get(existing_entry, "metadata", Dict()), "global_min_mz") && existing_entry["metadata"]["global_min_mz"] !== nothing
+                        msi_data.global_max_mz = entry["metadata"]["global_max_mz"]
-                        println("Injecting cached m/z range to skip Pass 1...")
-                        msi_data.global_min_mz = existing_entry["metadata"]["global_min_mz"]
-                        msi_data.global_max_mz = existing_entry["metadata"]["global_max_mz"]
                     else
                         precompute_analytics(msi_data)
                     end
                 end
 
-                plotdata, plotlayout, xSpectraMz, ySpectraMz = sumSpectrumPlot(msi_data, selected_folder_main)
+                local mask_path_for_plot::Union{String, Nothing} = nothing
+                if maskEnabled && get(entry, "has_mask", false)
+                    mask_path_for_plot = get(entry, "mask_path", "")
+                    if !isfile(mask_path_for_plot)
+                        @warn "Mask not found for plotting: $(mask_path_for_plot). Plotting without mask."
+                        mask_path_for_plot = nothing
+                    end
+                end
+
+                plotdata, plotlayout, xSpectraMz, ySpectraMz = sumSpectrumPlot(msi_data, selected_folder_main, mask_path=mask_path_for_plot)
                 selectedTab = "tab2"
                 fTime = time()
                 eTime = round(fTime - sTime, digits=3)
                 msg = "Total plot loaded in $(eTime) seconds"
+                log_memory_usage("Sum Plot Generated", msi_data)
             catch e
                 msg = "Could not generate total spectrum plot: $e"
                 warning_msg = true
@@ -878,6 +974,10 @@ include("./julia_imzML_visual.jl")
                 btnPlotDisable = false
                 btnSpectraDisable = false
                 btnStartDisable = false
+                GC.gc()
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
+                end
             end
         end
     end
@@ -899,7 +999,17 @@ include("./julia_imzML_visual.jl")
             try
                 sTime = time()
                 registry = load_registry(registry_path)
-                target_path = registry[selected_folder_main]["source_path"]
+                
+                # Add error handling for registry access
+                if !haskey(registry, selected_folder_main)
+                    msg = "Dataset '$selected_folder_main' not found in registry."
+                    warning_msg = true
+                    return
+                end
+                
+                entry = registry[selected_folder_main]
+                target_path = entry["source_path"]
+
                 if target_path == "unknown (manually added)"
                     msg = "Dataset selected contained no route."
                     warning_msg = true
@@ -910,25 +1020,62 @@ include("./julia_imzML_visual.jl")
                     msg = "Reloading $(basename(target_path)) for analysis..."
                     full_route = target_path
                     msi_data = OpenMSIData(target_path)
-
+                    if haskey(get(entry, "metadata", Dict()), "global_min_mz") && entry["metadata"]["global_min_mz"] !== nothing
-                    existing_entry = get(registry, selected_folder_main, nothing)
+                        msi_data.global_min_mz = entry["metadata"]["global_min_mz"]
-                    if existing_entry !== nothing && haskey(get(existing_entry, "metadata", Dict()), "global_min_mz") && existing_entry["metadata"]["global_min_mz"] !== nothing
+                        msi_data.global_max_mz = entry["metadata"]["global_max_mz"]
-                        println("Injecting cached m/z range to skip Pass 1...")
-                        msi_data.global_min_mz = existing_entry["metadata"]["global_min_mz"]
-                        msi_data.global_max_mz = existing_entry["metadata"]["global_max_mz"]
                     else
                         precompute_analytics(msi_data)
                     end
                 end
 
-                y = yCoord < 0 ? abs(yCoord) : yCoord
+                local mask_path_for_plot::Union{String, Nothing} = nothing
-                plotdata, plotlayout, xSpectraMz, ySpectraMz = xySpectrumPlot(msi_data, xCoord, y, imgWidth, imgHeight, selected_folder_main)
+                if maskEnabled && get(entry, "has_mask", false)
-                xCoord = plotlayout.title == "Spectrum #$(xCoord)" ? xCoord : clamp(xCoord, 1, imgWidth)
+                    mask_path_for_plot = get(entry, "mask_path", "")
-                yCoord = plotlayout.title == "Spectrum #$(xCoord)" ? 0 : -clamp(y, 1, imgHeight)
+                    if !isfile(mask_path_for_plot)
+                        @warn "Mask not found for plotting: $(mask_path_for_plot). Plotting without mask."
+                        mask_path_for_plot = nothing
+                    end
+                end
+
+                # Convert to positive coordinates for processing
+                y_positive = yCoord < 0 ? abs(yCoord) : yCoord
+                plotdata, plotlayout, xSpectraMz, ySpectraMz = xySpectrumPlot(msi_data, xCoord, y_positive, imgWidth, imgHeight, selected_folder_main, mask_path=mask_path_for_plot)
+                
+                # Update coordinates based on actual plot title
+                # Extract title text from the Dict safely
+                actual_title = if plotlayout.title isa Dict && haskey(plotlayout.title, :text)
+                    plotlayout.title[:text]
+                elseif plotlayout.title isa Dict && haskey(plotlayout.title, "text")
+                    plotlayout.title["text"]
+                else
+                    string(plotlayout.title)  # Fallback
+                end
+                
+                if occursin("Masked Spectrum at", actual_title)
+                    # Extract coordinates from masked spectrum title
+                    coords_match = match(r"Masked Spectrum at \((\d+), (\d+)\)", actual_title)
+                    if coords_match !== nothing
+                        xCoord = parse(Int, coords_match.captures[1])
+                        yCoord = -parse(Int, coords_match.captures[2])  # Negative for display
+                    end
+                elseif occursin("Spectrum at", actual_title)
+                    # Extract coordinates from regular spectrum title
+                    coords_match = match(r"Spectrum at \((\d+), (\d+)\)", actual_title)
+                    if coords_match !== nothing
+                        xCoord = parse(Int, coords_match.captures[1])
+                        yCoord = -parse(Int, coords_match.captures[2])  # Negative for display
+                    end
+                else
+                    # For non-imaging data or fallback, just clamp the coordinates
+                    xCoord = clamp(xCoord, 1, imgWidth)
+                    yCoord = yCoord < 0 ? yCoord : -clamp(yCoord, 1, imgHeight)
+                end
+                
                 selectedTab = "tab2"
                 fTime = time()
                 eTime = round(fTime - sTime, digits=3)
                 msg = "Plot loaded in $(eTime) seconds"
+                log_memory_usage("XY Plot Generated", msi_data)
             catch e
                 msg = "Could not retrieve spectrum: $e"
                 warning_msg = true
@@ -938,6 +1085,10 @@ include("./julia_imzML_visual.jl")
                 btnPlotDisable = false
                 btnSpectraDisable = false
                 btnStartDisable = false
+                GC.gc()
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
+                end
             end
         end
     end
@@ -1213,6 +1364,13 @@ include("./julia_imzML_visual.jl")
 
     # This handler will now correctly load the first image from the newly selected folder.
     @onchange selected_folder_main begin
+        msi_data = nothing
+        log_memory_usage("Folder Changed (msi_data cleared)", msi_data)
+        GC.gc()
+        if Sys.islinux()
+            ccall(:malloc_trim, Int32, (Int32,), 0)
+        end
+
         if !isempty(selected_folder_main)
             folder_path = joinpath("public", selected_folder_main)
             if !isdir(folder_path) 
@@ -1226,6 +1384,7 @@ include("./julia_imzML_visual.jl")
                 plotlayoutImg = layoutImg
                 plotdataImgT = [traceImg]
                 plotlayoutImgT = layoutImg
+                imgWidth, imgHeight = 0, 0
                 return
             end
 
@@ -1236,7 +1395,8 @@ include("./julia_imzML_visual.jl")
             if !isempty(msi_bmp)
                 current_msi = first(msi_bmp)
                 imgInt = "/$(selected_folder_main)/$(current_msi)"
-                plotdataImg, plotlayoutImg, _, _ = loadImgPlot(imgInt)
+                plotdataImg, plotlayoutImg, w, h = loadImgPlot(imgInt)
+                imgWidth, imgHeight = w, h
                 text_nmass = replace(current_msi, r"MSI_|.bmp" => "")
                 msgimg = "<i>m/z</i>: $(replace(text_nmass, "_" => "."))"
                 if !isempty(col_msi_png)
@@ -1260,7 +1420,11 @@ include("./julia_imzML_visual.jl")
             if !isempty(triq_bmp)
                 current_triq = first(triq_bmp)
                 imgIntT = "/$(selected_folder_main)/$(current_triq)"
-                plotdataImgT, plotlayoutImgT, _, _ = loadImgPlot(imgIntT)
+                plotdataImgT, plotlayoutImgT, w, h = loadImgPlot(imgIntT)
+                # If no MSI image was loaded, dimensions from TrIQ image are used.
+                if isempty(msi_bmp)
+                    imgWidth, imgHeight = w, h
+                end
                 text_nmass = replace(current_triq, r"TrIQ_|.bmp" => "")
                 msgtriq = "TrIQ <i>m/z</i>: $(replace(text_nmass, "_" => "."))"
                 if !isempty(col_triq_png)
@@ -1276,6 +1440,10 @@ include("./julia_imzML_visual.jl")
                 plotdataImgT = [traceImg]
                 plotlayoutImgT = layoutImg
             end
+
+            if isempty(msi_bmp) && isempty(triq_bmp)
+                imgWidth, imgHeight = 0, 0
+            end
         end
     end
 
@@ -1391,90 +1559,130 @@ include("./julia_imzML_visual.jl")
     
     # 3d plot
     @onbutton image3dPlot begin
-        msg="Image 3D plot selected"
+        msg = "Image 3D plot selected"
-        cleaned_imgInt=replace(imgInt, r"\?.*" => "")
+        cleaned_imgInt = replace(imgInt, r"\?.*" => "")
-        cleaned_imgInt=lstrip(cleaned_imgInt, '/')
+        cleaned_imgInt = lstrip(cleaned_imgInt, '/')
-        var=joinpath( "./public", cleaned_imgInt )
+        var = joinpath("./public", cleaned_imgInt)
 
         if !isfile(var)
-            msg="Image could not be 3d plotted"
+            msg = "Image could not be 3d plotted"
-            warning_msg=true
+            warning_msg = true
             return
         end
 
-        progressPlot=true
+        progressPlot = true
-        btnPlotDisable=true
+        btnPlotDisable = true
-        btnStartDisable=true
+        btnStartDisable = true
-        btnSpectraDisable=true
+        btnSpectraDisable = true
 
         @async begin
             try
-                sTime=time()
+                # --- Get Mask Path ---
-                plotdata3d, plotlayout3d=loadSurfacePlot(imgInt)
+                local mask_path_for_plot::Union{String, Nothing} = nothing
-                GC.gc() # Trigger garbage collection
+                if maskEnabled && !isempty(selected_folder_main)
-                if Sys.islinux()
+                    registry = load_registry(registry_path)
-                    ccall(:malloc_trim, Int32, (Int32,), 0) # Ensure julia returns the freed memory to OS
+                    entry = get(registry, selected_folder_main, nothing)
+                    if entry !== nothing && get(entry, "has_mask", false)
+                        mask_path_candidate = get(entry, "mask_path", "")
+                        if isfile(mask_path_candidate)
+                            mask_path_for_plot = mask_path_candidate
+                        else
+                            @warn "Mask enabled but file not found: $(mask_path_candidate). Plotting without mask."
+                        end
+                    end
                 end
-                selectedTab="tab4"
+                # ---
-                fTime=time()
+
-                eTime=round(fTime-sTime,digits=3)
+                sTime = time()
-                msg="Plot loaded in $(eTime) seconds"
+                if mask_path_for_plot !== nothing
+                    plotdata3d, plotlayout3d = loadSurfacePlot(imgInt, mask_path_for_plot)
+                else
+                    plotdata3d, plotlayout3d = loadSurfacePlot(imgInt)
+                end
+
+                selectedTab = "tab4"
+                fTime = time()
+                eTime = round(fTime - sTime, digits=3)
+                msg = "Plot loaded in $(eTime) seconds"
+                log_memory_usage("Mean Plot Generated", msi_data)
             catch e
-                msg="Failed to load and process image: $e" 
+                msg = "Failed to load and process image: $e"
-                warning_msg=true
+                warning_msg = true
+                @error "3D plot generation failed" exception=(e, catch_backtrace())
             finally
                 progressPlot=false
                 btnPlotDisable=false
                 btnStartDisable=false
-                if msi_data !== nothing
+                btnSpectraDisable=false
-                    # We enable coord search and spectra plot creation
+                SpectraEnabled=true
-                    btnSpectraDisable=false
+                GC.gc()
-                    SpectraEnabled=true
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
                 end
             end
         end
-    end    # 3d plot for TrIQ
+    end
     
     @onbutton triq3dPlot begin
-        msg="TrIQ 3D plot selected"
+        msg = "TrIQ 3D plot selected"
-        cleaned_imgIntT=replace(imgIntT, r"\?.*" => "")
+        cleaned_imgIntT = replace(imgIntT, r"\?.*" => "")
-        cleaned_imgIntT=lstrip(cleaned_imgIntT, '/')
+        cleaned_imgIntT = lstrip(cleaned_imgIntT, '/')
-        var=joinpath( "./public", cleaned_imgIntT )
+        var = joinpath("./public", cleaned_imgIntT)
 
         if !isfile(var)
-            msg="Image could not be 3d plotted"
+            msg = "Image could not be 3d plotted"
-            warning_msg=true
+            warning_msg = true
             return
         end
 
-        progressPlot=true
+        progressPlot = true
-        btnPlotDisable=true
+        btnPlotDisable = true
-        btnStartDisable=true
+        btnStartDisable = true
-        btnSpectraDisable=true
+        btnSpectraDisable = true
 
         @async begin
             try
-                sTime=time()
+                # --- Get Mask Path ---
-                plotdata3d, plotlayout3d=loadSurfacePlot(imgIntT)
+                local mask_path_for_plot::Union{String, Nothing} = nothing
-                GC.gc() # Trigger garbage collection
+                if maskEnabled[] && !isempty(selected_folder_main)
-                if Sys.islinux()
+                    registry = load_registry(registry_path)
-                    ccall(:malloc_trim, Int32, (Int32,), 0) # Ensure julia returns the freed memory to OS
+                    entry = get(registry, selected_folder_main, nothing)
+                    if entry !== nothing && get(entry, "has_mask", false)
+                        mask_path_candidate = get(entry, "mask_path", "")
+                        if isfile(mask_path_candidate)
+                            mask_path_for_plot = mask_path_candidate
+                        else
+                            @warn "Mask enabled but file not found: $(mask_path_candidate). Plotting without mask."
+                        end
+                    end
                 end
-                selectedTab="tab4"
+                # ---
-                fTime=time()
+
-                eTime=round(fTime-sTime,digits=3)
+                sTime = time()
-                msg="Plot loaded in $(eTime) seconds"
+                if mask_path_for_plot !== nothing
+                    plotdata3d, plotlayout3d = loadSurfacePlot(imgIntT, mask_path_for_plot)
+                else
+                    plotdata3d, plotlayout3d = loadSurfacePlot(imgIntT)
+                end
+
+                selectedTab = "tab4"
+                fTime = time()
+                eTime = round(fTime - sTime, digits=3)
+                msg = "Plot loaded in $(eTime) seconds"
+                log_memory_usage("Mean Plot Generated", msi_data)
             catch e
-                msg="Failed to load and process image: $e" 
+                msg = "Failed to load and process image: $e"
-                warning_msg=true
+                warning_msg = true
+                @error "3D TrIQ plot generation failed" exception=(e, catch_backtrace())
             finally
                 progressPlot=false
                 btnPlotDisable=false
                 btnStartDisable=false
-                if msi_data !== nothing
+                btnSpectraDisable=false
-                    # We enable coord search and spectra plot creation
+                SpectraEnabled=true
-                    btnSpectraDisable=false
+                GC.gc()
-                    SpectraEnabled=true
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
                 end
             end
         end
@@ -1517,10 +1725,11 @@ include("./julia_imzML_visual.jl")
                 progressPlot=false
                 btnPlotDisable=false
                 btnStartDisable=false
-                if msi_data !== nothing
+                btnSpectraDisable=false
-                    # We enable coord search and spectra plot creation
+                SpectraEnabled=true
-                    btnSpectraDisable=false
+                GC.gc()
-                    SpectraEnabled=true
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
                 end
             end
         end
@@ -1562,10 +1771,11 @@ include("./julia_imzML_visual.jl")
                 progressPlot=false
                 btnPlotDisable=false
                 btnStartDisable=false
-                if msi_data !== nothing
+                btnSpectraDisable=false
-                    # We enable coord search and spectra plot creation
+                SpectraEnabled=true
-                    btnSpectraDisable=false
+                GC.gc()
-                    SpectraEnabled=true
+                if Sys.islinux()
+                    ccall(:malloc_trim, Int32, (Int32,), 0)
                 end
             end
         end
@@ -1579,7 +1789,7 @@ include("./julia_imzML_visual.jl")
     @onchange Nmass begin
         if !isempty(xSpectraMz)
             # Main spectrum trace
-            traceSpectra = PlotlyBase.scatter(
+            traceSpectra = PlotlyBase.stem(
                 x=xSpectraMz, 
                 y=ySpectraMz, 
                 marker=attr(size=1, color="blue", opacity=0.5), 
@@ -1651,13 +1861,13 @@ include("./julia_imzML_visual.jl")
 
     @onchange xCoord, yCoord begin 
         if selectedTab == "tab1"
-            plotdataImgT = filter(trace -> !(get(trace, :name, "") in ["Line X", "Line Y"]), plotdataImgT)
+            main_trace = plotdataImgT[1]  # The heatmap/image trace
             trace1, trace2 = crossLinesPlot(xCoord, yCoord, imgWidth, -imgHeight)
-            plotdataImgT = append!(plotdataImgT, [trace1, trace2])
+            plotdataImgT = [main_trace, trace1, trace2]  # Fresh array every time
         elseif selectedTab == "tab0"
-            plotdataImg = filter(trace -> !(get(trace, :name, "") in ["Line X", "Line Y", "Optical"]), plotdataImg)
+            main_trace = plotdataImg[1]  # The heatmap/image trace
             trace1, trace2 = crossLinesPlot(xCoord, yCoord, imgWidth, -imgHeight)
-            plotdataImg = append!(plotdataImg, [trace1, trace2])
+            plotdataImg = [main_trace, trace1, trace2]
         end
     end
 
@@ -1769,6 +1979,7 @@ include("./julia_imzML_visual.jl")
                 end
             end
         end
+        log_memory_usage("App Ready", msi_data)
         warmup_init()
     end
 

julia_imzML_visual.jl

@@ -1,32 +1,67 @@
-# == Search functions ==
+# julia_imzML_visual.jl
-# Functions that recieve a list to update, and the current direction both as string for
+
-# searching in the directory the position the list is going
+"""
+    increment_image(current_image, image_list)
+
+Finds the next image in a list.
+
+# Arguments
+- `current_image`: The current image file name.
+- `image_list`: The list of available image file names.
+
+# Returns
+- The file name of the next image, or the last image if the current one is the last or not found.
+- `nothing` if `image_list` is empty.
+"""
 function increment_image(current_image, image_list)
     if isempty(image_list)
         return nothing
     end
-    current_index=findfirst(isequal(current_image), image_list)
+    current_index = findfirst(isequal(current_image), image_list)
-    if current_index==nothing || current_index==length(image_list) || current_image ===""
+    if current_index === nothing || current_index == length(image_list) || current_image == ""
-        return image_list[length(image_list)]  # Return the current image if it's the last one or not found
+        return image_list[end]  # Return the last image if current is not found or is the last
     else
         return image_list[current_index + 1]  # Move to the next image
     end
 end
 
+"""
+    decrement_image(current_image, image_list)
+
+Finds the previous image in a list.
+
+# Arguments
+- `current_image`: The current image file name.
+- `image_list`: The list of available image file names.
+
+# Returns
+- The file name of the previous image, or the first image if the current one is the first or not found.
+- `nothing` if `image_list` is empty.
+"""
 function decrement_image(current_image, image_list)
     if isempty(image_list)
         return nothing
     end
-    current_index=findfirst(isequal(current_image), image_list)
+    current_index = findfirst(isequal(current_image), image_list)
-    if current_index==nothing || current_index==1 || current_image===""
+    if current_index === nothing || current_index == 1 || current_image == ""
-        return image_list[1]  # Return the current image if it's the first one or not found
+        return image_list[1]  # Return the first image if current is not found or is the first
     else
         return image_list[current_index - 1]  # Move to the previous image
     end
 end
 
-## Plot Image functions
+"""
-# Downsample an image matrix to a maximum dimension while preserving aspect ratio
+    downsample_image(img_matrix, max_dim::Int)
+
+Downsamples an image matrix to a maximum dimension while preserving aspect ratio.
+
+# Arguments
+- `img_matrix`: The image matrix to downsample.
+- `max_dim`: The maximum dimension (width or height) for the downsampled image.
+
+# Returns
+- The downsampled image matrix.
+"""
 function downsample_image(img_matrix, max_dim::Int)
     h, w = size(img_matrix)
     if h <= max_dim && w <= max_dim
@@ -46,26 +81,37 @@ function downsample_image(img_matrix, max_dim::Int)
     return imresize(img_matrix, (new_h, new_w))
 end
 
-# loadImgPlot recieves the local directory of the image as a string,
+"""
-# returns the layout and data for the heatmap plotly plot
+    loadImgPlot(interfaceImg::String)
-# this function loads the image into a plot
+
+Loads an image and creates a Plotly heatmap.
+
+# Arguments
+- `interfaceImg`: The path to the image file relative to the "public" directory.
+
+# Returns
+- `plotdata`: A vector containing the Plotly trace.
+- `plotlayout`: The Plotly layout for the plot.
+- `width`: The width of the loaded image.
+- `height`: The height of the loaded image.
+"""
 function loadImgPlot(interfaceImg::String)
     # Load the image
-    cleaned_img=replace(interfaceImg, r"\?.*" => "")
+    cleaned_img = replace(interfaceImg, r"\?.*" => "")
-    cleaned_img=lstrip(cleaned_img, '/')
+    cleaned_img = lstrip(cleaned_img, '/')
-    var=joinpath("./public", cleaned_img)
+    var = joinpath("./public", cleaned_img)
-    img=load(var)
+    img = load(var)
     # Convert to grayscale
-    img_gray=Gray.(img)
+    img_gray = Gray.(img)
-    img_array=Array(img_gray)
+    img_array = Array(img_gray)
-    elevation=Float32.(Array(img_array)) ./ 255.0
+    elevation = Float32.(Array(img_array)) ./ 255.0
     # Get the X, Y coordinates of the image
-    height, width=size(img_array)
+    height, width = size(img_array)
-    X=collect(1:width)
+    X = 1:width
-    Y=collect(1:height)
+    Y = 1:height
 
     # Create the layout
-    layout=PlotlyBase.Layout(
+    layout = PlotlyBase.Layout(
         title=PlotlyBase.attr(
             text="",
             font=PlotlyBase.attr(
@@ -83,11 +129,11 @@ function loadImgPlot(interfaceImg::String)
             visible=false,
             range=[-height, 0]
         ),
-        margin=attr(l=0,r=0,t=0,b=0,pad=0)
+        margin=attr(l=0, r=0, t=0, b=0, pad=0)
     )
 
     # Create the trace for the image
-    trace=PlotlyBase.heatmap(
+    trace = PlotlyBase.heatmap(
         z=elevation,
         x=X,
         y=-Y,
@@ -113,16 +159,29 @@ function loadImgPlot(interfaceImg::String)
         )
     )
 
-    plotdata=[trace]
+    plotdata = [trace]
-    plotlayout=layout
+    plotlayout = layout
     return plotdata, plotlayout, width, height
 end
 
-# loadImgPlot recieves the local directory of the image as a string, the local directory o the overlay image
+"""
-# and the transparency its required to have. Returns the layout and data for the heatmap plotly plot
+    loadImgPlot(interfaceImg::String, overlayImg::String, imgTrans::Float64)
-# this function loads the image into a plot
+
+Loads a main image and overlays a second image on top, creating a Plotly heatmap.
+
+# Arguments
+- `interfaceImg`: Path to the main image file.
+- `overlayImg`: Path to the overlay image file.
+- `imgTrans`: Transparency level for the overlay image (0.0 to 1.0).
+
+# Returns
+- `plotdata`: A vector containing the Plotly trace for the main image.
+- `plotlayout`: The Plotly layout, including the overlay image.
+- `width`: The width of the main image.
+- `height`: The height of the main image.
+"""
 function loadImgPlot(interfaceImg::String, overlayImg::String, imgTrans::Float64)
-    timestamp=string(time_ns())
+    timestamp = string(time_ns())
     # Load the main image
     cleaned_img = replace(interfaceImg, r"\?.*" => "")
     cleaned_img = lstrip(cleaned_img, '/')
@@ -134,8 +193,8 @@ function loadImgPlot(interfaceImg::String, overlayImg::String, imgTrans::Float64
     elevation = Float32.(Array(img_array)) ./ 255.0
     # Get the X, Y coordinates of the image
     height, width = size(img_array)
-    X = collect(1:width)
+    X = 1:width
-    Y = collect(1:height)
+    Y = 1:height
 
     # Create the layout with overlay image
     layoutImg = PlotlyBase.Layout(
@@ -147,40 +206,40 @@ function loadImgPlot(interfaceImg::String, overlayImg::String, imgTrans::Float64
                 color="black"
             )
         ),
-        images = [attr(
+        images=[attr(
-            source = "$(overlayImg)?t=$(timestamp)",
+            source="$(overlayImg)?t=$(timestamp)",
-            xref = "x",
+            xref="x",
-            yref = "y",
+            yref="y",
-            x = 0,
+            x=0,
-            y = 0,
+            y=0,
-            sizex = width,
+            sizex=width,
-            sizey = -height,
+            sizey=-height,
-            sizing = "stretch",
+            sizing="stretch",
-            opacity = imgTrans,
+            opacity=imgTrans,
-            layer = "above"  # Place the overlay image in the foreground
+            layer="above"  # Place the overlay image in the foreground
         )],
-        xaxis = PlotlyBase.attr(
+        xaxis=PlotlyBase.attr(
-            visible = false,
+            visible=false,
-            scaleanchor = "y",
+            scaleanchor="y",
-            range = [0, width]
+            range=[0, width]
         ),
-        yaxis = PlotlyBase.attr(
+        yaxis=PlotlyBase.attr(
-            visible = false,
+            visible=false,
-            range = [-height,0]
+            range=[-height, 0]
         ),
-        margin = attr(l = 0, r = 0, t = 0, b = 0, pad = 0)
+        margin=attr(l=0, r=0, t=0, b=0, pad=0)
     )
 
     # Create the trace for the main image
     trace = PlotlyBase.heatmap(
-        z = elevation,
+        z=elevation,
-        x = X,
+        x=X,
-        y = -Y,
+        y=-Y,
-        name = "",
+        name="",
-        hoverinfo = "x+y",
+        hoverinfo="x+y",
-        showlegend = false,
+        showlegend=false,
-        colorscale = "Viridis",
+        colorscale="Viridis",
-        showscale = false
+        showscale=false
     )
 
     plotdata = [trace]
@@ -188,10 +247,18 @@ function loadImgPlot(interfaceImg::String, overlayImg::String, imgTrans::Float64
     return plotdata, plotlayout, width, height
 end
 
-# loadContourPlot recieves the local directory of the image as a string,
+"""
-# returns the layout and data for the contour plotly plot
+    loadContourPlot(interfaceImg::String)
-# this function loads the image and applies a gaussian filter 
+
-# to smoothen it and loads it into a plot
+Loads an image, smooths it, and creates a Plotly contour plot.
+
+# Arguments
+- `interfaceImg`: Path to the image file.
+
+# Returns
+- `plotdata`: A vector containing the Plotly contour trace.
+- `plotlayout`: The Plotly layout for the plot.
+"""
 function loadContourPlot(interfaceImg::String)
     # Load the image
     cleaned_img=replace(interfaceImg, r"\?.*" => "")
@@ -258,10 +325,18 @@ function loadContourPlot(interfaceImg::String)
     return plotdata, plotlayout
 end
 
-# loadSurfacePlot recieves the local directory of the image as a string,
+"""
-# returns the layout and data for the surface plotly plot
+    loadSurfacePlot(interfaceImg::String)
-# this function loads the image and applies a gaussian filter 
+
-# to smoothen it and loads it into a 3D plot
+Loads an image, smooths it, and creates a 3D Plotly surface plot.
+
+# Arguments
+- `interfaceImg`: Path to the image file.
+
+# Returns
+- `plotdata`: A vector containing the Plotly surface trace.
+- `plotlayout`: The Plotly layout for the 3D plot.
+"""
 function loadSurfacePlot(interfaceImg::String)
     # Load the image
     cleaned_img=replace(interfaceImg, r"\?.*" => "")
@@ -347,28 +422,51 @@ function loadSurfacePlot(interfaceImg::String)
     return plotdata, plotlayout
 end
 
-# This function recieves the x and y coords currently selected, and the dimentions of
+"""
-# the image to create two traces that will display in a cross section
+    crossLinesPlot(x, y, maxwidth, maxheight)
+
+Creates two line traces for a crosshair indicator on a plot.
+
+# Arguments
+- `x`: The x-coordinate of the crosshair center.
+- `y`: The y-coordinate of the crosshair center.
+- `maxwidth`: The width of the plot area.
+- `maxheight`: The height of the plot area.
+
+# Returns
+- `trace1`: The horizontal line trace.
+- `trace2`: The vertical line trace.
+"""
 function crossLinesPlot(x, y, maxwidth, maxheight)
     # Define the coordinates for the two lines
-    l1_x=[0, maxwidth]
+    l1_x = [0, maxwidth]
-    l1_y=[y, y]
+    l1_y = [y, y]
-    l2_x=[x, x]
+    l2_x = [x, x]
-    l2_y=[0, maxheight]
+    l2_y = [0, maxheight]
 
     # Create the line traces
-    trace1=PlotlyBase.scatter(x=l1_x, y=l1_y, mode="lines",line=attr(color="red", width=0.5),name="Line X",showlegend=false)
+    trace1 = PlotlyBase.scatter(x=l1_x, y=l1_y, mode="lines", line=attr(color="red", width=0.5), name="Line X", showlegend=false)
-    trace2=PlotlyBase.scatter(x=l2_x, y=l2_y, mode="lines",line=attr(color="red", width=0.5),name="Line Y",showlegend=false)
+    trace2 = PlotlyBase.scatter(x=l2_x, y=l2_y, mode="lines", line=attr(color="red", width=0.5), name="Line Y", showlegend=false)
 
     return trace1, trace2
 end
 
-# This function is used for giving colorbar values a visual format
+"""
-# that shortens long values giving them scientific notation
+    log_tick_formatter(values::Vector{Float64})
+
+Formats a vector of numbers into strings with a custom scientific notation for use as tick labels.
+For example, 1000 becomes "100x10¹" and 0.01 becomes "1.0x10⁻²".
+
+# Arguments
+- `values`: A vector of `Float64` values to format.
+
+# Returns
+- A vector of formatted strings.
+"""
 function log_tick_formatter(values::Vector{Float64})
     # Initialize exponents dictionary
-    exponents=zeros(Int, length(values))
+    exponents = zeros(Int, length(values))
-    formValues=zeros(Float64, length(values))
+    formValues = zeros(Float64, length(values))
     for i in 1:length(values)
         value = values[i]
         if value >= 1000  # positive formatting for notation
@@ -382,78 +480,36 @@ function log_tick_formatter(values::Vector{Float64})
                 exponents[i] -= 1
             end
         end
-        formValues[i]=value
+        formValues[i] = value
     end
     return map((v, e) -> e == 0 ? "$(round(v, sigdigits=2))" : "$(round(v, sigdigits=2))x10" * Makie.UnicodeFun.to_superscript(e), formValues, exponents)
 end
 
-function generate_colorbar_image(slice_data::AbstractMatrix, color_levels::Int, output_path::String; use_triq::Bool=false, triq_prob::Float64=0.98)
+"""
-    # 1. Determine bounds based on whether TrIQ is used
+    meanSpectrumPlot(data::MSIData, dataset_name::String="")
-    min_val, max_val = if use_triq
+
-        MSI_src.get_outlier_thres(slice_data, triq_prob)
+Generates a plot of the mean spectrum from MSIData.
+
+# Arguments
+- `data`: The `MSIData` object.
+- `dataset_name`: Optional name of the dataset for the plot title.
+
+# Returns
+- `plotdata`: A vector containing the Plotly trace.
+- `plotlayout`: The Plotly layout for the plot.
+- `xSpectraMz`: The m/z values of the spectrum.
+- `ySpectraMz`: The intensity values of the spectrum.
+"""
+function meanSpectrumPlot(data::MSIData, dataset_name::String=""; mask_path::Union{String, Nothing}=nothing)
+    # Determine base title based on mask usage
+    base_title = if mask_path !== nothing
+        "Masked Average Spectrum"
     else
-        extrema(slice_data)
+        "Average Spectrum"
     end
     
-    # 2. Replicate the tick calculation logic from plot_slices
+    title_text = isempty(dataset_name) ? base_title : "$base_title for: $dataset_name"
-    bins = color_levels
-    levels = range(min_val, stop=max_val, length=bins + 1)
-    level_range = levels[end] - levels[1]
     
-    if level_range == 0
-        levels = range(min_val - 0.1, stop=max_val + 0.1, length=bins + 1)
-        level_range = 0.2
-    end
-    
-    exponent = level_range > 0 ? floor(log10(level_range)) / 3 : 0
-    scale = 10^(3 * exponent)
-    scaled_levels = levels ./ scale
-    
-    format_num = level_range > 0 ? floor(log10(level_range)) % 3 : 0
-    labels = if format_num == 0
-        [ @sprintf("%3.2f", lvl) for lvl in scaled_levels]
-    elseif format_num == 1
-        [ @sprintf("%3.2f", lvl) for lvl in scaled_levels]
-    else
-        [ @sprintf("%3.2f", lvl) for lvl in scaled_levels]
-    end
-    
-    divisors = 2:7
-    remainders = (bins - 1) .% divisors
-    best_divisor = divisors[findlast(x -> x == minimum(remainders), remainders)]
-    tick_indices = round.(Int, range(1, stop=bins + 1, length=best_divisor + 1))
-
-    if !(1 in tick_indices)
-        pushfirst!(tick_indices, 1)
-    end
-    if !((bins + 1) in tick_indices)
-        push!(tick_indices, bins + 1)
-    end
-    unique!(sort!(tick_indices))
-    
-    tick_positions = levels[tick_indices]
-    tick_labels = labels[tick_indices]
-
-    # 3. Create and save the colorbar image
-    fig = Figure(size=(150, 250))
-    Colorbar(fig[1, 1], 
-        colormap=cgrad(:viridis, bins, categorical=true),
-        # limits=(min_val, max_val),
-        limits=(levels[1], levels[end]),
-        label=(scale == 1 ? "Intensity" : "Intensity ×10^$(round(Int, 3 * exponent))"),
-        ticks=(tick_positions, tick_labels),
-        labelsize=20,
-        ticklabelsize=16
-    )
-    save(output_path, fig)
-end
-
-# meanSpectrumPlot recieves the local directory of the image as a string,
-# returns the layout and data for the surface plotly plot
-# this function loads the spectra data and makes a mean to display
-# its values in the spectrum plot
-function meanSpectrumPlot(data::MSIData, dataset_name::String="")
-    title_text = isempty(dataset_name) ? "Average Spectrum Plot" : "Average Spectrum for: $dataset_name"
     layout = PlotlyBase.Layout(
         title=PlotlyBase.attr(
             text=title_text,
@@ -477,42 +533,76 @@ function meanSpectrumPlot(data::MSIData, dataset_name::String="")
     )
 
     # Use the new, efficient function from the backend
-    xSpectraMz, ySpectraMz = get_average_spectrum(data)
+    xSpectraMz, ySpectraMz = get_average_spectrum(data, mask_path=mask_path)
 
-    if isempty(xSpectraMz)
+    if isempty(xSpectraMz) || isempty(ySpectraMz)
         @warn "Average spectrum is empty."
-        trace = PlotlyBase.scatter(x=Float64[], y=Float64[])
+        trace = PlotlyBase.stem(x=Float64[], y=Float64[])
+        # Update title to indicate empty spectrum
+        layout.title.text = "Empty " * layout.title.text
     else
-        trace = PlotlyBase.scatter(x=xSpectraMz, y=ySpectraMz, marker=attr(size=1, color="blue", opacity=0.5), name="Average", hoverinfo="x",hovertemplate="<b>m/z</b>: %{x:.4f}<extra></extra>")
+        trace = PlotlyBase.stem(x=xSpectraMz, y=ySpectraMz, marker=attr(size=1, color="blue", opacity=0.5), name="Average", hoverinfo="x", hovertemplate="<b>m/z</b>: %{x:.4f}<extra></extra>")
     end
 
-
     plotdata = [trace]
     plotlayout = layout
     return plotdata, plotlayout, xSpectraMz, ySpectraMz
 end
 
-function xySpectrumPlot(data::MSIData, xCoord::Int, yCoord::Int, imgWidth::Int, imgHeight::Int, dataset_name::String="")
+"""
+    xySpectrumPlot(data::MSIData, xCoord::Int, yCoord::Int, imgWidth::Int, imgHeight::Int, dataset_name::String="")
+
+Generates a plot for the spectrum at a specific coordinate (for imaging data) or index (for non-imaging data).
+
+# Arguments
+- `data`: The `MSIData` object.
+- `xCoord`: The x-coordinate or spectrum index.
+- `yCoord`: The y-coordinate (used for imaging data).
+- `imgWidth`: The width of the MSI image.
+- `imgHeight`: The height of the MSI image.
+- `dataset_name`: Optional name of the dataset for the plot title.
+
+# Returns
+- `plotdata`: A vector containing the Plotly trace.
+- `plotlayout`: The Plotly layout for the plot.
+- `mz`: The m/z values of the spectrum.
+- `intensity`: The intensity values of the spectrum.
+"""
+function xySpectrumPlot(data::MSIData, xCoord::Int, yCoord::Int, imgWidth::Int, imgHeight::Int, dataset_name::String=""; mask_path::Union{String, Nothing}=nothing)
     local mz::AbstractVector, intensity::AbstractVector
     local plot_title::String
 
     is_imaging = data.source isa ImzMLSource
 
     if is_imaging
-        # For imaging data, use (X, Y) coordinates
         x = clamp(xCoord, 1, imgWidth)
         y = clamp(yCoord, 1, imgHeight)
 
-        process_spectrum(data, Int(x), Int(y)) do recieved_mz, recieved_intensity
+        if mask_path !== nothing
-            mz = recieved_mz
+            mask_matrix = load_and_prepare_mask(mask_path, (imgWidth, imgHeight))
-            intensity = recieved_intensity
+            if !mask_matrix[y, x]
+                @warn "Coordinate ($x, $y) is outside the specified mask."
+                mz, intensity = Float64[], Float64[]
+                base_title = "Spectrum Outside Mask at ($x, $y)"
+            else
+                process_spectrum(data, Int(x), Int(y)) do recieved_mz, recieved_intensity
+                    mz = recieved_mz
+                    intensity = recieved_intensity
+                end
+                base_title = "Masked Spectrum at ($x, $y)"
+            end
+        else
+            process_spectrum(data, Int(x), Int(y)) do recieved_mz, recieved_intensity
+                mz = recieved_mz
+                intensity = recieved_intensity
+            end
+            base_title = "Spectrum at ($x, $y)"
         end
-        base_title = "Spectrum at ($x, $y)"
     else
         # For non-imaging data, treat xCoord as the spectrum index
         index = clamp(xCoord, 1, length(data.spectra_metadata))
 
-        process_spectrum(data, index)  do recieved_mz, recieved_intensity
+        process_spectrum(data, index) do recieved_mz, recieved_intensity
             mz = recieved_mz
             intensity = recieved_intensity
         end
@@ -546,7 +636,7 @@ function xySpectrumPlot(data::MSIData, xCoord::Int, yCoord::Int, imgWidth::Int,
     # Downsample for plotting performance
     mz_down, int_down = MSI_src.downsample_spectrum(mz, intensity)
 
-    trace = PlotlyBase.scatter(x=mz_down, y=int_down, marker=attr(size=1, color="blue", opacity=0.5), name="Spectrum", hoverinfo="x", hovertemplate="<b>m/z</b>: %{x:.4f}<extra></extra>")
+    trace = PlotlyBase.stem(x=mz_down, y=int_down, marker=attr(size=1, color="blue", opacity=0.5), name="Spectrum", hoverinfo="x", hovertemplate="<b>m/z</b>: %{x:.4f}<extra></extra>")
 
     plotdata = [trace]
     plotlayout = layout
@@ -554,8 +644,31 @@ function xySpectrumPlot(data::MSIData, xCoord::Int, yCoord::Int, imgWidth::Int,
     return plotdata, plotlayout, mz, intensity
 end
 
-function sumSpectrumPlot(data::MSIData, dataset_name::String="")
+"""
-    title_text = isempty(dataset_name) ? "Total Spectrum Plot" : "Total Spectrum for: $dataset_name"
+    sumSpectrumPlot(data::MSIData, dataset_name::String="")
+
+Generates a plot of the total (summed) spectrum from MSIData.
+
+# Arguments
+- `data`: The `MSIData` object.
+- `dataset_name`: Optional name of the dataset for the plot title.
+
+# Returns
+- `plotdata`: A vector containing the Plotly trace.
+- `plotlayout`: The Plotly layout for the plot.
+- `xSpectraMz`: The m/z values of the spectrum.
+- `ySpectraMz`: The intensity values of the spectrum.
+"""
+function sumSpectrumPlot(data::MSIData, dataset_name::String=""; mask_path::Union{String, Nothing}=nothing)
+    # Determine base title based on mask usage
+    base_title = if mask_path !== nothing
+        "Masked Total Spectrum"
+    else
+        "Total Spectrum"
+    end
+    
+    title_text = isempty(dataset_name) ? base_title : "$base_title for: $dataset_name"
+    
     layout = PlotlyBase.Layout(
         title=PlotlyBase.attr(
             text=title_text,
@@ -579,13 +692,15 @@ function sumSpectrumPlot(data::MSIData, dataset_name::String="")
     )
 
     # Use the get_total_spectrum function from the backend
-    xSpectraMz, ySpectraMz = get_total_spectrum(data)
+    xSpectraMz, ySpectraMz, num_spectra = get_total_spectrum(data, mask_path=mask_path)
 
-    if isempty(xSpectraMz)
+    if isempty(xSpectraMz) || isempty(ySpectraMz)
         @warn "Total spectrum is empty."
-        trace = PlotlyBase.scatter(x=Float64[], y=Float64[])
+        trace = PlotlyBase.stem(x=Float64[], y=Float64[])
+        # Update title to indicate empty spectrum
+        layout.title.text = "Empty " * layout.title.text
     else
-        trace = PlotlyBase.scatter(x=xSpectraMz, y=ySpectraMz, marker=attr(size=1, color="blue", opacity=0.5), name="Total", hoverinfo="x",hovertemplate="<b>m/z</b>: %{x:.4f}<extra></extra>")
+        trace = PlotlyBase.stem(x=xSpectraMz, y=ySpectraMz, marker=attr(size=1, color="blue", opacity=0.5), name="Total", hoverinfo="x", hovertemplate="<b>m/z</b>: %{x:.4f}<extra></extra>")
     end
 
     plotdata = [trace]
@@ -593,30 +708,32 @@ function sumSpectrumPlot(data::MSIData, dataset_name::String="")
     return plotdata, plotlayout, xSpectraMz, ySpectraMz
 end
 
+"""
+    warmup_init()
+
+Performs pre-compilation of key functions at application startup to reduce first-use latency.
+This is run asynchronously and should not block application startup.
+"""
 function warmup_init()
     @async begin
         println("Pre-compiling functions at startup...")
 
-        # Create a dummy MSIData object to be used for pre-compilation
+        # Pre-compile image processing and plotting functions
-        # dummy_source = ImzMLSource("dummy.ibd", Float32, Float32)
+        try
-        # dummy_meta = MSI_src.SpectrumMetadata(0,0,"",MSI_src.UNKNOWN, MSI_src.SpectrumAsset(Float32,false,0,0,:mz), MSI_src.SpectrumAsset(Float32,false,0,0,:intensity))
+            TrIQ(zeros(10, 10), 256, 0.98)
-        # dummy_msi_data = MSIData(dummy_source, [dummy_meta], (1,1), zeros(Int,1,1), 0)
+        catch
-
+        end
-        # Pre-compile functions from btnSearch
+        try
-        # try OpenMSIData("dummy.imzML") catch end
+            quantize_intensity(zeros(10, 10), 256)
-        # try precompute_analytics(dummy_msi_data) catch end
+        catch
-
+        end
-        # Pre-compile functions from mainProcess
-        # try get_mz_slice(dummy_msi_data, 1.0, 1.0) catch end
-        try TrIQ(zeros(10,10), 256, 0.98) catch end
-        try quantize_intensity(zeros(10,10), 256) catch end
 
         dummy_bmp_path = joinpath("public", "dummy.bmp")
         dummy_png_path = joinpath("public", "dummy.png")
         try
             save_bitmap(dummy_bmp_path, zeros(UInt8, 10, 10), ViridisPalette)
             loadImgPlot("/dummy.bmp")
-            generate_colorbar_image(zeros(10,10), 256, dummy_png_path)
+            generate_colorbar_image(zeros(10, 10), 256, dummy_png_path)
         catch e
             @warn "Pre-compilation step failed (this is expected if dummy files can't be created/read)"
         finally
@@ -628,18 +745,41 @@ function warmup_init()
     end
 end
 
+"""
+    load_registry(registry_path)
+
+Loads the dataset registry from a JSON file.
+
+# Arguments
+- `registry_path`: Path to the `registry.json` file.
+
+# Returns
+- A dictionary containing the registry data. Returns an empty dictionary if the file doesn't exist or fails to parse.
+"""
 function load_registry(registry_path)
     if isfile(registry_path)
         try
-            return JSON.parsefile(registry_path, dicttype=Dict{String, Any})
+            return JSON.parsefile(registry_path, dicttype=Dict{String,Any})
         catch e
             @error "Failed to parse registry.json: $e"
-            return Dict{String, Any}()
+            return Dict{String,Any}()
         end
     end
-    return Dict{String, Any}()
+    return Dict{String,Any}()
 end
 
+"""
+    extract_metadata(msi_data::MSIData, source_path::String)
+
+Extracts key metadata from an `MSIData` object for display.
+
+# Arguments
+- `msi_data`: The `MSIData` object.
+- `source_path`: The path to the source data file.
+
+# Returns
+- A dictionary containing summary statistics and metadata.
+"""
 function extract_metadata(msi_data::MSIData, source_path::String)
     df = msi_data.spectrum_stats_df
     if df === nothing
@@ -686,16 +826,37 @@ function extract_metadata(msi_data::MSIData, source_path::String)
     )
 end
 
+"""
+    update_registry(registry_path, dataset_name, source_path, metadata=nothing, is_imzML=false)
+
+Adds or updates an entry in the dataset registry JSON file.
+
+# Arguments
+- `registry_path`: Path to the `registry.json` file.
+- `dataset_name`: The name of the dataset.
+- `source_path`: The path to the source data file.
+- `metadata`: Optional dictionary of metadata to store.
+- `is_imzML`: Boolean indicating if the source is an imzML file.
+"""
 function update_registry(registry_path, dataset_name, source_path, metadata=nothing, is_imzML=false)
     registry = load_registry(registry_path)
-    entry = Dict{String, Any}( # Explicitly type the dictionary to allow mixed value types
+    
-        "source_path" => source_path, 
+    # Get existing entry if it exists, otherwise create new one
-        "processed_date" => string(now()),
+    existing_entry = get(registry, dataset_name, Dict{String,Any}())
-        "is_imzML" => is_imzML
+    
-    )
+    # Start with existing data and update only the basic fields
+    entry = copy(existing_entry)
+    entry["source_path"] = source_path
+    entry["processed_date"] = string(now())
+    entry["is_imzML"] = is_imzML
+    
+    # Only update metadata if provided
     if metadata !== nothing
         entry["metadata"] = metadata
     end
+    
+    # Note: has_mask and mask_path are preserved from existing_entry if they exist
+    
     registry[dataset_name] = entry
 
     try
@@ -707,7 +868,25 @@ function update_registry(registry_path, dataset_name, source_path, metadata=noth
     end
 end
 
-function process_file_safely(file_path, masses, params, progress_message_ref, overall_progress_ref)
+"""
+    process_file_safely(file_path, masses, params, progress_message_ref, overall_progress_ref)
+
+Safely processes a single MSI data file, generating and saving m/z slices.
+
+This function handles loading data, generating slices, saving results as bitmaps,
+and updating the data registry. It includes error handling and memory cleanup.
+
+# Arguments
+- `file_path`: Path to the `.imzML` file.
+- `masses`: A vector of m/z values to generate slices for.
+- `params`: A structure or dictionary containing processing parameters.
+- `progress_message_ref`: A reference to update with progress messages.
+- `overall_progress_ref`: A reference to update with overall progress.
+
+# Returns
+- A tuple `(success::Bool, message::String)`.
+"""
+function process_file_safely(file_path, masses, params, progress_message_ref, overall_progress_ref; use_mask::Bool=false)
     local_msi_data = nothing
     dataset_name = replace(basename(file_path), r"\.imzML$"i => "")
     output_dir = joinpath("public", dataset_name)
@@ -722,15 +901,39 @@ function process_file_safely(file_path, masses, params, progress_message_ref, ov
             return (false, "Skipped: Not an imzML file")
         end
 
-        # --- Generate Slices (this will call precompute_analytics if needed) ---
+        # --- Get mask path and load mask_matrix only if use_mask is true ---
-        progress_message_ref = "Generating $(length(masses)) slices for $(dataset_name)..."
+        local mask_matrix_for_triq::Union{BitMatrix, Nothing} = nothing
-        slice_dict = get_multiple_mz_slices(local_msi_data, masses, params.tolerance)
+        mask_path = nothing
+        if use_mask
+            registry = load_registry(params.registry)
+            entry = get(registry, dataset_name, nothing)
+            if entry !== nothing && get(entry, "has_mask", false)
+                mask_path_candidate = get(entry, "mask_path", "")
+                if isfile(mask_path_candidate)
+                    mask_path = mask_path_candidate
+                    # Load the mask matrix here
+                    mask_matrix_for_triq = load_and_prepare_mask(mask_path, (local_msi_data.image_dims[1], local_msi_data.image_dims[2]))
+                    println("DEBUG: Using mask: $(mask_path) with dimensions $(size(mask_matrix_for_triq))")
+                else
+                    @warn "Mask enabled but file not found: $(mask_path_candidate). Clearing invalid mask entry."
+                    # Clear invalid mask entry
+                    update_registry_mask_fields(params.registry, dataset_name, false, "")
+                    mask_path = nothing
+                end
+            else
+                @warn "Mask enabled but no valid mask entry found for: $(dataset_name)"
+            end
+        end
 
-        # --- Extract metadata *after* it has been computed ---
+        # --- Generate Slices ---
+        progress_message_ref = "Generating $(length(masses)) slices for $(dataset_name)..."
+        slice_dict = get_multiple_mz_slices(local_msi_data, masses, params.tolerance, mask_path=mask_path)
+
+        # --- Extract metadata ---
         metadata = extract_metadata(local_msi_data, file_path)
 
         # --- Save Slices ---
-        mkpath(output_dir)  # Ensure output directory exists
+        mkpath(output_dir)
 
         for (mass_idx, mass) in enumerate(masses)
             progress_message_ref = "File $(params.fileIdx)/$(params.nFiles): Saving slice for m/z=$mass"
@@ -744,7 +947,7 @@ function process_file_safely(file_path, masses, params, progress_message_ref, ov
                 sliceQuant = zeros(UInt8, size(slice))
                 @warn "No intensity data for m/z = $mass in $(dataset_name)"
             else
-                sliceQuant = params.triqE ? TrIQ(slice, params.colorL, params.triqP) : quantize_intensity(slice, params.colorL)
+                sliceQuant = params.triqE ? TrIQ(slice, params.colorL, params.triqP, mask_matrix=mask_matrix_for_triq) : quantize_intensity(slice, params.colorL, mask_matrix=mask_matrix_for_triq)
                 if params.medianF
                     sliceQuant = round.(UInt8, median_filter(sliceQuant))
                 end
@@ -752,7 +955,8 @@ function process_file_safely(file_path, masses, params, progress_message_ref, ov
 
             save_bitmap(joinpath(output_dir, bitmap_filename), sliceQuant, ViridisPalette)
             if !all(iszero, slice)
-                generate_colorbar_image(slice, params.colorL, joinpath(output_dir, colorbar_filename); use_triq=params.triqE, triq_prob=params.triqP)
+                generate_colorbar_image(slice, params.colorL, joinpath(output_dir, colorbar_filename); 
+                    use_triq=params.triqE, triq_prob=params.triqP, mask_path=mask_path)
             end
         end
 
@@ -774,3 +978,149 @@ function process_file_safely(file_path, masses, params, progress_message_ref, ov
         end
     end
 end
+
+function update_registry_mask_fields(registry_path, dataset_name, has_mask, mask_path)
+    registry = load_registry(registry_path)
+    
+    if haskey(registry, dataset_name)
+        registry[dataset_name]["has_mask"] = has_mask
+        registry[dataset_name]["mask_path"] = mask_path
+    else
+        # Create new entry if it doesn't exist
+        registry[dataset_name] = Dict{String,Any}(
+            "has_mask" => has_mask,
+            "mask_path" => mask_path,
+            "source_path" => "",
+            "processed_date" => string(now()),
+            "is_imzML" => false
+        )
+    end
+
+    try
+        open(registry_path, "w") do f
+            JSON.print(f, registry, 4)
+        end
+    catch e
+        @error "Failed to write to registry.json: $e"
+    end
+end
+
+"""
+    loadSurfacePlot(interfaceImg::String, mask_path::String)
+
+Loads an image, smooths it, applies a mask, crops the data to the masked region, and creates a 3D Plotly surface plot.
+
+# Arguments
+- `interfaceImg`: Path to the image file.
+- `mask_path`: Path to a mask file. Areas outside the mask will be removed and the plot axes will be cropped to the masked region.
+
+# Returns
+- `plotdata`: A vector containing the Plotly surface trace.
+- `plotlayout`: The Plotly layout for the 3D plot.
+"""
+function loadSurfacePlot(interfaceImg::String, mask_path::String)
+    # Load the image
+    cleaned_img = replace(interfaceImg, r"\?.*" => "")
+    cleaned_img = lstrip(cleaned_img, '/')
+    var = joinpath("./public", cleaned_img)
+    img = load(var)
+    img_gray = Gray.(img) # Convert to grayscale
+    img_array = Array(img_gray)
+    elevation = Float32.(Array(img_array)) ./ 255.0 # Normalize between 0 and 1
+
+    # Smooth the image
+    sigma = 3.0
+    kernel = Kernel.gaussian(sigma)
+    elevation_smoothed = imfilter(elevation, kernel)
+
+    # --- APPLY MASK and CROP DATA ---
+    mask_matrix = load_and_prepare_mask(mask_path, (size(elevation_smoothed, 2), size(elevation_smoothed, 1)))
+    elevation_smoothed[.!mask_matrix] .= NaN32
+
+    non_nan_indices = findall(!isnan, elevation_smoothed)
+    if isempty(non_nan_indices)
+        # Return an empty plot if mask removes everything
+        return [PlotlyBase.surface()], PlotlyBase.Layout(title="Empty plot: Mask covered all data")
+    end
+
+    row_indices = [idx[1] for idx in non_nan_indices]
+    col_indices = [idx[2] for idx in non_nan_indices]
+    
+    y_range = minimum(row_indices):maximum(row_indices)
+    x_range = minimum(col_indices):maximum(col_indices)
+
+    cropped_elevation = elevation_smoothed[y_range, x_range]
+    # ---
+
+    # --- DOWNSAMPLE CROPPED DATA ---
+    cropped_elevation = downsample_image(cropped_elevation, 256)
+    # ---
+
+    # Create the X, Y meshgrid coordinates for the CROPPED data
+    x_coords = x_range
+    y_coords = y_range
+    X = repeat(reshape(x_coords, 1, length(x_coords)), length(y_coords), 1)
+    Y = repeat(reshape(y_coords, length(y_coords), 1), 1, length(x_coords))
+
+    # Define tick values and text for colorbars, ignoring NaNs
+    non_nan_values = filter(!isnan, cropped_elevation)
+    min_val = isempty(non_nan_values) ? 0.0 : minimum(non_nan_values)
+    max_val = isempty(non_nan_values) ? 1.0 : maximum(non_nan_values)
+    tickV = range(min_val, stop=max_val, length=8)
+    tickT = log_tick_formatter(collect(tickV))
+
+    # Transpose the elevation_smoothed array if Y axis is longer than X axis to fix chopping
+    cropped_elevation = transpose(cropped_elevation)
+    if size(cropped_elevation, 1) < size(cropped_elevation, 2)
+        Y = -Y
+    else
+        X = -X
+    end
+
+    # Calculate the number of ticks and aspect ratio for the 3d plot
+    x_nticks = min(20, length(x_coords))
+    y_nticks = min(20, length(y_coords))
+    z_nticks = 5
+    aspect_ratio = attr(x=1, y=length(y_coords) / length(x_coords), z=0.5)
+    
+    # Define the layout for the 3D plot
+    layout3D = PlotlyBase.Layout(
+        title=PlotlyBase.attr(
+            text="3D surface plot of $cleaned_img (masked, cropped, downsampled)",
+            font=PlotlyBase.attr(
+                family="Roboto, Lato, sans-serif",
+                size=18,
+                color="black"
+            )
+        ),
+        scene=attr(
+            xaxis_nticks=x_nticks,
+            yaxis_nticks=y_nticks,
+            zaxis_nticks=z_nticks,
+            camera=attr(eye=attr(x=0, y=1, z=0.5)),
+            aspectratio=aspect_ratio
+        ),
+        margin=attr(l=0, r=0, t=120, b=0, pad=0)
+    )
+
+    trace3D = PlotlyBase.surface(
+        x=X[1, :],
+        y=Y[:, 1],
+        z=cropped_elevation,
+        contours_z=attr(
+            show=true,
+            usecolormap=true,
+            highlightcolor="limegreen",
+            project_z=true
+        ),
+        colorscale="Viridis",
+        colorbar=attr(
+            tickvals=tickV,
+            ticktext=tickT,
+            nticks=8
+        )
+    )
+    plotdata = [trace3D]
+    plotlayout = layout3D
+    return plotdata, plotlayout
+end

mask.jl

@@ -9,14 +9,12 @@ using Statistics, NaturalSort, LinearAlgebra, StipplePlotly
 using Base.Filesystem: mv
 
 using MSI_src
-using .MSI_src: MSIData
+using .MSI_src: MSIData, process_image_pipeline
 
 # Plot Handling
 include("./julia_imzML_visual.jl")
 
 # Image Processing Pipeline
-include("src/ImageProcessing.jl")
-using .ImageProcessing
 using ImageBinarization
 
 function load_and_binarize_mask(path)
@@ -46,7 +44,7 @@ function alter_image(img_path, otsu_scale, noise_size_percent, hole_size_percent
         gray_img = Float32.(ensure_grayscale(original_img))
 
         binary, noise_removed, holes_filled, smoothed =
-            ImageProcessing.process_image_pipeline(gray_img;
+            process_image_pipeline(gray_img;
                 otsu_scale=otsu_scale, noise_size_percent=noise_size_percent,
                 hole_size_percent=hole_size_percent, smoothing=smoothing_level)
 
@@ -250,7 +248,6 @@ end
             is_browsing_slices = true
             is_editing_mask = false
             folder_path = joinpath("public", selected_folder_main)
-            println("Selected folder: $selected_folder_main")
             if !isdir(folder_path) 
                 imgInt = ""
                 msgimg = "Folder not found."
@@ -675,14 +672,16 @@ end
                 reg_path = abspath(joinpath(@__DIR__, "public", "registry.json"))
                 registry = isfile(reg_path) ? JSON.parsefile(reg_path) : Dict{String, Any}()
 
+                full_mask_path = abspath(final_path)
+                
                 # Update the registry entry
                 if haskey(registry, selected_folder_main)
-                    registry[selected_folder_main]["mask_path"] = "/css/masks/$(final_mask_name)"
+                    registry[selected_folder_main]["mask_path"] = full_mask_path
                     registry[selected_folder_main]["has_mask"] = true
                 else
                     # Create a new entry if folder doesn't exist in registry
                     registry[selected_folder_main] = Dict(
-                        "mask_path" => "/css/masks/$(final_mask_name)",
+                        "mask_path" => full_mask_path,
                         "has_mask" => true,
                         "is_imzML" => true,
                         "processed_date" => string(Dates.now())

mask.jl.html

@@ -1,60 +1,77 @@
 <header id="header">
     <img src="/css/LABI_logo.png" alt="Labi Logo Icon" id="imgLogo">
     <div>
-      <h4>JuliaMSI - Mask Editor&nbsp;</h4>
+        <h4>JuliaMSI - Mask Editor&nbsp;</h4>
     </div>
 </header>
 
 <div id="extDivStyle" class="row col-12 q-pa-xl">
-<div class="row col-4">
+    <div class="row col-4">
-    <!-- Left Panel: Controls -->
+        <!-- Left Panel: Controls -->
-    <div class="st-col col-12 st-module q-pa-md">
+        <div class="st-col col-12 st-module q-pa-md">
-        <div class="text-h6">Mask Generation Workflow</div>
+            <div class="text-h6">Mask Generation Workflow</div>
 
-        <!-- Step 1: Select Slice -->
+            <!-- Step 1: Select Slice -->
-        <div class="text-subtitle1 q-mt-md">Step 1: Select Slice</div>
+            <div class="text-subtitle1 q-mt-md">Step 1: Select Slice</div>
-        <div class="row items-center">
+            <div class="row items-center">
-            <q-select v-model="selected_folder_main" :options="image_available_folders" label="Select Dataset" class="q-ma-sm col" standout="custom-standout" :disable="is_editing_mask"></q-select>
+                <q-select v-model="selected_folder_main" :options="image_available_folders" label="Select Dataset"
-        </div>
+                    class="q-ma-sm col" standout="custom-standout" :disable="is_editing_mask"></q-select>
-        <p class="text-center" v-html="msgimg"></p>
+            </div>
+            <p class="text-center" v-html="msgimg"></p>
 
-        <!-- Step 2: Provide Mask Input -->
+            <!-- Step 2: Provide Mask Input -->
-        <div class="text-subtitle1 q-mt-md">Step 2: Create/Upload Mask</div>
+            <div class="text-subtitle1 q-mt-md">Step 2: Create/Upload Mask</div>
-        <div class="row justify-around">
+            <div class="row justify-around">
-            <q-btn class="q-ma-sm btn-style" v-on:click="btn_use_slice_as_mask = true" label="Use Current Slice" :disable="!is_browsing_slices || !imgInt"></q-btn>
+                <q-btn class="q-ma-sm btn-style" v-on:click="btn_use_slice_as_mask = true" label="Use Current Slice"
-            <q-btn class="q-ma-sm btn-style" v-on:click="btn_upload_mask = true" label="Upload to Create Mask" :disable="!is_browsing_slices || !imgInt"></q-btn>
+                    :disable="!is_browsing_slices || !imgInt"></q-btn>
-            <q-btn class="q-ma-sm btn-style" v-on:click="btn_change_slice = true" label="Change Mask Blueprint" :disable="!is_editing_mask"></q-btn>
+                <q-btn class="q-ma-sm btn-style" v-on:click="btn_upload_mask = true" label="Upload to Create Mask"
-        </div>
+                    :disable="!is_browsing_slices || !imgInt"></q-btn>
+                <q-btn class="q-ma-sm btn-style" v-on:click="btn_change_slice = true" label="Change Mask Blueprint"
+                    :disable="!is_editing_mask"></q-btn>
+            </div>
 
-        <!-- Step 3: Automated Processing & Overlay -->
+            <!-- Step 3: Automated Processing & Overlay -->
-        <div class="text-subtitle1 q-mt-md">Step 3: Adjust & Verify</div>
+            <div class="text-subtitle1 q-mt-md">Step 3: Adjust & Verify</div>
-        <div class="q-mt-sm">
+            <div class="q-mt-sm">
-          <p class="text-caption text-center">Adjust automated processing for the initial mask:</p>
+                <p class="text-caption text-center">Adjust automated processing for the initial mask:</p>
-          <q-slider color="primary" label-always v-model="otsu_scale" :label-value="'Otsu Scale: ' + otsu_scale.toFixed(2)" :step="0.01" :min="0.1" :max="2.0" :disable="!is_editing_mask"></q-slider>
+                <q-slider color="primary" label-always v-model="otsu_scale"
-          <q-slider color="primary" label-always v-model="noise_size_percent" :label-value="'Noise Size: ' + (noise_size_percent * 100).toFixed(3) + '%'" :step="0.001" :min="0.001" :max="0.5" :disable="!is_editing_mask"></q-slider>
+                    :label-value="'Otsu Scale: ' + otsu_scale.toFixed(2)" :step="0.01" :min="0.1" :max="2.0"
-          <q-slider color="primary" label-always v-model="hole_size_percent" :label-value="'Hole Size: ' + (hole_size_percent * 100).toFixed(4) + '%'" :step="0.0005" :min="0.0005" :max="0.5" :disable="!is_editing_mask"></q-slider>
+                    :disable="!is_editing_mask"></q-slider>
-          <q-slider color="primary" label-always v-model="smoothing_level" :label-value="'Smoothing: ' + smoothing_level + 'px'" :step="2" :min="1" :max="9" :disable="!is_editing_mask"></q-slider>
+                <q-slider color="primary" label-always v-model="noise_size_percent"
-        </div>
+                    :label-value="'Noise Size: ' + (noise_size_percent * 100).toFixed(3) + '%'" :step="0.001"
-        <div class="q-mt-md">
+                    :min="0.001" :max="0.5" :disable="!is_editing_mask"></q-slider>
-            <p class="text-caption text-center">Adjust slice overlay transparency:</p>
+                <q-slider color="primary" label-always v-model="hole_size_percent"
-            <q-slider color="black" v-model="imgTrans" :min="0.0" :max="1.0" :step="0.05" label-always :label-value="'Transparency: ' + imgTrans.toFixed(2)" :disable="!is_editing_mask" />
+                    :label-value="'Hole Size: ' + (hole_size_percent * 100).toFixed(4) + '%'" :step="0.0005"
-        </div>
+                    :min="0.0005" :max="0.5" :disable="!is_editing_mask"></q-slider>
+                <q-slider color="primary" label-always v-model="smoothing_level"
+                    :label-value="'Smoothing: ' + smoothing_level + 'px'" :step="2" :min="1" :max="9"
+                    :disable="!is_editing_mask"></q-slider>
+            </div>
+            <div class="q-mt-md">
+                <p class="text-caption text-center">Adjust slice overlay transparency:</p>
+                <q-slider color="black" v-model="imgTrans" :min="0.0" :max="1.0" :step="0.05" label-always
+                    :label-value="'Transparency: ' + imgTrans.toFixed(2)" :disable="!is_editing_mask" />
+            </div>
 
-        <p class="q-mt-md" :class="{'text-negative': mask_editor_warning}">{{ mask_editor_message }}</p>
+            <p class="q-mt-md" :class="{'text-negative': mask_editor_warning}">{{ mask_editor_message }}</p>
-        <div class="row">
+            <div class="row">
-            <q-btn  :loading="progress" class="q-ma-sm btn-style" :disable="!is_editing_mask" v-on:click="btn_save_final_mask = !btn_save_final_mask" padding="lg" icon="save" label="Save Final Mask"/>
+                <q-btn :loading="progress" class="q-ma-sm btn-style" :disable="!is_editing_mask"
-            <q-btn class="q-ma-sm btn-style" icon="arrow_back" label="Return" href="/" ></q-btn>
+                    v-on:click="btn_save_final_mask = !btn_save_final_mask" padding="lg" icon="save"
+                    label="Save Final Mask" />
+                <q-btn class="q-ma-sm btn-style" icon="arrow_back" label="Return" href="/"></q-btn>
+            </div>
         </div>
     </div>
-</div>
 
     <div class="row col-6">
         <!-- Right Panel: Image Displays -->
         <div id="intDivStyle-right" class="st-col col-12 st-module">
             <div class="text-h6">Mask Preview</div>
             <div class="row items-center">
-                <q-btn icon="arrow_back" class="q-my-sm btn-style" v-on:click="btn_img_minus=true" :disable="is_editing_mask"></q-btn>
+                <q-btn icon="arrow_back" class="q-my-sm btn-style" v-on:click="btn_img_minus=true"
-                <q-btn icon="arrow_forward" class="q-my-sm on-right btn-style" v-on:click="btn_img_plus=true" :disable="is_editing_mask"></q-btn>
+                    :disable="is_editing_mask"></q-btn>
+                <q-btn icon="arrow_forward" class="q-my-sm on-right btn-style" v-on:click="btn_img_plus=true"
+                    :disable="is_editing_mask"></q-btn>
             </div>
 
             <div v-if="!show_verification_plot" class="text-center text-grey q-pa-xl">
@@ -64,9 +81,11 @@
             </div>
 
             <div v-if="show_verification_plot">
-                <plotly :data="plotdata_verify" :layout="plotlayout_verify" class="q-pa-none q-ma-none pixelated-plot"></plotly>
+                <plotly :data="plotdata_verify" :layout="plotlayout_verify" class="q-pa-none q-ma-none pixelated-plot">
+                </plotly>
                 <div class="q-mt-md text-center">
-                    <q-btn color="primary" class="btn-style" label="Edit Manually" v-on:click="btn_edit_manually=true" />
+                    <q-btn color="primary" class="btn-style" label="Edit Manually"
+                        v-on:click="btn_edit_manually=true" />
                 </div>
             </div>
             <p class="text-center" v-html="msgimg"></p>
@@ -87,29 +106,39 @@
             <div class="q-pa-md q-gutter-sm row justify-center items-center">
                 <!-- Tool Selection -->
                 <q-btn-group>
-                    <q-btn label="Brush" v-on:click="current_tool = 'brush'" :color="current_tool === 'brush' ? 'primary' : 'white'" text-color="black"/>
+                    <q-btn label="Brush" v-on:click="current_tool = 'brush'"
-                    <q-btn label="Eraser" v-on:click="current_tool = 'eraser'" :color="current_tool === 'eraser' ? 'primary' : 'white'" text-color="black"/>
+                        :color="current_tool === 'brush' ? 'primary' : 'white'" text-color="black" />
-                    <q-btn label="Bucket" v-on:click="current_tool = 'bucket'" :color="current_tool === 'bucket' ? 'primary' : 'white'" text-color="black"/>
+                    <q-btn label="Eraser" v-on:click="current_tool = 'eraser'"
-                    <q-btn label="Drag" v-on:click="current_tool = 'drag'" :color="current_tool === 'drag' ? 'primary' : 'white'" text-color="black"/>
+                        :color="current_tool === 'eraser' ? 'primary' : 'white'" text-color="black" />
+                    <q-btn label="Bucket" v-on:click="current_tool = 'bucket'"
+                        :color="current_tool === 'bucket' ? 'primary' : 'white'" text-color="black" />
+                    <q-btn label="Drag" v-on:click="current_tool = 'drag'"
+                        :color="current_tool === 'drag' ? 'primary' : 'white'" text-color="black" />
                 </q-btn-group>
 
                 <q-separator vertical inset />
 
                 <!-- Tool Settings -->
-                <q-input dense filled label="Brush Size" type="number" v-model.number="brush_size" style="max-width: 120px" :min="1"/>
+                <q-input dense filled label="Brush Size" type="number" v-model.number="brush_size"
-                <q-input dense filled label="Zoom" type="number" step="0.1" v-model.number="editor_scale" style="max-width: 120px" :min="0.1"/>
+                    style="max-width: 120px" :min="1" />
+                <q-input dense filled label="Zoom" type="number" step="0.1" v-model.number="editor_scale"
+                    style="max-width: 120px" :min="0.1" />
 
                 <q-separator vertical inset />
 
                 <!-- Image Manipulation -->
-                <q-input dense filled label="Rotate (deg)" type="number" v-model.number="rotate_degrees" style="max-width: 150px" :step="90"/>
+                <q-input dense filled label="Rotate (deg)" type="number" v-model.number="rotate_degrees"
-                <q-select dense filled label="Flip" v-model="flip_direction" :options="['horizontal', 'vertical']" style="max-width: 150px"/>
+                    style="max-width: 150px" :step="90" />
-                <q-btn label="Apply Flip" v-on:click="btn_flip_mask = !btn_flip_mask"/>
+                <q-select dense filled label="Flip" v-model="flip_direction" :options="['horizontal', 'vertical']"
+                    style="max-width: 150px" />
+                <q-btn label="Apply Flip" v-on:click="btn_flip_mask = !btn_flip_mask" />
             </div>
 
             <!-- Drag Tool Controls -->
-            <div v-if="current_tool === 'drag'" class="q-pa-sm q-gutter-sm row justify-center items-center" style="border: 1px solid #ccc; border-radius: 4px;">
+            <div v-if="current_tool === 'drag'" class="q-pa-sm q-gutter-sm row justify-center items-center"
-                <q-input dense filled label="Move Pixels" type="number" v-model.number="move_pixels" style="max-width: 120px" :min="1"/>
+                style="border: 1px solid #ccc; border-radius: 4px;">
+                <q-input dense filled label="Move Pixels" type="number" v-model.number="move_pixels"
+                    style="max-width: 120px" :min="1" />
                 <q-btn icon="arrow_upward" v-on:click="move_mask_payload = { direction: 'up', t: Date.now() }" />
                 <q-btn icon="arrow_downward" v-on:click="move_mask_payload = { direction: 'down', t: Date.now() }" />
                 <q-btn icon="arrow_back" v-on:click="move_mask_payload = { direction: 'left', t: Date.now() }" />
@@ -118,23 +147,18 @@
 
             <!-- Interactive Canvas -->
             <div class="row justify-center q-pa-md">
-              <div style="position: relative; display: inline-block; overflow: auto; max-width: 100%;"> <!-- Added for scrolling very large images -->
+                <div style="position: relative; display: inline-block; overflow: auto; max-width: 100%;">
-                  <canvas id="maskCanvas" 
+                    <!-- Added for scrolling very large images -->
-                          :width="canvas_width" 
+                    <canvas id="maskCanvas" :width="canvas_width" :height="canvas_height" :style="{
-                          :height="canvas_height"
-                          :style="{
                               width: canvas_width * editor_scale + 'px', 
                               height: canvas_height * editor_scale + 'px',
                               border: '2px solid #ccc',
                               cursor: current_tool === 'drag' ? 'move' : 'crosshair',
                               imageRendering: 'pixelated'
-                          }"
+                          }" v-on:mousedown.prevent="startDrawing" v-on:mousemove.prevent="draw"
-                          v-on:mousedown.prevent="startDrawing"
+                        v-on:mouseup.prevent="stopDrawing()" v-on:mouseleave.prevent="stopDrawing()">
-                          v-on:mousemove.prevent="draw" 
+                    </canvas>
-                          v-on:mouseup.prevent="stopDrawing()"
+                </div>
-                          v-on:mouseleave.prevent="stopDrawing()">
-                  </canvas>
-              </div>
             </div>
         </q-card-section>
 

src/ImageProcessing.jl

@@ -1,16 +1,80 @@
-module ImageProcessing
+# src/ImageProcessing
 
 using Images
 using ImageBinarization
 using ImageMorphology
 using ImageComponentAnalysis
+using Colors # For converting to grayscale
 
 export process_image_pipeline
+export load_and_prepare_mask # Export the new function
+
+"""
+    load_and_prepare_mask(mask_path::String, target_dims::Tuple{Int, Int})
+
+Loads a PNG image mask, converts it to a binary (Boolean) matrix,
+and resizes it to the specified `target_dims`. White pixels in the mask
+are considered `true` (part of the ROI), and black pixels are `false`.
+
+# Arguments
+- `mask_path`: Absolute path to the PNG mask file.
+- `target_dims`: A tuple `(width, height)` representing the desired output dimensions.
+
+# Returns
+- A `BitMatrix` of `target_dims` where `true` indicates the ROI.
+"""
+function load_and_prepare_mask(mask_path::String, target_dims::Tuple{Int, Int})
+    if !isfile(mask_path)
+        error("Mask file not found: $(mask_path)")
+    end
+
+    # Load the image
+    img = Images.load(mask_path)
+
+    # Convert to grayscale if it's a color image
+    gray_img = Gray.(img)
+
+    # Binarize using Otsu's method - white regions become 'true'
+    binary_img = binarize(gray_img, Otsu())
+
+    # Resize to target dimensions
+    resized_img = imresize(binary_img, (target_dims[2], target_dims[1]))
+
+    # Ensure the output is a BitMatrix, as imresize can change the type
+    return resized_img .> 0.5
+end
 
 # ===================================================================
 # CORE PROCESSING PIPELINE
 # ===================================================================
 
+"""
+    process_image_pipeline(gray_img; otsu_scale=1.0, noise_size_percent=0.1, hole_size_percent=0.05, smoothing=2)
+
+Applies a multi-step image processing pipeline to a grayscale image to segment regions of interest.
+
+The pipeline consists of:
+1.  **Binarization**: An adjusted Otsu's threshold is used to create a binary image.
+2.  **Noise Removal**: Small white regions (noise) are removed using an area opening operation.
+3.  **Hole Filling**: Small black regions (holes) within larger objects are filled.
+4.  **Edge Smoothing**: The edges of the final regions are smoothed using a morphological closing operation.
+
+# Arguments
+- `gray_img`: The input grayscale image (`Matrix{<:Gray}`).
+
+# Keyword Arguments
+- `otsu_scale`: A factor to scale the automatically determined Otsu threshold. Values > 1.0 make the threshold stricter (less white), < 1.0 make it more lenient (more white). Default: `1.0`.
+- `noise_size_percent`: The percentage of the total image area used as a threshold to remove small white noise components. Default: `0.1`.
+- `hole_size_percent`: The percentage of the total image area used as a threshold to fill black holes in white components. Default: `0.05`.
+- `smoothing`: The size of the kernel for the final edge smoothing (closing) operation. Default: `2`.
+
+# Returns
+- A tuple containing four images representing the intermediate steps of the pipeline:
+    1.  `binary_img`: The result of the initial binarization.
+    2.  `noise_removed_img`: The image after noise removal.
+    3.  `holes_filled_img`: The image after filling holes.
+    4.  `smoothed_img`: The final smoothed image.
+"""
 function process_image_pipeline(gray_img;
                                 otsu_scale=1.0,
                                 noise_size_percent=0.1,
@@ -43,5 +107,3 @@ function process_image_pipeline(gray_img;
     # --- Return all intermediate steps for visualization ---
     return binary_img, noise_removed_img, holes_filled_img, smoothed_img
 end
-
-end

src/MSIData.jl

@@ -143,6 +143,45 @@ mutable struct MSIData
     end
 end
 
+"""
+    get_masked_spectrum_indices(data::MSIData, mask_matrix::BitMatrix) -> Set{Int}
+
+Converts a 2D boolean mask matrix (e.g., loaded from a PNG) into a `Set` of linear
+spectrum indices that fall within the `true` regions of the mask.
+
+# Arguments
+- `data`: The `MSIData` object containing the `coordinate_map`.
+- `mask_matrix`: A `BitMatrix` where `true` indicates a pixel is part of the ROI.
+
+# Returns
+- A `Set{Int}` containing the linear indices of spectra within the mask.
+"""
+function get_masked_spectrum_indices(data::MSIData, mask_matrix::BitMatrix)
+    if data.coordinate_map === nothing
+        error("Coordinate map not available. Cannot apply mask to non-imaging data.")
+    end
+    
+    width, height = data.image_dims
+    mask_height, mask_width = size(mask_matrix)
+
+    if mask_width != width || mask_height != height
+        error("Mask dimensions ($(mask_width)x$(mask_height)) do not match image dimensions ($(width)x$(height)).")
+    end
+
+    masked_indices = Set{Int}()
+    for y in 1:height
+        for x in 1:width
+            if mask_matrix[y, x] # If this pixel is part of the mask
+                idx = data.coordinate_map[x, y]
+                if idx != 0 # Ensure there's an actual spectrum at this coordinate
+                    push!(masked_indices, idx)
+                end
+            end
+        end
+    end
+    return masked_indices
+end
+
 
 # --- Internal function for reading binary data --- #
 
@@ -513,7 +552,7 @@ It uses a fast, two-pass approach:
 This function is highly optimized for `.imzML` by leveraging direct binary
 reading and optimized binning logic. It is called by `get_total_spectrum`.
 """
-function get_total_spectrum_imzml(msi_data::MSIData; num_bins::Int=2000)
+function get_total_spectrum_imzml(msi_data::MSIData; num_bins::Int=2000, masked_indices::Union{Set{Int}, Nothing}=nothing)
     println("Calculating total spectrum for imzML (2-pass method)...")
     total_start_time = time_ns()
 
@@ -529,7 +568,7 @@ function get_total_spectrum_imzml(msi_data::MSIData; num_bins::Int=2000)
         println("  Pass 1: Finding global m/z range...")
         g_min_mz = Inf
         g_max_mz = -Inf
-        _iterate_spectra_fast(msi_data) do idx, mz, _
+        _iterate_spectra_fast(msi_data, masked_indices === nothing ? nothing : collect(masked_indices)) do idx, mz, _
             if !isempty(mz)
                 local_min, local_max = extrema(mz)
                 g_min_mz = min(g_min_mz, local_min)
@@ -539,7 +578,7 @@ function get_total_spectrum_imzml(msi_data::MSIData; num_bins::Int=2000)
         pass1_duration = (time_ns() - pass1_start_time) / 1e9
         if !isfinite(g_min_mz)
             @warn "Could not determine a valid m/z range for imzML. All spectra might be empty."
-            return (Float64[], Float64[])
+            return (Float64[], Float64[], 0)
         end
         
         # Use and cache the result
@@ -560,7 +599,9 @@ function get_total_spectrum_imzml(msi_data::MSIData; num_bins::Int=2000)
     # 3. Second Pass: Optimized binning
     pass2_start_time = time_ns()
     println("  Pass 2: Summing intensities into $num_bins bins...")
-    _iterate_spectra_fast(msi_data) do idx, mz, intensity
+    num_spectra_processed = 0
+    _iterate_spectra_fast(msi_data, masked_indices === nothing ? nothing : collect(masked_indices)) do idx, mz, intensity
+        num_spectra_processed += 1
         if isempty(mz)
             return
         end
@@ -594,7 +635,7 @@ function get_total_spectrum_imzml(msi_data::MSIData; num_bins::Int=2000)
     println("----------------------------------------\n")
 
     println("Total spectrum calculation complete for imzML.")
-    return (collect(mz_bins), intensity_sum)
+    return (collect(mz_bins), intensity_sum, num_spectra_processed)
 end
 
 """
@@ -608,7 +649,7 @@ It uses a two-pass approach analogous to the `imzML` implementation:
 
 This function is called by `get_total_spectrum`.
 """
-function get_total_spectrum_mzml(msi_data::MSIData; num_bins::Int=2000)
+function get_total_spectrum_mzml(msi_data::MSIData; num_bins::Int=2000, masked_indices::Union{Set{Int}, Nothing}=nothing)
     println("Calculating total spectrum for mzML (2-pass method)...")
     total_start_time = time_ns()
 
@@ -624,7 +665,7 @@ function get_total_spectrum_mzml(msi_data::MSIData; num_bins::Int=2000)
         println("  Pass 1: Finding global m/z range...")
         g_min_mz = Inf
         g_max_mz = -Inf
-        _iterate_spectra_fast(msi_data) do idx, mz, intensity
+        _iterate_spectra_fast(msi_data, masked_indices === nothing ? nothing : collect(masked_indices)) do idx, mz, intensity
             if !isempty(mz)
                 local_min, local_max = extrema(mz)
                 g_min_mz = min(g_min_mz, local_min)
@@ -635,7 +676,7 @@ function get_total_spectrum_mzml(msi_data::MSIData; num_bins::Int=2000)
 
         if !isfinite(g_min_mz)
             @warn "Could not determine a valid m/z range for mzML. All spectra might be empty."
-            return (Float64[], Float64[])
+            return (Float64[], Float64[], 0)
         end
 
         # Use and cache the result
@@ -658,7 +699,9 @@ function get_total_spectrum_mzml(msi_data::MSIData; num_bins::Int=2000)
     inv_bin_step = 1.0 / bin_step  # Precompute reciprocal
     min_mz = global_min_mz
 
-    _iterate_spectra_fast(msi_data) do idx, mz, intensity
+    num_spectra_processed = 0
+    _iterate_spectra_fast(msi_data, masked_indices === nothing ? nothing : collect(masked_indices)) do idx, mz, intensity
+        num_spectra_processed += 1
         if isempty(mz)
             return
         end
@@ -687,7 +730,7 @@ function get_total_spectrum_mzml(msi_data::MSIData; num_bins::Int=2000)
     println("----------------------------------------\n")
 
     println("Total spectrum calculation complete for mzML.")
-    return (collect(mz_bins), intensity_sum)
+    return (collect(mz_bins), intensity_sum, num_spectra_processed)
 end
 
 """
@@ -699,11 +742,18 @@ This function dispatches to a specialized implementation based on the file type
 
 Returns a tuple containing two vectors: the binned m/z axis and the summed intensities.
 """
-function get_total_spectrum(msi_data::MSIData; num_bins::Int=2000)
+function get_total_spectrum(msi_data::MSIData; num_bins::Int=2000, mask_path::Union{String, Nothing}=nothing)::Tuple{Vector{Float64}, Vector{Float64}, Int}
+    local masked_indices::Union{Set{Int}, Nothing} = nothing
+    if mask_path !== nothing
+        mask_matrix = load_and_prepare_mask(mask_path, msi_data.image_dims)
+        masked_indices = get_masked_spectrum_indices(msi_data, mask_matrix)
+        println("Calculating total spectrum for masked region (mask from: $(mask_path))")
+    end
+
     if msi_data.source isa ImzMLSource
-        return get_total_spectrum_imzml(msi_data, num_bins=num_bins)
+        return get_total_spectrum_imzml(msi_data, num_bins=num_bins, masked_indices=masked_indices)
     else # MzMLSource
-        return get_total_spectrum_mzml(msi_data, num_bins=num_bins)
+        return get_total_spectrum_mzml(msi_data, num_bins=num_bins, masked_indices=masked_indices)
     end
 end
 
@@ -715,14 +765,15 @@ This is effectively the total spectrum divided by the number of spectra.
 
 Returns a tuple containing two vectors: the binned m/z axis and the averaged intensities.
 """
-function get_average_spectrum(msi_data::MSIData; num_bins::Int=2000)
+function get_average_spectrum(msi_data::MSIData; num_bins::Int=2000, mask_path::Union{String, Nothing}=nothing)::Tuple{Vector{Float64}, Vector{Float64}}
-    mz_bins, intensity_sum = get_total_spectrum(msi_data, num_bins=num_bins)
+    mz_bins, intensity_sum, num_spectra_processed = get_total_spectrum(msi_data, num_bins=num_bins, mask_path=mask_path)
 
-    if isempty(intensity_sum)
+    if isempty(intensity_sum) || num_spectra_processed == 0
-        return (mz_bins, intensity_sum)
+        @warn "Cannot calculate average spectrum: no spectra were processed."
+        return (mz_bins, Float64[])
     end
-    num_spectra = length(msi_data.spectra_metadata)
+
-    average_intensity = intensity_sum ./ num_spectra
+    average_intensity = intensity_sum ./ num_spectra_processed
     
     return (mz_bins, average_intensity)
 end
@@ -855,13 +906,16 @@ is significantly faster for bulk processing than reading spectra one by one.
 - `data`: The `MSIData` object.
 - `source`: The `ImzMLSource`.
 """
-function _iterate_uncompressed_fast(f::Function, data::MSIData, source::ImzMLSource)
+function _iterate_uncompressed_fast(f::Function, data::MSIData, source::ImzMLSource, indices_to_iterate::Union{AbstractVector{Int}, Nothing})
     # Optimized path for uncompressed data using buffer reuse
     max_points = maximum(meta -> meta.mz_asset.encoded_length, data.spectra_metadata)
     mz_buffer = Vector{source.mz_format}(undef, max_points)
     int_buffer = Vector{source.intensity_format}(undef, max_points)
 
-    for i in 1:length(data.spectra_metadata)
+    # Determine which indices to iterate over
+    spectrum_indices = (indices_to_iterate === nothing) ? (1:length(data.spectra_metadata)) : indices_to_iterate
+
+    for i in spectrum_indices
         meta = data.spectra_metadata[i]
         nPoints = meta.mz_asset.encoded_length
 
@@ -904,10 +958,14 @@ and will be slower and allocate more memory than `_iterate_uncompressed_fast`.
 - `data`: The `MSIData` object.
 - `source`: The `ImzMLSource`.
 """
-function _iterate_compressed_fast(f::Function, data::MSIData, source::ImzMLSource)
+function _iterate_compressed_fast(f::Function, data::MSIData, source::ImzMLSource, indices_to_iterate::Union{AbstractVector{Int}, Nothing})
     # Path for datasets containing at least one compressed spectrum.
     # This path reads and decompresses each spectrum individually.
-    for i in 1:length(data.spectra_metadata)
+    
+    # Determine which indices to iterate over
+    spectrum_indices = (indices_to_iterate === nothing) ? (1:length(data.spectra_metadata)) : indices_to_iterate
+
+    for i in spectrum_indices
         meta = data.spectra_metadata[i]
         
         if meta.mz_asset.encoded_length == 0 && meta.int_asset.encoded_length == 0
@@ -937,7 +995,7 @@ This function acts as a dispatcher:
 - If all spectra are uncompressed, it uses a highly optimized path that reuses pre-allocated 
   buffers to minimize memory allocations and overhead.
 """
-function _iterate_spectra_fast_impl(f::Function, data::MSIData, source::ImzMLSource)
+function _iterate_spectra_fast_impl(f::Function, data::MSIData, source::ImzMLSource, indices_to_iterate::Union{AbstractVector{Int}, Nothing})
     if isempty(data.spectra_metadata)
         return
     end
@@ -947,9 +1005,9 @@ function _iterate_spectra_fast_impl(f::Function, data::MSIData, source::ImzMLSou
                         data.spectra_metadata)
     
     if any_compressed
-        _iterate_compressed_fast(f, data, source)
+        _iterate_compressed_fast(f, data, source, indices_to_iterate)
     else
-        _iterate_uncompressed_fast(f, data, source)
+        _iterate_uncompressed_fast(f, data, source, indices_to_iterate)
     end
 end
 
@@ -961,12 +1019,16 @@ through each spectrum, decodes the Base64 data on the fly, and calls the
 provided function. It bypasses the `GetSpectrum` cache to avoid storing
 all decoded spectra in memory.
 """
-function _iterate_spectra_fast_impl(f::Function, data::MSIData, source::MzMLSource)
+function _iterate_spectra_fast_impl(f::Function, data::MSIData, source::MzMLSource, indices_to_iterate::Union{AbstractVector{Int}, Nothing})
     # This implementation is for mzML. It iterates through each spectrum,
     # decodes it, and then calls the function `f`. It's less performant than
     # the imzML version because we cannot pre-allocate buffers of a known size,
     # but it is correct and still faster than using GetSpectrum due to bypassing the cache.
-    for i in 1:length(data.spectra_metadata)
+    
+    # Determine which indices to iterate over
+    spectrum_indices = (indices_to_iterate === nothing) ? (1:length(data.spectra_metadata)) : indices_to_iterate
+
+    for i in spectrum_indices
         meta = data.spectra_metadata[i]
         
         mz = read_binary_vector(source.file_handle, meta.mz_asset)
@@ -986,7 +1048,7 @@ It dispatches to a specialized implementation based on the data source type
 - `f`: A function to call for each spectrum, with signature `f(index, mz, intensity)`.
 - `data`: The `MSIData` object.
 """
-function _iterate_spectra_fast(f::Function, data::MSIData)
+function _iterate_spectra_fast(f::Function, data::MSIData, indices_to_iterate::Union{AbstractVector{Int}, Nothing}=nothing)
     # Dispatch to the correct implementation based on the source type
-    _iterate_spectra_fast_impl(f, data, data.source)
+    _iterate_spectra_fast_impl(f, data, data.source, indices_to_iterate)
 end

src/MSI_src.jl

@@ -14,7 +14,10 @@ export OpenMSIData,
        get_average_spectrum, 
        LoadMzml, 
        precompute_analytics, 
-       process_spectrum
+       process_spectrum,
+       generate_colorbar_image,
+       process_image_pipeline,
+       load_and_prepare_mask
 
 # Export the public Preprocessing API
 export FeatureMatrix,
@@ -43,6 +46,7 @@ include("mzML.jl")
 include("imzML.jl")
 include("MzmlConverter.jl")
 include("Preprocessing.jl")
+include("ImageProcessing.jl")
 
 
 # --- Main Entry Point --- #

src/imzML.jl

@@ -744,10 +744,17 @@ This is a performant function that iterates through spectra once.
 # Returns
 - A `Matrix{Float64}` representing the intensity slice.
 """
-function get_mz_slice(data::MSIData, mass::Real, tolerance::Real)
+function get_mz_slice(data::MSIData, mass::Real, tolerance::Real; mask_path::Union{String, Nothing}=nothing)
     width, height = data.image_dims
     slice_matrix = zeros(Float64, height, width)
 
+    local masked_indices::Union{Set{Int}, Nothing} = nothing
+    if mask_path !== nothing
+        mask_matrix = load_and_prepare_mask(mask_path, (width, height))
+        masked_indices = get_masked_spectrum_indices(data, mask_matrix)
+        println("Applying mask from: $(mask_path), found $(length(masked_indices)) spectra in ROI.")
+    end
+
     # INTELLIGENT LOADING: Ensure analytics are computed for filtering.
     if data.spectrum_stats_df === nothing || !hasproperty(data.spectrum_stats_df, :MinMZ)
         println("Per-spectrum metadata not found. Running one-time analytics computation...")
@@ -761,7 +768,9 @@ function get_mz_slice(data::MSIData, mass::Real, tolerance::Real)
 
     # 1. Find all candidate spectra first for efficient filtering
     candidate_indices = Set{Int}()
-    for i in 1:length(data.spectra_metadata)
+    indices_to_check = masked_indices === nothing ? (1:length(data.spectra_metadata)) : masked_indices
+    
+    for i in indices_to_check
         spec_min_mz = stats_df.MinMZ[i]
         spec_max_mz = stats_df.MaxMZ[i]
         if target_max >= spec_min_mz && target_min <= spec_max_mz
@@ -769,20 +778,17 @@ function get_mz_slice(data::MSIData, mass::Real, tolerance::Real)
         end
     end
 
-    println("Found $(length(candidate_indices)) candidate spectra (filtered from $(length(data.spectra_metadata)))")
+    println("Found $(length(candidate_indices)) candidate spectra (filtered from $(length(indices_to_check)) initial spectra)")
 
     # 2. Iterate using the optimized, low-allocation iterator
     results_count = 0
-    _iterate_spectra_fast(data) do idx, mz_array, intensity_array
+    _iterate_spectra_fast(data, collect(candidate_indices)) do idx, mz_array, intensity_array
-        # Process only the spectra that are candidates
+        meta = data.spectra_metadata[idx]
-        if idx in candidate_indices
+        intensity = find_mass(mz_array, intensity_array, mass, tolerance)
-            intensity = find_mass(mz_array, intensity_array, mass, tolerance)
+        if intensity > 0.0
-            if intensity > 0.0
+            if 1 <= meta.x <= width && 1 <= meta.y <= height
-                meta = data.spectra_metadata[idx]
+                slice_matrix[meta.y, meta.x] = intensity
-                if 1 <= meta.x <= width && 1 <= meta.y <= height
+                results_count += 1
-                    slice_matrix[meta.y, meta.x] = intensity
-                    results_count += 1
-                end
             end
         end
     end
@@ -802,7 +808,7 @@ This is a highly performant function that iterates through the full dataset only
 # Returns
 - A `Dict{Real, Matrix{Float64}}` mapping each mass to its intensity slice matrix.
 """
-function get_multiple_mz_slices(data::MSIData, masses::AbstractVector{<:Real}, tolerance::Real)
+function get_multiple_mz_slices(data::MSIData, masses::AbstractVector{<:Real}, tolerance::Real; mask_path::Union{String, Nothing}=nothing)
     width, height = data.image_dims
     
     # 1. Initialize a dictionary to hold the output slice matrices
@@ -811,6 +817,13 @@ function get_multiple_mz_slices(data::MSIData, masses::AbstractVector{<:Real}, t
         slice_dict[mass] = zeros(Float64, height, width)
     end
 
+    local masked_indices::Union{Set{Int}, Nothing} = nothing
+    if mask_path !== nothing
+        mask_matrix = load_and_prepare_mask(mask_path, (width, height))
+        masked_indices = get_masked_spectrum_indices(data, mask_matrix)
+        println("Applying mask from: $(mask_path), found $(length(masked_indices)) spectra in ROI.")
+    end
+
     # 2. Ensure analytics are computed for filtering.
     if data.spectrum_stats_df === nothing || !hasproperty(data.spectrum_stats_df, :MinMZ)
         println("Per-spectrum metadata not found. Running one-time analytics computation...")
@@ -820,12 +833,13 @@ function get_multiple_mz_slices(data::MSIData, masses::AbstractVector{<:Real}, t
     println("Filtering candidate spectra for $(length(masses)) m/z values...")
     stats_df = data.spectrum_stats_df
     candidate_indices = Set{Int}()
+    indices_to_check = masked_indices === nothing ? (1:length(data.spectra_metadata)) : masked_indices
 
     # 3. Find all spectra that could contain *any* of the requested masses.
     for mass in masses
         target_min = mass - tolerance
         target_max = mass + tolerance
-        for i in 1:length(data.spectra_metadata)
+        for i in indices_to_check
             # If already a candidate, no need to check again
             if i in candidate_indices
                 continue
@@ -841,20 +855,17 @@ function get_multiple_mz_slices(data::MSIData, masses::AbstractVector{<:Real}, t
     println("Found $(length(candidate_indices)) total candidate spectra.")
 
     # 4. Iterate through the data a single time using the optimized iterator.
-    _iterate_spectra_fast(data) do idx, mz_array, intensity_array
+    _iterate_spectra_fast(data, collect(candidate_indices)) do idx, mz_array, intensity_array
-        # Process only the spectra that are candidates
+        meta = data.spectra_metadata[idx]
-        if idx in candidate_indices
+        # For this single spectrum, check all masses of interest
-            meta = data.spectra_metadata[idx]
+        for mass in masses
-            # For this single spectrum, check all masses of interest
+            # Check if this spectrum's range actually covers the current mass
-            for mass in masses
+            # This is a finer-grained check than the initial filtering
-                # Check if this spectrum's range actually covers the current mass
+            if !isempty(mz_array) && (mass + tolerance) >= first(mz_array) && (mass - tolerance) <= last(mz_array)
-                # This is a finer-grained check than the initial filtering
+                intensity = find_mass(mz_array, intensity_array, mass, tolerance)
-                if !isempty(mz_array) && (mass + tolerance) >= first(mz_array) && (mass - tolerance) <= last(mz_array)
+                if intensity > 0.0
-                    intensity = find_mass(mz_array, intensity_array, mass, tolerance)
+                    if 1 <= meta.x <= width && 1 <= meta.y <= height
-                    if intensity > 0.0
+                        slice_dict[mass][meta.y, meta.x] = intensity
-                        if 1 <= meta.x <= width && 1 <= meta.y <= height
-                            slice_dict[mass][meta.y, meta.x] = intensity
-                        end
                     end
                 end
             end
@@ -878,30 +889,15 @@ Generates and saves a plot of a single image slice for a given m/z value.
 This function closely imitates the logic of the original `GetSlice` but uses
 the modern `MSIData` access patterns and robust peak finding.
 """
-function plot_slice(msi_data::MSIData, mass::Real, tolerance::Real, output_dir::String; stage_name="slice_mz_$(mass)", bins=256)
+function plot_slice(msi_data::MSIData, mass::Real, tolerance::Real, output_dir::String; 
+                    stage_name="slice_mz_$(mass)", bins=256, mask_path::Union{String, Nothing}=nothing)
     
-    # 1. Create an empty image for the slice, with dimensions matching plotting expectations
+    # 1. Generate the slice, applying the mask if provided.
-    width, height = msi_data.image_dims
-    slice_matrix = zeros(Float64, height, width)
-
-    # 2. Iterate through each spectrum to build the slice
     println("Generating slice for m/z $mass...")
-    _iterate_spectra_fast(msi_data) do spec_idx, mz_array, intensity_array
+    slice_matrix = get_mz_slice(msi_data, mass, tolerance, mask_path=mask_path)
-        meta = msi_data.spectra_metadata[spec_idx]
-        
-        # Find the peak intensity using the modern, robust find_mass
-        intensity = find_mass(mz_array, intensity_array, mass, tolerance)
-        
-        if intensity > 0.0
-            # Populate the matrix using (y, x) indexing
-            if 1 <= meta.x <= width && 1 <= meta.y <= height
-                slice_matrix[meta.y, meta.x] = intensity
-            end
-        end
-    end
     println("Slice generation complete.")
 
-    # 3. Plot the resulting slice matrix using CairoMakie
+    # 2. Plot the resulting slice matrix using CairoMakie
     println("Plotting slice...")
     
     fig = Figure(size = (600, 500))
@@ -925,7 +921,7 @@ function plot_slice(msi_data::MSIData, mass::Real, tolerance::Real, output_dir::
     Colorbar(fig[1, 2], hm, label="Intensity")
     colgap!(fig.layout, 5)
 
-    # 4. Save the plot
+    # 3. Save the plot
     mkpath(output_dir)
     filename = "$(stage_name).png"
     save_path = joinpath(output_dir, filename)
@@ -957,9 +953,9 @@ is used to exclude outliers before normalization.
 # Returns
 - A tuple `(low, high)` representing the calculated lower and upper intensity bounds.
 """
-function get_outlier_thres(img::AbstractMatrix{<:Real}, prob::Real=0.98)
+function get_outlier_thres(img::AbstractVector{<:Real}, prob::Real=0.98)
     # DO NOT filter zeros. Use all pixel values like R does.
-    int_values = vec(img)
+    int_values = img
     low = minimum(int_values)
     upp = maximum(int_values)
 
@@ -1001,6 +997,10 @@ function get_outlier_thres(img::AbstractMatrix{<:Real}, prob::Real=0.98)
     return (low, actual_threshold)
 end
 
+function get_outlier_thres(img::AbstractMatrix{<:Real}, prob::Real=0.98)
+    return get_outlier_thres(vec(img), prob)
+end
+
 """
     set_pixel_depth(img, bounds, depth)
 
@@ -1063,11 +1063,26 @@ within these bounds.
 # Returns
 - A new image matrix with intensities quantized to the specified depth within the TrIQ bounds.
 """
-function TrIQ(pixMap::AbstractMatrix{<:Real}, depth::Integer, prob::Real=0.98)
+function TrIQ(pixMap::AbstractMatrix{<:Real}, depth::Integer, prob::Real=0.98; mask_matrix::Union{BitMatrix, Nothing}=nothing)
-    # Compute new dynamic range
+    local values_for_thres
-    bounds = get_outlier_thres(pixMap, prob)
+    if mask_matrix !== nothing
+        # Extract values only from the masked region for threshold calculation
+        values_for_thres = pixMap[mask_matrix]
+        # Filter out only zeros created by the mask.
+        filter!(x -> x != 0, values_for_thres)
+    else
+        values_for_thres = pixMap
+    end
 
-    # Set intensity dynamic range
+    # If the relevant values are empty or all zero, return a zero matrix
+    if isempty(values_for_thres) || all(iszero, values_for_thres)
+        bounds = (0.0, 0.0)
+    else
+        # Compute new dynamic range from the (potentially filtered) values
+        bounds = get_outlier_thres(values_for_thres, prob)
+    end
+    
+    # Set intensity dynamic range for the *entire* pixMap, using bounds derived from masked data
     return set_pixel_depth(pixMap, bounds, depth)
 end
 
@@ -1107,8 +1122,17 @@ Linearly scales the intensity values in a slice to a specified number of levels.
 The output is an array of `UInt8` values.
 This is a modernized version of `IntQuantCl`.
 """
-function quantize_intensity(slice::AbstractMatrix{<:Real}, levels::Integer=256)
+function quantize_intensity(slice::AbstractMatrix{<:Real}, levels::Integer=256; mask_matrix::Union{BitMatrix, Nothing}=nothing)
-    max_val = maximum(slice)
+    local max_val
+    if mask_matrix !== nothing
+        masked_slice = slice[mask_matrix]
+        # Filter out zeros that might have been introduced by masking, if any
+        filter!(x -> x != 0, masked_slice)
+        max_val = isempty(masked_slice) ? 0.0 : maximum(masked_slice)
+    else
+        max_val = maximum(slice)
+    end
+
     if max_val <= 0
         return zeros(UInt8, size(slice))
     end
@@ -1208,7 +1232,7 @@ function display_statistics(slices::AbstractArray{<:AbstractMatrix{<:Real}, 2},
     stats_to_calc = Dict{String, Function}("Mean" => mean)
     all_dfs = Dict{String, DataFrame}()
 
-    for (stat_name, stat_func) in stats_to_to_calc
+    for (stat_name, stat_func) in stats_to_calc
         # Pre-allocate matrix for better performance
         stat_matrix = zeros(Float64, n_files, n_masses)
         
@@ -1498,3 +1522,80 @@ end
 
 # Viridis color palette (256 colors) - defined as constant
 const ViridisPalette = generate_palette(ColorSchemes.viridis)
+
+function generate_colorbar_image(slice_data::AbstractMatrix, color_levels::Int, output_path::String; use_triq::Bool=false, triq_prob::Float64=0.98, mask_path::Union{String, Nothing}=nothing)
+    # 1. Determine bounds based on whether TrIQ is used and if a mask is applied
+    local data_for_bounds
+    if mask_path !== nothing
+        height, width = size(slice_data)
+        mask_matrix = load_and_prepare_mask(mask_path, (width, height))
+        # Extract only the non-zero values within the mask to accurately calculate the range
+        data_for_bounds = slice_data[mask_matrix]
+        filter!(x -> x > 0, data_for_bounds)
+    else
+        data_for_bounds = vec(slice_data)
+    end
+
+    # Handle cases where data_for_bounds might be empty after filtering
+    if isempty(data_for_bounds)
+        min_val, max_val = 0.0, 1.0 # Default range if no data in ROI
+    else
+        min_val, max_val = if use_triq
+            get_outlier_thres(data_for_bounds, triq_prob)
+        else
+            extrema(data_for_bounds)
+        end
+    end
+
+    # 2. Replicate the tick calculation logic from plot_slices
+    bins = color_levels
+    levels = range(min_val, stop=max_val, length=bins + 1)
+    level_range = levels[end] - levels[1]
+    
+    if level_range == 0
+        levels = range(min_val - 0.1, stop=max_val + 0.1, length=bins + 1)
+        level_range = 0.2
+    end
+    
+    exponent = level_range > 0 ? floor(log10(level_range)) / 3 : 0
+    scale = 10^(3 * exponent)
+    scaled_levels = levels ./ scale
+    
+    format_num = level_range > 0 ? floor(log10(level_range)) % 3 : 0
+    labels = if format_num == 0
+        [ @sprintf("%3.2f", lvl) for lvl in scaled_levels]
+    elseif format_num == 1
+        [ @sprintf("%3.2f", lvl) for lvl in scaled_levels]
+    else
+        [ @sprintf("%3.2f", lvl) for lvl in scaled_levels]
+    end
+    
+    divisors = 2:7
+    remainders = (bins - 1) .% divisors
+    best_divisor = divisors[findlast(x -> x == minimum(remainders), remainders)]
+    tick_indices = round.(Int, range(1, stop=bins + 1, length=best_divisor + 1))
+
+    if !(1 in tick_indices)
+        pushfirst!(tick_indices, 1)
+    end
+    if !((bins + 1) in tick_indices)
+        push!(tick_indices, bins + 1)
+    end
+    unique!(sort!(tick_indices))
+    
+    tick_positions = levels[tick_indices]
+    tick_labels = labels[tick_indices]
+
+    # 3. Create and save the colorbar image
+    fig = Figure(size=(150, 250))
+    Colorbar(fig[1, 1], 
+        colormap=cgrad(:viridis, bins, categorical=true),
+        # limits=(min_val, max_val),
+        limits=(levels[1], levels[end]),
+        label=(scale == 1 ? "Intensity" : "Intensity ×10^$(round(Int, 3 * exponent))"),
+        ticks=(tick_positions, tick_labels),
+        labelsize=20,
+        ticklabelsize=16
+    )
+    save(output_path, fig)
+end