Spaces:
Sleeping
Sleeping
import streamlit as st | |
from PIL import Image | |
from ultralytics import YOLO | |
# Load the YOLO model | |
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 | |