671 lines
24 KiB
Python
671 lines
24 KiB
Python
import os
|
||
import uuid
|
||
import requests
|
||
from flask import Flask, request, jsonify
|
||
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 QUEUE_NAME, DEFAULT_TIMEOUT, KEY_QUEUE_LOCK_NAME, acquire_lock
|
||
|
||
|
||
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('换泳装')
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
if queue["queue_running"] or queue["queue_pending"]:
|
||
return None, "cur gpu is busy"
|
||
|
||
with open('/home/szlc/code/ComfyUI/change_cloth/change_new.json', '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.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('生成写真')
|
||
queue = requests.get("http://localhost:8188/queue").json()
|
||
if queue["queue_running"] or queue["queue_pending"]:
|
||
return None, "cur gpu is busy"
|
||
|
||
with open('/home/szlc/code/ComfyUI/change_cloth/xiezhen_girl.json', '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."
|
||
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."
|
||
|
||
#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
|
||
|
||
kua = False
|
||
if human_json_data['跨'] == '有':
|
||
kua = True
|
||
|
||
|
||
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:
|
||
print(f'不用第二部 no2{no2}')
|
||
else:
|
||
print(f'要第二部 no2{no2} 先搞第一步')
|
||
# thigh_visible = is_thigh_visible(f"/home/szlc/code/ComfyUI/input/{human_name}")
|
||
if kua:
|
||
#脱衣服
|
||
takeoff_file_name = takeoff_cloth_first(human_name)
|
||
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('/home/szlc/code/ComfyUI/output', takeoff_file_name)
|
||
takeoff_file_path_name_input = os.path.join('/home/szlc/code/ComfyUI/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('/home/szlc/code/ComfyUI/output', generate_name)
|
||
generate_name_file_path_name_input = os.path.join('/home/szlc/code/ComfyUI/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('开始换衣服')
|
||
with open('/home/szlc/code/ComfyUI/change_cloth/change_new.json', '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"]
|
||
|
||
|
||
# 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, 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('/home/szlc/code/ComfyUI/change_cloth/static/imgs', filename)
|
||
|
||
# 解码并保存图片
|
||
with open(filepath, 'wb') as f:
|
||
f.write(base64.b64decode(data))
|
||
|
||
input_filepath = os.path.join('/home/szlc/code/ComfyUI/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'/home/szlc/code/ComfyUI/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'/home/szlc/code/ComfyUI/output/{out_human_name}')
|
||
out_human_jpg_name = out_human_name.replace(".png", ".jpg")
|
||
out_human_jpg_path_name = f'/home/szlc/code/ComfyUI/change_cloth/static/imgs/{out_human_jpg_name}'
|
||
out_human_image.save(out_human_jpg_path_name, quality=95)
|
||
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'/home/szlc/code/ComfyUI/output/{out_put_name}')
|
||
jpg_name = out_put_name.replace(".png", ".jpg")
|
||
jpg_path_name = f'/home/szlc/code/ComfyUI/change_cloth/static/imgs/{jpg_name}'
|
||
image.save(jpg_path_name, quality=95)
|
||
|
||
|
||
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('/home/szlc/code/ComfyUI/change_cloth/static/imgs', filename)
|
||
|
||
with open(filepath, "wb") as f:
|
||
for chunk in response.iter_content(1024):
|
||
f.write(chunk)
|
||
|
||
input_filepath = os.path.join('/home/szlc/code/ComfyUI/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')
|
||
|
||
no2 = data['no2']
|
||
|
||
try:
|
||
return process_change_cloth(human_filename, cloth_filename, output_format, cloth_url, human_url, no2)
|
||
except:
|
||
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
|
||
|
||
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:18888/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:18888/static/imgs/{cloth_filename}"
|
||
|
||
data["human_url"] = human_url
|
||
data["cloth_url"] = cloth_url
|
||
|
||
no2 = data.get('no2')
|
||
if not no2:
|
||
data['no2'] = False
|
||
|
||
return queueCall(data)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
app.run(host="0.0.0.0", port=8888, debug=True) |