rigela_utils/
screen.rs

1/*
2 * Copyright (c) 2024. The RigelA open source project team and
3 * its contributors reserve all rights.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 * http://www.apache.org/licenses/LICENSE-2.0
9 * Unless required by applicable law or agreed to in writing, software distributed under the
10 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 * See the License for the specific language governing permissions and limitations under the License.
12 */
13
14use png::{AdaptiveFilterType, BitDepth, ColorType, Compression, Encoder};
15use std::io::Write;
16use win_wrap::{
17    common::HWND,
18    ext::ToBytesExt,
19    graphic::{
20        bit_blt, create_compatible_bitmap, create_compatible_dc, delete_dc, delete_object, get_dc,
21        get_di_bits, release_dc, select_object, BITMAPFILEHEADER, BITMAPINFOHEADER, BI_RGB,
22        DIB_RGB_COLORS, RGBQUAD, SRCCOPY,
23    },
24};
25
26/**
27图像压缩格式。
28*/
29pub enum ImageCompressionFormat {
30    BMP,
31    PNG,
32}
33
34/**
35屏幕截图，返回未压缩的位图数据。
36`h_wnd` 窗口句柄，0表示整个屏幕。
37`left` 要截取的左上角x坐标。
38`top` 要截取的左上角y坐标。
39`width` 要截取的宽度。
40`height` 要截取的高度。
41*/
42pub fn snapshot(
43    h_wnd: Option<HWND>,
44    left: i32,
45    top: i32,
46    with: i32,
47    height: i32,
48) -> Option<(Vec<Vec<RGBQUAD>>, BITMAPINFOHEADER, Option<Vec<RGBQUAD>>)> {
49    let h_dc = get_dc(h_wnd);
50    if h_dc.is_invalid() {
51        return None;
52    }
53    let h_mem_dc = create_compatible_dc(Some(h_dc));
54    if h_mem_dc.is_invalid() {
55        release_dc(h_wnd, h_dc);
56        return None;
57    }
58    let h_bm = create_compatible_bitmap(h_dc, with, height);
59    if h_bm.is_invalid() {
60        delete_dc(h_mem_dc);
61        release_dc(h_wnd, h_dc);
62    }
63    let h_old_obj = select_object(h_mem_dc, h_bm.into());
64    let res = bit_blt(h_mem_dc, 0, 0, with, height, Some(h_dc), left, top, SRCCOPY);
65    let data = if res {
66        let (pixels, mut bm_header, color_table) =
67            get_di_bits(h_dc, h_bm, 0, height as u32, DIB_RGB_COLORS);
68        bm_header.biCompression = BI_RGB.0;
69        Some((pixels, bm_header, color_table))
70    } else {
71        None
72    };
73    if !h_old_obj.is_invalid() {
74        select_object(h_mem_dc, h_old_obj);
75    }
76    delete_object(h_bm.into());
77    delete_dc(h_mem_dc);
78    release_dc(h_wnd, h_dc);
79    data
80}
81
82/**
83屏幕截图，返回二进制bytes数据。
84`h_wnd` 窗口句柄，0表示整个屏幕。
85`left` 要截取的左上角x坐标。
86`top` 要截取的左上角y坐标。
87`width` 要截取的宽度。
88`height` 要截取的高度。
89`format` 图像格式。
90*/
91pub fn snapshot_bytes(
92    h_wnd: Option<HWND>,
93    left: i32,
94    top: i32,
95    width: i32,
96    height: i32,
97    format: ImageCompressionFormat,
98) -> Option<Vec<u8>> {
99    let Some((mut pixels, bm_header, color_table)) = snapshot(h_wnd, left, top, width, height)
100    else {
101        return None;
102    };
103    match format {
104        ImageCompressionFormat::BMP => {
105            let offset = bm_header.biSize
106                + size_of::<BITMAPFILEHEADER>() as u32
107                + bm_header.biClrUsed * size_of::<RGBQUAD>() as u32;
108            let header = BITMAPFILEHEADER {
109                bfType: 0x4d42,
110                bfSize: bm_header.biSizeImage + offset,
111                bfReserved1: 0,
112                bfReserved2: 0,
113                bfOffBits: offset,
114            };
115            let mut v = Vec::<u8>::with_capacity((bm_header.biSize + offset) as usize);
116            v.write(header.to_bytes()).unwrap();
117            v.write(bm_header.to_bytes()).unwrap();
118            if let Some(color_table) = color_table {
119                for i in color_table.iter() {
120                    v.write(i.to_bytes()).unwrap();
121                }
122            }
123            for i in pixels.iter() {
124                for j in i.iter() {
125                    v.write(j.to_bytes()).unwrap();
126                }
127            }
128            Some(v)
129        }
130        ImageCompressionFormat::PNG => {
131            let mut v2 = Vec::with_capacity((bm_header.biSize + bm_header.biSizeImage) as usize);
132            let mut encoder =
133                Encoder::new(&mut v2, bm_header.biWidth as u32, bm_header.biHeight as u32);
134            encoder.set_color(ColorType::Rgba);
135            encoder.set_depth(BitDepth::Eight);
136            encoder.set_compression(Compression::Best);
137            encoder.set_adaptive_filter(AdaptiveFilterType::Adaptive);
138            if let Ok(mut w) = encoder.write_header() {
139                let mut buf = Vec::with_capacity(bm_header.biSizeImage as usize);
140                if bm_header.biHeight > 0 {
141                    // 高度如果是正数则图像是倒立的，需要翻转
142                    pixels.reverse();
143                }
144                for i in pixels.iter() {
145                    for j in i.iter() {
146                        buf.write(j.to_bytes()).unwrap();
147                    }
148                }
149                w.write_image_data(&buf).unwrap();
150                w.finish().unwrap();
151            }
152            Some(v2)
153        }
154    }
155}
156
157#[cfg(test)]
158mod test_screen {
159    use std::{fs::OpenOptions, io::Write};
160
161    use win_wrap::common::get_desktop_window;
162
163    use crate::screen::{snapshot_bytes, ImageCompressionFormat};
164
165    #[test]
166    fn main() {
167        let Some(data) = snapshot_bytes(
168            Some(get_desktop_window()),
169            20,
170            20,
171            1300,
172            620,
173            ImageCompressionFormat::PNG,
174        ) else {
175            return;
176        };
177        OpenOptions::new()
178            .create(true)
179            .write(true)
180            .open("screen.png")
181            .unwrap()
182            .write(&data)
183            .unwrap();
184    }
185}
rigela_utils/screen.rs

rigela_utils/
screen.rs
