dmxcapture/dmx.py

from PyQt5.QtCore import QObject, QThread, pyqtSignal
import socket
import struct
import numpy as np
import subprocess
import time
import os

class DMXRecorder(QObject):
    frame_ready = pyqtSignal(np.ndarray)
    
    def __init__(self):
        super().__init__()
        self.thread = None
        self.recording_thread = None

    def start_recording(self):
        self.thread = QThread()
        self.recording_thread = RecordingThread()
        self.recording_thread.moveToThread(self.thread)

        self.thread.started.connect(self.recording_thread.run)
        self.recording_thread.finished.connect(self.thread.quit)
        self.recording_thread.finished.connect(self.recording_thread.deleteLater)
        self.thread.finished.connect(self.thread.deleteLater)

        # Connect frame_ready signal
        self.recording_thread.frame_ready.connect(self.frame_ready)

        self.thread.start()

    def stop_recording(self):
        if self.recording_thread:
            self.recording_thread.stop()
        if self.thread:
            self.thread.quit()
            self.thread.wait()

class RecordingThread(QObject):
    finished = pyqtSignal()
    frame_ready = pyqtSignal(np.ndarray)

    def __init__(self):
        super().__init__()
        self.running = False
        self.socket = None
        self.ffmpeg_process = None

    def run(self):
        self.running = True

        # Set up UDP socket to listen for Art-Net DMX data
        UDP_IP = "0.0.0.0"
        UDP_PORT = 6454

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind((UDP_IP, UDP_PORT))

        # Set socket to non-blocking
        self.socket.setblocking(False)

        # Prepare for recording
        expected_universes = set(range(4))  # Universes 0 to 3 (adjust as needed)
        num_universes = len(expected_universes)
        universe_size = 512

        # Define frame dimensions
        frame_width = 64  # Adjust as needed (must divide evenly into total data size)
        total_pixels = num_universes * universe_size
        frame_height = total_pixels // frame_width

        # Frame rate and timing
        frame_rate = 30  # frames per second
        frame_interval = 1.0 / frame_rate
        last_frame_time = time.time()

        # Adjust the path to the ffmpeg executable
        ffmpeg_executable = os.path.join(os.path.dirname(__file__), 'bin', 'ffmpeg')

        # For Windows, add .exe extension
        if os.name == 'nt':
            ffmpeg_executable += '.exe'

        # Ensure the ffmpeg executable is executable
        # For Unix-like systems, set the execute permission if needed
        if not os.access(ffmpeg_executable, os.X_OK):
            if os.name != 'nt':
                os.chmod(ffmpeg_executable, 0o755)
            else:
                # On Windows, os.access may not check execute permissions accurately
                if not os.path.isfile(ffmpeg_executable):
                    print(f"ffmpeg executable not found at: {ffmpeg_executable}")
                    self.finished.emit()
                    return

        # Set up ffmpeg process using the local executable path
        ffmpeg_cmd = [
            ffmpeg_executable,
            '-y',
            '-f', 'rawvideo',
            '-pix_fmt', 'gray',
            '-s', f'{frame_width}x{frame_height}',
            '-r', str(frame_rate),
            '-i', '-',  # Input from stdin
            '-c:v', 'ffv1',
            'dmx_video.mkv'
        ]

        try:
            self.ffmpeg_process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE)
        except Exception as e:
            print(f"Error starting ffmpeg: {e}")
            self.finished.emit()
            return

        universes_data = {}
        last_received = {}

        try:
            while self.running:
                current_time = time.time()

                # Receive data
                try:
                    data, addr = self.socket.recvfrom(2048)
                    if data.startswith(b'Art-Net'):
                        op_code = struct.unpack('<H', data[8:10])[0]
                        # OpCode for ArtDMX is 0x5000
                        if op_code == 0x5000:
                            # Parse universe
                            universe = struct.unpack('<H', data[14:16])[0]
                            # Get length
                            length = struct.unpack('>H', data[16:18])[0]
                            dmx_data = data[18:18+length]
                            # Store the data for this universe
                            universes_data[universe] = dmx_data
                            last_received[universe] = current_time
                except BlockingIOError:
                    # No data received
                    pass
                except Exception as e:
                    print(f"Error receiving data: {e}")

                # Check if it's time to write a frame
                if current_time - last_frame_time >= frame_interval:
                    # Assemble frame data
                    frame_data = b''
                    for universe in sorted(expected_universes):
                        dmx_data = universes_data.get(universe, b'\x00' * universe_size)
                        if len(dmx_data) < universe_size:
                            dmx_data += b'\x00' * (universe_size - len(dmx_data))
                        frame_data += dmx_data
                    # Convert frame_data to numpy array
                    frame_array = np.frombuffer(frame_data, dtype=np.uint8)
                    # Reshape to image dimensions
                    frame_array = frame_array.reshape((frame_height, frame_width))
                    # Write frame to ffmpeg stdin
                    if self.ffmpeg_process.stdin:
                        try:
                            self.ffmpeg_process.stdin.write(frame_array.tobytes())
                        except BrokenPipeError:
                            print("FFmpeg process has terminated unexpectedly.")
                            self.running = False
                    # Emit frame for visualization
                    self.frame_ready.emit(frame_array)
                    last_frame_time = current_time

                time.sleep(0.001)  # Sleep briefly to avoid maxing out CPU
        finally:
            # Clean up
            if self.ffmpeg_process and self.ffmpeg_process.stdin:
                try:
                    self.ffmpeg_process.stdin.close()
                    self.ffmpeg_process.wait()
                except Exception as e:
                    print(f"Error closing ffmpeg process: {e}")
            self.socket.close()
            self.finished.emit()

    def stop(self):
        self.running = False