Update app.py

app.py

@@ -1,81 +1,75 @@
 import cv2
 import numpy as np
-import gradio as gr
+import matplotlib.pyplot as plt
 import pywt
 from skimage import exposure
+import gradio as gr
 from PIL import Image
+from io import BytesIO
 
-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)
+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) / 255.0
+    img_norm = img.astype(np.float32) / 65535.0
 
-    # 1. Wavelet Decomposition
-    coeffs = pywt.wavedec2(img_norm, 'bior1.3', level=3)
+    # 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
 
-    # 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)
+    # Reconstruction
+    recon = pywt.waverec2([cA3, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)], 'bior1.3')
+    recon = np.clip(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)
+    # 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)
 
-    # 5. Gamma Correction
-    p5, p95 = np.percentile(img_clahe, (5, 95))
+    # Gamma correction
+    p5, p95 = np.percentile(clahe_img, (5, 95))
     gamma = 0.7 if (p95 - p5) < 0.3 else 0.9
-    img_gamma = exposure.adjust_gamma(img_clahe, gamma=gamma)
+    gamma_img = exposure.adjust_gamma(clahe_img, 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
+    # 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)
 
-def process_image(input_img):
-    # Convert Gradio input (PIL Image) to NumPy array
-    input_img = np.array(input_img)
+    # Prepare outputs
+    original_display = (np.clip(img/np.percentile(img, 99.5), 0, 1)*255).astype(np.uint8)
     
-    # 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
+    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)
 
-# Gradio Interface
-with gr.Blocks() as demo:
-    gr.Markdown("# X-ray Enhancement (MUSICA®-Style)")
-    gr.Markdown("Upload a TIFF, JPEG, or PNG X-ray image to enhance it using multi-scale wavelet processing and adaptive contrast adjustment.")
-    
-    with gr.Row():
-        input_image = gr.Image(label="Upload X-ray (TIFF or JPEG/PNG)")
-        output_original = gr.Image(label="Original Image")
-        output_enhanced = gr.Image(label="Enhanced Image")
-    
-    submit_button = gr.Button("Enhance")
-    submit_button.click(
-        fn=process_image,
-        inputs=input_image,
-        outputs=[output_original, output_enhanced]
-    )
+    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"
+)
 
-demo.launch()
+interface.launch()