main.py

import tkinter as tk
from tkinter import filedialog, messagebox, Scale, IntVar, Frame
from tkinter.ttk import Progressbar
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import cv2
import csv
import os
import threading
import time
import tracemalloc
import numpy as np
from datetime import datetime
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as ski_psnr
from jpeg_encoder import JPEGEncoder
from jpeg_decoder import JPEGDecoder
import matplotlib.pyplot as plt


class ImageQualityAssessor:
    @staticmethod
    def calculate_psnr(original, compressed):
        return ski_psnr(original, compressed, data_range=1.0)

    @staticmethod
    def calculate_ssim(original, compressed):
        return ssim(original, compressed, channel_axis=2, data_range=1.0)

    @staticmethod
    def calculate_mse(original, compressed):
        return np.mean((original - compressed) ** 2)

    @staticmethod
    def calculate_mae(original, compressed):
        return np.mean(np.abs(original - compressed))

    @classmethod
    def assess_quality(cls, original, compressed):
        # Ensure both images have the same dimensions
        h, w = compressed.shape[:2]
        original = cv2.resize(original, (w, h))

        # Convert to float and scale to [0, 1]
        original = original.astype(np.float32) / 255.0
        compressed = compressed.astype(np.float32) / 255.0

        return {
            'PSNR': cls.calculate_psnr(original, compressed),
            'SSIM': cls.calculate_ssim(original, compressed),
            'MSE': cls.calculate_mse(original, compressed),
            'MAE': cls.calculate_mae(original, compressed)
        }


class CompressionApp:
    def __init__(self, master):
        self.master = master
        master.title("Compression/Decompression")
        self.quality_assessor = ImageQualityAssessor()
        self.quality_var = IntVar()
        self.quality_var.set(50)

        self.setup_ui()

    def setup_ui(self):

        self.setup_jpeg_frame()
        self.setup_quality_slider()
        self.setup_progress_bar()



    def setup_jpeg_frame(self):
        jpeg_frame = Frame(self.master)
        jpeg_frame.pack(pady=10)

        jpeg_encode_button = tk.Button(jpeg_frame, text="JPEG Encode File", command=self.open_file_to_jpeg_encode)
        jpeg_encode_button.pack(side=tk.LEFT, padx=5)

        jpeg_decode_button = tk.Button(jpeg_frame, text="JPEG Decode File", command=self.open_file_to_jpeg_decode)
        jpeg_decode_button.pack(side=tk.LEFT, padx=5)



    def setup_quality_slider(self):
        quality_frame = Frame(self.master)
        quality_frame.pack(pady=10)

        quality_slider = Scale(quality_frame, from_=1, to=100, orient=tk.HORIZONTAL,
                               label="Quality", variable=self.quality_var, length=200)
        quality_slider.pack()

    def setup_progress_bar(self):
        self.progress_bar = Progressbar(self.master, orient=tk.HORIZONTAL, length=300, mode='determinate')
        self.progress_bar.pack(pady=10)

        self.progress_label = tk.Label(self.master, text="")
        self.progress_label.pack(pady=5)

    def update_progress(self, value):
        self.master.after(0, self._update_progress_gui, value)

    def _update_progress_gui(self, value):
        self.progress_bar['value'] = value
        self.progress_label.config(text=f"{value:.2f}%")
    def save_metrics_to_csv(self, input_file, output_file, elapsed_time, memory_current, memory_peak, 
                            compression_ratio, size_reduction, quality_metrics):
        csv_file = 'performance_metrics.csv'
        file_exists = os.path.isfile(csv_file)
        
        with open(csv_file, mode='a', newline='') as file:
            writer = csv.writer(file)
            
            if not file_exists:
                writer.writerow(['Timestamp', 'Input File', 'Output File', 'Elapsed Time (s)', 
                                'Current Memory (MB)', 'Peak Memory (MB)', 'Compression Ratio', 
                                'Size Reduction (%)', 'PSNR', 'SSIM', 'MSE', 'MAE'])
            
            writer.writerow([
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                input_file,
                output_file,
                f"{elapsed_time:.4f}",
                f"{memory_current:.2f}",
                f"{memory_peak:.2f}",
                f"{compression_ratio:.2f}",
                f"{size_reduction:.2f}",
                f"{quality_metrics['PSNR']:.4f}",
                f"{quality_metrics['SSIM']:.4f}",
                f"{quality_metrics['MSE']:.4f}",
                f"{quality_metrics['MAE']:.4f}"
            ])


    def display_images(self, original, compressed, info_text, compr_placeholder):
        top = tk.Toplevel(self.master)
        top.title("Image Comparison")
        original = cv2.cvtColor(original, cv2.COLOR_BGR2RGB)
        compressed = cv2.cvtColor(compressed, cv2.COLOR_BGR2RGB)

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))

        ax1.imshow(original)
        ax1.set_title("Original Image")
        ax1.axis('off')

        ax2.imshow(compressed)
        ax2.set_title(f"{compr_placeholder} Image")
        ax2.axis('off')

        plt.tight_layout()

        canvas = FigureCanvasTkAgg(fig, master=top)
        canvas_widget = canvas.get_tk_widget()
        canvas_widget.pack()

        text_widget = tk.Text(top, height=10, wrap=tk.WORD)
        text_widget.pack(padx=10, pady=10, fill=tk.BOTH, expand=True)
        text_widget.insert(tk.END, info_text)
        text_widget.config(state=tk.DISABLED)

        close_button = tk.Button(top, text="Close", command=top.destroy)
        close_button.pack(pady=10)

        canvas.draw()

    def open_file_to_jpeg_encode(self):
        quality = self.quality_var.get()

        if file_path := filedialog.askopenfilename(
                filetypes=[("BMP Files", "*.bmp"), ("JPEG", "*.jpeg"), ("jpg", "*.jpg"), ("png", "*.png")]
        ):
            save_path = filedialog.asksaveasfilename(
                defaultextension=".jpg", filetypes=[("jpg", "*.jpg")]
            )
            if save_path:
                self.progress_bar['value'] = 0
                self.progress_label.config(text="0.00%")

                def encode_thread():
                    start_time = time.time()

                    tracemalloc.start()
                    encoder = JPEGEncoder(quality, progress_callback=self.update_progress)
                    image = encoder.load_image(file_path)
                    if image is None:
                        return
                    prepared_image = encoder.prepare_image(image)
                    prepared_image = cv2.cvtColor(prepared_image, cv2.COLOR_BGR2YCrCb)
                    encoder.encode(prepared_image, save_path)

                    current, peak = tracemalloc.get_traced_memory()
                    tracemalloc.stop()
                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    memory_current = current / 1024 / 1024
                    memory_peak = peak / 1024 / 1024

                    compression_ratio = os.path.getsize(file_path) / os.path.getsize(save_path)
                    quality_metrics = self.quality_assessor.assess_quality(image, cv2.imread(save_path))
                    metrics = "\n".join([f"{key}: {value:.4f}" for key, value in quality_metrics.items()])
                    size_percentage = (1 - 1 / compression_ratio) * 100
                    self.save_metrics_to_csv(file_path, save_path, elapsed_time, memory_current, memory_peak, 
                                         compression_ratio, size_percentage, quality_metrics)
                    self.master.after(0, lambda: self.display_images(image, cv2.imread(save_path),
                                                                     f"Image encoded and saved as {save_path}\n"
                                                                     f"Time taken: {elapsed_time:.4f} seconds\n"
                                                                     f"Current memory usage: {memory_current:.2f} MB\n"
                                                                     f"Peak memory usage: {memory_peak:.2f} MB\n"
                                                                     f"The {save_path} file is {compression_ratio:.2f} times smaller than the {file_path} file.\n"
                                                                     f"This is a {size_percentage:.2f}% reduction in size.\n"
                                                                     f"\nQuality Metrics:\n{metrics}", "Decompressed ", ))

                threading.Thread(target=encode_thread, daemon=True).start()

    def open_file_to_jpeg_decode(self):
        if file_path := filedialog.askopenfilename(filetypes=[("JPEG", "*.jpeg"), ("jpg", "*.jpg")]):
            save_path = filedialog.asksaveasfilename(defaultextension=".jpeg", filetypes=[("JPEG Files", "*.jpeg")])
            if save_path:
                self.progress_bar['value'] = 0
                self.progress_label.config(text="0.00%")

                def decode_thread():
                    start_time = time.time()
                    decoder = JPEGDecoder(progress_callback=self.update_progress)
                    decoder.decode(file_path, save_path)
                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    original = cv2.imread(file_path)
                    self.master.after(0, lambda: self.display_images(original, cv2.imread(save_path),
                                                                     f"Image decoded and saved as {save_path}\nTime taken: {elapsed_time:.4f} seconds",
                                                                     "Decompressed "))

                threading.Thread(target=decode_thread, daemon=True).start()



if __name__ == "__main__":
    root = tk.Tk()
    app = CompressionApp(root)
    root.mainloop()