app.py

"""

Image File Compare Application

Supports TIFF, FITS and common image formats.

Provides threshold-based difference detection with scaling options.

Deployable on Hugging Face Spaces (Streamlit).

"""

import io
import streamlit as st
import numpy as np
from PIL import Image
from astropy.io import fits
import cv2
from pathlib import Path


# --- Helper Functions ---

def load_image_from_upload(uploaded_file) -> np.ndarray:
    """Load an image from a Streamlit uploaded file object."""
    name = uploaded_file.name.lower()

    if name.endswith((".fits", ".fit", ".fts")):
        return _load_fits_from_bytes(uploaded_file.getvalue())
    elif name.endswith((".tif", ".tiff")):
        return _load_tiff_from_bytes(uploaded_file.getvalue())
    else:
        img = Image.open(uploaded_file)
        return np.array(img).astype(np.float64)


def _load_fits_from_bytes(data_bytes: bytes) -> np.ndarray:
    """Load FITS from bytes."""
    with fits.open(io.BytesIO(data_bytes)) as hdul:
        for hdu in hdul:
            if hdu.data is not None:
                data = hdu.data.astype(np.float64)
                return _normalize_fits_data(data)
    raise ValueError("No image data found in FITS file")


def _normalize_fits_data(data: np.ndarray) -> np.ndarray:
    """Normalize FITS data to 0-255 range."""
    if data.ndim == 2:
        dmin, dmax = np.nanmin(data), np.nanmax(data)
        if dmax - dmin > 0:
            data = (data - dmin) / (dmax - dmin) * 255.0
        else:
            data = np.zeros_like(data)
    elif data.ndim == 3:
        if data.shape[0] in (1, 3, 4):
            data = np.moveaxis(data, 0, -1)
        dmin, dmax = np.nanmin(data), np.nanmax(data)
        if dmax - dmin > 0:
            data = (data - dmin) / (dmax - dmin) * 255.0
        else:
            data = np.zeros_like(data)
    return data


def _load_tiff_from_bytes(data_bytes: bytes) -> np.ndarray:
    """Load TIFF from bytes."""
    img = Image.open(io.BytesIO(data_bytes))
    data = np.array(img).astype(np.float64)
    if data.max() > 255:
        dmin, dmax = data.min(), data.max()
        if dmax - dmin > 0:
            data = (data - dmin) / (dmax - dmin) * 255.0
    return data


def align_images(img_ref: np.ndarray, img_to_align: np.ndarray,

                 method: str = "ECC (intensity-based)") -> tuple:
    """

    Align img_to_align to img_ref to minimize differences.

    Returns (aligned_image, info_dict).



    Methods:

      - "Feature-based (ORB)": Uses ORB keypoints + homography

      - "ECC (intensity-based)": Uses Enhanced Correlation Coefficient for sub-pixel alignment

      - "Phase correlation (translation only)": Handles pure translation/shift

    """
    # Convert to grayscale uint8 for alignment computation
    def to_gray_u8(img):
        if img.ndim == 3:
            gray = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2GRAY)
        else:
            gray = img.astype(np.uint8)
        return gray

    ref_gray = to_gray_u8(img_ref)
    align_gray = to_gray_u8(img_to_align)
    h, w = ref_gray.shape[:2]

    info = {"method": method, "success": False, "details": ""}

    if method == "Feature-based (ORB)":
        aligned, info = _align_feature_based(img_to_align, ref_gray, align_gray, h, w, info)
    elif method == "ECC (intensity-based)":
        aligned, info = _align_ecc(img_to_align, ref_gray, align_gray, h, w, info)
    elif method == "Phase correlation (translation only)":
        aligned, info = _align_phase_correlation(img_to_align, ref_gray, align_gray, h, w, info)
    else:
        aligned = img_to_align
        info["details"] = "Unknown method"

    return aligned, info


def _align_feature_based(img_to_align, ref_gray, align_gray, h, w, info):
    """Align using ORB feature detection + homography."""
    orb = cv2.ORB_create(nfeatures=5000)
    kp1, des1 = orb.detectAndCompute(ref_gray, None)
    kp2, des2 = orb.detectAndCompute(align_gray, None)

    if des1 is None or des2 is None or len(des1) < 10 or len(des2) < 10:
        info["details"] = "Not enough features detected"
        return img_to_align, info

    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False)
    matches = bf.knnMatch(des2, des1, k=2)

    # Lowe's ratio test
    good_matches = []
    for m_pair in matches:
        if len(m_pair) == 2:
            m, n = m_pair
            if m.distance < 0.75 * n.distance:
                good_matches.append(m)

    if len(good_matches) < 10:
        info["details"] = f"Only {len(good_matches)} good matches found (need >= 10)"
        return img_to_align, info

    src_pts = np.float32([kp2[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
    dst_pts = np.float32([kp1[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

    M, mask_inliers = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

    if M is None:
        info["details"] = "Homography estimation failed"
        return img_to_align, info

    inliers = int(mask_inliers.sum()) if mask_inliers is not None else 0
    info["success"] = True
    info["details"] = f"{len(good_matches)} matches, {inliers} inliers"

    if img_to_align.ndim == 2:
        aligned = cv2.warpPerspective(img_to_align, M, (w, h),
                                      flags=cv2.INTER_LINEAR,
                                      borderMode=cv2.BORDER_REFLECT)
    else:
        aligned = cv2.warpPerspective(img_to_align, M, (w, h),
                                      flags=cv2.INTER_LINEAR,
                                      borderMode=cv2.BORDER_REFLECT)

    return aligned, info


def _align_ecc(img_to_align, ref_gray, align_gray, h, w, info):
    """Align using Enhanced Correlation Coefficient (ECC) — handles rotation + translation."""
    # Use affine (6 DOF: rotation, translation, scale, shear)
    warp_matrix = np.eye(2, 3, dtype=np.float32)

    criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 200, 1e-6)

    try:
        # Downscale for initial estimate if images are large
        scale = 1.0
        if max(h, w) > 1000:
            scale = 500.0 / max(h, w)
            ref_small = cv2.resize(ref_gray, None, fx=scale, fy=scale)
            align_small = cv2.resize(align_gray, None, fx=scale, fy=scale)

            warp_small = np.eye(2, 3, dtype=np.float32)
            _, warp_small = cv2.findTransformECC(
                ref_small, align_small, warp_small, cv2.MOTION_AFFINE, criteria
            )

            # Scale the translation back
            warp_matrix = warp_small.copy()
            warp_matrix[0, 2] /= scale
            warp_matrix[1, 2] /= scale

        # Refine at full resolution
        _, warp_matrix = cv2.findTransformECC(
            ref_gray, align_gray, warp_matrix, cv2.MOTION_AFFINE, criteria
        )

        info["success"] = True
        # Extract rotation angle from the matrix
        angle = np.degrees(np.arctan2(warp_matrix[1, 0], warp_matrix[0, 0]))
        tx, ty = warp_matrix[0, 2], warp_matrix[1, 2]
        info["details"] = f"Rotation: {angle:.3f}°, Translation: ({tx:.1f}, {ty:.1f}) px"

    except cv2.error as e:
        info["details"] = f"ECC failed to converge: {str(e)}"
        return img_to_align, info

    if img_to_align.ndim == 2:
        aligned = cv2.warpAffine(img_to_align, warp_matrix, (w, h),
                                 flags=cv2.INTER_LINEAR,
                                 borderMode=cv2.BORDER_REFLECT)
    else:
        aligned = cv2.warpAffine(img_to_align, warp_matrix, (w, h),
                                 flags=cv2.INTER_LINEAR,
                                 borderMode=cv2.BORDER_REFLECT)

    return aligned, info


def _align_phase_correlation(img_to_align, ref_gray, align_gray, h, w, info):
    """Align using phase correlation — handles pure translation/shift."""
    ref_f = ref_gray.astype(np.float32)
    align_f = align_gray.astype(np.float32)

    shift, response = cv2.phaseCorrelate(ref_f, align_f)

    tx, ty = shift  # (x, y) shift
    info["success"] = True
    info["details"] = f"Translation: ({tx:.2f}, {ty:.2f}) px, confidence: {response:.4f}"

    M = np.float32([[1, 0, tx], [0, 1, ty]])

    if img_to_align.ndim == 2:
        aligned = cv2.warpAffine(img_to_align, M, (w, h),
                                 flags=cv2.INTER_LINEAR,
                                 borderMode=cv2.BORDER_REFLECT)
    else:
        aligned = cv2.warpAffine(img_to_align, M, (w, h),
                                 flags=cv2.INTER_LINEAR,
                                 borderMode=cv2.BORDER_REFLECT)

    return aligned, info


def resize_image(image: np.ndarray, target_shape: tuple, method: str) -> np.ndarray:
    """Resize image to target shape using specified interpolation method."""
    interpolation_methods = {
        "Nearest": cv2.INTER_NEAREST,
        "Bilinear": cv2.INTER_LINEAR,
        "Bicubic": cv2.INTER_CUBIC,
        "Lanczos": cv2.INTER_LANCZOS4,
        "Area": cv2.INTER_AREA,
    }
    interp = interpolation_methods.get(method, cv2.INTER_LINEAR)
    target_h, target_w = target_shape[:2]
    resized = cv2.resize(image, (target_w, target_h), interpolation=interp)
    return resized


def compute_difference(img1: np.ndarray, img2: np.ndarray, threshold: float):
    """

    Compute difference between two images.

    Returns:

        diff_image: absolute difference (float64)

        mask_image: binary mask where diff > threshold (uint8, 0 or 255)

        stats: dictionary with comparison statistics

    """
    img1_proc = img1.copy()
    img2_proc = img2.copy()

    # Strip alpha channel if present
    if img1_proc.ndim == 3 and img1_proc.shape[2] == 4:
        img1_proc = img1_proc[:, :, :3]
    if img2_proc.ndim == 3 and img2_proc.shape[2] == 4:
        img2_proc = img2_proc[:, :, :3]

    # If one is grayscale and other is color, convert both to grayscale
    if img1_proc.ndim != img2_proc.ndim:
        if img1_proc.ndim == 3:
            img1_proc = cv2.cvtColor(img1_proc.astype(np.uint8), cv2.COLOR_RGB2GRAY).astype(np.float64)
        if img2_proc.ndim == 3:
            img2_proc = cv2.cvtColor(img2_proc.astype(np.uint8), cv2.COLOR_RGB2GRAY).astype(np.float64)

    diff = np.abs(img1_proc - img2_proc)

    # For multi-channel, take the max across channels for the mask
    if diff.ndim == 3:
        diff_for_mask = np.max(diff, axis=2)
    else:
        diff_for_mask = diff

    mask = (diff_for_mask > threshold).astype(np.uint8) * 255

    # Statistics
    total_pixels = mask.shape[0] * mask.shape[1]
    diff_pixels = int(np.count_nonzero(mask))
    diff_percentage = (diff_pixels / total_pixels) * 100

    stats = {
        "total_pixels": total_pixels,
        "different_pixels": diff_pixels,
        "difference_percentage": diff_percentage,
        "max_difference": float(np.max(diff)),
        "mean_difference": float(np.mean(diff)),
    }

    return diff, mask, stats


def to_display_image(data: np.ndarray) -> np.ndarray:
    """Convert array to displayable uint8 image."""
    if data.ndim == 2:
        dmin, dmax = data.min(), data.max()
        if dmax - dmin > 0:
            normalized = ((data - dmin) / (dmax - dmin) * 255).astype(np.uint8)
        else:
            normalized = np.zeros_like(data, dtype=np.uint8)
        return normalized
    else:
        return np.clip(data, 0, 255).astype(np.uint8)


def create_colored_mask(mask: np.ndarray, diff: np.ndarray) -> np.ndarray:
    """Create a colored overlay showing differences.

    Green = no difference, Red = difference detected."""
    h, w = mask.shape[:2]
    colored = np.zeros((h, w, 3), dtype=np.uint8)
    colored[:, :, 1] = 100  # slight green background

    if diff.ndim == 3:
        diff_gray = np.max(diff, axis=2)
    else:
        diff_gray = diff

    diff_normalized = np.clip(diff_gray / max(diff_gray.max(), 1) * 255, 0, 255).astype(np.uint8)

    colored[mask > 0, 0] = 255
    colored[mask > 0, 1] = 0
    colored[mask > 0, 2] = diff_normalized[mask > 0]

    return colored


def image_to_png_bytes(img_array: np.ndarray) -> bytes:
    """Convert numpy array to PNG bytes for download."""
    img = Image.fromarray(img_array)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def image_to_tiff_bytes(img_array: np.ndarray) -> bytes:
    """Convert numpy array to TIFF bytes for download."""
    img = Image.fromarray(img_array)
    buf = io.BytesIO()
    img.save(buf, format="TIFF")
    return buf.getvalue()


# --- Streamlit App ---

def main():
    st.set_page_config(page_title="Image Compare Tool", layout="wide")
    st.title("🔍 Image File Compare")
    st.markdown("Compare two images with threshold-based difference detection. "
                "Supports **TIFF**, **FITS**, and common image formats.")

    # --- Sidebar Controls ---
    st.sidebar.header("⚙️ Settings")

    threshold = st.sidebar.slider(
        "Difference Threshold",
        min_value=0.0,
        max_value=255.0,
        value=10.0,
        step=0.5,
        help="Pixel differences below this threshold are ignored."
    )

    scaling_method = st.sidebar.selectbox(
        "Scaling Method (if sizes differ)",
        ["Bilinear", "Nearest", "Bicubic", "Lanczos", "Area"],
        help="Interpolation method used when resizing images to match dimensions."
    )

    scale_target = st.sidebar.radio(
        "Scale which image?",
        ["Scale Image 2 to match Image 1", "Scale Image 1 to match Image 2"],
        help="Choose which image gets resized when dimensions differ."
    )

    st.sidebar.markdown("---")
    st.sidebar.header("🔄 Auto-Alignment")

    enable_alignment = st.sidebar.checkbox(
        "Enable auto-alignment",
        value=False,
        help="Automatically align Image 2 to Image 1 to compensate for rotation/translation before comparison."
    )

    alignment_method = "ECC (intensity-based)"
    if enable_alignment:
        alignment_method = st.sidebar.selectbox(
            "Alignment method",
            ["ECC (intensity-based)", "Feature-based (ORB)", "Phase correlation (translation only)"],
            help="ECC: best for small rotations/shifts (sub-pixel). "
                 "ORB: best for larger rotations with texture. "
                 "Phase correlation: translation/shift only."
        )

    st.sidebar.markdown("---")
    st.sidebar.header("📊 Display Options")

    show_enhanced_diff = st.sidebar.checkbox("Enhanced difference (amplified)", value=True,
                                             help="Amplify small differences for better visibility.")
    amplify_factor = 1.0
    if show_enhanced_diff:
        amplify_factor = st.sidebar.slider("Amplification factor", 1.0, 50.0, 10.0, 1.0)

    st.sidebar.markdown("---")
    st.sidebar.header("💾 Download Format")
    download_format = st.sidebar.selectbox("Output format", ["PNG", "TIFF"])

    # --- Help & Credits ---
    st.sidebar.markdown("---")

    with st.sidebar.expander("❓ Help", expanded=False):
        st.markdown("""

## Overview



**Image File Compare** is a tool designed for **astronomical image comparison**. It helps astronomers and astrophotographers compare star field images captured by different telescopes, at different times, or with different processing pipelines.



By overlaying and differencing two images of the same region of sky, you can identify:

- **New or missing objects** (transients, variable stars, asteroids, novae)

- **Alignment/registration errors** between exposures

- **Processing artifacts** introduced by different reduction pipelines

- **Changes over time** in stellar fields (proper motion, variability)



---



## Parameters & Settings



### Difference Threshold (0–255)

Controls the sensitivity of the comparison. Pixel-level differences **below** this value are treated as "no change" and ignored.

- **Low threshold (0–5):** Very sensitive — shows noise-level differences, sensor read noise, and thermal artifacts.

- **Medium threshold (5–20):** Good default — filters out minor noise while catching real changes.

- **High threshold (20+):** Only shows large differences — useful for finding bright transients or major changes.



### Scaling Method

When two images have **different pixel dimensions** (e.g., different instruments or plate scales), one image must be resized to match. Choose the interpolation method:

- **Bilinear:** Good general-purpose method. Fast, smooth results.

- **Nearest:** No interpolation — preserves exact pixel values. Best for integer data.

- **Bicubic:** Smoother than bilinear. Slightly sharper results.

- **Lanczos:** Highest quality. Best for downsampling. Preserves fine detail.

- **Area:** Best for shrinking images. Uses pixel area relation.



### Scale Which Image?

Choose which image gets resized when dimensions differ. Typically you keep your **reference image** (Image 1) at its native resolution and scale Image 2 to match.



---



## Auto-Alignment



When images have rotational or translational offsets (common when comparing across different telescopes, mounts, or epochs), enable auto-alignment to register Image 2 to Image 1 before comparison.



### Alignment Methods



- **ECC (intensity-based):** Enhanced Correlation Coefficient method. Works by maximizing the correlation between pixel intensities. Best for **small rotations** (< 5°) and sub-pixel translations. Very accurate but can fail on large offsets.



- **Feature-based (ORB):** Detects keypoint features (stars, galaxies) in both images and matches them to compute a geometric transform (homography). Best for **larger rotations** and images with many point sources. Works well for star fields.



- **Phase correlation (translation only):** Computes the translational shift between images using Fourier analysis. Only corrects X/Y offset — does **not** handle rotation. Very fast, good for dithered exposures from the same telescope.



---



## Display Options



### Enhanced Difference (Amplified)

When enabled, the difference image is multiplied by an amplification factor to make subtle differences visible. Without this, a 1-count difference on a 0–255 scale would be nearly invisible.



### Amplification Factor (1–50×)

How much to boost the difference image. Higher values make faint differences more visible but can saturate bright differences.



---



## Download Format

Choose whether output images are saved as:

- **PNG:** Lossless, 8-bit per channel. Compatible with all viewers.

- **TIFF:** Lossless, widely used in astronomical imaging pipelines.



---



## Output Images



- **Actual Difference Image:** Absolute pixel-by-pixel difference between the two images (optionally amplified).

- **Binary Mask:** White pixels where the difference exceeds the threshold, black elsewhere. Useful for counting changed regions.

- **Colored Difference Overlay:** Red marks pixels exceeding the threshold, dark green shows matching regions. Intensity indicates the magnitude of difference.

- **Overlay on Original:** The colored overlay blended onto Image 1 at adjustable opacity, helping you localize differences in context.

""")

    st.sidebar.markdown("---")
    st.sidebar.markdown(
        "<div style='text-align: center; color: gray; font-size: 0.85em;'>"
        "Created by <b>Andy Kong</b>"
        "</div>",
        unsafe_allow_html=True,
    )

    # --- Image Input ---
    st.header("📁 Upload Images")

    col1, col2 = st.columns(2)
    with col1:
        file1 = st.file_uploader("Image 1", type=["tif", "tiff", "fits", "fit", "fts",
                                                    "png", "jpg", "jpeg", "bmp"])
    with col2:
        file2 = st.file_uploader("Image 2", type=["tif", "tiff", "fits", "fit", "fts",
                                                    "png", "jpg", "jpeg", "bmp"])

    img1_data = None
    img2_data = None
    img1_name = ""
    img2_name = ""

    if file1 and file2:
        img1_name = file1.name
        img2_name = file2.name
        try:
            img1_data = load_image_from_upload(file1)
            img2_data = load_image_from_upload(file2)
        except Exception as e:
            st.error(f"Error loading images: {e}")

    # --- Comparison ---
    if img1_data is not None and img2_data is not None:
        st.markdown("---")
        st.header("📐 Image Information")

        col1, col2 = st.columns(2)
        with col1:
            st.markdown(f"**Image 1:** `{img1_name}`")
            st.markdown(f"Shape: `{img1_data.shape}` | "
                        f"Range: [{img1_data.min():.1f}, {img1_data.max():.1f}]")
        with col2:
            st.markdown(f"**Image 2:** `{img2_name}`")
            st.markdown(f"Shape: `{img2_data.shape}` | "
                        f"Range: [{img2_data.min():.1f}, {img2_data.max():.1f}]")

        # Handle size mismatch
        if img1_data.shape[:2] != img2_data.shape[:2]:
            st.warning(f"⚠️ Image sizes differ! Image 1: {img1_data.shape[:2]}, "
                       f"Image 2: {img2_data.shape[:2]}. "
                       f"Applying **{scaling_method}** scaling.")

            if "Image 2" in scale_target:
                img2_data = resize_image(img2_data, img1_data.shape[:2], scaling_method)
            else:
                img1_data = resize_image(img1_data, img2_data.shape[:2], scaling_method)

        # Handle channel mismatch
        if img1_data.ndim != img2_data.ndim:
            st.info("Channel mismatch detected. Converting both to grayscale for comparison.")
            if img1_data.ndim == 3:
                img1_data = cv2.cvtColor(img1_data.astype(np.uint8),
                                         cv2.COLOR_RGB2GRAY).astype(np.float64)
            if img2_data.ndim == 3:
                img2_data = cv2.cvtColor(img2_data.astype(np.uint8),
                                         cv2.COLOR_RGB2GRAY).astype(np.float64)

        # Auto-alignment
        if enable_alignment:
            with st.spinner(f"Aligning images using {alignment_method}..."):
                img2_data, align_info = align_images(img1_data, img2_data, alignment_method)
            if align_info["success"]:
                st.success(f"✅ Alignment successful — {align_info['details']}")
            else:
                st.warning(f"⚠️ Alignment issue — {align_info['details']}")

        # Compute difference
        diff, mask, stats = compute_difference(img1_data, img2_data, threshold)

        # --- Results ---
        st.markdown("---")
        st.header("📊 Comparison Results")

        # Stats
        stat_cols = st.columns(4)
        stat_cols[0].metric("Different Pixels", f"{stats['different_pixels']:,}")
        stat_cols[1].metric("Difference %", f"{stats['difference_percentage']:.2f}%")
        stat_cols[2].metric("Max Difference", f"{stats['max_difference']:.1f}")
        stat_cols[3].metric("Mean Difference", f"{stats['mean_difference']:.2f}")

        # Determine download format
        if download_format == "PNG":
            ext = "png"
            mime = "image/png"
            convert_fn = image_to_png_bytes
        else:
            ext = "tiff"
            mime = "image/tiff"
            convert_fn = image_to_tiff_bytes

        st.markdown("---")

        # Display images
        st.subheader("🖼️ Source Images")
        col1, col2 = st.columns(2)
        with col1:
            st.markdown(f"**Image 1:** `{img1_name}`")
            st.image(to_display_image(img1_data), use_column_width=True)
        with col2:
            st.markdown(f"**Image 2:** `{img2_name}`")
            st.image(to_display_image(img2_data), use_column_width=True)

        st.markdown("---")

        # Prepare output images
        diff_display = diff.copy()
        if show_enhanced_diff:
            diff_display = diff_display * amplify_factor
        diff_display_uint8 = to_display_image(diff_display)
        mask_display = mask  # already uint8

        # --- Difference Image with inline download ---
        header_col, btn_col = st.columns([4, 1])
        with header_col:
            st.subheader("🔎 Actual Difference Image")
        with btn_col:
            st.download_button(
                label=f"⬇️ .{ext}",
                data=convert_fn(diff_display_uint8),
                file_name=f"difference.{ext}",
                mime=mime,
                key="dl_diff",
            )
        st.image(diff_display_uint8, use_column_width=True)

        st.markdown("---")

        # --- Binary Mask with inline download ---
        header_col, btn_col = st.columns([4, 1])
        with header_col:
            st.subheader("🎭 Binary Mask")
        with btn_col:
            st.download_button(
                label=f"⬇️ .{ext}",
                data=convert_fn(mask_display),
                file_name=f"mask.{ext}",
                mime=mime,
                key="dl_mask",
            )
        st.image(mask_display, use_column_width=True,
                 caption="White = difference > threshold")

        st.markdown("---")

        # --- Colored Overlay with inline download ---
        colored_mask = create_colored_mask(mask, diff)
        header_col, btn_col = st.columns([4, 1])
        with header_col:
            st.subheader("🎨 Colored Difference Overlay")
        with btn_col:
            st.download_button(
                label=f"⬇️ .{ext}",
                data=convert_fn(colored_mask),
                file_name=f"colored_overlay.{ext}",
                mime=mime,
                key="dl_colored",
            )
        st.image(colored_mask, use_column_width=True,
                 caption="Red = pixels exceeding threshold | Dark green = within threshold")

        st.markdown("---")

        # --- Overlay on Original with inline download ---
        overlay_alpha = st.slider("Overlay opacity", 0.0, 1.0, 0.4, 0.05)

        base_img = to_display_image(img1_data)
        if base_img.ndim == 2:
            base_img = cv2.cvtColor(base_img, cv2.COLOR_GRAY2RGB)
        elif base_img.ndim == 3 and base_img.shape[2] == 4:
            base_img = cv2.cvtColor(base_img, cv2.COLOR_RGBA2RGB)

        overlay = cv2.addWeighted(base_img, 1 - overlay_alpha,
                                  colored_mask, overlay_alpha, 0)

        header_col, btn_col = st.columns([4, 1])
        with header_col:
            st.subheader("📷 Overlay on Original")
        with btn_col:
            st.download_button(
                label=f"⬇️ .{ext}",
                data=convert_fn(overlay),
                file_name=f"overlay_on_original.{ext}",
                mime=mime,
                key="dl_overlay",
            )
        st.image(overlay, use_column_width=True,
                 caption="Image 1 with difference overlay")

    else:
        st.info("👆 Please upload two images to compare.")


if __name__ == "__main__":
    main()