Create app.py

app.py

@@ -0,0 +1,46 @@
+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()