From b38df7850224325082226caa47f658cb7b8f6474 Mon Sep 17 00:00:00 2001 From: Richard Abrich Date: Mon, 30 Dec 2024 13:58:16 -0500 Subject: [PATCH] add capture._linux --- openadapt/capture/__init__.py | 8 ++- openadapt/capture/_linux.py | 119 ++++++++++++++++++++++++++++++++++ openadapt/capture/_macos.py | 5 +- openadapt/capture/_windows.py | 6 +- 4 files changed, 127 insertions(+), 11 deletions(-) create mode 100644 openadapt/capture/_linux.py diff --git a/openadapt/capture/__init__.py b/openadapt/capture/__init__.py index 91ae7b98c..dd572551a 100644 --- a/openadapt/capture/__init__.py +++ b/openadapt/capture/__init__.py @@ -9,13 +9,15 @@ if sys.platform == "darwin": from . import _macos as impl - device = impl.Capture() elif sys.platform == "win32": from . import _windows as impl - device = impl.Capture() +elif sys.platform == "linux": + from . import _linux as impl else: print(f"WARNING: openadapt.capture is not yet supported on {sys.platform=}") - device = None + impl = None + +device = impl.Capture() if impl else None def get_capture() -> impl.Capture: diff --git a/openadapt/capture/_linux.py b/openadapt/capture/_linux.py new file mode 100644 index 000000000..8da5f4a77 --- /dev/null +++ b/openadapt/capture/_linux.py @@ -0,0 +1,119 @@ +"""Allows for capturing the screen and audio on Linux. + +usage: see bottom of file +""" + +import os +import subprocess +from datetime import datetime +from sys import platform +import wave +import pyaudio + +from openadapt.config import CAPTURE_DIR_PATH + + +class Capture: + """Capture the screen, audio, and camera on Linux.""" + + def __init__(self) -> None: + """Initialize the capture object.""" + assert platform == "linux", platform + + self.is_recording = False + self.audio_out = None + self.video_out = None + self.audio_stream = None + self.audio_frames = [] + + # Initialize PyAudio + self.audio = pyaudio.PyAudio() + + def start(self, audio: bool = True, camera: bool = False) -> None: + """Start capturing the screen, audio, and camera. + + Args: + audio (bool, optional): Whether to capture audio (default: True). + camera (bool, optional): Whether to capture the camera (default: False). + """ + if self.is_recording: + raise RuntimeError("Recording is already in progress") + + self.is_recording = True + capture_dir = CAPTURE_DIR_PATH + if not os.path.exists(capture_dir): + os.mkdir(capture_dir) + + # Start video capture using ffmpeg + video_filename = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + ".mp4" + self.video_out = os.path.join(capture_dir, video_filename) + self._start_video_capture() + + # Start audio capture + if audio: + audio_filename = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + ".wav" + self.audio_out = os.path.join(capture_dir, audio_filename) + self._start_audio_capture() + + def _start_video_capture(self) -> None: + """Start capturing the screen using ffmpeg.""" + cmd = [ + "ffmpeg", + "-f", "x11grab", # Capture X11 display + "-video_size", "1920x1080", # Set resolution + "-framerate", "30", # Set frame rate + "-i", ":0.0", # Capture from display 0 + "-c:v", "libx264", # Video codec + "-preset", "ultrafast", # Speed/quality tradeoff + "-y", self.video_out # Output file + ] + self.video_proc = subprocess.Popen(cmd) + + def _start_audio_capture(self) -> None: + """Start capturing audio using PyAudio.""" + self.audio_stream = self.audio.open( + format=pyaudio.paInt16, + channels=2, + rate=44100, + input=True, + frames_per_buffer=1024, + stream_callback=self._audio_callback + ) + self.audio_frames = [] + self.audio_stream.start_stream() + + def _audio_callback(self, in_data: bytes, frame_count: int, time_info: dict, status: int) -> tuple: + """Callback function to process audio data.""" + self.audio_frames.append(in_data) + return (None, pyaudio.paContinue) + + def stop(self) -> None: + """Stop capturing the screen, audio, and camera.""" + if self.is_recording: + # Stop the video capture + self.video_proc.terminate() + + # Stop audio capture + if self.audio_stream: + self.audio_stream.stop_stream() + self.audio_stream.close() + self.audio.terminate() + self.save_audio() + + self.is_recording = False + + def save_audio(self) -> None: + """Save the captured audio to a WAV file.""" + if self.audio_out: + with wave.open(self.audio_out, "wb") as wf: + wf.setnchannels(2) + wf.setsampwidth(self.audio.get_sample_size(pyaudio.paInt16)) + wf.setframerate(44100) + wf.writeframes(b"".join(self.audio_frames)) + + +if __name__ == "__main__": + capture = Capture() + capture.start(audio=True, camera=False) + input("Press enter to stop") + capture.stop() diff --git a/openadapt/capture/_macos.py b/openadapt/capture/_macos.py index 65529910c..d2c00e7d4 100644 --- a/openadapt/capture/_macos.py +++ b/openadapt/capture/_macos.py @@ -23,10 +23,7 @@ class Capture: def __init__(self) -> None: """Initialize the capture object.""" - if platform != "darwin": - raise NotImplementedError( - "This is the macOS implementation, please use the Windows version" - ) + assert platform == "darwin", platform objc.options.structs_indexable = True diff --git a/openadapt/capture/_windows.py b/openadapt/capture/_windows.py index ab400c950..ad09e48b1 100644 --- a/openadapt/capture/_windows.py +++ b/openadapt/capture/_windows.py @@ -21,10 +21,8 @@ def __init__(self, pid: int = 0) -> None: pid (int, optional): The process ID of the window to capture. Defaults to 0 (the entire screen) """ - if platform != "win32": - raise NotImplementedError( - "This is the Windows implementation, please use the macOS version" - ) + assert platform == "win32", platform + self.is_recording = False self.video_out = None self.audio_out = None