import streamlit as st
import cv2
import numpy as np
import base64
from io import BytesIO
from PIL import Image

def process_image(image):
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
    _, thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
    return thresh

def detect_obstacle(thresh_img):
    ground_part = thresh_img[50:, :]
    air_part = thresh_img[:50, :]
    
    ground_nonzero = np.count_nonzero(ground_part)
    air_nonzero = np.count_nonzero(air_part)
    
    if ground_nonzero > 0:
        return "ground"
    elif air_nonzero > 0:
        return "air"
    return None

st.title('障碍物检测应用')

uploaded_file = st.file_uploader("选择一张图片", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption='上传的图片', use_column_width=True)
    
    if st.button('检测障碍物'):
        processed_img = process_image(image)
        obstacle = detect_obstacle(processed_img)
        
        if obstacle:
            st.success(f'检测到的障碍物类型: {obstacle}')
        else:
            st.info('未检测到障碍物')
        
        st.image(processed_img, caption='处理后的图片', use_column_width=True)