764 lines
28 KiB
Python
764 lines
28 KiB
Python
import os
|
||
import uuid
|
||
import requests
|
||
from flask import Flask, request, jsonify, Response
|
||
from urllib.parse import urlparse
|
||
import json
|
||
import base64
|
||
import time
|
||
from PIL import Image
|
||
from datetime import datetime
|
||
from volcenginesdkarkruntime import Ark
|
||
import shutil
|
||
import oss2
|
||
import redis
|
||
# from check_img_body import is_thigh_visible
|
||
import logging
|
||
from config import *
|
||
import websockets
|
||
import asyncio
|
||
import traceback
|
||
|
||
|
||
APP_ROOT = os.path.dirname(os.path.abspath(__file__))
|
||
STATIC_FOLDER = os.path.join(APP_ROOT, 'static')
|
||
|
||
# 创建 logger
|
||
logger = logging.getLogger(__name__)
|
||
logger.setLevel(logging.INFO) # 设置 logger 的级别
|
||
|
||
# 获取当前日期和时间
|
||
now = datetime.now()
|
||
|
||
# 格式化为字符串（例如：2023-10-25 14:30:45）
|
||
date_time_str = now.strftime("%Y-%m-%d_%H:%M:%S")
|
||
|
||
# 创建文件 handler
|
||
file_handler = logging.FileHandler(f'/var/log/fuyan/change_app_{date_time_str}.log')
|
||
file_handler.setLevel(logging.INFO)
|
||
file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||
file_handler.setFormatter(file_formatter)
|
||
|
||
# 创建控制台 handler
|
||
console_handler = logging.StreamHandler()
|
||
console_handler.setLevel(logging.INFO)
|
||
console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
||
console_handler.setFormatter(console_formatter)
|
||
|
||
# 添加 handlers 到 logger
|
||
logger.addHandler(file_handler)
|
||
logger.addHandler(console_handler)
|
||
|
||
from config import REDIS_HOST
|
||
from config import REDIS_PORT
|
||
from config import REDIS_DB
|
||
|
||
# 创建 Redis 连接池
|
||
redis_pool = redis.ConnectionPool(
|
||
host=REDIS_HOST,
|
||
port=REDIS_PORT,
|
||
db=REDIS_DB,
|
||
max_connections=2000 # 根据实际情况调整
|
||
)
|
||
|
||
def get_redis_conn():
|
||
"""获取 Redis 连接"""
|
||
return redis.Redis(connection_pool=redis_pool)
|
||
|
||
app = Flask(__name__)
|
||
|
||
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 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": "图片是一件服装的照片，返回为json格式， 第一个字段'服装长度', 图片中的服装穿到正常人身上的长度会覆盖到人体的哪个部位, 只要答案, 选项有(胸、腰、跨、大腿、膝盖、小腿、脚踝、拖地、难以辨认). 第二个字段'衣袖'，只要答案，选项有(短袖、长袖)"},
|
||
],
|
||
}
|
||
],
|
||
|
||
)
|
||
|
||
text = response.choices[0].message.content
|
||
print(text)
|
||
return text
|
||
|
||
def GetHumanDesDesc(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": "图片是一个人物的照片，返回为json格式， 第一个字段'性别', 判断一下性别，只要答案, 选项有(男、女). 第二个字段'身高完整度'，判断一下照片中人物的身高完整度（用小数表示）"},
|
||
],
|
||
}
|
||
],
|
||
|
||
)
|
||
|
||
text = response.choices[0].message.content
|
||
print(text)
|
||
return text
|
||
|
||
def takeoff_cloth_first(human_name, is_girl):
|
||
print(f'换泳装 {change_with_add_cloth}')
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
if queue["queue_running"] or queue["queue_pending"]:
|
||
return None, "cur gpu is busy"
|
||
|
||
with open(f'{APP_ROOT}/{change_with_add_cloth}', 'r', encoding='utf-8') as file:
|
||
prompt_text = file.read()
|
||
|
||
prompt = json.loads(prompt_text)
|
||
|
||
prompt["96"]["inputs"]["cloth_len"] = '拖地'
|
||
prompt["96"]["inputs"]["cloth_short"] = True
|
||
|
||
# prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)
|
||
|
||
#input cloth img
|
||
if is_girl:
|
||
prompt["22"]["inputs"]["image"] = 'girl_cloth.jpg'
|
||
else:
|
||
prompt["22"]["inputs"]["image"] = 'man_cloth_new.jpg'
|
||
# prompt["22"]["inputs"]["image"] = 'man_cloth.jpg'
|
||
|
||
#input human img
|
||
prompt["61"]["inputs"]["image"] = human_name
|
||
|
||
#out put name
|
||
out_img_name = str(uuid.uuid4())[:8]
|
||
prompt["102"]["inputs"]["filename_prefix"] = out_img_name
|
||
|
||
p = {"prompt": prompt}
|
||
|
||
data = json.dumps(p).encode('utf-8')
|
||
|
||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||
prompt_id = response.json()["prompt_id"]
|
||
|
||
|
||
# 2. 轮询队列，直到任务完成
|
||
while True:
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
# print(queue)
|
||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||
break # 队列为空，任务已完成
|
||
time.sleep(0.5) # 避免频繁请求
|
||
|
||
# 3. 从历史记录中获取结果
|
||
history = requests.get("http://localhost:8188/history").json()
|
||
# print("History:", history)
|
||
outputs = history[prompt_id]["outputs"]
|
||
for out in outputs:
|
||
if 'images' in outputs[out]:
|
||
for out_img in outputs[out]['images']:
|
||
if 'filename' in out_img:
|
||
out_img_file_name = out_img['filename']
|
||
if out_img_name in out_img_file_name:
|
||
return out_img_file_name
|
||
return None
|
||
|
||
def generate_from_face(human_name, sex_girl):
|
||
print(f'生成写真 {xiezhen_name}')
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
if queue["queue_running"] or queue["queue_pending"]:
|
||
return None, "cur gpu is busy"
|
||
|
||
with open(f'{APP_ROOT}/{xiezhen_name}', 'r', encoding='utf-8') as file:
|
||
prompt_text = file.read()
|
||
|
||
prompt = json.loads(prompt_text)
|
||
|
||
|
||
prompt["93"]["inputs"]["image"] = human_name
|
||
|
||
if sex_girl:
|
||
prompt["6"]["inputs"]["text"] = "asian girl,full body shot(1.9),complete figure,front-facing,model pose,hands are at the sides of the body,smile,wearing black tube top and micro skirt,simple ,white background,32k,high detail,and face is illuminated by soft side light and natural light. The sunlight spills over the white wall behind."
|
||
else:
|
||
prompt["6"]["inputs"]["text"] = "asian man,full body shot(1.9),complete figure,front-facing,model pose,hands are at the sides of the body,smile,wearing black tight shortst,simple white background,32k,high detail,and face is illuminated by soft side light and natural light. The sunlight spills over the white wall behind."
|
||
|
||
#out put name
|
||
out_img_name = str(uuid.uuid4())[:8]
|
||
prompt["65"]["inputs"]["filename_prefix"] = out_img_name
|
||
|
||
p = {"prompt": prompt}
|
||
|
||
data = json.dumps(p).encode('utf-8')
|
||
|
||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||
prompt_id = response.json()["prompt_id"]
|
||
|
||
|
||
# 2. 轮询队列，直到任务完成
|
||
while True:
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
# print(queue)
|
||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||
break # 队列为空，任务已完成
|
||
time.sleep(0.5) # 避免频繁请求
|
||
|
||
# 3. 从历史记录中获取结果
|
||
history = requests.get("http://localhost:8188/history").json()
|
||
# print("History:", history)
|
||
outputs = history[prompt_id]["outputs"]
|
||
for out in outputs:
|
||
if 'images' in outputs[out]:
|
||
for out_img in outputs[out]['images']:
|
||
if 'filename' in out_img:
|
||
out_img_file_name = out_img['filename']
|
||
if out_img_name in out_img_file_name:
|
||
return out_img_file_name
|
||
return None
|
||
|
||
def change(human_name, cloth_name, c_width, c_height, cloth_url, human_url, no2):
|
||
human_json_str = GetHumanDesDesc(human_url)
|
||
print(f"human_json_str {human_json_str}")
|
||
human_json_data = json.loads(human_json_str)
|
||
sex_type = human_json_data['性别']
|
||
sex_girl = True
|
||
if '男' in sex_type:
|
||
sex_girl = False
|
||
|
||
sgwzd = human_json_data['身高完整度']
|
||
|
||
json_str = GetPicDesc(cloth_url)
|
||
json_data = json.loads(json_str)
|
||
cloth_len = json_data['服装长度']
|
||
cloth_short = True
|
||
if '长袖' in json_data['衣袖']:
|
||
cloth_short = False
|
||
if cloth_len == '难以辨认':
|
||
return None, None, sex_girl, "get image type error"
|
||
|
||
if no2 == True:
|
||
print(f'不用第二部 no2{no2}')
|
||
else:
|
||
print(f'要第二部 no2{no2} 先搞第一步')
|
||
if sgwzd >= 0.55:
|
||
#脱衣服
|
||
takeoff_file_name = takeoff_cloth_first(human_name, sex_girl)
|
||
if takeoff_file_name == None:
|
||
return None, None, sex_girl, f"takeoff_cloth_first error {human_name}"
|
||
else:
|
||
human_name = takeoff_file_name
|
||
takeoff_file_path_name = os.path.join(f'{APP_ROOT}/../output', takeoff_file_name)
|
||
takeoff_file_path_name_input = os.path.join(f'{APP_ROOT}/../input', takeoff_file_name)
|
||
shutil.copy(takeoff_file_path_name, takeoff_file_path_name_input)
|
||
else:
|
||
#生成写真
|
||
generate_name = generate_from_face(human_name, sex_girl)
|
||
if generate_name == None:
|
||
return None, None, sex_girl, f"xiezhen error {human_name}"
|
||
else:
|
||
human_name = generate_name
|
||
generate_name_file_path_name = os.path.join(f'{APP_ROOT}/../output', generate_name)
|
||
generate_name_file_path_name_input = os.path.join(f'{APP_ROOT}/../input', generate_name)
|
||
shutil.copy(generate_name_file_path_name, generate_name_file_path_name_input)
|
||
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
if queue["queue_running"] or queue["queue_pending"]:
|
||
return None, None, sex_girl, "cur gpu is busy, Shou not happen"
|
||
|
||
print(f'开始换衣服 {change_only_name}')
|
||
with open(f'{APP_ROOT}/{change_only_name}', 'r', encoding='utf-8') as file:
|
||
prompt_text = file.read()
|
||
|
||
prompt = json.loads(prompt_text)
|
||
|
||
prompt["96"]["inputs"]["cloth_len"] = cloth_len
|
||
prompt["96"]["inputs"]["cloth_short"] = cloth_short
|
||
|
||
# prompt["99"]["inputs"]["width"] = int((c_width/c_height) * 1024)
|
||
|
||
#input cloth img
|
||
prompt["22"]["inputs"]["image"] = cloth_name
|
||
|
||
#input human img
|
||
prompt["61"]["inputs"]["image"] = human_name
|
||
|
||
#out put name
|
||
out_img_name = str(uuid.uuid4())[:8]
|
||
prompt["102"]["inputs"]["filename_prefix"] = out_img_name
|
||
|
||
p = {"prompt": prompt}
|
||
|
||
data = json.dumps(p).encode('utf-8')
|
||
|
||
response = requests.post("http://localhost:8188/prompt", data=data)
|
||
prompt_id = response.json()["prompt_id"]
|
||
|
||
|
||
async def listen_node_output(prompt_id, target_node_id):
|
||
uri = "ws://localhost:8188/ws"
|
||
async with websockets.connect(uri) as websocket:
|
||
await websocket.send(json.dumps({"prompt_id": prompt_id}))
|
||
|
||
while True:
|
||
message = await websocket.recv()
|
||
|
||
# 1. 处理二进制数据
|
||
if isinstance(message, bytes):
|
||
try:
|
||
message = message.decode("utf-8")
|
||
except UnicodeDecodeError:
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] 收到二进制数据 (长度: {len(message)} bytes)")
|
||
continue
|
||
|
||
# 2. 解析JSON
|
||
try:
|
||
data = json.loads(message)
|
||
print(f"cur data {data}")
|
||
except json.JSONDecodeError:
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] 非JSON数据: {message[:100]}...")
|
||
continue
|
||
|
||
if data.get("type") == "progress":
|
||
pdata = data.get('data', {})
|
||
print(f"cur progress {pdata}")
|
||
# if value >= max:
|
||
# return f"finished value:{value} max:{max}"
|
||
|
||
if data.get("type") == "progress_state":
|
||
nodes = data.get('data', {}).get('nodes', {})
|
||
for node in nodes:
|
||
if node.get("state", None) != 'finished':
|
||
print(f"progress_state node:{node}")
|
||
|
||
# 3. 输出运行状态
|
||
if data.get("type") == "status":
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] 系统状态: {data.get('data', {}).get('status', {})}")
|
||
status = data.get('data', {}).get('status', None)
|
||
print(f"statue:{status}")
|
||
if status:
|
||
exec_info = status.get('exec_info', None)
|
||
print(f"exec_info:{exec_info}")
|
||
if exec_info:
|
||
queue_remaining = exec_info.get('queue_remaining', None)
|
||
print(f"queue_remaining:{queue_remaining}")
|
||
if queue_remaining == 0:
|
||
return "process end"
|
||
continue
|
||
|
||
if data.get("type") == "executing":
|
||
node_id = data.get("data", {}).get("node")
|
||
progress = data.get("data", {}).get("progress", 0)
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] 正在执行节点 {node_id} (进度: {progress:.0%})")
|
||
|
||
# 4. 目标节点完成
|
||
if data.get("type") == "executed" and data.get("data", {}).get("node") == target_node_id:
|
||
output = data.get("data", {}).get("output")
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] 节点 {target_node_id} 完成!")
|
||
return output
|
||
|
||
|
||
# 调用方式保持不变
|
||
output = asyncio.run(listen_node_output(prompt_id, "104"))
|
||
print("最终输出:", output)
|
||
|
||
while True:
|
||
# response = requests.post("http://localhost:8188/interrupt")
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
# print(queue)
|
||
if not queue["queue_running"] and not queue["queue_pending"]:
|
||
break # 队列为空，任务已完成
|
||
time.sleep(0.5) # 避免频繁请求
|
||
|
||
# 3. 从历史记录中获取结果
|
||
history = requests.get("http://localhost:8188/history").json()
|
||
# print("History:", history)
|
||
outputs = history[prompt_id]["outputs"]
|
||
for out in outputs:
|
||
if 'images' in outputs[out]:
|
||
for out_img in outputs[out]['images']:
|
||
if 'filename' in out_img:
|
||
out_img_file_name = out_img['filename']
|
||
if out_img_name in out_img_file_name:
|
||
return out_img_file_name, human_name, sex_girl,"success"
|
||
return None, human_name, sex_girl, "can not find output image"
|
||
|
||
def save_base64_image(base64_str, prefix):
|
||
"""
|
||
将base64字符串保存为图片文件
|
||
:param base64_str: 带有类型标识的base64字符串
|
||
:param prefix: 文件名前缀
|
||
:return: 保存的文件路径
|
||
"""
|
||
try:
|
||
# 分离base64头部和实际数据
|
||
header, data = base64_str.split(',', 1)
|
||
|
||
# 从头部获取文件扩展名
|
||
file_ext = ''
|
||
if 'image/png' in header:
|
||
file_ext = '.png'
|
||
elif 'image/jpeg' in header:
|
||
file_ext = '.jpg'
|
||
elif 'image/jpg' in header:
|
||
file_ext = '.jpg'
|
||
elif 'image/gif' in header:
|
||
file_ext = '.gif'
|
||
else:
|
||
file_ext = '.png' # 默认使用png
|
||
|
||
# 生成唯一文件名
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
unique_id = str(uuid.uuid4())[:8]
|
||
filename = f"{prefix}_{timestamp}_{unique_id}{file_ext}"
|
||
filepath = os.path.join(f'{APP_ROOT}/static/imgs', filename)
|
||
|
||
# 解码并保存图片
|
||
with open(filepath, 'wb') as f:
|
||
f.write(base64.b64decode(data))
|
||
|
||
input_filepath = os.path.join(f'{APP_ROOT}/../input', filename)
|
||
# 解码并保存图片
|
||
with open(input_filepath, 'wb') as f:
|
||
f.write(base64.b64decode(data))
|
||
|
||
return filename
|
||
|
||
except Exception as e:
|
||
print(f"Error saving image: {str(e)}")
|
||
return None
|
||
|
||
def get_image_dimensions(image_path):
|
||
try:
|
||
with Image.open(image_path) as img:
|
||
width, height = img.size
|
||
return width, height
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
return None
|
||
|
||
def image_to_base64(file_path, mime_type=None):
|
||
if mime_type is None:
|
||
extension = file_path.split('.')[-1].lower()
|
||
mime_types = {
|
||
'jpg': 'image/jpeg',
|
||
'jpeg': 'image/jpeg',
|
||
'png': 'image/png',
|
||
'gif': 'image/gif',
|
||
'webp': 'image/webp',
|
||
'bmp': 'image/bmp'
|
||
}
|
||
mime_type = mime_types.get(extension, 'application/octet-stream')
|
||
|
||
# 读取文件内容并编码为Base64
|
||
with open(file_path, 'rb') as image_file:
|
||
encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
|
||
|
||
# 组合成Data URI格式
|
||
return f"data:{mime_type};base64,{encoded_string}"
|
||
|
||
def upload_to_oss(image_path, object_name=None):
|
||
# 配置信息（替换为你的实际信息）
|
||
access_key_id = 'LTAI5tB9t2RH6f1drSLvVLLZ'
|
||
access_key_secret = '91uzPI1RAFHN7n3Y6TJDGFP8w0dG1R'
|
||
endpoint = 'oss-cn-beijing.aliyuncs.com' # 替换为你的Endpoint
|
||
bucket_name = 'llyz'
|
||
|
||
# 创建Bucket实例
|
||
auth = oss2.Auth(access_key_id, access_key_secret)
|
||
bucket = oss2.Bucket(auth, endpoint, bucket_name)
|
||
|
||
# 如果没有指定OSS文件名，则使用本地文件名
|
||
if object_name is None:
|
||
object_name = image_path.split('/')[-1] # 取本地文件名
|
||
|
||
try:
|
||
# 上传文件
|
||
bucket.put_object_from_file(object_name, image_path)
|
||
|
||
# 获取文件URL（有效期默认10年）
|
||
url = bucket.sign_url('GET', object_name, 3600 * 24 * 365 * 10)
|
||
|
||
# 或者使用公共读Bucket的URL（如果Bucket是公共读权限）
|
||
# url = f"https://{bucket_name}.{endpoint}/{object_name}"
|
||
|
||
print(f"文件上传成功，URL: {url}")
|
||
return url
|
||
except Exception as e:
|
||
print(f"上传失败: {str(e)}")
|
||
return None
|
||
|
||
|
||
def process_change_cloth(human_filename, cloth_filename, output_format, img_url, human_url, no2):
|
||
w,h = get_image_dimensions(f'{APP_ROOT}/../input/{human_filename}')
|
||
|
||
out_put_name, out_human_name, is_girl, msg = change(human_filename, cloth_filename, w, h, img_url, human_url, no2)
|
||
if out_put_name == None:
|
||
print(f'Failed to change cloth {msg}')
|
||
return jsonify({"ret":-1, 'msg': f'Failed to change cloth {msg}'}), 200
|
||
|
||
if no2:
|
||
out_human_https_url = human_url
|
||
else:
|
||
out_human_image = Image.open(f'{APP_ROOT}/../output/{out_human_name}')
|
||
out_human_jpg_name = out_human_name.replace(".png", ".jpg")
|
||
out_human_jpg_path_name = f'{APP_ROOT}/static/imgs/{out_human_jpg_name}'
|
||
out_human_image.save(out_human_jpg_path_name, quality=90)
|
||
upload_to_oss(out_human_jpg_path_name, out_human_jpg_name) # 第二个参数可选，指定OSS上的路径
|
||
out_human_https_url = f'https://llyz.oss-cn-beijing.aliyuncs.com/{out_human_jpg_name}'
|
||
print(f"生成的第一步图片 HTTPS URL: {out_human_https_url}")
|
||
|
||
|
||
image = Image.open(f'{APP_ROOT}/../output/{out_put_name}')
|
||
jpg_name = out_put_name.replace(".png", ".jpg")
|
||
jpg_path_name = f'{APP_ROOT}/static/imgs/{jpg_name}'
|
||
image.save(jpg_path_name, quality=90)
|
||
|
||
|
||
upload_to_oss(jpg_path_name, jpg_name) # 第二个参数可选，指定OSS上的路径
|
||
https_url = f'https://llyz.oss-cn-beijing.aliyuncs.com/{jpg_name}'
|
||
print(f"生成的HTTPS URL: {https_url}")
|
||
|
||
if 'base64' in output_format:
|
||
return jsonify({
|
||
"ret":0,
|
||
"state": 0,
|
||
"msg":"success",
|
||
"first_step_data":image_to_base64(out_human_jpg_path_name),
|
||
"is_girl":is_girl,
|
||
"data":image_to_base64(jpg_path_name)
|
||
})
|
||
else:
|
||
return jsonify({
|
||
"ret":0,
|
||
"state": 0,
|
||
"msg":"success",
|
||
"first_url":out_human_https_url,
|
||
"is_girl":is_girl,
|
||
"url":https_url
|
||
})
|
||
|
||
|
||
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}"
|
||
filepath = os.path.join(f'{APP_ROOT}/static/imgs', filename)
|
||
|
||
with open(filepath, "wb") as f:
|
||
for chunk in response.iter_content(1024):
|
||
f.write(chunk)
|
||
|
||
input_filepath = os.path.join(f'{APP_ROOT}/../input', filename)
|
||
shutil.copy(filepath, input_filepath)
|
||
|
||
return filename
|
||
except Exception as e:
|
||
print(f"Error saving image from URL: {e}")
|
||
return None
|
||
|
||
|
||
@app.route('/do_change_cloth', methods=['POST'])
|
||
def do_change_cloth():
|
||
"""从 URL 下载图片"""
|
||
data = request.json
|
||
|
||
human_url = data.get("human_url")
|
||
|
||
human_filename = save_image_from_url(human_url)
|
||
if not human_filename:
|
||
return jsonify({"state":-1, "error": "Failed to download or save human image"}), 500
|
||
|
||
cloth_url = data.get("cloth_url")
|
||
cloth_filename = save_image_from_url(cloth_url)
|
||
if not cloth_filename:
|
||
return jsonify({"state":-1, "error": "Failed to download or save image"}), 500
|
||
|
||
output_format = data.get('output_format')
|
||
|
||
def check_type(var):
|
||
if isinstance(var, bool):
|
||
print(f"{var} 是布尔值no2")
|
||
elif isinstance(var, str):
|
||
print(f"'{var}' 是字符串no2")
|
||
else:
|
||
print(f"既不是字符串也不是布尔值no2，实际类型: {type(var)}")
|
||
|
||
no2 = data['no2']
|
||
check_type(no2)
|
||
|
||
try:
|
||
return process_change_cloth(human_filename, cloth_filename, output_format, cloth_url, human_url, no2)
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
print(f"错误详情：{e}")
|
||
return jsonify({
|
||
"ret":-1,
|
||
"state": -1,
|
||
"msg":"process_change_cloth error"
|
||
})
|
||
|
||
get_redis_conn().set(QUEUE_NAME, "[]")
|
||
|
||
def pushStr2Queue(key, data_str):
|
||
redis_conn = get_redis_conn()
|
||
lock = acquire_lock(redis_conn, KEY_QUEUE_LOCK_NAME)
|
||
if lock:
|
||
try:
|
||
redis_conn.set(key, data_str, ex=180)
|
||
key_queue_str = redis_conn.get(QUEUE_NAME).decode("utf-8")
|
||
key_queue = json.loads(key_queue_str)
|
||
key_queue.append(key)
|
||
logger.info(f"pushStr2Queue queue len:{key_queue}")
|
||
key_queue_str = json.dumps(key_queue)
|
||
redis_conn.set(QUEUE_NAME, key_queue_str)
|
||
finally:
|
||
lock.release()
|
||
else:
|
||
logger.error(f"pushStr2Queue get lock error {data_str}")
|
||
|
||
def queueCall(data):
|
||
key = str(uuid.uuid4())
|
||
redis_conn = get_redis_conn()
|
||
|
||
task_data = {}
|
||
task_data['key'] = key
|
||
task_data['api'] = "/do_change_cloth"
|
||
task_data['request'] = data
|
||
logger.info(f"request do_change_cloth img:{data['human_url'][:64]}, output_format:{data['output_format']}")
|
||
task_data_str = json.dumps(task_data)
|
||
pushStr2Queue(key, task_data_str)
|
||
start_time = datetime.now()
|
||
result_key = f"result_{key}"
|
||
result_status_key = f"result_status_code_{key}"
|
||
while (datetime.now() - start_time).seconds < DEFAULT_TIMEOUT:
|
||
if redis_conn.exists(result_key):
|
||
result_str = redis_conn.get(result_key)
|
||
return result_str, int(redis_conn.get(result_status_key)), {'Content-Type': 'application/json'}
|
||
time.sleep(0.1)
|
||
|
||
logger.error(f"request hairColor time out ")
|
||
return jsonify({
|
||
'msg': f'Timeout after time out'
|
||
, "state":-1, "data":""
|
||
}), 408
|
||
|
||
|
||
@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({"state":-1, "error": "Missing 'human_url' parameter"}), 400
|
||
|
||
cloth_url = data.get("cloth_url")
|
||
if not cloth_url:
|
||
return jsonify({"state":-1, "error": "Missing 'cloth_url' parameter"}), 400
|
||
|
||
output_format = data.get('output_format')
|
||
if not output_format:
|
||
return jsonify({"state":-1, "error": "Missing 'output_format' parameter"}), 500
|
||
|
||
no2 = data.get('no2')
|
||
if not no2:
|
||
data['no2'] = False
|
||
|
||
print(f"change_cloth input data:{data}")
|
||
|
||
return queueCall(data)
|
||
|
||
@app.route('/change_cloth_base64', methods=['POST'])
|
||
def change_cloth_base64():
|
||
# 获取参数
|
||
data = request.get_json()
|
||
if not data:
|
||
return jsonify({"ret":-1, "state":-1, 'msg': 'No JSON data provided'}), 400
|
||
|
||
human_img = data.get('human_img')
|
||
cloth_img = data.get('cloth_img')
|
||
output_format = data.get('output_format')
|
||
if not output_format:
|
||
return jsonify({"state":-1, "error": "Missing 'output_format' parameter"}), 500
|
||
|
||
if not human_img or not cloth_img:
|
||
return jsonify({"ret":-1, 'msg': 'Both human_img and cloth_img are required'}), 400
|
||
|
||
human_filename = save_base64_image(human_img, 'human')
|
||
if not human_filename:
|
||
return jsonify({"ret":-1, 'msg': 'Failed to save human image'}), 500
|
||
data['human_img'] = None
|
||
|
||
human_url = f"http://112.126.94.241:{base64_test_port}/static/imgs/{human_filename}"
|
||
|
||
# 保存服装图片
|
||
cloth_filename = save_base64_image(cloth_img, 'cloth')
|
||
if not cloth_filename:
|
||
return jsonify({"ret":-1, 'msg': 'Failed to save cloth image'}), 500
|
||
data['cloth_img'] = None
|
||
|
||
cloth_url = f"http://112.126.94.241:{base64_test_port}/static/imgs/{cloth_filename}"
|
||
|
||
data["human_url"] = human_url
|
||
data["cloth_url"] = cloth_url
|
||
|
||
no2 = data.get('no2')
|
||
if not no2:
|
||
data['no2'] = False
|
||
|
||
# 在内部调用第二个HTTP请求
|
||
try:
|
||
# 调用第二个API（可以是外部服务或自己的另一个端点）
|
||
response = requests.post(f'http://112.126.94.241:{base64_test_port}/do_change_cloth', json=data)
|
||
|
||
return Response(
|
||
response=response.content,
|
||
status=response.status_code,
|
||
headers=dict(response.headers)
|
||
)
|
||
except requests.exceptions.RequestException as e:
|
||
data['second_api_error'] = str(e)
|
||
|
||
return jsonify("error"), 500
|
||
|
||
|
||
if __name__ == '__main__':
|
||
app.run(host="0.0.0.0", port=8888, debug=True) |