Spaces:
Paused
Paused
| """ | |
| OralGuard AI - Hugging Face Spaces Deployment | |
| Advanced AI Diagnostic System for Oral Cancer Detection | |
| """ | |
| import streamlit as st | |
| import torch | |
| from PIL import Image | |
| import pandas as pd | |
| import plotly.graph_objects as go | |
| import plotly.express as px | |
| from datetime import datetime | |
| import os | |
| import os | |
| from pathlib import Path | |
| from model_utils import GoogleDriveModelLoader, get_model_display_name, get_magnification_display_name, get_model_size_info | |
| from predict_utils import OralGuardPredictor | |
| # CRITICAL: Create cache directory | |
| cache_dir = Path("./models_cache") | |
| cache_dir.mkdir(parents=True, exist_ok=True) | |
| os.chmod(cache_dir, 0o755) # Ensure write permissions | |
| # Page configuration | |
| st.set_page_config( | |
| page_title="OralGuard AI", | |
| page_icon="🦷", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| # Custom CSS | |
| st.markdown(""" | |
| <style> | |
| .main-header { | |
| font-size: 2.5rem; | |
| font-weight: bold; | |
| color: #1E88E5; | |
| text-align: center; | |
| padding: 1rem; | |
| } | |
| .sub-header { | |
| text-align: center; | |
| color: #666; | |
| font-size: 1.1rem; | |
| margin-bottom: 2rem; | |
| } | |
| .stAlert { | |
| margin-top: 1rem; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| # Initialize session state | |
| if 'model_loader' not in st.session_state: | |
| st.session_state.model_loader = None | |
| if 'predictor' not in st.session_state: | |
| st.session_state.predictor = None | |
| if 'prediction_history' not in st.session_state: | |
| st.session_state.prediction_history = [] | |
| def check_gdrive_configuration(): | |
| """Check if Google Drive file IDs are configured""" | |
| try: | |
| required_keys = [ | |
| 'EFFICIENTNET_100X_ID', 'EFFICIENTNET_400X_ID', | |
| 'VGG16_100X_ID', 'VGG16_400X_ID', | |
| 'VIT_100X_ID', 'VIT_400X_ID' | |
| ] | |
| missing_keys = [] | |
| for key in required_keys: | |
| if not st.secrets.get(key) and not os.getenv(key): | |
| missing_keys.append(key) | |
| return len(missing_keys) == 0, missing_keys | |
| except: | |
| return False, required_keys | |
| def initialize_models(): | |
| """Initialize model loader and predictor""" | |
| if st.session_state.model_loader is None: | |
| try: | |
| st.session_state.model_loader = GoogleDriveModelLoader() | |
| st.session_state.predictor = OralGuardPredictor(st.session_state.model_loader) | |
| return True | |
| except Exception as e: | |
| st.error(f"❌ Error initializing models: {str(e)}") | |
| return False | |
| return True | |
| def display_prediction_result(result, image=None): | |
| """Display prediction result in a beautiful format""" | |
| # Determine styling based on prediction | |
| if result['prediction'] == 'Normal': | |
| emoji = '✅' | |
| color = '#4CAF50' | |
| else: | |
| emoji = '⚠️' | |
| color = '#F44336' | |
| # Create columns | |
| if image: | |
| col1, col2 = st.columns([1, 2]) | |
| with col1: | |
| st.image(image, use_container_width=True, caption="Input Image") | |
| else: | |
| col2 = st.container() | |
| with col2: | |
| st.markdown(f""" | |
| <div style='padding: 2rem; border-radius: 10px; background: linear-gradient(135deg, {color}22 0%, {color}11 100%); border-left: 5px solid {color};'> | |
| <h2 style='color: {color};'>{emoji} {result['prediction']}</h2> | |
| <h3>Confidence: {result['confidence']*100:.2f}%</h3> | |
| </div> | |
| """, unsafe_allow_html=True) | |
| # Probability chart | |
| fig = go.Figure(data=[ | |
| go.Bar( | |
| x=['Normal', 'OSCC'], | |
| y=[result['probabilities']['Normal']*100, result['probabilities']['OSCC']*100], | |
| marker_color=['#4CAF50', '#F44336'], | |
| text=[f"{result['probabilities']['Normal']*100:.1f}%", | |
| f"{result['probabilities']['OSCC']*100:.1f}%"], | |
| textposition='auto', | |
| ) | |
| ]) | |
| fig.update_layout( | |
| title="Class Probabilities", | |
| yaxis_title="Probability (%)", | |
| height=300, | |
| showlegend=False | |
| ) | |
| st.plotly_chart(fig, use_container_width=True) | |
| # Detailed metrics | |
| st.markdown("### 📊 Detailed Results") | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| st.metric("Prediction", result['prediction']) | |
| st.metric("Confidence", f"{result['confidence']*100:.2f}%") | |
| with col2: | |
| st.metric("Normal Probability", f"{result['probabilities']['Normal']*100:.2f}%") | |
| st.metric("OSCC Probability", f"{result['probabilities']['OSCC']*100:.2f}%") | |
| with col3: | |
| st.metric("Inference Time", f"{result['inference_time']:.3f}s") | |
| st.metric("Model Used", result['model']) | |
| # Individual model results for ensemble | |
| if 'individual_results' in result: | |
| st.markdown("### 🤖 Individual Model Predictions") | |
| individual_data = [] | |
| for model_name, res in result['individual_results'].items(): | |
| individual_data.append({ | |
| 'Model': get_model_display_name(model_name), | |
| 'Prediction': res['prediction'], | |
| 'Confidence': f"{res['confidence']*100:.2f}%", | |
| 'Normal Prob': f"{res['probabilities']['Normal']*100:.2f}%", | |
| 'OSCC Prob': f"{res['probabilities']['OSCC']*100:.2f}%" | |
| }) | |
| st.dataframe(pd.DataFrame(individual_data), use_container_width=True) | |
| def single_image_prediction(): | |
| """Single image prediction interface""" | |
| st.markdown("## 🖼️ Single Image Analysis") | |
| col1, col2 = st.columns([2, 1]) | |
| with col1: | |
| uploaded_file = st.file_uploader( | |
| "Upload Histopathology Image", | |
| type=['png', 'jpg', 'jpeg'], | |
| help="Upload a histopathology image of oral tissue" | |
| ) | |
| with col2: | |
| st.markdown("### ⚙️ Configuration") | |
| model_type = st.selectbox( | |
| "Select Model", | |
| ['ensemble', 'efficientnet', 'vgg16', 'vit'], | |
| format_func=get_model_display_name, | |
| help="Choose which model to use for prediction" | |
| ) | |
| magnification = st.selectbox( | |
| "Magnification Level", | |
| ['100x', '400x'], | |
| format_func=get_magnification_display_name, | |
| help="Select the magnification level of your image" | |
| ) | |
| # Show expected download size | |
| if model_type != 'ensemble': | |
| sizes = get_model_size_info() | |
| st.info(f"📦 Download size: {sizes[model_type]}") | |
| else: | |
| st.info(f"📦 Download size: ~2.94 GB (all 3 models)") | |
| if uploaded_file is not None: | |
| # Display image FIRST | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_container_width=True) | |
| # Predict button | |
| if st.button("🔬 Analyze Image", key="single_predict", type="primary"): | |
| with st.spinner("🔄 Analyzing image..."): | |
| try: | |
| # Make prediction | |
| if model_type == 'ensemble': | |
| result = st.session_state.predictor.predict_ensemble(image, magnification) | |
| else: | |
| result = st.session_state.predictor.predict_single_model( | |
| image, model_type, magnification | |
| ) | |
| # Display result (pass image for display) | |
| display_prediction_result(result, image) | |
| # Add to history | |
| st.session_state.prediction_history.append({ | |
| 'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'), | |
| 'prediction': result['prediction'], | |
| 'confidence': result['confidence'], | |
| 'model': result['model'] | |
| }) | |
| st.success("✅ Analysis complete!") | |
| except Exception as e: | |
| st.error(f"❌ Error during prediction: {str(e)}") | |
| import traceback | |
| st.code(traceback.format_exc()) | |
| def batch_image_prediction(): | |
| """Batch image prediction interface""" | |
| st.markdown("## 📁 Batch Image Analysis") | |
| col1, col2 = st.columns([2, 1]) | |
| with col1: | |
| uploaded_files = st.file_uploader( | |
| "Upload Multiple Histopathology Images", | |
| type=['png', 'jpg', 'jpeg'], | |
| accept_multiple_files=True, | |
| help="Upload multiple histopathology images for batch processing" | |
| ) | |
| with col2: | |
| st.markdown("### ⚙️ Configuration") | |
| model_type = st.selectbox( | |
| "Select Model", | |
| ['ensemble', 'efficientnet', 'vgg16', 'vit'], | |
| format_func=get_model_display_name, | |
| key="batch_model" | |
| ) | |
| magnification = st.selectbox( | |
| "Magnification Level", | |
| ['100x', '400x'], | |
| format_func=get_magnification_display_name, | |
| key="batch_mag" | |
| ) | |
| if uploaded_files: | |
| st.info(f"📊 {len(uploaded_files)} images uploaded") | |
| # Display thumbnails | |
| cols = st.columns(min(len(uploaded_files), 5)) | |
| for idx, file in enumerate(uploaded_files[:5]): | |
| with cols[idx]: | |
| img = Image.open(file) | |
| st.image(img, use_container_width=True) | |
| if len(uploaded_files) > 5: | |
| st.info(f"... and {len(uploaded_files) - 5} more images") | |
| if st.button("🔬 Analyze All Images", key="batch_predict", type="primary"): | |
| # Load images | |
| images = [Image.open(file) for file in uploaded_files] | |
| with st.spinner("🔄 Analyzing images..."): | |
| try: | |
| # Batch prediction | |
| results = st.session_state.predictor.predict_batch( | |
| images, | |
| model_type, | |
| magnification | |
| ) | |
| # Display summary | |
| st.success(f"✅ Analysis complete! Processed {len(results)} images") | |
| # Summary statistics | |
| successful = [r for r in results if r.get('status') == 'success'] | |
| normal_count = sum(1 for r in successful if r['prediction'] == 'Normal') | |
| oscc_count = sum(1 for r in successful if r['prediction'] == 'OSCC') | |
| avg_confidence = sum(r['confidence'] for r in successful) / len(successful) if successful else 0 | |
| col1, col2, col3, col4 = st.columns(4) | |
| with col1: | |
| st.metric("Total Images", len(results)) | |
| with col2: | |
| st.metric("Normal", normal_count) | |
| with col3: | |
| st.metric("OSCC", oscc_count) | |
| with col4: | |
| st.metric("Avg Confidence", f"{avg_confidence*100:.1f}%") | |
| # Results table | |
| st.markdown("### 📋 Detailed Results") | |
| results_data = [] | |
| for idx, result in enumerate(successful): | |
| results_data.append({ | |
| 'Image': f"Image_{idx+1}", | |
| 'Prediction': result['prediction'], | |
| 'Confidence': f"{result['confidence']*100:.2f}%", | |
| 'Normal Prob': f"{result['probabilities']['Normal']*100:.2f}%", | |
| 'OSCC Prob': f"{result['probabilities']['OSCC']*100:.2f}%", | |
| 'Time (s)': f"{result['inference_time']:.3f}" | |
| }) | |
| df = pd.DataFrame(results_data) | |
| st.dataframe(df, use_container_width=True) | |
| # Download results | |
| csv = df.to_csv(index=False) | |
| st.download_button( | |
| "📥 Download Results (CSV)", | |
| csv, | |
| "oralguard_batch_results.csv", | |
| "text/csv", | |
| key='download-csv' | |
| ) | |
| # Visualization | |
| fig = px.pie( | |
| values=[normal_count, oscc_count], | |
| names=['Normal', 'OSCC'], | |
| title="Prediction Distribution", | |
| color_discrete_sequence=['#4CAF50', '#F44336'] | |
| ) | |
| st.plotly_chart(fig, use_container_width=True) | |
| except Exception as e: | |
| st.error(f"❌ Error during batch prediction: {str(e)}") | |
| def about_page(): | |
| """About page""" | |
| st.markdown(""" | |
| ## About OralGuard AI | |
| **OralGuard AI** is a production-ready AI diagnostic system designed to address Pakistan's oral cancer crisis. | |
| ### 🎯 Key Features | |
| - **90%+ Accuracy** on international benchmarks (Mendeley oral pathology dataset) | |
| - **On-Demand Model Loading** - Models download only when selected | |
| - **Ensemble Architecture** combining EfficientNet-B5, VGG16, and ViT-B/16 | |
| - **Dual Magnification** support (100× and 400×) | |
| - **Batch Processing** for multiple images | |
| ### 🏥 Clinical Impact | |
| - Reduces diagnostic turnaround from **5-14 days to same-day** preliminary screening | |
| - Addresses Pakistan's **severe pathologist shortage** | |
| - Enables **specialist-level diagnostics** in remote areas | |
| - **Cost-effective**: Significantly lower per-case cost | |
| ### 🔬 Technical Specifications | |
| - **Models**: EfficientNet-B5 (340 MB), VGG16-BN (1.6 GB), ViT-B/16 (1.0 GB) | |
| - **Image Size**: 224×224 pixels | |
| - **Classes**: Normal, OSCC (Oral Squamous Cell Carcinoma) | |
| - **Deployment**: Hugging Face Spaces with Google Drive integration | |
| ### 📊 Performance Metrics | |
| - **100× Magnification**: 84.91% accuracy (ViT-B/16) | |
| - **400× Magnification**: 81.43% accuracy (ViT-B/16) | |
| - **Ensemble**: Further improved performance | |
| ### ⚖️ Disclaimer | |
| This system is designed for **preliminary screening and clinical decision support**. | |
| All diagnoses should be confirmed by qualified pathologists. | |
| Not a replacement for professional medical judgment. | |
| --- | |
| **Version**: 1.0.0 | **Last Updated**: October 2024 | |
| """) | |
| def main(): | |
| """Main application""" | |
| # Header | |
| st.markdown("<h1 class='main-header'>🦷 OralGuard AI</h1>", unsafe_allow_html=True) | |
| st.markdown("<p class='sub-header'>Advanced AI Diagnostic System for Oral Cancer Detection</p>", unsafe_allow_html=True) | |
| # Check Google Drive configuration | |
| configured, missing = check_gdrive_configuration() | |
| if not configured: | |
| st.error("❌ Google Drive file IDs not configured!") | |
| st.markdown(""" | |
| ### 🔧 Configuration Required | |
| Please configure Google Drive file IDs in Hugging Face Secrets: | |
| 1. Go to your Space Settings → Repository secrets | |
| 2. Add the following secrets with your Google Drive file IDs: | |
| """) | |
| for key in missing: | |
| st.code(f"{key} = your_google_drive_file_id") | |
| st.info("📖 See README.md for detailed setup instructions") | |
| st.stop() | |
| # Initialize models | |
| if not initialize_models(): | |
| st.stop() | |
| # Sidebar | |
| with st.sidebar: | |
| st.markdown("### 🎯 Navigation") | |
| page = st.radio( | |
| "Select Mode", | |
| ["🖼️ Single Image", "📁 Batch Analysis", "ℹ️ About"], | |
| label_visibility="collapsed" | |
| ) | |
| st.markdown("---") | |
| st.markdown("### ⚙️ System Status") | |
| if torch.cuda.is_available(): | |
| st.success("✅ GPU Available") | |
| st.info(f"🖥️ {torch.cuda.get_device_name(0)}") | |
| else: | |
| st.warning("⚠️ Running on CPU") | |
| loaded_models = len(st.session_state.model_loader.loaded_models) | |
| st.metric("Loaded Models", loaded_models) | |
| if loaded_models > 0: | |
| if st.button("🗑️ Clear Model Cache"): | |
| st.session_state.model_loader.clear_cache() | |
| st.rerun() | |
| st.markdown("---") | |
| # Model size info | |
| with st.expander("📦 Model Sizes"): | |
| sizes = get_model_size_info() | |
| st.write(f"**EfficientNet-B5:** {sizes['efficientnet']}") | |
| st.write(f"**VGG16-BN:** {sizes['vgg16']}") | |
| st.write(f"**ViT-B/16:** {sizes['vit']}") | |
| st.write(f"**Total (one magnification):** {sizes['total_per_magnification']}") | |
| # Main content | |
| if "🖼️ Single Image" in page: | |
| single_image_prediction() | |
| elif "📁 Batch Analysis" in page: | |
| batch_image_prediction() | |
| elif "ℹ️ About" in page: | |
| about_page() | |
| if __name__ == "__main__": | |
| main() |