Genarrative/server-rs/crates/platform-image/tests/generated_asset_sheets.rs

use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
use image::{DynamicImage, ImageFormat, Rgba, RgbaImage};
use platform_image::DownloadedImage;
use platform_image::generated_asset_sheets::{
    GeneratedAssetSheetAlphaOptions, GeneratedAssetSheetPersistInput,
    GeneratedAssetSheetPersistPrompt, GeneratedAssetSheetPromptInput,
    apply_generated_asset_sheet_alpha_with_options, apply_generated_asset_sheet_green_screen_alpha,
    build_generated_asset_sheet_prompt, crop_generated_asset_sheet_view_edge_matte,
    crop_generated_asset_sheet_view_edge_matte_with_options,
    prepare_generated_asset_sheet_put_request,
    remove_generated_asset_sheet_green_screen_background_bytes, slice_generated_asset_sheet,
    slice_generated_asset_sheet_two_items_per_row,
};
use platform_oss::LegacyAssetPrefix;

fn encode_image(image: RgbaImage) -> Vec<u8> {
    let mut encoded = std::io::Cursor::new(Vec::new());
    DynamicImage::ImageRgba8(image)
        .write_to(&mut encoded, ImageFormat::Png)
        .expect("image should encode");
    encoded.into_inner()
}

fn build_test_sheet(width: u32, height: u32) -> DownloadedImage {
    let mut sheet = RgbaImage::new(width, height);
    for row in 0..height / 100 {
        for col in 0..width / 100 {
            let row_u8 = row as u8;
            let col_u8 = col as u8;
            let color = Rgba([
                32u8.saturating_add(row_u8.saturating_mul(40)),
                24u8.saturating_add(col_u8.saturating_mul(36)),
                210u8.saturating_sub(row_u8.saturating_mul(30)),
                255,
            ]);
            for y in row * 100..(row + 1) * 100 {
                for x in col * 100..(col + 1) * 100 {
                    sheet.put_pixel(x, y, color);
                }
            }
        }
    }

    DownloadedImage {
        bytes: encode_image(sheet),
        mime_type: "image/png".to_string(),
        extension: "png".to_string(),
    }
}

#[test]
fn generated_asset_sheet_prompt_uses_default_rows_and_special_instruction() {
    let item_names = vec!["草莓".to_string(), "苹果".to_string()];

    let prompt = build_generated_asset_sheet_prompt(&GeneratedAssetSheetPromptInput {
        subject_text: "水果题材的抓大鹅 2D 物品素材",
        item_names: &item_names,
        grid_size: 5,
        item_name_prompt_template: None,
        special_prompt: None,
    })
    .expect("prompt should build");

    assert!(prompt.contains("5行*5列"));
    assert!(prompt.contains("第1行：草莓 的 5 个不同视图"));
    assert!(prompt.contains("第2行：苹果 的 5 个不同视图"));
    assert!(prompt.contains("每个物品生成 5 个不同视图"));
}

#[test]
fn generated_asset_sheet_prompt_allows_custom_row_template_and_special_prompt() {
    let item_names = vec!["草莓".to_string()];

    let prompt = build_generated_asset_sheet_prompt(&GeneratedAssetSheetPromptInput {
        subject_text: "水果题材的抓大鹅 2D 物品素材",
        item_names: &item_names,
        grid_size: 5,
        item_name_prompt_template: Some("第{row_index}行是 {item_name}，共 {view_count} 个视图"),
        special_prompt: Some("每个物品要生成五个不同视图：正面、左前、右前、俯视、背面。"),
    })
    .expect("prompt should build");

    assert!(prompt.contains("第1行是 草莓，共 5 个视图"));
    assert!(prompt.contains("每个物品要生成五个不同视图"));
}

#[test]
fn generated_asset_sheet_prompt_rejects_zero_grid_size() {
    let item_names = vec!["草莓".to_string()];

    let error = build_generated_asset_sheet_prompt(&GeneratedAssetSheetPromptInput {
        subject_text: "水果题材的抓大鹅 2D 物品素材",
        item_names: &item_names,
        grid_size: 0,
        item_name_prompt_template: None,
        special_prompt: None,
    })
    .expect_err("grid size 0 should be rejected");

    assert_eq!(error.provider(), "generated-asset-sheets");
}

#[test]
fn generated_asset_sheet_slices_by_requested_grid_size() {
    let item_names = vec!["樱桃".to_string(), "苹果".to_string()];
    let image = build_test_sheet(500, 500);

    let slices = slice_generated_asset_sheet(&image, &item_names, 5).expect("sheet should slice");

    assert_eq!(slices.len(), 2);
    assert_eq!(slices[0].len(), 5);
    assert_eq!(slices[1].len(), 5);
}

#[test]
fn generated_asset_sheet_two_items_per_row_slices_match3d_layout() {
    let item_names = vec![
        "苹果".to_string(),
        "香蕉".to_string(),
        "葡萄".to_string(),
        "草莓".to_string(),
    ];
    let image = build_test_sheet(1000, 1000);

    let slices = slice_generated_asset_sheet_two_items_per_row(&image, &item_names, 10, 5)
        .expect("sheet should slice");

    assert_eq!(slices.len(), 4);
    assert!(slices.iter().all(|views| views.len() == 5));
}

#[test]
fn generated_asset_sheet_green_screen_alpha_removes_green_background() {
    let mut sheet = RgbaImage::from_pixel(20, 20, Rgba([0, 255, 0, 255]));
    for y in 6..14 {
        for x in 6..14 {
            sheet.put_pixel(x, y, Rgba([220, 40, 40, 255]));
        }
    }

    let cleaned =
        apply_generated_asset_sheet_green_screen_alpha(DynamicImage::ImageRgba8(sheet)).to_rgba8();

    assert_eq!(cleaned.get_pixel(0, 0).0[3], 0);
    assert_eq!(cleaned.get_pixel(10, 10).0[3], 255);
}

#[test]
fn generated_asset_sheet_green_screen_background_bytes_returns_png() {
    let mut sheet = RgbaImage::from_pixel(3, 3, Rgba([0, 255, 0, 255]));
    sheet.put_pixel(1, 1, Rgba([240, 40, 40, 255]));

    let output = remove_generated_asset_sheet_green_screen_background_bytes(
        encode_image(sheet).as_slice(),
        GeneratedAssetSheetAlphaOptions {
            remove_muted_green_screen_background: true,
            detect_internal_holes: true,
            internal_hole_min_pixels: 16,
            ..GeneratedAssetSheetAlphaOptions::green_screen()
        },
    )
    .expect("green screen bytes should process");
    let image = image::load_from_memory(output.as_slice())
        .expect("output PNG should decode")
        .to_rgba8();

    assert_eq!(image.get_pixel(0, 0).0[3], 0);
    assert_eq!(image.get_pixel(1, 1).0[3], 255);
}

#[test]
fn generated_asset_sheet_muted_green_alpha_requires_explicit_option() {
    let mut sheet = RgbaImage::from_pixel(20, 20, Rgba([59, 131, 67, 255]));
    for y in 7..13 {
        for x in 7..13 {
            sheet.put_pixel(x, y, Rgba([220, 40, 40, 255]));
        }
    }

    let default_cleaned =
        apply_generated_asset_sheet_green_screen_alpha(DynamicImage::ImageRgba8(sheet.clone()))
            .to_rgba8();
    let muted_cleaned = apply_generated_asset_sheet_alpha_with_options(
        DynamicImage::ImageRgba8(sheet),
        GeneratedAssetSheetAlphaOptions {
            remove_muted_green_screen_background: true,
            ..GeneratedAssetSheetAlphaOptions::green_screen()
        },
    )
    .to_rgba8();

    assert!(
        default_cleaned.get_pixel(0, 0).0[3] > 0,
        "默认绿幕图集清理不应把暗绿背景直接清透明"
    );
    assert_eq!(muted_cleaned.get_pixel(0, 0).0[3], 0);
    assert_eq!(muted_cleaned.get_pixel(10, 10).0[3], 255);
}

#[test]
fn generated_asset_sheet_magenta_key_preserves_green_white_and_disconnected_key_subject() {
    let mut sheet = RgbaImage::from_pixel(28, 28, Rgba([255, 0, 255, 255]));
    for y in 6..22 {
        for x in 6..14 {
            sheet.put_pixel(x, y, Rgba([64, 188, 74, 255]));
        }
    }
    for y in 6..22 {
        for x in 14..22 {
            sheet.put_pixel(x, y, Rgba([244, 244, 236, 255]));
        }
    }
    for y in 12..16 {
        for x in 12..16 {
            sheet.put_pixel(x, y, Rgba([255, 0, 255, 255]));
        }
    }

    let cleaned = apply_generated_asset_sheet_alpha_with_options(
        DynamicImage::ImageRgba8(sheet),
        GeneratedAssetSheetAlphaOptions::jump_hop_magenta_screen(),
    )
    .to_rgba8();

    assert_eq!(cleaned.get_pixel(0, 0).0[3], 0);
    assert_eq!(cleaned.get_pixel(8, 8).0[3], 255);
    assert_eq!(cleaned.get_pixel(18, 8).0[3], 255);
    assert_eq!(
        cleaned.get_pixel(13, 13).0[3],
        255,
        "非边缘连通的 key 色像素不应被当成背景清掉"
    );
}

#[test]
fn generated_asset_sheet_magenta_edge_matte_does_not_remove_white_subject() {
    let mut sheet = RgbaImage::from_pixel(24, 24, Rgba([0, 0, 0, 0]));
    for y in 2..22 {
        for x in 2..22 {
            sheet.put_pixel(x, y, Rgba([246, 246, 240, 255]));
        }
    }
    for y in 0..24 {
        sheet.put_pixel(0, y, Rgba([255, 0, 255, 255]));
        sheet.put_pixel(23, y, Rgba([255, 0, 255, 255]));
    }

    let cleaned = crop_generated_asset_sheet_view_edge_matte_with_options(
        DynamicImage::ImageRgba8(sheet),
        GeneratedAssetSheetAlphaOptions::jump_hop_magenta_screen(),
    )
    .to_rgba8();

    assert_eq!(cleaned.get_pixel(1, 1).0[3], 255);
    assert!(
        cleaned
            .pixels()
            .any(|pixel| pixel.0 == [246, 246, 240, 255])
    );
}

#[test]
fn generated_asset_sheet_magenta_alpha_defringes_pink_halo() {
    let mut sheet = RgbaImage::from_pixel(24, 24, Rgba([255, 0, 255, 255]));
    for y in 7..17 {
        for x in 7..17 {
            sheet.put_pixel(x, y, Rgba([198, 170, 120, 255]));
        }
    }
    for y in 6..18 {
        sheet.put_pixel(6, y, Rgba([226, 26, 218, 220]));
        sheet.put_pixel(17, y, Rgba([226, 26, 218, 220]));
    }
    for x in 6..18 {
        sheet.put_pixel(x, 6, Rgba([226, 26, 218, 220]));
        sheet.put_pixel(x, 17, Rgba([226, 26, 218, 220]));
    }

    let cleaned = apply_generated_asset_sheet_alpha_with_options(
        DynamicImage::ImageRgba8(sheet),
        GeneratedAssetSheetAlphaOptions::jump_hop_magenta_screen(),
    )
    .to_rgba8();
    let edge = cleaned.get_pixel(6, 12).0;

    assert_eq!(cleaned.get_pixel(0, 0).0[3], 0);
    assert_eq!(cleaned.get_pixel(12, 12).0, [198, 170, 120, 255]);
    if edge[3] > 0 {
        assert!(
            edge[0].saturating_sub(edge[1]) <= 76,
            "红色 key 通道残留过强：{edge:?}"
        );
        assert!(
            edge[2].saturating_sub(edge[1]) <= 76,
            "蓝色 key 通道残留过强：{edge:?}"
        );
    }
}

#[test]
fn generated_asset_sheet_magenta_edge_matte_defringes_bottom_shadow() {
    let mut sheet = RgbaImage::from_pixel(32, 32, Rgba([0, 0, 0, 0]));
    for y in 8..18 {
        for x in 10..22 {
            sheet.put_pixel(x, y, Rgba([202, 176, 126, 255]));
        }
    }
    for y in 18..22 {
        for x in 9..23 {
            sheet.put_pixel(x, y, Rgba([224, 30, 220, 186]));
        }
    }

    let cleaned = crop_generated_asset_sheet_view_edge_matte_with_options(
        DynamicImage::ImageRgba8(sheet),
        GeneratedAssetSheetAlphaOptions::jump_hop_magenta_screen(),
    )
    .to_rgba8();

    assert!(
        cleaned
            .pixels()
            .any(|pixel| pixel.0 == [202, 176, 126, 255])
    );
    assert!(
        !cleaned.pixels().any(|pixel| {
            let [red, green, blue, alpha] = pixel.0;
            alpha > 0 && red > 200 && blue > 200 && green < 96
        }),
        "底部洋红残影应被删除或去彩边"
    );
}

#[test]
fn generated_asset_sheet_view_edge_matte_trims_transparent_border() {
    let mut sheet = RgbaImage::from_pixel(20, 20, Rgba([0, 0, 0, 0]));
    for y in 4..16 {
        for x in 4..16 {
            sheet.put_pixel(x, y, Rgba([220, 40, 40, 255]));
        }
    }

    let cropped =
        crop_generated_asset_sheet_view_edge_matte(DynamicImage::ImageRgba8(sheet)).to_rgba8();

    assert_eq!(cropped.width(), 12);
    assert_eq!(cropped.height(), 12);
    assert_eq!(cropped.get_pixel(0, 0).0[3], 255);
}

#[test]
fn generated_asset_sheet_prepare_put_request_packs_prompt_metadata() {
    let request = prepare_generated_asset_sheet_put_request(GeneratedAssetSheetPersistInput {
        prefix: LegacyAssetPrefix::Match3DAssets,
        owner_user_id: "user-1".to_string(),
        session_id: "session-1".to_string(),
        profile_id: "profile-1".to_string(),
        path_segments: vec!["items".to_string(), "view".to_string()],
        file_name: "view-01.png".to_string(),
        content_type: "image/png".to_string(),
        bytes: b"sheet-bytes".to_vec(),
        asset_kind: "match3d_item_image_view".to_string(),
        source_job_id: Some("task-1".to_string()),
        generated_at_micros: 123,
        grid_size: 5,
        row_index: 1,
        view_index: 2,
        prompt: GeneratedAssetSheetPersistPrompt {
            sheet_prompt: Some("sheet prompt".to_string()),
            item_name_prompt: Some("item prompt".to_string()),
            special_prompt: Some("special prompt".to_string()),
        },
    })
    .expect("request should prepare");

    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-prompt-encoding"),
        Some(&"utf8-base64".to_string())
    );
    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-grid-size"),
        Some(&"5".to_string())
    );
    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-row-index"),
        Some(&"1".to_string())
    );
    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-view-index"),
        Some(&"2".to_string())
    );
    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-prompt-b64"),
        Some(&BASE64_STANDARD.encode("sheet prompt"))
    );
    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-item-name-prompt-b64"),
        Some(&BASE64_STANDARD.encode("item prompt"))
    );
    assert_eq!(
        request
            .metadata
            .get("x-oss-meta-generated-asset-sheet-special-prompt-b64"),
        Some(&BASE64_STANDARD.encode("special prompt"))
    );
}