Enhance KMeans analysis visualization using Plotly and update requirements

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
 import pandas as pd
 import numpy as np
-import io
-import base64
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
 
 import vlai_template
 
@@ -66,116 +66,143 @@ def standardize_data(X):
     """Standardize data to have mean=0 and std=1"""
     return (X - np.mean(X, axis=0)) / np.std(X, axis=0)
 
-def create_simple_plot_svg(X_pca, wine_types, labels, centroids_pca, k, explained_var):
-    """Create a simple SVG plot"""
-    width, height = 1400, 700
-    margin = 80
+def create_plotly_visualization(X_pca, wine_types, labels, centroids_pca, k, explained_var):
+    """Create a plotly visualization with two subplots"""
     
-    # Split into two plots
-    plot_width = (width - 3 * margin) // 2
-    plot_height = height - 2 * margin
-    
-    # Normalize data to plot coordinates
-    x_min, x_max = X_pca[:, 0].min(), X_pca[:, 0].max()
-    y_min, y_max = X_pca[:, 1].min(), X_pca[:, 1].max()
-    
-    def norm_x(x, plot_x_offset=0):
-        return plot_x_offset + margin + (x - x_min) / (x_max - x_min) * plot_width
-    
-    def norm_y(y):
-        return margin + (1 - (y - y_min) / (y_max - y_min)) * plot_height
-    
-    svg_content = f'''
-    <div class="plot-container" style="width: 100%; overflow-x: auto;">
-    <svg width="{width}" height="{height}" viewBox="0 0 {width} {height}" 
-         style="width: 100%; height: auto; min-width: 800px;" xmlns="http://www.w3.org/2000/svg">
-        <defs>
-            <style>
-                .title {{ font: bold 20px sans-serif; text-anchor: middle; }}
-                .axis-label {{ font: 16px sans-serif; text-anchor: middle; }}
-                .legend {{ font: 14px sans-serif; }}
-            </style>
-        </defs>
-        
-        <!-- Plot 1: Original Wine Types -->
-        <g>
-            <rect x="{margin}" y="{margin}" width="{plot_width}" height="{plot_height}" 
-                  fill="none" stroke="black" stroke-width="1"/>
-            <text x="{margin + plot_width//2}" y="{margin - 10}" class="title">
-                Original Data (Red vs White Wine)
-            </text>
-            <text x="{margin + plot_width//2}" y="{height - 10}" class="axis-label">
-                PC1 ({explained_var[0]:.1%} variance)
-            </text>
-            <text x="{margin - 30}" y="{margin + plot_height//2}" class="axis-label" 
-                  transform="rotate(-90, {margin - 30}, {margin + plot_height//2})">
-                PC2 ({explained_var[1]:.1%} variance)
-            </text>
-    '''
-    
-    # Plot original data points
-    for i, (point, wine_type) in enumerate(zip(X_pca, wine_types)):
-        color = "#d62728" if wine_type == "red" else "#1f77b4"  # red or blue
-        x, y = norm_x(point[0]), norm_y(point[1])
-        svg_content += f'<circle cx="{x}" cy="{y}" r="2" fill="{color}" opacity="0.6"/>\n'
+    # Create subplots
+    fig = make_subplots(
+        rows=1, cols=2,
+        subplot_titles=(
+            "Original Data (Red vs White Wine)",
+            f"After KMeans Clustering (K={k})"
+        ),
+        horizontal_spacing=0.1
+    )
     
-    # Plot 1 legend
-    svg_content += f'''
-            <circle cx="{margin + 10}" cy="{margin + 20}" r="4" fill="#d62728"/>
-            <text x="{margin + 20}" y="{margin + 25}" class="legend">Red Wine</text>
-            <circle cx="{margin + 10}" cy="{margin + 40}" r="4" fill="#1f77b4"/>
-            <text x="{margin + 20}" y="{margin + 45}" class="legend">White Wine</text>
-        </g>
-        
-        <!-- Plot 2: KMeans Clusters -->
-        <g>
-    '''
+    # Plot 1: Original Wine Types
+    red_mask = np.array(wine_types) == "red"
+    white_mask = np.array(wine_types) == "white"
     
-    plot2_x_offset = plot_width + margin
+    # Add red wine points
+    if np.any(red_mask):
+        fig.add_trace(
+            go.Scatter(
+                x=X_pca[red_mask, 0],
+                y=X_pca[red_mask, 1],
+                mode='markers',
+                marker=dict(color='#d62728', size=4, opacity=0.6),
+                name='Red Wine',
+                showlegend=True
+            ),
+            row=1, col=1
+        )
     
-    svg_content += f'''
-            <rect x="{margin + plot2_x_offset}" y="{margin}" width="{plot_width}" height="{plot_height}" 
-                  fill="none" stroke="black" stroke-width="1"/>
-            <text x="{margin + plot2_x_offset + plot_width//2}" y="{margin - 10}" class="title">
-                After KMeans Clustering (K={k})
-            </text>
-            <text x="{margin + plot2_x_offset + plot_width//2}" y="{height - 10}" class="axis-label">
-                PC1 ({explained_var[0]:.1%} variance)
-            </text>
-            <text x="{margin + plot2_x_offset - 30}" y="{margin + plot_height//2}" class="axis-label" 
-                  transform="rotate(-90, {margin + plot2_x_offset - 30}, {margin + plot_height//2})">
-                PC2 ({explained_var[1]:.1%} variance)
-            </text>
-    '''
+    # Add white wine points
+    if np.any(white_mask):
+        fig.add_trace(
+            go.Scatter(
+                x=X_pca[white_mask, 0],
+                y=X_pca[white_mask, 1],
+                mode='markers',
+                marker=dict(color='#1f77b4', size=4, opacity=0.6),
+                name='White Wine',
+                showlegend=True
+            ),
+            row=1, col=1
+        )
     
-    # Colors for clusters
+    # Plot 2: KMeans Clusters
     cluster_colors = ["#ff7f0e", "#2ca02c", "#d62728", "#9467bd", "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22"]
     
-    # Plot clustered data points
-    for i, (point, label) in enumerate(zip(X_pca, labels)):
-        color = cluster_colors[label % len(cluster_colors)]
-        x, y = norm_x(point[0], plot2_x_offset), norm_y(point[1])
-        svg_content += f'<circle cx="{x}" cy="{y}" r="2" fill="{color}" opacity="0.6"/>\n'
+    # Add cluster points
+    for i in range(k):
+        cluster_mask = labels == i
+        if np.any(cluster_mask):
+            fig.add_trace(
+                go.Scatter(
+                    x=X_pca[cluster_mask, 0],
+                    y=X_pca[cluster_mask, 1],
+                    mode='markers',
+                    marker=dict(color=cluster_colors[i % len(cluster_colors)], size=4, opacity=0.6),
+                    name=f'Cluster {i}',
+                    showlegend=True
+                ),
+                row=1, col=2
+            )
     
-    # Plot centroids
-    for i, centroid in enumerate(centroids_pca):
-        x, y = norm_x(centroid[0], plot2_x_offset), norm_y(centroid[1])
-        svg_content += f'<g><circle cx="{x}" cy="{y}" r="6" fill="black" stroke="white" stroke-width="2"/>'
-        svg_content += f'<text x="{x}" y="{y+2}" text-anchor="middle" class="legend" fill="white">{i}</text></g>\n'
+    # Add centroids
+    fig.add_trace(
+        go.Scatter(
+            x=centroids_pca[:, 0],
+            y=centroids_pca[:, 1],
+            mode='markers+text',
+            marker=dict(color='black', size=12, line=dict(color='white', width=2)),
+            text=[str(i) for i in range(k)],
+            textfont=dict(color='white', size=10),
+            textposition="middle center",
+            name='Centroids',
+            showlegend=True
+        ),
+        row=1, col=2
+    )
     
-    # Plot 2 legend
-    for i in range(k):
-        color = cluster_colors[i % len(cluster_colors)]
-        svg_content += f'<circle cx="{margin + plot2_x_offset + 10}" cy="{margin + 20 + i*20}" r="4" fill="{color}"/>\n'
-        svg_content += f'<text x="{margin + plot2_x_offset + 20}" y="{margin + 25 + i*20}" class="legend">Cluster {i}</text>\n'
+    # Update layout
+    fig.update_layout(
+        title="Wine Quality Dataset - KMeans Clustering with PCA",
+        title_x=0.5,
+        height=600,
+        plot_bgcolor='white',
+        paper_bgcolor='white',
+        font=dict(size=12),
+        showlegend=True,
+        legend=dict(
+            orientation="h",
+            yanchor="bottom",
+            y=-0.2,
+            xanchor="center",
+            x=0.5
+        )
+    )
     
-    svg_content += '''
-        </g>
-    </svg>
-    </div>
-    '''
+    # Update x and y axes labels and styling
+    fig.update_xaxes(
+        title_text=f"PC1 ({explained_var[0]:.1%} variance)",
+        showgrid=True,
+        gridcolor='lightgray',
+        gridwidth=1,
+        zeroline=True,
+        zerolinecolor='lightgray',
+        row=1, col=1
+    )
+    fig.update_xaxes(
+        title_text=f"PC1 ({explained_var[0]:.1%} variance)",
+        showgrid=True,
+        gridcolor='lightgray',
+        gridwidth=1,
+        zeroline=True,
+        zerolinecolor='lightgray',
+        row=1, col=2
+    )
+    fig.update_yaxes(
+        title_text=f"PC2 ({explained_var[1]:.1%} variance)",
+        showgrid=True,
+        gridcolor='lightgray',
+        gridwidth=1,
+        zeroline=True,
+        zerolinecolor='lightgray',
+        row=1, col=1
+    )
+    fig.update_yaxes(
+        title_text=f"PC2 ({explained_var[1]:.1%} variance)",
+        showgrid=True,
+        gridcolor='lightgray',
+        gridwidth=1,
+        zeroline=True,
+        zerolinecolor='lightgray',
+        row=1, col=2
+    )
     
-    return svg_content
+    return fig
 
 def run_kmeans_analysis(n_clusters, random_state):
     """Main function to run the KMeans analysis"""
@@ -202,12 +229,27 @@ def run_kmeans_analysis(n_clusters, random_state):
         
         # Create plot
         wine_types = df['wine_type'].tolist() if 'wine_type' in df.columns else ['unknown'] * len(df)
-        plot_svg = create_simple_plot_svg(X_pca, wine_types, labels, centroids_pca, n_clusters, explained_var)
+        plot_fig = create_plotly_visualization(X_pca, wine_types, labels, centroids_pca, n_clusters, explained_var)
         
-        return plot_svg
+        return plot_fig
         
     except Exception as e:
-        return f"<p style='color: red;'>Error: {str(e)}</p>"
+        # Return an empty plotly figure with error message
+        fig = go.Figure()
+        fig.add_annotation(
+            text=f"Error: {str(e)}",
+            xref="paper", yref="paper",
+            x=0.5, y=0.5,
+            showarrow=False,
+            font=dict(color="red", size=16)
+        )
+        fig.update_layout(
+            plot_bgcolor='white',
+            paper_bgcolor='white',
+            xaxis=dict(visible=False),
+            yaxis=dict(visible=False)
+        )
+        return fig
 
 # ────────────────────────────  Main  ─────────────────────────
 with gr.Blocks(theme='gstaff/sketch', css=vlai_template.custom_css, title="Wine Quality KMeans Demo") as demo:
@@ -245,7 +287,7 @@ with gr.Blocks(theme='gstaff/sketch', css=vlai_template.custom_css, title="Wine
             """)
 
         with gr.Column(scale=7):
-            output_plot = gr.HTML(label="📈 PCA Visualization & KMeans Results")
+            output_plot = gr.Plot(label="📈 PCA Visualization & KMeans Results")
 
     run_btn.click(
         run_kmeans_analysis,

requirements.txt

@@ -1,3 +1,4 @@
 gradio==5.38.0
 pandas
-numpy
+numpy
+plotly