Spaces:
Sleeping
Sleeping
File size: 2,708 Bytes
d68f082 d17f16b d68f082 b4c66ae d17f16b d68f082 d17f16b d68f082 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import streamlit as st
from PIL import Image
import torch
from torchvision import transforms
import torch.nn.functional as F
# Load the trained model
MODEL_PATH = "resnet_model.pth" # Update with your actual model path
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load(MODEL_PATH, map_location=device)
model.eval()
# Define the image transformation pipeline
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])
])
# Streamlit UI
st.title("Saliva Disease Detection App")
st.subheader("Predict Streptococcal infection vs No Streptococcal infection from saliva images")
# Initialize session state for managing the uploaded file
if "uploaded_file" not in st.session_state:
st.session_state["uploaded_file"] = None
# File uploader
uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"], key="file_uploader")
if uploaded_file is not None:
st.session_state["uploaded_file"] = uploaded_file
# If a file has been uploaded, process and predict
if st.session_state["uploaded_file"] is not None:
image = Image.open(st.session_state["uploaded_file"])
st.image(image, caption="Uploaded Image", use_container_width=True)
# Preprocess the image
input_image = transform(image).unsqueeze(0).to(device)
# Perform prediction
with torch.no_grad():
outputs = model(input_image)
probabilities = F.softmax(outputs, dim=1) # Convert to probabilities
_, predicted_class = torch.max(outputs, 1)
# Map predicted class to labels
#class_names = ['Not_Streptococcosis', 'Streptococcosis']
class_names = ['Not_Streptococcal_Infection', 'Streptococcal_Infection']
predicted_label = class_names[predicted_class.item()]
predicted_probability = probabilities[0][predicted_class.item()].item() * 100 # Convert to percentage
# Display the result
st.write("### Prediction Result:")
if predicted_label == "Streptococcal_Infection":
st.error(f"The sample is predicted as **{predicted_label}** with **{predicted_probability:.2f}%** probability.")
else:
st.success(f"The sample is predicted as **{predicted_label}** with **{predicted_probability:.2f}%** probability.")
# Show probabilities for all classes
st.write("### Class Probabilities:")
for idx, class_name in enumerate(class_names):
st.write(f"- **{class_name}**: {probabilities[0][idx].item() * 100:.2f}%")
# Button to reset the file uploader
if st.button("Upload Another Image"):
st.session_state["uploaded_file"] = None
st.rerun()
|