import gradio as gr
import cv2
from transforms import *
import numpy as np

# Wrapper function to call the correct transform
def apply_transform(image, domain, transform):
    img, img_type = load_image(image)

    result = None  # Initialize result to avoid UnboundLocalError

    if domain == 'Spatial':
        if transform == 'Histogram Equalization':
            result = histogram_equalization(img)
        elif transform == 'Sobel Edge Detection':
            result = sobel_edge_detection(img)
        elif transform == 'Gaussian Blur':
            result = gaussian_blur(img)
        elif transform == 'Laplacian of Gaussian':
            result = laplacian_of_gaussian(img)
        elif transform == 'Median Filter':
            result = median_filter(img)
        else:
            return "Invalid spatial transform selection."
    
    elif domain == 'Frequency':
        if transform == 'Fourier Transform':
            result = fourier_transform(img)
            # Normalize the Fourier transform result to fit within 0-255 for visualization
            result = np.abs(result)
            result = 255 * result / np.max(result)
            result = result.astype(np.uint8)
        elif transform == 'Discrete Cosine Transform':
            result = discrete_cosine_transform(img)
            result = 255 * result / np.max(result)
            result = result.astype(np.uint8)
        elif transform == 'High-Pass Filter':
            result = high_pass_filter(img)
        elif transform == 'Low-Pass Filter':
            result = low_pass_filter(img)
        elif transform == 'Wavelet Transform':
            result = wavelet_transform(img)
            result = 255 * result / np.max(result)  # Normalize wavelet transform output
            result = result.astype(np.uint8)
        else:
            return "Invalid frequency transform selection."
    
    else:
        return "Invalid domain selection."

    # Ensure that result has been assigned
    if result is None:
        return "An error occurred while processing the image."

    return result



# Gradio Interface
def interface():
    image_input = gr.Image(type="numpy", label="Input Image")  # Updated input
    domain_input = gr.Dropdown(choices=["Spatial", "Frequency"], label="Transform Domain")
    transform_input = gr.Dropdown(
        choices=[
            "Histogram Equalization", "Sobel Edge Detection", "Gaussian Blur",
            "Laplacian of Gaussian", "Median Filter",
            "Fourier Transform", "Discrete Cosine Transform", "High-Pass Filter",
            "Low-Pass Filter", "Wavelet Transform"
        ], 
        label="Transform Type"
    )

    outputs = gr.Image(type="numpy", label="Transformed Image")  # Updated output

    gr.Interface(
        fn=apply_transform, 
        inputs=[image_input, domain_input, transform_input], 
        outputs=outputs, 
        title="TransformSuite: Image Filtering & Transforms",
        description="Choose between spatial or frequency domain transforms and apply various filters to the input image."
    ).launch()

if __name__ == "__main__":
    interface()