91 lines
3.4 KiB
Python
91 lines
3.4 KiB
Python
import os
|
||
from dotenv import load_dotenv
|
||
from generate_layout import call_deepseek
|
||
import json
|
||
from colorama import init, Fore, Back, Style
|
||
|
||
# 初始化colorama
|
||
init(autoreset=True)
|
||
|
||
# 加载环境变量
|
||
load_dotenv()
|
||
|
||
def llm_user_analysis(user_input):
|
||
"""
|
||
使用DeepSeek动态分析用户输入，提取海报设计参数
|
||
"""
|
||
if not user_input:
|
||
user_input = "端午节海报，包含背景、活动亮点和图标"
|
||
|
||
print(f"{Fore.CYAN}🔍 正在分析用户输入: {Style.BRIGHT}{user_input}{Style.RESET_ALL}")
|
||
|
||
# 构建分析提示词
|
||
analysis_prompt = f"""
|
||
请分析以下用户输入的海报需求，提取关键信息并以JSON格式返回：
|
||
|
||
用户输入：{user_input}
|
||
|
||
请严格按照以下JSON格式返回：
|
||
{{
|
||
"analyzed_prompt": "原始用户输入",
|
||
"keywords": ["提取的关键词1", "关键词2", "关键词3"],
|
||
"width": 1080,
|
||
"height": 1920,
|
||
"batch_size": 2,
|
||
"poster_type": "海报类型（如：节日海报、活动海报、产品海报等）",
|
||
"main_theme": "主要主题",
|
||
"style_preference": "风格偏好（如：现代、传统、简约等）",
|
||
"color_preference": "颜色偏好（如：暖色调、冷色调、单色等）"
|
||
}}
|
||
|
||
注意：
|
||
1. keywords应该包含3-5个最重要的关键词
|
||
2. 根据用户输入推断合适的海报类型
|
||
3. 尺寸默认为1080x1920，除非用户明确指定
|
||
4. batch_size根据需求调整，通常为1-4
|
||
5. 分析用户的风格和颜色偏好
|
||
"""
|
||
|
||
system_prompt = "你是一个专业的设计需求分析师，擅长从用户描述中提取海报设计的关键参数。请严格按照JSON格式返回结果，确保输出的JSON格式正确且完整。"
|
||
|
||
try:
|
||
result, _ = call_deepseek(prompt=analysis_prompt, system_prompt=system_prompt, temperature=0.3)
|
||
|
||
# 解析JSON
|
||
json_str = result.strip()
|
||
if "```json" in json_str:
|
||
json_str = json_str.split("```json")[1].split("```")[0].strip()
|
||
elif json_str.startswith("```") and json_str.endswith("```"):
|
||
json_str = json_str[3:-3].strip()
|
||
|
||
analysis_result = json.loads(json_str)
|
||
|
||
print(f"{Fore.GREEN}✅ 分析完成! {Style.RESET_ALL}")
|
||
print(f"{Fore.YELLOW}📊 主题: {Style.BRIGHT}{analysis_result.get('main_theme', '未知')}{Style.RESET_ALL}")
|
||
print(f"{Fore.YELLOW}🎨 风格: {analysis_result.get('style_preference', '未设置')}")
|
||
print(f"{Fore.YELLOW}🔖 关键词: {', '.join(analysis_result.get('keywords', []))}{Style.RESET_ALL}")
|
||
|
||
return analysis_result
|
||
|
||
except Exception as e:
|
||
print(f"{Fore.RED}❌ 分析失败: {str(e)}{Style.RESET_ALL}")
|
||
# 返回默认值
|
||
return {
|
||
"analyzed_prompt": user_input,
|
||
"keywords": ["海报", "设计", "活动"],
|
||
"width": 1080,
|
||
"height": 1920,
|
||
"batch_size": 2,
|
||
"poster_type": "通用海报",
|
||
"main_theme": "默认主题",
|
||
"style_preference": "现代",
|
||
"color_preference": "暖色调"
|
||
}
|
||
|
||
if __name__ == "__main__":
|
||
# 测试函数
|
||
test_input = "春节海报，红色背景，包含灯笼和烟花，现代风格"
|
||
result = llm_user_analysis(test_input)
|
||
print(f"\n{Fore.GREEN}测试结果:")
|
||
print(json.dumps(result, indent=2, ensure_ascii=False))
|