nicekate · September 21, 2024 06:13
diff --git a/app.py b/app.py
 import os
 import base64
 import streamlit as st
 from mistralai import Mistral
 from PIL import Image
 import io
 from dotenv import load_dotenv

 # 加载 .env 文件
 load_dotenv()

 # 初始化会话状态
 if 'base64_images' not in st.session_state:
    st.session_state.base64_images = []
 if 'result' not in st.session_state:
    st.session_state.result = ""
 if 'conversation' not in st.session_state:
    st.session_state.conversation = 1
 if 'prompt' not in st.session_state:
    st.session_state.prompt = ""

 def perform_analysis(prompt, uploaded_files):
    if not uploaded_files:
        st.error("请先上传图片。")
        return

    if len(uploaded_files) > 8:
        st.error("最多只能上传8张图片。")
        return

    if not prompt:
        st.error("请先选择或输入提示词。")
        return

    base64_images = []
    for uploaded_file in uploaded_files:
        # 修改编码方式，使用压缩后的图片
        img = Image.open(uploaded_file)
        if img.mode in ("RGBA", "P"):
            img = img.convert("RGB")  # 转换为 RGB 模式
        img.thumbnail((300, 300))
        buffer = io.BytesIO()
        img.save(buffer, format="JPEG")
        base64_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
        base64_images.append(base64_image)
        
        # 显示图片
        st.image(img, caption='上传的图片')
    
    st.session_state.base64_images = base64_images
    
    api_key = os.getenv("MISTRAL_API_KEY")
    if not api_key:
        st.error("错误：未找到MISTRAL_API_KEY。请确保在 .env 文件中设置了该变量。")
        return
    
    client = Mistral(api_key=api_key)
    
    content = [
        {"type": "text", "text": prompt}
    ]
    for img_str in base64_images:
        content.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{img_str}"})
    
    messages = [
        {
            "role": "user",
            "content": content
        }
    ]
    
    try:
        chat_response = client.chat.complete(
            model="pixtral-12b-2409",
            messages=messages
        )
        result = chat_response.choices[0].message.content
    except Exception as e:
        result = f"分析过程中出错：{str(e)}"
    
    st.session_state.result = result

 def upload_and_analyze():
    # 使用 Sidebar 代替左侧栏
    with st.sidebar:
        st.header("设置")
        
        # 新增"新对话"按钮
        if st.button("新对话"):
            st.session_state.base64_images = []
            st.session_state.result = ""
            st.session_state.conversation += 1  # 增加会话计数
            st.session_state.prompt = ""
            st.rerun()  # 使用 st.rerun() 替代 st.experimental_rerun()

        uploaded_files = st.file_uploader(
            "上传图片",
            type=["jpg", "jpeg", "png"],
            accept_multiple_files=True,
            key=f"file_uploader_{st.session_state.conversation}"
        )
        
        # 显示预设提示词按钮，每行三个
        st.markdown("### 预设提示词")
        prompts = {
            "图像描述": "请描述这张图片的内容。",
            "图像推理": "请推理这张图片中发生了什么。",
            "解释图片": "请解释这张图片中的内容。",
            "信息提取": "请从这张图片中提取关键信息。",
            "内容转录（OCR）": "请转录这张图片中的文本内容。",
            "有趣故事": "请根据这些图片创作一个有趣的故事。",
            "数据分析": "请分析这些图表或数据可视化图像，并提供关键洞察。",
            "趋势识别": "请识别并解释这些图表中的主要趋势。",
            "数据比较": "请比较这些图表中的数据，并指出主要差异。",
            "Extract Content and Generate Table": "Please extract content from this image and generate a table.",
            "生成网页": "请根据这张图片生成一个html网页(内联css、js)。",
            "解决这个问题": "请根据这张图片，解决这个问题。",
        }
        prompt_keys = list(prompts.keys())
        for i in range(0, len(prompt_keys), 3):
            cols = st.columns(3)
            for j in range(3):
                if i + j < len(prompt_keys):
                    key = prompt_keys[i + j]
                    value = prompts[key]
                    if cols[j].button(key, key=f"preset_{key}_{st.session_state.conversation}"):
                        perform_analysis(value, uploaded_files)
        
        # 自定义提示词输入框
        st.markdown("### 自定义提示词")
        custom_prompt = st.text_input(
            "请输入自定义提示词：",
            key=f"custom_prompt_{st.session_state.conversation}"
        )
        if custom_prompt:
            if st.button("确认", key=f"confirm_{st.session_state.conversation}"):
                perform_analysis(custom_prompt, uploaded_files)
    
    # 主区域显示
    st.markdown("## 上传的图片")
    if st.session_state.base64_images:
        for img_str in st.session_state.base64_images:
            st.image(f"data:image/jpeg;base64,{img_str}", caption='上传的图片')
    
    st.markdown("## 分析结果")
    st.text_area("分析结果", st.session_state.result, height=600)
    if st.session_state.result:
        st.download_button(
            label="下载结果",
            data=st.session_state.result,
            file_name="分析结果.md",
            mime="text/markdown"
        )

 # 修改主界面
 if __name__ == "__main__":
    st.title("图片分析器")
    upload_and_analyze()
	import os
	import base64
	import streamlit as st
	from mistralai import Mistral
	from PIL import Image
	import io
	from dotenv import load_dotenv

	# 加载 .env 文件
	load_dotenv()

	# 初始化会话状态
	if 'base64_images' not in st.session_state:
	st.session_state.base64_images = []
	if 'result' not in st.session_state:
	st.session_state.result = ""
	if 'conversation' not in st.session_state:
	st.session_state.conversation = 1
	if 'prompt' not in st.session_state:
	st.session_state.prompt = ""

	def perform_analysis(prompt, uploaded_files):
	if not uploaded_files:
	st.error("请先上传图片。")
	return

	if len(uploaded_files) > 8:
	st.error("最多只能上传8张图片。")
	return

	if not prompt:
	st.error("请先选择或输入提示词。")
	return

	base64_images = []
	for uploaded_file in uploaded_files:
	# 修改编码方式，使用压缩后的图片
	img = Image.open(uploaded_file)
	if img.mode in ("RGBA", "P"):
	img = img.convert("RGB") # 转换为 RGB 模式
	img.thumbnail((300, 300))
	buffer = io.BytesIO()
	img.save(buffer, format="JPEG")
	base64_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
	base64_images.append(base64_image)

	# 显示图片
	st.image(img, caption='上传的图片')

	st.session_state.base64_images = base64_images

	api_key = os.getenv("MISTRAL_API_KEY")
	if not api_key:
	st.error("错误：未找到MISTRAL_API_KEY。请确保在 .env 文件中设置了该变量。")
	return

	client = Mistral(api_key=api_key)

	content = [
	{"type": "text", "text": prompt}
	]
	for img_str in base64_images:
	content.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{img_str}"})

	messages = [
	{
	"role": "user",
	"content": content
	}
	]

	try:
	chat_response = client.chat.complete(
	model="pixtral-12b-2409",
	messages=messages
	)
	result = chat_response.choices[0].message.content
	except Exception as e:
	result = f"分析过程中出错：{str(e)}"

	st.session_state.result = result

	def upload_and_analyze():
	# 使用 Sidebar 代替左侧栏
	with st.sidebar:
	st.header("设置")

	# 新增"新对话"按钮
	if st.button("新对话"):
	st.session_state.base64_images = []
	st.session_state.result = ""
	st.session_state.conversation += 1 # 增加会话计数
	st.session_state.prompt = ""
	st.rerun() # 使用 st.rerun() 替代 st.experimental_rerun()

	uploaded_files = st.file_uploader(
	"上传图片",
	type=["jpg", "jpeg", "png"],
	accept_multiple_files=True,
	key=f"file_uploader_{st.session_state.conversation}"
	)

	# 显示预设提示词按钮，每行三个
	st.markdown("### 预设提示词")
	prompts = {
	"图像描述": "请描述这张图片的内容。",
	"图像推理": "请推理这张图片中发生了什么。",
	"解释图片": "请解释这张图片中的内容。",
	"信息提取": "请从这张图片中提取关键信息。",
	"内容转录（OCR）": "请转录这张图片中的文本内容。",
	"有趣故事": "请根据这些图片创作一个有趣的故事。",
	"数据分析": "请分析这些图表或数据可视化图像，并提供关键洞察。",
	"趋势识别": "请识别并解释这些图表中的主要趋势。",
	"数据比较": "请比较这些图表中的数据，并指出主要差异。",
	"Extract Content and Generate Table": "Please extract content from this image and generate a table.",
	"生成网页": "请根据这张图片生成一个html网页(内联css、js)。",
	"解决这个问题": "请根据这张图片，解决这个问题。",
	}
	prompt_keys = list(prompts.keys())
	for i in range(0, len(prompt_keys), 3):
	cols = st.columns(3)
	for j in range(3):
	if i + j < len(prompt_keys):
	key = prompt_keys[i + j]
	value = prompts[key]
	if cols[j].button(key, key=f"preset_{key}_{st.session_state.conversation}"):
	perform_analysis(value, uploaded_files)

	# 自定义提示词输入框
	st.markdown("### 自定义提示词")
	custom_prompt = st.text_input(
	"请输入自定义提示词：",
	key=f"custom_prompt_{st.session_state.conversation}"
	)
	if custom_prompt:
	if st.button("确认", key=f"confirm_{st.session_state.conversation}"):
	perform_analysis(custom_prompt, uploaded_files)

	# 主区域显示
	st.markdown("## 上传的图片")
	if st.session_state.base64_images:
	for img_str in st.session_state.base64_images:
	st.image(f"data:image/jpeg;base64,{img_str}", caption='上传的图片')

	st.markdown("## 分析结果")
	st.text_area("分析结果", st.session_state.result, height=600)
	if st.session_state.result:
	st.download_button(
	label="下载结果",
	data=st.session_state.result,
	file_name="分析结果.md",
	mime="text/markdown"
	)

	# 修改主界面
	if __name__ == "__main__":
	st.title("图片分析器")
	upload_and_analyze()