src/vis_utils.py

import gradio as gr
import pandas as pd
import re
import os
import json
import yaml
import matplotlib.pyplot as plt
import seaborn as sns
import plotnine as p9

from about import *
global data_component, filter_component

def get_method_color(method):
    return color_dict.get(method, 'black')  # If method is not in color_dict, use black

def draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title):
    df = pd.read_csv(CSV_RESULT_PATH)
    # Filter the dataframe based on selected methods
    filtered_df = df[df['method_name'].isin(methods_selected)]

    def get_method_color(method):
        return color_dict.get(method.upper(), 'black')

    # Add a new column to the dataframe for the color
    filtered_df['color'] = filtered_df['method_name'].apply(get_method_color)

    adjust_text_dict = {
        'expand_text': (1.15, 1.4), 'expand_points': (1.15, 1.25), 'expand_objects': (1.05, 1.5),
        'expand_align': (1.05, 1.2), 'autoalign': 'xy', 'va': 'center', 'ha': 'center',
        'force_text': (.0, 1.), 'force_objects': (.0, 1.),
        'lim': 500000, 'precision': 1., 'avoid_points': True, 'avoid_text': True
    }

    # Create the scatter plot using plotnine (ggplot)
    g = (p9.ggplot(data=filtered_df,
                   mapping=p9.aes(x=x_metric,  # Use the selected x_metric
                                  y=y_metric,  # Use the selected y_metric
                                  color='color',  # Use the dynamically generated color
                                  label='method_names'))  # Label each point by the method name
         + p9.geom_point(size=3)  # Add points with no jitter, set point size
         + p9.geom_text(nudge_y=0.02, size=8)  # Add method names as labels, nudge slightly above the points
         + p9.labs(title=title, x=f"{x_metric}", y=f"{y_metric}")  # Dynamic labels for X and Y axes
         + p9.scale_color_identity()  # Use colors directly from the dataframe
         + p9.theme(legend_position='none', 
                    figure_size=(8, 8),  # Set figure size
                    axis_text=p9.element_text(size=10),   
                    axis_title_x=p9.element_text(size=12),
                    axis_title_y=p9.element_text(size=12))
    )

    # Save the plot as an image
    save_path = "./plot_images"  # Ensure this folder exists or adjust the path
    os.makedirs(save_path, exist_ok=True)  # Create directory if it doesn't exist
    filename = os.path.join(save_path, title.replace(" ", "_") + "_Similarity_Scatter.png")
    
    g.save(filename=filename, dpi=400)
    
    return filename

def benchmark_plot(benchmark_type, methods_selected, x_metric, y_metric):
    if benchmark_type == 'flexible':
        # Use general visualizer logic
        return general_visualizer_plot(methods_selected, x_metric=x_metric, y_metric=y_metric)
    elif benchmark_type == 'similarity':
        title = f"{x_metric} vs {y_metric}"
        return draw_scatter_plot_similarity(methods_selected, x_metric, y_metric, title)
    elif benchmark_type == 'Benchmark 3':
        return benchmark_3_plot(x_metric, y_metric)
    elif benchmark_type == 'Benchmark 4':
        return benchmark_4_plot(x_metric, y_metric)
    else:
        return "Invalid benchmark type selected."


def get_baseline_df(selected_methods, selected_metrics):
    df = pd.read_csv(CSV_RESULT_PATH)
    present_columns = ["method_name"] + selected_metrics
    df = df[df['method_name'].isin(selected_methods)][present_columns]
    return df

def general_visualizer(methods_selected, x_metric, y_metric):
    df = pd.read_csv(CSV_RESULT_PATH)
    filtered_df = df[df['method_name'].isin(methods_selected)]

    # Create a Seaborn lineplot with method as hue
    plt.figure(figsize=(10, 8))  # Increase figure size
    sns.lineplot(
        data=filtered_df, 
        x=x_metric, 
        y=y_metric, 
        hue="method_name",  # Different colors for different methods
        marker="o",  # Add markers to the line plot
    )
    
    # Add labels and title
    plt.xlabel(x_metric)
    plt.ylabel(y_metric)
    plt.title(f'{y_metric} vs {x_metric} for selected methods')
    plt.grid(True)
    
    # Save the plot to display it in Gradio
    plot_path = "plot.png"
    plt.savefig(plot_path)
    plt.close()
    
    return plot_path