plot_comparison.py

import pandas as pd
import matplotlib.pyplot as plt
import os
def plot_rd_curves(fractal_csv, wavelet_csv, output_dir, metrics):
    """
    Generates and saves Rate-Distortion (PSNR or SSIM vs. Bitrate) plots
    to compare fractal and wavelet compression performance.
    """
    try:
        # The main.py output contains fractal data
        df_fractal = pd.read_csv(fractal_csv)
        # The wavelet_test.py output contains wavelet data
        df_wavelet = pd.read_csv(wavelet_csv)
    except FileNotFoundError as e:
        print(f"Error: Could not find a metrics file. {e}")
        print("Please run 'python main.py' and 'python wavelet_test.py' first to generate the CSV files.")
        return

    os.makedirs(output_dir, exist_ok=True)
    print(f"Plots will be saved to: {output_dir}")

    # Get the unique datasets from both files
    datasets = pd.concat([df_fractal['dataset'], df_wavelet['dataset']]).unique()

    for metric in metrics:
        print(f"\n--- Generating {metric.upper()} vs. Bitrate plots ---")
        for dataset in datasets:
            plt.figure(figsize=(10, 7))

            # Filter data for the current dataset
            fractal_data = df_fractal[df_fractal['dataset'] == dataset]
            wavelet_data = df_wavelet[df_wavelet['dataset'] == dataset]

            # Plotting
            plt.scatter(fractal_data['bitrate'], fractal_data[metric], label='Fractal', marker='o')
            plt.scatter(wavelet_data['bitrate'], wavelet_data[metric], label='Wavelet', marker='x')

            # Formatting the plot
            metric_name = metric.upper()
            y_label = f'{metric_name} (dB)' if metric_name == 'PSNR' else metric_name
            plt.title(f'{metric_name} vs. Bitrate Comparison for "{dataset}" Dataset')
            plt.xlabel('Bitrate (bits per pixel)')
            plt.ylabel(y_label)
            plt.grid(True, which='both', linestyle='--')
            plt.legend()

            # Save the plot
            plot_filename = os.path.join(output_dir, f'{dataset}_{metric}_comparison.png')
            plt.savefig(plot_filename)
            print(f"Saved plot: {plot_filename}")
            plt.close()

if __name__ == "__main__":
    # Assumes main.py has been run to generate fractal metrics
    fractal_metrics_file = 'results/metrics/fractal_cuda_parallel.csv'
    wavelet_metrics_file = 'results/metrics/wavelet.csv'
    plot_output_dir = 'results/plots/'
    metrics_to_plot = ['psnr', 'ssim']
    plot_rd_curves(fractal_metrics_file, wavelet_metrics_file, plot_output_dir, metrics=metrics_to_plot)