Janaalharbii's picture
Update app.py
73f6d3a verified
import streamlit as st
from PIL import Image
from ultralytics import YOLO
# Load the YOLO model
@st.cache_resource
def load_model():
model = YOLO('best (5).pt') # Use the path to your trained model
return model
# Prediction function
def predict(model, image):
results = model(image)
return results
# Load and display the header image
header_image = Image.open('historicalPIC.png')
st.image(header_image, use_column_width=True)
# Streamlit UI
st.title("Historical Places")
# File uploader for users to upload images
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
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)
# Convert the image to RGB format if needed
if image.mode != 'RGB':
image = image.convert('RGB')
# Load the model
model = load_model()
# Run YOLO prediction
results = predict(model, image)
# Access class names
class_names = model.names
# Process and display results
detected_objects = []
for result in results:
# Iterate over each detected object
for bbox in result.boxes:
x1, y1, x2, y2 = bbox.xyxy[0].tolist()
conf = bbox.conf.item()
cls = int(bbox.cls.item())
detected_objects.append(f"Detected {class_names[cls]}")
# Display detected objects' names and confidence scores
st.subheader("Detection Results")
if detected_objects:
for obj in detected_objects:
st.write(obj)
else:
st.write("No objects detected.")
# Render the image with bounding boxes
if results:
try:
results.render() # Modify the image with bounding boxes
img_with_boxes = Image.fromarray(results.imgs[0])
st.image(img_with_boxes, caption='Detected Objects', use_column_width=True)
except Exception:
pass # Ignore any errors in rendering