import streamlit as st
import numpy as np
from PIL import Image
import joblib
from flask import Flask, request, jsonify
from streamlit_webrtc import webrtc_streamer, VideoTransformerBase
import cv2
import threading

# 加载训练好的模型
model_file_path = 'skin_cancer_model.pkl'
model = joblib.load(model_file_path)

# 图像预处理函数
def preprocess_image(image):
    size = (224, 224)  # 模型输入尺寸
    image = image.resize(size)
    image_array = np.array(image) / 255.0
    image_flattened = image_array.flatten()
    return [image_flattened]

# 定义视频转换器类，处理视频帧
class VideoTransformer(VideoTransformerBase):
    def transform(self, frame):
        img = frame.to_ndarray(format="bgr24")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 预处理图像
        img_resized = cv2.resize(img, (224, 224))
        img_resized = img_resized.astype("float32") / 255.0
        img_resized = np.expand_dims(img_resized, axis=0)

        # 模型预测
        prediction = model.predict(img_resized)
        predicted_class = np.argmax(prediction, axis=1)[0]

        # 在 Streamlit 中显示预测结果
        st.write(f"Skin Cancer Prediction (Video): {predicted_class}")

        return img

# Flask 应用设置
app = Flask(__name__)

# 接收图像并通过模型分析
@app.route('/analyze', methods=['POST'])
def analyze():
    file = request.files['image']
    image = Image.open(file.stream)
    processed_image = preprocess_image(image)
    prediction = model.predict(processed_image)[0]
    class_names = {0: 'Benign', 1: 'Malignant'}
    result = class_names[prediction]
    return jsonify({"result": result})

# 运行 Flask 应用
def run_flask():
    app.run(host='0.0.0.0', port=8000)

# 主程序 - Streamlit UI 和交互逻辑
def main():
    st.title("Skin Diagnosis")

    st.write(
        "Please note that this app is not 100% accurate. If your result happens to be malignant, please contact a medical professional for further instructions.")

    # 实时视频诊断
    st.header("Real-Time Video Diagnosis")
    webrtc_streamer(key="example", video_processor_factory=VideoTransformer)

    # 图像上传与预测
    st.header("Image Upload Diagnosis")
    st.write("**Upload an image to detect if it's benign or malignant:**")

    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)

        # 预处理图像
        processed_image = preprocess_image(image)

        # 模型预测
        prediction = model.predict([processed_image])[0]
        class_names = {0: 'Benign', 1: 'Malignant'}
        result = class_names[prediction]

        st.write(f"Skin Cancer Prediction (Image): {result}")

    # 阳光暴露选项
    st.write("**Have you had a large amount of sun exposure recently?**")
    result = st.button("Yes")
    if result:
        st.write(
            "Please note that high levels of sun exposure can increase the risk of skin cancer. Be cautious and consider regular skin check-ups.")

# 使用线程并行运行 Flask 和 Streamlit
if __name__ == '__main__':
    # 启动 Flask 服务器线程
    flask_thread = threading.Thread(target=run_flask)
    flask_thread.daemon = True
    flask_thread.start()

    # 运行 Streamlit 前端
    main()