app.py

import streamlit as st
import tensorflow as tf
from PIL import Image
import numpy as np
import gdown

# Google Drive link to the model file
google_drive_link = "https://drive.google.com/drive/folders/1H2Yh06nKlWqWqbc9loCpbebMUFMvvS6T"

# Download the model file
gdown.download(google_drive_link, "efficientnetb4_model.h5", quiet=False)

# Load the model
model_path = "efficientnetb4_model.h5"
model = tf.keras.models.load_model(model_path)

# Streamlit app
st.title("Malware Image Prediction App")

# Upload image through Streamlit
uploaded_file = st.file_uploader("Choose an image...", type="jpg")

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image.", use_column_width=True)

    # Preprocess the image for prediction
    img_array = tf.keras.preprocessing.image.img_to_array(image)
    img_array = tf.image.resize(img_array, (75, 75))
    img_array = tf.expand_dims(img_array, 0)  # Create batch axis

    # Make prediction
    predictions = model.predict(img_array)
    predicted_class = tf.argmax(predictions[0]).numpy()

    # Display the prediction
    st.write("Prediction:")
    st.write(f"Class: {predicted_class}")

    # Display the probability distribution
    st.write("Probability Distribution:")
    st.bar_chart(predictions[0])