app.py

#!/usr/bin/env python3
"""
zeroFire - Fire Detection Classification App
AI-powered fire detection using ConvNeXt transfer learning
"""

import streamlit as st
import torch
import torch.nn.functional as F
from PIL import Image
import numpy as np
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import pandas as pd
import sys
import os
import time
from io import BytesIO
import base64

# Add utils to path
sys.path.append('utils')
from model_utils import load_model, FireDetectionClassifier
from data_utils import get_inference_transform, prepare_image_for_inference, check_data_directory

# Page Configuration
st.set_page_config(
    page_title="🔥 zeroFire - Fire Detection System",
    page_icon="🔥",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Beautiful UI
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%);
        padding: 2rem;
        border-radius: 15px;
        text-align: center;
        color: white;
        margin-bottom: 2rem;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .main-header h1 {
        font-size: 3rem;
        margin: 0;
        font-weight: bold;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
    }
    
    .main-header p {
        font-size: 1.2rem;
        margin: 0.5rem 0 0 0;
        opacity: 0.9;
    }
    
    .upload-section {
        background: linear-gradient(135deg, #a8e6cf 0%, #74b9ff 100%);
        color: white;
        padding: 15px 20px;
        border-radius: 15px;
        text-align: center;
        margin-bottom: 20px;
        box-shadow: 0 6px 20px rgba(168, 230, 207, 0.3);
        height: 100px;
        display: flex;
        flex-direction: column;
        justify-content: center;
    }
    
    .upload-section h3 {
        font-size: 1.3rem;
        margin: 0 0 5px 0;
    }
    
    .upload-section p {
        font-size: 0.9rem;
        margin: 0;
        opacity: 0.9;
    }
    
    .result-fire {
        background: linear-gradient(135deg, #dc3545 0%, #c82333 100%);
        color: white;
        padding: 15px 20px;
        border-radius: 15px;
        text-align: center;
        margin: 0 0 20px 0;
        box-shadow: 0 6px 20px rgba(220, 53, 69, 0.3);
        height: 100px;
        display: flex;
        flex-direction: column;
        justify-content: center;
    }
    
    .result-fire h2 {
        font-size: 1.2rem;
        margin: 0 0 5px 0;
        line-height: 1.2;
    }
    
    .result-fire p {
        font-size: 0.85rem;
        margin: 0;
        font-weight: bold;
    }
    
    .result-no-fire {
        background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
        color: white;
        padding: 15px 20px;
        border-radius: 15px;
        text-align: center;
        margin: 0 0 20px 0;
        box-shadow: 0 6px 20px rgba(40, 167, 69, 0.3);
        height: 100px;
        display: flex;
        flex-direction: column;
        justify-content: center;
    }
    
    .result-no-fire h2 {
        font-size: 1.2rem;
        margin: 0 0 5px 0;
        line-height: 1.2;
    }
    
    .result-no-fire p {
        font-size: 0.85rem;
        margin: 0;
        font-weight: bold;
    }
    
    .metric-card {
        background: linear-gradient(135deg, #ff9ff3 0%, #f368e0 100%);
        padding: 15px;
        border-radius: 10px;
        text-align: center;
        color: white;
        margin: 10px 0;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
    }
    
    .info-card {
        background: linear-gradient(135deg, #ff9ff3 0%, #f368e0 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin: 15px 0;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .info-card-fire {
        background: linear-gradient(135deg, #ff6b6b 0%, #ee5a52 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin: 15px 0;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    
    .stButton > button {
        background: linear-gradient(45deg, #74b9ff, #0984e3);
        color: white;
        border: none;
        border-radius: 10px;
        padding: 12px 24px;
        font-weight: bold;
        font-size: 1.1rem;
        transition: all 0.3s ease;
        box-shadow: 0 4px 15px rgba(0,0,0,0.2);
        width: 100%;
    }
    
    .stButton > button:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(0,0,0,0.3);
    }
    
    .sidebar .stSelectbox > div > div {
        background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%);
        color: white;
    }
    
    /* Constrain image height to prevent scrolling */
    .stImage > img {
        max-height: 400px;
        width: auto;
        object-fit: contain;
    }
    
    /* Mobile-friendly camera and tabs - Fixed styling */
    .stTabs {
        margin-top: 10px;
    }
    
    .stTabs > div > div > div > div {
        padding: 8px 12px;
        font-weight: bold;
        border: none;
        background: transparent;
        color: #333;
        border-radius: 8px;
        margin-right: 5px;
        transition: all 0.3s ease;
    }
    
    .stTabs > div > div > div > div[aria-selected="true"] {
        background: linear-gradient(135deg, #74b9ff 0%, #0984e3 100%);
        color: white !important;
        border-radius: 8px;
        box-shadow: 0 2px 8px rgba(116, 185, 255, 0.3);
    }
    
    .stTabs > div > div > div > div:hover {
        background: rgba(116, 185, 255, 0.1);
        color: #0984e3;
    }
    
    /* Fix tab underline */
    .stTabs > div > div > div {
        border-bottom: none;
    }
    
    /* Camera input styling */
    .stCamera > div {
        border-radius: 15px;
        overflow: hidden;
        box-shadow: 0 4px 15px rgba(0,0,0,0.1);
    }
    
    /* Mobile responsive adjustments */
    @media (max-width: 768px) {
        .main-header h1 {
            font-size: 2rem;
        }
        
        .main-header p {
            font-size: 1rem;
        }
        
        .upload-section {
            height: 80px;
        }
        
        .upload-section h3 {
            font-size: 1.1rem;
        }
        
        .result-fire h2, .result-no-fire h2 {
            font-size: 1rem;
        }
        
        .result-fire, .result-no-fire {
            height: 80px;
        }
        
        .stImage > img {
            max-height: 300px;
        }
        
        /* Mobile-specific file uploader improvements */
        .stFileUploader > div > div > div > div {
            padding: 15px;
            border-radius: 10px;
            border: 2px dashed #74b9ff;
            background: rgba(116, 185, 255, 0.05);
            text-align: center;
        }
        
        .stFileUploader > div > div > div > div:hover {
            border-color: #0984e3;
            background: rgba(116, 185, 255, 0.1);
        }
        
        /* Make tabs more touch-friendly on mobile */
        .stTabs > div > div > div > div {
            padding: 12px 16px;
            font-size: 1rem;
            min-height: 44px;
        }
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def load_fire_model():
    """Load the trained fire detection model"""
    model_path = 'models/fire_detection_classifier.pth'
    
    if not os.path.exists(model_path):
        return None, "Model not found. Please train the model first."
    
    try:
        model, model_info = load_model(model_path, device='cpu')
        return model, model_info
    except Exception as e:
        return None, f"Error loading model: {str(e)}"

def get_prediction(image, model, transform):
    """Get prediction from the model"""
    try:
        # Prepare image
        input_tensor = prepare_image_for_inference(image, transform)
        
        # Get prediction
        with torch.no_grad():
            model.eval()
            outputs = model(input_tensor)
            probabilities = F.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, 1)
            
            # Convert to numpy
            predicted_class = predicted.item()
            confidence_score = confidence.item()
            all_probs = probabilities.squeeze().cpu().numpy()
            
            return predicted_class, confidence_score, all_probs
    
    except Exception as e:
        st.error(f"Error during prediction: {str(e)}")
        return None, None, None

def create_confidence_chart(probabilities, class_names):
    """Create confidence chart using Plotly"""
    fig = go.Figure(data=[
        go.Bar(
            x=class_names,
            y=probabilities,
            marker_color=['#dc3545', '#28a745'],
            text=[f'{p:.1%}' for p in probabilities],
            textposition='auto',
        )
    ])
    
    fig.update_layout(
        title="Fire Detection Confidence",
        xaxis_title="Prediction",
        yaxis_title="Confidence",
        yaxis=dict(range=[0, 1]),
        showlegend=False,
        height=400,
        template="plotly_white"
    )
    
    return fig

def create_safety_metrics_chart(predicted_class, confidence):
    """Create safety metrics visualization"""
    if predicted_class == 0:  # Fire
        danger_level = confidence * 100
        safety_level = (1 - confidence) * 100
        primary_color = '#dc3545'
        status = "FIRE DETECTED"
    else:  # No Fire
        danger_level = (1 - confidence) * 100
        safety_level = confidence * 100
        primary_color = '#28a745'
        status = "NO FIRE"
    
    fig = go.Figure(go.Indicator(
        mode = "gauge+number+delta",
        value = danger_level,
        domain = {'x': [0, 1], 'y': [0, 1]},
        title = {'text': "Fire Risk Level"},
        delta = {'reference': 50},
        gauge = {'axis': {'range': [None, 100]},
                'bar': {'color': primary_color},
                'steps' : [
                    {'range': [0, 25], 'color': "lightgray"},
                    {'range': [25, 50], 'color': "yellow"},
                    {'range': [50, 75], 'color': "orange"},
                    {'range': [75, 100], 'color': "red"}],
                'threshold' : {'line': {'color': "red", 'width': 4},
                              'thickness': 0.75, 'value': 90}}))
    
    fig.update_layout(height=400)
    return fig

def analyze_fire_risk(predicted_class, confidence):
    """Analyze fire risk and provide recommendations"""
    if predicted_class == 0:  # Fire detected
        risk_level = confidence * 100
        
        if risk_level >= 90:
            return {
                'level': 'CRITICAL',
                'color': '#dc3545',
                'icon': '🚨',
                'message': 'IMMEDIATE ACTION REQUIRED',
                'recommendations': [
                    'Activate fire suppression system immediately',
                    'Evacuate the area',
                    'Call emergency services UAE 997',
                    'Shut down affected equipment if safe to do so',
                    'Monitor surrounding areas for spread'
                ]
            }
        elif risk_level >= 75:
            return {
                'level': 'HIGH',
                'color': '#fd7e14',
                'icon': '⚠️',
                'message': 'HIGH FIRE RISK DETECTED',
                'recommendations': [
                    'Investigate the area immediately',
                    'Prepare fire suppression systems',
                    'Alert security personnel',
                    'Consider equipment shutdown',
                    'Increase monitoring frequency'
                ]
            }
        else:
            return {
                'level': 'MODERATE',
                'color': '#ffc107',
                'icon': '🔶',
                'message': 'POSSIBLE FIRE DETECTED',
                'recommendations': [
                    'Verify with additional sensors',
                    'Send personnel to investigate',
                    'Check equipment temperatures',
                    'Review recent maintenance logs',
                    'Maintain heightened awareness'
                ]
            }
    else:  # No fire
        safety_level = confidence * 100
        
        if safety_level >= 95:
            return {
                'level': 'SAFE',
                'color': '#28a745',
                'icon': '✅',
                'message': 'NORMAL OPERATION',
                'recommendations': [
                    'Continue normal operations',
                    'Maintain regular monitoring',
                    'Keep fire suppression systems ready',
                    'Perform scheduled maintenance',
                    'Review safety protocols periodically'
                ]
            }
        else:
            return {
                'level': 'CAUTION',
                'color': '#17a2b8',
                'icon': '⚠️',
                'message': 'MONITOR CLOSELY',
                'recommendations': [
                    'Increase monitoring frequency',
                    'Check for unusual conditions',
                    'Verify sensor functionality',
                    'Review environmental factors',
                    'Maintain readiness for action'
                ]
            }

# display_fire_safety_checklist function removed - content moved to right column

def main():
    """Main application function"""
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>🔥 zeroFire</h1>
        <p>AI-Powered Fire Detection System for Data Centers</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.markdown("### 🔥 Fire Detection System")
        st.markdown("---")
        
        # Model status
        model, model_info = load_fire_model()
        
        if model is None:
            st.error("❌ Model not available")
            st.info("Train the model first using: `python train_fire_detection.py`")
            return
        else:
            st.success("✅ Model loaded successfully")
            if isinstance(model_info, dict):
                accuracy = model_info.get('best_acc', 'Unknown')
                if accuracy != 'Unknown':
                    # Format accuracy to 2 decimal places
                    accuracy_formatted = f"{float(accuracy):.2f}%"
                    st.info(f"📊 Model: ConvNeXt Large")
                    st.info(f"🎯 Accuracy: {accuracy_formatted}")
                    st.info(f"🔄 Transfer Learning: Food→Fire")
                    st.info(f"⚡ Precision: High-recall optimized")
                else:
                    st.info("📊 Model Accuracy: Unknown")
        
        st.markdown("---")
        
        # Settings
        st.markdown("### ⚙️ Settings")
        confidence_threshold = st.slider(
            "Confidence Threshold",
            min_value=0.0,
            max_value=1.0,
            value=0.5,
            step=0.05,
            help="Minimum confidence required for fire detection"
        )
        
        show_details = st.checkbox("Show detailed analysis", value=True)
        
        st.markdown("---")
        st.markdown("### 📊 Data Center Status")
        check_data_directory('data')
    
    # Main content
    col1, col2 = st.columns([2, 1])
    
    with col1:
        # Upload section
        st.markdown("""
        <div class="upload-section">
            <h3>📸 Fire Detection Input</h3>
            <p>Upload an image or take a photo to detect fire or smoke in your data center</p>
        </div>
        """, unsafe_allow_html=True)
        
        # Create tabs for different input methods
        tab1, tab2 = st.tabs(["📱 Mobile Camera", "📁 Upload File"])
        
        image = None
        image_source = None
        
        with tab1:
            st.markdown("""
            <div style="background: linear-gradient(135deg, #e8f5e8 0%, #c8e6c9 100%); 
                       padding: 15px; border-radius: 10px; margin: 10px 0;
                       border: 2px solid #81c784; text-align: center;">
                <h4 style="color: #2e7d32; margin: 0;">📱 Best for Mobile Users</h4>
                <p style="color: #388e3c; margin: 5px 0; font-size: 14px;">
                    Click "Browse files" → "Camera" to open your phone's camera app
                </p>
            </div>
            """, unsafe_allow_html=True)
            
            # Mobile-optimized file uploader
            uploaded_file = st.file_uploader(
                "📱 Tap to Take Photo or Upload Image",
                type=['jpg', 'jpeg', 'png', 'bmp', 'tiff'],
                help="On mobile: Tap 'Browse files' then select 'Camera' to take a photo, or 'Photo Library' to upload existing images",
                key="mobile_uploader"
            )
            
            if uploaded_file is not None:
                image = Image.open(uploaded_file)
                image_source = "mobile"
        
        with tab2:
            st.markdown("""
            <div style="background: linear-gradient(135deg, #e3f2fd 0%, #bbdefb 100%); 
                       padding: 15px; border-radius: 10px; margin: 10px 0;
                       border: 2px solid #64b5f6; text-align: center;">
                <h4 style="color: #1976d2; margin: 0;">💻 Desktop Upload</h4>
                <p style="color: #1976d2; margin: 5px 0; font-size: 14px;">
                    Select an image file from your computer
                </p>
            </div>
            """, unsafe_allow_html=True)
            
            # Desktop file uploader
            desktop_file = st.file_uploader(
                "Choose an image file...",
                type=['jpg', 'jpeg', 'png', 'bmp', 'tiff'],
                help="Upload an image from your computer for fire detection",
                key="desktop_uploader"
            )
            
            if desktop_file is not None:
                image = Image.open(desktop_file)
                image_source = "desktop"
        
        # Process the image (whether uploaded or from camera)
        if image is not None:
            # Display the image with appropriate caption
            caption_map = {
                "mobile": "📱 Mobile Photo",
                "desktop": "💻 Desktop Image",
                "camera": "📷 Camera Photo"
            }
            caption = caption_map.get(image_source, "Uploaded Image")
            st.image(image, caption=caption, use_column_width=True)
            
            # Get prediction
            transform = get_inference_transform()
            predicted_class, confidence_score, all_probs = get_prediction(image, model, transform)
            
            if predicted_class is not None:
                class_names = ['Fire', 'No Fire']
                predicted_label = class_names[predicted_class]
                
                # Fire detection result moved to right column
                
                # Technical details (keep only this in left column)
                with st.expander("🔬 Technical Details"):
                    st.markdown(f"""
                    **Prediction Details:**
                    - Input Source: {image_source.title()}
                    - Predicted Class: {predicted_label}
                    - Confidence Score: {confidence_score:.4f}
                    - Fire Probability: {all_probs[0]:.4f}
                    - No-Fire Probability: {all_probs[1]:.4f}
                    - Threshold: {confidence_threshold:.2f}
                    """)
    
    with col2:
        # Display fire detection result first
        if 'predicted_class' in locals():
            # Display fire detection result box
            if predicted_class == 0:  # Fire
                st.markdown(f"""
                <div class="result-fire">
                    <h2>🚨 FIRE DETECTED - Confidence: {confidence_score:.1%}</h2>
                    <p>IMMEDIATE ACTION REQUIRED</p>
                </div>
                """, unsafe_allow_html=True)
            else:  # No Fire
                st.markdown(f"""
                <div class="result-no-fire">
                    <h2>✅ NO FIRE DETECTED - Confidence: {confidence_score:.1%}</h2>
                    <p>Normal Operation</p>
                </div>
                """, unsafe_allow_html=True)
        
        # Quick metrics
        st.markdown("### 📊 Quick Metrics")
        
        if 'predicted_class' in locals():
            # Create three columns for metrics in a row
            metric_col1, metric_col2, metric_col3 = st.columns(3)
            
            with metric_col1:
                # Fire risk gauge
                risk_percentage = all_probs[0] * 100
                st.metric(
                    label="Fire Risk",
                    value=f"{risk_percentage:.1f}%"
                )
            
            with metric_col2:
                # Safety score
                safety_score = all_probs[1] * 100
                st.metric(
                    label="Safety Score",
                    value=f"{safety_score:.1f}%"
                )
            
            with metric_col3:
                # Model status instead of redundant confidence
                status = "FIRE ALERT" if predicted_class == 0 else "NORMAL"
                st.metric(
                    label="Status",
                    value=status
                )
            
                         # Charts removed to eliminate confusion and redundancy
            
            # Detailed analysis moved from left column
            if show_details:
                st.markdown("### 📋 Detailed Analysis")
                
                analysis = analyze_fire_risk(predicted_class, confidence_score)
                
                # Use red styling only for fire detection
                card_class = "info-card-fire" if predicted_class == 0 else "info-card"
                
                st.markdown(f"""
                <div class="{card_class}">
                    <h3>{analysis['icon']} Risk Level: {analysis['level']}</h3>
                    <p><strong>{analysis['message']}</strong></p>
                </div>
                """, unsafe_allow_html=True)
                
                st.markdown("#### 🎯 Recommended Actions:")
                for rec in analysis['recommendations']:
                    st.markdown(f"- {rec}")
                
                # Fire Safety Checklist moved here
                st.markdown("---")
                st.markdown("""
                <div class="info-card">
                    <h3>🔥 Fire Safety Checklist</h3>
                    <p>Essential fire safety measures for data centers:</p>
                </div>
                """, unsafe_allow_html=True)
                
                safety_items = [
                    "🔥 **Fire Detection Systems** - Smoke, heat, and flame detectors",
                    "💨 **Suppression Systems** - Clean agent, water mist, or CO2",
                    "🌡️ **Temperature Monitoring** - Continuous thermal monitoring",
                    "⚡ **Electrical Safety** - Arc fault and ground fault protection",
                    "🚪 **Emergency Exits** - Clear and well-marked escape routes",
                    "📋 **Emergency Procedures** - Staff training and evacuation plans",
                    "🔧 **Equipment Maintenance** - Regular inspection and testing",
                    "📞 **Emergency Contacts** - Quick access to fire department",
                    "🎯 **Response Plans** - Pre-defined actions for different scenarios",
                    "📊 **Documentation** - Incident logging and safety records"
                ]
                
                for item in safety_items:
                    st.markdown(f"- {item}")
        
        # Additional info moved to bottom
    
    # Footer
    st.markdown("---")
    
    # Emergency Contacts only (Fire Safety Checklist moved to right column)
    st.markdown("""
    <div class="info-card" style="padding: 15px;">
        <h3 style="margin: 0 0 10px 0; text-align: center;">🚨 Emergency Contacts</h3>
        <div style="display: flex; justify-content: space-around; flex-wrap: wrap; gap: 15px;">
            <span><strong>Fire Dept:</strong> UAE 997</span>
            <span><strong>Security:</strong> [Your Number]</span>
            <span><strong>Facilities:</strong> [Your Number]</span>
            <span><strong>IT Ops:</strong> [Your Number]</span>
        </div>
    </div>
    """, unsafe_allow_html=True)
    
    st.markdown("---")
    st.markdown("""
    <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #a8e6cf 0%, #74b9ff 100%); border-radius: 15px; color: white;">
        <p>🔥 zeroFire - AI-Powered Fire Detection System</p>
        <p>Protecting your data center with advanced machine learning</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()