import os os.environ["STREAMLIT_HOME"] = os.getcwd() os.environ["STREAMLIT_RUNTIME_METRICS_ENABLED"] = "false" os.makedirs(".streamlit", exist_ok=True) import streamlit as st import torch import joblib import numpy as np import random from PIL import Image from transformers import AutoTokenizer, AutoModel, ViTModel, ViTImageProcessor # CPU device only device = torch.device("cpu") # Define Swahili VQA Model class SwahiliVQAModel(torch.nn.Module): def __init__(self, num_answers): super().__init__() self.vision_encoder = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k') self.text_encoder = AutoModel.from_pretrained("benjamin/roberta-base-wechsel-swahili") self.fusion = torch.nn.Sequential( torch.nn.Linear(768 + 768, 512), torch.nn.ReLU(), torch.nn.Dropout(0.3), torch.nn.LayerNorm(512) ) self.classifier = torch.nn.Linear(512, num_answers) def forward(self, image, input_ids, attention_mask): vision_outputs = self.vision_encoder(pixel_values=image) image_feats = vision_outputs.last_hidden_state[:, 0, :] text_outputs = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask) text_feats = text_outputs.last_hidden_state[:, 0, :] combined = torch.cat([image_feats, text_feats], dim=1) fused = self.fusion(combined) return self.classifier(fused) # Load label encoder le = joblib.load("Vit_3895_label_encoder_best.pkl") # Load model weights normally โ€” no override model = SwahiliVQAModel(num_answers=len(le.classes_)).to(device) state_dict = torch.load("Vit_3895_best_model_epoch25.pth", map_location=device) model.load_state_dict(state_dict) model.eval() # Load tokenizer and processor tokenizer = AutoTokenizer.from_pretrained("benjamin/roberta-base-wechsel-swahili") vit_processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k') # Streamlit UI st.set_page_config(page_title="Swahili VQA", layout="wide") st.title("๐Ÿฆœ Swahili Visual Question Answering (VQA)") uploaded_image = st.file_uploader("๐Ÿ“‚ Pakia picha hapa:", type=["jpg", "jpeg", "png"]) def generate_random_color(): return f"rgb({random.randint(150, 255)}, {random.randint(80, 200)}, {random.randint(80, 200)})" col1, col2 = st.columns([1, 2], gap="large") with col1: if uploaded_image: st.image(uploaded_image, caption="Picha Iliyopakiwa", use_container_width=True) st.markdown("
", unsafe_allow_html=True) with col2: st.markdown("
", unsafe_allow_html=True) question = st.text_input("๐Ÿ’ฌAndika swali lako hapa:", key="question_input") submit_button = st.button("๐Ÿ“ฉTuma") st.markdown("
", unsafe_allow_html=True) if submit_button and uploaded_image and question: with st.spinner("๐Ÿ” Inachakata jibu..."): image = Image.open(uploaded_image).convert("RGB") image_tensor = vit_processor(images=image, return_tensors="pt")["pixel_values"] inputs = tokenizer(question, max_length=128, padding="max_length", truncation=True, return_tensors="pt") input_ids = inputs["input_ids"] attention_mask = inputs["attention_mask"] with torch.no_grad(): logits = model(image_tensor, input_ids, attention_mask) probs = torch.softmax(logits, dim=1) top_probs, top_indices = torch.topk(probs, 5) decoded_answers = le.inverse_transform(top_indices.cpu().numpy()[0]) results = [ {"answer": ans, "confidence": round(prob * 100, 2)} for ans, prob in zip(decoded_answers, top_probs[0].tolist()) ] results = sorted(results, key=lambda x: x["confidence"], reverse=True) st.subheader("Majibu Yanayowezekana:") max_confidence = max(result["confidence"] for result in results) for i, pred in enumerate(results): bar_width = (pred["confidence"] / max_confidence) * 70 color = generate_random_color() st.markdown( f"""
{pred['answer']}
{pred['confidence']}%
""", unsafe_allow_html=True ) else: st.info("๐Ÿ“ฅ Pakia picha na andika swali kisha bonyeza Tuma ili kupata jibu.")