import gradio as gr
import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim

def preprocess_image(image, blur_value):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # Apply Gaussian blur to reduce noise
    blurred = cv2.GaussianBlur(gray, (blur_value, blur_value), 0)
    return blurred

def compare_images(image1, image2, blur_value, technique):
    # Preprocess images
    gray1 = preprocess_image(image1, blur_value)
    gray2 = preprocess_image(image2, blur_value)
    
    # Compute SSIM between the two images
    score, diff = ssim(gray1, gray2, full=True)
    diff = (diff * 255).astype("uint8")
    
    if technique == "Adaptive Threshold":
        _, thresh = cv2.threshold(diff, 30, 255, cv2.THRESH_BINARY_INV)
    elif technique == "Otsu's Threshold":
        _, thresh = cv2.threshold(diff, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    else:  # Default to simple binary threshold
        _, thresh = cv2.threshold(diff, 50, 255, cv2.THRESH_BINARY)
    
    # Find contours of differences
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Filter out small noise using contour area threshold
    filtered_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 500]
    
    # Create a mask to isolate only the significant added object
    mask = np.zeros_like(image1)
    cv2.drawContours(mask, filtered_contours, -1, (255, 255, 255), thickness=cv2.FILLED)
    
    # Apply the mask to highlight the object added in the second image
    highlighted = cv2.bitwise_and(image2, mask)
    
    # Show the raw difference in magenta
    diff_colored = cv2.merge([np.zeros_like(diff), diff, diff])
    
    return highlighted, diff_colored

demo = gr.Interface(
    fn=compare_images,
    inputs=[
        gr.Image(type="numpy", label="Image Without Object"),
        gr.Image(type="numpy", label="Image With Object"),
        gr.Slider(minimum=1, maximum=15, step=2, value=5, label="Gaussian Blur"),
        gr.Dropdown(["Adaptive Threshold", "Otsu's Threshold", "Simple Binary"], label="Thresholding Technique", value="Adaptive Threshold")
    ],
    outputs=[
        gr.Image(type="numpy", label="Highlighted Differences"),
        gr.Image(type="numpy", label="Raw Difference (Magenta)")
    ],
    title="Object Difference Highlighter",
    description="Upload two images: one without an object and one with an object. The app will highlight only the newly added object and show the real differences in magenta."
)

demo.launch()