diff --git a/Cargo.lock b/Cargo.lock index 388a1ca..6dae1f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4079,7 +4079,6 @@ checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" name = "vibetty" version = "0.3.0" dependencies = [ - "ab_glyph", "anyhow", "axum", "axum-extra", @@ -4092,7 +4091,6 @@ dependencies = [ "futures-util", "hanconv", "image", - "imageproc", "log", "pty-process", "ratatui", @@ -4102,16 +4100,28 @@ dependencies = [ "serde", "serde_json", "strip-ansi-escapes", - "thiserror 2.0.18", - "tiny-skia", "tokio", "tower-http 0.5.2", "tui-term", "uuid", + "vibetty-screenshot", "vt100", "wav_io", ] +[[package]] +name = "vibetty-screenshot" +version = "0.2.0" +source = "git+https://github.com/second-state/vibetty-screenshot?tag=v0.2.0#3f6a2e69da8db29253b9f619bd71b29865d563f3" +dependencies = [ + "ab_glyph", + "image", + "imageproc", + "thiserror 2.0.18", + "tiny-skia", + "vt100", +] + [[package]] name = "vt100" version = "0.16.2" diff --git a/Cargo.toml b/Cargo.toml index 86ce7b3..a40c31d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,8 +43,5 @@ tui-term = "0.3" vt100 = "0.16" # Screenshot / Image rendering -ab_glyph = "0.2" +vibetty-screenshot = { git = "https://github.com/second-state/vibetty-screenshot", tag = "v0.2.0" } image = "0.25" -imageproc = "0.26" -tiny-skia = "0.12" -thiserror = "2.0" diff --git a/src/bin/font_render_test.rs b/src/bin/font_render_test.rs deleted file mode 100644 index 3b94639..0000000 --- a/src/bin/font_render_test.rs +++ /dev/null @@ -1,32 +0,0 @@ -//! Test program for font rendering - -// Reuse the main screenshot module -#[path = "../screenshot/mod.rs"] -mod screenshot; - -fn main() -> Result<(), Box> { - println!("Testing font rendering..."); - - // Create a simple terminal screen with colored text - let mut parser = vt100::Parser::new(10, 60, 8096); - - parser.process(b"\x1b[31mRed Text\x1b[0m\r\n"); - parser.process(b"\x1b[32mGreen Text\x1b[0m\r\n"); - parser.process(b"\x1b[33mYellow Text\x1b[0m\r\n"); - parser.process(b"\x1b[34mBlue Text\x1b[0m\r\n"); - parser.process(b"\x1b[1mBold Text\x1b[0m\r\n"); - parser.process(b"Normal text with some content\r\n"); - parser.process(b"The quick brown fox jumps over the lazy dog\r\n"); - - // Save screenshot using the library function - let config = screenshot::ScreenshotConfig { - title: Some("Font Render Test".to_string()), - ..Default::default() - }; - - screenshot::save_screen_png(parser.screen(), "font_render_test.png", &config)?; - - println!("✓ Screenshot saved to font_render_test.png"); - - Ok(()) -} diff --git a/src/config.rs b/src/config.rs index db242a0..bfe6096 100644 --- a/src/config.rs +++ b/src/config.rs @@ -39,9 +39,20 @@ pub struct Args { /// Command to execute on PTY start (e.g., -- bash -l) #[arg(last = true)] pub command: Vec, + + /// Image format for screen rendering (png or jpeg) + #[arg(short = 'f', long, default_value = "jpeg", value_name = "FORMAT")] + pub image_format: String, } impl Args { + pub fn image_format(&self) -> crate::protocol::ImageFormat { + match self.image_format.to_lowercase().as_str() { + "jpeg" | "jpg" => crate::protocol::ImageFormat::Jpeg, + _ => crate::protocol::ImageFormat::Png, + } + } + pub fn asr_config(&self) -> AsrConfig { // 如果指定了配置文件，从文件读取 if let Some(path) = &self.asr_config_path { diff --git a/src/main.rs b/src/main.rs index 9cf4465..c053b2b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,7 +17,7 @@ mod terminal; mod ui; -pub mod screenshot; +pub use vibetty_screenshot as screenshot; use config::{Args, AsrConfig}; @@ -145,11 +145,14 @@ async fn main() { let (screenshot_tx, screenshot_rx) = tokio::sync::mpsc::channel(4); + let image_format = args.image_format(); + let state = ws::AppState { tx: tx.clone(), cli_tx, web_vosk_tx, screenshot_tx: screenshot_tx.clone(), + image_format, }; let listener = tokio::net::TcpListener::bind(&args.bind_addr) @@ -165,7 +168,7 @@ async fn main() { .route("/setup", get(static_page::setup_handler)) .route("/vosk", get(static_page::vosk_handler)) .route("/ws", get(ws::ws_handler)) - .route("/screenshot.jpeg", get(ws::screenshot_handler)) + .route("/screenshot", get(ws::screenshot_handler)) .route("/api/change-dir", post(static_page::change_dir_handler)) .route("/vosk_ws", get(ws::web_vosk_ws_handler)) .nest_service( @@ -225,6 +228,7 @@ async fn main() { &mut tui, &mut ui_title, &server_url, + image_format, ) .await; match r { diff --git a/src/protocol.rs b/src/protocol.rs index 3464d66..2cc55a4 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -144,7 +144,7 @@ pub struct NotificationData { // ========== 辅助类型 ========== /// 图片格式 -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Copy, Clone, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum ImageFormat { Png, diff --git a/src/screenshot/canvas.rs b/src/screenshot/canvas.rs deleted file mode 100644 index c2b4c29..0000000 --- a/src/screenshot/canvas.rs +++ /dev/null @@ -1,156 +0,0 @@ -//! Canvas for rendering terminal content to images - -use ab_glyph::{Font, PxScale}; -use image::{ImageBuffer, Rgba, RgbaImage}; -use imageproc::drawing::draw_text_mut; -use tiny_skia::{Color, Paint, Pixmap, Rect, Transform}; - -/// Canvas for drawing shapes and text -pub struct Canvas { - background: Pixmap, - text_layer: ImageBuffer, Vec>, - char_width: u32, - char_height: u32, -} - -impl Canvas { - /// Create a new canvas with the given dimensions - pub fn new(width: u32, height: u32) -> Result { - let background = - Pixmap::new(width, height).ok_or_else(|| "Failed to create pixmap".to_string())?; - - let text_layer = ImageBuffer::new(width, height); - - Ok(Self { - background, - text_layer, - char_width: 8, - char_height: 16, - }) - } - - /// Set character size for text rendering - pub fn set_char_size(&mut self, width: u32, height: u32) { - self.char_width = width; - self.char_height = height; - } - - /// Get character width - #[allow(dead_code)] - pub fn char_width(&self) -> u32 { - self.char_width - } - - /// Get character height - #[allow(dead_code)] - pub fn char_height(&self) -> u32 { - self.char_height - } - - /// Fill the entire canvas with a color - pub fn fill(&mut self, color: [u8; 4]) { - let color = Color::from_rgba8(color[0], color[1], color[2], color[3]); - self.background.fill(color); - } - - /// Fill a rectangle with a color - pub fn fill_rect(&mut self, x: i32, y: i32, width: u32, height: u32, color: [u8; 4]) { - if let Some(rect) = Rect::from_xywh(x as f32, y as f32, width as f32, height as f32) { - let mut paint = Paint::default(); - paint.set_color(Color::from_rgba8(color[0], color[1], color[2], color[3])); - self.background - .fill_rect(rect, &paint, Transform::identity(), None); - } - } - - /// Draw a title bar background at the top - pub fn draw_title_bar(&mut self, _title: &str, _padding: u32) { - let height = 32; - let bg = [40, 40, 45, 255]; - self.fill_rect(0, 0, self.width(), height, bg); - self.fill_rect(0, height as i32 - 2, self.width(), 2, [60, 60, 65, 255]); - } - - /// Draw text at the specified position (simple placeholder) - #[allow(dead_code)] - pub fn draw_text( - &mut self, - text: &str, - x: i32, - y: i32, - color: [u8; 4], - _font: Option<&()>, - _font_size: f32, - ) { - self.draw_text_simple(text, x, y, color); - } - - /// Draw text at the specified position (simple version without font) - #[allow(dead_code)] - pub fn draw_text_simple(&mut self, text: &str, x: i32, y: i32, color: [u8; 4]) { - for (i, ch) in text.chars().enumerate() { - let px_x = x + i as i32 * 8; - let px_y = y; - if !ch.is_whitespace() { - self.fill_rect(px_x, px_y, 6, 10, color); - } - } - } - - /// Draw text using ab_glyph font - pub fn draw_text_with_font( - &mut self, - text: &str, - x: i32, - y: i32, - color: [u8; 4], - font: &F, - scale: PxScale, - ) { - let rgba = Rgba(color); - draw_text_mut(&mut self.text_layer, rgba, x, y, scale, font, text); - } - - /// Get the canvas width - pub fn width(&self) -> u32 { - self.background.width() - } - - /// Get the canvas height - #[allow(dead_code)] - pub fn height(&self) -> u32 { - self.background.height() - } - - /// Convert the canvas to a final image - pub fn into_image(self) -> Result { - let mut final_image = RgbaImage::from_raw( - self.background.width(), - self.background.height(), - self.background.data().to_vec(), - ) - .ok_or_else(|| "Failed to create image from raw data".to_string())?; - - // Blend text layer on top (fast integer alpha blending) - for (dst, src) in final_image.pixels_mut().zip(self.text_layer.pixels()) { - let a = src[3] as u32; - if a == 0 { - continue; - } - if a == 255 { - dst[0] = src[0]; - dst[1] = src[1]; - dst[2] = src[2]; - dst[3] = 255; - } else { - let inv_a = 255 - a; - dst[0] = ((src[0] as u32 * a + dst[0] as u32 * inv_a + 128) >> 8) as u8; - dst[1] = ((src[1] as u32 * a + dst[1] as u32 * inv_a + 128) >> 8) as u8; - dst[2] = ((src[2] as u32 * a + dst[2] as u32 * inv_a + 128) >> 8) as u8; - dst[3] = 255; - } - } - - Ok(final_image) - } -} diff --git a/src/screenshot/font.rs b/src/screenshot/font.rs deleted file mode 100644 index d00238c..0000000 --- a/src/screenshot/font.rs +++ /dev/null @@ -1,33 +0,0 @@ -//! Font loading utilities - -use ab_glyph::{FontArc, PxScale}; -use std::sync::LazyLock; - -/// Global cached font, parsed only once -static FONT: LazyLock = LazyLock::new(|| { - FontArc::try_from_slice(include_bytes!("../../assets/SarasaMonoSC-Light.ttf")) - .expect("Embedded font is valid") -}); - -/// Font data container with ab_glyph FontArc -pub struct FontData { - /// Regular font - pub font: FontArc, - /// Font scale for rendering - pub scale: PxScale, -} - -impl FontData { - /// Create a new FontData with the given font size - pub fn new(font_size: f32) -> Self { - Self { - font: FONT.clone(), - scale: PxScale::from(font_size), - } - } -} - -/// Load font with specific size -pub fn load_font_with_size(size: f32) -> Result { - Ok(FontData::new(size)) -} diff --git a/src/screenshot/mod.rs b/src/screenshot/mod.rs deleted file mode 100644 index 09a30b6..0000000 --- a/src/screenshot/mod.rs +++ /dev/null @@ -1,189 +0,0 @@ -//! Terminal screenshot functionality -//! -//! Renders a vt100::Screen to an image file. - -mod canvas; -mod font; -mod theme; -mod utils; - -pub use canvas::Canvas; -pub use font::{FontData, load_font_with_size}; -pub use theme::Theme; - -use ab_glyph::Font; -use image::ImageError; -use thiserror::Error; - -#[derive(Error, Debug)] -#[allow(clippy::enum_variant_names)] -pub enum ScreenshotError { - #[error("Failed to load font: {0}")] - #[allow(dead_code)] - FontLoadError(String), - - #[error("Canvas error: {0}")] - CanvasError(String), - - #[error("Image error: {0}")] - ImageError(#[from] ImageError), - - #[error("IO error: {0}")] - IoError(#[from] std::io::Error), -} - -/// Configuration for screenshot generation -pub struct ScreenshotConfig { - /// Font size in points - pub font_size: f32, - - /// Padding around the content (in pixels) - pub padding: u32, - - /// Background color (R, G, B, A) - pub background_color: [u8; 4], - - /// Whether to show window decorations - pub show_decorations: bool, - - /// Window title - pub title: Option, -} - -impl Default for ScreenshotConfig { - fn default() -> Self { - Self { - font_size: 14.0, - padding: 16, - background_color: [30, 30, 30, 255], - show_decorations: true, - title: None, - } - } -} - -/// Capture a vt100::Screen as an image -pub fn capture_screen( - screen: &vt100::Screen, - config: &ScreenshotConfig, -) -> Result { - // Try to load font, fallback to built-in if not available - let font_data = - load_font_with_size(config.font_size).unwrap_or_else(|_| FontData::new(config.font_size)); - - let theme = Theme::default(); - - // Calculate character dimensions from the font - let units_per_em = font_data.font.units_per_em().unwrap_or(2048.0); - let space_id = font_data.font.glyph_id(' '); - let advance_unscaled = font_data.font.h_advance_unscaled(space_id); - let char_width = ((advance_unscaled / units_per_em) * font_data.scale.x).round() as u32; - - // Use proper line height = ascent + descent - let ascent = - (font_data.font.ascent_unscaled() / units_per_em * font_data.scale.y).round() as u32; - let descent = - (font_data.font.descent_unscaled() / units_per_em * font_data.scale.y).round() as u32; - let char_height = ascent + descent; - - let (rows, cols) = screen.size(); - let padding = config.padding; - let title_height = if config.show_decorations { 32 } else { 0 }; - - // Find the last row with actual content (skip trailing empty rows) - let mut last_content_row = 0; - 'b: for row in (0..rows).rev() { - for col in 0..cols { - if let Some(cell) = screen.cell(row, col) - && cell.has_contents() - && cell.contents() != " " - { - last_content_row = row; - break 'b; - } - } - } - - // Only render rows up to the last one with content - let actual_rows = last_content_row + 1; - - let image_width = cols as u32 * char_width + padding * 2; - let image_height = actual_rows as u32 * char_height + title_height + padding * 2; - - let mut canvas = Canvas::new(image_width, image_height) - .map_err(|e| ScreenshotError::CanvasError(e.to_string()))?; - - canvas.set_char_size(char_width, char_height); - - // Fill background - canvas.fill(config.background_color); - - // Draw title bar if decorations are enabled - if config.show_decorations { - let title = config.title.as_deref().unwrap_or("Terminal"); - canvas.draw_title_bar(title, config.padding); - - // Draw title with proper font - let title_x = (padding + 8) as i32; - let title_y = 10; - canvas.draw_text_with_font( - title, - title_x, - title_y, - [220, 220, 220, 255], - &font_data.font, - font_data.scale, - ); - } - - // Draw terminal content (only up to last_content_row) - for row in 0..actual_rows { - for col in 0..cols { - if let Some(cell) = screen.cell(row, col) { - let x = padding + col as u32 * char_width; - let y = title_height + padding + row as u32 * char_height; - - // Draw background color - let bg = cell.bgcolor(); - if bg != vt100::Color::Default { - let color = theme.color_to_rgba(bg); - let w = if cell.is_wide() { - char_width * 2 - } else { - char_width - }; - canvas.fill_rect(x as i32, y as i32, w, char_height, color); - } - - // Draw text - imageproc's draw_text_mut y param is the top of the text - if cell.has_contents() && !cell.is_wide_continuation() { - let fg = cell.fgcolor(); - let fg_color = theme.get_foreground(fg, cell.bold(), cell.dim()); - canvas.draw_text_with_font( - cell.contents(), - x as i32, - y as i32, - fg_color, - &font_data.font, - font_data.scale, - ); - } - } - } - } - - canvas - .into_image() - .map_err(|e| ScreenshotError::CanvasError(e.to_string())) -} - -/// Save a vt100::Screen to a PNG file -pub fn save_screen_png( - screen: &vt100::Screen, - path: &str, - config: &ScreenshotConfig, -) -> Result<(), ScreenshotError> { - let image = capture_screen(screen, config)?; - image.save(path)?; - Ok(()) -} diff --git a/src/screenshot/theme.rs b/src/screenshot/theme.rs deleted file mode 100644 index c7e7280..0000000 --- a/src/screenshot/theme.rs +++ /dev/null @@ -1,96 +0,0 @@ -//! Color theme for terminal rendering - -/// Theme for terminal colors -#[derive(Clone, Debug)] -pub struct Theme { - /// Standard ANSI 16 colors - pub colors: [[u8; 4]; 16], -} - -impl Default for Theme { - fn default() -> Self { - Self { - colors: [ - // Normal colors - [0, 0, 0, 255], // Black - [194, 54, 33, 255], // Red - [37, 188, 36, 255], // Green - [173, 173, 39, 255], // Yellow - [73, 46, 225, 255], // Blue - [211, 56, 211, 255], // Magenta - [51, 187, 200, 255], // Cyan - [203, 204, 205, 255], // White - // Bright colors - [128, 128, 128, 255], // Bright Black (Gray) - [255, 85, 85, 255], // Bright Red - [85, 255, 85, 255], // Bright Green - [255, 255, 85, 255], // Bright Yellow - [85, 85, 255, 255], // Bright Blue - [255, 85, 255, 255], // Bright Magenta - [85, 255, 255, 255], // Bright Cyan - [255, 255, 255, 255], // Bright White - ], - } - } -} - -impl Theme { - /// Convert a vt100 Color to RGBA - pub fn color_to_rgba(&self, color: vt100::Color) -> [u8; 4] { - match color { - vt100::Color::Default => [229, 229, 229, 255], - vt100::Color::Idx(i) => { - if i < 16 { - self.colors[i as usize] - } else if i >= 232 { - // Grayscale - let shade = (i - 232) * 10 + 8; - [shade, shade, shade, 255] - } else { - // 216-color cube - let i = i - 16; - let r = (i / 36) * 51; - let g = ((i / 6) % 6) * 51; - let b = (i % 6) * 51; - [r, g, b, 255] - } - } - vt100::Color::Rgb(r, g, b) => [r, g, b, 255], - } - } - - /// Get foreground color with bold and dim attributes - pub fn get_foreground(&self, color: vt100::Color, bold: bool, dim: bool) -> [u8; 4] { - let rgba = if bold { - match color { - vt100::Color::Default => self.colors[15], // bright white - vt100::Color::Idx(i) if i < 8 => { - // Map normal color to bright variant (index + 8) - self.colors[i as usize + 8] - } - _ => self.color_to_rgba(color), - } - } else { - self.color_to_rgba(color) - }; - - if dim { - // Dim: blend 50% toward background (30,30,30) - let bg = [30u8, 30, 30]; - [ - (rgba[0] as u16 / 2 + bg[0] as u16 / 2) as u8, - (rgba[1] as u16 / 2 + bg[1] as u16 / 2) as u8, - (rgba[2] as u16 / 2 + bg[2] as u16 / 2) as u8, - rgba[3], - ] - } else { - rgba - } - } - - /// Get background color - #[allow(dead_code)] - pub fn get_background(&self, color: vt100::Color) -> [u8; 4] { - self.color_to_rgba(color) - } -} diff --git a/src/screenshot/utils.rs b/src/screenshot/utils.rs deleted file mode 100644 index e6656d2..0000000 --- a/src/screenshot/utils.rs +++ /dev/null @@ -1,21 +0,0 @@ -//! Utility functions for screenshot rendering - -/// Calculate character size based on font size -#[allow(dead_code)] -pub fn calculate_char_size(font_size: f32) -> (u32, u32) { - let width = (font_size * 0.6) as u32; - let height = (font_size * 1.2) as u32; - (width, height) -} - -/// Clamp a value between min and max -#[allow(dead_code)] -pub fn clamp(value: T, min: T, max: T) -> T { - if value < min { - min - } else if value > max { - max - } else { - value - } -} diff --git a/src/ws.rs b/src/ws.rs index 1944fe7..9e9dd1d 100644 --- a/src/ws.rs +++ b/src/ws.rs @@ -7,6 +7,7 @@ use axum::{ }, response::IntoResponse, }; +use image::ImageEncoder; use tokio::sync::{broadcast, mpsc}; use vt100::Callbacks; @@ -65,6 +66,7 @@ pub struct AppState { pub cli_tx: ClientTx, pub web_vosk_tx: Option, pub screenshot_tx: ScreenshotTx, + pub image_format: crate::protocol::ImageFormat, } fn t2s>(s: S) -> String { @@ -152,19 +154,6 @@ impl ASRInterface { } fn send_screen(tx: &ServerTx, screen: Arc) { - // if let Ok(jpeg) = render_screen_to_jpeg(&screen) { - // let chunk_size = IMAGE_CHUNK_SIZE; - // let total_chunks = jpeg.len().div_ceil(chunk_size); - // for (i, chunk) in jpeg.chunks(chunk_size).enumerate() { - // let is_last = i == total_chunks - 1; - // let _ = tx.send(ServerMessage::screen_image_chunk( - // crate::protocol::ImageFormat::Jpeg, - // is_last, - // chunk.to_vec(), - // )); - // } - // } - let _ = tx.send(ServerMessage::Screen(screen)); } @@ -190,6 +179,7 @@ pub async fn run_command( tui: &mut crate::ui::TuiTerminal, ui_title: &mut String, ui_footer: &str, + image_format: crate::protocol::ImageFormat, ) -> anyhow::Result { let dir_path = current_dir .as_ref() @@ -283,7 +273,8 @@ pub async fn run_command( TerminalEvent::ScreenGetter(getter) => { let screen = vt_parser.screen().clone(); let mut window_scrollback = 0; - let result = render_screen_to_jpeg(&screen, None, &mut window_scrollback); + let result = + render_screen_to_image(&screen, None, &mut window_scrollback, image_format); let jpeg = match result { Ok(data) => Ok(data), @@ -518,33 +509,31 @@ async fn send_screen_to_client( screen: &vt100::Screen, window_size: Option<(u16, u16)>, // (width, height) window_h_offset: &mut u16, + format: crate::protocol::ImageFormat, ) -> anyhow::Result<()> { - let jpeg = render_screen_to_jpeg(screen, window_size, window_h_offset)?; - if jpeg.is_empty() { + let image_data = render_screen_to_image(screen, window_size, window_h_offset, format)?; + if image_data.is_empty() { state .cli_tx .send(ClientMessage::ScrollDown) .await .map_err(|e| { anyhow::anyhow!( - "Failed to send ScrollDown message to cli_tx after empty JPEG: {}", + "Failed to send ScrollDown message to cli_tx after empty image: {}", e ) })?; } log::debug!( - "Sending screen JPEG to client, size: {} KB", - jpeg.len() / 1024 + "Sending screen image to client, format: {:?}, size: {} KB", + format, + image_data.len() / 1024 ); let chunk_size = IMAGE_CHUNK_SIZE; - let total_chunks = jpeg.len().div_ceil(chunk_size); - for (i, chunk) in jpeg.chunks(chunk_size).enumerate() { + let total_chunks = image_data.len().div_ceil(chunk_size); + for (i, chunk) in image_data.chunks(chunk_size).enumerate() { let is_last = i == total_chunks - 1; - let msg = ServerMessage::screen_image_chunk( - crate::protocol::ImageFormat::Jpeg, - is_last, - chunk.to_vec(), - ); + let msg = ServerMessage::screen_image_chunk(format, is_last, chunk.to_vec()); let data = msg.to_msgpack()?; socket.send(Message::Binary(data.into())).await?; } @@ -585,6 +574,7 @@ async fn handle_client_message( screen, Some((width, height)), window_h_offset, + state.image_format, ) .await?; } else { @@ -607,6 +597,7 @@ async fn handle_client_message( screen, Some((width, height)), window_h_offset, + state.image_format, ) .await?; } else { @@ -676,7 +667,7 @@ async fn handle_socket(mut socket: WebSocket, state: AppState, params: WsParams) }, ServerMessage::Screen(screen_) => { screen = Some(screen_.clone()); - if let Err(e) = send_screen_to_client(&state, &mut socket, screen_, Some(window_size), &mut window_h_offset).await { + if let Err(e) = send_screen_to_client(&state, &mut socket, screen_, Some(window_size), &mut window_h_offset, state.image_format).await { log::error!("Failed to send screen to client: {}", e); } continue; @@ -873,11 +864,12 @@ async fn handle_web_vosk_socket(mut socket: WebSocket, vosk_tx: WebVoskTx) { } } -/// Render a vt100 screen to JPEG bytes -fn render_screen_to_jpeg( +/// Render a vt100 screen to image bytes (JPEG or PNG) +fn render_screen_to_image( screen: &vt100::Screen, window_size: Option<(u16, u16)>, // (width, height) window_h_offset: &mut u16, + format: crate::protocol::ImageFormat, ) -> anyhow::Result> { let config = crate::screenshot::ScreenshotConfig { show_decorations: false, @@ -887,12 +879,12 @@ fn render_screen_to_jpeg( let image = crate::screenshot::capture_screen(screen, &config) .map_err(|e| anyhow::anyhow!("Failed to capture screen: {}", e))?; - // Convert RGBA to RGB for JPEG - let mut rgb_image = image::DynamicImage::ImageRgba8(image).to_rgb8(); + let mut dyn_image = image::DynamicImage::ImageRgba8(image); + let is_png = matches!(format, crate::protocol::ImageFormat::Png); if let Some((width, height)) = window_size { - let orig_width = rgb_image.width(); - let orig_height = rgb_image.height(); + let orig_width = dyn_image.width(); + let orig_height = dyn_image.height(); log::debug!( "Original image size: {}x{}, requested window size: {}x{}", orig_width, @@ -903,12 +895,12 @@ fn render_screen_to_jpeg( let scale = width as f32 / orig_width as f32; let new_height = (orig_height as f32 * scale).round() as u32; - rgb_image = image::imageops::resize( - &rgb_image, + dyn_image = image::DynamicImage::ImageRgba8(image::imageops::resize( + &dyn_image, width as u32, new_height, image::imageops::FilterType::Lanczos3, - ); + )); // 根据垂直偏移截取 let crop_height = height as u32; @@ -927,14 +919,22 @@ fn render_screen_to_jpeg( crop_height - new_height ); let padding_top = crop_height - new_height; - let mut padded = image::RgbImage::new(width as u32, crop_height); - // 填充背景色（深色背景） - for pixel in padded.pixels_mut() { - *pixel = image::Rgb([30, 30, 30]); + + if is_png { + let mut padded = image::RgbaImage::new(width as u32, crop_height); + for pixel in padded.pixels_mut() { + *pixel = image::Rgba([30, 30, 30, 255]); + } + image::imageops::overlay(&mut padded, &dyn_image.to_rgba8(), 0, padding_top as i64); + dyn_image = image::DynamicImage::ImageRgba8(padded); + } else { + let mut padded = image::RgbImage::new(width as u32, crop_height); + for pixel in padded.pixels_mut() { + *pixel = image::Rgb([30, 30, 30]); + } + image::imageops::overlay(&mut padded, &dyn_image.to_rgb8(), 0, padding_top as i64); + dyn_image = image::DynamicImage::ImageRgb8(padded); } - // 把原图放在底部 - image::imageops::overlay(&mut padded, &rgb_image, 0, padding_top as i64); - rgb_image = padded; } else if y_offset + crop_height > new_height { // 偏移量过大，调整到对齐底部 log::warn!( @@ -947,24 +947,46 @@ fn render_screen_to_jpeg( return Ok(Vec::new()); } else { // 正常截取 - rgb_image = - image::imageops::crop(&mut rgb_image, 0, y_offset, width as u32, crop_height) - .to_image(); + dyn_image = + image::imageops::crop(&mut dyn_image, 0, y_offset, width as u32, crop_height) + .to_image() + .into(); } } - let mut jpeg_buf = std::io::Cursor::new(Vec::new()); - let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg_buf, 100); - encoder - .encode( - rgb_image.as_raw(), - rgb_image.width(), - rgb_image.height(), - image::ExtendedColorType::Rgb8, - ) - .map_err(|e| anyhow::anyhow!("Failed to encode JPEG: {}", e))?; + // Encode based on format + let mut buf = std::io::Cursor::new(Vec::new()); + match format { + crate::protocol::ImageFormat::Png => { + let encoder = image::codecs::png::PngEncoder::new(&mut buf); + encoder + .write_image( + dyn_image.as_bytes(), + dyn_image.width(), + dyn_image.height(), + image::ExtendedColorType::from(dyn_image.color()), + ) + .map_err(|e| anyhow::anyhow!("Failed to encode PNG: {}", e))?; + } + crate::protocol::ImageFormat::Jpeg => { + // Convert to RGB for JPEG + let rgb_image = dyn_image.to_rgb8(); + let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, 100); + encoder + .encode( + rgb_image.as_raw(), + rgb_image.width(), + rgb_image.height(), + image::ExtendedColorType::Rgb8, + ) + .map_err(|e| anyhow::anyhow!("Failed to encode JPEG: {}", e))?; + } + _ => { + return Err(anyhow::anyhow!("Unsupported image format: {:?}", format)); + } + } - Ok(jpeg_buf.into_inner()) + Ok(buf.into_inner()) } /// HTTP handler for /screenshot.jpeg