Update app.py

app.py

@@ -26,29 +26,37 @@ API_URL = (
 
 # --- LOAD MODELS ---
 def load_models():
-    vgg_model = load_model(
-        'vgg16_focal_unfreeze_more.keras',
-        custom_objects={'BinaryFocalCrossentropy': BinaryFocalCrossentropy}
-    )
-    def focal_loss_fixed(gamma=2., alpha=.25):
-        import tensorflow.keras.backend as K
-        def loss_fn(y_true, y_pred):
-            eps = K.epsilon()
-            y_pred = K.clip(y_pred, eps, 1. - eps)
-            ce = -y_true * K.log(y_pred)
-            w = alpha * K.pow(1 - y_pred, gamma)
-            return K.mean(w * ce, axis=-1)
-        return loss_fn
-    xce_model = load_model(
-        'severity_post_tta.keras',
-        custom_objects={'focal_loss_fixed': focal_loss_fixed()}
-    )
-    rf_model = joblib.load('ensemble_rf_model.pkl')
-    xgb_model = joblib.load('ensemble_xgb_model.pkl')
-    lr_model = joblib.load('wildfire_logistic_model_synthetic.joblib')
-    return vgg_model, xce_model, rf_model, xgb_model, lr_model
+    try:
+        vgg_model = load_model(
+            'vgg16_focal_unfreeze_more.keras',
+            custom_objects={'BinaryFocalCrossentropy': BinaryFocalCrossentropy}
+        )
+        
+        def focal_loss_fixed(gamma=2., alpha=.25):
+            import tensorflow.keras.backend as K
+            def loss_fn(y_true, y_pred):
+                eps = K.epsilon()
+                y_pred = K.clip(y_pred, eps, 1. - eps)
+                ce = -y_true * K.log(y_pred)
+                w = alpha * K.pow(1 - y_pred, gamma)
+                return K.mean(w * ce, axis=-1)
+            return loss_fn
+        
+        xce_model = load_model(
+            'severity_post_tta.keras',
+            custom_objects={'focal_loss_fixed': focal_loss_fixed()}
+        )
+        rf_model = joblib.load('ensemble_rf_model.pkl')
+        xgb_model = joblib.load('ensemble_xgb_model.pkl')
+        lr_model = joblib.load('wildfire_logistic_model_synthetic.joblib')
+        return vgg_model, xce_model, rf_model, xgb_model, lr_model
+    except Exception as e:
+        raise gr.Error(f"Model loading failed: {str(e)}")
 
-vgg_model, xception_model, rf_model, xgb_model, lr_model = load_models()
+try:
+    vgg_model, xception_model, rf_model, xgb_model, lr_model = load_models()
+except Exception as e:
+    print(f"Initial model loading failed: {str(e)}")
 
 # --- RULES & TEMPLATES ---
 target_map = {0: 'mild', 1: 'moderate', 2: 'severe'}
@@ -167,57 +175,53 @@ def pipeline(image):
     )
 
 # --- GRADIO APP ---
-with gr.Blocks(css=""" 
-  .gradio-container {
-    background-color: #f5f7fa !important;
-  }
-  .gradio-textbox textarea {
-    background-color: #ffffff !important;
-    border: 1px solid #cbd2d9 !important;
-    border-radius: 8px !important;
-    padding: 12px !important;
-    font-size: 1rem !important;
-    color: #333333 !important;
-    min-height: 3em !important;
-  }
-  .gradio-accordion {
-    background-color: #ffffff !important;
-    border: 1px solid #cbd2d9 !important;
-    border-radius: 8px !important;
-    padding: 8px !important;
-  }
-  .gradio-button {
-    background-color: #0072ce !important;
-    color: white !important;
-    border-radius: 6px !important;
-    padding: 8px 16px !important;
-    font-weight: 600 !important;
-  }
-  .gradio-button:hover {
-    background-color: #005bb5 !important;
-  }
-  .gradio-markdown h1, .gradio-markdown h2 {
-    color: #1f2937 !important;
-    margin-bottom: 0.5em !important;
-  }
-""") as demo:
-    gr.Markdown("# Wildfire Detection & Management Assistant")
-    gr.Markdown("Upload a forest image from Pakistan; the system will detect fire, assess severity, analyze weather trends, and provide in-depth recommendations.")
+custom_css = """
+#component-0 { max-width: 800px; margin: 0 auto; }
+.gradio-container { background: #f0f4f7; }
+#upload-wildfire-image { min-height: 300px; }
+.panel {
+    background: white !important;
+    border-radius: 12px !important;
+    padding: 20px !important;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+}
+.status-box { 
+    background: #fff3e6 !important; 
+    border: 1px solid #ffd8b3 !important;
+}
+.dark-red { color: #cc0000 !important; }
+.green { color: #008000 !important; }
+"""
+
+with gr.Blocks(css=custom_css) as demo:
+    gr.Markdown("# 🔥 Wildfire Detection & Management Assistant")
     
-    with gr.Row():
-        inp = gr.Image(type="numpy", label="Upload Wildfire Image")
-        with gr.Column():
-            status = gr.Textbox(label="Fire Status", interactive=False)
+    with gr.Row(variant="panel"):
+        with gr.Column(scale=2):
+            inp = gr.Image(type="numpy", label="Satellite Image", elem_id="upload-wildfire-image")
+        with gr.Column(scale=1):
+            status = gr.Textbox(label="Fire Status", interactive=False, elem_classes="status-box")
             severity = gr.Textbox(label="Severity Level", interactive=False)
             trend = gr.Textbox(label="Weather Trend", interactive=False)
     
     with gr.Accordion("📋 Detailed Recommendations", open=False):
-        rec_box = gr.Markdown(label="Recommendations")
+        rec_box = gr.Markdown()
     
-    btn = gr.Button("Analyze")
-    btn.click(fn=pipeline, inputs=inp, outputs=[status, severity, trend, rec_box])
+    btn = gr.Button("Analyze", variant="primary")
+    btn.click(
+        fn=pipeline,
+        inputs=inp,
+        outputs=[status, severity, trend, rec_box],
+        api_name="analyze"
+    )
     
-    gr.HTML("<p style='font-size:0.8em; color:#666;'>© 2025 ForestAI Labs</p>")
+    gr.Markdown("---")
+    gr.HTML("<div style='text-align: center; color: #666;'>© 2025 ForestAI Labs</div>")
+
+def handle_errors(inputs, outputs):
+    for output in outputs:
+        if isinstance(output, Exception):
+            raise gr.Error("Analysis failed. Please check the input and try again.")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()  # Fixed the space typo in launch()