implement support for display capture

Author: modelflat

Cargo.lock

@@ -21,7 +21,7 @@ Alternatively, you can install suitable wheel from [releases page](https://githu
 ```python
 from zbl import Capture
 
-with Capture('visual studio code') as cap:
+with Capture(window_name='visual studio code') as cap:
     frame = next(cap.frames())
     print(frame.shape)
 ```
@@ -31,7 +31,7 @@ The snippet above will capture a window which title contains the string `visual
 To run an example using OpenCV's `highgui`:
 
 1. Install `opencv-python`
-2. Run `python -m zbl '<full or partial window name, case insensitive>'`
+2. Run `python -m zbl --window-name '<full or partial window name, case insensitive>'`
 
 ## Rust
 

README.md

@@ -1,6 +1,6 @@
 [package]
 name = "zbl"
-version = "0.0.4"
+version = "0.1.0"
 edition = "2021"
 
 [lib]
@@ -13,27 +13,29 @@ lazy_static = "1"
 version = "0.43"
 features = [
     "Foundation",
+    "Graphics",
+    "Graphics_Capture",
+    "Graphics_DirectX",
+    "Graphics_DirectX_Direct3D11",
+    "UI",
     "Win32_Foundation",
-    "Win32_UI_HiDpi",
-    "Win32_UI_WindowsAndMessaging",
-    "Win32_UI_Accessibility",
-    "Win32_System_Console",
+    "Win32_Graphics_Direct3D",
+    "Win32_Graphics_Direct3D11",
     "Win32_Graphics_Dwm",
-    "Win32_Graphics_Gdi",
     "Win32_Graphics_Dxgi",
     "Win32_Graphics_Dxgi_Common",
-    "Win32_Graphics_Direct3D",
-    "Win32_Graphics_Direct3D11",
+    "Win32_Graphics_Gdi",
+    "Win32_System_Console",
     "Win32_System_WinRT",
     "Win32_System_WinRT_Direct3D11",
     "Win32_System_WinRT_Graphics_Capture",
-    "UI",
-    "Graphics",
-    "Graphics_Capture",
-    "Graphics_DirectX",
-    "Graphics_DirectX_Direct3D11",
+    "Win32_UI_Accessibility",
+    "Win32_UI_HiDpi",
+    "Win32_UI_Shell",
+    "Win32_UI_Shell_Common",
+    "Win32_UI_WindowsAndMessaging",
 ]
 
 [dev-dependencies]
 clap = { version = "4", features = ["derive"] }
-opencv = "0.77"
+opencv = "0.82"

zbl/Cargo.toml

@@ -2,23 +2,35 @@ use std::time::Instant;
 
 use clap::Parser;
 use opencv::{highgui, prelude::*};
-use zbl::{init, Capture, Frame, Window};
+use zbl::{display::Display, Capturable, Capture, Frame, Window};
 
 #[derive(Parser, Debug)]
 #[clap(version)]
 struct Args {
     #[clap(long)]
-    window: String,
+    window_name: Option<String>,
+    #[clap(long)]
+    display_id: Option<usize>,
 }
 
 fn main() {
-    init();
+    zbl::init();
 
     let args = Args::parse();
-    let window = Window::find_first(&args.window).expect("failed to find window");
-    let mut capturer = Capture::new(window).expect("failed to initialize capture");
 
-    capturer.start().expect("failed to start capture");
+    let mut target = if let Some(window_name) = args.window_name {
+        let window = Window::find_first(&window_name).expect("failed to find window");
+        Box::new(window) as Box<dyn Capturable>
+    } else if let Some(display_id) = args.display_id {
+        let display = Display::find_by_id(display_id).expect("failed to find display");
+        Box::new(display) as Box<dyn Capturable>
+    } else {
+        panic!("either --window-name or --display-id should be set!");
+    };
+
+    let mut capture = Capture::new(target, true).expect("failed to initialize capture");
+
+    capture.start().expect("failed to start capture");
 
     highgui::named_window("Test", highgui::WINDOW_AUTOSIZE).expect("failed to setup opencv window");
 
@@ -28,7 +40,7 @@ fn main() {
     let mut tt = 0f32;
     loop {
         let t = Instant::now();
-        if let Some(Frame { texture, ptr }) = capturer.grab().expect("failed to get frame") {
+        if let Some(Frame { texture, ptr }) = capture.grab().expect("failed to get frame") {
             let mat = unsafe {
                 Mat::new_size_with_data(
                     opencv::core::Size::new(texture.desc.Width as i32, texture.desc.Height as i32),
@@ -56,5 +68,5 @@ fn main() {
         }
     }
 
-    capturer.stop().expect("failed to stop capture");
+    capture.stop().expect("failed to stop capture");
 }

zbl/examples/trivial.rs

@@ -1,12 +1,5 @@
-use std::{
-    collections::HashMap,
-    sync::{
-        mpsc::{sync_channel, Receiver, SyncSender, TryRecvError, TrySendError},
-        RwLock,
-    },
-};
+use std::sync::mpsc::{sync_channel, Receiver, TryRecvError, TrySendError};
 
-use lazy_static::lazy_static;
 use windows::{
     core::{IInspectable, Interface, Result},
     Foundation::TypedEventHandler,
@@ -15,62 +8,18 @@ use windows::{
         DirectX::{Direct3D11::IDirect3DDevice, DirectXPixelFormat},
         SizeInt32,
     },
-    Win32::{
-        Foundation::HWND,
-        Graphics::Direct3D11::{
-            ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D, D3D11_BOX,
-            D3D11_MAPPED_SUBRESOURCE, D3D11_TEXTURE2D_DESC,
-        },
-        UI::{
-            Accessibility::{SetWinEventHook, UnhookWinEvent, HWINEVENTHOOK},
-            WindowsAndMessaging::{EVENT_OBJECT_DESTROY, WINEVENT_OUTOFCONTEXT},
-        },
+    Win32::Graphics::Direct3D11::{
+        ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D, D3D11_BOX, D3D11_MAPPED_SUBRESOURCE,
+        D3D11_TEXTURE2D_DESC,
     },
 };
 
 use crate::{
     staging_texture::StagingTexture,
     util::{create_d3d_device, create_direct3d_device, get_dxgi_interface_from_object},
-    window::Window,
+    Capturable,
 };
 
-lazy_static! {
-    static ref OBJECT_DESTROYED_USER_DATA: RwLock<HashMap<isize, (isize, SyncSender<()>)>> =
-        Default::default();
-}
-
-extern "system" fn object_destroyed_cb(
-    this: HWINEVENTHOOK,
-    _: u32,
-    handle: HWND,
-    id_object: i32,
-    id_child: i32,
-    _: u32,
-    _: u32,
-) {
-    if id_object == 0 && id_child == 0 && handle != HWND::default() {
-        let has_been_closed = if let Ok(handles) = OBJECT_DESTROYED_USER_DATA.read() {
-            if let Some((window_handle, tx)) = handles.get(&this.0) {
-                if *window_handle == handle.0 {
-                    tx.send(()).ok();
-                    true
-                } else {
-                    false
-                }
-            } else {
-                false
-            }
-        } else {
-            // TODO is that correct?
-            true
-        };
-
-        if has_been_closed {
-            unsafe { UnhookWinEvent(this) };
-        }
-    }
-}
-
 pub struct Frame<'a> {
     pub texture: &'a StagingTexture,
     pub ptr: D3D11_MAPPED_SUBRESOURCE,
@@ -80,9 +29,9 @@ pub struct Capture {
     device: ID3D11Device,
     direct3d_device: IDirect3DDevice,
     context: ID3D11DeviceContext,
-    window: Window,
-    window_box: D3D11_BOX,
-    window_closed_signal: Receiver<()>,
+    capturable: Box<dyn Capturable>,
+    capture_box: D3D11_BOX,
+    capture_done_signal: Receiver<()>,
     frame_pool: Direct3D11CaptureFramePool,
     frame_source: Receiver<Option<Direct3D11CaptureFrame>>,
     session: GraphicsCaptureSession,
@@ -92,51 +41,31 @@ pub struct Capture {
 }
 
 impl Capture {
-    /// Create a new capture of `window`. This will initialize D3D11 devices, context, and Windows.Graphics.Capture's
+    /// Create a new capture. This will initialize D3D11 devices, context, and Windows.Graphics.Capture's
     /// frame pool / capture session.
     ///
     /// Note that this will not start capturing yet. Call `start()` to actually start receiving frames.
-    pub fn new(window: Window) -> Result<Self> {
-        let d3d_device = create_d3d_device()?;
-        let d3d_context = {
+    pub fn new(capturable: Box<dyn Capturable>, capture_cursor: bool) -> Result<Self> {
+        let device = create_d3d_device()?;
+        let context = unsafe {
             let mut d3d_context = None;
-            unsafe { d3d_device.GetImmediateContext(&mut d3d_context) };
+            device.GetImmediateContext(&mut d3d_context);
             d3d_context.expect("failed to create d3d_context")
         };
-        let device = create_direct3d_device(&d3d_device)?;
+        let direct3d_device = create_direct3d_device(&device)?;
 
-        let capture_item = window.create_capture_item()?;
+        let capture_item = capturable.create_capture_item()?;
         let capture_item_size = capture_item.Size()?;
 
         let frame_pool = Direct3D11CaptureFramePool::CreateFreeThreaded(
-            &device,
+            &direct3d_device,
             DirectXPixelFormat::B8G8R8A8UIntNormalized,
             1,
             capture_item_size,
         )?;
 
         let session = frame_pool.CreateCaptureSession(&capture_item)?;
-        session.SetIsCursorCaptureEnabled(false)?;
-
-        let (sender, window_closed_signal) = sync_channel(1);
-        let hook_id = unsafe {
-            SetWinEventHook(
-                EVENT_OBJECT_DESTROY,
-                EVENT_OBJECT_DESTROY,
-                None,
-                Some(object_destroyed_cb),
-                // TODO filtering by process id does not always catch the moment when the window is closed
-                // why? aren't windows bound to their process ids?
-                // moreover, for explorer windows even that does not work.
-                // need some more realiable and simpler way to track window closing
-                0,
-                0,
-                WINEVENT_OUTOFCONTEXT,
-            )
-        };
-        if let Ok(mut handles) = OBJECT_DESTROYED_USER_DATA.write() {
-            handles.insert(hook_id.0, (window.handle.0, sender));
-        }
+        session.SetIsCursorCaptureEnabled(capture_cursor)?;
 
         let (sender, receiver) = sync_channel(1 << 5);
         frame_pool.FrameArrived(
@@ -160,15 +89,16 @@ impl Capture {
             ),
         )?;
 
-        let window_box = window.get_client_box();
+        let capture_box = capturable.get_client_box()?;
+        let capture_done_signal = capturable.get_close_notification_channel();
 
         Ok(Self {
-            device: d3d_device,
-            direct3d_device: device,
-            context: d3d_context,
-            window,
-            window_box,
-            window_closed_signal,
+            device,
+            direct3d_device,
+            context,
+            capturable,
+            capture_box,
+            capture_done_signal,
             frame_pool,
             frame_source: receiver,
             session,
@@ -178,9 +108,9 @@ impl Capture {
         })
     }
 
-    /// Get attached window.
-    pub fn window(&self) -> &Window {
-        &self.window
+    /// Get attached capturable.
+    pub fn capturable(&self) -> &Box<dyn Capturable> {
+        &self.capturable
     }
 
     /// Start capturing frames.
@@ -195,7 +125,7 @@ impl Capture {
     ///
     /// Returns:
     /// * `Ok(Some(...))` if there is a frame and it's been successfully captured;
-    /// * `Ok(None)` if no frames can be received from the application (i.e. when the window was closed).
+    /// * `Ok(None)` if no frames can be received (e.g. when the window was closed).
     /// * `Err(...)` if an error has occured while capturing a frame.
     pub fn grab(&mut self) -> Result<Option<Frame>> {
         if self.grab_next()? {
@@ -223,9 +153,9 @@ impl Capture {
     }
 
     fn recreate_frame_pool(&mut self) -> Result<()> {
-        let capture_item = self.window.create_capture_item()?;
+        let capture_item = self.capturable.create_capture_item()?;
         let capture_item_size = capture_item.Size()?;
-        self.window_box = self.window.get_client_box();
+        self.capture_box = self.capturable.get_client_box()?;
         self.frame_pool.Recreate(
             &self.direct3d_device,
             DirectXPixelFormat::B8G8R8A8UIntNormalized,
@@ -245,7 +175,7 @@ impl Capture {
                 Err(TryRecvError::Empty) => {
                     // TODO busy loop? so uncivilized
                     if let Ok(()) | Err(TryRecvError::Disconnected) =
-                        self.window_closed_signal.try_recv()
+                        self.capture_done_signal.try_recv()
                     {
                         self.stop()?;
                         return Ok(false);
@@ -267,8 +197,8 @@ impl Capture {
             self.recreate_frame_pool()?;
             let new_staging_texture = StagingTexture::new(
                 &self.device,
-                self.window_box.right - self.window_box.left,
-                self.window_box.bottom - self.window_box.top,
+                self.capture_box.right - self.capture_box.left,
+                self.capture_box.bottom - self.capture_box.top,
                 desc.Format,
             )?;
             self.staging_texture = Some(new_staging_texture);
@@ -286,7 +216,7 @@ impl Capture {
                 0,
                 Some(&copy_src),
                 0,
-                Some(&self.window_box as *const _),
+                Some(&self.capture_box as *const _),
             );
         }