janis/grisly
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

grim commandline option support

Browse source
This commit is contained in:
Janis 2024-11-27 07:19:34 +01:00
parent 28a882852d
commit c82b13251d
2 changed files with 410 additions and 109 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
Cargo.toml
View file

@ -6,7 +6,7 @@ edition = "2021"
[dependencies]
anyhow = "1.0.93"
bytemuck = "1.20.0"
clap = "4.5.21"
clap = { version = "4.5.21", features = ["cargo", "derive"] }
env_logger = "0.11.5"
image = "0.25.5"
libc = "0.2.165"

483
src/main.rs
View file

@ -1,10 +1,14 @@
use std::{
collections::HashMap,
default,
ffi::CString,
io::BufWriter,
os::fd::{AsFd, AsRawFd, FromRawFd},
};
use clap::{
arg, builder::TypedValueParser, command, value_parser, ArgAction, Parser,
};
use image::{math::Rect, GenericImage, GenericImageView};
use rand::Rng;
use wayland_client::{
@ -20,8 +24,12 @@ use wayland_client::{
Connection, Dispatch, Proxy, QueueHandle, WEnum,
};
use wayland_protocols::xdg::xdg_output::zv1::client::{
zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1,
use wayland_protocols::{
wp::input_timestamps::zv1::client::__interfaces::zwp_input_timestamps_manager_v1_interface,
xdg::xdg_output::zv1::client::{
zxdg_output_manager_v1::ZxdgOutputManagerV1,
zxdg_output_v1::ZxdgOutputV1,
},
};
use wayland_protocols_wlr::screencopy::v1::client::{
zwlr_screencopy_frame_v1::{self, ZwlrScreencopyFrameV1},
@ -43,11 +51,15 @@ enum FrameData {
Copying {
buffer: WlBuffer,
info: BufferInfo,
rect: Rect,
offset: usize,
scale: f32,
},
Copied {
info: BufferInfo,
rect: Rect,
offset: usize,
scale: f32,
},
Done,
}
@ -70,6 +82,7 @@ struct Output {
logical_position: Option<(i32, i32)>,
physical_position: Option<(i32, i32)>,
transform: Option<Transform>,
name: Option<String>,
scale: Option<i32>,
frame: Frame,
}
@ -89,10 +102,16 @@ impl Output {
output_data: U,
xdg_manager: &ZxdgOutputManagerV1,
screencopy: &ZwlrScreencopyManagerV1,
with_cursor: bool,
) -> Self {
let output = registry.bind(name, version, &qh, output_data);
let xdg_output = xdg_manager.get_xdg_output(&output, &qh, output.id());
let frame = screencopy.capture_output(0, &output, &qh, output.id());
let frame = screencopy.capture_output(
with_cursor as i32,
&output,
&qh,
output.id(),
);
Output {
output,
xdg_output,
@ -102,6 +121,7 @@ impl Output {
scale: Some(1),
logical_position: None,
physical_size: None,
name: None,
frame: Frame {
frame,
inverted_y: false,
@ -113,6 +133,7 @@ impl Output {
#[derive(Debug)]
struct App {
with_cursor: bool,
screencopy_mgr: ZwlrScreencopyManagerV1,
xdg_manager: ZxdgOutputManagerV1,
outputs: HashMap<ObjectId, Output>,
@ -170,11 +191,19 @@ impl Dispatch<ZwlrScreencopyFrameV1, ObjectId> for App {
if let FrameData::Copying {
buffer,
info,
rect,
offset,
} = core::mem::replace(&mut output.frame.data, FrameData::Done)
scale,
} =
core::mem::replace(&mut output.frame.data, FrameData::Done)
{
buffer.destroy();
output.frame.data = FrameData::Copied { info, offset };
output.frame.data = FrameData::Copied {
info,
offset,
rect,
scale,
};
}
}
zwlr_screencopy_frame_v1::Event::Failed => {
@ -272,6 +301,9 @@ impl Dispatch<wl_output::WlOutput, u32> for App {
wl_output::Event::Scale { factor } => {
output.scale = Some(factor);
}
wl_output::Event::Name { name } => {
output.name = Some(name);
}
_ => {}
}
}
@ -332,21 +364,25 @@ impl Dispatch<wl_registry::WlRegistry, GlobalListContents> for App {
name,
&state.xdg_manager,
&state.screencopy_mgr,
state.with_cursor,
);
state.outputs.insert(output.output.id(), output);
}
wl_registry::Event::GlobalRemove { name } => globals.with_list(|globals| {
if globals
.iter()
.find(|global| global.name == name)
.map(|global| &global.interface == wl_output::WlOutput::interface().name)
== Some(true)
wl_registry::Event::GlobalRemove { name } => {
globals.with_list(|globals| {
if globals.iter().find(|global| global.name == name).map(
|global| {
&global.interface
== wl_output::WlOutput::interface().name
},
) == Some(true)
{
state
.outputs
.retain(|_, output| output.output.data() != Some(&name));
state.outputs.retain(|_, output| {
output.output.data() != Some(&name)
});
}
})
}
}),
_ => {}
}
}
@ -400,7 +436,9 @@ impl Shm {
)
} {
libc::MAP_FAILED => Err(std::io::Error::last_os_error()),
ptr => Ok(unsafe { core::slice::from_raw_parts_mut(ptr as *mut _, self.size) }),
ptr => Ok(unsafe {
core::slice::from_raw_parts_mut(ptr as *mut _, self.size)
}),
}
}
}
@ -422,7 +460,154 @@ impl Drop for Shm {
// wl_output gives a transform which is applied to buffer contents.
//
#[derive(clap::Parser)]
#[command(name = "grisly")]
#[command(version = "1.0")]
#[command(
about = "grim compatible screenshot tool for wlroots-based wayland compositors."
)]
struct Config {
/// Specify the region to capture.
#[arg(short = 'g', value_parser = RectParser)]
geometry: Option<Rect>,
/// Use the physical size of the buffers, not the logical size of the outputs.
///
/// Will in most cases be the same as the physical size of the monitor, ignoring any fractional scaling via xdg protocols.
#[arg(long)]
physical_size: bool,
/// Ignore any transforms (flipping, rotating) applied by the compositors.
///
/// This means that the resulting screenshot will look as if all the monitors were all in their normal orientation.
#[arg(long)]
ignore_transforms: bool,
/// Include the cursor in screenshots.
#[arg(short = 'c')]
with_cursor: bool,
/// Specify the wayland output (monitor) to capture from.
#[arg(long, short = 'o')]
output: Option<String>,
/// Specify the output file format. One of png, jpg, ppm, webp, farbfeld. Defaults to png.
#[arg(short = 't', value_parser = ImageTypeParser, default_value = "png")]
kind: ImageType,
/// Specify the quality of the output file from 0 to 100. Defaults to 80.
#[arg(short = 'q', default_value = "80")]
jpeg_quality: u32,
/// For PNG: specify the compression level: fast, best, default.
#[arg(short = 'l', default_value = "default", value_parser = CompressionTypeParser)]
png_level: image::codecs::png::CompressionType,
}
#[derive(Clone, Copy, Debug, Default)]
enum ImageType {
#[default]
Png,
Jpeg,
Ppm,
WebP,
Farbfeld,
}
#[derive(Clone, Copy, Debug)]
struct CompressionTypeParser;
impl TypedValueParser for CompressionTypeParser {
type Value = image::codecs::png::CompressionType;
fn parse_ref(
&self,
_cmd: &clap::Command,
_arg: Option<&clap::Arg>,
value: &std::ffi::OsStr,
) -> Result<Self::Value, clap::Error> {
let value = value
.to_str()
.ok_or(clap::Error::new(clap::error::ErrorKind::InvalidUtf8))?;
match value {
"best" => Ok(image::codecs::png::CompressionType::Best),
"fast" => Ok(image::codecs::png::CompressionType::Fast),
"default" => Ok(image::codecs::png::CompressionType::Default),
_ => Err(clap::Error::new(clap::error::ErrorKind::ValueValidation)),
}
}
}
#[derive(Clone, Copy, Debug)]
struct ImageTypeParser;
impl TypedValueParser for ImageTypeParser {
type Value = ImageType;
fn parse_ref(
&self,
_cmd: &clap::Command,
_arg: Option<&clap::Arg>,
value: &std::ffi::OsStr,
) -> Result<Self::Value, clap::Error> {
let value = value
.to_str()
.ok_or(clap::Error::new(clap::error::ErrorKind::InvalidUtf8))?;
match value {
"png" => Ok(ImageType::Png),
"jpg" => Ok(ImageType::Jpeg),
"ppm" => Ok(ImageType::Ppm),
"webp" => Ok(ImageType::WebP),
"farbfeld" => Ok(ImageType::Farbfeld),
_ => Err(clap::Error::new(clap::error::ErrorKind::ValueValidation)),
}
}
}
#[derive(Clone, Copy, Debug)]
struct RectParser;
impl TypedValueParser for RectParser {
type Value = image::math::Rect;
fn parse_ref(
&self,
_cmd: &clap::Command,
_arg: Option<&clap::Arg>,
value: &std::ffi::OsStr,
) -> Result<Self::Value, clap::Error> {
let value = value
.to_str()
.ok_or(clap::Error::new(clap::error::ErrorKind::InvalidUtf8))?;
let (pos, size) = value
.split_once(' ')
.ok_or(clap::Error::new(clap::error::ErrorKind::ValueValidation))?;
let (x, y) = pos
.split_once(',')
.ok_or(clap::Error::new(clap::error::ErrorKind::ValueValidation))?;
let (width, height) = size
.split_once('x')
.ok_or(clap::Error::new(clap::error::ErrorKind::ValueValidation))?;
let x = x.parse::<u32>().map_err(|_| {
clap::Error::new(clap::error::ErrorKind::ValueValidation)
})?;
let y = y.parse::<u32>().map_err(|_| {
clap::Error::new(clap::error::ErrorKind::ValueValidation)
})?;
let width = width.parse::<u32>().map_err(|_| {
clap::Error::new(clap::error::ErrorKind::ValueValidation)
})?;
let height = height.parse::<u32>().map_err(|_| {
clap::Error::new(clap::error::ErrorKind::ValueValidation)
})?;
Ok(Rect {
x,
y,
width,
height,
})
}
}
fn main() {
let config = Config::parse();
let conn = Connection::connect_to_env().unwrap();
let (globals, mut queue) = registry_queue_init(&conn).unwrap();
@ -438,7 +623,9 @@ fn main() {
let outputs = globals.contents().with_list(|list| {
list.iter()
.filter(|global| &global.interface == wl_output::WlOutput::interface().name)
.filter(|global| {
&global.interface == wl_output::WlOutput::interface().name
})
.map(|global| {
let output = Output::bind(
globals.registry(),
@ -448,6 +635,7 @@ fn main() {
global.name,
&xdg_manager,
&screencopy_mgr,
config.with_cursor,
);
(output.output.id(), output)
@ -462,6 +650,7 @@ fn main() {
xdg_manager,
screencopy_mgr,
shm,
with_cursor: config.with_cursor,
};
while !app
@ -473,20 +662,67 @@ fn main() {
}
let mut total_size = 0;
let mut union_rect = Rect {
x: 0,
y: 0,
width: 0,
height: 0,
};
let buffers = app
.outputs
.iter_mut()
.filter(|(_, output)| {
if let Some(output_name) = config.output.as_ref() {
output.name.as_ref() == Some(output_name)
} else {
true
}
})
.map(|(_, output)| {
let (px, py) = output.physical_position.unwrap_or((0, 0));
let (lx, ly) = output.logical_position.unwrap_or((0, 0));
// image destination
let (x, y) = (lx - px, ly - py);
let FrameData::Buffer(info) =
core::mem::replace(&mut output.frame.data, FrameData::Done)
else {
unreachable!();
};
let (width, height) = if config.physical_size {
(info.width as i32, info.height as i32)
} else {
output
.logical_size
.unwrap_or((info.width as i32, info.height as i32))
};
let scale = width as f32 / info.width as f32;
let rect = image::math::Rect {
x: x as u32,
y: y as u32,
width: width as u32,
height: height as u32,
};
(rect, info, output, scale)
})
.filter(|(rect, _, _, _)| {
if let Some(ref geom) = config.geometry {
rect_overlaps(*rect, *geom)
} else {
true
}
})
.map(|(rect, info, output, scale)| {
let size = info.height * info.stride;
let offset = total_size;
total_size += size;
union_rect = rect_union(union_rect, rect);
(output, info, offset)
(output, info, rect, offset, scale)
})
.collect::<Vec<_>>();
@ -497,7 +733,7 @@ fn main() {
.shm
.create_pool(shm.fd.as_fd(), total_size as i32, &qh, ());
for (output, info, offset) in buffers {
for (output, info, rect, offset, scale) in buffers {
let buffer = pool.create_buffer(
offset as i32,
info.width as i32,
@ -513,25 +749,61 @@ fn main() {
output.frame.data = FrameData::Copying {
buffer,
info,
rect,
scale,
offset: offset as usize,
};
}
while !app
.outputs
.iter()
.all(|(_, output)| matches!(output.frame.data, FrameData::Copied { .. }))
.values()
.all(|output| !matches!(output.frame.data, FrameData::Copying { .. }))
{
_ = queue.roundtrip(&mut app).unwrap();
}
fn rect_union(lhs: Rect, rhs: Rect) -> Rect {
struct AABB {
ax: u32,
ay: u32,
bx: u32,
by: u32,
}
fn rect_overlaps(lhs: Rect, rhs: Rect) -> bool {
let intersection = rect_intersection(lhs, rhs);
!(intersection.width == 0 || intersection.height == 0)
}
fn rect_intersection(lhs: Rect, rhs: Rect) -> Rect {
let lhs = AABB {
ax: lhs.x,
ay: lhs.y,
bx: lhs.x + lhs.width,
by: lhs.y + lhs.height,
};
let rhs = AABB {
ax: rhs.x,
ay: rhs.y,
bx: rhs.x + rhs.width,
by: rhs.y + rhs.height,
};
let intersection = AABB {
ax: lhs.ax.max(rhs.ax),
ay: lhs.ay.max(rhs.ay),
bx: lhs.bx.min(rhs.bx),
by: lhs.by.min(rhs.by),
};
Rect {
x: intersection.ax,
y: intersection.ay,
width: intersection.bx.saturating_sub(intersection.ax),
height: intersection.by.saturating_sub(intersection.ay),
}
}
fn rect_union(lhs: Rect, rhs: Rect) -> Rect {
let lhs = AABB {
ax: lhs.x,
ay: lhs.y,
@ -559,77 +831,99 @@ fn main() {
}
}
let mut image_rect = Rect {
x: 0,
y: 0,
width: 0,
height: 0,
};
let buffers = app
.outputs
.values_mut()
.map(|output| {
let FrameData::Copied { info, offset, .. } =
core::mem::replace(&mut output.frame.data, FrameData::Done)
else {
unreachable!();
};
let (px, py) = output.physical_position.unwrap_or((0, 0));
let (lx, ly) = output.logical_position.unwrap_or((0, 0));
// image destination
let (x, y) = (lx - px, ly - py);
let (width, height) = output
.logical_size
.unwrap_or((info.width as i32, info.height as i32));
let rect = image::math::Rect {
x: x as u32,
y: y as u32,
width: width as u32,
height: height as u32,
};
println!("buffer: {:?}", info);
println!(" - invert-y: {}", output.frame.inverted_y);
println!("destination rect: {:?}", rect);
println!(" - transform: {:?}", output.transform);
image_rect = rect_union(image_rect, rect);
(
info,
offset,
rect,
output.frame.inverted_y,
output.transform.unwrap_or(wl_output::Transform::Normal),
)
})
.collect::<Vec<_>>();
let image_rect = config.geometry.unwrap_or(union_rect);
let mut image = image::RgbaImage::new(image_rect.width, image_rect.height);
let geom = config.geometry.unwrap_or(image_rect);
println!("image: {:?}", image_rect);
println!(" - geom: {:?}", config.geometry);
let now = std::time::Instant::now();
for (info, offset, rect, inverted_y, transform) in buffers {
app.outputs
.values_mut()
.filter_map(|output| {
match core::mem::replace(&mut output.frame.data, FrameData::Done) {
FrameData::Copied {
info,
rect,
offset,
scale,
} => Some((output, info, rect, offset, scale)),
_ => None,
}
})
.for_each(|(output, info, rect, offset, scale)| {
let transform = if config.ignore_transforms {
wl_output::Transform::Normal
} else {
output.transform.unwrap_or(wl_output::Transform::Normal)
};
let size = info.stride * info.height;
// get intersection of selected geometry and this output.
// output rect `rect` is already in 'global space'
// view() needs buffer-local space coordinates
// which are the diff between rect and section_rect
let section_rect = rect_intersection(geom, rect);
let section = Rect {
x: section_rect.x - rect.x,
y: section_rect.y - rect.y,
width: section_rect.width,
height: section_rect.height,
};
let view = BufferView {
info,
inverted_y,
inverted_y: output.frame.inverted_y,
transform,
crop: section,
bytes: &buffer[offset..][..size as usize],
};
if rect.width != view.width() || rect.height != view.height() {
let sampled = SampledView {
view: &view,
// println!("buffer: {:?}", info);
// println!(" - invert-y: {}", output.frame.inverted_y);
// println!("destination rect: {:?}", rect);
// println!(" - section_rect: {:?}", section_rect);
// println!(" - section: {:?}", section);
// println!(" - transform: {:?}", output.transform);
// println!("view: {:?}", view.crop);
if scale != 1.0 {
let view = ScaledView {
view: &BufferView {
info,
inverted_y: output.frame.inverted_y,
transform,
crop: Rect {
x: 0,
y: 0,
width: info.width,
height: info.height,
},
bytes: &buffer[offset..][..size as usize],
},
crop: section,
dimensions: (rect.width, rect.height),
};
image.copy_from(&sampled, rect.x, rect.y).unwrap();
image
.copy_from(
&view,
section_rect.x - image_rect.x,
section_rect.y - image_rect.y,
)
.unwrap();
} else {
image.copy_from(&view, rect.x, rect.y).unwrap();
}
image
.copy_from(
&view,
section_rect.x - image_rect.x,
section_rect.y - image_rect.y,
)
.unwrap();
}
});
println!("copying: {}s", now.elapsed().as_secs_f64());
let file = std::fs::File::create("out.png").unwrap();
@ -645,21 +939,22 @@ fn main() {
println!("encoding: {}s", now.elapsed().as_secs_f64());
}
struct SampledView<'a, V: GenericImageView> {
struct ScaledView<'a, V: GenericImageView> {
view: &'a V,
crop: Rect,
dimensions: (u32, u32),
}
impl<'a, V: GenericImageView> GenericImageView for SampledView<'a, V> {
impl<'a, V: GenericImageView> GenericImageView for ScaledView<'a, V> {
type Pixel = V::Pixel;
fn dimensions(&self) -> (u32, u32) {
self.dimensions
(self.crop.width, self.crop.height)
}
fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel {
let u = x as f32 / self.dimensions.0 as f32;
let v = y as f32 / self.dimensions.1 as f32;
let u = (x + self.crop.x) as f32 / self.dimensions.0 as f32;
let v = (y + self.crop.y) as f32 / self.dimensions.1 as f32;
image::imageops::sample_bilinear(self.view, u, v).unwrap()
}
@ -668,6 +963,7 @@ impl<'a, V: GenericImageView> GenericImageView for SampledView<'a, V> {
struct BufferView<'a> {
info: BufferInfo,
inverted_y: bool,
crop: Rect,
transform: wl_output::Transform,
bytes: &'a [u8],
}
@ -678,6 +974,8 @@ impl BufferView<'_> {
let width = width - 1;
let height = height - 1;
let (x, y) = (x + self.crop.x, y + self.crop.y);
let (x, y) = match self.transform {
Transform::Normal => (x, y),
Transform::_90 => (y, x),
@ -701,12 +999,14 @@ impl BufferView<'_> {
}
fn transformed_dimensions(&self) -> (u32, u32) {
match self.transform {
Transform::Flipped180 | Transform::_180 | Transform::Flipped | Transform::Normal => {
(self.info.width, self.info.height)
}
Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => {
(self.info.height, self.info.width)
}
Transform::Flipped180
| Transform::_180
| Transform::Flipped
| Transform::Normal => (self.crop.width, self.crop.height),
Transform::_90
| Transform::_270
| Transform::Flipped90
| Transform::Flipped270 => (self.crop.height, self.crop.width),
_ => unreachable!(),
}
}
@ -714,7 +1014,8 @@ impl BufferView<'_> {
// apply transform
let (x, y) = self.transformed_pixel_position(x, y);
self.info.stride as usize * y as usize + (x as usize * self.pixel_stride() as usize)
self.info.stride as usize * y as usize
+ (x as usize * self.pixel_stride() as usize)
}
fn pixel_stride(&self) -> u32 {