use std::{error::Error, path::Path};
use image::{Rgba, RgbaImage};
use ndarray::{s, ArrayBase, ArrayView, Dim, OwnedRepr};
use windows::Win32::Graphics::Direct3D11::{
ID3D11DeviceContext, ID3D11Texture2D, D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ,
};
use crate::settings::ColorFormat;
#[derive(thiserror::Error, Eq, PartialEq, Clone, Copy, Debug)]
pub enum FrameError {
#[error("Graphics Capture API Is Not Supported")]
InvalidSize,
}
pub struct Frame<'a> {
texture: ID3D11Texture2D,
frame_surface: ID3D11Texture2D,
context: &'a ID3D11DeviceContext,
width: u32,
height: u32,
color_format: ColorFormat,
}
impl<'a> Frame<'a> {
#[must_use]
pub const fn new(
texture: ID3D11Texture2D,
frame_surface: ID3D11Texture2D,
context: &'a ID3D11DeviceContext,
width: u32,
height: u32,
color_format: ColorFormat,
) -> Self {
Self {
texture,
frame_surface,
context,
width,
height,
color_format,
}
}
#[must_use]
pub const fn width(&self) -> u32 {
self.width
}
#[must_use]
pub const fn height(&self) -> u32 {
self.height
}
pub fn buffer(&mut self) -> Result<FrameBuffer, Box<dyn Error + Send + Sync>> {
unsafe {
self.context
.CopyResource(&self.texture, &self.frame_surface);
};
let mut mapped_resource = D3D11_MAPPED_SUBRESOURCE::default();
unsafe {
self.context.Map(
&self.texture,
0,
D3D11_MAP_READ,
0,
Some(&mut mapped_resource),
)?;
};
let mapped_frame_data = unsafe {
std::slice::from_raw_parts(
mapped_resource.pData as *const u8,
(self.height * mapped_resource.RowPitch) as usize,
)
};
let frame_buffer = FrameBuffer::new(mapped_frame_data, self.width, self.height);
Ok(frame_buffer)
}
pub fn save_as_image<T: AsRef<Path>>(
&mut self,
path: T,
) -> Result<(), Box<dyn Error + Send + Sync>> {
let buffer = self.buffer()?;
let nopadding_buffer = buffer.as_raw_nopadding_buffer()?;
let (height, width, _) = nopadding_buffer.dim();
let mut rgba_image: RgbaImage = RgbaImage::new(width as u32, height as u32);
if self.color_format == ColorFormat::Rgba8 {
for y in 0..height {
for x in 0..width {
let r = nopadding_buffer[(y, x, 0)];
let g = nopadding_buffer[(y, x, 1)];
let b = nopadding_buffer[(y, x, 2)];
let a = nopadding_buffer[(y, x, 3)];
rgba_image.put_pixel(x as u32, y as u32, Rgba([r, g, b, a]));
}
}
} else {
for y in 0..height {
for x in 0..width {
let b = nopadding_buffer[(y, x, 0)];
let g = nopadding_buffer[(y, x, 1)];
let r = nopadding_buffer[(y, x, 2)];
let a = nopadding_buffer[(y, x, 3)];
rgba_image.put_pixel(x as u32, y as u32, Rgba([r, g, b, a]));
}
}
}
rgba_image.save(path)?;
Ok(())
}
}
pub struct FrameBuffer<'a> {
raw_buffer: &'a [u8],
width: u32,
height: u32,
}
impl<'a> FrameBuffer<'a> {
#[must_use]
pub const fn new(raw_buffer: &'a [u8], width: u32, height: u32) -> Self {
Self {
raw_buffer,
width,
height,
}
}
#[must_use]
pub const fn width(&self) -> u32 {
self.width
}
#[must_use]
pub const fn height(&self) -> u32 {
self.height
}
#[must_use]
pub const fn has_padding(&self) -> bool {
let raw_buffer = self.as_raw_buffer();
self.width as usize * 4 != raw_buffer.len() / self.height as usize
}
#[must_use]
pub const fn as_raw_buffer(&self) -> &'a [u8] {
self.raw_buffer
}
#[allow(clippy::type_complexity)]
pub fn as_raw_nopadding_buffer(
&self,
) -> Result<ArrayBase<OwnedRepr<u8>, Dim<[usize; 3]>>, Box<dyn Error + Send + Sync>> {
let row_pitch = self.raw_buffer.len() / self.height as usize;
let array =
ArrayView::from_shape((self.height as usize, row_pitch), self.raw_buffer)?.to_owned();
if self.width as usize * 4 == self.raw_buffer.len() / self.height as usize {
let array = array.into_shape((self.height as usize, self.width as usize, 4))?;
Ok(array)
} else {
let array = array
.slice_move(s![.., ..self.width as usize * 4])
.into_shape((self.height as usize, self.width as usize, 4))?;
Ok(array)
}
}
}