SD_LarkBot/utils/image_io.py

import json
import os
from utils.logger import logger

def save_image(img, message_id, img_key):
    # 确保目录存在，如果不存在则创建
    output_directory = f"data/input_images/{message_id}/{img_key}"
    os.makedirs(output_directory, exist_ok=True)

    # 生成图片文件路径
    image_path = os.path.join(output_directory, f"{img_key}.jpg")

    # 保存图片
    with open(image_path, 'wb') as f:
        f.write(img)

def get_img_key_from_input_images(message_id):
    input_directory = f"data/input_images/{message_id}"
    img_key = []
    for file in os.listdir(input_directory):
        if not file.endswith(".json"):
            img_key.append(file)
    return img_key

def create_image_path(message_id, img_key):
    # 确保目录存在，如果不存在则创建
    output_directory = f"data/input_images/{message_id}/{img_key}"
    os.makedirs(output_directory, exist_ok=True)

def get_img_file(img_path):
    with open(img_path, 'rb') as f:
        img = f.read()
    return img