Update app.py

app.py

@@ -1,75 +1,125 @@
 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
+import matplotlib.pyplot as plt
 
 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.")
-    
+    # Read file content
+    try:
+        img = cv2.imread(file.name, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_GRAYSCALE)
+        if img is None:
+            raise ValueError("Invalid or corrupted TIFF file")
+    except Exception as e:
+        raise gr.Error(f"Error reading file: {str(e)}")
+
     # Normalize to [0, 1]
     img_norm = img.astype(np.float32) / 65535.0
 
-    # Wavelet decomposition
+    # Check dimensions for wavelet transform
+    if img.shape[0] % 8 != 0 or img.shape[1] % 8 != 0:
+        raise gr.Error("Image dimensions must be divisible by 8 for wavelet processing")
+
     try:
+        # Wavelet decomposition
         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)
+        cA3, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1) = coeffs
+
+        # Processing coefficients
+        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)
+
+    except Exception as e:
+        raise gr.Error(f"Processing error: {str(e)}")
 
     # Prepare outputs
-    original_display = (np.clip(img/np.percentile(img, 99.5), 0, 1)*255).astype(np.uint8)
+    original_display = (np.clip(img / np.percentile(img, 99.5), 0, 1) * 255).astype(np.uint8)
+    
+    # Create histogram plot
+    fig, ax = plt.subplots()
+    ax.hist(sharp.ravel(), bins=256, range=(0, 255))
+    ax.set_title("Enhanced Histogram")
+    ax.set_xlabel("Pixel Value")
+    ax.set_ylabel("Frequency")
     
+    # Convert plot to PIL Image
     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()
+    plt.close(fig)
     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()
+with gr.Blocks(title="MUSICA Enhancement", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🖼️ MUSICA X-ray Image Enhancement")
+    gr.Markdown("Upload a 16-bit grayscale TIFF for wavelet-based enhancement")
+    
+    with gr.Row():
+        with gr.Column():
+            file_input = gr.File(
+                label="Input TIFF",
+                file_types=["tif", "tiff"],
+                height=100
+            )
+            submit_btn = gr.Button("Process", variant="primary")
+        
+        with gr.Column():
+            original_output = gr.Image(
+                label="Original (Clipped)",
+                height=400,
+                type="numpy"
+            )
+            
+    with gr.Row():
+        enhanced_output = gr.Image(
+            label="Enhanced Result",
+            type="numpy",
+            height=400
+        )
+        hist_output = gr.Image(
+            label="Histogram",
+            type="pil",
+            height=400
+        )
+
+    submit_btn.click(
+        process_tiff,
+        inputs=file_input,
+        outputs=[original_output, enhanced_output, hist_output]
+    )
+
+    gr.Examples(
+        examples=[["sample.tif"]],
+        inputs=file_input,
+        outputs=[original_output, enhanced_output, hist_output],
+        fn=process_tiff
+    )
+
+if __name__ == "__main__":
+    demo.launch()