fix eval and plotting refactor

.gitignore

@@ -3,3 +3,5 @@
 temp/
 *.png
 *.pdf
+*.hash
+*.npz

README.md

@@ -7,7 +7,7 @@
 		<a href="https://keras.io/"><img src="https://img.shields.io/badge/Keras-EE4C2C?logo=keras&logoColor=white" alt="Keras"></a>
 	</p>
 		<h3>
-			<span style="display:inline-block; margin: 0 20px;">
+			<span style="display:inline-block; margin: 0 40px;">
 				<a href="https://example.com/tristan-stevens">Tristan Stevens</a>
 			</span>
 			<span style="display:inline-block; margin: 0 20px;">

eval.py

@@ -2,16 +2,17 @@ import warnings
 from glob import glob
 from pathlib import Path
 
-import matplotlib.pyplot as plt
 import numpy as np
 import torch
 import tyro
 from PIL import Image
 from scipy.ndimage import binary_erosion, distance_transform_edt
 from scipy.stats import ks_2samp
+from zea import log
 from zea.io_lib import load_image
 
 import fid_score
+from plots import plot_metrics
 
 
 def calculate_fid_score(denoised_image_dirs, ground_truth_dir):
@@ -207,64 +208,38 @@ def calculate_final_score(aggregates):
         return 0
 
 
-def plot_metrics(metrics, limits, out_path):
-    plt.style.use("seaborn-v0_8-darkgrid")
-    fig, axes = plt.subplots(1, len(metrics), figsize=(7.2, 2.7), dpi=600)
-    colors = ["#0057b7", "#ffb300", "#008744", "#d62d20"]
-    # Arrow direction: ↑ for up, ↓ for down
-    metric_labels = {
-        "CNR": r"CNR $\uparrow$",
-        "gCNR": r"gCNR $\uparrow$",
-        "KS_A": r"KS$_{septum}$ $\downarrow$",
-        "KS_B": r"KS$_{ventricle}$ $\uparrow$",
-    }
-    for idx, (ax, (name, values)) in enumerate(zip(axes, metrics.items())):
-        ax.hist(
-            values,
-            bins=30,
-            color=colors[idx % len(colors)],
-            alpha=0.85,
-            edgecolor="black",
-            linewidth=0.7,
-        )
-        ax.set_xlabel(metric_labels.get(name, name), fontsize=11)
-        ax.set_ylabel("Count", fontsize=10)
-        # Draw limits
-        if name in limits:
-            for lim in limits[name]:
-                ax.axvline(lim, color="crimson", linestyle="--", lw=1.2)
-        ax.spines["top"].set_visible(False)
-        ax.spines["right"].set_visible(False)
-        ax.tick_params(axis="both", which="major", labelsize=9)
-    fig.tight_layout(pad=1.5)
-    fig.savefig(out_path, bbox_inches="tight", dpi=600)
-    plt.close(fig)
-
-
-def main(folder: str, roi_folder: str, reference_folder: str):
+def main(folder: str, noisy_folder: str, roi_folder: str, reference_folder: str):
+    """Evaluate the dehazing algorithm.
+
+    Args:
+        folder (str): Path to the folder containing the dehazed images.
+            Used for evaluating all metrics.
+        noisy_folder (str): Path to the folder containing the noisy images.
+            Only used for KS statistics.
+        roi_folder (str): Path to the folder containing the ROI images.
+            Used for contrast and KS statistic metrics.
+        reference_folder (str): Path to the folder containing the reference images.
+            Used only for FID calculation.
+    """
     folder = Path(folder)
+    noisy_folder = Path(noisy_folder)
     roi_folder = Path(roi_folder)
     reference_folder = Path(reference_folder)
 
     folder_files = set(f.name for f in folder.glob("*.png"))
+    noisy_files = set(f.name for f in noisy_folder.glob("*.png"))
     roi_files = set(f.name for f in roi_folder.glob("*.png"))
-    ref_files = set(f.name for f in reference_folder.glob("*.png"))
 
     print(f"Found {len(folder_files)} .png files in output folder: {folder}")
+    print(f"Found {len(noisy_files)} .png files in noisy folder: {noisy_folder}")
     print(f"Found {len(roi_files)} .png files in ROI folder: {roi_folder}")
-    print(f"Found {len(ref_files)} .png files in reference folder: {reference_folder}")
 
     # Find intersection of filenames
-    common_files = sorted(folder_files & roi_files & ref_files)
-    print(f"Found {len(common_files)} images present in all folders.")
-    if len(common_files) == 0:
-        print("No matching images found in all folders. Check your folder contents.")
-        print(f"Output folder files: {sorted(folder_files)}")
-        print(f"ROI folder files: {sorted(roi_files)}")
-        print(f"Reference folder files: {sorted(ref_files)}")
-        assert len(common_files) > 0, (
-            "No matching .png files in all folders. Cannot proceed."
-        )
+    common_files = sorted(folder_files & roi_files & noisy_files)
+    print(f"Found {len(common_files)} matching images in noisy/dehazed/roi folders")
+    assert len(common_files) > 0, (
+        "No matching .png files in all folders. Cannot proceed."
+    )
 
     metrics = {"CNR": [], "gCNR": [], "KS_A": [], "KS_B": []}
     limits = {
@@ -275,28 +250,26 @@ def main(folder: str, roi_folder: str, reference_folder: str):
     }
 
     for name in common_files:
-        our_path = folder / name
+        dehazed_path = folder / name
+        noisy_path = noisy_folder / name
         roi_path = roi_folder / name
-        ref_path = reference_folder / name
-
-        assert our_path.exists(), f"Missing file in output folder: {our_path}"
-        assert roi_path.exists(), f"Missing file in ROI folder: {roi_path}"
-        assert ref_path.exists(), f"Missing file in reference folder: {ref_path}"
 
         try:
-            img = np.array(load_image(str(our_path)))
-            img_ref = np.array(load_image(str(ref_path)))
+            img_dehazed = np.array(load_image(str(dehazed_path)))
+            img_noisy = np.array(load_image(str(noisy_path)))
         except Exception as e:
             print(f"Error loading image {name}: {e}")
             continue
 
         # CNR/gCNR
-        cnr_gcnr = calculate_cnr_gcnr(img, str(roi_path))
+        cnr_gcnr = calculate_cnr_gcnr(img_dehazed, str(roi_path))
         metrics["CNR"].append(cnr_gcnr[0][0])
         metrics["gCNR"].append(cnr_gcnr[0][1])
 
         # KS statistics
-        ks_a, _, ks_b, _ = calculate_ks_statistics(img_ref, img, str(roi_path))
+        ks_a, _, ks_b, _ = calculate_ks_statistics(
+            img_noisy, img_dehazed, str(roi_path)
+        )
         metrics["KS_A"].append(ks_a)
         metrics["KS_B"].append(ks_b)
 
@@ -308,8 +281,13 @@ def main(folder: str, roi_folder: str, reference_folder: str):
     for k, (mean, std, minv, maxv) in stats.items():
         print(f"{k}: mean={mean:.3f}, std={std:.3f}, min={minv:.3f}, max={maxv:.3f}")
 
-    plot_metrics(metrics, limits, str(folder / "contrast_metrics.png"))
-    print(f"Saved metrics plot to {folder / 'contrast_metrics.png'}")
+    fig = plot_metrics(metrics, limits, "contrast_metrics.png")
+
+    path = Path("contrast_metrics.png")
+    save_kwargs = {"bbox_inches": "tight", "dpi": 300}
+    fig.savefig(path, **save_kwargs)
+    fig.savefig(path.with_suffix(".pdf"), **save_kwargs)
+    log.success(f"Metrics plot saved to {log.yellow(path)}")
 
     # Compute FID
     fid_image_paths = [str(folder / name) for name in common_files]

fid_score.py

@@ -88,6 +88,7 @@ parser.add_argument(
 )
 
 IMAGE_EXTENSIONS = {"bmp", "jpg", "jpeg", "pgm", "png", "ppm", "tif", "tiff", "webp"}
+TEMP_DIR = pathlib.Path("temp")
 
 
 class ImagePathDataset(torch.utils.data.Dataset):
@@ -277,7 +278,7 @@ def compute_statistics_of_path(path, model, batch_size, dims, device, num_worker
 
 
 def _fid_cache_paths():
-    tmp_dir = pathlib.Path("tmp")
+    tmp_dir = TEMP_DIR
     tmp_dir.mkdir(exist_ok=True)
     stats_path = tmp_dir / "fid_stats.npz"
     hash_path = tmp_dir / "fid_stats.hash"
@@ -391,7 +392,7 @@ def calculate_fid_with_cached_ground_truth(
                 continue
         return hash_md5.hexdigest()
 
-    tmp_dir = pathlib.Path("tmp")
+    tmp_dir = TEMP_DIR
     tmp_dir.mkdir(exist_ok=True)
     stats_path = tmp_dir / "fid_stats.npz"
     hash_path = tmp_dir / "fid_stats.hash"

main.py

@@ -1,9 +1,6 @@
 import copy
-import os
 from pathlib import Path
 
-os.environ["KERAS_BACKEND"] = "jax"
-
 import jax
 import keras
 import matplotlib.pyplot as plt
@@ -12,10 +9,8 @@ import scipy
 import tyro
 import zea
 from keras import ops
-from matplotlib.patches import PathPatch
-from matplotlib.path import Path as pltPath
 from PIL import Image
-from skimage import filters, measure, morphology
+from skimage import filters, morphology
 from zea import Config, init_device, log
 from zea.internal.operators import Operator
 from zea.models.diffusion import (
@@ -25,7 +20,8 @@ from zea.models.diffusion import (
 )
 from zea.tensor_ops import L2
 from zea.utils import translate
-from zea.visualize import plot_image_grid
+
+from plots import plot_batch_with_named_masks, plot_dehazed_results
 
 
 def L1(x):
@@ -476,186 +472,6 @@ def run(
     return hazy_images, pred_tissue_images, pred_haze_images, masks_out
 
 
-def add_shape_from_mask(ax, mask, **kwargs):
-    """add a shape to axis from mask array.
-
-    Args:
-        ax (plt.ax): matplotlib axis
-        mask (ndarray): numpy array with non-zero
-            shape defining the region of interest.
-    Kwargs:
-        edgecolor (str): color of the shape's edge
-        facecolor (str): color of the shape's face
-        linewidth (int): width of the shape's edge
-
-    Returns:
-        plt.ax: matplotlib axis with shape added
-    """
-    # Pad mask to ensure edge contours are found
-    padded_mask = np.pad(mask, pad_width=1, mode="constant", constant_values=0)
-    contours = measure.find_contours(padded_mask, 0.5)
-    patches = []
-    for contour in contours:
-        # Remove padding offset
-        contour -= 1
-        path = pltPath(contour[:, ::-1])
-        patch = PathPatch(path, **kwargs)
-        patches.append(ax.add_patch(patch))
-    return patches
-
-
-def plot_batch_with_named_masks(
-    images, masks_dict, mask_colors=None, titles=None, **kwargs
-):
-    """
-    Plot batch of images in rows, each column overlays a different mask from the dict.
-    Mask labels are shown as column titles. If mask name is 'per_pixel_omega', show it
-    directly with inferno colormap (no overlay).
-
-    Args:
-        images: np.ndarray, shape (batch, height, width, channels)
-        masks_dict: dict of {name: mask}, each mask shape  (batch, height, width, channels)
-        mask_colors: dict of {name: color} or None (default colors used)
-    """
-    mask_names = list(masks_dict.keys())
-    batch_size = images.shape[0]
-    default_colors = ["red", "green", "#33aaff", "yellow", "magenta", "cyan"]
-    mask_colors = mask_colors or {
-        name: default_colors[i % len(default_colors)]
-        for i, name in enumerate(mask_names)
-    }
-
-    # Prepare images for each column
-    columns = []
-    cmaps = []
-    for name in mask_names:
-        if name == "per_pixel_omega":
-            mask_np = np.array(masks_dict[name])
-            columns.append(np.squeeze(mask_np))
-            cmaps.append(["inferno"] * batch_size)
-        else:
-            columns.append(np.squeeze(images))
-            cmaps.append(["gray"] * batch_size)
-
-    # Stack columns: shape (num_columns, batch, ...)
-    all_images = np.stack(columns, axis=0)  # (num_columns, batch, ...)
-    # Rearrange to (batch, num_columns, ...)
-    all_images = (
-        np.transpose(all_images, (1, 0, 2, 3, 4))
-        if all_images.ndim == 5
-        else np.transpose(all_images, (1, 0, 2, 3))
-    )
-    # Flatten to (batch * num_columns, ...)
-    all_images = all_images.reshape(batch_size * len(mask_names), *images.shape[1:])
-
-    # Flatten cmaps for plot_image_grid in the same order as images
-    flat_cmaps = []
-    for row in range(batch_size):
-        for col in range(len(mask_names)):
-            flat_cmaps.append(cmaps[col][row])
-
-    fig, _ = plot_image_grid(
-        all_images,
-        ncols=len(mask_names),
-        remove_axis=False,
-        cmap=flat_cmaps,
-        figsize=(8, 3.3),
-        **kwargs,
-    )
-
-    # Overlay masks for non-per_pixel_omega columns
-    for col_idx, name in enumerate(mask_names):
-        if name == "per_pixel_omega":
-            continue
-        mask_np = np.array(masks_dict[name])
-        axes = fig.axes[col_idx : batch_size * len(mask_names) : len(mask_names)]
-        for ax, mask_img in zip(axes, mask_np):
-            add_shape_from_mask(
-                ax, mask_img.squeeze(), color=mask_colors[name], alpha=0.3
-            )
-
-    # Add column titles
-    row_idx = 0
-    if titles is None:
-        titles = mask_names
-    for col_idx, name in enumerate(titles):
-        ax_idx = row_idx * len(mask_names) + col_idx
-        fig.axes[ax_idx].set_title(name, fontsize=9, color="white")
-        fig.axes[ax_idx].set_facecolor("black")
-
-    # Add colorbar for per_pixel_omega if present
-    if "per_pixel_omega" in mask_names:
-        col_idx = mask_names.index("per_pixel_omega")
-        axes = fig.axes[col_idx : batch_size * len(mask_names) : len(mask_names)]
-
-        # Get vertical bounds of the subplot column
-        top_ax = axes[0]
-        bottom_ax = axes[-1]
-        top_pos = top_ax.get_position()
-        bottom_pos = bottom_ax.get_position()
-
-        full_y0 = bottom_pos.y0
-        full_y1 = top_pos.y1
-        full_height = full_y1 - full_y0
-
-        # Manually shrink to 80% of full height and center vertically
-        scale = 0.8
-        height = full_height * scale
-        y0 = full_y0 + (full_height - height) / 2
-
-        x0 = top_pos.x1 + 0.015  # Horizontal position to the right
-        width = 0.015  # Thin bar
-
-        # Add colorbar axis
-        cax = fig.add_axes([x0, y0, width, height])
-
-        im = axes[0].get_images()[0] if axes[0].get_images() else None
-        cbar = fig.colorbar(im, cax=cax)
-        cbar.set_label(r"Guidance weighting \mathbf{p}")
-        cbar.ax.yaxis.set_major_locator(plt.MaxNLocator(nbins=6))
-        cbar.ax.yaxis.set_tick_params(labelsize=7)
-        cbar.ax.yaxis.label.set_size(8)
-
-    return fig
-
-
-def plot_dehazed_results(
-    hazy_images,
-    pred_tissue_images,
-    pred_haze_images,
-    diffusion_model,
-    titles=("Hazy", "Dehazed", "Haze"),
-):
-    """Create and save visualization with optional mask overlays."""
-
-    # Create the processed image stack using the helper function
-    input_shape = diffusion_model.input_shape
-    stack_images = ops.stack(
-        [
-            hazy_images,
-            pred_tissue_images,
-            pred_haze_images,
-        ]
-    )
-    stack_images = ops.reshape(stack_images, (-1, input_shape[0], input_shape[1]))
-
-    # Define labels based on what we're showing
-    fig, _ = plot_image_grid(
-        stack_images,
-        ncols=len(hazy_images),
-        remove_axis=False,
-        vmin=0,
-        vmax=255,
-    )
-    # Set labels and styling
-    for i, ax in enumerate(fig.axes):
-        if i % len(hazy_images) == 0:
-            label = titles[(i // len(hazy_images)) % len(titles)]
-            ax.set_ylabel(label, fontsize=12)
-
-    return fig
-
-
 def main(
     input_folder: str = "./assets",
     output_folder: str = "./temp",

plots.py

@@ -0,0 +1,254 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from keras import ops
+from matplotlib.patches import PathPatch
+from matplotlib.path import Path as pltPath
+from skimage import measure
+from zea.visualize import plot_image_grid
+
+
+def add_shape_from_mask(ax, mask, **kwargs):
+    """add a shape to axis from mask array.
+
+    Args:
+        ax (plt.ax): matplotlib axis
+        mask (ndarray): numpy array with non-zero
+            shape defining the region of interest.
+    Kwargs:
+        edgecolor (str): color of the shape's edge
+        facecolor (str): color of the shape's face
+        linewidth (int): width of the shape's edge
+
+    Returns:
+        plt.ax: matplotlib axis with shape added
+    """
+    # Pad mask to ensure edge contours are found
+    padded_mask = np.pad(mask, pad_width=1, mode="constant", constant_values=0)
+    contours = measure.find_contours(padded_mask, 0.5)
+    patches = []
+    for contour in contours:
+        # Remove padding offset
+        contour -= 1
+        path = pltPath(contour[:, ::-1])
+        patch = PathPatch(path, **kwargs)
+        patches.append(ax.add_patch(patch))
+    return patches
+
+
+def plot_batch_with_named_masks(
+    images, masks_dict, mask_colors=None, titles=None, **kwargs
+):
+    """
+    Plot batch of images in rows, each column overlays a different mask from the dict.
+    Mask labels are shown as column titles. If mask name is 'per_pixel_omega', show it
+    directly with inferno colormap (no overlay).
+
+    Args:
+        images: np.ndarray, shape (batch, height, width, channels)
+        masks_dict: dict of {name: mask}, each mask shape  (batch, height, width, channels)
+        mask_colors: dict of {name: color} or None (default colors used)
+    """
+    mask_names = list(masks_dict.keys())
+    batch_size = images.shape[0]
+    default_colors = ["red", "green", "#33aaff", "yellow", "magenta", "cyan"]
+    mask_colors = mask_colors or {
+        name: default_colors[i % len(default_colors)]
+        for i, name in enumerate(mask_names)
+    }
+
+    # Prepare images for each column
+    columns = []
+    cmaps = []
+    for name in mask_names:
+        if name == "per_pixel_omega":
+            mask_np = np.array(masks_dict[name])
+            columns.append(np.squeeze(mask_np))
+            cmaps.append(["inferno"] * batch_size)
+        else:
+            columns.append(np.squeeze(images))
+            cmaps.append(["gray"] * batch_size)
+
+    # Stack columns: shape (num_columns, batch, ...)
+    all_images = np.stack(columns, axis=0)  # (num_columns, batch, ...)
+    # Rearrange to (batch, num_columns, ...)
+    all_images = (
+        np.transpose(all_images, (1, 0, 2, 3, 4))
+        if all_images.ndim == 5
+        else np.transpose(all_images, (1, 0, 2, 3))
+    )
+    # Flatten to (batch * num_columns, ...)
+    all_images = all_images.reshape(batch_size * len(mask_names), *images.shape[1:])
+
+    # Flatten cmaps for plot_image_grid in the same order as images
+    flat_cmaps = []
+    for row in range(batch_size):
+        for col in range(len(mask_names)):
+            flat_cmaps.append(cmaps[col][row])
+
+    fig, _ = plot_image_grid(
+        all_images,
+        ncols=len(mask_names),
+        remove_axis=False,
+        cmap=flat_cmaps,
+        figsize=(8, 3.3),
+        **kwargs,
+    )
+
+    # Overlay masks for non-per_pixel_omega columns
+    for col_idx, name in enumerate(mask_names):
+        if name == "per_pixel_omega":
+            continue
+        mask_np = np.array(masks_dict[name])
+        axes = fig.axes[col_idx : batch_size * len(mask_names) : len(mask_names)]
+        for ax, mask_img in zip(axes, mask_np):
+            add_shape_from_mask(
+                ax, mask_img.squeeze(), color=mask_colors[name], alpha=0.3
+            )
+
+    # Add column titles
+    row_idx = 0
+    if titles is None:
+        titles = mask_names
+    for col_idx, name in enumerate(titles):
+        ax_idx = row_idx * len(mask_names) + col_idx
+        fig.axes[ax_idx].set_title(name, fontsize=9, color="white")
+        fig.axes[ax_idx].set_facecolor("black")
+
+    # Add colorbar for per_pixel_omega if present
+    if "per_pixel_omega" in mask_names:
+        col_idx = mask_names.index("per_pixel_omega")
+        axes = fig.axes[col_idx : batch_size * len(mask_names) : len(mask_names)]
+
+        # Get vertical bounds of the subplot column
+        top_ax = axes[0]
+        bottom_ax = axes[-1]
+        top_pos = top_ax.get_position()
+        bottom_pos = bottom_ax.get_position()
+
+        full_y0 = bottom_pos.y0
+        full_y1 = top_pos.y1
+        full_height = full_y1 - full_y0
+
+        # Manually shrink to 80% of full height and center vertically
+        scale = 0.8
+        height = full_height * scale
+        y0 = full_y0 + (full_height - height) / 2
+
+        x0 = top_pos.x1 + 0.015  # Horizontal position to the right
+        width = 0.015  # Thin bar
+
+        # Add colorbar axis
+        cax = fig.add_axes([x0, y0, width, height])
+
+        im = axes[0].get_images()[0] if axes[0].get_images() else None
+        cbar = fig.colorbar(im, cax=cax)
+        cbar.set_label(r"Guidance weighting \mathbf{p}")
+        cbar.ax.yaxis.set_major_locator(plt.MaxNLocator(nbins=6))
+        cbar.ax.yaxis.set_tick_params(labelsize=7)
+        cbar.ax.yaxis.label.set_size(8)
+
+    return fig
+
+
+def plot_dehazed_results(
+    hazy_images,
+    pred_tissue_images,
+    pred_haze_images,
+    diffusion_model,
+    titles=("Hazy", "Dehazed", "Haze"),
+):
+    """Create and save visualization with optional mask overlays."""
+
+    # Create the processed image stack using the helper function
+    input_shape = diffusion_model.input_shape
+    stack_images = ops.stack(
+        [
+            hazy_images,
+            pred_tissue_images,
+            pred_haze_images,
+        ]
+    )
+    stack_images = ops.reshape(stack_images, (-1, input_shape[0], input_shape[1]))
+
+    # Define labels based on what we're showing
+    fig, _ = plot_image_grid(
+        stack_images,
+        ncols=len(hazy_images),
+        remove_axis=False,
+        vmin=0,
+        vmax=255,
+    )
+    # Set labels and styling
+    for i, ax in enumerate(fig.axes):
+        if i % len(hazy_images) == 0:
+            label = titles[(i // len(hazy_images)) % len(titles)]
+            ax.set_ylabel(label, fontsize=12)
+
+    return fig
+
+
+def plot_metrics(metrics, limits, out_path):
+    plt.style.use("seaborn-v0_8-darkgrid")
+    fig, axes = plt.subplots(1, len(metrics), figsize=(7.2, 2.7), dpi=600)
+    colors = ["#0057b7", "#ffb300", "#008744", "#d62d20"]
+    metric_labels = {
+        "CNR": r"CNR $\uparrow$",
+        "gCNR": r"gCNR $\uparrow$",
+        "KS_A": r"KS$_{septum}$ $\downarrow$",
+        "KS_B": r"KS$_{ventricle}$ $\uparrow$",
+    }
+    # For legend handles
+    legend_handles = []
+    import matplotlib.lines as mlines
+
+    min_style = {
+        "color": "crimson",
+        "linestyle": "--",
+        "lw": 1.2,
+        "marker": "o",
+        "markersize": 5,
+    }
+    max_style = {
+        "color": "crimson",
+        "linestyle": ":",
+        "lw": 1.2,
+        "marker": "s",
+        "markersize": 5,
+    }
+    for idx, (ax, (name, values)) in enumerate(zip(axes, metrics.items())):
+        ax.hist(
+            values,
+            bins=30,
+            color=colors[idx % len(colors)],
+            alpha=0.85,
+            edgecolor="black",
+            linewidth=0.7,
+        )
+        ax.set_xlabel(metric_labels.get(name, name), fontsize=11)
+        if idx == 0:
+            ax.set_ylabel("Count", fontsize=10)
+        # Draw limits and collect legend handles only once
+        if name in limits:
+            lims = limits[name]
+            if len(legend_handles) == 0:
+                # Only add legend handles for the first metric
+                min_handle = mlines.Line2D([], [], **min_style, label="min score")
+                max_handle = mlines.Line2D([], [], **max_style, label="max score")
+                legend_handles.extend([min_handle, max_handle])
+            if len(lims) > 0:
+                ax.axvline(lims[0], **min_style)
+            if len(lims) > 1:
+                ax.axvline(lims[1], **max_style)
+        ax.spines["top"].set_visible(False)
+        ax.spines["right"].set_visible(False)
+        ax.tick_params(axis="both", which="major", labelsize=9)
+    # Place legend above all subplots
+    fig.legend(
+        handles=legend_handles,
+        loc="upper center",
+        ncol=2,
+        fontsize=10,
+        frameon=False,
+        bbox_to_anchor=(0.5, 1.02),
+    )
+    return fig