diff --git a/bertopic/_bertopic.py b/bertopic/_bertopic.py index cfafb58a..dc10b973 100644 --- a/bertopic/_bertopic.py +++ b/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 = "Representative Images per Topic", + 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, diff --git a/bertopic/plotting/__init__.py b/bertopic/plotting/__init__.py index a79a747c..8c64e016 100644 --- a/bertopic/plotting/__init__.py +++ b/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", diff --git a/bertopic/plotting/_images.py b/bertopic/plotting/_images.py new file mode 100644 index 00000000..bc13ceb4 --- /dev/null +++ b/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 = "Representative Images per Topic", + 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}
{name}"}, + ], + ) + 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}
{names[0]}", + "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 diff --git a/docs/api/plotting/representative_images.md b/docs/api/plotting/representative_images.md new file mode 100644 index 00000000..86d64bf0 --- /dev/null +++ b/docs/api/plotting/representative_images.md @@ -0,0 +1,3 @@ +# `Representative Images` + +::: bertopic.plotting._images.visualize_representative_images diff --git a/docs/getting_started/multimodal/multimodal.md b/docs/getting_started/multimodal/multimodal.md index 4ed9b767..26df88c5 100644 --- a/docs/getting_started/multimodal/multimodal.md +++ b/docs/getting_started/multimodal/multimodal.md @@ -74,6 +74,12 @@ HTML(df.to_html(formatters={'Visual_Aspect': image_formatter}, escape=False))

+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))



+ +You can also step through the representative images of each topic, at their original resolution, with a slider: + +```python +topic_model.visualize_representative_images() +``` diff --git a/mkdocs.yml b/mkdocs.yml index 26a823df..07ca6cee 100644 --- a/mkdocs.yml +++ b/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 diff --git a/tests/test_plotting/test_images.py b/tests/test_plotting/test_images.py new file mode 100644 index 00000000..3c0086c9 --- /dev/null +++ b/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()