app.py

import gradio as gr
import numpy as np
import cv2
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from huggingface_hub import hf_hub_download
import os
import tempfile
from datetime import datetime
import seaborn as sns
from tensorflow.keras.models import load_model

# Load model directly
model_path = hf_hub_download(repo_id="vanhai123/brain-tumor-unet", filename="best_unet_model.h5")
model = load_model(model_path, compile=False)

# Set styles
plt.style.use('seaborn-v0_8')
sns.set_palette("husl")

def predict_image(model, image_path):
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    image = cv2.resize(image, (256, 256))
    image = image / 255.0
    input_image = np.expand_dims(np.expand_dims(image, axis=0), axis=-1)
    prediction = model.predict(input_image)[0, :, :, 0]
    return prediction

def calculate_metrics(prediction, threshold=0.5):
    binary_pred = (prediction > threshold).astype(np.uint8)
    tumor_pixels = np.sum(binary_pred)
    total_pixels = binary_pred.shape[0] * binary_pred.shape[1]
    tumor_percentage = (tumor_pixels / total_pixels) * 100
    confidence = np.mean(prediction[binary_pred == 1]) if tumor_pixels > 0 else 0
    return {
        "tumor_pixels": int(tumor_pixels),
        "tumor_percentage": round(tumor_percentage, 2),
        "confidence_score": round(confidence * 100, 2),
        "image_size": f"{binary_pred.shape[1]}x{binary_pred.shape[0]}"
    }

def create_overlay_visualization(original, prediction, threshold=0.5):
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('🧠 Brain Tumor Segmentation Analysis', fontsize=20, fontweight='bold', y=0.98)

    axes[0, 0].imshow(original, cmap='gray')
    axes[0, 0].set_title('📷 Original MRI Scan')
    axes[0, 0].axis('off')

    im1 = axes[0, 1].imshow(prediction, cmap='hot', alpha=0.8)
    axes[0, 1].set_title('🔥 Prediction Heatmap')
    axes[0, 1].axis('off')
    plt.colorbar(im1, ax=axes[0, 1], fraction=0.046, pad=0.04)

    binary_mask = (prediction > threshold).astype(np.uint8)
    axes[1, 0].imshow(binary_mask, cmap='Reds', alpha=0.8)
    axes[1, 0].set_title('🎯 Binary Segmentation')
    axes[1, 0].axis('off')

    overlay = original.copy()
    if len(overlay.shape) == 2:
        overlay = np.stack([overlay, overlay, overlay], axis=-1)
    tumor_mask = prediction > threshold
    overlay[tumor_mask, 0] = np.clip(overlay[tumor_mask, 0] + 0.3, 0, 1)
    overlay[tumor_mask, 1] = np.clip(overlay[tumor_mask, 1] - 0.2, 0, 1)
    overlay[tumor_mask, 2] = np.clip(overlay[tumor_mask, 2] - 0.2, 0, 1)
    axes[1, 1].imshow(overlay)
    axes[1, 1].set_title('✨ Overlay Visualization')
    axes[1, 1].axis('off')
    red_patch = mpatches.Patch(color='red', alpha=0.7, label='Tumor Region')
    axes[1, 1].legend(handles=[red_patch], loc='upper right')
    plt.tight_layout()
    return fig

def generate_report(metrics):
    time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    status = "POSITIVE" if metrics['tumor_percentage'] > 0.5 else "NEGATIVE"
    return f"""
# 🧠 Medical Report

**Date:** {time}  
**Tumor Detected:** {status}  
**Confidence:** {metrics['confidence_score']}%  
**Tumor Area:** {metrics['tumor_percentage']}%

> This report is auto-generated by AI and is for research use only.
"""

def predict_and_display(image):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
        image.save(tmp.name)
        prediction = predict_image(model, tmp.name)
    original = cv2.imread(tmp.name, cv2.IMREAD_GRAYSCALE)
    original = cv2.resize(original, (256, 256)) / 255.0
    metrics = calculate_metrics(prediction)
    fig1 = create_overlay_visualization(original, prediction)
    report = generate_report(metrics)
    return fig1, report

def app():
    with gr.Blocks(title="Brain Tumor Analyzer") as demo:
        gr.Markdown("# 🧠 Brain Tumor Segmentation AI")
        img = gr.Image(type="pil", label="Upload MRI Image")
        btn = gr.Button("Analyze")
        out1 = gr.Plot(label="Prediction Visualization")
        out2 = gr.Markdown(label="Report")
        btn.click(fn=predict_and_display, inputs=img, outputs=[out1, out2])
    return demo

if __name__ == "__main__":
    interface = app()
    interface.launch(server_name="0.0.0.0", server_port=7860, share=True)