| import streamlit as st |
| import cv2 |
| import numpy as np |
| import numpy |
| import torch |
| from torchvision import transforms |
| from PIL import Image |
|
|
| |
| model = torch.hub.load('ultralytics/yolov5', 'yolov5s') |
| classes = model.names |
|
|
| def object_detection(img_path): |
| """ |
| 对图像进行物体识别。 |
| """ |
| results = model(img_path) |
| return results.pandas().xyxy[0] |
|
|
| def show_results(results): |
| """ |
| 在 Streamlit 中显示物体识别结果。 |
| """ |
| for index, row in results.iterrows(): |
| left, top, right, bottom = row['xmin'], row['ymin'], row['xmax'], row['ymax'] |
| label = classes[int(row['class'])] |
| st.write(f"物体：{label}, 位置：({left}, {top}) - ({right}, {bottom})") |
|
|
| if __name__ == "__main__": |
| if st.button("选择图片"): |
| uploaded_file = st.file_uploader("Choose an image file", type=["png", "jpg", "jpeg"]) |
| if uploaded_file is not None: |
| img = Image.open(uploaded_file) |
| img_array = np.array(img) |
| results = object_detection(img_array) |
| show_results(results) |
