import numpy as np
import gradio as gr
from skimage import io, color
from numpy.linalg import norm
from PIL import Image

def svd_compress_color(image, k):
    """Compress the color image using SVD by keeping only the top k singular values for each channel."""
    # Initialize an empty array for the compressed image
    compressed_image = np.zeros_like(image, dtype=np.float64)
    
    # Process each color channel separately
    for channel in range(3):  # Assuming RGB
        U, S, Vt = np.linalg.svd(image[:, :, channel], full_matrices=False)
        compressed_channel = np.dot(U[:, :k], np.dot(np.diag(S[:k]), Vt[:k, :]))
        compressed_image[:, :, channel] = compressed_channel
    
    # Clip values to valid range [0, 255] and convert to uint8
    compressed_image = np.clip(compressed_image, 0, 255)
    return compressed_image.astype(np.uint8)

def process_image_color(image, k):
    """Process the uploaded image, compress it using SVD for each color channel, and return the result."""
    # Convert PIL Image to NumPy array
    image_np = np.array(image)
    
    # Ensure the image has three color channels
    if image_np.ndim != 3 or image_np.shape[2] != 3:
        raise ValueError("Input image must be a color image with 3 channels (RGB).")
    
    # Compress the image
    compressed_image = svd_compress_color(image_np, k)
    
    # Convert compressed image back to PIL Image for Gradio output
    compressed_image_pil = Image.fromarray(compressed_image)
    
    return compressed_image_pil

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

gr_interface.launch()