KeCheng/BaiHu/Util/CommonUtil.py

import json
import base64
import os.path
import uuid
from datetime import datetime
import requests
# pip install pillow
from PIL import Image
import io
import urllib.parse
import urllib.request


# 是否启用代理
def set_http_proxy(proxy):
    if proxy == None:  # Use system default setting
        proxy_support = urllib.request.ProxyHandler()
    elif proxy == '':  # Don't use any proxy
        proxy_support = urllib.request.ProxyHandler({})
    else:  # Use proxy
        proxy_support = urllib.request.ProxyHandler({'http': '%s' % proxy, 'https': '%s' % proxy})
    opener = urllib.request.build_opener(proxy_support)
    urllib.request.install_opener(opener)


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


# 文生图，完成一个极小的图片创建，防止img_2_img暴显存
def txt_2_img_blank(webui_address):
    txt2img_url = 'http://' + webui_address + '/sdapi/v1/txt2img'
    data = {"prompt": "a dog", "negative_prompt": "", "seed": -1, "steps": 1, "width": 1, "height": 1, "cfg_scale": 1,
            "override_settings": {
                "sd_model_checkpoint": "majicMix sombre 麦橘唯美_v2.0.safetensors [796bc0cc14]",
                "sd_vae": "Automatic"
            }}
    submit_post(txt2img_url, data)


# 将图片一切为四
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