import gradio as gr
import cv2
import numpy as np
import matplotlib.pyplot as plt

def brighten_image(image):
    # Convert the image from RGB to Grayscale
    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Apply CLAHE (with reduced intensity)
    clahe = cv2.createCLAHE(clipLimit=0.3, tileGridSize=(16, 16))  # Reduced clipLimit and increased tileGridSize
    cl1 = clahe.apply(image)

    # Create a mask for darker regions (e.g., craters)
    _, mask = cv2.threshold(image, 100, 255, cv2.THRESH_BINARY_INV)

    # Apply CLAHE only to the masked regions (darker areas)
    brightened_craters = cv2.bitwise_and(cl1, cl1, mask=mask)

    # Combine brightened craters with the original image
    combined = cv2.addWeighted(brightened_craters, 1.2, cl1, 1, 0)

    # Further brighten the final image by increasing alpha (contrast) and beta (brightness)
    brightened_image = cv2.convertScaleAbs(combined, alpha=1.1, beta=40)

    # Gradually brighten the dark spots without completely removing them
    dark_spots_mask = cv2.inRange(brightened_image, 0, 120)

    # Increase the brightness of only the dark areas
    brightened_dark_spots = cv2.addWeighted(brightened_image, 1.0, dark_spots_mask, 0.5, 50)

    # Combine the brightened dark spots with the rest of the image
    final_image = cv2.bitwise_or(brightened_image, brightened_dark_spots)

    # Return the enhanced image
    return brightened_image

# Gradio interface
def process_image(input_image):
    # Convert PIL image to OpenCV format
    input_image = np.array(input_image)
    
    # Brighten the image
    enhanced_image = brighten_image(input_image)
    
    # Convert OpenCV format back to RGB for display
    enhanced_image_rgb = cv2.cvtColor(enhanced_image, cv2.COLOR_GRAY2RGB)
    
    return enhanced_image_rgb

# Create Gradio interface
interface = gr.Interface(
    fn=process_image,  # Function to process the image
    inputs=gr.Image(type="pil"),  # Image input from the user
    outputs="image",  # Display the output image
    title="Moon PSR Image Enhancement",  # Title of the app
    description="Upload your PSR moon image to brighten dark areas using CLAHE."
)

# Launch the Gradio app
interface.launch()