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

def load_image(image_path):
    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: Unable to load image at {image_path}")
        return None
    return image
def filter_sky_contours(mask, image_height, min_area=8000, max_area=None, height_ratio=0.2):
    """
    Filter contours that are likely to be the sky based on area and position.
    
    Parameters:
    - mask: Binary mask where the sky is white and the rest is black.
    - image_height: Height of the original image.
    - min_area: Minimum area of a contour to be considered as sky.
    - max_area: Maximum area of a contour to be considered as sky.
    - height_ratio: The ratio of the image height where we expect the sky to be located.
    
    Returns:
    - sky_mask: Refined binary mask with filtered sky regions.
    """
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    sky_mask = np.zeros_like(mask)
    for contour in contours:
        area = cv2.contourArea(contour)
        x, y, w, h = cv2.boundingRect(contour)
        aspect_ratio = w / float(h)

        convex_hull_area = cv2.contourArea(cv2.convexHull(contour))
        if convex_hull_area > 0:  # Ensure the denominator is not zero
            smoothness = area / convex_hull_area
        else:
            smoothness = 1  # Or some other default value that makes sense for your application

        if (area > min_area) and (max_area is None or area < max_area) and (y < image_height * height_ratio):
            if aspect_ratio > 1 and smoothness > 0.5:  # Adjust thresholds as necessary
                cv2.drawContours(sky_mask, [contour], -1, (255), thickness=cv2.FILLED)
    return sky_mask
def detect_edges(image):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Apply bilateral filter to preserve edges while reducing noise
    bilateral = cv2.bilateralFilter(gray, d=9, sigmaColor=75, sigmaSpace=75)
    
    # Use a combination of Sobel filters to find gradients
    sobelx = cv2.Sobel(bilateral, cv2.CV_64F, 1, 0, ksize=3)
    sobely = cv2.Sobel(bilateral, cv2.CV_64F, 0, 1, ksize=3)
    
    # Combine the gradients
    gradient_magnitude = cv2.magnitude(sobelx, sobely)
    gradient_magnitude = np.uint8(gradient_magnitude)
    
    # Apply thresholding to get binary result
    _, edge_binary = cv2.threshold(gradient_magnitude, 20, 255, cv2.THRESH_BINARY)
    
    # Optionally, you can apply dilation followed by erosion to close gaps
    kernel = np.ones((3,3), np.uint8)
    edge_dilated = cv2.dilate(edge_binary, kernel, iterations=1)
    edge_processed = cv2.erode(edge_dilated, kernel, iterations=1)
    
    return edge_processed



def adaptive_threshold_sky(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    return cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                 cv2.THRESH_BINARY,21,2)

def refine_mask(mask):
    kernel = np.ones((35,35), np.uint8)
    return cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

def segment_sky(image):
    # Detect edges.
    edges = detect_edges(image)

    # Invert the edges.
    edges_inv = cv2.bitwise_not(edges)

    # Create a color mask using adaptive thresholding.
    color_mask = adaptive_threshold_sky(image)

    # Combine the color mask with the inverted edges to get an initial sky mask.
    combined_mask = cv2.bitwise_and(color_mask, edges_inv)

    # Filter contours that are likely to be the sky and refine the mask.
    refined_mask_contour = filter_sky_contours(combined_mask, image.shape[0])

    # Optionally, you can further refine the mask with morphological operations if needed.
    refined_mask_morph = refine_mask(refined_mask_contour)

    # Apply the final refined mask to segment the sky.
    segmented_sky = cv2.bitwise_and(image, image, mask=refined_mask_morph)
    contours, _ = cv2.findContours(refined_mask_morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return segmented_sky, contours


def display_results_with_contours(original, segmented, edges, contours):
    # Create a copy of the original image to draw contours on.
    image_with_contours = original.copy()
    cv2.drawContours(image_with_contours, contours, -1, (0, 255, 0), 3)

    plt.figure(figsize=(12, 8))

    plt.subplot(1, 4, 1)
    plt.imshow(cv2.cvtColor(original, cv2.COLOR_BGR2RGB))
    plt.title('Original Image')

    plt.subplot(1, 4, 2)
    plt.imshow(edges, cmap='gray')
    plt.title('Edges')

    plt.subplot(1, 4, 3)
    plt.imshow(cv2.cvtColor(segmented, cv2.COLOR_BGR2RGB))
    plt.title('Segmented Sky')

    plt.subplot(1, 4, 4)
    plt.imshow(cv2.cvtColor(image_with_contours, cv2.COLOR_BGR2RGB))
    plt.title('Contours')

    plt.show()

def process_and_display(image):
    # Convert the uploaded PIL image to an OpenCV format
    image = np.array(image)
    image = image[:, :, :3]  # Ensure the image is in RGB format
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    # Segment the sky
    segmented_sky, _ = segment_sky(image)

    # Prepare the segmented sky image for display in Gradio
    segmented_rgb = cv2.cvtColor(segmented_sky, cv2.COLOR_BGR2RGB)

    return segmented_rgb

# Set up the Gradio interface
iface = gr.Interface(
    fn=process_and_display,
    inputs=gr.Image(),
    outputs=gr.Image(label="Segmented Sky"),
    title="Sky Pixel Identification",
    description="Upload an image to see the segmented sky."
)

# Run the app
iface.launch()