Changes for v0.21

libs/k21/Cargo.toml

@@ -26,8 +26,6 @@ tempfile = "3.8.0"
 axum = "0.7.4"
 reqwest = { version = "0.11", features = ["json", "blocking"] }
 
-
-
 [target.'cfg(target_os = "windows")'.dependencies]
 windows = { version = "0.58", features = [
   "Graphics_Imaging",
@@ -38,4 +36,5 @@ windows = { version = "0.58", features = [
 
 [target.'cfg(target_os = "macos")'.dependencies]
 libc = "=0.2.164"
-cidre = { git = "https://github.com/yury/cidre", rev = "efb9e060c6f8edc48551365c2e80d3e8c6887433", features = ["ns", "cv", "vn"] }
+cidre = { git = "https://github.com/yury/cidre", rev = "efb9e060c6f8edc48551365c2e80d3e8c6887433", features = ["ns", "cv", "vn"] }
+# cidre = { git = "https://github.com/yury/cidre" }

libs/k21/src/capture/mod.rs

@@ -0,0 +1,12 @@
+mod utils;
+pub use utils::capture;
+pub use utils::spawn_screenshot_task;
+pub use utils::capture_with_stdout;
+pub use utils::handle_captured_frames;
+
+mod screen_record;
+pub use screen_record::ScreenCapturer;
+
+
+mod types;
+pub use types::ScreenCaptureConfig;

libs/k21/src/screen_capture/screen_record.rs → libs/k21/src/capture/screen_record.rs

@@ -2,24 +2,26 @@ use image::DynamicImage;
 use openh264::encoder::Encoder;
 use std::path::Path;
 use xcap::Monitor;
-
+use anyhow::Result;
 pub struct ScreenCapturer {
-    monitor_id: u32,
     encoder: Encoder,
     buf: Vec<u8>,
     frame_count: u32,
 }
 
 impl ScreenCapturer {
-    pub fn new(monitor_id: u32) -> Self {
+    pub fn new() -> Self {
         Self {
-            monitor_id,
             encoder: Encoder::new().unwrap(),
             buf: Vec::new(),
             frame_count: 0,
         }
     }
 
+    pub fn is_buf_empty(&self) -> bool {
+        self.buf.len() == 0
+    }
+
     pub fn frame(&mut self, image: &DynamicImage) {
         use openh264::formats::*;
         let frame = image.to_rgb8();
@@ -50,12 +52,7 @@ impl ScreenCapturer {
         use minimp4::Mp4Muxer;
         use std::io::{Cursor, Read, Seek, SeekFrom};
 
-        let monitor = Monitor::all()
-            .unwrap()
-            .into_iter()
-            .find(|m| m.id() == self.monitor_id)
-            .ok_or_else(|| anyhow::anyhow!("Monitor not found"))
-            .unwrap();
+        let monitor = get_primary_monitor();
 
         let mut video_buffer = Cursor::new(Vec::new());
         let mut mp4muxer = Mp4Muxer::new(&mut video_buffer);
@@ -65,6 +62,7 @@ impl ScreenCapturer {
             false,
             "Screen capturer",
         );
+
         mp4muxer.write_video_with_fps(&self.buf, fps as u32);
         mp4muxer.close();
 
@@ -82,3 +80,39 @@ impl ScreenCapturer {
         self.frame_count = 0;
     }
 }
+
+fn get_monitor(monitor_id: u32) -> Monitor {
+    Monitor::all()
+        .unwrap()
+        .into_iter()
+        .find(|m| m.id() == monitor_id)
+        .ok_or_else(|| anyhow::anyhow!("Monitor not found"))
+        .unwrap()
+}
+
+fn get_primary_monitor_id() -> u32 {
+    Monitor::all()
+        .unwrap()
+        .iter()
+        .find(|m| m.is_primary())
+        .unwrap()
+        .id()
+}
+
+pub fn get_primary_monitor() -> Monitor {
+    get_monitor(get_primary_monitor_id())
+}
+
+pub async fn get_screenshot() -> Result<DynamicImage> {
+    let image = std::thread::spawn(move || -> Result<DynamicImage> {
+        let monitor = get_primary_monitor();
+        let image = monitor
+            .capture_image()
+            .map_err(anyhow::Error::from)
+            .map(DynamicImage::ImageRgba8)?;
+        Ok(image)
+    })
+    .join()
+    .unwrap()?;
+    Ok(image)
+}

libs/k21/src/capture/types.rs

@@ -0,0 +1,50 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ScreenCaptureConfig {
+    pub fps: f32,
+    pub video_chunk_duration_in_seconds: u64,
+    pub save_screenshot: bool,
+    pub save_video: bool,
+    pub record_length_in_seconds: u64,
+    pub output_dir_video: Option<String>,
+    pub output_dir_screenshot: Option<String>,
+}
+
+impl Default for ScreenCaptureConfig {
+    fn default() -> Self {
+        Self {
+            fps: 1.0,
+            video_chunk_duration_in_seconds: 60,
+            save_screenshot: false,
+            save_video: false,
+            record_length_in_seconds: 1,
+            output_dir_video: None,
+            output_dir_screenshot: None,
+        }
+    }
+}
+
+impl ScreenCaptureConfig {
+    pub fn new(
+        fps: f32,
+        video_chunk_duration_in_seconds: u64,
+        save_screenshot: bool,
+        save_video: bool,
+        record_length_in_seconds: u64,
+        output_dir_video: Option<String>,
+        output_dir_screenshot: Option<String>,
+    ) -> Self {
+        let config: ScreenCaptureConfig = Self {
+            fps,
+            video_chunk_duration_in_seconds,
+            record_length_in_seconds,
+            save_screenshot,
+            save_video,
+            output_dir_video,
+            output_dir_screenshot,
+            ..Default::default()
+        };
+        config
+    }
+}

libs/k21/src/capture/utils.rs

@@ -0,0 +1,222 @@
+use anyhow::Result;
+use image::DynamicImage;
+
+use std::time::{Duration, Instant};
+use tokio::io::{self, AsyncWriteExt};
+use tokio::sync::broadcast::channel;
+
+use crate::common::to_verified_path;
+use crate::capture::screen_record;
+use super::screen_record::get_screenshot;
+use tokio::sync::watch;
+use super::ScreenCaptureConfig;
+
+pub async fn capture(config: ScreenCaptureConfig) -> Result<()> {
+    capture_with_stdout(config, false).await
+}
+
+pub async fn capture_with_stdout(mut config: ScreenCaptureConfig, stdout: bool) -> Result<()> {
+    if config.save_video {
+        config.output_dir_video = Some(match &config.output_dir_video {
+            Some(path) => to_verified_path(path)?.to_string_lossy().to_string(),
+            None => std::env::current_dir()?.to_string_lossy().to_string(), 
+        });
+    }
+
+    log::info!("Starting capture at {} fps", config.fps);
+
+    let (screenshot_tx, mut screenshot_rx) = channel(512);
+    let (close_tx, close_rx) = watch::channel(false);
+
+    let screenshot_task = spawn_screenshot_task(
+        &config,
+        screenshot_tx,
+        close_tx,
+    );
+
+    let _ = handle_captured_frames(
+        &config,
+        stdout,
+        &mut screenshot_rx,
+        close_rx,
+    ).await;
+
+    log::info!("Exiting...");
+    let _ = screenshot_task.await;
+
+    Ok(())
+}
+
+pub fn spawn_screenshot_task(
+    config: &ScreenCaptureConfig,
+    screenshot_tx: tokio::sync::broadcast::Sender<(u64, DynamicImage)>,
+    close_tx: tokio::sync::watch::Sender<bool>
+) -> tokio::task::JoinHandle<()> {
+    tokio::task::spawn({
+        let interval = Duration::from_secs_f32(1.0 / config.fps);
+        let total_frames_to_process = config.record_length_in_seconds * config.fps as u64;
+        let live_capture = config.record_length_in_seconds == 0;
+
+        async move {
+            let mut frame_counter: u64 = 1;
+            while live_capture || frame_counter <= total_frames_to_process {
+                let capture_start = Instant::now();
+                match get_screenshot().await {
+                    Ok(image) => {
+                        // Use try_send to avoid blocking if receiver is slow
+                        if let Err(e) = screenshot_tx.send((frame_counter, image)) {
+                            log::error!("Failed to send screenshot: {}", e);
+                            break;
+                        }
+                    },
+                    Err(e) => {
+                        log::error!("Failed to capture screenshot: {}", e);
+                        // Continue to next iteration instead of breaking
+                        tokio::time::sleep(interval).await;
+                        continue;
+                    }
+                }
+
+                let capture_duration = capture_start.elapsed();
+                frame_counter += 1;
+
+                if let Some(diff) = interval.checked_sub(capture_duration) {
+                    log::debug!("Sleeping for {:?}", diff);
+                    tokio::time::sleep(diff).await;
+                } else {
+                    log::warn!(
+                        "Capture took longer than expected: {:?}, will not sleep",
+                        capture_duration
+                    );
+                }
+            }
+            let _ = close_tx.send(true);
+            log::debug!("Screenshot task completed after {} frames", frame_counter - 1);
+        }
+    })
+}
+
+pub async fn handle_captured_frames(
+    config: &ScreenCaptureConfig,
+    stdout: bool,
+    screenshot_rx: &mut tokio::sync::broadcast::Receiver<(u64, DynamicImage)>,
+    close_rx: tokio::sync::watch::Receiver<bool>
+) -> Result<()> {
+    let screen_record = &mut screen_record::ScreenCapturer::new();
+    let mut chunk_number = 0;
+
+    // Handle frames
+    save_or_send_captured_frames(
+        config,
+        stdout,
+        screen_record,
+        screenshot_rx,
+        close_rx,
+        &mut chunk_number,
+    ).await;
+
+    // Save final video chunk if needed
+    if config.save_video && !screen_record.is_buf_empty() {
+        save_video_chunk(
+            screen_record, 
+            &mut chunk_number, 
+            config.fps, 
+            config.output_dir_video.as_ref().unwrap()
+        );
+    }
+
+    Ok(())
+}
+
+async fn save_or_send_captured_frames(
+    config: &ScreenCaptureConfig,
+    stdout: bool,
+    screen_record: &mut screen_record::ScreenCapturer,
+    screenshot_rx: &mut tokio::sync::broadcast::Receiver<(u64, DynamicImage)>,
+    mut close_rx: tokio::sync::watch::Receiver<bool>,
+    chunk_number: &mut u64,
+) {
+    let total_fps_in_chunk = config.fps as u64 * config.video_chunk_duration_in_seconds;
+
+    loop {
+        tokio::select! {
+            Ok((frame_number, image)) = screenshot_rx.recv() => {
+                if stdout {
+                    send_frame_to_stdout(frame_number, &image).await;
+                }
+
+                // record the frame
+                if config.save_video {
+                    screen_record.frame(&image);
+                    log::info!("frame {}", frame_number);
+
+                    if frame_number % total_fps_in_chunk == 0 {
+                        log::info!(
+                            "frame {}, total_fps_in_chunk {}",
+                            frame_number,
+                            total_fps_in_chunk
+                        );
+                        save_video_chunk(screen_record, chunk_number, config.fps, config.output_dir_video.as_ref().unwrap());
+                    }
+                }
+
+                // save screenshot to disk
+                if config.save_screenshot {
+                    if let Some(output_dir) = &config.output_dir_screenshot {
+                        save_screenshot(frame_number, image.clone(), output_dir);
+                    } else {
+                        log::warn!("Screenshot saving enabled but no output directory specified");
+                    }
+                }
+            }
+
+            Ok(_) = close_rx.changed() => {
+                if *close_rx.borrow() {
+                    log::debug!("Screenshot channel closed, stopping OCR processing");
+                    break;
+                }
+            }
+        }
+    }
+
+    if config.save_screenshot {
+        if let Some(output_dir) = &config.output_dir_screenshot {
+            log::info!("Total screenshots saved in directory: {}", 
+                    output_dir);
+        }
+    }
+}
+
+async fn send_frame_to_stdout(frame_number: u64, image: &DynamicImage) {
+    let rgb = image.to_rgb8();
+    let data = rgb.as_raw();
+    let mut stdout = io::stdout();
+
+    log::info!("Sending frame {}, len {}", frame_number, data.len());
+
+    // send frame & size of raw image data
+    stdout.write_all(&frame_number.to_le_bytes()).await.unwrap(); // Send frame number
+    stdout.write_all(&rgb.width().to_le_bytes()).await.unwrap(); // Send width
+    stdout.write_all(&rgb.height().to_le_bytes()).await.unwrap(); // Send height
+    stdout.write_all(&data.len().to_le_bytes()).await.unwrap(); // Send data size
+    stdout.write_all(&data).await.unwrap(); // Send frame data
+    stdout.flush().await.unwrap(); // Ensure it's sent
+}
+
+fn save_video_chunk(screen_record: &mut screen_record::ScreenCapturer, chunk_number: &mut u64, fps: f32, output_dir_video: &str) {
+    // save video chunk to disk with unique name using the provided output directory
+    let path = std::path::PathBuf::from(output_dir_video).join(format!("output-{}.mp4", chunk_number));
+    screen_record.save(&path, fps);
+    *chunk_number += 1;
+}
+
+fn save_screenshot(frame_number: u64, image: DynamicImage, output_dir: &str) {
+    let output_dir = std::path::PathBuf::from(output_dir);
+    tokio::task::spawn(async move {
+        let path = output_dir.join(format!("screenshot-{}.png", frame_number));
+        match image.save_with_format(&path, image::ImageFormat::Png) {
+            Ok(_) => log::info!("Saved screenshot to {}", path.display()),
+            Err(e) => log::error!("Failed to save screenshot: {}", e),
+        }
+    });
+}