app.py

import cv2
import numpy as np
import matplotlib.pyplot as plt
import pywt
from skimage import exposure
import gradio as gr
from PIL import Image
from io import BytesIO

def process_tiff(file):
    # Load 16-bit TIFF
    img = cv2.imread(file.name, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise gr.Error("Could not read TIFF file. Ensure it's 16-bit grayscale.")
    
    # Normalize to [0, 1]
    img_norm = img.astype(np.float32) / 65535.0

    # Wavelet decomposition
    try:
        coeffs = pywt.wavedec2(img_norm, 'bior1.3', level=3)
    except ValueError as e:
        raise gr.Error(f"Image dimensions must be divisible by 8. {str(e)}")
    
    # Processing coefficients (same as original)
    cA3, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1) = coeffs
    cD1 = pywt.threshold(cD1, 0.05*np.max(cD1), 'soft')
    cD2 = pywt.threshold(cD2, 0.07*np.max(cD2), 'soft')
    cH1 *= 1.2; cV1 *= 1.2

    # Reconstruction
    recon = pywt.waverec2([cA3, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)], 'bior1.3')
    recon = np.clip(recon, 0, 1)

    # CLAHE
    entropy = -np.sum(recon * np.log2(recon + 1e-7))
    clahe_img = exposure.equalize_adapthist(recon, clip_limit=0.02 if entropy >7 else 0.05, kernel_size=64)

    # Gamma correction
    p5, p95 = np.percentile(clahe_img, (5, 95))
    gamma = 0.7 if (p95 - p5) < 0.3 else 0.9
    gamma_img = exposure.adjust_gamma(clahe_img, gamma)

    # Sharpening
    sharp = cv2.detailEnhance(cv2.cvtColor((gamma_img*255).astype(np.uint8), cv2.COLOR_GRAY2BGR), 
                            sigma_s=12, sigma_r=0.15)
    sharp = cv2.cvtColor(sharp, cv2.COLOR_BGR2GRAY)

    # Prepare outputs
    original_display = (np.clip(img/np.percentile(img, 99.5), 0, 1)*255).astype(np.uint8)
    
    buf = BytesIO()
    plt.hist(sharp.ravel(), bins=256, range=(0, 255))
    plt.title("Enhanced Histogram")
    plt.savefig(buf, format='png', bbox_inches='tight')
    plt.close()
    hist_img = Image.open(buf)

    return original_display, sharp, hist_img

# Create Gradio interface
interface = gr.Interface(
    fn=process_tiff,
    inputs=gr.File(label="Upload TIFF", file_types=[".tif", ".tiff"]),
    outputs=[
        gr.Image(label="Original (Clipped)"),
        gr.Image(label="Enhanced Image"),
        gr.Image(label="Histogram")
    ],
    title="MUSICA Image Enhancement",
    description="Upload 16-bit grayscale TIFF for enhancement using wavelet-based processing",
    allow_flagging="never"
)

interface.launch()