app.py

#!/usr/bin/env python3
"""
ACCEPTIN - Telecom Site Quality Classification App
AI-powered telecom site inspection 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
from utils.model_utils import load_model, TelecomClassifier
from utils.data_utils import get_inference_transform, prepare_image_for_inference, check_data_directory

# Page Configuration
st.set_page_config(
    page_title="📡 ACCEPTIN - Telecom Site Inspector",
    page_icon="📡",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for Beautiful UI
st.markdown("""
<style>
    .main-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 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, #56ab2f 0%, #a8e6cf 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin-bottom: 20px;
        box-shadow: 0 6px 20px rgba(86, 171, 47, 0.3);
    }
    
    .result-good {
        background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin: 20px 0;
        box-shadow: 0 6px 20px rgba(40, 167, 69, 0.3);
    }
    
    .result-bad {
        background: linear-gradient(135deg, #dc3545 0%, #e83e8c 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        text-align: center;
        margin: 20px 0;
        box-shadow: 0 6px 20px rgba(220, 53, 69, 0.3);
    }
    
    .metric-card {
        background: linear-gradient(135deg, #4facfe 0%, #00f2fe 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, #667eea 0%, #764ba2 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, #667eea, #764ba2);
        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, #667eea 0%, #764ba2 100%);
        color: white;
    }
</style>
""", unsafe_allow_html=True)

@st.cache_resource
def load_telecom_model():
    """Load the telecom classification model"""
    model_path = 'models/telecom_classifier.pth'
    
    if not os.path.exists(model_path):
        return None, "Model not found. Please ensure the model is in the models/ directory."
    
    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, confidence_threshold=0.70):
    """Get prediction from the model with confidence threshold"""
    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()
            
            # Check confidence threshold
            if confidence_score < confidence_threshold:
                return None, confidence_score, all_probs  # Return confidence for display
            
            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="Classification Confidence",
        xaxis_title="Site Quality",
        yaxis_title="Confidence",
        yaxis=dict(range=[0, 1]),
        showlegend=False,
        height=400,
        template="plotly_white"
    )
    
    return fig

def create_quality_metrics_chart(predicted_class, confidence):
    """Create quality metrics visualization"""
    if predicted_class == 1:  # Good
        quality_score = confidence * 100
        color = '#28a745'
        status = 'ACCEPTED'
    else:  # Bad
        quality_score = (1 - confidence) * 100
        color = '#dc3545'
        status = 'REJECTED'
    
    fig = go.Figure(go.Indicator(
        mode="gauge+number+delta",
        value=quality_score,
        domain={'x': [0, 1], 'y': [0, 1]},
        title={'text': f"Quality Score<br><span style='font-size:0.8em;color:{color}'>{status}</span>"},
        delta={'reference': 80},
        gauge={
            'axis': {'range': [None, 100]},
            'bar': {'color': color},
            'steps': [
                {'range': [0, 50], 'color': "lightgray"},
                {'range': [50, 80], 'color': "yellow"},
                {'range': [80, 100], 'color': "lightgreen"}
            ],
            'threshold': {
                'line': {'color': "red", 'width': 4},
                'thickness': 0.75,
                'value': 90
            }
        }
    ))
    
    fig.update_layout(height=400)
    return fig

def analyze_site_quality(predicted_class, confidence):
    """Analyze site quality and provide detailed feedback"""
    class_names = ['Bad', 'Good']
    predicted_label = class_names[predicted_class]
    
    if predicted_class == 1:  # Good site
        analysis = {
            'status': 'ACCEPTED ✅',
            'color': 'result-good',
            'icon': '✅',
            'message': 'Site installation meets quality standards',
            'details': [
                '✅ Cable assembly appears properly organized',
                '✅ Equipment installation looks correct',
                '✅ Overall site organization is acceptable',
                '✅ No obvious safety violations detected'
            ],
            'recommendations': [
                '🔍 Verify all labels are clearly readable',
                '🔧 Double-check all card installations',
                '📋 Complete final inspection checklist',
                '📸 Document final installation state'
            ]
        }
    else:  # Bad site
        analysis = {
            'status': 'REJECTED ❌',
            'color': 'result-bad',
            'icon': '❌',
            'message': 'Site installation requires attention',
            'details': [
                '❌ Cable organization may need improvement',
                '❌ Equipment installation issues detected',
                '❌ Site organization below standards',
                '❌ Potential safety or quality concerns'
            ],
            'recommendations': [
                '🔧 Reorganize cable routing and bundling',
                '🔍 Check all card installations and seating',
                '🏷️ Verify all labels are present and readable',
                '⚠️ Address any safety violations',
                '📋 Complete corrective actions before acceptance'
            ]
        }
    
    analysis['confidence'] = confidence
    analysis['predicted_label'] = predicted_label
    
    return analysis

def display_inspection_checklist():
    """Display telecom site inspection checklist"""
    st.markdown("""
    <div class="info-card">
        <h3>📋 Telecom Site Inspection Checklist</h3>
        <p>Use this checklist to ensure comprehensive site evaluation:</p>
    </div>
    """, unsafe_allow_html=True)
    
    checklist_items = {
        "Card Installation": [
            "All required cards present",
            "Cards properly seated and secured",
            "No missing or damaged cards",
            "Card configurations correct"
        ],
        "Labeling": [
            "All equipment properly labeled",
            "Labels clearly readable",
            "Label placement follows standards",
            "No missing identification tags"
        ],
        "Safety & Organization": [
            "Safety covers properly installed",
            "Grounding connections secure",
            "Warning signs present where required",
            "Overall rack organization acceptable"
        ]
    }
    
    # Create three columns for better organization
    col1, col2, col3 = st.columns(3)
    
    with col1:
        st.subheader("🔍 Card Installation")
        for item in checklist_items["Card Installation"]:
            st.write(f"• {item}")
    
    with col2:
        st.subheader("🔍 Labeling")
        for item in checklist_items["Labeling"]:
            st.write(f"• {item}")
    
    with col3:
        st.subheader("🔍 Safety & Organization")
        for item in checklist_items["Safety & Organization"]:
            st.write(f"• {item}")

def main():
    """Main application function"""
    # Custom CSS for beautiful gradient colors
    st.markdown("""
    <style>
    /* Beautiful gradient background */
    .stApp {
        background: linear-gradient(135deg, #f8f9fa 0%, #e3f2fd 25%, #f3e5f5 50%, #e8f5e8 75%, #fff3e0 100%);
        background-attachment: fixed;
    }
    
    /* Header styling with gradient */
    .main-header {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 2rem;
        border-radius: 15px;
        margin-bottom: 2rem;
        box-shadow: 0 8px 32px rgba(102, 126, 234, 0.2);
        text-align: center;
        border: 1px solid rgba(255, 255, 255, 0.2);
    }
    
    .main-header h1 {
        color: white;
        font-size: 3rem;
        margin-bottom: 0.5rem;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
        font-weight: 700;
    }
    
    .main-header p {
        color: rgba(255, 255, 255, 0.9);
        font-size: 1.2rem;
        margin: 0;
        font-weight: 300;
    }
    
    /* Sidebar styling */
    .css-1d391kg {
        background: linear-gradient(180deg, #f8f9fa 0%, #e3f2fd 100%);
        border-right: 1px solid rgba(102, 126, 234, 0.1);
    }
    
    /* Card styling */
    .stCard {
        background: linear-gradient(135deg, rgba(255,255,255,0.95) 0%, rgba(248,249,250,0.9) 100%);
        border-radius: 15px;
        padding: 1.5rem;
        margin: 1rem 0;
        box-shadow: 0 4px 20px rgba(102, 126, 234, 0.1);
        border: 1px solid rgba(102, 126, 234, 0.1);
        backdrop-filter: blur(10px);
    }
    
    /* Upload section styling */
    .upload-section {
        background: linear-gradient(135deg, rgba(227, 242, 253, 0.4) 0%, rgba(243, 229, 245, 0.4) 100%);
        padding: 1.5rem;
        border-radius: 12px;
        border: 1px solid rgba(102, 126, 234, 0.2);
        box-shadow: 0 2px 10px rgba(102, 126, 234, 0.05);
    }
    
    /* Success message styling */
    .success-message {
        background: linear-gradient(135deg, #4caf50 0%, #66bb6a 100%);
        color: white;
        padding: 1rem;
        border-radius: 10px;
        margin: 1rem 0;
        border: none;
        box-shadow: 0 2px 8px rgba(76, 175, 80, 0.3);
    }
    
    /* Warning message styling */
    .warning-message {
        background: linear-gradient(135deg, #ff9800 0%, #ffb74d 100%);
        color: white;
        padding: 1rem;
        border-radius: 10px;
        margin: 1rem 0;
        border: none;
        box-shadow: 0 2px 8px rgba(255, 152, 0, 0.3);
    }
    
    /* Error message styling */
    .error-message {
        background: linear-gradient(135deg, #f44336 0%, #ef5350 100%);
        color: white;
        padding: 1rem;
        border-radius: 10px;
        margin: 1rem 0;
        border: none;
        box-shadow: 0 2px 8px rgba(244, 67, 54, 0.3);
    }
    
    /* Button styling */
    .stButton > button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        border-radius: 8px;
        padding: 0.75rem 1.5rem;
        font-weight: 600;
        transition: all 0.3s ease;
        box-shadow: 0 2px 8px rgba(102, 126, 234, 0.3);
    }
    
    .stButton > button:hover {
        background: linear-gradient(135deg, #764ba2 0%, #667eea 100%);
        transform: translateY(-2px);
        box-shadow: 0 4px 16px rgba(102, 126, 234, 0.4);
    }
    
    /* Tab styling */
    .stTabs [data-baseweb="tab-list"] {
        background: linear-gradient(135deg, rgba(248, 249, 250, 0.8) 0%, rgba(227, 242, 253, 0.8) 100%);
        border-radius: 10px;
        border: 1px solid rgba(102, 126, 234, 0.1);
    }
    
    /* Chart container styling */
    .stPlotlyChart {
        background: rgba(255,255,255,0.95);
        border-radius: 12px;
        padding: 1.5rem;
        margin: 1rem 0;
        border: 1px solid rgba(102, 126, 234, 0.1);
        box-shadow: 0 4px 16px rgba(102, 126, 234, 0.08);
    }
    
    /* Column alignment fix */
    .row-widget.stHorizontal {
        align-items: flex-start !important;
    }
    
    /* Perfect horizontal alignment for main content boxes */
    .main-content-row {
        display: flex;
        align-items: stretch;
        justify-content: space-between;
        width: 100%;
        gap: 1rem;
    }
    
    /* Ensure equal height and perfect alignment */
    .stColumn {
        display: flex;
        flex-direction: column;
        align-items: stretch;
        flex: 1;
    }
    
    /* Force same height for content boxes */
    .upload-section, .status-box {
        min-height: 250px;
        display: flex;
        flex-direction: column;
        justify-content: flex-start;
        height: 100%;
    }
    
    /* Ensure both columns are exactly the same height */
    .stHorizontal > div {
        height: 100% !important;
        display: flex !important;
        flex-direction: column !important;
    }
    
    /* Perfect alignment for all content sections */
    .stMarkdown, .stImage, .stWarning, .stSuccess {
        margin-top: 0 !important;
        margin-bottom: 0 !important;
    }
    
    /* Ensure consistent spacing */
    .stColumn > div {
        padding-top: 0 !important;
        padding-bottom: 0 !important;
    }
    
    /* Force exact alignment */
    .main-content-row .stColumn {
        align-items: stretch !important;
        justify-content: flex-start !important;
    }
    
    /* Remove repetitive status box */
    .status-box {
        display: none !important;
    }
    
    /* Text styling for better readability */
    .stMarkdown {
        color: #2c3e50;
        line-height: 1.6;
    }
    
    /* Selectbox styling */
    .stSelectbox > div > div {
        background: rgba(255, 255, 255, 0.9);
        border: 1px solid rgba(102, 126, 234, 0.2);
        border-radius: 8px;
    }
    
    /* File uploader styling */
    .stFileUploader > div {
        background: rgba(255, 255, 255, 0.9);
        border: 2px dashed rgba(102, 126, 234, 0.3);
        border-radius: 10px;
        padding: 1rem;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Header
    st.markdown("""
    <div class="main-header">
        <h1>📡 ACCEPTIN</h1>
        <p>AI-Powered Telecom Data Center Site Quality Inspector</p>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar
    st.sidebar.title("🛠️ Controls")
    
    # Load model
    model, model_info = load_telecom_model()
    
    if model is None:
        st.error(f"❌ {model_info}")
        st.info("Please train the model first using: `python train_telecom.py`")
        return
    
    # Model info in sidebar
    st.sidebar.success("✅ Model loaded successfully")
    if isinstance(model_info, dict):
        accuracy = model_info.get('best_acc', 'Unknown')
        if isinstance(accuracy, (int, float)):
            st.sidebar.write(f"**Accuracy:** {accuracy:.2f}%")
        else:
            st.sidebar.write(f"**Accuracy:** {accuracy}")
        st.sidebar.write(f"**Architecture:** ConvNeXt 197M Parameters")
        
        # Add concise About text below Architecture
        st.sidebar.markdown("---")
        st.sidebar.markdown("**About ACCEPTIN:**")
        st.sidebar.markdown("AI-powered telecom site quality inspector using computer vision to classify installations as good/bad with confidence scores.")
    
    # Main content - 2 equal columns with perfect alignment
    st.markdown('<div class="main-content-row">', unsafe_allow_html=True)
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("""
        <div class="upload-section">
            <h3 style='color: #006400;'>📤 Upload or Capture Telecom/IT Datacenter Image</h3>
            <p style="color: #000000; font-weight: 500;">Upload an image or take a photo of the telecom/IT datacenter for quality inspection</p>
        </div>
   """, unsafe_allow_html=True)
        
        # File upload option
        uploaded_file = st.file_uploader(
            "Choose an image...",
            type=['jpg', 'jpeg', 'png', 'bmp', 'tiff'],
            help="Upload a clear image of the telecom/IT datacenter installation"
        )
        
        # Camera option (hidden but functional)
        st.markdown("""
        <style>
        [data-testid="stCameraInput"] {
            display: none !important;
        }
        </style>
        """, unsafe_allow_html=True)
        
        camera_photo = st.camera_input("Or take a photo with camera")
        
        image = None
        if uploaded_file is not None:
            image = Image.open(uploaded_file)
        elif camera_photo is not None:
            image = Image.open(camera_photo)
        
        if image is not None:
            # Create two columns for image and warning/status
            img_col, msg_col = st.columns([2, 1])
            
            with img_col:
                # Display uploaded or captured image
                st.image(image, caption="Telecom Site Image", use_container_width=True)
            
            with msg_col:
                with st.spinner("Analyzing site quality..."):
                    # Get prediction
                    transform = get_inference_transform()
                    predicted_class, confidence, probabilities = get_prediction(
                        image, model, transform
                    )
                    if predicted_class is not None:
                        class_names = ['Bad', 'Good', 'Non Telecom']
                        predicted_label = class_names[predicted_class]
                        if predicted_label == 'Non Telecom':
                            st.warning("⚠️ This image does not appear to be a telecom site. Please upload a valid telecom site photo.")
                            st.session_state.prediction_results = None
                        else:
                            # Store results in session state
                            st.session_state.prediction_results = {
                                'predicted_class': predicted_class,
                                'confidence': confidence,
                                'probabilities': probabilities,
                                'analysis': analyze_site_quality(predicted_class, confidence)
                            }
                    elif confidence is not None and confidence < 0.70:
                        st.warning(f"⚠️ Low confidence prediction ({confidence:.1%}). Please upload a higher quality image with better lighting and clearer focus for more accurate analysis.")
                        st.session_state.prediction_results = None
            
            # Move "AI Analysis complete!" message below the image
            if hasattr(st.session_state, 'prediction_results') and st.session_state.prediction_results is not None:
                st.success("✅ AI Analysis complete!")
    
    with col2:
        if hasattr(st.session_state, 'prediction_results') and st.session_state.prediction_results is not None:
            results = st.session_state.prediction_results
            analysis = results['analysis']
            
            # Display status box in middle column
            st.markdown(f"""
            <div class="{analysis['color']}">
                <h2>{analysis['icon']} {analysis['status']}</h2>
                <h3>{analysis['message']}</h3>
                <p><strong>Confidence:</strong> {analysis['confidence']:.1%}</p>
            </div>
            """, unsafe_allow_html=True)
    

    
    st.markdown('</div>', unsafe_allow_html=True)
    
    # Detailed analysis section - moved back to below the green box
    if hasattr(st.session_state, 'prediction_results') and st.session_state.prediction_results is not None:
        with col2:
            st.markdown("---")
            st.markdown("""
            <div style="text-align: center;">
                <h2>📊 Detailed Analysis</h2>
            </div>
            """, unsafe_allow_html=True)
            
            analysis = st.session_state.prediction_results['analysis']
            
            # Create two columns for Quality Assessment and Recommendations
            qa_col, rec_col = st.columns(2)
            
            with qa_col:
                st.markdown("""
                <div style="text-align: left;">
                    <h3>🔍 Quality Assessment</h3>
                </div>
                """, unsafe_allow_html=True)
                for detail in analysis['details']:
                    st.write(detail)
            
            with rec_col:
                st.markdown("""
                <div style="text-align: left;">
                    <h3>💡 Recommendations</h3>
                </div>
                """, unsafe_allow_html=True)
                for recommendation in analysis['recommendations']:
                    st.write(recommendation)
    
    # Tabs for additional features
    st.markdown("---")
    tab1, tab2 = st.tabs(["📋 Inspection Checklist", "ℹ️ About"])
    
    with tab1:
        display_inspection_checklist()
    
    with tab2:
        st.header("ℹ️ About ACCEPTIN")
        st.markdown("""
        **ACCEPTIN** is an AI-powered telecom site quality inspection system that uses computer vision 
        to automatically classify telecom installations as "good" or "bad" based on visual criteria.
        
        ### 🎯 Key Features:
        - **Transfer Learning**: Leverages pre-trained ConvNeXt model (197M parameters)
        - **Binary Classification**: Classifies sites as good/bad with confidence scores
        - **Quality Assessment**: Evaluates cable assembly, card installation, and labeling
        - **Real-time Analysis**: Instant feedback on site quality
        
        ### 🔧 Technical Details:
        - **Model**: ConvNeXt Large with custom classification head
        - **Training**: Transfer learning from food detection model
        - **Input**: 224x224 RGB images
        - **Output**: Binary classification with confidence scores
        
        ### 📊 Quality Criteria:
        - Cable assembly and routing
        - Card installation and labeling
        - General organization and safety
        - Compliance with telecom standards
        
        ### 🚀 Usage:
        1. Upload telecom site image
        2. Click "Analyze Site Quality"
        3. Review results and recommendations
        4. Use inspection checklist for verification
        """)

if __name__ == "__main__":
    main()