154 lines
4.9 KiB
Python
154 lines
4.9 KiB
Python
import os
|
||
import uuid
|
||
import requests
|
||
from flask import Flask, request, jsonify
|
||
from urllib.parse import urlparse
|
||
from volcenginesdkarkruntime import Ark
|
||
import json
|
||
|
||
app = Flask(__name__)
|
||
|
||
# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
|
||
# 初始化Ark客户端，从环境变量中读取您的API Key
|
||
client = Ark(
|
||
# 此为默认路径，您可根据业务所在地域进行配置
|
||
base_url="https://ark.cn-beijing.volces.com/api/v3",
|
||
# 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
|
||
# api_key=os.environ.get("ARK_API_KEY"),
|
||
api_key='14fc0280-fc65-462d-ac2d-50178c0212e3'
|
||
)
|
||
|
||
def parse_facial_data(data_str):
|
||
print(f'原始字符串 {data_str}')
|
||
# 去除多余的标记和换行符
|
||
json_str_clean = data_str.strip('```json\n').strip('```').strip()
|
||
|
||
# 转换为Python字典
|
||
data = json.loads(json_str_clean)
|
||
s = data['面部年龄']
|
||
s = s.replace('岁', '')
|
||
parts = s.split('-')
|
||
start_age = int(parts[0])
|
||
end_age = int(parts[1])
|
||
data['start_age'] = start_age
|
||
data['end_age'] = end_age
|
||
return data
|
||
|
||
def GetPicDesc(img_url):
|
||
response = client.chat.completions.create(
|
||
# 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
|
||
model="doubao-1.5-vision-pro-250328",
|
||
messages=[
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": img_url
|
||
},
|
||
},
|
||
{"type": "text", "text": "告诉我图片中的服装长度, 只要答案, 选项有(腰部以上、腰、跨、大腿上部1/3、大腿中部、膝盖、小腿、小腿上部1/3、脚踝、拖地、难以辨认)"},
|
||
],
|
||
}
|
||
],
|
||
|
||
)
|
||
|
||
text = response.choices[0].message.content
|
||
result = {}
|
||
|
||
try:
|
||
result = parse_facial_data(text)
|
||
except:
|
||
print(f'Error {text}')
|
||
|
||
return result
|
||
|
||
# ComfyUI 的 input 目录路径（根据实际修改）
|
||
COMFYUI_INPUT_DIR = "/path/to/ComfyUI/input"
|
||
os.makedirs(COMFYUI_INPUT_DIR, exist_ok=True)
|
||
|
||
def get_file_extension(url_or_filename):
|
||
"""从 URL 或文件名中提取扩展名（如 .jpg、.png）"""
|
||
# 处理 URL 情况（如 https://example.com/image.jpg?width=200）
|
||
if url_or_filename.startswith(('http://', 'https://')):
|
||
parsed = urlparse(url_or_filename)
|
||
path = parsed.path
|
||
else:
|
||
path = url_or_filename
|
||
|
||
# 提取扩展名（转换为小写，去掉问号后的参数）
|
||
ext = os.path.splitext(path)[1].lower().split('?')[0]
|
||
return ext if ext else '.png' # 默认 PNG（如果无法提取）
|
||
|
||
def save_image_from_url(image_url):
|
||
"""从 URL 下载图片并保留原始格式"""
|
||
try:
|
||
response = requests.get(image_url, stream=True)
|
||
response.raise_for_status()
|
||
|
||
# 获取原始图片格式
|
||
ext = get_file_extension(image_url)
|
||
filename = f"{uuid.uuid4()}{ext}"
|
||
save_path = os.path.join(COMFYUI_INPUT_DIR, filename)
|
||
|
||
with open(save_path, "wb") as f:
|
||
for chunk in response.iter_content(1024):
|
||
f.write(chunk)
|
||
|
||
return filename
|
||
except Exception as e:
|
||
print(f"Error saving image from URL: {e}")
|
||
return None
|
||
|
||
@app.route('/change_cloth', methods=['POST'])
|
||
def change_cloth():
|
||
"""从 URL 下载图片"""
|
||
data = request.json
|
||
human_url = data.get("human_url")
|
||
|
||
if not human_url:
|
||
return jsonify({"error": "Missing 'human_url' parameter"}), 400
|
||
|
||
human_filename = save_image_from_url(human_url)
|
||
if not human_filename:
|
||
return jsonify({"error": "Failed to download or save human image"}), 500
|
||
|
||
cloth_url = data.get("cloth_url")
|
||
|
||
if not cloth_url:
|
||
return jsonify({"error": "Missing 'cloth_url' parameter"}), 400
|
||
|
||
cloth_filename = save_image_from_url(cloth_url)
|
||
if not cloth_filename:
|
||
return jsonify({"error": "Failed to download or save image"}), 500
|
||
|
||
|
||
|
||
return jsonify({"filename": cloth_filename})
|
||
|
||
@app.route('/upload_image', methods=['POST'])
|
||
def upload_image():
|
||
"""直接上传图片（保留原始格式）"""
|
||
if 'file' not in request.files:
|
||
return jsonify({"error": "No file uploaded"}), 400
|
||
|
||
file = request.files['file']
|
||
if file.filename == '':
|
||
return jsonify({"error": "Empty filename"}), 400
|
||
|
||
# 获取原始文件扩展名
|
||
ext = os.path.splitext(file.filename)[1].lower()
|
||
if not ext: # 如果无扩展名，默认 PNG
|
||
ext = '.png'
|
||
|
||
# 生成唯一文件名（保留原始格式）
|
||
filename = f"{uuid.uuid4()}{ext}"
|
||
save_path = os.path.join(COMFYUI_INPUT_DIR, filename)
|
||
|
||
file.save(save_path)
|
||
return jsonify({"filename": filename})
|
||
|
||
if __name__ == '__main__':
|
||
app.run(host="0.0.0.0", port=5000, debug=True) |