Added Seaborn SNS Heatmap, SNS Pairplot nodes.

2025-02-28 23:38:22 -08:00
parent f207bcbcdd
commit 1c4de170f9
8 changed files with 202 additions and 5 deletions
@@ -207,8 +207,10 @@ from .modules.pytorch_wrapper.pt_where import PtWhere
 from .modules.pytorch_wrapper.pt_zeros import PtZeros
 from .modules.seaborn_wrapper.sns_bar import SNSBar
 from .modules.seaborn_wrapper.sns_boxplot import SNSBoxplot
+from .modules.seaborn_wrapper.sns_heatmap import SNSHeatmap
 from .modules.seaborn_wrapper.sns_histogram import SNSHistogram
 from .modules.seaborn_wrapper.sns_line import SNSLine
+from .modules.seaborn_wrapper.sns_pairplot import SNSPairplot
 from .modules.seaborn_wrapper.sns_scatter import SNSScatter
 T = TypeVar("T")

@@ -426,8 +428,10 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
    "PyKvStringCreate": PyKvStringCreate,
    "SNSBar": SNSBar,
    "SNSBoxplot": SNSBoxplot,
+    "SNSHeatmap": SNSHeatmap,
    "SNSHistogram": SNSHistogram,
    "SNSLine": SNSLine,
+    "SNSPairplot": SNSPairplot,
    "SNSScatter": SNSScatter,
 }

@@ -645,8 +649,10 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
    "PyKvStringCreate": "Py Kv String Create",
    "SNSBar": "SNS Bar Chart",
    "SNSBoxplot": "SNS Boxplot",
+    "SNSHeatmap": "SNS Heatmap",
    "SNSHistogram": "SNS Histogram",
    "SNSLine": "SNS Line Plot",
+    "SNSPairplot": "SNS Pairplot",
    "SNSScatter": "SNS Scatter Plot",
 }

@@ -140,8 +140,10 @@
 | [MPL Scatter Plot](mpl_scatter.md) |     Generates a scatter plot from a pandas DataFrame. |
 | [SNS Bar Chart](sns_bar.md) | Generates a bar chart from a Pandas DataFrame using Seaborn. |
 | [SNS Boxplot](sns_boxplot.md) | Generates a boxplot from a Pandas DataFrame using Seaborn. |
+| [SNS Heatmap](sns_heatmap.md) | Generates a heatmap from a Pandas DataFrame using Seaborn. |
 | [SNS Histogram](sns_histogram.md) | Generates a histogram from a Pandas DataFrame using Seaborn. |
 | [SNS Line Plot](sns_line.md) | Generates a line plot from a pandas DataFrame using Seaborn. |
+| [SNS Pairplot](sns_pairplot.md) | Generates a pairplot from a Pandas DataFrame using Seaborn. |
 | [SNS Scatter Plot](sns_scatter.md) | Generates a scatter plot from a pandas DataFrame using Seaborn. |
 ## PyTorch wrapper - Arithmetic operations
 | Node | Description |
@@ -0,0 +1,19 @@
+# SNS Heatmap
+Generates a heatmap from a Pandas DataFrame using Seaborn.
+
+## Input
+| Name | Data type |
+|---|---|
+| dataframe | DataFrame |
+| title | String |
+| palette | String |
+
+## Output
+| Data type |
+|---|
+| Image |
+
+<HR>
+Category: Plot
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
@@ -0,0 +1,23 @@
+# SNS Pairplot
+Generates a pairplot from a Pandas DataFrame using Seaborn.
+
+Specify a column name for the hue parameter to color the data points based on that column.
+
+## Input
+| Name | Data type |
+|---|---|
+| dataframe | DataFrame |
+| hue_column_name | String |
+| title | String |
+| style | String |
+| palette | String |
+
+## Output
+| Data type |
+|---|
+| Image |
+
+<HR>
+Category: Plot
+
+ComfyUI Data Analysis Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.
@@ -0,0 +1,67 @@
+from typing import Any, Dict
+from matplotlib import pyplot as plt
+import pandas as pd
+import seaborn as sns
+from .util import plot_post_steps
+
+
+class SNSHeatmap:
+    """
+    SNS Heatmap:
+    Generates a heatmap from a Pandas DataFrame using Seaborn.
+
+    category: Plot
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the `bar_chart` function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "dataframe": ("DATAFRAME", {}),
+                "title": ("STRING", {"default": ""}),
+                "palette": ("STRING", {"default": ""})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("IMAGE",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Data Analysis"
+
+    def f(self,
+         dataframe: pd.DataFrame,
+         title: str,
+         palette: str=None) -> tuple:
+        """
+        Generates a heatmap from a Pandas DataFrame using Seaborn.
+
+        Args:
+            dataframe (DataFrame): The DataFrame.
+            title (str): The title of the plot.
+            palette (str): Palette of the plot
+        Returns:
+            tuple: A tuple containing the image tensor representation of the plot.
+        """
+        # set_up_sns_style_palette(style, palette) won't work for heatmap
+
+        # Create the plot
+        fig, ax = plt.subplots()
+        corr = dataframe.corr()
+        if palette:
+            sns.heatmap(corr, annot=True, ax=ax, cmap=sns.color_palette(palette, as_cmap=True))
+        else:
+            sns.heatmap(corr, annot=True, ax=ax)
+
+        return plot_post_steps(
+            fig,
+            ax,
+            plt,
+            title,
+            None,
+            None,
+            None)
@@ -0,0 +1,79 @@
+from typing import Any, Dict
+from matplotlib import pyplot as plt
+import pandas as pd
+import seaborn as sns
+from .util import plot_post_steps
+from .util import set_up_sns_style_palette
+
+
+class SNSPairplot:
+    """
+    SNS Pairplot:
+    Generates a pairplot from a Pandas DataFrame using Seaborn.
+
+    Specify a column name for the hue parameter to color the data points based on that column.
+
+    category: Plot
+    """
+    
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the input types for the `bar_chart` function.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required input types.
+        """
+        return {
+            "required": {
+                "dataframe": ("DATAFRAME", {}),
+                "hue_column_name": ("STRING", {"default": ""}),
+                "title": ("STRING", {"default": ""}),
+                "style": ("STRING", {"default": ""}),
+                "palette": ("STRING", {"default": ""})
+            }
+        }
+
+    RETURN_TYPES: tuple = ("IMAGE",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Data Analysis"
+
+    def f(self,
+                  dataframe: pd.DataFrame,
+                  hue_column_name: str,
+                  title: str,
+                  style: str=None,
+                  palette: str=None) -> tuple:
+        """
+        Generates a pairplot from a pandas DataFrame.
+
+        Args:
+            dataframe (DataFrame): The DataFrame.
+            hue_column_name (str): The category (variable) column name.
+            title (str): The title of the plot.
+            style (str): Style of the plot
+            palette (str): Palette of the plot
+        Returns:
+            tuple: A tuple containing the image tensor representation of the plot.
+        """
+        df = dataframe
+        set_up_sns_style_palette(style, palette)
+
+        # Create the plot
+        if hue_column_name:
+            pp = sns.pairplot(df, hue=hue_column_name)
+        else:
+            pp = sns.pairplot(df)
+        
+        pp.figure.suptitle(title, y=0.95)
+        pp.figure.tight_layout(rect=[0, 0, 1, 0.95]) 
+        fig = pp.figure
+
+        return plot_post_steps(
+            fig,
+            None,
+            plt,
+            title,
+            None,
+            None,
+            None)
@@ -89,10 +89,11 @@ def plot_post_steps(
        ax.xaxis.set_major_locator(mticker.MaxNLocator(integer=True))

    # Set title and labels
-    ax.set_title(title)
-    if x_axis_label:
-        ax.set_xlabel(x_axis_label)
-    ax.set_ylabel(y_axis_label)
+    if ax:
+        ax.set_title(title)
+        if x_axis_label:
+            ax.set_xlabel(x_axis_label)
+        ax.set_ylabel(y_axis_label)

    # Save the plot to a buffer
    buf = io.BytesIO()
@@ -1,7 +1,7 @@
 [project]
 name = "data-analysis"
 description = "Data analysis custom modules for ComfyUI - Use Pandas & Matplotlib within ComfyUI"
-version = "1.4.3"
+version = "1.4.4"
 license = {file = "LICENSE"}

 [project.urls]