import cv2
import numpy as np
import gradio as gr
import pywt
from skimage import exposure
from PIL import Image

def musica_enhancement(img):
    # Convert PIL Image to NumPy array
    img = np.array(img)
    
    # Convert from RGB to grayscale if needed
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    
    # Normalize to [0, 1]
    img_norm = img.astype(np.float32) / 255.0

    # 1. Wavelet Decomposition
    coeffs = pywt.wavedec2(img_norm, 'bior1.3', level=3)
    cA3, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1) = coeffs

    # 2. Frequency Band Processing
    cD1 = pywt.threshold(cD1, 0.05*np.max(cD1), mode='soft')
    cD2 = pywt.threshold(cD2, 0.07*np.max(cD2), mode='soft')
    cH1 = cH1 * 1.2
    cV1 = cV1 * 1.2

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

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

    # 5. Gamma Correction
    p5, p95 = np.percentile(img_clahe, (5, 95))
    gamma = 0.7 if (p95 - p5) < 0.3 else 0.9
    img_gamma = exposure.adjust_gamma(img_clahe, gamma=gamma)

    # 6. Edge Enhancement
    img_gamma_8bit = (img_gamma * 255).astype(np.uint8)
    img_bgr = cv2.cvtColor(img_gamma_8bit, cv2.COLOR_GRAY2BGR)
    img_sharp = cv2.detailEnhance(img_bgr, sigma_s=12, sigma_r=0.15)
    img_sharp = cv2.cvtColor(img_sharp, cv2.COLOR_BGR2RGB)
    
    return img_sharp

def process_image(input_img):
    # Convert Gradio input (PIL Image) to NumPy array
    input_img = np.array(input_img)
    
    # Enhance the image
    enhanced_img = musica_enhancement(input_img)
    
    # Convert enhanced image back to PIL for display
    enhanced_img = Image.fromarray(enhanced_img)
    
    return input_img, enhanced_img

# Gradio Interface
demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(label="Upload X-ray (TIFF or JPEG/PNG)"),
    outputs=[
        gr.Image(label="Original Image"), 
        gr.Image(label="Enhanced Image")
    ],
    title="X-ray Enhancement (MUSICA®-Style)",
    description="Upload a TIFF, JPEG, or PNG X-ray image to enhance it using multi-scale wavelet processing and adaptive contrast adjustment."
)

demo.launch()