Spaces:
Runtime error
Runtime error
File size: 3,708 Bytes
7bd5908 8745d6b eb9abdb fcb3849 7bd5908 fcb3849 ad50413 fcb3849 0a0334e 572cf00 fcb3849 e44ae6b 572cf00 e44ae6b 572cf00 e44ae6b fcb3849 e44ae6b fcb3849 e44ae6b eb9abdb e44ae6b fcb3849 572cf00 fcb3849 f4d65e2 0a0334e d73a810 0a0334e eb9abdb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import os
import io
from PIL import Image,ImageDraw
from transformers import AutoImageProcessor, AutoModelForObjectDetection
import streamlit as st
import torch
import requests
#def prettier(results):
# for item in results:
# score = round(item['score'], 3)
# label = item['label'] # Use square brackets to access the 'label' key
# location = [round(value, 2) for value in item['box'].values()]
# print(f'Detected {label} with confidence {score} at location {location}')
#def prettify_results(results):
# for item in results:
# score = round(item['score'].item(), 3)
# label = model.config.id2label[item['label']] # Get label from id2label mapping in model config
# box = [round(coord, 2) for coord in item['box']]
# st.write(f'Detected {label} with confidence {score} at location {box}')
# Function to process uploaded image and prepare input for model
def input_image_setup(uploaded_file):
if uploaded_file is not None:
bytes_data = uploaded_file.getvalue()
image = Image.open(io.BytesIO(bytes_data)) # Convert bytes data to PIL image
return image
else:
raise FileNotFoundError("No file uploaded")
#Streamlit App
st.set_page_config(page_title="Image Detection")
st.header("Object Detection Application")
#Select your model
models = ["facebook/detr-resnet-50", "ciasimbaya/ObjectDetection", "hustvl/yolos-tiny"] # List of supported models
model_name = st.selectbox("Select model", models)
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForObjectDetection.from_pretrained(model_name)
#Upload an image
uploaded_file = st.file_uploader("choose an image...", type=["jpg","jpeg","png"])
image=""
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image.", use_column_width=True)
submit = st.button("Detect Objects ")
"""if submit:
image_data=input_image_setup(uploaded_file)
st.subheader("The response is..")
#process with model
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
# model predicts bounding boxes and corresponding COCO classes
logits = outputs.logits
bboxes = outputs.pred_boxes
# print results
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
# prettify_results(results)
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
box = [round(i, 2) for i in box.tolist()]
st.write(
f"Detected {model.config.id2label[label.item()]} with confidence "
f"{round(score.item(), 3)} at location {box}"
)
"""
if submit:
image_data = input_image_setup(uploaded_file)
st.subheader("The response is..")
inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
bboxes = outputs.pred_boxes
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
# Draw bounding boxes on the image
drawn_image = image.copy()
draw = ImageDraw.Draw(drawn_image)
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
box = [int(i) for i in box.tolist()]
draw.rectangle(box, outline="red", width=2)
label_text = f"{model.config.id2label[label.item()]} ({round(score.item(), 2)})"
draw.text((box[0], box[1]), label_text, fill="red")
st.image(drawn_image, caption="Detected Objects", use_column_width=True) |