Update app.py

app.py

@@ -7,19 +7,16 @@ import torch.nn as nn
 import os
 import onnx
 import onnxruntime as ort
+import numpy as np
 
 # === Model Path ===
 MODEL_PATH = "bangla_disaster_model.pth"
+ONNX_PATH = "bangla_disaster_model.onnx"
 
-# Check if model exists
-if not os.path.exists(MODEL_PATH):
-    st.error("❌ Model file not found. Please ensure bangla_disaster_model.pth is uploaded.")
-    st.stop()
-
-# Global class list
+# === Class Labels ===
 classes = ['HYD', 'MET', 'FD', 'EQ', 'OTHD']
 
-# === Model Setup ===
+# === Model Architecture (used only for export) ===
 class MultimodalBanglaClassifier(nn.Module):
     def __init__(self, text_model_name='sagorsarker/bangla-bert-base', num_classes=5):
         super(MultimodalBanglaClassifier, self).__init__()
@@ -50,60 +47,51 @@ class MultimodalBanglaClassifier(nn.Module):
         fused = self.transformer_fusion(fused).squeeze(1)
         return self.classifier(fused)
 
-# === ONNX Export Helper ===
+# === ONNX Export ===
 def export_to_onnx_if_needed(model):
-    onnx_path = "bangla_disaster_model.onnx"
-    if os.path.exists(onnx_path):
+    if os.path.exists(ONNX_PATH):
         return
     dummy_input_ids = torch.randint(0, 30522, (1, 128), dtype=torch.long)
     dummy_attention_mask = torch.ones((1, 128), dtype=torch.long)
-    dummy_image = torch.randn(1, 3, 224, 224, dtype=torch.float)
+    dummy_image = torch.randn(1, 3, 224, 224)
     torch.onnx.export(
         model,
         (dummy_input_ids, dummy_attention_mask, dummy_image),
-        onnx_path,
+        ONNX_PATH,
         input_names=["input_ids", "attention_mask", "image"],
         output_names=["output"],
-        dynamic_axes={
-            "input_ids": {0: "batch"},
-            "attention_mask": {0: "batch"},
-            "image": {0: "batch"},
-            "output": {0: "batch"}
-        },
+        dynamic_axes={"input_ids": {0: "batch"}, "attention_mask": {0: "batch"}, "image": {0: "batch"}, "output": {0: "batch"}},
         opset_version=14,
         do_constant_folding=True
     )
 
+# === Load Model and Tokenizer for Exporting Only ===
 @st.cache_resource
 def load_model_and_tokenizer():
     model = MultimodalBanglaClassifier()
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
+    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
     model.eval()
     tokenizer = AutoTokenizer.from_pretrained("sagorsarker/bangla-bert-base")
     export_to_onnx_if_needed(model)
-    return model, tokenizer
+    return tokenizer
 
+# === Load ONNX Session ===
 @st.cache_resource
 def load_onnx_session():
-    return ort.InferenceSession("bangla_disaster_model.onnx")
+    return ort.InferenceSession(ONNX_PATH)
 
+# === ONNX Prediction ===
 def predict_with_onnx(session, tokenizer, image, caption):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225])
     ])
     image_tensor = transform(image).unsqueeze(0).numpy()
-
-    encoded = tokenizer(
-        caption,
-        padding='max_length',
-        truncation=True,
-        max_length=128,
-        return_tensors='pt'
-    )
-    input_ids = encoded['input_ids'].numpy()
-    attention_mask = encoded['attention_mask'].numpy()
+    encoded = tokenizer(caption, padding="max_length", truncation=True, max_length=128, return_tensors="pt")
+    input_ids = encoded["input_ids"].numpy()
+    attention_mask = encoded["attention_mask"].numpy()
 
     inputs = {
         "input_ids": input_ids,
@@ -112,45 +100,46 @@ def predict_with_onnx(session, tokenizer, image, caption):
     }
     outputs = session.run(None, inputs)
     logits = outputs[0]
-    pred_class = logits.argmax(axis=1).item()
-    confidence_scores = logits[0].tolist()
-    return classes[pred_class], confidence_scores
 
+    # Softmax to get probabilities
+    exp_logits = np.exp(logits - np.max(logits, axis=1, keepdims=True))
+    probs = exp_logits / np.sum(exp_logits, axis=1, keepdims=True)
+    pred_class = np.argmax(probs, axis=1)[0]
+    return classes[pred_class], probs[0].tolist()
+
+# === Bangla Labels ===
 def get_bangla_response(class_name):
     responses = {
         'HYD': "🌊 এটি একটি জলসম্পর্কিত দুর্যোগ (Hydrological Disaster)। সতর্ক থাকুন!",
         'MET': "🌪️ এটি একটি আবহাওয়া সংক্রান্ত দুর্যোগ (Meteorological Disaster)। সাবধানে থাকুন!",
         'FD': "🔥 আগুন লেগেছে! এটি একটি অগ্নিদুর্ঘটনা (Fire Disaster)। দ্রুত ব্যবস্থা নিন!",
-        'EQ': "🌍 ভুমিকম্প শনাক্ত হয়েছে (Earthquake)! নিরাপদ স্থানে যান!",
+        'EQ': "🌍 ভূমিকম্প শনাক্ত ���য়েছে (Earthquake)! নিরাপদ স্থানে যান!",
         'OTHD': "😌 এটা কোনো দুর্যোগ নয়। চিন্তার কিছু নেই!"
     }
-    return responses.get(class_name, "🤔 শ্রেণিবিন্যাস করা যায়নি!")
+    return responses.get(class_name, "🤔 শ্রেণিবিন্যাস করা যায়নি।")
 
 # === Streamlit UI ===
 st.set_page_config(page_title="Bangla Disaster Classifier", layout="centered")
 st.title("🌪️🇧🇩 Bangla Disaster Classifier")
 st.markdown("এই অ্যাপটি একটি multimodal deep learning মডেল ব্যবহার করে ছবির সাথে বাংলা ক্যাপশন বিশ্লেষণ করে দুর্যোগ শনাক্ত করে।")
 
-model, tokenizer = load_model_and_tokenizer()
+tokenizer = load_model_and_tokenizer()
 onnx_session = load_onnx_session()
 
 uploaded_file = st.file_uploader(
     "🖼️ একটি দুর্যোগের ছবি আপলোড করুন",
     type=['jpg', 'png', 'jpeg'],
-    accept_multiple_files=False,
     key="disaster_image_uploader",
     help="ছবি আপলোড করতে এখানে ক্লিক করুন অথবা drag & drop করুন"
 )
 
-if uploaded_file is not None:
+if uploaded_file:
     st.success(f"✅ ছবি আপলোড সফল: {uploaded_file.name}")
 else:
     st.info("📁 অনুগ্রহ করে একটি ছবি আপলোড করুন")
 
 caption = st.text_area("✍️ বাংলায় একটি ক্যাপশন লিখুন", "")
 
-st.caption("🎯 পূর্বাভাস মোড: উচ্চ নির্ভুলতা (High Accuracy)")
-
 col1, col2 = st.columns([1, 1])
 submit = col1.button("🔍 পূর্বাভাস দিন")
 clear = col2.button("🧹 রিসেট করুন")
@@ -163,16 +152,15 @@ if submit and uploaded_file and caption:
     st.image(img, caption="আপলোড করা ছবি", use_container_width=True)
 
     with st.spinner("🧠 মডেল পূর্বাভাস দিচ্ছে... (Model processing...)"):
-        progress_bar = st.progress(0, text="ছবি প্রক্রিয়াকরণ... (Processing image...)")
-        progress_bar.progress(50, text="বিশ্লেষণ চলছে... (Running inference...)")
-        
+        progress_bar = st.progress(0, text="প্রক্রিয়াকরণ শুরু হচ্ছে...")
+
+        progress_bar.progress(50, text="বিশ্লেষণ চলছে...")
         prediction, probs = predict_with_onnx(onnx_session, tokenizer, img, caption)
+        progress_bar.progress(100, text="✅ বিশ্লেষণ সম্পন্ন!")
 
-        progress_bar.progress(100, text="✅ সম্পূর্ণ! (Complete!)")
     progress_bar.empty()
 
     st.markdown(f"### ✅ পূর্বাভাস: {get_bangla_response(prediction)}")
-
     col1, col2 = st.columns([2, 1])
     with col1:
         st.markdown(f"#### 📊 সম্ভাব্যতা: **{probs[classes.index(prediction)]:.2%}**")
@@ -188,6 +176,6 @@ if submit and uploaded_file and caption:
             'OTHD': 'কোনো দুর্যোগ নয়'
         }
         for i, class_code in enumerate(classes):
-            percentage = probs[i] * 100
-            st.write(f"**{class_names[class_code]}**: {percentage:.1f}%")
-            st.progress(probs[i])
+            percentage = probs[i]
+            st.write(f"**{class_names[class_code]}**: {percentage:.1%}")
+            st.progress(min(max(percentage, 0.0), 1.0))  # ensure range [0.0, 1.0]