Spaces:
Sleeping
Sleeping
import os | |
import cv2 | |
import streamlit as st | |
from ultralytics import YOLO | |
from streamlit_image_comparison import image_comparison | |
from utils import save_uploadedfile, apply_masks | |
model = YOLO('models/last.pt') | |
st.title("Image Segmentation with YOLOv8: A Web Integration") | |
st.subheader("Implementing for Image Segmentation and Object Detection") | |
st.write("Image segmentation is a critical task in computer vision that involves dividing an image into multiple " | |
"segments or regions. YOLOv8 is a state-of-the-art deep learning model that can be used for " | |
"image segmentation and object detection. In this web app, we will implement an interesting example " | |
"using YOLOv8 for image segmentation. We can simply drop an image, View the identified segment " | |
"a piece, a whole and the distribution of the identified segments. The essence of this application " | |
"is to build a practical understanding and implementation of the powerful and light YOLOv8 " | |
"for image segmentation and object detection") | |
url = "https://github.com/AkanimohOD19A/img-segmentation" | |
link = f'<a href="{url}">This sample app was heavily based on code shared by Akan Daniel. Huge credits to him.</a>' | |
st.markdown(link, unsafe_allow_html=True) | |
st.divider() | |
st.markdown('') | |
st.markdown('##### Segmented Pieces') | |
## Placeholder Image | |
parent_media_path = "media-directory" | |
img_file = './media-directory/example_input.png' | |
## Application States | |
APPLICATION_MODE = st.sidebar.selectbox("Our Options", | |
["Take Picture", "Upload Picture"] | |
) | |
## captured_picture Image | |
if APPLICATION_MODE == "Take Picture": | |
st.sidebar.write( | |
""" | |
A computer aided application that segments your input image, built on | |
the powerful YOLOv8 instance segmentation algorithm developed by *ultralytics*. | |
Simply take a captured_picture and it gets segmentated in real time. | |
""" | |
) | |
picture = st.camera_input("Take a picture") | |
st.markdown('') | |
if picture: | |
st.sidebar.divider() | |
st.sidebar.image(picture, caption="captured_picture") | |
if st.button("Segment!"): | |
img_file = os.path.join(parent_media_path, "captured_picture.jpg") | |
save_uploadedfile(picture, save_path=img_file) | |
st.sidebar.success("Saved File") | |
st.write("Click on **Clear photo** to retake picture") | |
st.divider() | |
elif APPLICATION_MODE == "Upload Picture": | |
st.sidebar.write( | |
""" | |
A computer aided application that segments your input image, built on | |
the powerful YOLOv8 object detection algorithm developed by *ultralytics*. | |
Simply drop your image and it gets segmentated in real time. | |
""" | |
) | |
st.sidebar.divider() | |
uploaded_file = st.sidebar.file_uploader("Drop a JPG/PNG file", accept_multiple_files=False, type=['jpg', 'png']) | |
if uploaded_file is not None: | |
img_file = os.path.join(parent_media_path, "uploaded_image.jpg") | |
save_uploadedfile(uploaded_file, save_path=os.path.join(parent_media_path, "uploaded_image.jpg")) | |
file_details = {"FileName": uploaded_file.name, "FileType": uploaded_file.type} | |
st.sidebar.success("File saved successfully") | |
print(f"File saved successfully to {os.path.abspath(img_file)}") | |
else: | |
st.sidebar.write("You are using a placeholder image, Upload your Image (.jpg for now) to explore") | |
results = model(img_file) | |
img = cv2.imread(img_file) | |
for result in results: | |
# segmentation | |
masks = [] if not result.masks else result.masks.data.cpu().numpy() | |
numCols = len(masks) | |
if numCols > 0: | |
cols = st.columns(numCols) | |
print(f"Number of instances found: {numCols}") | |
else: | |
st.warning("Unable to id Distinct items - Please retry with a clearer Image") | |
img = apply_masks(img, masks) | |
st.markdown('') | |
st.markdown('##### Slider of Uploaded Image and Segments') | |
image_comparison( | |
img1=img_file, | |
img2=img, | |
label1="Actual Image", | |
label2="Segmented Image", | |
width=700, | |
starting_position=50, | |
show_labels=True, | |
make_responsive=True, | |
in_memory=True | |
) | |
st.sidebar.divider() | |
st.sidebar.markdown('') | |
st.sidebar.markdown('#### Distribution of identified items') | |
st.markdown('') | |
st.markdown('') | |
st.markdown('') | |
st.markdown('') | |
st.markdown('') | |
st.markdown('') | |
st.sidebar.divider() | |