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()