Benjy's picture
Update app.py
7dbe229 verified
import streamlit as st
import cv2
from ultralytics import YOLO
from PIL import Image
import numpy as np
# Load the YOLOv8 model
model = YOLO('yolov8n.pt')
# Streamlit app
def main():
st.title("Object Detection - General Use")
st.write("This is a general use object detection space using YOLOv8. For more complex projects, video, or real-time object detection, can be implemented.")
st.header("Upload an Image")
# Upload image
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Read the uploaded image
image = Image.open(uploaded_file)
image_array = np.array(image)
# Perform object detection
results = model.predict(image_array)
# Display the detected objects
for result in results:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
cv2.rectangle(image_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
conf = box.conf[0]
cls = int(box.cls[0])
label = f"{model.names[cls]} ({conf:.2f})"
cv2.putText(image_array, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Convert the image array back to PIL Image
image = Image.fromarray(image_array)
# Display the image with detected objects
st.image(image, caption='Detected Objects')
if not uploaded_file:
# Display example image
st.header("Example Result")
example_image = Image.open("example.jpeg")
st.image(example_image, caption='Cars are picked out in green and the person is distinguished from the motorcycle', use_column_width=True)
if __name__ == "__main__":
main()