Brain_Tumor_Yolo / upload_view.py
davron04's picture
added requirements file
c9a2598
raw
history blame
2.75 kB
import streamlit as st
import numpy as np
import cv2 as cv
import supervision as sv
from PIL import Image
class Upload_View:
def __init__(self, app, model):
self.app = app
self.model = model
def show(self):
col1_back, col2_back = st.columns([0.2, 0.8])
with col1_back:
if st.button("Back", key='upload_back', icon = ':material/arrow_back:', type = 'primary'):
self.app.change_page("Main")
st.markdown("<h1 style='text-align: center;'>🧠 Upload MRI or CT scan image</h1>", unsafe_allow_html=True)
st.divider()
col1_img, col2_img = st.columns(2,gap='medium')
with col1_img:
st.header('Input')
uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
if uploaded_file is not None:
st.divider()
st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
with col2_img:
st.header('Output')
if uploaded_file is not None:
# Convert uploaded image to OpenCV format
image = Image.open(uploaded_file).convert("RGB")
image_np = np.array(image) # Convert to NumPy array (H, W, C)
image_bgr = cv.cvtColor(image_np, cv.COLOR_RGB2BGR) # Convert RGB to BGR for OpenCV
# Run YOLO inference
with st.spinner("Making the predictions... ⏳"):
results = self.model(image_bgr)[0]
# Convert YOLO results into Supervision Detections format
detections = sv.Detections.from_ultralytics(results)
# Generate labels with confidence scores
labels = [
f"{results.names[class_id]}: {confidence*100:.2f}%"
for class_id, confidence in zip(detections.class_id, detections.confidence)
]
# Annotate image with bounding boxes & labels
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(image_np, detections)
annotated_image = label_annotator.annotate(annotated_image, detections, labels=labels)
# Add tumor count and confidence description
num_tumors = len(detections) # Total number of detected tumors
# Display the annotated image
st.image(annotated_image, caption="Tumor Detection Result", use_container_width=True)
# Display tumor count and confidence information using Streamlit text element
st.subheader(f"Number of Tumors Detected: {num_tumors}")