ake178178/objectDetect

1import streamlit as st2from PIL import Image3import torchvision.transforms as T4from torchvision.models.detection import fasterrcnn_resnet50_fpn5import torch6 7# 载入一个预训练的 Faster R-CNN 模型8model = fasterrcnn_resnet50_fpn(pretrained=True)9model.eval()10 11# 设置图片转换12transform = T.Compose([13    T.ToTensor(), 14])15 16def detect_objects(image):17    # 转换图片并添加批次维度18    img_tensor = transform(image).unsqueeze(0)19    with torch.no_grad():20        predictions = model(img_tensor)21    22    # 返回预测结果23    return predictions[0]24 25def main():26    st.title("物体识别与距离估计")27    file_uploader = st.file_uploader("上传图片", type=["png", "jpg", "jpeg"])28    29    if file_uploader is not None:30        image = Image.open(file_uploader)31        st.image(image, caption="上传的图片", use_column_width=True)32        33        # 运行物体识别34        predictions = detect_objects(image)35        36        # 显示结果37        for i, (box, score, label) in enumerate(zip(predictions['boxes'], predictions['scores'], predictions['labels'])):38            if score > 0.5:  # 筛选置信度大于 0.5 的预测结果39                st.write(f"物体 {i + 1}: 类别 {label}, 置信度 {score:.2f}")40                # 简单的距离估计：基于物体的大小41                area = (box[2] - box[0]) * (box[3] - box[1])42                distance = 2000 / area.sqrt()  # 假设计算，不是真实世界的精确测量43                st.write(f"估计距离: {distance:.2f} 米")44 45if __name__ == "__main__":46    main()47