Rust でスクリーンキャプチャするプログラムを書いたので残しておく。
実装方法は Google 検索などで色々調査したが、動かないコードが結構多かった為、これから似たようなものを作る人には参考になると思う。
GDI は使っていないので、Chrome のような画面もキャプチャできる。
Windows.Graphics.Capture を使用している。
キャプチャした結果はテクスチャ(ID3D11Texture2D)として取得する事になるが、その画素値を取得して png で保存する処理まで書いている。
capture.rs
use std::slice::from_raw_parts;
use windows::{core::{factory, IInspectable, Interface, HSTRING}, Foundation::TypedEventHandler, Graphics::{Capture::{Direct3D11CaptureFrame, Direct3D11CaptureFramePool, GraphicsCaptureItem, GraphicsCaptureSession}, DirectX::{Direct3D11::IDirect3DDevice, DirectXPixelFormat}, Imaging::{BitmapAlphaMode, BitmapEncoder, BitmapPixelFormat}}, Storage::{CreationCollisionOption, FileAccessMode, StorageFolder}, Win32::{Foundation::HWND, Graphics::{Direct3D::D3D_DRIVER_TYPE_HARDWARE, Direct3D11::{D3D11CreateDevice, ID3D11Device, ID3D11Texture2D, D3D11_CPU_ACCESS_READ, D3D11_CREATE_DEVICE_BGRA_SUPPORT, D3D11_MAPPED_SUBRESOURCE, D3D11_MAP_READ, D3D11_SDK_VERSION, D3D11_TEXTURE2D_DESC, D3D11_USAGE_STAGING}, Dxgi::IDXGIDevice}, System::WinRT::{Direct3D11::{CreateDirect3D11DeviceFromDXGIDevice, IDirect3DDxgiInterfaceAccess}, Graphics::Capture::IGraphicsCaptureItemInterop}}};
pub fn create_d3d11_device() -> Result<ID3D11Device, Box<dyn std::error::Error + Send + Sync>> {
let mut d3d11_device = None;
match unsafe { D3D11CreateDevice(None, D3D_DRIVER_TYPE_HARDWARE, None, D3D11_CREATE_DEVICE_BGRA_SUPPORT, None, D3D11_SDK_VERSION, Some(&mut d3d11_device), None, None) } {
Ok(_) => {
Ok(d3d11_device.unwrap())
},
Err(e) => {
Err(Box::new(e))
},
}
}
pub fn create_d3d11_device_from_dxgi_device(d3d11_device: &ID3D11Device) -> Result<IDirect3DDevice, Box<dyn std::error::Error + Send + Sync>> {
let dxgi: IDXGIDevice = match d3d11_device.cast() {
Ok(dxgi) => {
dxgi
},
Err(e) => {
return Err(Box::new(e));
},
};
let inspectable = match unsafe { CreateDirect3D11DeviceFromDXGIDevice(&dxgi) } {
Ok(inspectable) => {
inspectable
},
Err(e) => {
return Err(Box::new(e));
},
};
let d3d_device: IDirect3DDevice = match inspectable.cast() {
Ok(d3d_device) => {
d3d_device
},
Err(e) => {
return Err(Box::new(e));
},
};
Ok(d3d_device)
}
pub fn create_capture_item_for_window(hwnd: HWND) -> Result<GraphicsCaptureItem, Box<dyn std::error::Error + Send + Sync>> {
let interop = match factory::<GraphicsCaptureItem, IGraphicsCaptureItemInterop>() {
Ok(interop) => {
interop
},
Err(e) => {
return Err(Box::new(e));
},
};
let item: GraphicsCaptureItem = match unsafe { interop.CreateForWindow(hwnd) } {
Ok(item) => {
item
},
Err(e) => {
return Err(Box::new(e));
},
};
Ok(item)
}
pub fn create_free_threaded(d3d_device: &IDirect3DDevice, item: &GraphicsCaptureItem) -> Result<Direct3D11CaptureFramePool, Box<dyn std::error::Error + Send + Sync>> {
match item.Size() {
Ok(size) => {
match Direct3D11CaptureFramePool::CreateFreeThreaded(d3d_device, DirectXPixelFormat::B8G8R8A8UIntNormalized, 1, size) {
Ok(frame_pool) => {
Ok(frame_pool)
},
Err(e) => {
return Err(Box::new(e));
},
}
},
Err(e) => {
return Err(Box::new(e));
},
}
}
pub fn start_capture<F>(item: &GraphicsCaptureItem, frame_pool: &Direct3D11CaptureFramePool, handler: F) -> Result<GraphicsCaptureSession, Box<dyn std::error::Error + Send + Sync>>
where
F: FnMut(&Option<Direct3D11CaptureFramePool>, &Option<IInspectable>) -> windows::core::Result<()> + Send + 'static
{
let handler_stub = TypedEventHandler::<Direct3D11CaptureFramePool, IInspectable>::new(handler);
if let Err(e) = frame_pool.FrameArrived(&handler_stub) {
return Err(Box::new(e));
};
match frame_pool.CreateCaptureSession(item) {
Ok(capture_session) => {
if let Err(e) = capture_session.StartCapture() {
Err(Box::new(e))
} else {
Ok(capture_session)
}
},
Err(e) => {
Err(Box::new(e))
},
}
}
pub fn get_texture2d(frame: &Direct3D11CaptureFrame) -> Result<ID3D11Texture2D, Box<dyn std::error::Error + Send + Sync>> {
match frame.Surface() {
Ok(surface) => {
match surface.cast::<IDirect3DDxgiInterfaceAccess>() {
Ok(access) => {
match unsafe { access.GetInterface() } {
Ok(texture) => {
Ok(texture)
},
Err(e) => {
return Err(Box::new(e));
},
}
},
Err(e) => {
return Err(Box::new(e));
},
}
},
Err(e) => {
return Err(Box::new(e));
},
}
}
pub fn get_bits(d3d11_device: &ID3D11Device, texture: &ID3D11Texture2D) -> Result<(Vec<u8>, D3D11_TEXTURE2D_DESC), Box<dyn std::error::Error + Send + Sync>> {
let mut desc = D3D11_TEXTURE2D_DESC::default();
unsafe {
texture.GetDesc(&mut desc);
}
// 以下の 4 つのプロパティを設定するのがポイント
desc.Usage = D3D11_USAGE_STAGING;
desc.BindFlags = 0;
desc.CPUAccessFlags |= D3D11_CPU_ACCESS_READ.0 as u32;
desc.MiscFlags = 0;
let context = match unsafe { d3d11_device.GetImmediateContext() } {
Ok(context) => {
context
},
Err(e) => {
return Err(Box::new(e));
},
};
// オリジナルのテクスチャは CPU に転送できないため、コピーを作る必要がある
let mut texture_copy = None;
if let Err(e) = unsafe { d3d11_device.CreateTexture2D(&desc, None, Some(&mut texture_copy)) } {
return Err(Box::new(e));
}
let texture_copy = texture_copy.unwrap();
unsafe {
context.CopyResource(&texture_copy, texture);
}
let mut resource = D3D11_MAPPED_SUBRESOURCE::default();
let slice: &[u8] = match unsafe { context.Map(&texture_copy, 0, D3D11_MAP_READ, 0, Some(&mut resource)) } {
Ok(_) => {
unsafe {
from_raw_parts(resource.pData as *const u8, (desc.Height * resource.RowPitch) as usize)
}
},
Err(e) => {
return Err(Box::new(e));
},
};
// slice には余分なデータが含まれているため、必要な部分だけを取得する
// フォーマットは B8G8R8A8 なので、一つの画素は 4 バイトになる
let mut bits = vec![0; (desc.Width * desc.Height * 4) as usize];
for row in 0..desc.Height {
let data_begin = (row * (desc.Width * 4)) as usize;
let data_end = ((row + 1) * (desc.Width * 4)) as usize;
let slice_begin = (row * resource.RowPitch) as usize;
let slice_end = slice_begin + (desc.Width * 4) as usize;
bits[data_begin..data_end].copy_from_slice(&slice[slice_begin..slice_end]);
}
Ok((bits, desc))
}
pub fn save_to_png_file(bits: &Vec<u8>, width: u32, height: u32) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let current_path = match std::env::current_dir() {
Ok(dir) => {
dir.to_string_lossy().to_string()
},
Err(e) => {
return Err(Box::new(e));
},
};
let current_folder = StorageFolder::GetFolderFromPathAsync(&HSTRING::from(¤t_path))?.get()?;
let file = current_folder.CreateFileAsync(&HSTRING::from("キャプチャ.png"), CreationCollisionOption::ReplaceExisting)?.get()?;
let stream = file.OpenAsync(FileAccessMode::ReadWrite)?.get()?;
let encoder = BitmapEncoder::CreateAsync(BitmapEncoder::PngEncoderId()?, &stream)?.get()?;
encoder.SetPixelData(
BitmapPixelFormat::Bgra8,
BitmapAlphaMode::Premultiplied,
width,
height,
1.0,
1.0,
bits,
)?;
encoder.FlushAsync()?.get()?;
Ok(())
}
// もしキャプチャした画像を start_x, start_y 座標から width, height の範囲で切り抜きたい場合はこれを使う
pub fn rectangle(src: &Vec<u8>, src_width: u32, start_x: u32, start_y: u32, width: u32, height: u32) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>> {
let mut rect = vec![0; (width * height * 4) as usize];
for row in 0..height {
let data_begin = (row * (width * 4)) as usize;
let data_end = ((row + 1) * (width * 4)) as usize;
let slice_begin = ((start_y + row) * (src_width * 4) + start_x * 4) as usize;
let slice_end = slice_begin + (width * 4) as usize;
rect[data_begin..data_end].copy_from_slice(&src[slice_begin..slice_end]);
}
Ok(rect)
}
使い方は以下。
簡略化の為エラー処理はしていないので流用する際は注意。
main.rs
use windows::Win32::UI::WindowsAndMessaging::GetForegroundWindow;
mod capture;
fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let d3d11_device = capture::create_d3d11_device()?;
let d3d_device = capture::create_d3d11_device_from_dxgi_device(&d3d11_device)?;
let item = capture::create_capture_item_for_window(unsafe { GetForegroundWindow() })?;
let frame_pool = capture::create_free_threaded(&d3d_device, &item)?;
let _capture_session = capture::start_capture(&item, &frame_pool, move |frame_pool, _inspectable| {
let frame = frame_pool.as_ref().unwrap().TryGetNextFrame()?;
let _size = frame.ContentSize()?;
let texture = capture::get_texture2d(&frame).unwrap();
let (bits, desc) = capture::get_bits(&d3d11_device, &texture).unwrap();
capture::save_to_png_file(&bits, desc.Width, desc.Height).unwrap();
frame_pool.as_ref().unwrap().Close()?;
Ok(())
})?;
std::thread::sleep(std::time::Duration::from_millis(1000));
Ok(())
}
Cargo.toml
[dependencies]
windows = { version = "0.58.0", features = [
"Foundation",
"Win32_Foundation",
"Win32_UI_WindowsAndMessaging",
"Win32_System_WinRT_Direct3D11",
"Win32_System_WinRT_Graphics_Capture",
"Win32_Graphics_Gdi",
"Win32_Graphics_Direct3D",
"Win32_Graphics_Direct3D11",
"Win32_Graphics_Dxgi",
"Win32_Graphics_Dxgi_Common",
"Storage",
"Storage_Streams",
"Graphics_Imaging",
"Graphics_Capture",
"Graphics_DirectX",
"Graphics_DirectX_Direct3D11"
] }