import gradio as gr
import cv2
import numpy as np
from PIL import Image, ImageOps

def extract_components(image, component):
    image = np.array(image)
    
    if component == "Red Channel":
        red_channel = image.copy()
        red_channel[:, :, 1] = 0  # Zero out the green channel
        red_channel[:, :, 2] = 0  # Zero out the blue channel
        return Image.fromarray(red_channel)
    
    elif component == "Green Channel":
        green_channel = image.copy()
        green_channel[:, :, 0] = 0  # Zero out the red channel
        green_channel[:, :, 2] = 0  # Zero out the blue channel
        return Image.fromarray(green_channel)
    
    elif component == "Blue Channel":
        blue_channel = image.copy()
        blue_channel[:, :, 0] = 0  # Zero out the red channel
        blue_channel[:, :, 1] = 0  # Zero out the green channel
        return Image.fromarray(blue_channel)
    
    elif component == "Grayscale":
        grayscale = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        return Image.fromarray(grayscale)
    
    elif component == "Edges":
        grayscale = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        edges = cv2.Canny(grayscale, 100, 200)
        return Image.fromarray(edges)
    
    elif component == "Original":
        return Image.fromarray(image)
    
    else:
        return Image.fromarray(image)

# Define the Gradio interface
iface = gr.Interface(
    fn=extract_components,
    inputs=[
        gr.Image(label="Input Image"),
        gr.Radio(["Red Channel", "Green Channel", "Blue Channel", "Grayscale", "Edges", "Original"], label="Select Component")
    ],
    outputs=gr.Image(label="Output Image"),
    title="Image to Components",
    description="Extract different components from an image (e.g., RGB channels, grayscale, edges)",
    theme='NoCrypt/miku'
)

# Launch the app
iface.launch(share=True)