Update app.py

app.py

@@ -10,14 +10,23 @@ import requests
 import io
 from datetime import datetime
 import time
+import logging
 
-# Configure page
-st.set_page_config(
-    page_title="Satellite Classification Dashboard",
-    page_icon="🛰️",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
+# Set up logging to help debug issues
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Configure page - move this to the very top and add error handling
+try:
+    st.set_page_config(
+        page_title="Satellite Classification Dashboard",
+        page_icon="🛰️",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+except Exception as e:
+    logger.error(f"Error setting page config: {e}")
+    # Continue without custom config if it fails
 
 # Custom CSS for better styling
 st.markdown("""
@@ -51,6 +60,10 @@ st.markdown("""
         color: white;
         font-size: 1.2rem;
     }
+    .stAlert > div {
+        padding: 10px;
+        border-radius: 5px;
+    }
 </style>
 """, unsafe_allow_html=True)
 
@@ -134,38 +147,77 @@ MODEL_METRICS = {
 
 @st.cache_resource
 def load_model(model_name):
-    """Load model from HuggingFace with caching"""
+    """Load model from HuggingFace with caching and better error handling"""
     try:
-        with st.spinner(f'Loading {model_name}...'):
-            url = MODEL_CONFIGS[model_name]["url"]
-            response = requests.get(url)
-            response.raise_for_status()
+        logger.info(f"Loading model: {model_name}")
+        url = MODEL_CONFIGS[model_name]["url"]
+        
+        # Add timeout and better error handling
+        response = requests.get(url, timeout=60, stream=True)
+        response.raise_for_status()
+        
+        # Check if response is actually a Keras model
+        if len(response.content) < 1000:  # Too small to be a model
+            st.error(f"Model {model_name} download failed - file too small")
+            return None
             
-            model_bytes = io.BytesIO(response.content)
+        model_bytes = io.BytesIO(response.content)
+        
+        # Try to load the model with error handling
+        try:
             model = tf.keras.models.load_model(model_bytes)
+            logger.info(f"Successfully loaded model: {model_name}")
             return model
+        except Exception as load_error:
+            st.error(f"Error loading Keras model {model_name}: {str(load_error)}")
+            return None
+            
+    except requests.exceptions.Timeout:
+        st.error(f"Timeout loading {model_name}. Please try again.")
+        return None
+    except requests.exceptions.RequestException as e:
+        st.error(f"Network error loading {model_name}: {str(e)}")
+        return None
     except Exception as e:
-        st.error(f"Error loading {model_name}: {str(e)}")
+        st.error(f"Unexpected error loading {model_name}: {str(e)}")
+        logger.error(f"Error loading {model_name}: {str(e)}")
         return None
 
 def preprocess_image(image, target_size=(224, 224)):
-    """Preprocess image for model prediction"""
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    image = image.resize(target_size)
-    image_array = np.array(image) / 255.0
-    return np.expand_dims(image_array, axis=0)
+    """Preprocess image for model prediction with error handling"""
+    try:
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        image = image.resize(target_size)
+        image_array = np.array(image) / 255.0
+        return np.expand_dims(image_array, axis=0)
+    except Exception as e:
+        st.error(f"Error preprocessing image: {str(e)}")
+        return None
 
 def predict_with_model(model, image, model_name):
-    """Make prediction with a specific model"""
+    """Make prediction with a specific model with better error handling"""
+    if model is None:
+        return None
+        
     try:
         start_time = time.time()
         predictions = model.predict(image, verbose=0)
         inference_time = (time.time() - start_time) * 1000  # Convert to ms
         
+        # Validate predictions
+        if predictions is None or len(predictions) == 0:
+            st.error(f"No predictions returned from {model_name}")
+            return None
+            
         predicted_class = np.argmax(predictions[0])
         confidence = np.max(predictions[0]) * 100
         
+        # Validate class prediction
+        if predicted_class not in CLASS_NAMES:
+            st.error(f"Invalid class prediction from {model_name}: {predicted_class}")
+            return None
+        
         return {
             'class': predicted_class,
             'class_name': CLASS_NAMES[predicted_class],
@@ -175,6 +227,7 @@ def predict_with_model(model, image, model_name):
         }
     except Exception as e:
         st.error(f"Prediction error with {model_name}: {str(e)}")
+        logger.error(f"Prediction error with {model_name}: {str(e)}")
         return None
 
 def recommend_best_model(image_predictions):
@@ -199,402 +252,540 @@ def recommend_best_model(image_predictions):
     return "EfficientNetB0"
 
 def create_metrics_comparison():
-    """Create comprehensive metrics comparison dashboard"""
-    
-    # Create subplots
-    fig = make_subplots(
-        rows=2, cols=2,
-        subplot_titles=('Accuracy Comparison', 'Model Size vs Inference Time', 
-                       'Performance Metrics Radar', 'Training Efficiency'),
-        specs=[[{"type": "bar"}, {"type": "scatter"}],
-               [{"type": "scatterpolar"}, {"type": "bar"}]]
-    )
-    
-    models = list(MODEL_METRICS.keys())
-    
-    # 1. Accuracy Comparison Bar Chart
-    accuracies = [MODEL_METRICS[model]["accuracy"] for model in models]
-    fig.add_trace(
-        go.Bar(x=models, y=accuracies, name="Accuracy", 
-               marker_color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']),
-        row=1, col=1
-    )
-    
-    # 2. Model Size vs Inference Time Scatter
-    sizes = [MODEL_METRICS[model]["model_size"] for model in models]
-    times = [MODEL_METRICS[model]["inference_time"] for model in models]
-    fig.add_trace(
-        go.Scatter(x=sizes, y=times, mode='markers+text', 
-                   text=models, textposition="top center",
-                   marker=dict(size=15, color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']),
-                   name="Size vs Speed"),
-        row=1, col=2
-    )
-    
-    # 3. Radar Chart for Performance Metrics
-    metrics = ['accuracy', 'precision', 'recall', 'f1_score']
-    for i, model in enumerate(models):
-        values = [MODEL_METRICS[model][metric] for metric in metrics]
+    """Create comprehensive metrics comparison dashboard with error handling"""
+    try:
+        # Create subplots
+        fig = make_subplots(
+            rows=2, cols=2,
+            subplot_titles=('Accuracy Comparison', 'Model Size vs Inference Time', 
+                           'Performance Metrics Radar', 'Training Efficiency'),
+            specs=[[{"type": "bar"}, {"type": "scatter"}],
+                   [{"type": "scatterpolar"}, {"type": "bar"}]]
+        )
+        
+        models = list(MODEL_METRICS.keys())
+        
+        # 1. Accuracy Comparison Bar Chart
+        accuracies = [MODEL_METRICS[model]["accuracy"] for model in models]
         fig.add_trace(
-            go.Scatterpolar(r=values, theta=metrics, fill='toself', 
-                           name=model, opacity=0.7),
-            row=2, col=1
+            go.Bar(x=models, y=accuracies, name="Accuracy", 
+                   marker_color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']),
+            row=1, col=1
         )
-    
-    # 4. Training Time Comparison
-    training_times = [MODEL_METRICS[model]["training_time"] for model in models]
-    fig.add_trace(
-        go.Bar(x=models, y=training_times, name="Training Time",
-               marker_color=['#9467bd', '#8c564b', '#e377c2', '#7f7f7f']),
-        row=2, col=2
-    )
-    
-    # Update layout
-    fig.update_layout(height=800, showlegend=True, 
-                     title_text="Comprehensive Model Comparison Dashboard")
-    fig.update_xaxes(title_text="Models", row=1, col=1)
-    fig.update_yaxes(title_text="Accuracy (%)", row=1, col=1)
-    fig.update_xaxes(title_text="Model Size (MB)", row=1, col=2)
-    fig.update_yaxes(title_text="Inference Time (ms)", row=1, col=2)
-    fig.update_xaxes(title_text="Models", row=2, col=2)
-    fig.update_yaxes(title_text="Training Time (minutes)", row=2, col=2)
-    
-    return fig
-
-def create_class_distribution_chart():
-    """Create class distribution visualization"""
-    classes = list(CLASS_NAMES.values())
-    samples = [7500 if cls != 'Debris' else 15000 for cls in classes]
-    percentages = [8.33 if cls != 'Debris' else 16.67 for cls in classes]
-    
-    fig = go.Figure()
-    fig.add_trace(go.Bar(
-        x=classes,
-        y=samples,
-        text=[f'{s} ({p}%)' for s, p in zip(samples, percentages)],
-        textposition='auto',
-        marker_color=['#ff6b6b' if cls == 'Debris' else '#4ecdc4' for cls in classes]
-    ))
-    
-    fig.update_layout(
-        title="Class Distribution in Training Dataset",
-        xaxis_title="Satellite Classes",
-        yaxis_title="Number of Samples",
-        height=400
-    )
-    
-    return fig
-
-# Main App
-def main():
-    # Header
-    st.markdown('<h1 class="main-header">🛰️ Satellite Classification Dashboard</h1>', 
-                unsafe_allow_html=True)
-    
-    # Sidebar
-    st.sidebar.title("Navigation")
-    page = st.sidebar.selectbox("Choose a page", 
-                               ["🏠 Home", "📊 Model Comparison", "🔍 Image Classification", 
-                                "📈 Performance Analytics", "ℹ️ About Models"])
-    
-    if page == "🏠 Home":
-        st.markdown("## Welcome to the Satellite Classification System")
         
-        col1, col2 = st.columns(2)
+        # 2. Model Size vs Inference Time Scatter
+        sizes = [MODEL_METRICS[model]["model_size"] for model in models]
+        times = [MODEL_METRICS[model]["inference_time"] for model in models]
+        fig.add_trace(
+            go.Scatter(x=sizes, y=times, mode='markers+text', 
+                       text=models, textposition="top center",
+                       marker=dict(size=15, color=['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']),
+                       name="Size vs Speed"),
+            row=1, col=2
+        )
         
-        with col1:
-            st.markdown("### 🎯 System Overview")
-            st.write("""
-            This dashboard provides comprehensive satellite classification using 4 different
-            deep learning models. Upload satellite images to classify them into 11 different
-            categories including various satellites and space debris.
-            """)
-            
-            st.markdown("### 🛰️ Supported Classes")
-            for i, (class_id, class_name) in enumerate(CLASS_NAMES.items()):
-                if i < 6:  # First column
-                    st.write(f"• **{class_name}**")
+        # 3. Radar Chart for Performance Metrics
+        metrics = ['accuracy', 'precision', 'recall', 'f1_score']
+        colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']
+        for i, model in enumerate(models):
+            values = [MODEL_METRICS[model][metric] for metric in metrics]
+            fig.add_trace(
+                go.Scatterpolar(r=values, theta=metrics, fill='toself', 
+                               name=model, opacity=0.7, line_color=colors[i]),
+                row=2, col=1
+            )
         
-        with col2:
-            st.markdown("### 🤖 Available Models")
-            st.write("""
-            - **Custom CNN**: Tailored architecture for satellite imagery
-            - **MobileNetV2**: Lightweight and fast inference
-            - **EfficientNetB0**: Best accuracy-efficiency balance
-            - **DenseNet121**: Complex pattern recognition
-            """)
-            
-            st.markdown("### 📊 Class Distribution")
-            for i, (class_id, class_name) in enumerate(CLASS_NAMES.items()):
-                if i >= 6:  # Second column
-                    st.write(f"• **{class_name}**")
+        # 4. Training Time Comparison
+        training_times = [MODEL_METRICS[model]["training_time"] for model in models]
+        fig.add_trace(
+            go.Bar(x=models, y=training_times, name="Training Time",
+                   marker_color=['#9467bd', '#8c564b', '#e377c2', '#7f7f7f']),
+            row=2, col=2
+        )
         
-        # Class distribution chart
-        st.plotly_chart(create_class_distribution_chart(), use_container_width=True)
-    
-    elif page == "📊 Model Comparison":
-        st.markdown("## 📊 Model Performance Comparison")
+        # Update layout
+        fig.update_layout(height=800, showlegend=True, 
+                         title_text="Comprehensive Model Comparison Dashboard")
+        fig.update_xaxes(title_text="Models", row=1, col=1)
+        fig.update_yaxes(title_text="Accuracy (%)", row=1, col=1)
+        fig.update_xaxes(title_text="Model Size (MB)", row=1, col=2)
+        fig.update_yaxes(title_text="Inference Time (ms)", row=1, col=2)
+        fig.update_xaxes(title_text="Models", row=2, col=2)
+        fig.update_yaxes(title_text="Training Time (minutes)", row=2, col=2)
         
-        # Metrics table
-        st.markdown("### Performance Metrics Summary")
-        df_metrics = pd.DataFrame(MODEL_METRICS).T
-        st.dataframe(df_metrics.style.highlight_max(axis=0), use_container_width=True)
+        return fig
+    except Exception as e:
+        st.error(f"Error creating metrics comparison chart: {str(e)}")
+        return None
+
+def create_class_distribution_chart():
+    """Create class distribution visualization with error handling"""
+    try:
+        classes = list(CLASS_NAMES.values())
+        samples = [7500 if cls != 'Debris' else 15000 for cls in classes]
+        percentages = [8.33 if cls != 'Debris' else 16.67 for cls in classes]
         
-        # Comprehensive comparison chart
-        st.plotly_chart(create_metrics_comparison(), use_container_width=True)
+        fig = go.Figure()
+        fig.add_trace(go.Bar(
+            x=classes,
+            y=samples,
+            text=[f'{s} ({p:.1f}%)' for s, p in zip(samples, percentages)],
+            textposition='auto',
+            marker_color=['#ff6b6b' if cls == 'Debris' else '#4ecdc4' for cls in classes]
+        ))
         
-        # Model recommendations
-        st.markdown("### 🎯 Model Selection Guide")
+        fig.update_layout(
+            title="Class Distribution in Training Dataset",
+            xaxis_title="Satellite Classes",
+            yaxis_title="Number of Samples",
+            height=400,
+            xaxis_tickangle=-45
+        )
         
-        col1, col2 = st.columns(2)
+        return fig
+    except Exception as e:
+        st.error(f"Error creating class distribution chart: {str(e)}")
+        return None
+
+def create_confusion_matrix_heatmap():
+    """Create a sample confusion matrix heatmap for demonstration"""
+    try:
+        # Sample confusion matrix data (you would replace this with actual data)
+        classes = list(CLASS_NAMES.values())
+        np.random.seed(42)  # For reproducible demo data
         
-        with col1:
-            st.markdown("#### 🏆 Best for Accuracy")
-            st.success("**EfficientNetB0** - 96.4% accuracy")
-            
-            st.markdown("#### ⚡ Best for Speed")
-            st.info("**MobileNetV2** - 18ms inference time")
+        # Create a realistic-looking confusion matrix
+        confusion_matrix = np.random.randint(0, 100, size=(len(classes), len(classes)))
+        # Make diagonal elements higher (correct predictions)
+        np.fill_diagonal(confusion_matrix, np.random.randint(400, 500, size=len(classes)))
         
-        with col2:
-            st.markdown("#### 💾 Most Lightweight")
-            st.info("**MobileNetV2** - 8.7MB model size")
-            
-            st.markdown("#### 🎯 Best Overall Balance")
-            st.warning("**EfficientNetB0** - High accuracy + efficiency")
-    
-    elif page == "🔍 Image Classification":
-        st.markdown("## 🔍 Image Classification")
+        fig = go.Figure(data=go.Heatmap(
+            z=confusion_matrix,
+            x=classes,
+            y=classes,
+            colorscale='Blues',
+            showscale=True
+        ))
         
-        uploaded_file = st.file_uploader(
-            "Upload a satellite image", 
-            type=['png', 'jpg', 'jpeg'],
-            help="Upload an image of a satellite or space object for classification"
+        fig.update_layout(
+            title="Sample Confusion Matrix (Demo Data)",
+            xaxis_title="Predicted Class",
+            yaxis_title="True Class",
+            height=600
         )
         
-        if uploaded_file is not None:
-            # Display uploaded image
-            image = Image.open(uploaded_file)
+        return fig
+    except Exception as e:
+        st.error(f"Error creating confusion matrix: {str(e)}")
+        return None
+
+# Main App
+def main():
+    try:
+        # Header
+        st.markdown('<h1 class="main-header">🛰️ Satellite Classification Dashboard</h1>', 
+                    unsafe_allow_html=True)
+        
+        # Sidebar
+        st.sidebar.title("Navigation")
+        page = st.sidebar.selectbox("Choose a page", 
+                                   ["🏠 Home", "📊 Model Comparison", "🔍 Image Classification", 
+                                    "📈 Performance Analytics", "ℹ️ About Models"])
+        
+        # Add sidebar information
+        st.sidebar.markdown("---")
+        st.sidebar.markdown("### System Info")
+        st.sidebar.info(f"Total Classes: {len(CLASS_NAMES)}")
+        st.sidebar.info(f"Available Models: {len(MODEL_CONFIGS)}")
+        st.sidebar.info("Built with Streamlit & TensorFlow")
+        
+        if page == "🏠 Home":
+            st.markdown("## Welcome to the Satellite Classification System")
             
-            col1, col2 = st.columns([1, 2])
+            col1, col2 = st.columns(2)
             
             with col1:
-                st.image(image, caption="Uploaded Image", use_container_width=True)
+                st.markdown("### 🎯 System Overview")
+                st.write("""
+                This dashboard provides comprehensive satellite classification using 4 different
+                deep learning models. Upload satellite images to classify them into 11 different
+                categories including various satellites and space debris.
+                """)
+                
+                st.markdown("### 🛰️ Supported Classes")
+                for i, (class_id, class_name) in enumerate(CLASS_NAMES.items()):
+                    if i < 6:  # First column
+                        st.write(f"• **{class_name}**")
             
             with col2:
-                st.markdown("### Image Details")
-                st.write(f"**Filename:** {uploaded_file.name}")
-                st.write(f"**Size:** {image.size}")
-                st.write(f"**Mode:** {image.mode}")
-            
-            # Model selection
-            selected_models = st.multiselect(
-                "Select models for prediction",
-                list(MODEL_CONFIGS.keys()),
-                default=["EfficientNetB0", "Custom CNN"]
-            )
+                st.markdown("### 🤖 Available Models")
+                st.write("""
+                - **Custom CNN**: Tailored architecture for satellite imagery
+                - **MobileNetV2**: Lightweight and fast inference
+                - **EfficientNetB0**: Best accuracy-efficiency balance
+                - **DenseNet121**: Complex pattern recognition
+                """)
+                
+                st.markdown("### 📊 Remaining Classes")
+                for i, (class_id, class_name) in enumerate(CLASS_NAMES.items()):
+                    if i >= 6:  # Second column
+                        st.write(f"• **{class_name}**")
             
-            if st.button("🚀 Classify Image", type="primary"):
-                if not selected_models:
-                    st.warning("Please select at least one model.")
-                    return
+            # Class distribution chart
+            chart = create_class_distribution_chart()
+            if chart:
+                st.plotly_chart(chart, use_container_width=True)
                 
-                # Preprocess image
-                processed_image = preprocess_image(image)
+            # Quick start guide
+            st.markdown("### 🚀 Quick Start Guide")
+            st.markdown("""
+            1. Navigate to **🔍 Image Classification** to upload and classify satellite images
+            2. Check **📊 Model Comparison** to compare different model performances
+            3. Explore **📈 Performance Analytics** for detailed metrics
+            4. Read **ℹ️ About Models** to understand each model's capabilities
+            """)
+        
+        elif page == "📊 Model Comparison":
+            st.markdown("## 📊 Model Performance Comparison")
+            
+            # Metrics table
+            st.markdown("### Performance Metrics Summary")
+            df_metrics = pd.DataFrame(MODEL_METRICS).T
+            st.dataframe(df_metrics.style.highlight_max(axis=0), use_container_width=True)
+            
+            # Comprehensive comparison chart
+            chart = create_metrics_comparison()
+            if chart:
+                st.plotly_chart(chart, use_container_width=True)
+            
+            # Model recommendations
+            st.markdown("### 🎯 Model Selection Guide")
+            
+            col1, col2 = st.columns(2)
+            
+            with col1:
+                st.markdown("#### 🏆 Best for Accuracy")
+                st.success("**EfficientNetB0** - 96.4% accuracy")
                 
-                # Store predictions
-                predictions = {}
+                st.markdown("#### ⚡ Best for Speed")
+                st.info("**MobileNetV2** - 18ms inference time")
+            
+            with col2:
+                st.markdown("#### 💾 Most Lightweight")
+                st.info("**MobileNetV2** - 8.7MB model size")
                 
-                # Create progress bar
-                progress_bar = st.progress(0)
-                status_text = st.empty()
+                st.markdown("#### 🎯 Best Overall Balance")
+                st.warning("**EfficientNetB0** - High accuracy + efficiency")
                 
-                # Make predictions with selected models
-                for i, model_name in enumerate(selected_models):
-                    status_text.text(f'Loading {model_name}...')
-                    model = load_model(model_name)
-                    
-                    if model:
-                        status_text.text(f'Predicting with {model_name}...')
-                        pred = predict_with_model(model, processed_image, model_name)
-                        predictions[model_name] = pred
+            # Model rankings
+            st.markdown("### 🏅 Model Rankings")
+            
+            # Calculate overall scores
+            rankings = []
+            for model_name, metrics in MODEL_METRICS.items():
+                # Weighted score: accuracy (40%), speed (30%), size (30%)
+                score = (metrics['accuracy'] * 0.4 + 
+                        (100 - metrics['inference_time']) * 0.3 + 
+                        (50 - metrics['model_size']) * 0.3)
+                rankings.append({'Model': model_name, 'Overall Score': round(score, 1)})
+            
+            rankings_df = pd.DataFrame(rankings).sort_values('Overall Score', ascending=False)
+            st.dataframe(rankings_df, use_container_width=True)
+        
+        elif page == "🔍 Image Classification":
+            st.markdown("## 🔍 Image Classification")
+            
+            # Instructions
+            st.info("""
+            📋 **Instructions:**
+            1. Upload a satellite or space object image (PNG, JPG, or JPEG)
+            2. Select one or more models for classification
+            3. Click 'Classify Image' to get predictions
+            4. View results, confidence scores, and recommendations
+            """)
+            
+            uploaded_file = st.file_uploader(
+                "Upload a satellite image", 
+                type=['png', 'jpg', 'jpeg'],
+                help="Upload an image of a satellite or space object for classification"
+            )
+            
+            if uploaded_file is not None:
+                try:
+                    # Display uploaded image
+                    image = Image.open(uploaded_file)
                     
-                    progress_bar.progress((i + 1) / len(selected_models))
-                
-                status_text.empty()
-                progress_bar.empty()
-                
-                # Display results
-                if predictions:
-                    # Get recommendation
-                    recommended_model = recommend_best_model(predictions)
+                    col1, col2 = st.columns([1, 2])
                     
-                    st.markdown("### 🎯 Prediction Results")
+                    with col1:
+                        st.image(image, caption="Uploaded Image", use_container_width=True)
                     
-                    # Show recommendation
-                    st.markdown(f"""
-                    <div class="prediction-box">
-                        <h3>🏆 Recommended Model: {recommended_model}</h3>
-                        <p>Based on confidence and model performance</p>
-                    </div>
-                    """, unsafe_allow_html=True)
+                    with col2:
+                        st.markdown("### Image Details")
+                        st.write(f"**Filename:** {uploaded_file.name}")
+                        st.write(f"**Size:** {image.size}")
+                        st.write(f"**Mode:** {image.mode}")
+                        st.write(f"**File Size:** {len(uploaded_file.getvalue())} bytes")
                     
-                    # Results table
-                    results_data = []
-                    for model_name, pred in predictions.items():
-                        if pred:
-                            results_data.append({
-                                'Model': model_name,
-                                'Predicted Class': pred['class_name'],
-                                'Confidence (%)': f"{pred['confidence']:.1f}%",
-                                'Inference Time (ms)': f"{pred['inference_time']:.1f}",
-                                'Recommended': '🏆' if model_name == recommended_model else ''
-                            })
+                    # Model selection
+                    st.markdown("### Select Models for Classification")
+                    selected_models = st.multiselect(
+                        "Choose models to run predictions with:",
+                        list(MODEL_CONFIGS.keys()),
+                        default=["EfficientNetB0"],  # Start with just one model to avoid timeouts
+                        help="Select one or more models. More models = longer processing time."
+                    )
                     
-                    if results_data:
-                        df_results = pd.DataFrame(results_data)
-                        st.dataframe(df_results, use_container_width=True)
+                    if st.button("🚀 Classify Image", type="primary"):
+                        if not selected_models:
+                            st.warning("Please select at least one model.")
+                            return
                         
-                        # Confidence comparison
-                        st.markdown("### 📊 Confidence Comparison")
-                        confidences = [pred['confidence'] for pred in predictions.values() if pred]
-                        model_names = [name for name, pred in predictions.items() if pred]
+                        # Preprocess image
+                        processed_image = preprocess_image(image)
+                        if processed_image is None:
+                            st.error("Failed to preprocess image")
+                            return
                         
-                        fig_conf = go.Figure()
-                        fig_conf.add_trace(go.Bar(
-                            x=model_names,
-                            y=confidences,
-                            marker_color=['gold' if name == recommended_model else 'lightblue' 
-                                        for name in model_names]
-                        ))
-                        fig_conf.update_layout(
-                            title="Prediction Confidence by Model",
-                            xaxis_title="Models",
-                            yaxis_title="Confidence (%)",
-                            height=400
-                        )
-                        st.plotly_chart(fig_conf, use_container_width=True)
+                        # Store predictions
+                        predictions = {}
                         
-                        # Probability distribution for recommended model
-                        if recommended_model in predictions and predictions[recommended_model]:
-                            st.markdown(f"### 🔍 Detailed Probabilities - {recommended_model}")
-                            probs = predictions[recommended_model]['probabilities']
-                            prob_df = pd.DataFrame({
-                                'Class': [CLASS_NAMES[i] for i in range(len(probs))],
-                                'Probability': probs * 100
-                            }).sort_values('Probability', ascending=False)
+                        # Create progress bar
+                        progress_bar = st.progress(0)
+                        status_text = st.empty()
+                        
+                        # Make predictions with selected models
+                        for i, model_name in enumerate(selected_models):
+                            try:
+                                status_text.text(f'Loading {model_name}... ({i+1}/{len(selected_models)})')
+                                model = load_model(model_name)
+                                
+                                if model:
+                                    status_text.text(f'Predicting with {model_name}... ({i+1}/{len(selected_models)})')
+                                    pred = predict_with_model(model, processed_image, model_name)
+                                    if pred:
+                                        predictions[model_name] = pred
+                                    else:
+                                        st.warning(f"Failed to get prediction from {model_name}")
+                                else:
+                                    st.warning(f"Failed to load {model_name}")
+                                
+                            except Exception as e:
+                                st.error(f"Error processing {model_name}: {str(e)}")
+                                logger.error(f"Error processing {model_name}: {str(e)}")
                             
-                            fig_prob = px.bar(
-                                prob_df.head(5), 
-                                x='Probability', 
-                                y='Class',
-                                orientation='h',
-                                title=f"Top 5 Class Probabilities - {recommended_model}"
-                            )
-                            st.plotly_chart(fig_prob, use_container_width=True)
-    
-    elif page == "📈 Performance Analytics":
-        st.markdown("## 📈 Performance Analytics")
-        
-        # Performance overview
-        col1, col2, col3, col4 = st.columns(4)
-        
-        with col1:
-            st.metric("Best Accuracy", "96.4%", "EfficientNetB0")
-        with col2:
-            st.metric("Fastest Inference", "18ms", "MobileNetV2")
-        with col3:
-            st.metric("Smallest Model", "8.7MB", "MobileNetV2")
-        with col4:
-            st.metric("Total Classes", "11", "Satellites + Debris")
-        
-        # Detailed analytics
-        tab1, tab2, tab3 = st.tabs(["Accuracy Analysis", "Efficiency Metrics", "Model Comparison"])
+                            progress_bar.progress((i + 1) / len(selected_models))
+                        
+                        status_text.empty()
+                        progress_bar.empty()
+                        
+                        # Display results
+                        if predictions:
+                            # Get recommendation
+                            recommended_model = recommend_best_model(predictions)
+                            
+                            st.markdown("### 🎯 Prediction Results")
+                            
+                            # Show recommendation
+                            st.markdown(f"""
+                            <div class="prediction-box">
+                                <h3>🏆 Recommended Model: {recommended_model}</h3>
+                                <p>Based on confidence and model performance</p>
+                            </div>
+                            """, unsafe_allow_html=True)
+                            
+                            # Results table
+                            results_data = []
+                            for model_name, pred in predictions.items():
+                                if pred:
+                                    results_data.append({
+                                        'Model': model_name,
+                                        'Predicted Class': pred['class_name'],
+                                        'Confidence (%)': f"{pred['confidence']:.1f}%",
+                                        'Inference Time (ms)': f"{pred['inference_time']:.1f}",
+                                        'Recommended': '🏆' if model_name == recommended_model else ''
+                                    })
+                            
+                            if results_data:
+                                df_results = pd.DataFrame(results_data)
+                                st.dataframe(df_results, use_container_width=True)
+                                
+                                # Confidence comparison
+                                if len(predictions) > 1:
+                                    st.markdown("### 📊 Confidence Comparison")
+                                    confidences = [pred['confidence'] for pred in predictions.values() if pred]
+                                    model_names = [name for name, pred in predictions.items() if pred]
+                                    
+                                    try:
+                                        fig_conf = go.Figure()
+                                        fig_conf.add_trace(go.Bar(
+                                            x=model_names,
+                                            y=confidences,
+                                            marker_color=['gold' if name == recommended_model else 'lightblue' 
+                                                        for name in model_names]
+                                        ))
+                                        fig_conf.update_layout(
+                                            title="Prediction Confidence by Model",
+                                            xaxis_title="Models",
+                                            yaxis_title="Confidence (%)",
+                                            height=400
+                                        )
+                                        st.plotly_chart(fig_conf, use_container_width=True)
+                                    except Exception as e:
+                                        st.warning(f"Could not create confidence chart: {str(e)}")
+                                
+                                # Probability distribution for recommended model
+                                if recommended_model in predictions and predictions[recommended_model]:
+                                    try:
+                                        st.markdown(f"### 🔍 Detailed Probabilities - {recommended_model}")
+                                        probs = predictions[recommended_model]['probabilities']
+                                        prob_df = pd.DataFrame({
+                                            'Class': [CLASS_NAMES[i] for i in range(len(probs))],
+                                            'Probability': probs * 100
+                                        }).sort_values('Probability', ascending=False)
+                                        
+                                        fig_prob = px.bar(
+                                            prob_df.head(5), 
+                                            x='Probability', 
+                                            y='Class',
+                                            orientation='h',
+                                            title=f"Top 5 Class Probabilities - {recommended_model}",
+                                            color='Probability',
+                                            color_continuous_scale='viridis'
+                                        )
+                                        st.plotly_chart(fig_prob, use_container_width=True)
+                                    except Exception as e:
+                                        st.warning(f"Could not create probability chart: {str(e)}")
+                        else:
+                            st.error("No successful predictions were made. Please try again with different models.")
+                
+                except Exception as e:
+                    st.error(f"Error processing uploaded image: {str(e)}")
+                    logger.error(f"Error processing uploaded image: {str(e)}")
         
-        with tab1:
-            # Accuracy breakdown
-            models = list(MODEL_METRICS.keys())
-            metrics_list = ['accuracy', 'precision', 'recall', 'f1_score']
+        elif page == "📈 Performance Analytics":
+            st.markdown("## 📈 Performance Analytics")
             
-            for metric in metrics_list:
-                values = [MODEL_METRICS[model][metric] for model in models]
-                fig = go.Figure()
-                fig.add_trace(go.Bar(x=models, y=values, name=metric.title()))
-                fig.update_layout(title=f"{metric.title()} Comparison", height=300)
-                st.plotly_chart(fig, use_container_width=True)
-        
-        with tab2:
-            # Efficiency metrics
-            col1, col2 = st.columns(2)
+            # Performance overview
+            col1, col2, col3, col4 = st.columns(4)
             
             with col1:
-                # Inference time
-                times = [MODEL_METRICS[model]["inference_time"] for model in models]
-                fig_time = go.Figure()
-                fig_time.add_trace(go.Bar(x=models, y=times, 
-                                        marker_color=['red' if t > 40 else 'green' for t in times]))
-                fig_time.update_layout(title="Inference Time (ms)", height=400)
-                st.plotly_chart(fig_time, use_container_width=True)
-            
+                st.metric("Best Accuracy", "96.4%", "EfficientNetB0")
             with col2:
-                # Model size
-                sizes = [MODEL_METRICS[model]["model_size"] for model in models]
-                fig_size = go.Figure()
-                fig_size.add_trace(go.Bar(x=models, y=sizes,
-                                        marker_color=['red' if s > 25 else 'green' for s in sizes]))
-                fig_size.update_layout(title="Model Size (MB)", height=400)
-                st.plotly_chart(fig_size, use_container_width=True)
-        
-        with tab3:
-            # Side-by-side comparison
-            comparison_data = []
-            for model in models:
-                metrics = MODEL_METRICS[model]
-                comparison_data.append({
-                    'Model': model,
-                    'Accuracy (%)': metrics['accuracy'],
-                    'Inference Time (ms)': metrics['inference_time'],
-                    'Model Size (MB)': metrics['model_size'],
-                    'Training Time (min)': metrics['training_time'],
-                    'Efficiency Score': round(metrics['accuracy'] / (metrics['inference_time'] * 0.1 + metrics['model_size'] * 0.1), 2)
-                })
+                st.metric("Fastest Inference", "18ms", "MobileNetV2")
+            with col3:
+                st.metric("Smallest Model", "8.7MB", "MobileNetV2")
+            with col4:
+                st.metric("Total Classes", "11", "Satellites + Debris")
             
-            df_comparison = pd.DataFrame(comparison_data)
-            st.dataframe(df_comparison.style.highlight_max(axis=0, subset=['Accuracy (%)', 'Efficiency Score'])
-                                          .highlight_min(axis=0, subset=['Inference Time (ms)', 'Model Size (MB)', 'Training Time (min)']), 
-                        use_container_width=True)
-    
-    elif page == "ℹ️ About Models":
-        st.markdown("## ℹ️ Model Information")
-        
-        for model_name, config in MODEL_CONFIGS.items():
-            with st.expander(f"📋 {model_name}", expanded=False):
-                col1, col2 = st.columns(2)
-                
-                with col1:
-                    st.markdown("### Description")
-                    st.write(config["description"])
-                    
-                    st.markdown("### Input Shape")
-                    st.code(f"{config['input_shape']}")
-                    
-                    st.markdown("### Model URL")
-                    st.code(config["url"])
-                
-                with col2:
-                    st.markdown("### Strengths")
-                    for strength in config["strengths"]:
-                        st.write(f"• {strength}")
+            # Detailed analytics
+            tab1, tab2, tab3, tab4 = st.tabs(["Accuracy Analysis", "Efficiency Metrics", "Model Comparison", "Confusion Matrix"])
+            
+            with tab1:
+                try:
+                    # Accuracy breakdown
+                    models = list(MODEL_METRICS.keys())
+                    metrics_list = ['accuracy', 'precision', 'recall', 'f1_score']
                     
-                    st.markdown("### Best Use Cases")
-                    for use_case in config["best_for"]:
-                        st.write(f"• {use_case}")
+                    for metric in metrics_list:
+                        values = [MODEL_METRICS[model][metric] for model in models]
+                        fig = go.Figure()
+                        fig.add_trace(go.Bar(
+                            x=models, 
+                            y=values, 
+                            name=metric.title(),
+                            marker_color='lightblue',
+                            text=[f'{v:.1f}%' for v in values],
+                            textposition='auto'
+                        ))
+                        fig.update_layout(
+                            title=f"{metric.title()} Comparison", 
+                            height=300,
+                            yaxis_title=f"{metric.title()} (%)"
+                        )
+                        st.plotly_chart(fig, use_container_width=True)
+                except Exception as e:
+                    st.error(f"Error creating accuracy charts: {str(e)}")
+            
+            with tab2:
+                try:
+                    # Efficiency metrics
+                    col1, col2 = st.columns(2)
                     
-                    # Performance summary
-                    metrics = MODEL_METRICS[model_name]
-                    st.markdown("### Key Metrics")
-                    st.write(f"**Accuracy:** {metrics['accuracy']}%")
-                    st.write(f"**Inference Time:** {metrics['inference_time']}ms")
-                    st.write(f"**Model Size:** {metrics['model_size']}MB")
+                    with col1:
+                        # Inference time
+                        times = [MODEL_METRICS[model]["inference_time"] for model in models]
+                        fig_time = go.Figure()
+                                                fig_time.add_trace(go.Bar(
+                            x=models, 
+                            y=times,
+                            name="Inference Time",
+                            marker_color='orange',
+                            text=[f'{t:.1f} ms' for t in times],
+                            textposition='auto'
+                        ))
+                        fig_time.update_layout(
+                            title="Inference Time per Model",
+                            yaxis_title="Time (ms)",
+                            height=300
+                        )
+                        st.plotly_chart(fig_time, use_container_width=True)
+
+                    with col2:
+                        # Model sizes
+                        sizes = [MODEL_METRICS[model]["model_size"] for model in models]
+                        fig_size = go.Figure()
+                        fig_size.add_trace(go.Bar(
+                            x=models,
+                            y=sizes,
+                            name="Model Size",
+                            marker_color='green',
+                            text=[f'{s:.1f} MB' for s in sizes],
+                            textposition='auto'
+                        ))
+                        fig_size.update_layout(
+                            title="Model Size per Model",
+                            yaxis_title="Size (MB)",
+                            height=300
+                        )
+                        st.plotly_chart(fig_size, use_container_width=True)
+
+                except Exception as e:
+                    st.error(f"Error displaying efficiency metrics: {str(e)}")
+
+            with tab3:
+                # Reuse full comparison dashboard
+                comp_fig = create_metrics_comparison()
+                if comp_fig:
+                    st.plotly_chart(comp_fig, use_container_width=True)
+
+            with tab4:
+                # Display the confusion matrix
+                cm_fig = create_confusion_matrix_heatmap()
+                if cm_fig:
+                    st.plotly_chart(cm_fig, use_container_width=True)
+
+        elif page == "ℹ️ About Models":
+            st.markdown("## ℹ️ Model Details and Use Cases")
+
+            for model_name, config in MODEL_CONFIGS.items():
+                with st.expander(f"🔍 {model_name}"):
+                    st.markdown(f"<div class='model-card'><h4>{model_name}</h4>", unsafe_allow_html=True)
+                    st.markdown(f"**Description:** {config['description']}")
+                    st.markdown(f"**Input Shape:** {config['input_shape']}")
+                    st.markdown("**Strengths:**")
+                    for s in config['strengths']:
+                        st.markdown(f"• {s}")
+                    st.markdown("**Best For:**")
+                    for use in config['best_for']:
+                        st.markdown(f"• {use}")
+                    st.markdown("</div>", unsafe_allow_html=True)
+
+    except Exception as e:
+        st.error(f"An unexpected error occurred: {str(e)}")
+        logger.error(f"Main app error: {str(e)}")
+
 
+# Run the app
 if __name__ == "__main__":
-    main()
+    main()