Add visualize_representative_images for image-only topic models

bertopic/_bertopic.py

@@ -2691,6 +2691,70 @@ def visualize_documents(
             height=height,
         )
 
+    def visualize_representative_images(
+        self,
+        topics: List[int] | None = None,
+        aspect: str = "Visual_Aspect",
+        custom_labels: Union[bool, str] = False,
+        title: str = "<b>Representative Images per Topic</b>",
+        width: int = 1200,
+        height: int = 800,
+    ) -> "go.Figure":
+        """Visualize the representative images of each topic with a slider for topic selection.
+
+        When BERTopic is fitted on images with a `VisualRepresentation` aspect, each topic
+        is represented by a collage of its most representative images. This visualization
+        shows those collages at their original resolution, one topic at a time, with a
+        slider to step through the topics.
+
+        Arguments:
+            topics: A selection of topics to visualize.
+                    For example, if you want to visualize only topics 1 through 5:
+                    `topics = [1, 2, 3, 4, 5]`.
+            aspect: The name of the topic aspect that contains the representative images
+                    as created by `bertopic.representation.VisualRepresentation`.
+            custom_labels: If bool, whether to use custom topic labels that were defined using
+                           `topic_model.set_topic_labels`.
+                           If `str`, it uses labels from other aspects, e.g., "Aspect1".
+            title: Title of the plot.
+            width: The width of the figure.
+            height: The height of the figure.
+
+        Examples:
+        To visualize the representative images of each topic, make sure to fit BERTopic
+        with a `VisualRepresentation` aspect first:
+
+        ```python
+        from bertopic import BERTopic
+        from bertopic.representation import VisualRepresentation
+
+        # Additional representation of topics as a collage of images
+        representation_model = {"Visual_Aspect": VisualRepresentation()}
+        topic_model = BERTopic(embedding_model="clip-ViT-B-32", representation_model=representation_model)
+        topics, probs = topic_model.fit_transform(documents=None, images=images)
+
+        # Run the visualization
+        topic_model.visualize_representative_images()
+        ```
+
+        Or if you want to save the resulting figure:
+
+        ```python
+        fig = topic_model.visualize_representative_images()
+        fig.write_html("path/to/file.html")
+        ```
+        """
+        check_is_fitted(self)
+        return plotting.visualize_representative_images(
+            self,
+            topics=topics,
+            aspect=aspect,
+            custom_labels=custom_labels,
+            title=title,
+            width=width,
+            height=height,
+        )
+
     def visualize_document_datamap(
         self,
         docs: List[str] | None = None,

bertopic/plotting/__init__.py

@@ -6,6 +6,7 @@
 from ._hierarchy import visualize_hierarchy
 from ._datamap import visualize_document_datamap
 from ._distribution import visualize_distribution
+from ._images import visualize_representative_images
 from ._topics_over_time import visualize_topics_over_time
 from ._topics_per_class import visualize_topics_per_class
 from ._hierarchical_documents import visualize_hierarchical_documents
@@ -21,6 +22,7 @@
     "visualize_heatmap",
     "visualize_hierarchical_documents",
     "visualize_hierarchy",
+    "visualize_representative_images",
     "visualize_term_rank",
     "visualize_topics",
     "visualize_topics_over_time",

bertopic/plotting/_images.py

@@ -0,0 +1,124 @@
+import numpy as np
+import plotly.graph_objects as go
+
+from typing import List, Union
+
+
+def visualize_representative_images(
+    topic_model,
+    topics: List[int] | None = None,
+    aspect: str = "Visual_Aspect",
+    custom_labels: Union[bool, str] = False,
+    title: str = "<b>Representative Images per Topic</b>",
+    width: int = 1200,
+    height: int = 800,
+):
+    """Visualize the representative images of each topic with a slider for topic selection.
+
+    When BERTopic is fitted on images with a `VisualRepresentation` aspect, each topic
+    is represented by a collage of its most representative images. This visualization
+    shows those collages at their original resolution, one topic at a time, with a
+    slider to step through the topics.
+
+    Arguments:
+        topic_model: A fitted BERTopic instance.
+        topics: A selection of topics to visualize.
+                For example, if you want to visualize only topics 1 through 5:
+                `topics = [1, 2, 3, 4, 5]`.
+        aspect: The name of the topic aspect that contains the representative images
+                as created by `bertopic.representation.VisualRepresentation`.
+        custom_labels: If bool, whether to use custom topic labels that were defined using
+                       `topic_model.set_topic_labels`.
+                       If `str`, it uses labels from other aspects, e.g., "Aspect1".
+        title: Title of the plot.
+        width: The width of the figure.
+        height: The height of the figure.
+
+    Examples:
+    To visualize the representative images of each topic, make sure to fit BERTopic
+    with a `VisualRepresentation` aspect first:
+
+    ```python
+    from bertopic import BERTopic
+    from bertopic.representation import VisualRepresentation
+
+    # Additional representation of topics as a collage of images
+    representation_model = {"Visual_Aspect": VisualRepresentation()}
+    topic_model = BERTopic(embedding_model="clip-ViT-B-32", representation_model=representation_model)
+    topics, probs = topic_model.fit_transform(documents=None, images=images)
+
+    # Run the visualization
+    topic_model.visualize_representative_images()
+    ```
+
+    Or if you want to save the resulting figure:
+
+    ```python
+    fig = topic_model.visualize_representative_images()
+    fig.write_html("path/to/file.html")
+    ```
+    """
+    if aspect not in topic_model.topic_aspects_:
+        raise ValueError(
+            f"The aspect '{aspect}' could not be found in `topic_model.topic_aspects_`. "
+            "Make sure to fit BERTopic with a `VisualRepresentation` aspect, for example: "
+            '`BERTopic(representation_model={"Visual_Aspect": VisualRepresentation()})`.'
+        )
+    image_topics = topic_model.topic_aspects_[aspect]
+
+    # Select topics based on top_n and topics args
+    freq_df = topic_model.get_topic_freq()
+    freq_df = freq_df.loc[freq_df.Topic != -1, :]
+    if topics is None:
+        topics = sorted(freq_df.Topic.to_list())
+    topics = [topic for topic in topics if topic in image_topics]
+    if not topics:
+        raise ValueError(f"None of the selected topics have images in the '{aspect}' aspect.")
+
+    # Prepare topic names
+    if isinstance(custom_labels, str):
+        names = [[[str(topic), None]] + topic_model.topic_aspects_[custom_labels][topic] for topic in topics]
+        names = ["_".join([label[0] for label in labels[:4]]) for labels in names]
+        names = [label if len(label) < 30 else label[:27] + "..." for label in names]
+    elif topic_model.custom_labels_ is not None and custom_labels:
+        names = [topic_model.custom_labels_[topic + topic_model._outliers] for topic in topics]
+    else:
+        names = [f"{topic}_" + "_".join([word for word, value in topic_model.get_topic(topic)][:3]) for topic in topics]
+
+    # Visualize
+    fig = go.Figure()
+    for index, topic in enumerate(topics):
+        fig.add_trace(go.Image(z=np.asarray(image_topics[topic]), hoverinfo="skip", visible=index == 0))
+
+    # Create a slider for topic selection
+    steps = [
+        dict(
+            label=f"Topic {topic}",
+            method="update",
+            args=[
+                {"visible": [index == i for i in range(len(topics))]},
+                {"title.text": f"{title}<br><sup>{name}</sup>"},
+            ],
+        )
+        for index, (topic, name) in enumerate(zip(topics, names))
+    ]
+    sliders = [dict(active=0, pad={"t": 50}, steps=steps)]
+
+    # Stylize layout
+    fig.update_layout(
+        template="simple_white",
+        title={
+            "text": f"{title}<br><sup>{names[0]}</sup>",
+            "x": 0.5,
+            "xanchor": "center",
+            "yanchor": "top",
+            "font": dict(size=22, color="Black"),
+        },
+        width=width,
+        height=height,
+        sliders=sliders,
+    )
+
+    fig.update_xaxes(visible=False)
+    fig.update_yaxes(visible=False)
+    return fig

docs/api/plotting/representative_images.md

@@ -0,0 +1,3 @@
+# `Representative Images`
+
+::: bertopic.plotting._images.visualize_representative_images

docs/getting_started/multimodal/multimodal.md

@@ -74,6 +74,12 @@ HTML(df.to_html(formatters={'Visual_Aspect': image_formatter}, escape=False))
 <img src="images_and_text.jpg">
 <br><br>
 
+You can also step through the representative images of each topic, at their original resolution, with a slider:
+
+```python
+topic_model.visualize_representative_images()
+```
+
 !!! Tip
     In the example above, we are clustering the documents but since you have
     images, you might want to cluster those or cluster an aggregation of both
@@ -188,3 +194,9 @@ HTML(df.to_html(formatters={'Visual_Aspect': image_formatter}, escape=False))
 <br><br>
 <img src="images_only.jpg">
 <br><br>
+
+You can also step through the representative images of each topic, at their original resolution, with a slider:
+
+```python
+topic_model.visualize_representative_images()
+```

mkdocs.yml

@@ -72,6 +72,7 @@ nav:
         - Hierarchical topics: api/plotting/hierarchy.md
         - Distribution: api/plotting/distribution.md
         - Heatmap: api/plotting/heatmap.md
+        - Representative Images: api/plotting/representative_images.md
         - Term Scores: api/plotting/term.md
         - Topics: api/plotting/topics.md
         - Topics per Class: api/plotting/topics_per_class.md

tests/test_plotting/test_images.py

@@ -0,0 +1,61 @@
+import copy
+
+import numpy as np
+import pytest
+
+
+def add_image_aspect(topic_model):
+    """Add an image per topic, mimicking the output of `VisualRepresentation`."""
+    topics = set(topic_model.topics_)
+    images = {topic: np.random.randint(0, 255, size=(60, 90, 3), dtype=np.uint8) for topic in topics}
+    topic_model.topic_aspects_["Visual_Aspect"] = images
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        ("kmeans_pca_topic_model"),
+        ("base_topic_model"),
+        ("custom_topic_model"),
+        ("merged_topic_model"),
+        ("reduced_topic_model"),
+    ],
+)
+def test_representative_images(model, request):
+    topic_model = copy.deepcopy(request.getfixturevalue(model))
+    add_image_aspect(topic_model)
+    topics = set(topic_model.topics_)
+    if -1 in topics:
+        topics.remove(-1)
+
+    fig = topic_model.visualize_representative_images()
+    fig_dict = fig.to_dict()
+
+    assert len(fig_dict["data"]) == len(topics)
+    assert [trace["visible"] for trace in fig_dict["data"]].count(True) == 1
+    assert fig_dict["data"][0]["visible"]
+
+    for slider in fig_dict["layout"]["sliders"]:
+        assert len(slider["steps"]) == len(topics)
+        for step in slider["steps"]:
+            assert int(step["label"].split(" ")[-1]) != -1
+
+
+def test_representative_images_subset(base_topic_model):
+    topic_model = copy.deepcopy(base_topic_model)
+    add_image_aspect(topic_model)
+    selected = sorted(set(topic_model.topics_) - {-1})[:2]
+
+    fig = topic_model.visualize_representative_images(topics=selected)
+    fig_dict = fig.to_dict()
+
+    assert len(fig_dict["data"]) == len(selected)
+    for slider in fig_dict["layout"]["sliders"]:
+        assert [step["label"] for step in slider["steps"]] == [f"Topic {topic}" for topic in selected]
+
+
+def test_representative_images_missing_aspect(base_topic_model):
+    topic_model = copy.deepcopy(base_topic_model)
+
+    with pytest.raises(ValueError, match="Visual_Aspect"):
+        topic_model.visualize_representative_images()