try-streamlit / app.py
ake178178's picture
Update app.py
2762962 verified
import streamlit as st
from PIL import Image
import torch
# 载入一个预训练的物体识别模型,这里使用 YOLOv5
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
def detect_objects(image):
# 对图像进行处理并通过模型进行预测
results = model(image)
# 解析结果,获取检测到的物体和置信度
data = results.pandas().xyxy[0]
return results, data
def main():
st.title("物体识别应用")
# 创建一个文件上传器,用户可以上传图片
uploaded_file = st.file_uploader("请选择一张图片进行物体识别", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# 将上传的文件转换为图像
image = Image.open(uploaded_file)
# 显示原始图片
st.image(image, caption='上传的图片', use_column_width=True)
# 检测图片中的物体
results, data = detect_objects(image)
# 显示结果
st.write("检测结果:")
st.write(data)
# 显示带有检测框的图片
st.image(results.render()[0], caption='检测结果', use_column_width=True)
if __name__ == "__main__":
main()