ake178178's picture
Update app.py
caa944e verified
raw
history blame
1.35 kB
# 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()