jupyter/viewers.py

class ImageSliceViewer3D:  # pylint: disable=R0903
    """
    Allows thumbing through a given volume's slices interactively.

    Parameters
    ----------
    volume : np.ndarray
        3D volume to be sliced
    show : bool, optional
        Whether to display the viewer immediately, by default True
    """

    def __init__(self, volume, show=True):
        self.app = dash.Dash(__name__, update_title=None)
        slicer = VolumeSlicer(self.app, volume)
        slicer.graph.config["scrollZoom"] = False  # pylint: disable=E1101
        self.app.layout = html.Div(
            children=[slicer.graph, slicer.slider, *slicer.stores]
        )
        if show is True:
            self.show()

    def show(self):
        """
        Allows thumbing through multiple volumes' slices interactively.
        You can

        Parameters
        ----------
        volume : np.ndarray
            3D volume to be sliced
        show : bool, optional
            Whether to display the viewer immediately, by default True
        """
        # Run the app
        self.app.run_server(debug=True)

class MultiVolumeSliceViewer:
    """
    _summary_
    """

    def __init__(self, volumes, show=True):
        self.app = dash.Dash(
            __name__, external_stylesheets=[dbc.themes.BOOTSTRAP], update_title=None
        )
        self.volumes = volumes
        self.slicers = []

        # Verify all volumes have the same slice count
        slice_counts = [volume.shape[0] for volume in volumes]
        if len(set(slice_counts)) != 1:
            raise ValueError("All volumes must have the same number of slices")
        self.slice_max = slice_counts[0] - 1  # Max index for slices

        # Container for all slicer elements
        slicer_elements = []

        for i, volume in enumerate(volumes):
            slicer = VolumeSlicer(self.app, volume, axis=0)
            slicer.graph.config["scrollZoom"] = False  # pylint: disable=E1101
            self.slicers.append(slicer)

            setpos_store = Store(
                id={"context": "app", "scene": slicer.scene_id, "name": "setpos"}
            )

            # Append slicer elements for each volume
            slicer_elements.append(
                html.Div(
                    [
                        html.H3(f"Volume {i+1}", style={"text-align": "center"}),
                        slicer.graph,
                        html.Div(
                            slicer.slider,
                            id={"type": "slicer-slider", "index": i},
                            style={"display": "none"},
                        ),
                        setpos_store,
                        *slicer.stores,
                    ],
                    id=f"slicer-{i}",
                    style={
                        "width": f"{100 / len(volumes)}%",
                        "display": "inline-block",
                    },
                )
            )  # Adjust width based on the number of volumes

        # Create a unified slider
        unified_slider = Slider(
            id="unified-slider",
            min=0,
            max=self.slice_max,
            step=1,
            value=self.slice_max // 2,
            marks={
                i: str(i)
                for i in range(0, self.slice_max + 1, max(1, self.slice_max // 10))
            },
        )

        # Create a horizontal layout of slicers
        self.slicer_container = html.Div(
            children=slicer_elements,
            id="slicers",
            style={"display": "flex", "flex-direction": "row", "align-items": "start"},
        )

        self.app.layout = html.Div(
            children=[
                self.slicer_container,
                html.Div(
                    [unified_slider],
                    id="unified-slider-container",
                    style={"display": "none", "width": "100%", "text-align": "center"},
                ),
                self.create_toggle_switch(),
                html.Div(id="info"),
            ]
        )

        self.add_sync_feature()

        if show:
            self.show()

    def show(self):
        """
        Shows the viewer if not already displayed.
        """
        # Run the app
        self.app.run_server(debug=True)

    def create_toggle_switch(self) -> dbc.Container:
        """
        Creates the toggle for switching between unified and individual slicer views.

        Returns
        -------
        dbc.Container
            Container with button.
        """
        return dbc.Container(
            [
                dbc.Row(
                    [
                        dbc.Col(
                            [
                                dbc.Label("Toggle Switch"),
                                dbc.Checklist(
                                    options=[
                                        {"label": "Enable", "value": 1},
                                    ],
                                    value=[],
                                    id="toggle-switch",
                                    switch=True,
                                ),
                            ],
                            width=3,
                        ),
                    ],
                    justify="center",
                    align="center",
                    style={"margin-top": "50px"},
                ),
                dbc.Row(
                    [
                        dbc.Col(html.Div(id="toggle-output"), width=3),
                    ],
                    justify="center",
                    align="center",
                    style={"margin-top": "20px"},
                ),
            ],
            fluid=True,
        )

    def add_sync_feature(self):
        """
        Adds the synchronization feature to the viewer.
        """

        @self.app.callback(
            Output("unified-slider-container", "style"),
            [Input("toggle-switch", "value")],
        )
        def toggle_unified_slider_visibility(enabled):
            if 1 in enabled:
                return {"width": "100%", "text-align": "center"}
            return {"display": "none"}

        @self.app.callback(
            Output({"type": "slicer-slider", "index": ALL}, "style"),
            [Input("toggle-switch", "value")],
        )
        def update_slicers_visibility(enabled):
            if 1 in enabled:
                # Hide individual sliders
                return [{"display": "none"} for _ in self.slicers]
            # Show individual sliders
            return [{"display": "block"} for _ in self.slicers]

        @self.app.callback(
            Output("toggle-output", "children"), [Input("toggle-switch", "value")]
        )
        def update_output(value):
            if 1 in value:
                return "Switch is ON"
            return "Switch is OFF"

        @self.app.callback(
            Output({"context": "app", "scene": ALL, "name": "setpos"}, "data"),
            [Input("unified-slider", "value")],
            [State("toggle-switch", "value")],
        )
        def sync_slicers_from_unified(value, toggle_value):
            if toggle_value and 1 in toggle_value:
                return [(None, None, value)] * len(self.slicers)
            return [dash.no_update] * len(self.slicers)

        @self.app.callback(
            Output("unified-slider", "value"),
            [Input(slicer.state.id, "data") for slicer in self.slicers],
            [State("toggle-switch", "value")],
        )
        def sync_unified_slider(*args):
            toggle_value = args[-1]
            if not toggle_value or 1 not in toggle_value:
                return dash.no_update

            ctx = dash.callback_context
            if not ctx.triggered:
                return dash.no_update

            triggered_state = next(
                (state for state in args[:-1] if state["index_changed"]), None
            )
            if triggered_state:
                return triggered_state["index"]
            return dash.no_update