windows_capture/
frame.rs

1use std::{
2    fs::{self},
3    io,
4    path::Path,
5    ptr, slice,
6};
7
8use rayon::iter::{IntoParallelIterator, ParallelIterator};
9use windows::{
10    Foundation::TimeSpan,
11    Graphics::DirectX::Direct3D11::IDirect3DSurface,
12    Win32::Graphics::{
13        Direct3D11::{
14            ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D, D3D11_BOX, D3D11_CPU_ACCESS_READ,
15            D3D11_CPU_ACCESS_WRITE, D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ_WRITE,
16            D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING,
17        },
18        Dxgi::Common::{DXGI_FORMAT, DXGI_SAMPLE_DESC},
19    },
20};
21
22use crate::{
23    encoder::{self, ImageEncoder},
24    settings::ColorFormat,
25};
26
27#[derive(thiserror::Error, Debug)]
28pub enum Error {
29    #[error("Invalid box size")]
30    InvalidSize,
31    #[error("This color format is not supported for saving as image")]
32    UnsupportedFormat,
33    #[error("Failed to encode image buffer to image bytes with specified format: {0}")]
34    ImageEncoderError(#[from] encoder::ImageEncoderError),
35    #[error("IO error: {0}")]
36    IoError(#[from] io::Error),
37    #[error("Windows API error: {0}")]
38    WindowsError(#[from] windows::core::Error),
39}
40
41#[derive(Eq, PartialEq, Clone, Copy, Debug)]
42pub enum ImageFormat {
43    Jpeg,
44    Png,
45    Gif,
46    Tiff,
47    Bmp,
48    JpegXr,
49}
50
51/// Represents a frame captured from a graphics capture item.
52///
53/// # Example
54/// ```ignore
55/// // Get frame from capture the session
56/// let mut buffer = frame.buffer()?;
57/// buffer.save_as_image("screenshot.png", ImageFormat::Png)?;
58/// ```
59pub struct Frame<'a> {
60    d3d_device: &'a ID3D11Device,
61    frame_surface: IDirect3DSurface,
62    frame_texture: ID3D11Texture2D,
63    time: TimeSpan,
64    context: &'a ID3D11DeviceContext,
65    buffer: &'a mut Vec<u8>,
66    width: u32,
67    height: u32,
68    color_format: ColorFormat,
69}
70
71impl<'a> Frame<'a> {
72    /// Create a new Frame.
73    ///
74    /// # Arguments
75    ///
76    /// * `d3d_device` - The ID3D11Device used for creating the frame.
77    /// * `frame_surface` - The IDirect3DSurface representing the frame surface.
78    /// * `frame_texture` - The ID3D11Texture2D representing the frame texture.
79    /// * `time` - The TimeSpan representing the frame time.
80    /// * `context` - The ID3D11DeviceContext used for copying the texture.
81    /// * `buffer` - The mutable Vec<u8> representing the frame buffer.
82    /// * `width` - The width of the frame.
83    /// * `height` - The height of the frame.
84    /// * `color_format` - The ColorFormat of the frame.
85    ///
86    /// # Returns
87    ///
88    /// A new Frame instance.
89    #[allow(clippy::too_many_arguments)]
90    #[must_use]
91    #[inline]
92    pub fn new(
93        d3d_device: &'a ID3D11Device,
94        frame_surface: IDirect3DSurface,
95        frame_texture: ID3D11Texture2D,
96        time: TimeSpan,
97        context: &'a ID3D11DeviceContext,
98        buffer: &'a mut Vec<u8>,
99        width: u32,
100        height: u32,
101        color_format: ColorFormat,
102    ) -> Self {
103        Self {
104            d3d_device,
105            frame_surface,
106            frame_texture,
107            time,
108            context,
109            buffer,
110            width,
111            height,
112            color_format,
113        }
114    }
115
116    /// Get the width of the frame.
117    ///
118    /// # Returns
119    ///
120    /// The width of the frame.
121    #[must_use]
122    #[inline]
123    pub const fn width(&self) -> u32 {
124        self.width
125    }
126
127    /// Get the height of the frame.
128    ///
129    /// # Returns
130    ///
131    /// The height of the frame.
132    #[must_use]
133    #[inline]
134    pub const fn height(&self) -> u32 {
135        self.height
136    }
137
138    /// Get the time of the frame.
139    ///
140    /// # Returns
141    ///
142    /// The time of the frame.
143    #[must_use]
144    #[inline]
145    pub const fn timespan(&self) -> TimeSpan {
146        self.time
147    }
148
149    /// Get the color format of the frame.
150    ///
151    /// # Returns
152    ///
153    /// The color format of the frame.
154    #[must_use]
155    #[inline]
156    pub const fn color_format(&self) -> ColorFormat {
157        self.color_format
158    }
159
160    /// Get the raw surface of the frame.
161    ///
162    /// # Returns
163    ///
164    /// The IDirect3DSurface representing the raw surface of the frame.
165    ///
166    /// # Safety
167    ///
168    /// This method is unsafe because it returns a raw pointer to the IDirect3DSurface.
169    #[allow(clippy::missing_safety_doc)]
170    #[must_use]
171    #[inline]
172    pub const unsafe fn as_raw_surface(&self) -> &IDirect3DSurface {
173        &self.frame_surface
174    }
175
176    /// Get the raw texture of the frame.
177    ///
178    /// # Returns
179    ///
180    /// The ID3D11Texture2D representing the raw texture of the frame.
181    #[allow(clippy::missing_safety_doc)]
182    #[must_use]
183    #[inline]
184    pub const unsafe fn as_raw_texture(&self) -> &ID3D11Texture2D {
185        &self.frame_texture
186    }
187
188    /// Get the frame buffer.
189    ///
190    /// # Returns
191    ///
192    /// The FrameBuffer containing the frame data.
193    #[inline]
194    pub fn buffer(&mut self) -> Result<FrameBuffer, Error> {
195        // Texture Settings
196        let texture_desc = D3D11_TEXTURE2D_DESC {
197            Width: self.width,
198            Height: self.height,
199            MipLevels: 1,
200            ArraySize: 1,
201            Format: DXGI_FORMAT(self.color_format as i32),
202            SampleDesc: DXGI_SAMPLE_DESC {
203                Count: 1,
204                Quality: 0,
205            },
206            Usage: D3D11_USAGE_STAGING,
207            BindFlags: 0,
208            CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32 | D3D11_CPU_ACCESS_WRITE.0 as u32,
209            MiscFlags: 0,
210        };
211
212        // Create a texture that CPU can read
213        let mut texture = None;
214        unsafe {
215            self.d3d_device
216                .CreateTexture2D(&texture_desc, None, Some(&mut texture))?;
217        };
218
219        let texture = texture.unwrap();
220
221        // Copy the real texture to copy texture
222        unsafe {
223            self.context.CopyResource(&texture, &self.frame_texture);
224        };
225
226        // Map the texture to enable CPU access
227        let mut mapped_resource = D3D11_MAPPED_SUBRESOURCE::default();
228        unsafe {
229            self.context.Map(
230                &texture,
231                0,
232                D3D11_MAP_READ_WRITE,
233                0,
234                Some(&mut mapped_resource),
235            )?;
236        };
237
238        // Get the mapped resource data slice
239        let mapped_frame_data = unsafe {
240            slice::from_raw_parts_mut(
241                mapped_resource.pData.cast(),
242                (self.height * mapped_resource.RowPitch) as usize,
243            )
244        };
245
246        // Create frame buffer from slice
247        let frame_buffer = FrameBuffer::new(
248            mapped_frame_data,
249            self.buffer,
250            self.width,
251            self.height,
252            mapped_resource.RowPitch,
253            mapped_resource.DepthPitch,
254            self.color_format,
255        );
256
257        Ok(frame_buffer)
258    }
259
260    /// Get a cropped frame buffer.
261    ///
262    /// # Arguments
263    ///
264    /// * `start_width` - The starting width of the cropped frame.
265    /// * `start_height` - The starting height of the cropped frame.
266    /// * `end_width` - The ending width of the cropped frame.
267    /// * `end_height` - The ending height of the cropped frame.
268    ///
269    /// # Returns
270    ///
271    /// The FrameBuffer containing the cropped frame data.
272    #[inline]
273    pub fn buffer_crop(
274        &mut self,
275        start_width: u32,
276        start_height: u32,
277        end_width: u32,
278        end_height: u32,
279    ) -> Result<FrameBuffer, Error> {
280        if start_width >= end_width || start_height >= end_height {
281            return Err(Error::InvalidSize);
282        }
283
284        let texture_width = end_width - start_width;
285        let texture_height = end_height - start_height;
286
287        // Texture Settings
288        let texture_desc = D3D11_TEXTURE2D_DESC {
289            Width: texture_width,
290            Height: texture_height,
291            MipLevels: 1,
292            ArraySize: 1,
293            Format: DXGI_FORMAT(self.color_format as i32),
294            SampleDesc: DXGI_SAMPLE_DESC {
295                Count: 1,
296                Quality: 0,
297            },
298            Usage: D3D11_USAGE_STAGING,
299            BindFlags: 0,
300            CPUAccessFlags: D3D11_CPU_ACCESS_READ.0 as u32 | D3D11_CPU_ACCESS_WRITE.0 as u32,
301            MiscFlags: 0,
302        };
303
304        // Create a texture that CPU can read
305        let mut texture = None;
306        unsafe {
307            self.d3d_device
308                .CreateTexture2D(&texture_desc, None, Some(&mut texture))?;
309        };
310        let texture = texture.unwrap();
311
312        // Box Settings
313        let resource_box = D3D11_BOX {
314            left: start_width,
315            top: start_height,
316            front: 0,
317            right: end_width,
318            bottom: end_height,
319            back: 1,
320        };
321
322        // Copy the real texture to copy texture
323        unsafe {
324            self.context.CopySubresourceRegion(
325                &texture,
326                0,
327                0,
328                0,
329                0,
330                &self.frame_texture,
331                0,
332                Some(&resource_box),
333            );
334        };
335
336        // Map the texture to enable CPU access
337        let mut mapped_resource = D3D11_MAPPED_SUBRESOURCE::default();
338        unsafe {
339            self.context.Map(
340                &texture,
341                0,
342                D3D11_MAP_READ_WRITE,
343                0,
344                Some(&mut mapped_resource),
345            )?;
346        };
347
348        // Get the mapped resource data slice
349        let mapped_frame_data = unsafe {
350            slice::from_raw_parts_mut(
351                mapped_resource.pData.cast(),
352                (texture_height * mapped_resource.RowPitch) as usize,
353            )
354        };
355
356        // Create frame buffer from slice
357        let frame_buffer = FrameBuffer::new(
358            mapped_frame_data,
359            self.buffer,
360            texture_width,
361            texture_height,
362            mapped_resource.RowPitch,
363            mapped_resource.DepthPitch,
364            self.color_format,
365        );
366
367        Ok(frame_buffer)
368    }
369
370    /// Save the frame buffer as an image to the specified path.
371    ///
372    /// # Arguments
373    ///
374    /// * `path` - The path where the image will be saved.
375    /// * `format` - The ImageFormat of the saved image.
376    ///
377    /// # Returns
378    ///
379    /// An empty Result if successful, or an Error if there was an issue saving the image.
380    #[inline]
381    pub fn save_as_image<T: AsRef<Path>>(
382        &mut self,
383        path: T,
384        format: ImageFormat,
385    ) -> Result<(), Error> {
386        let mut frame_buffer = self.buffer()?;
387
388        frame_buffer.save_as_image(path, format)?;
389
390        Ok(())
391    }
392}
393
394/// Represents a frame buffer containing pixel data.
395///
396/// # Example
397/// ```ignore
398/// // Get frame from the capture session
399/// let mut buffer = frame.buffer()?;
400/// buffer.save_as_image("screenshot.png", ImageFormat::Png)?;
401/// ```
402pub struct FrameBuffer<'a> {
403    raw_buffer: &'a mut [u8],
404    buffer: &'a mut Vec<u8>,
405    width: u32,
406    height: u32,
407    row_pitch: u32,
408    depth_pitch: u32,
409    color_format: ColorFormat,
410}
411
412impl<'a> FrameBuffer<'a> {
413    /// Create a new Frame Buffer.
414    ///
415    /// # Arguments
416    ///
417    /// * `raw_buffer` - A mutable reference to the raw pixel data buffer.
418    /// * `buffer` - A mutable reference to the buffer used for copying pixel data without padding.
419    /// * `width` - The width of the frame buffer.
420    /// * `height` - The height of the frame buffer.
421    /// * `row_pitch` - The row pitch of the frame buffer.
422    /// * `depth_pitch` - The depth pitch of the frame buffer.
423    /// * `color_format` - The color format of the frame buffer.
424    ///
425    /// # Returns
426    ///
427    /// A new `FrameBuffer` instance.
428    #[must_use]
429    #[inline]
430    pub fn new(
431        raw_buffer: &'a mut [u8],
432        buffer: &'a mut Vec<u8>,
433        width: u32,
434        height: u32,
435        row_pitch: u32,
436        depth_pitch: u32,
437        color_format: ColorFormat,
438    ) -> Self {
439        Self {
440            raw_buffer,
441            buffer,
442            width,
443            height,
444            row_pitch,
445            depth_pitch,
446            color_format,
447        }
448    }
449
450    /// Get the width of the frame buffer.
451    #[must_use]
452    #[inline]
453    pub const fn width(&self) -> u32 {
454        self.width
455    }
456
457    /// Get the height of the frame buffer.
458    #[must_use]
459    #[inline]
460    pub const fn height(&self) -> u32 {
461        self.height
462    }
463
464    /// Get the row pitch of the frame buffer.
465    #[must_use]
466    #[inline]
467    pub const fn row_pitch(&self) -> u32 {
468        self.row_pitch
469    }
470
471    /// Get the depth pitch of the frame buffer.
472    #[must_use]
473    #[inline]
474    pub const fn depth_pitch(&self) -> u32 {
475        self.depth_pitch
476    }
477
478    /// Check if the buffer has padding.
479    #[must_use]
480    #[inline]
481    pub const fn has_padding(&self) -> bool {
482        self.width * 4 != self.row_pitch
483    }
484
485    /// Get the raw pixel data with possible padding.
486    #[must_use]
487    #[inline]
488    pub fn as_raw_buffer(&mut self) -> &mut [u8] {
489        self.raw_buffer
490    }
491
492    /// Get the raw pixel data without padding.
493    ///
494    /// # Returns
495    ///
496    /// A mutable reference to the buffer containing pixel data without padding.
497    #[inline]
498    pub fn as_nopadding_buffer(&mut self) -> Result<&mut [u8], Error> {
499        if !self.has_padding() {
500            return Ok(self.raw_buffer);
501        }
502
503        let multiplyer = match self.color_format {
504            ColorFormat::Rgba16F => 8,
505            ColorFormat::Rgba8 => 4,
506            ColorFormat::Bgra8 => 4,
507        };
508
509        let frame_size = (self.width * self.height * multiplyer) as usize;
510        if self.buffer.capacity() < frame_size {
511            self.buffer.resize(frame_size, 0);
512        }
513
514        let width_size = (self.width * multiplyer) as usize;
515        let buffer_address = self.buffer.as_mut_ptr() as isize;
516        (0..self.height).into_par_iter().for_each(|y| {
517            let index = (y * self.row_pitch) as usize;
518            let ptr = buffer_address as *mut u8;
519
520            unsafe {
521                ptr::copy_nonoverlapping(
522                    self.raw_buffer.as_ptr().add(index),
523                    ptr.add(y as usize * width_size),
524                    width_size,
525                );
526            }
527        });
528
529        Ok(&mut self.buffer[0..frame_size])
530    }
531
532    /// Save the frame buffer as an image to the specified path.
533    ///
534    /// # Arguments
535    ///
536    /// * `path` - The path where the image will be saved.
537    /// * `format` - The image format to use for saving.
538    ///
539    /// # Returns
540    ///
541    /// An `Ok` result if the image is successfully saved, or an `Err` result if there was an error.
542    #[inline]
543    pub fn save_as_image<T: AsRef<Path>>(
544        &mut self,
545        path: T,
546        format: ImageFormat,
547    ) -> Result<(), Error> {
548        let width = self.width;
549        let height = self.height;
550
551        let bytes = ImageEncoder::new(format, self.color_format).encode(
552            self.as_nopadding_buffer()?,
553            width,
554            height,
555        )?;
556
557        fs::write(path, bytes)?;
558
559        Ok(())
560    }
561}
windows_capture/frame.rs

windows_capture/
frame.rs
