app.py

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 create_dramatic_magenta(image1, diff):
    """Create a more dramatic magenta overlay to highlight differences"""
    # Create a more intense magenta by boosting the red and blue channels
    diff_colored = cv2.absdiff(image1, diff)
    
    # Normalize to enhance contrast
    diff_normalized = cv2.normalize(diff_colored, None, 0, 255, cv2.NORM_MINMAX)
    
    # Amplify the red channel for more dramatic magenta
    diff_normalized[:, :, 0] = 0  # Remove blue
    diff_normalized[:, :, 1] = 0  # Remove green
    diff_normalized[:, :, 2] = np.clip(diff_normalized[:, :, 2] * 2, 0, 255)  # Boost red
    
    # Create more dramatic overlay with higher contrast
    overlay = cv2.addWeighted(image1, 0.5, diff_normalized, 0.8, 0)
    
    return overlay

def background_subtraction(image1, image2):
    subtractor = cv2.createBackgroundSubtractorMOG2()
    fgmask1 = subtractor.apply(image1)
    fgmask2 = subtractor.apply(image2)
    diff = cv2.absdiff(fgmask1, fgmask2)
    
    # Create a binary mask
    _, mask = cv2.threshold(diff, 30, 255, cv2.THRESH_BINARY)
    
    # Create highlighted differences
    highlighted = cv2.bitwise_and(image2, image2, mask=mask)
    
    # Create raw difference overlay with dramatic magenta
    raw_overlay = create_dramatic_magenta(image1, image2)
    
    # Create a blended image
    blended = cv2.addWeighted(image1, 0.5, image2, 0.5, 0)
    
    # Create a composite using the mask
    composite = image1.copy()
    composite[mask > 0] = image2[mask > 0]
    
    # Create final difference overlay with dramatic magenta
    final_overlay = create_dramatic_magenta(image1, composite)
    
    return blended, raw_overlay, highlighted, mask, composite, final_overlay

def optical_flow(image1, image2):
    gray1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)
    flow = cv2.calcOpticalFlowFarneback(gray1, gray2, None, 0.5, 3, 15, 3, 5, 1.2, 0)
    mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
    hsv = np.zeros_like(image1)
    hsv[..., 1] = 255
    hsv[..., 0] = ang * 180 / np.pi / 2
    hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
    
    # Create mask from magnitude
    mask = cv2.threshold(hsv[..., 2], 30, 255, cv2.THRESH_BINARY)[1].astype(np.uint8)
    
    # Create highlighted differences
    highlighted = cv2.bitwise_and(image2, image2, mask=mask)
    
    # Create raw difference overlay with dramatic magenta
    raw_overlay = create_dramatic_magenta(image1, image2)
    
    # Create a blended image
    blended = cv2.addWeighted(image1, 0.5, image2, 0.5, 0)
    
    # Create a composite using the mask
    composite = image1.copy()
    composite[mask > 0] = image2[mask > 0]
    
    # Create final difference overlay with dramatic magenta
    final_overlay = create_dramatic_magenta(image1, composite)
    
    return blended, raw_overlay, highlighted, mask, composite, final_overlay

def feature_matching(image1, image2):
    # Use SSIM as a fallback for feature matching since the original implementation doesn't give us a good mask
    return compare_ssim(image1, image2, 5, "Adaptive Threshold", 30)

def compare_ssim(image1, image2, blur_value, technique, threshold_value):
    gray1 = preprocess_image(image1, blur_value)
    gray2 = preprocess_image(image2, blur_value)
    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:
        _, thresh = cv2.threshold(diff, threshold_value, 255, cv2.THRESH_BINARY)
    
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    filtered_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 500]
    mask = np.zeros_like(gray1, dtype=np.uint8)
    cv2.drawContours(mask, filtered_contours, -1, 255, thickness=cv2.FILLED)
    
    # Create highlighted differences
    highlighted = cv2.bitwise_and(image2, image2, mask=mask)
    
    # Create raw difference overlay with dramatic magenta
    raw_overlay = create_dramatic_magenta(image1, image2)
    
    # Create a blended image
    blended = cv2.addWeighted(image1, 0.5, image2, 0.5, 0)
    
    # Create a composite using the mask
    composite = image1.copy()
    mask_3channel = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
    masked_obj = cv2.bitwise_and(image2, mask_3channel)
    masked_bg = cv2.bitwise_and(image1, cv2.bitwise_not(mask_3channel))
    composite = cv2.add(masked_bg, masked_obj)
    
    # Create final difference overlay with dramatic magenta
    final_overlay = create_dramatic_magenta(image1, composite)
    
    return blended, raw_overlay, highlighted, mask, composite, final_overlay

def compare_images(image1, image2, blur_value, technique, threshold_value, method):
    if method == "Background Subtraction":
        return background_subtraction(image1, image2)
    elif method == "Optical Flow":
        return optical_flow(image1, image2)
    elif method == "Feature Matching":
        return feature_matching(image1, image2)
    else:  # SSIM
        return compare_ssim(image1, image2, blur_value, technique, threshold_value)

def update_threshold_visibility(technique):
    return gr.update(visible=(technique == "Simple Binary"))

with gr.Blocks() as demo:
    gr.Markdown("# Object Difference Highlighter\nUpload 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 overlayed on the original image.")
    
    with gr.Row():
        img1 = gr.Image(type="numpy", label="Image Without Object (Scene)")
        img2 = gr.Image(type="numpy", label="Image With Object")
    
    blur_slider = gr.Slider(minimum=1, maximum=15, step=2, value=5, label="Gaussian Blur")
    technique_dropdown = gr.Dropdown(["Adaptive Threshold", "Otsu's Threshold", "Simple Binary"], label="Thresholding Technique", value="Adaptive Threshold", interactive=True)
    threshold_slider = gr.Slider(minimum=0, maximum=255, step=1, value=50, label="Threshold Value", visible=False)
    method_dropdown = gr.Dropdown(["SSIM", "Background Subtraction", "Optical Flow", "Feature Matching"], label="Comparison Method", value="SSIM", interactive=True)
    
    technique_dropdown.change(update_threshold_visibility, inputs=[technique_dropdown], outputs=[threshold_slider])
    
    # Row 1 - Blend and Raw Difference
    with gr.Row():
        output1 = gr.Image(type="numpy", label="Blended Image")
        output2 = gr.Image(type="numpy", label="Raw Difference Overlay (Magenta)")
    
    # Row 2 - Algorithmic Differences and Mask
    with gr.Row():
        output3 = gr.Image(type="numpy", label="Highlighted Differences")
        output4 = gr.Image(type="numpy", label="Black & White Mask")
    
    # Row 3 - Composite and Final Difference
    with gr.Row():
        output5 = gr.Image(type="numpy", label="Composite (Scene + Masked Object)")
        output6 = gr.Image(type="numpy", label="Final Difference Overlay (Magenta)")
    
    btn = gr.Button("Process")
    btn.click(compare_images, inputs=[img1, img2, blur_slider, technique_dropdown, threshold_slider, method_dropdown], outputs=[output1, output2, output3, output4, output5, output6])

demo.launch()