# app.py

import streamlit as st
import torch
from PIL import Image
import camera  # 引入拍照功能
import numpy as np
import cv2

# 物体识别函数
def detect_objects(image_path):
    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # 使用YOLOv5模型
    img = Image.open(image_path)
    results = model(img)
    return results

def main():
    st.title("摄像头拍照并进行物体识别")

    # 拍照
    if st.button('拍照'):
        camera.take_picture()
        st.write("照片已拍摄并保存")

    # 物体识别
    if st.button('物体识别'):
        st.write("正在进行物体识别...")
        image_path = "captured_image.jpg"
        
        results = detect_objects(image_path)
        # 显示原始图片
        st.image(image_path, caption='原始图片', use_column_width=True)
        
        # 显示识别的结果
        results.render()  # 在图片上绘制检测到的物体
        detected_img = Image.fromarray(results.imgs[0])

        st.image(detected_img, caption='物体识别结果', use_column_width=True)

        # 显示从左到右的物体列表
        st.write("识别到的物体：")
        objects = results.pandas().xyxy[0]['name'].tolist()
        objects_sorted_by_x = sorted(objects)
        st.write(objects_sorted_by_x)

if __name__ == '__main__':
    main()