4
6

Rust で Windows の画面をキャプチャする

Last updated at Posted at 2024-08-23

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(&current_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"
] }
4
6
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
4
6