import gradio as gr
import cv2
import numpy as np
import cupy as cp

# CUDA kernel for custom filter (example sharpening)
custom_filter_kernel = cp.array([[0, -1, 0],
                                 [-1, 5, -1],
                                 [0, -1, 0]], dtype=cp.float32)

def edge_detection_cuda(image):
    img_gpu = cp.asarray(image, dtype=cp.uint8)
    gray_gpu = cp.asarray(cv2.cvtColor(image, cv2.COLOR_RGB2GRAY))
    edges_gpu = cv2.Canny(cp.asnumpy(gray_gpu), 100, 200)
    return cv2.cvtColor(edges_gpu, cv2.COLOR_GRAY2RGB)

def gaussian_blur_cuda(image):
    img_gpu = cp.asarray(image, dtype=cp.float32)
    blurred_gpu = cv2.GaussianBlur(cp.asnumpy(img_gpu), (15, 15), 0)
    return blurred_gpu.astype(np.uint8)

def custom_filter_cuda(image):
    img_gpu = cp.asarray(image, dtype=cp.float32)
    img_gpu = cv2.filter2D(cp.asnumpy(img_gpu), -1, cp.asnumpy(custom_filter_kernel))
    return img_gpu.astype(np.uint8)

def process_image(image, operation):
    if image is None:
        return None, None, None

    edge_result = None
    blur_result = None
    custom_result = None

    if operation == "All":
        # Run operations in parallel using CUDA Streams
        stream1 = cp.cuda.Stream()
        stream2 = cp.cuda.Stream()
        stream3 = cp.cuda.Stream()

        with stream1:
            edge_result = edge_detection_cuda(image)
        with stream2:
            blur_result = gaussian_blur_cuda(image)
        with stream3:
            custom_result = custom_filter_cuda(image)

        stream1.synchronize()
        stream2.synchronize()
        stream3.synchronize()

    elif operation == "Edge Detection":
        edge_result = edge_detection_cuda(image)
    elif operation == "Gaussian Blur":
        blur_result = gaussian_blur_cuda(image)
    elif operation == "Custom Filter":
        custom_result = custom_filter_cuda(image)

    return edge_result, blur_result, custom_result

with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Parallel Image Processing with CUDA (GPU)")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy", label="Upload Image")
            process_btn = gr.Button("Process Image", variant="primary")
        with gr.Column():
            operation = gr.Radio(["All", "Edge Detection", "Gaussian Blur", "Custom Filter"],
                                 label="Choose Operation", value="All")

    with gr.Row():
        edge_output = gr.Image(label="Edge Detection Result")
        blur_output = gr.Image(label="Gaussian Blur Result")
        custom_output = gr.Image(label="Custom Filter Result")

    process_btn.click(process_image, inputs=[image_input, operation],
                      outputs=[edge_output, blur_output, custom_output])

# Launch app
demo.launch()