Genarrative/server-rs/crates/api-server/src/editor_green_screen.rs

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; 12] = [
    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,
    },
    // 中明度低饱和柔雾绿：填补安全弧的绿色段。浅粉彩绿（淡薄荷绿）会被露肤角色的皮肤规则剔除，
    // 故绿色改用中明度以过 RGB 分离（Rule 3）；色相远离皮肤橙调，过色调投影（Rule 2）。
    EditorScreenBackgroundColor {
        label: "灰竹绿",
        hex: "#A0BBA0",
        red: 160,
        green: 187,
        blue: 160,
    },
];

pub(crate) fn default_editor_screen_background_color() -> EditorScreenBackgroundColor {
    EDITOR_SCREEN_BACKGROUND_COLORS[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 =
    "角色主体及其服饰、道具的颜色必须与背景色明显区分，不得带与背景色相同或相近的描边、投影或反光";

fn editor_screen_background_color_prompt(color: EditorScreenBackgroundColor) -> String {
    format!(
        "单一纯色背景 {} {} / RGB({},{},{})",
        color.label, color.hex, color.red, color.green, color.blue
    )
}

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 { .. }
        | GeneratedAssetSheetError::RawConnectedComponentLimitExceeded { .. }
        | GeneratedAssetSheetError::OutputSliceLimitExceeded { .. }
        | GeneratedAssetSheetError::TotalCropPixelLimitExceeded { .. }
        | GeneratedAssetSheetError::MergeCandidateLimitExceeded { .. } => 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());
    }
}