KeCheng/BaiHu/Util/CommonUtil.py

import json
import base64
import os.path
import uuid
from datetime import datetime
import requests
from PIL import Image
import io

def submit_post(url: str, data: dict):
    """
    Submit a POST request to the given URL with the given data.
    """
    return requests.post(url, data=json.dumps(data))


def save_encoded_image(b64_image: str, output_path: str):
    """
    Save the given image to the given output path.
    """
    with open(output_path, "wb") as image_file:
        image_file.write(base64.b64decode(b64_image))


def encode_image(image_path):
    with open(image_path, "rb") as i:
        b64 = base64.b64encode(i.read())
    return b64.decode("utf-8")


# 打开一个JSON文件用于读取
def getJson(json_path):
    with open(json_path, 'r', encoding='utf-8') as f:
        # 从文件中加载JSON数据到字典
        data = json.load(f)
        return data


def printf(str):
    # 获取当前时间
    now = datetime.now()
    # 格式化时间
    current_time = now.strftime("%Y-%m-%d %H:%M:%S")
    # 打印当前时间
    print(current_time + ' ' + str)


# 获取图片的Base64编码
def get_image_base64_str(image_path):
    with open(image_path, 'rb') as image_file:
        # 使用Pillow库打开图像文件
        image = Image.open(image_file)
        # 转换图像为RGB模式，因为JPEG不支持透明度
        image = image.convert('RGB')
        # 创建一个字节流对象
        buffer = io.BytesIO()
        # 将图像保存到字节流中，格式为JPEG
        image.save(buffer, format='JPEG')
        # 移动到字节流的开始位置
        buffer.seek(0)
        # 将字节流中的数据编码为base64
        encoded_image = base64.b64encode(buffer.read())
    # 将编码后的图像数据转换为字符串，并构建完整的URL格式
    base64_image_url = 'data:image/jpeg;base64,' + encoded_image.decode('utf-8')
    return base64_image_url


# 将图片一切为四
def split_4_image(image_path, out_path):
    res = []
    # 获取图片的宽度
    img = Image.open(image_path)
    width, height = img.size

    new_width, new_height = width // 2, height // 2

    # 左上角的图片
    top_left = img.crop((0, 0, new_width, new_height))
    f_name = str(uuid.uuid4()) + '.png'
    fi = os.path.join(out_path, f_name)
    res.append(f_name)
    top_left.save(fi)

    # 右上角的图片
    top_right = img.crop((new_width, 0, width, new_height))
    f_name = str(uuid.uuid4()) + '.png'
    fi = os.path.join(out_path, f_name)
    res.append(f_name)
    top_right.save(fi)

    # 左下角的图片
    bottom_left = img.crop((0, new_height, new_width, height))
    f_name = str(uuid.uuid4()) + '.png'
    fi = os.path.join(out_path, f_name)
    res.append(f_name)
    bottom_left.save(fi)

    # 右下角的图片
    bottom_right = img.crop((new_width, new_height, width, height))
    f_name = str(uuid.uuid4()) + '.png'
    fi = os.path.join(out_path, f_name)
    res.append(f_name)
    bottom_right.save(fi)

    return res


# 图生图函数
def img2img(url,model_id, prompt_id, input_source_img, output_file):
    # JSON文件模板
    templet_file = r'../JSON/' + str(model_id) + '_' + str(prompt_id) + '.json'
    # 模型图片
    input_refer_img = '../Image/' + str(model_id) + '/' + str(prompt_id) + '.jpg'

    # 读取JSON文件
    with open(templet_file, 'r', encoding='utf-8') as file:
        payload = json.load(file)

    # 用户输入的图片
    source_img = encode_image(input_source_img)
    # 模型固定的参考图片
    refer_img = encode_image(input_refer_img)

    # 在ControlNet的第一个图片中使用用户输入的图片，用此图片的人脸进行换脸操作
    payload["alwayson_scripts"]["ControlNet"]["args"][0]["input_image"] = source_img
    # 图生图，风格参考图是输入
    payload["init_images"] = [refer_img]
    payload["alwayson_scripts"]["ControlNet"]["args"][1]["input_image"] = refer_img
    # 调用SD服务
    response = requests.post(url=f'{url}/sdapi/v1/img2img', json=payload)
    if response.status_code != 200:
        print(response.text)
        exit(0)
    r = response.json()
    # 保存图片
    save_encoded_image(response.json()['images'][0], output_file)