lhk229
308d440816
新抠图算法（轻量版）上线，生图背景色同步改为多颜色+自动决策 --------- Co-authored-by: kdletters <kdletters@qq.com> Reviewed-on: https://git.genarrative.world/git/GenarrativeAI/Genarrative/pulls/73 Reviewed-by: kdletters <kdletters@qq.com> Co-authored-by: Linghong <ink29535@proton.me> Co-committed-by: Linghong <ink29535@proton.me>
284 lines
9.2 KiB
Rust
284 lines
9.2 KiB
Rust
use axum::http::StatusCode;
|
||
use platform_image::generated_asset_sheets::{
|
||
GeneratedAssetSheetAlphaOptions, GeneratedAssetSheetError, GeneratedAssetSheetKeyColor,
|
||
remove_generated_asset_sheet_green_screen_background_bytes,
|
||
};
|
||
use serde_json::json;
|
||
|
||
use crate::{http_error::AppError, openai_image_generation::DownloadedOpenAiImage};
|
||
|
||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||
pub(crate) struct EditorScreenBackgroundColor {
|
||
pub(crate) label: &'static str,
|
||
pub(crate) hex: &'static str,
|
||
pub(crate) red: u8,
|
||
pub(crate) green: u8,
|
||
pub(crate) blue: u8,
|
||
}
|
||
|
||
pub(crate) const EDITOR_SCREEN_BACKGROUND_COLORS: [EditorScreenBackgroundColor; 11] = [
|
||
EditorScreenBackgroundColor {
|
||
label: "浅雾蓝",
|
||
hex: "#CFEFFF",
|
||
red: 207,
|
||
green: 239,
|
||
blue: 255,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "浅钢蓝",
|
||
hex: "#B0C2E0",
|
||
red: 176,
|
||
green: 194,
|
||
blue: 224,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "暖浅桃色",
|
||
hex: "#FFD6C2",
|
||
red: 255,
|
||
green: 214,
|
||
blue: 194,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "淡薰衣草紫",
|
||
hex: "#E6D8FF",
|
||
red: 230,
|
||
green: 216,
|
||
blue: 255,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "浅粉灰",
|
||
hex: "#F4D8E8",
|
||
red: 244,
|
||
green: 216,
|
||
blue: 232,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "中度天蓝",
|
||
hex: "#7FB3FF",
|
||
red: 127,
|
||
green: 179,
|
||
blue: 255,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "浅柠黄",
|
||
hex: "#FFF2A8",
|
||
red: 255,
|
||
green: 242,
|
||
blue: 168,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "淡薄荷绿",
|
||
hex: "#CFFFE1",
|
||
red: 207,
|
||
green: 255,
|
||
blue: 225,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "浅中性灰",
|
||
hex: "#D8DEE8",
|
||
red: 216,
|
||
green: 222,
|
||
blue: 232,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "淡灰紫",
|
||
hex: "#D8D2E8",
|
||
red: 216,
|
||
green: 210,
|
||
blue: 232,
|
||
},
|
||
EditorScreenBackgroundColor {
|
||
label: "高对比浅青",
|
||
hex: "#A8F7F0",
|
||
red: 168,
|
||
green: 247,
|
||
blue: 240,
|
||
},
|
||
];
|
||
|
||
pub(crate) fn default_editor_screen_background_color() -> EditorScreenBackgroundColor {
|
||
EDITOR_SCREEN_BACKGROUND_COLORS[0]
|
||
}
|
||
|
||
pub(crate) fn legacy_editor_green_screen_background_color() -> EditorScreenBackgroundColor {
|
||
EditorScreenBackgroundColor {
|
||
label: "纯绿色",
|
||
hex: "#00FF00",
|
||
red: 0,
|
||
green: 255,
|
||
blue: 0,
|
||
}
|
||
}
|
||
|
||
pub(crate) fn parse_editor_screen_background_color(
|
||
value: Option<&str>,
|
||
) -> Result<EditorScreenBackgroundColor, AppError> {
|
||
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
|
||
return Ok(default_editor_screen_background_color());
|
||
};
|
||
let normalized_value = value.to_ascii_uppercase();
|
||
EDITOR_SCREEN_BACKGROUND_COLORS
|
||
.iter()
|
||
.copied()
|
||
.find(|option| option.hex == normalized_value)
|
||
.ok_or_else(|| {
|
||
AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({
|
||
"provider": "editor-screen-background",
|
||
"message": "screenColor 必须是编辑器支持的纯色背景色。",
|
||
"allowed": EDITOR_SCREEN_BACKGROUND_COLORS
|
||
.iter()
|
||
.map(|option| option.hex)
|
||
.collect::<Vec<_>>(),
|
||
}))
|
||
})
|
||
}
|
||
|
||
pub(crate) const EDITOR_GREEN_SCREEN_BACKGROUND_GUARDRAILS: &str =
|
||
"纯色背景必须平整无纹理、无渐变、无阴影、无地面、无环境、无道具";
|
||
pub(crate) const EDITOR_GREEN_SCREEN_ASSET_GUARDRAILS: &str =
|
||
"素材自身不要出现与背景色相同或相近的描边、底板、投影或反光";
|
||
pub(crate) const EDITOR_GREEN_SCREEN_CHARACTER_GUARDRAILS: &str =
|
||
"角色主体不得带与背景色相同或相近的描边、投影或反光";
|
||
pub(crate) const LEGACY_EDITOR_GREEN_SCREEN_CHARACTER_PROMPT_CLAUSE: &str = "背景固定为单一纯绿色 #00FF00 / RGB(0,255,0) 绿幕，只作为抠像底色；绿幕背景必须平整无纹理、无渐变、无阴影、无地面、无环境、无道具；角色主体不得带绿色描边、绿色投影或绿色反光";
|
||
|
||
fn editor_screen_background_color_prompt(color: EditorScreenBackgroundColor) -> String {
|
||
if color.hex == "#00FF00" {
|
||
return "单一纯绿色 #00FF00 / RGB(0,255,0) 绿幕".to_string();
|
||
}
|
||
format!(
|
||
"单一纯色背景 {} {} / RGB({},{},{})",
|
||
color.label, color.hex, color.red, color.green, color.blue
|
||
)
|
||
}
|
||
|
||
pub(crate) fn editor_green_screen_asset_prompt_clause(
|
||
color: EditorScreenBackgroundColor,
|
||
) -> String {
|
||
format!(
|
||
"背景必须是{}，且{}，方便扣除背景；{}",
|
||
editor_screen_background_color_prompt(color),
|
||
EDITOR_GREEN_SCREEN_BACKGROUND_GUARDRAILS,
|
||
EDITOR_GREEN_SCREEN_ASSET_GUARDRAILS
|
||
)
|
||
}
|
||
|
||
pub(crate) fn editor_green_screen_character_prompt_clause(
|
||
color: EditorScreenBackgroundColor,
|
||
) -> String {
|
||
format!(
|
||
"背景固定为{}，只作为抠像底色；{}；{}",
|
||
editor_screen_background_color_prompt(color),
|
||
EDITOR_GREEN_SCREEN_BACKGROUND_GUARDRAILS,
|
||
EDITOR_GREEN_SCREEN_CHARACTER_GUARDRAILS
|
||
)
|
||
}
|
||
|
||
pub(crate) fn editor_ui_design_asset_extraction_prompt(
|
||
color: EditorScreenBackgroundColor,
|
||
) -> String {
|
||
format!(
|
||
"仅提取被红色框框选的素材并整理成spritesheet，图集背景必须使用{}。{}，方便后续扣除背景；{}。",
|
||
editor_screen_background_color_prompt(color),
|
||
EDITOR_GREEN_SCREEN_BACKGROUND_GUARDRAILS,
|
||
EDITOR_GREEN_SCREEN_ASSET_GUARDRAILS
|
||
)
|
||
}
|
||
|
||
pub(crate) fn remove_editor_generated_green_screen_background(
|
||
image: &DownloadedOpenAiImage,
|
||
color: EditorScreenBackgroundColor,
|
||
) -> Result<DownloadedOpenAiImage, AppError> {
|
||
let key_color = GeneratedAssetSheetKeyColor {
|
||
red: color.red,
|
||
green: color.green,
|
||
blue: color.blue,
|
||
};
|
||
let bytes = remove_generated_asset_sheet_green_screen_background_bytes(
|
||
image.bytes.as_slice(),
|
||
GeneratedAssetSheetAlphaOptions {
|
||
key_color,
|
||
remove_near_white_background: key_color.is_green_screen(),
|
||
remove_disconnected_hard_key_background: true,
|
||
remove_muted_green_screen_background: key_color.is_green_screen(),
|
||
detect_internal_holes: true,
|
||
internal_hole_min_pixels: 16,
|
||
},
|
||
)
|
||
.map_err(map_editor_green_screen_error)?;
|
||
|
||
Ok(DownloadedOpenAiImage {
|
||
bytes,
|
||
mime_type: "image/png".to_string(),
|
||
extension: "png".to_string(),
|
||
})
|
||
}
|
||
|
||
fn map_editor_green_screen_error(error: GeneratedAssetSheetError) -> AppError {
|
||
let status = match error {
|
||
GeneratedAssetSheetError::DecodeImage { .. } => StatusCode::BAD_GATEWAY,
|
||
GeneratedAssetSheetError::InvalidRequest { .. } => StatusCode::BAD_REQUEST,
|
||
GeneratedAssetSheetError::EncodeImage { .. }
|
||
| GeneratedAssetSheetError::BuildHttpClient { .. }
|
||
| GeneratedAssetSheetError::Oss(_) => StatusCode::INTERNAL_SERVER_ERROR,
|
||
};
|
||
AppError::from_status(status).with_details(json!({
|
||
"provider": "editor-green-screen",
|
||
"message": error.message(),
|
||
}))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use std::io::Cursor;
|
||
|
||
use image::{DynamicImage, ImageFormat, Rgba, RgbaImage};
|
||
|
||
#[test]
|
||
fn editor_green_screen_postprocess_turns_selected_background_transparent() {
|
||
let mut source = RgbaImage::from_pixel(3, 3, Rgba([207, 239, 255, 255]));
|
||
source.put_pixel(1, 1, Rgba([240, 40, 40, 255]));
|
||
let mut bytes = Vec::new();
|
||
DynamicImage::ImageRgba8(source)
|
||
.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
|
||
.expect("test PNG should encode");
|
||
|
||
let output = remove_editor_generated_green_screen_background(
|
||
&DownloadedOpenAiImage {
|
||
bytes,
|
||
mime_type: "image/png".to_string(),
|
||
extension: "png".to_string(),
|
||
},
|
||
default_editor_screen_background_color(),
|
||
)
|
||
.expect("screen background should be removed");
|
||
let image = image::load_from_memory(output.bytes.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 editor_screen_background_color_parser_uses_default_and_rejects_unknown() {
|
||
assert_eq!(
|
||
parse_editor_screen_background_color(None).expect("default should parse"),
|
||
default_editor_screen_background_color()
|
||
);
|
||
assert_eq!(
|
||
parse_editor_screen_background_color(Some("#ffd6c2"))
|
||
.expect("known lowercase hex should parse")
|
||
.label,
|
||
"暖浅桃色"
|
||
);
|
||
assert_eq!(
|
||
parse_editor_screen_background_color(Some("#a8f7f0"))
|
||
.expect("new high-contrast cyan should parse")
|
||
.label,
|
||
"高对比浅青"
|
||
);
|
||
assert!(parse_editor_screen_background_color(Some("#00FF00")).is_err());
|
||
}
|
||
}
|