app.py

""" Streamlit UI for object detection with DETR. """

# Use a pipeline as a high-level helper
from transformers import pipeline
import streamlit as st
from PIL import Image
import pandas as pd

pipe = pipeline("object-detection", model="facebook/detr-resnet-101")

# Set the title
st.title("Vision Quest 2")

results = None
image = None

# Create a file uploader and set the upload type to images
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file:
    upload_image_button = st.button("Upload Image")
    if upload_image_button:
        with st.spinner("Uploading Image...")
            # Convert the image to a file object
            image = Image.open(uploaded_file)
            
            # Process the image through the pipeline
            results = pipe(image)

col1, col2 = st.columns(2)
if image and results:
    with col1:
        st.image(image, use_column_width=True)
    with col2:
        # Display the individual objects, the bounding boxes, and the confidence
        # And then display the total number of each type of object
        # Create a dataframe to hold the results
        df = pd.DataFrame(results)
        st.dataframe(df)