|
import streamlit as st |
|
from PIL import Image |
|
from ultralytics import YOLO |
|
|
|
|
|
model = YOLO('best.pt') |
|
|
|
|
|
st.title("Book Detection with YOLOv8") |
|
st.markdown(""" |
|
**Created by Siddharth Basale** |
|
|
|
Upload an image containing books, and this app will detect the number of books and generate an image with the bounding boxes. |
|
""") |
|
|
|
|
|
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) |
|
|
|
if uploaded_file is not None: |
|
|
|
image = Image.open(uploaded_file) |
|
st.image(image, caption="Uploaded Image", use_column_width=True) |
|
|
|
|
|
results = model(image) |
|
|
|
|
|
book_class_id = None |
|
for class_id, class_name in model.names.items(): |
|
if class_name == 'book': |
|
book_class_id = class_id |
|
break |
|
|
|
|
|
if book_class_id is not None: |
|
books_detected = len([r for r in results[0].boxes.data if int(r[-1]) == book_class_id]) |
|
else: |
|
books_detected = 0 |
|
|
|
st.write(f"Number of books detected: {books_detected}") |
|
|
|
|
|
annotated_image = results[0].plot() |
|
st.image(annotated_image, caption="Processed Image with Book Detection", use_column_width=True) |
|
|