Spaces:
Build error
Build error
File size: 1,671 Bytes
5a3267d | 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 | import torch
from pathlib import Path
from PIL import Image
import streamlit as st
from torchvision import transforms
# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s', pretrained=True)
# Define image transformation
transform = transforms.Compose([
transforms.ToTensor(),
])
def perform_inference(image):
# Apply transformation
input_image = transform(image).unsqueeze(0)
# Perform inference
with torch.no_grad():
results = model(input_image)
return results
def display_results(image, results):
# Display the image with bounding boxes
st.image(image, caption="Input Image", use_column_width=True)
# Access the bounding box coordinates and class labels
boxes = results.xyxy[0].cpu().numpy()[:, :-1]
class_labels = results.xyxy[0].cpu().numpy()[:, -1]
# Display bounding boxes on the image
for box, label in zip(boxes, class_labels):
st.rectangle(
xy=(box[0], box[1]),
width=box[2] - box[0],
height=box[3] - box[1],
color='red',
label=f'Class {int(label)}'
)
# Streamlit app
def main():
st.title("YOLOv5 Object Detection with Streamlit")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Load the image
image = Image.open(uploaded_file).convert('RGB')
# Perform inference
results = perform_inference(image)
# Display results
display_results(image, results)
# Save the result image
results.save(Path('output'))
if __name__ == "__main__":
main()
|