waycap_rs/capture/
video.rs

1use std::{
2    os::fd::{FromRawFd, OwnedFd, RawFd},
3    sync::{atomic::AtomicBool, mpsc, Arc},
4    time::Instant,
5};
6
7use pipewire::{
8    self as pw,
9    context::Context,
10    main_loop::MainLoop,
11    spa::{
12        buffer::{Data, DataType},
13        pod::{Property, PropertyFlags},
14        utils::{Choice, ChoiceEnum, ChoiceFlags, Direction},
15    },
16    stream::{Stream, StreamFlags, StreamState},
17};
18use pw::{properties::properties, spa};
19
20use ringbuf::{traits::Producer, HeapProd};
21use spa::pod::Pod;
22
23use crate::types::video_frame::RawVideoFrame;
24
25use super::Terminate;
26
27// Literally stole these by looking at what OBS uses
28// just magic numbers to me no clue what these are
29// but they enable DMA Buf so it is what it is
30const NVIDIA_MODIFIERS: &[i64] = &[
31    216172782120099856,
32    216172782120099857,
33    216172782120099858,
34    216172782120099859,
35    216172782120099860,
36    216172782120099861,
37    216172782128496656,
38    216172782128496657,
39    216172782128496658,
40    216172782128496659,
41    216172782128496660,
42    216172782128496661,
43    72057594037927935,
44];
45
46pub struct VideoCapture {
47    video_ready: Arc<AtomicBool>,
48    audio_ready: Arc<AtomicBool>,
49    use_nvidia_modifiers: bool,
50}
51
52#[derive(Clone, Copy, Default)]
53struct UserData {
54    video_format: spa::param::video::VideoInfoRaw,
55}
56
57impl VideoCapture {
58    pub fn new(
59        video_ready: Arc<AtomicBool>,
60        audio_ready: Arc<AtomicBool>,
61        use_nvidia_modifiers: bool,
62    ) -> Self {
63        Self {
64            video_ready,
65            audio_ready,
66            use_nvidia_modifiers,
67        }
68    }
69
70    #[allow(clippy::too_many_arguments)]
71    pub fn run(
72        &self,
73        pipewire_fd: RawFd,
74        stream_node: u32,
75        mut ringbuf_producer: HeapProd<RawVideoFrame>,
76        termination_recv: pw::channel::Receiver<Terminate>,
77        saving: Arc<AtomicBool>,
78        start_time: Instant,
79        resolution_negotiation_channel: mpsc::Sender<(u32, u32)>,
80    ) -> Result<(), pipewire::Error> {
81        let pw_loop = MainLoop::new(None)?;
82        let terminate_loop = pw_loop.clone();
83
84        let _recv = termination_recv.attach(pw_loop.loop_(), move |_| {
85            log::debug!("Terminating video capture loop");
86            terminate_loop.quit();
87        });
88
89        let pw_context = Context::new(&pw_loop)?;
90        let core = pw_context.connect_fd(unsafe { OwnedFd::from_raw_fd(pipewire_fd) }, None)?;
91
92        let data = UserData::default();
93
94        let _listener = core
95            .add_listener_local()
96            .info(|i| log::info!("VIDEO CORE:\n{0:#?}", i))
97            .error(|e, f, g, h| log::error!("{0},{1},{2},{3}", e, f, g, h))
98            .done(|d, _| log::info!("DONE: {0}", d))
99            .register();
100
101        // Set up video stream
102        let video_stream = Stream::new(
103            &core,
104            "waycap-video",
105            properties! {
106                *pw::keys::MEDIA_TYPE => "Video",
107                *pw::keys::MEDIA_CATEGORY => "Capture",
108                *pw::keys::MEDIA_ROLE => "Screen",
109            },
110        )?;
111
112        let ready_clone = Arc::clone(&self.video_ready);
113        let audio_ready_clone = Arc::clone(&self.audio_ready);
114        let _video_stream = video_stream
115            .add_local_listener_with_user_data(data)
116            .state_changed(move |_, _, old, new| {
117                log::info!("Video Stream State Changed: {0:?} -> {1:?}", old, new);
118                ready_clone.store(
119                    new == StreamState::Streaming,
120                    std::sync::atomic::Ordering::Release,
121                );
122            })
123            .param_changed(move |_, user_data, id, param| {
124                let Some(param) = param else {
125                    return;
126                };
127
128                if id != pw::spa::param::ParamType::Format.as_raw() {
129                    return;
130                }
131
132                let (media_type, media_subtype) =
133                    match pw::spa::param::format_utils::parse_format(param) {
134                        Ok(v) => v,
135                        Err(_) => return,
136                    };
137
138                if media_type != pw::spa::param::format::MediaType::Video
139                    || media_subtype != pw::spa::param::format::MediaSubtype::Raw
140                {
141                    return;
142                }
143
144                user_data
145                    .video_format
146                    .parse(param)
147                    .expect("Failed to parse param");
148
149                log::debug!(
150                    "  format: {} ({:?})",
151                    user_data.video_format.format().as_raw(),
152                    user_data.video_format.format()
153                );
154
155                resolution_negotiation_channel
156                    .send((
157                        user_data.video_format.size().width,
158                        user_data.video_format.size().height,
159                    ))
160                    .unwrap();
161
162                log::debug!(
163                    "  size: {}x{}",
164                    user_data.video_format.size().width,
165                    user_data.video_format.size().height
166                );
167                log::debug!(
168                    "  framerate: {}/{}",
169                    user_data.video_format.framerate().num,
170                    user_data.video_format.framerate().denom
171                );
172            })
173            .process(move |stream, udata| {
174                match stream.dequeue_buffer() {
175                    None => log::debug!("out of buffers"),
176                    Some(mut buffer) => {
177                        // Wait until audio is streaming before we try to process
178                        if !audio_ready_clone.load(std::sync::atomic::Ordering::Acquire)
179                            || saving.load(std::sync::atomic::Ordering::Acquire)
180                        {
181                            return;
182                        }
183
184                        let datas = buffer.datas_mut();
185                        if datas.is_empty() {
186                            return;
187                        }
188
189                        let time_us = start_time.elapsed().as_micros() as i64;
190
191                        // send frame data to encoder
192                        let data = &mut datas[0];
193
194                        let fd = Self::get_dmabuf_fd(data);
195
196                        if fd.is_some()
197                            && ringbuf_producer
198                                .try_push(RawVideoFrame {
199                                    data: Vec::new(),
200                                    timestamp: time_us,
201                                    dmabuf_fd: fd,
202                                    stride: data.chunk().stride(),
203                                    offset: data.chunk().offset(),
204                                    size: data.chunk().size(),
205                                    modifier: udata.video_format.modifier(),
206                                })
207                                .is_err()
208                        {
209                            log::error!(
210                                "Error sending video frame at: {:?}. Ring buf full?",
211                                time_us
212                            );
213                        }
214                    }
215                }
216            })
217            .register()?;
218
219        // TODO: Use features? Probably should not have runtime conditionals like this
220        let pw_obj = if self.use_nvidia_modifiers {
221            let nvidia_mod_property = Property {
222                key: pw::spa::param::format::FormatProperties::VideoModifier.as_raw(),
223                flags: PropertyFlags::empty(),
224                value: spa::pod::Value::Choice(spa::pod::ChoiceValue::Long(Choice::<i64>(
225                    ChoiceFlags::empty(),
226                    ChoiceEnum::<i64>::Enum {
227                        default: NVIDIA_MODIFIERS[0],
228                        alternatives: NVIDIA_MODIFIERS.to_vec(),
229                    },
230                ))),
231            };
232
233            Some(pw::spa::pod::object!(
234                pw::spa::utils::SpaTypes::ObjectParamFormat,
235                pw::spa::param::ParamType::EnumFormat,
236                pw::spa::pod::property!(
237                    pw::spa::param::format::FormatProperties::MediaType,
238                    Id,
239                    pw::spa::param::format::MediaType::Video
240                ),
241                pw::spa::pod::property!(
242                    pw::spa::param::format::FormatProperties::MediaSubtype,
243                    Id,
244                    pw::spa::param::format::MediaSubtype::Raw
245                ),
246                nvidia_mod_property,
247                pw::spa::pod::property!(
248                    pw::spa::param::format::FormatProperties::VideoFormat,
249                    Choice,
250                    Enum,
251                    Id,
252                    pw::spa::param::video::VideoFormat::NV12,
253                    pw::spa::param::video::VideoFormat::I420,
254                    pw::spa::param::video::VideoFormat::BGRA,
255                ),
256                pw::spa::pod::property!(
257                    pw::spa::param::format::FormatProperties::VideoSize,
258                    Choice,
259                    Range,
260                    Rectangle,
261                    pw::spa::utils::Rectangle {
262                        width: 2560,
263                        height: 1440
264                    }, // Default
265                    pw::spa::utils::Rectangle {
266                        width: 1,
267                        height: 1
268                    }, // Min
269                    pw::spa::utils::Rectangle {
270                        width: 4096,
271                        height: 4096
272                    } // Max
273                ),
274                pw::spa::pod::property!(
275                    pw::spa::param::format::FormatProperties::VideoFramerate,
276                    Choice,
277                    Range,
278                    Fraction,
279                    pw::spa::utils::Fraction { num: 240, denom: 1 }, // Default
280                    pw::spa::utils::Fraction { num: 0, denom: 1 },   // Min
281                    pw::spa::utils::Fraction { num: 244, denom: 1 }  // Max
282                ),
283            ))
284        } else {
285            Some(pw::spa::pod::object!(
286                pw::spa::utils::SpaTypes::ObjectParamFormat,
287                pw::spa::param::ParamType::EnumFormat,
288                pw::spa::pod::property!(
289                    pw::spa::param::format::FormatProperties::MediaType,
290                    Id,
291                    pw::spa::param::format::MediaType::Video
292                ),
293                pw::spa::pod::property!(
294                    pw::spa::param::format::FormatProperties::MediaSubtype,
295                    Id,
296                    pw::spa::param::format::MediaSubtype::Raw
297                ),
298                pw::spa::pod::property!(
299                    pw::spa::param::format::FormatProperties::VideoModifier,
300                    Long,
301                    0
302                ),
303                pw::spa::pod::property!(
304                    pw::spa::param::format::FormatProperties::VideoFormat,
305                    Choice,
306                    Enum,
307                    Id,
308                    pw::spa::param::video::VideoFormat::NV12,
309                    pw::spa::param::video::VideoFormat::I420,
310                    pw::spa::param::video::VideoFormat::BGRA,
311                ),
312                pw::spa::pod::property!(
313                    pw::spa::param::format::FormatProperties::VideoSize,
314                    Choice,
315                    Range,
316                    Rectangle,
317                    pw::spa::utils::Rectangle {
318                        width: 2560,
319                        height: 1440
320                    }, // Default
321                    pw::spa::utils::Rectangle {
322                        width: 1,
323                        height: 1
324                    }, // Min
325                    pw::spa::utils::Rectangle {
326                        width: 4096,
327                        height: 4096
328                    } // Max
329                ),
330                pw::spa::pod::property!(
331                    pw::spa::param::format::FormatProperties::VideoFramerate,
332                    Choice,
333                    Range,
334                    Fraction,
335                    pw::spa::utils::Fraction { num: 240, denom: 1 }, // Default
336                    pw::spa::utils::Fraction { num: 0, denom: 1 },   // Min
337                    pw::spa::utils::Fraction { num: 244, denom: 1 }  // Max
338                ),
339            ))
340        };
341
342        let video_spa_values: Vec<u8> = pw::spa::pod::serialize::PodSerializer::serialize(
343            std::io::Cursor::new(Vec::new()),
344            &pw::spa::pod::Value::Object(pw_obj.unwrap()),
345        )
346        .unwrap()
347        .0
348        .into_inner();
349        let mut video_params = [Pod::from_bytes(&video_spa_values).unwrap()];
350        video_stream.connect(
351            Direction::Input,
352            Some(stream_node),
353            StreamFlags::AUTOCONNECT | StreamFlags::MAP_BUFFERS,
354            &mut video_params,
355        )?;
356
357        log::debug!("Video Stream: {0:?}", video_stream);
358
359        pw_loop.run();
360        Ok(())
361    }
362
363    fn get_dmabuf_fd(data: &Data) -> Option<RawFd> {
364        let raw_data = data.as_raw();
365
366        if data.type_() == DataType::DmaBuf {
367            let fd = raw_data.fd;
368
369            if fd > 0 {
370                return Some(fd as i32);
371            }
372        }
373
374        None
375    }
376}
waycap_rs/capture/video.rs

waycap_rs/capture/
video.rs
