230 lines
7.8 KiB
Rust
230 lines
7.8 KiB
Rust
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::{
|
||
GeneratedAssetSheetPersistInput, GeneratedAssetSheetPersistPrompt,
|
||
GeneratedAssetSheetPromptInput, apply_generated_asset_sheet_green_screen_alpha,
|
||
build_generated_asset_sheet_prompt, crop_generated_asset_sheet_view_edge_matte,
|
||
prepare_generated_asset_sheet_put_request, 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_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"))
|
||
);
|
||
}
|