app.py

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

def random_black_or_white():
    """返回 (r,g,b,a) = 黑 or 白，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)，其中 (r,g,b) != (0,0,0) 且 != (255,255,255)。
    用于填充最后拼图时剩余空格，使之是“非黑非白”的纯色。
    """
    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 resize_to_64_multiple(img: Image.Image):
    """
    将单张 RGBA 图片就近缩放到 (w', h')，其中 w'、h' 均为 64 的倍数（至少64）。
    背景随机填充黑/白。原图保持居中（若有空余，四周即背景色；透明度保留）。
    """
    w, h = img.size
    # 找到最接近的 64 倍数（至少 64）
    w64 = max(64, round(w / 64) * 64)
    h64 = max(64, round(h / 64) * 64)

    # 算缩放比：保证原图能放入 w64*h64
    scale = min(w64 / w, h64 / h)
    new_w = int(w * scale)
    new_h = int(h * scale)

    # 随机黑或白做背景
    bg_color = random_black_or_white()
    background = Image.new("RGBA", (w64, h64), bg_color)
    
    # 缩放
    scaled = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
    
    # 居中贴到背景
    offset_x = (w64 - new_w) // 2
    offset_y = (h64 - new_h) // 2
    # 注意第三个参数 scaled：保持其透明度
    background.paste(scaled, (offset_x, offset_y), scaled)
    return background

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 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):
    """
    对剩余的 1~3 张图做“兼容性拼接”。
    1) 首先统一尺寸（同宽高）；
    2) 随机选择 (rows, cols) 布局（如 1x1/1x2/2x1/2x2 等）能容纳所有图；
    3) 将图随机放到网格单元，剩余格子用“非黑非白”的纯色填充；
    4) 最后若超出 2048，则缩小。
    """
    n = len(images_leftover)
    if n < 1 or n > 3:
        return None  # 保险
    
    # 统一每张图片的尺寸：按64倍数策略后，找出最大 w,h
    resized_list = []
    max_w = 0
    max_h = 0
    for img in images_leftover:
        rimg = resize_to_64_multiple(img)
        resized_list.append(rimg)
        if rimg.size[0] > max_w:
            max_w = rimg.size[0]
        if rimg.size[1] > max_h:
            max_h = rimg.size[1]
    
    # 再次处理，使它们都成为 (max_w, max_h)
    # （若小于max_w或max_h，就在背景再“居中贴图”）
    uniformed = []
    for rimg in resized_list:
        w, h = rimg.size
        if w == max_w and h == max_h:
            uniformed.append(rimg)
        else:
            bg = Image.new("RGBA", (max_w, max_h), rimg.getpixel((0,0)))
            # 取自身背景色（黑或白）进行填充，这样保持一致
            offx = (max_w - w)//2
            offy = (max_h - h)//2
            bg.paste(rimg, (offx, offy), rimg)
            uniformed.append(bg)
    
    # 现在 uniformed 每张都是 (max_w, max_h)，按 n ∈ [1,2,3]
    # 决定随机布局
    possible_layouts = []
    if n == 1:
        # 可以放在 (1x1), (1x2), (2x1), (2x2)
        possible_layouts = [(1,1), (1,2), (2,1), (2,2)]
    elif n == 2:
        # (1x2), (2x1), (2x2)
        possible_layouts = [(1,2), (2,1), (2,2)]
    else:  # n == 3
        # (2x2) 
        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 = []
    for r in range(rows):
        for c in range(cols):
            cells.append((r,c))
    
    random.shuffle(cells)  # 打乱单元格顺序
    # 将 n 张图放前 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:
        fill_col = random_non_black_white()
        rect = Image.new("RGBA", (max_w, max_h), fill_col)
        offset_x = c * max_w
        offset_y = r * max_h
        collage.paste(rect, (offset_x, offset_y), rect)
    
    return limit_2048(collage)

def process_images(uploaded_files):
    """
    1) 把文件读成 RGBA；
    2) 分成若干 4 张组 => each 2×2 拼接；
    3) 若最后有 1~3 张剩余，则调用 make_collage_leftover();
    4) 返回多张结果图(列表)
    """
    pil_images = []
    for f in uploaded_files:
        if f is not None:
            img = Image.open(f.name).convert("RGBA")
            pil_images.append(img)
    
    results = []
    # 每 4 张一组
    full_groups = len(pil_images) // 4
    leftover_count = len(pil_images) % 4
    
    # 处理完整的 4 张组
    idx = 0
    for _ in range(full_groups):
        group_4 = pil_images[idx : idx+4]
        idx += 4
        
        # 先统一尺寸 => 找 max_w, max_h
        temp = [resize_to_64_multiple(im) for im in group_4]
        # 再次统一(可能有不同64倍数)
        max_w = max([im.size[0] for im in temp])
        max_h = max([im.size[1] for im in temp])
        
        uniformed = []
        for rimg in temp:
            w, h = rimg.size
            if w == max_w and h == max_h:
                uniformed.append(rimg)
            else:
                bg = Image.new("RGBA", (max_w, max_h), rimg.getpixel((0,0)))
                offx = (max_w - w)//2
                offy = (max_h - h)//2
                bg.paste(rimg, (offx, offy), rimg)
                uniformed.append(bg)
        
        # 2x2 拼接
        collage_4 = make_collage_2x2(uniformed)
        results.append(collage_4)
    
    # 处理剩余 1~3 张
    if leftover_count > 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 (PNG 格式)，若无结果则返回 None。"""
    collages = process_images(uploaded_files)
    if not collages:
        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

with gr.Blocks() as demo:
    gr.Markdown("## 图片 2×2 拼接小工具（含随机填充、兼容不足4张）")
    gr.Markdown(
        "1. 一次可上传多张图片，每 4 张为一组严格 2×2 拼接；\n"
        "2. 若最后不足 4 张 (1~3)，会随机选择网格大小 (1x1,1x2,2x1,2x2)，并随机分配位置；\n"
        "   剩余空格用“非黑非白”的随机颜色填充；\n"
        "3. 每张图先按 64 的倍数就近缩放，空余处随机黑/白 (50% 概率)；\n"
        "4. 拼出的图若任一边超 2048，则等比例缩小到不超 2048；\n"
        "5. 保留 PNG 透明度，背景填充只在超出区域；\n"
        "6. 生成结果可预览，也可打包下载成 ZIP。"
    )
    with gr.Row():
        with gr.Column():
            input_files = gr.Files(label="上传图片（可多选）", file_types=["image"])
            btn_preview = gr.Button("生成预览")
            btn_zip = gr.Button("打包下载 ZIP")
        with gr.Column():
            gallery = gr.Gallery(label="拼接结果预览", columns=2)
            zipfile_output = gr.File(label="下载拼接结果 ZIP", visible=False, interactive=False)
    
    btn_preview.click(fn=process_images, inputs=[input_files], outputs=[gallery])
    btn_zip.click(fn=make_zip, inputs=[input_files], outputs=[zipfile_output])

demo.launch()