import streamlit as st
from PIL import Image
import torchvision.transforms as T
from torchvision.models.detection import fasterrcnn_resnet50_fpn
import torch

# 载入一个预训练的 Faster R-CNN 模型
model = fasterrcnn_resnet50_fpn(pretrained=True)
model.eval()

# 设置图片转换
transform = T.Compose([
    T.ToTensor(), 
])

def detect_objects(image):
    # 转换图片并添加批次维度
    img_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        predictions = model(img_tensor)
    
    # 返回预测结果
    return predictions[0]

def main():
    st.title("物体识别与距离估计")
    file_uploader = st.file_uploader("上传图片", type=["png", "jpg", "jpeg"])
    
    if file_uploader is not None:
        image = Image.open(file_uploader)
        st.image(image, caption="上传的图片", use_column_width=True)
        
        # 运行物体识别
        predictions = detect_objects(image)
        
        # 显示结果
        for i, (box, score, label) in enumerate(zip(predictions['boxes'], predictions['scores'], predictions['labels'])):
            if score > 0.5:  # 筛选置信度大于 0.5 的预测结果
                st.write(f"物体 {i + 1}: 类别 {label}, 置信度 {score:.2f}")
                # 简单的距离估计：基于物体的大小
                area = (box[2] - box[0]) * (box[3] - box[1])
                distance = 2000 / area.sqrt()  # 假设计算，不是真实世界的精确测量
                st.write(f"估计距离: {distance:.2f} 米")

if __name__ == "__main__":
    main()