app.py

import gradio as gr
from PIL import Image
import io
import zipfile
import random

def random_black_or_white():
    """返回 RGBA (黑 或 白)，50% 概率"""
    return (0, 0, 0, 255) if random.random() < 0.5 else (255, 255, 255, 255)

def random_non_black_white():
    """
    返回 (r,g,b,a)，且不是纯黑(0,0,0)也不是纯白(255,255,255)。
    用于最后不足4张时的空格颜色填充。
    """
    while True:
        r = random.randint(0, 255)
        g = random.randint(0, 255)
        b = random.randint(0, 255)
        if not (r == g == b == 0 or r == g == b == 255):
            return (r, g, b, 255)

def limit_2048(img: Image.Image):
    """若宽或高大于2048，则等比例缩小到不超过2048。"""
    w, h = img.size
    if w > 2048 or h > 2048:
        scale = min(2048 / w, 2048 / h)
        new_w = int(w * scale)
        new_h = int(h * scale)
        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
    return img

def resize_to_64_multiple(img: Image.Image):
    """
    将图像就近缩放到(w', h')，其中 w'、h'均为64倍数（至少64）。
    空余区域随机黑/白填充，保留透明度。
    """
    w, h = img.size
    w64 = max(64, round(w / 64) * 64)
    h64 = max(64, round(h / 64) * 64)
    
    scale = min(w64 / w, h64 / h)
    nw = int(w * scale)
    nh = int(h * scale)
    
    bg_color = random_black_or_white()  
    background = Image.new("RGBA", (w64, h64), bg_color)
    scaled = img.resize((nw, nh), Image.Resampling.LANCZOS)
    ox = (w64 - nw) // 2
    oy = (h64 - nh) // 2
    background.paste(scaled, (ox, oy), scaled)
    return background

def make_collage_2x2(images_4):
    """传入4张同尺寸RGBA，以2×2拼接，并限制在2048以内。"""
    w, h = images_4[0].size
    collage = Image.new("RGBA", (2*w, 2*h), (0, 0, 0, 255))
    collage.paste(images_4[0], (0, 0), images_4[0])
    collage.paste(images_4[1], (w, 0), images_4[1])
    collage.paste(images_4[2], (0, h), images_4[2])
    collage.paste(images_4[3], (w, h), images_4[3])
    return limit_2048(collage)

def make_collage_leftover(images_leftover):
    """
    针对不足4张(1~3)的情况，随机布局(1x1,1x2,2x1,2x2等)容纳这些图片，
    其余空格以“非黑非白”的纯色填充，并最终限制2048以内。
    """
    n = len(images_leftover)
    if n < 1 or n > 3:
        return None
    
    # 先做“64倍数随机黑/白填充”缩放
    resized_list = [resize_to_64_multiple(img) for img in images_leftover]
    # 找最大宽、高，统一尺寸
    max_w = max(im.size[0] for im in resized_list)
    max_h = max(im.size[1] for im in resized_list)
    
    # 再居中贴到背景 (max_w, max_h) 保证每张图一致
    uniformed = []
    for rimg in resized_list:
        w, h = rimg.size
        if (w, h) == (max_w, max_h):
            uniformed.append(rimg)
        else:
            # 用本身的背景色(左上像素)来填充
            bg_color = rimg.getpixel((0,0))
            bg = Image.new("RGBA", (max_w, max_h), bg_color)
            offx = (max_w - w) // 2
            offy = (max_h - h) // 2
            bg.paste(rimg, (offx, offy), rimg)
            uniformed.append(bg)
    
    # 选定随机布局
    #   n=1: (1x1),(1x2),(2x1),(2x2)
    #   n=2: (1x2),(2x1),(2x2)
    #   n=3: (2x2)
    possible_layouts = []
    if n == 1:
        possible_layouts = [(1,1), (1,2), (2,1), (2,2)]
    elif n == 2:
        possible_layouts = [(1,2), (2,1), (2,2)]
    else:  # n == 3
        possible_layouts = [(2,2)]
    
    rows, cols = random.choice(possible_layouts)
    
    big_w = cols * max_w
    big_h = rows * max_h
    collage = Image.new("RGBA", (big_w, big_h), (0,0,0,255))
    
    # 网格坐标
    cells = [(r, c) for r in range(rows) for c in range(cols)]
    random.shuffle(cells)
    
    # 把uniformed中的图贴到前n个格子
    for i, img_ in enumerate(uniformed):
        r, c = cells[i]
        offset_x = c * max_w
        offset_y = r * max_h
        collage.paste(img_, (offset_x, offset_y), img_)
    
    # 剩余格子(空格)用“非黑非白”的随机色填充
    leftover_cells = cells[n:]
    for (r, c) in leftover_cells:
        color_ = random_non_black_white()
        rect = Image.new("RGBA", (max_w, max_h), color_)
        collage.paste(rect, (c*max_w, r*max_h), rect)
    
    return limit_2048(collage)

def process_images(uploaded_files):
    """
    1) 分成若干4张组 => each 2x2；
    2) 对最后剩余(1~3张)，make_collage_leftover 处理。
    3) 返回所有结果图
    """
    pil_images = []
    for f in uploaded_files:
        if f is not None:
            # 保留透明
            img = Image.open(f.name).convert("RGBA")
            pil_images.append(img)
    
    results = []
    total = len(pil_images)
    groups_4 = total // 4
    leftover = total % 4
    
    idx = 0
    # 先拼满4张的组
    for _ in range(groups_4):
        group_4 = pil_images[idx:idx+4]
        idx += 4
        # 每张先resize
        resized_4 = [resize_to_64_multiple(im) for im in group_4]
        # 再统一max_w, max_h
        max_w = max(im.size[0] for im in resized_4)
        max_h = max(im.size[1] for im in resized_4)
        
        final_4 = []
        for rimg in resized_4:
            w, h = rimg.size
            if (w, h) == (max_w, max_h):
                final_4.append(rimg)
            else:
                # 补背景居中
                bg_color = rimg.getpixel((0,0))
                bg = Image.new("RGBA", (max_w, max_h), bg_color)
                offx = (max_w - w)//2
                offy = (max_h - h)//2
                bg.paste(rimg, (offx, offy), rimg)
                final_4.append(bg)
        
        collage_2x2 = make_collage_2x2(final_4)
        results.append(collage_2x2)
    
    # 再拼 leftover 1~3张
    if leftover > 0:
        leftover_images = pil_images[idx:]
        collage_left = make_collage_leftover(leftover_images)
        if collage_left is not None:
            results.append(collage_left)
    
    return results

def make_zip(uploaded_files):
    """把所有拼接结果打包成zip并返回给Gradio的File组件。"""
    collages = process_images(uploaded_files)
    # 若无生成任何拼图
    if not collages:
        # 返回 None 说明无法下载；会显示“无可下载内容”提示
        return None
    
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for i, img in enumerate(collages, start=1):
            img_bytes = io.BytesIO()
            img.save(img_bytes, format="PNG")
            img_bytes.seek(0)
            zf.writestr(f"collage_{i}.png", img_bytes.read())
    buf.seek(0)
    return buf

def on_zip_click(files):
    """
    用来返回 (zip_file_obj, message_str) 两个输出：
    1) zip_file_obj 要么是zip流，要么是None；
    2) message_str 用于提示结果或错误。
    """
    z = make_zip(files)
    if z is None:
        return (None, "无可下载内容 - 请检查是否上传了图片或剩余不足4张且无法拼图")
    else:
        return (z, "打包完成！可点击上方链接下载。")

with gr.Blocks() as demo:
    gr.Markdown("## 2×2 拼接小工具（支持最后不足4张、随机填充、保留透明）")
    with gr.Row():
        with gr.Column():
            file_input = gr.Files(label="上传多张图片", file_types=["image"])
            preview_btn = gr.Button("生成预览")
            zip_btn = gr.Button("打包下载 ZIP")
        with gr.Column():
            # 不使用 .style() 以兼容老Gradio
            gallery_out = gr.Gallery(label="拼接结果预览", columns=2)
            # 一开始就 visible=True，这样点击按钮后能马上显示下载链接
            zip_file_out = gr.File(label="点击下载打包结果", visible=True, interactive=False)
            msg_output = gr.Textbox(label="处理信息", interactive=False)
    
    # 生成预览
    preview_btn.click(
        fn=process_images,
        inputs=[file_input],
        outputs=[gallery_out]
    )
    
    # 打包下载ZIP，额外给一个文本提示
    zip_btn.click(
        fn=on_zip_click,
        inputs=[file_input],
        outputs=[zip_file_out, msg_output]
    )

demo.launch()