Update app.py

app.py

@@ -4,33 +4,48 @@ from skimage import io, color
 from numpy.linalg import norm
 from PIL import Image
 
-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 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(image, k):
-    """Process the uploaded image, compress it using SVD, and return the result."""
+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)
     
-    # Convert to grayscale
-    gray_image = color.rgb2gray(image_np)
+    # 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(gray_image, k)
+    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 * 255).astype(np.uint8))
+    compressed_image_pil = Image.fromarray(compressed_image)
     
     return compressed_image_pil
 
 # Gradio interface
-gr_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.")
+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()