Update WER plot with improved model metrics and visualization enhancements

comparison/wer_plot.py

@@ -7,110 +7,151 @@ import matplotlib.ticker as ticker
 
 # Data for the models
 models = [
-    ("Piper", 0.22, 0.09, "Ours"),
-    ("StyleTTS2", 0.18, 0.50, "Ours"),
-    ("HebTTS", 0.36, 25.44, "Open"),
-    ("LoTHM", 0.59, 84.75, "Open"),
-    ("MMS", 0.33, 0.21, "Open"),
-    ("SASPEECH", 0.22, 0.16, "Open"),
-    ("Robo-Shaul", 0.20, 1.58, "Open"),
-    ("Google", 0.20, 4.08, "Proprietary"),
-    ("OpenAI", 0.21, 1.60, "Proprietary"),
+    ("Piper", 0.11, 0.09, "Ours"),
+    ("StyleTTS2", 0.07, 0.50, "Ours"),
+    ("HebTTS", 0.24, 25.44, "Open"),
+    ("LoTHM", 0.49, 84.75, "Open"),
+    ("MMS", 0.20, 0.21, "Open"),
+    ("SASPEECH", 0.11, 0.16, "Open"),
+    ("Robo-Shaul", 0.08, 1.58, "Open"),
+    ("Google", 0.04, 4.08, "Proprietary"),
+    ("OpenAI", 0.05, 1.60, "Proprietary"),
 ]
 
 # Filter out models with None values for WER or RTF
 filtered = [m for m in models if m[1] is not None and m[2] is not None]
 
-# Create the figure and axes
-fig, ax = plt.subplots(figsize=(10, 6))
+# Create the figure and axes with better sizing
+fig, ax = plt.subplots(figsize=(12, 8))
+
+# Color mapping with fancy colors
+colors = {'Ours': '#e74c3c', 'Open': '#3498db', 'Proprietary': '#f39c12'}
+legend_elements = []
 
 # Plot each model
 for name, wer, rtf, category in filtered:
     # Determine color based on category
-    if name in ["Google", "OpenAI"]:
-        color = '#f4c285'
-    elif category == 'Ours':
-        color = 'red'
-    else:
-        color = 'blue'
-
+    color = colors[category]
+    
     # Determine size and weight for our models
-    size = 200 if category == 'Ours' else 100
+    size = 180  # Same size for all models
     weight = 'bold' if category == 'Ours' else 'normal'
-    weight = 'normal' if name in ["Google", "OpenAI"] else weight
+    edgewidth = 2 if category == 'Ours' else 1.5
 
     # Create label for the point
     label = f"Ours ({name})" if category == 'Ours' else name
 
     # Plot the scatter point
-    ax.scatter(rtf, wer, s=size, c=color, edgecolors='black', linewidths=1, zorder=3) # zorder to ensure points are above grid
+    scatter = ax.scatter(rtf, wer, s=size, c=color, edgecolors='black', 
+                        linewidths=edgewidth, zorder=3, alpha=0.8)
 
-    # Adjust text position for HebTTS to avoid overlap
+    # Adjust text position for each model
     if name == "HebTTS":
-        x_text = rtf * 0.85
+        x_text = rtf * 0.75
+        y_text = wer
         ha = 'right'
+        va = 'center'
     elif name == "Google":
-        x_text = rtf * 1.15  # right
-        y_text = wer - 0.01  # slightly down
+        x_text = rtf * 1.2
+        y_text = wer - 0.008
         ha = 'left'
+        va = 'center'
     elif name == 'LoTHM':
-        x_text = rtf * 0.90
+        x_text = rtf * 0.85
+        y_text = wer
+        ha = 'right'
+        va = 'center'
     elif name == "OpenAI":
-        x_text = rtf * 0.85  # left
-        y_text = wer - 0.01  # slightly down
+        x_text = rtf
+        y_text = wer - 0.012
+        ha = 'center'
+        va = 'top'
+    elif name == "Robo-Shaul":
+        x_text = rtf * 1.2
+        y_text = wer + 0.008
+        ha = 'left'
+        va = 'center'
+    elif name == "Piper":
+        x_text = rtf * 0.9
+        y_text = wer - 0.022
+        ha = 'left'
+        va = 'top'
+    elif name == "StyleTTS2":
+        x_text = rtf * 0.8
+        y_text = wer - 0.018
+        ha = 'center'
+        va = 'top'
+    elif name == "SASPEECH":
+        x_text = rtf * 1.2
+        y_text = wer + 0.008
+        ha = 'left'
+        va = 'center'
+    elif name == "MMS":
+        x_text = rtf * 0.85
+        y_text = wer + 0.015
         ha = 'right'
+        va = 'bottom'
     else:
         x_text = rtf * 1.15
+        y_text = wer
         ha = 'left'
+        va = 'center'
 
     # Add text label for each point
-    ax.text(x_text, wer, label, fontsize=22, ha=ha, va='center', color='black', weight=weight, zorder=4)
+    fontsize = 20 if category == 'Ours' else 22
+    ax.text(x_text, y_text, label, fontsize=fontsize, ha=ha, va=va, 
+            color='black', weight=weight, zorder=4)
 
 # Set x-axis to log scale and format it
 ax.set_xscale('log')
-ax.tick_params(axis='both', which='major', labelsize=16)
+ax.tick_params(axis='both', which='major', labelsize=14)
+ax.tick_params(axis='both', which='minor', labelsize=12)
 ax.xaxis.set_major_formatter(ticker.ScalarFormatter())
 ax.xaxis.get_major_formatter().set_scientific(False)
 ax.xaxis.get_major_formatter().set_useOffset(False)
 
-# Set axis labels
-ax.set_xlabel("RTF (lower is faster)", fontsize=22)
-ax.set_ylabel("WER (lower is more accurate)", fontsize=22)
+# Add minor ticks for better readability
+ax.xaxis.set_minor_locator(ticker.LogLocator(base=10.0, subs=np.arange(2, 10) * 0.1))
 
-# Remove grid lines
-ax.grid(False)
+# Set axis labels with larger font
+ax.set_xlabel("RTF (lower is faster)", fontsize=18, fontweight='bold')
+ax.set_ylabel("WER (lower is more accurate)", fontsize=18, fontweight='bold')
+
+# Remove title
+
+# Add subtle grid
+ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
+ax.set_axisbelow(True)
+
+# Create custom legend
+from matplotlib.patches import Patch
+legend_elements = [
+    plt.scatter([], [], c=colors['Ours'], s=150, edgecolors='black', 
+                linewidths=1.5, label='Our Models'),
+    plt.scatter([], [], c=colors['Open'], s=120, edgecolors='black', 
+                linewidths=1.5, label='Open Source'),
+    plt.scatter([], [], c=colors['Proprietary'], s=120, edgecolors='black', 
+                linewidths=1.5, label='Proprietary')
+]
+
+ax.legend(handles=legend_elements, loc='upper right', fontsize=14, 
+          frameon=True, fancybox=True, shadow=True)
 
 # Adjust layout to prevent labels from being cut off
 plt.tight_layout()
 
-# Extend x-axis limits by 20% to make space for labels/arrows
+# Extend x-axis limits by 30% to make space for labels
 x_min, x_max = ax.get_xlim()
-ax.set_xlim(x_min, x_max * 1.2)
-
-# --- Add Arrows for Direction of Improvement ---
-
-# Get current axis limits to position arrows relative to the plot
-x_lims = ax.get_xlim()
-y_lims = ax.get_ylim()
-
-# Position arrow in upper-right area of plot, pointing to bottom-left
-# Start point (upper-right area)
-arrow_start_x = x_lims[1] * 0.002  # 80% across the x-axis
-arrow_start_y = y_lims[1] * 0.55  # 85% up the y-axis
-
-# End point (much closer to create shorter arrow with steeper angle)
-arrow_end_x = x_lims[1] * 0.0006   # 60% across the x-axis (shorter horizontal distance)
-arrow_end_y = y_lims[1] * 0.43   # 60% up the y-axis (steeper vertical drop)
-
-# Draw arrow pointing from upper-right toward bottom-left
-ax.annotate('',
-            xy=(arrow_end_x, arrow_end_y),      # End point (arrow head)
-            xytext=(arrow_start_x, arrow_start_y), # Start point (arrow tail)
-            arrowprops=dict(facecolor='gray', shrink=0.05, width=0.5, headwidth=8, alpha=0.3),
-            annotation_clip=False,
-            zorder=1)  # Behind circles (3) and text (4)
-
-# Clear any existing title and save the figure
-plt.title("")
-plt.savefig("plot.png", dpi=1200)
+ax.set_xlim(x_min, x_max * 1.3)
+
+# Extend y-axis limits slightly for better spacing
+y_min, y_max = ax.get_ylim()
+ax.set_ylim(y_min - 0.01, y_max + 0.02)
+
+# Keep plot clean and simple
+
+# Remove figure caption
+
+# Save with high quality
+plt.savefig("plot.png", dpi=300, bbox_inches='tight', facecolor='white')
 plt.show()