import numpy as np
import gradio as gr
from skimage import color
from PIL import Image
from skimage.transform import resize

def resize_image(image, target_shape=(500, 500)):
    """Resize image to target shape."""
    return resize(image, target_shape, anti_aliasing=True)

def svd_compress(image, k):
    """Compress the image using SVD by keeping only the top k singular values."""
    U, S, Vt = np.linalg.svd(image, full_matrices=False)
    compressed_image = np.dot(U[:, :k], np.dot(np.diag(S[:k]), Vt[:k, :]))
    return compressed_image

def process_image(image, k):
    """Process the uploaded image, compress it using SVD, and return the result."""
    # Convert PIL Image to NumPy array
    image_np = np.array(image)
    
    # Convert to grayscale
    gray_image = color.rgb2gray(image_np)
    
    # Resize image to 500x500 for faster processing
    resized_image = resize_image(gray_image, target_shape=(500, 500))
    
    # Compress the image
    compressed_image = svd_compress(resized_image, k)
    
    # Convert compressed image back to PIL Image for Gradio output
    compressed_image_pil = Image.fromarray((compressed_image * 255).astype(np.uint8))
    
    return compressed_image_pil

# Gradio interface
gr.Interface(fn=process_image,
             inputs=[gr.Image(type="pil"), 
                     gr.Slider(1, 100, step=1, value=50, label="Compression Rank")],
             outputs=gr.Image(type="pil"),
             title="Interactive Image Compression using SVD",
             description="Upload an image and adjust the compression rank to see the compressed version."
).launch()