Make Easy UI for YOLO Models using Streamlit

SUMITH
4 min readJun 15, 2024

--

YOLO integration with Streamlit

Streamlit: Simplifying UI Design for YOLO Models

In today’s digital world, effective user interfaces (UIs) are critical for enhancing user experiences, especially in applications involving complex tasks like object detection. This blog post will walk you through a Streamlit application designed for YOLO (You Only Look Once) model-based object detection. This app allows users to upload images, videos, or URLs for analysis, leveraging the powerful YOLOv8 model for detecting objects within the uploaded media.

The application has a simple, intuitive design with a sidebar for navigation and a main area for content display. The sidebar offers three options: IMAGE, VIDEO, and URL LINKS, each corresponding to a different type of media upload and detection.

Designed UI

Overview of the UI

The UI consists of three main components:

  1. Sidebar Navigation: Allows users to switch between different functionalities like image, video, and URL link uploads.
  2. Upload Section: Enables users to upload files for object detection.
  3. Detection Button: Initiates the object detection process on the uploaded file running Yolov8 as the Detection model.

Here’s a breakdown of the UI elements:

  • Sidebar Menu: Contains navigation options (IMAGE, VIDEO, URL LINKS) and a button to trigger the detection process.
  • Main Content Area: Displays the uploaded media and the detection results.

Sidebar

The sidebar uses the option_menu component from streamlit_option_menu to provide a visually appealing and functional navigation menu.

Main Content Area

The main area changes dynamically based on the selected option from the sidebar. For the IMAGE option, users can upload an image, and upon clicking the “Detect” button, the YOLO model processes the image and displays the detected objects.

Install necessary libraries

  1. Streamlit
  2. ultralytics
  3. streamlit_option_menu

Python Script for running inference on the images

    model = YOLO("yolov8n.pt")
res = model.predict(file,conf=0.5,save=True)
box = res[0].boxes.xyxy.tolist()
res_plotted = res[0].plot()[:, :, ::-1]
st.image(res_plotted, caption='Text Detections',use_column_width=True)
st.write("Number of the Detections : "+str(len(box)))
Inferenced Image

Python Script to run YOLO model Inference on the Videos

In most of the online blogs or documentation you find that video is directly passed as the source to the yolo model for the inference, but I personally believe that processing the video frame by frame and plotting Bounding Boxes and Annotations is the best app.

Frame-by-Frame Processing

Advantages

Flexibility:

  • Easier to handle custom processing and analysis on individual frames.
  • Allows for more fine-grained control over the detection process.

Parallel Processing:

  • Frames can be processed in parallel, potentially speeding up the analysis if you have the necessary computational resources.

Error Handling:

  • Easier to handle errors or skips in processing. If a frame fails, you can move to the next without losing the entire video.

Debugging:

  • Simplifies debugging by allowing you to inspect and process frames individually.

Python Script

model = YOLO("yolov8n.pt")

cap = cv2.VideoCapture(inputfile)

# Retrieve video properties: width, height, and frames per second
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))

# Initialize video writer to save the output video with the specified properties
out = cv2.VideoWriter(outputfile, cv2.VideoWriter_fourcc(*"MJPG"), fps, (w, h))

while True:
# Read a frame from the video
ret, im0 = cap.read()
if not ret:
print("Video frame is empty or video processing has been successfully completed.")
break
res = model.predict(im0,conf=0.5,save=True)
res_plotted = res[0].plot()[:, :, ::-1]
out.write(res_plotted)
st.video(data=outputfile)
Video Output

To access the complete code for this project, visit the GitHub repository. https://github.com/sumith5/Yolo-Streamlit.git Feel free to clone the repository, customize it to fit your needs, and start building your own YOLO-based applications with ease.

Conclusion

Using Streamlit to build a UI for YOLO models simplifies the process of creating and deploying powerful object detection applications. Whether you’re working with images or videos, Streamlit’s intuitive interface and rich feature set allow you to focus on the core functionality of your project without getting bogged down by UI development.

--

--

SUMITH
0 Followers

Passionate AI and Computer Vision Engineer