Skip to main content

g2g_python/
source.rs

1//! `PySource`: a Python frame source hosted as a g2g [`SourceLoop`].
2//!
3//! Each tick hands the Python element a blank, writable buffer-protocol frame
4//! via `g2g_produce(buf, width, height, fmt, meta) -> bool`; the element fills
5//! it in place and returns `True`, or returns `False` to end the stream. Output
6//! caps are fixed by the source's properties, so negotiation is synchronous
7//! (`intercept_caps` returns them with no I/O). The same GIL-owning worker
8//! thread as the other hosts runs the calls.
9//!
10//! Under `cuda-frames` the source produces GPU-resident frames instead (M986):
11//! `g2g_produce_cuda(width, height, meta)` returns the two semi-planar planes as
12//! `__cuda_array_interface__` objects (a cupy / torch allocation), since this
13//! crate links no CUDA and cannot allocate device memory itself. The frame carries
14//! those objects as its keep-alive, so the memory outlives it, and this source
15//! stamps the timing on the way out either way.
16
17use core::future::{Future, Ready};
18use core::pin::Pin;
19
20use g2g_core::memory::SystemSlice;
21use g2g_core::runtime::SourceLoop;
22use g2g_core::{
23    Caps, ConfigureOutcome, Dim, Frame, FrameTiming, G2gError, MemoryDomain, OutputSink,
24    PipelinePacket, PropError, PropKind, PropValue, PropertySpec, Rate, RawVideoFormat,
25};
26
27use crate::format::{format_from_py, format_to_py, frame_bytes};
28
29/// A Python frame source hosted as a first-class g2g source.
30#[derive(Debug)]
31pub struct PySource {
32    module: String,
33    class: String,
34    /// Fixed output caps (concrete format / dims / rate).
35    caps: Caps,
36    /// Optional cap on frames produced; `None` runs until Python signals EOS.
37    num_buffers: Option<u64>,
38    /// Whether the hosted source produces GPU-resident CUDA surfaces
39    /// (`g2g_produce_cuda`) rather than filling blank System frames.
40    cuda_frames: bool,
41    configured: bool,
42    emitted: u64,
43    #[cfg(feature = "python")]
44    worker: Option<crate::host::PyWorker>,
45}
46
47impl PySource {
48    /// Host `class` from Python `module` as a frame source. Default caps are
49    /// RGBA 320x240 @ 30; override with [`with_caps`](Self::with_caps) or the
50    /// `format` / `width` / `height` / `framerate` properties.
51    pub fn new(module: impl Into<String>, class: impl Into<String>) -> Self {
52        Self {
53            module: module.into(),
54            class: class.into(),
55            caps: Caps::RawVideo {
56                format: RawVideoFormat::Rgba8,
57                width: Dim::Fixed(320),
58                height: Dim::Fixed(240),
59                framerate: Rate::Fixed(30),
60                interlace: g2g_core::Interlace::Any,
61            },
62            num_buffers: None,
63            cuda_frames: false,
64            configured: false,
65            emitted: 0,
66            #[cfg(feature = "python")]
67            worker: None,
68        }
69    }
70
71    /// Set the fixed output caps (must be fully fixed: concrete format/dims/rate).
72    pub fn with_caps(mut self, caps: Caps) -> Self {
73        self.caps = caps;
74        self
75    }
76
77    /// Stop after `n` frames (otherwise run until the Python source returns EOS).
78    pub fn with_num_buffers(mut self, n: u64) -> Self {
79        self.num_buffers = Some(n);
80        self
81    }
82
83    /// Host a source that produces GPU-resident CUDA surfaces through
84    /// `g2g_produce_cuda`, so the frames leave in the Cuda memory domain and the
85    /// caps must be semi-planar (NV12 / P010).
86    pub fn with_cuda_frames(mut self, on: bool) -> Self {
87        self.cuda_frames = on;
88        self
89    }
90
91    /// Count of frames pushed downstream. Useful in tests.
92    pub fn emitted_count(&self) -> u64 {
93        self.emitted
94    }
95
96    /// Mutate the RawVideo caps fields in place (no-op on non-RawVideo caps).
97    fn edit_caps(&mut self, f: impl FnOnce(&mut RawVideoFormat, &mut Dim, &mut Dim, &mut Rate)) {
98        if let Caps::RawVideo {
99            format,
100            width,
101            height,
102            framerate,
103            interlace: _,
104        } = &mut self.caps
105        {
106            f(format, width, height, framerate);
107        }
108    }
109
110    /// Ask the hosted source for its next frame: a blank System frame to fill, or
111    /// (under `cuda-frames`) a surface the Python side allocates itself.
112    #[cfg(feature = "python")]
113    async fn produce_one(&self, seq: u64, step: u64) -> Result<Option<Frame>, G2gError> {
114        let worker = self.worker.as_ref().ok_or(G2gError::NotConfigured)?;
115        let mut produced = if self.cuda_frames {
116            worker.run_produce_cuda(&self.caps).await?
117        } else {
118            worker
119                .run_produce(self.blank_frame(seq, step)?, &self.caps)
120                .await?
121        };
122        if let Some(frame) = &mut produced {
123            frame.timing = timing(seq, step);
124            frame.sequence = seq;
125        }
126        Ok(produced)
127    }
128
129    #[cfg(not(feature = "python"))]
130    async fn produce_one(&self, _seq: u64, _step: u64) -> Result<Option<Frame>, G2gError> {
131        Err(G2gError::UnsupportedDomain)
132    }
133
134    /// Allocate a zeroed frame of the output geometry, stamped at `seq` /
135    /// `pts_step_ns`, for the Python source to fill. Only the interpreter build
136    /// has anything to hand it to.
137    #[cfg_attr(not(feature = "python"), allow(dead_code))]
138    fn blank_frame(&self, seq: u64, pts_step_ns: u64) -> Result<Frame, G2gError> {
139        let Caps::RawVideo {
140            format,
141            width: Dim::Fixed(w),
142            height: Dim::Fixed(h),
143            ..
144        } = &self.caps
145        else {
146            return Err(G2gError::FixationFailed);
147        };
148        let bytes = vec![0u8; frame_bytes(*format, *w, *h)].into_boxed_slice();
149        Ok(Frame {
150            domain: MemoryDomain::System(SystemSlice::from_boxed(bytes)),
151            timing: timing(seq, pts_step_ns),
152            sequence: seq,
153            meta: Default::default(),
154        })
155    }
156}
157
158/// Timing for frame `seq` at `pts_step_ns` per frame.
159#[cfg_attr(not(feature = "python"), allow(dead_code))]
160fn timing(seq: u64, pts_step_ns: u64) -> FrameTiming {
161    let pts = seq.saturating_mul(pts_step_ns);
162    FrameTiming {
163        pts_ns: pts,
164        dts_ns: pts,
165        duration_ns: pts_step_ns,
166        capture_ns: pts,
167        arrival_ns: 0,
168        keyframe: false,
169    }
170}
171
172/// Per-frame PTS step from a fixed framerate, or 0 when the rate is not fixed.
173fn pts_step_ns(caps: &Caps) -> u64 {
174    match caps {
175        Caps::RawVideo {
176            framerate: Rate::Fixed(fps),
177            ..
178        } if *fps > 0 => 1_000_000_000u64 / u64::from(*fps),
179        _ => 0,
180    }
181}
182
183impl SourceLoop for PySource {
184    type RunFuture<'a>
185        = Pin<Box<dyn Future<Output = Result<u64, G2gError>> + 'a>>
186    where
187        Self: 'a;
188    type CapsFuture<'a>
189        = Ready<Result<Caps, G2gError>>
190    where
191        Self: 'a;
192
193    fn intercept_caps<'a>(&'a mut self) -> Self::CapsFuture<'a> {
194        core::future::ready(Ok(self.caps.clone()))
195    }
196
197    fn configure_pipeline(&mut self, absolute_caps: &Caps) -> Result<ConfigureOutcome, G2gError> {
198        absolute_caps.intersect(&self.caps)?;
199        self.caps = absolute_caps.clone();
200        #[cfg(feature = "python")]
201        {
202            if self.module.is_empty() || self.class.is_empty() {
203                return Err(G2gError::NotConfigured);
204            }
205            if self.worker.is_none() {
206                // Property forwarding to the hosted source instance is a
207                // follow-up; transforms (PyTransform) forward theirs today.
208                self.worker = Some(crate::host::PyWorker::spawn(
209                    &self.module,
210                    &self.class,
211                    false,
212                    &[],
213                )?);
214            }
215        }
216        self.configured = true;
217        Ok(ConfigureOutcome::Accepted)
218    }
219
220    fn run<'a>(&'a mut self, out: &'a mut dyn OutputSink) -> Self::RunFuture<'a> {
221        Box::pin(async move {
222            if !self.configured {
223                return Err(G2gError::NotConfigured);
224            }
225            let step = pts_step_ns(&self.caps);
226            let mut produced = 0u64;
227            loop {
228                if let Some(limit) = self.num_buffers {
229                    if produced >= limit {
230                        break;
231                    }
232                }
233                match self.produce_one(produced, step).await? {
234                    Some(frame) => {
235                        out.push(PipelinePacket::DataFrame(frame)).await?;
236                        produced += 1;
237                    }
238                    // Python source signalled end of stream.
239                    None => break,
240                }
241            }
242            self.emitted = produced;
243            out.push(PipelinePacket::Eos).await?;
244            Ok(produced)
245        })
246    }
247
248    /// The output caps are property-fixed, so the auto-plug parser can read them
249    /// without negotiation.
250    fn configured_output_caps(&self) -> Option<Caps> {
251        Some(self.caps.clone())
252    }
253
254    /// A GPU source emits into the Cuda domain, so a downstream GPU consumer takes
255    /// the frames with no upload and a CPU consumer gets a download spliced (M985).
256    fn output_memory(&self) -> g2g_core::memory::MemoryDomainKind {
257        if self.cuda_frames {
258            g2g_core::memory::MemoryDomainKind::Cuda
259        } else {
260            g2g_core::memory::MemoryDomainKind::System
261        }
262    }
263
264    fn properties(&self) -> &'static [PropertySpec] {
265        PYSOURCE_PROPS
266    }
267
268    fn set_property(&mut self, name: &str, value: PropValue) -> Result<(), PropError> {
269        match name {
270            "module" => {
271                self.module = value.as_str().ok_or(PropError::Type)?.to_string();
272                Ok(())
273            }
274            "class" => {
275                self.class = value.as_str().ok_or(PropError::Type)?.to_string();
276                Ok(())
277            }
278            "num-buffers" => {
279                let n = value.as_int().ok_or(PropError::Type)?;
280                self.num_buffers = if n < 0 { None } else { Some(n as u64) };
281                Ok(())
282            }
283            "cuda-frames" => {
284                self.cuda_frames = value.as_bool().ok_or(PropError::Type)?;
285                Ok(())
286            }
287            "format" => {
288                let f = format_from_py(value.as_str().ok_or(PropError::Type)?)
289                    .ok_or(PropError::Value)?;
290                self.edit_caps(|format, _, _, _| *format = f);
291                Ok(())
292            }
293            "width" => {
294                let w = value.as_uint().ok_or(PropError::Type)? as u32;
295                self.edit_caps(|_, width, _, _| *width = Dim::Fixed(w));
296                Ok(())
297            }
298            "height" => {
299                let h = value.as_uint().ok_or(PropError::Type)? as u32;
300                self.edit_caps(|_, _, height, _| *height = Dim::Fixed(h));
301                Ok(())
302            }
303            "framerate" => {
304                let (n, d) = value.as_fraction().ok_or(PropError::Type)?;
305                if d <= 0 || n <= 0 {
306                    return Err(PropError::Value);
307                }
308                let fps = (n / d) as u32;
309                self.edit_caps(|_, _, _, rate| *rate = Rate::Fixed(fps));
310                Ok(())
311            }
312            _ => Err(PropError::Unknown),
313        }
314    }
315
316    fn get_property(&self, name: &str) -> Option<PropValue> {
317        let raw = match &self.caps {
318            Caps::RawVideo {
319                format,
320                width,
321                height,
322                framerate,
323                interlace: _,
324            } => Some((format, width, height, framerate)),
325            _ => None,
326        };
327        match name {
328            "module" => Some(PropValue::Str(self.module.clone())),
329            "class" => Some(PropValue::Str(self.class.clone())),
330            "num-buffers" => Some(PropValue::Int(self.num_buffers.map_or(-1, |n| n as i64))),
331            "cuda-frames" => Some(PropValue::Bool(self.cuda_frames)),
332            "format" => raw.map(|(f, _, _, _)| PropValue::Str(format_to_py(*f).to_string())),
333            "width" => raw.and_then(|(_, w, _, _)| match w {
334                Dim::Fixed(v) => Some(PropValue::Uint(u64::from(*v))),
335                _ => None,
336            }),
337            "height" => raw.and_then(|(_, _, h, _)| match h {
338                Dim::Fixed(v) => Some(PropValue::Uint(u64::from(*v))),
339                _ => None,
340            }),
341            "framerate" => raw.and_then(|(_, _, _, r)| match r {
342                Rate::Fixed(fps) => Some(PropValue::Fraction(*fps as i32, 1)),
343                _ => None,
344            }),
345            _ => None,
346        }
347    }
348}
349
350/// `PySource`'s settable properties (the runtime / `gst-launch` face).
351static PYSOURCE_PROPS: &[PropertySpec] = &[
352    PropertySpec::new(
353        "module",
354        PropKind::Str,
355        "Python module to import (the source element)",
356    ),
357    PropertySpec::new(
358        "class",
359        PropKind::Str,
360        "class within the module to instantiate",
361    ),
362    PropertySpec::new(
363        "format",
364        PropKind::Str,
365        "output pixel format (RGBA | BGRA | NV12 | I420 | YUY2)",
366    )
367    .with_default("RGBA"),
368    PropertySpec::new("width", PropKind::Uint, "output width in pixels").with_default("320"),
369    PropertySpec::new("height", PropKind::Uint, "output height in pixels").with_default("240"),
370    PropertySpec::new("framerate", PropKind::Fraction, "output framerate").with_default("30/1"),
371    PropertySpec::new(
372        "num-buffers",
373        PropKind::Int,
374        "frames to produce, or -1 until EOS",
375    )
376    .with_default("-1"),
377    PropertySpec::new(
378        "cuda-frames",
379        PropKind::Bool,
380        "host a source that produces GPU-resident CUDA surfaces (needs g2g_produce_cuda, NV12 / P010)",
381    )
382    .with_default("false"),
383];