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, listen_address='0.0.0.0', port=6454):
        super().__init__()
        self.thread = None
        self.recording_thread = None
        self.listen_address = listen_address
        self.port = port

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

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

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

        self.thread.finished.connect(self.thread.deleteLater)
        self.recording_thread.finished.connect(self.recording_thread.deleteLater)
        self.thread.start()

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

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

    def __init__(self, listen_address='0.0.0.0', port=6454):
        super().__init__()
        self.running = False
        self.socket = None
        self.ffmpeg_process = None
        self.listen_address = listen_address
        self.port = port

    def run(self):
        self.running = True

        # Set up UDP socket to listen for Art-Net DMX data
        UDP_IP = self.listen_address
        UDP_PORT = self.port

        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
        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 = {}

        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
                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

class DMXPlayer(QObject):
    frame_ready = pyqtSignal(np.ndarray)
    playback_finished = pyqtSignal()

    def __init__(self, num_universes, transmit_address='255.255.255.255', port=6454):
        super().__init__()
        self.thread = None
        self.playback_thread = None
        self.num_universes = num_universes
        self.transmit_address = transmit_address
        self.port = port

    def start_playback(self):
        self.thread = QThread()
        self.playback_thread = PlaybackThread(self.num_universes, transmit_address=self.transmit_address, port=self.port)
        self.playback_thread.moveToThread(self.thread)

        self.thread.started.connect(self.playback_thread.run)
        self.playback_thread.finished.connect(self.thread.quit)

        # Connect signals
        self.playback_thread.frame_ready.connect(self.frame_ready)
        self.playback_thread.finished.connect(self.on_playback_thread_finished)

        self.thread.finished.connect(self.on_thread_finished)

        self.thread.start()

    def on_playback_thread_finished(self):
        self.playback_finished.emit()
        # Wait for the thread to exit before cleaning up
        self.thread.quit()
        self.thread.wait()
        self.playback_thread = None

    def on_thread_finished(self):
        self.thread = None

    def stop_playback(self):
        if self.playback_thread:
            self.playback_thread.stop()
            self.thread.quit()
            self.thread.wait()
            self.playback_thread = None
            self.thread = None

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

    def __init__(self, num_universes, transmit_address='255.255.255.255', port=6454):
        super().__init__()
        self.num_universes = num_universes
        self.running = False
        self.ffmpeg_process = None
        self.socket = None
        self.transmit_address = transmit_address
        self.port = port

    def run(self):
        self.running = True

        # Set up UDP socket to send Art-Net DMX data
        UDP_IP = self.transmit_address
        UDP_PORT = self.port

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # Enable broadcasting mode if necessary
        if UDP_IP == '255.255.255.255':
            self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)

        # 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
        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

        # Define frame dimensions (must match recording)
        frame_width = 64  # Must match recording settings
        universe_size = 512  # DMX universe size
        num_universes_in_file = 4  # Universes recorded (0 to 3)
        total_pixels_in_file = num_universes_in_file * universe_size
        frame_height_in_file = total_pixels_in_file // frame_width
        frame_rate = 30  # frames per second
        frame_interval = 1.0 / frame_rate

        # Compute playback dimensions based on number of universes to playback
        total_pixels_playback = self.num_universes * universe_size
        frame_height_playback = total_pixels_playback // frame_width

        # FFmpeg command to read frames
        ffmpeg_cmd = [
            ffmpeg_executable,
            '-i', 'dmx_video.mkv',  # Input file
            '-f', 'rawvideo',
            '-pix_fmt', 'gray',
            '-s', f'{frame_width}x{frame_height_in_file}',
            '-r', str(frame_rate),
            '-an',
            '-sn',
            '-'
        ]

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

        # Read frames from ffmpeg stdout
        frame_size = frame_width * frame_height_in_file
        sequence = 0

        try:
            while self.running:
                # Read frame data
                frame_data = self.ffmpeg_process.stdout.read(frame_size)
                if not frame_data or len(frame_data) < frame_size:
                    # End of file
                    break

                # Extract data for the specified number of universes
                total_pixels_playback_bytes = self.num_universes * universe_size
                frame_data_playback = frame_data[:total_pixels_playback_bytes]

                # Convert frame_data to numpy array
                frame_array = np.frombuffer(frame_data_playback, dtype=np.uint8)
                # Reshape to image dimensions
                frame_array = frame_array.reshape((frame_height_playback, frame_width))

                # Send Art-Net data
                # Split frame_data into universes
                expected_universes = sorted(range(self.num_universes))
                for idx, universe in enumerate(expected_universes):
                    start_idx = idx * universe_size
                    end_idx = start_idx + universe_size
                    universe_data = frame_data_playback[start_idx:end_idx]
                    # Construct ArtDMX packet
                    packet = self.build_artdmx_packet(sequence, universe, universe_data)
                    # Send packet over UDP
                    self.socket.sendto(packet, (UDP_IP, UDP_PORT))

                sequence = (sequence + 1) % 256
                # Emit frame for visualization
                self.frame_ready.emit(frame_array)

                # Sleep to synchronize to frame rate
                time.sleep(frame_interval)
        finally:
            if self.ffmpeg_process:
                self.ffmpeg_process.stdout.close()
                self.ffmpeg_process.wait()
            self.socket.close()
            self.finished.emit()

    def stop(self):
        self.running = False

    def build_artdmx_packet(self, sequence, universe, data):
        packet = bytearray()
        packet.extend(b'Art-Net\x00')  # ID
        packet.extend(struct.pack('<H', 0x5000))  # OpCode ArtDMX
        packet.extend(struct.pack('>H', 14))  # Protocol Version (High byte first)
        packet.extend(struct.pack('B', sequence))  # Sequence
        packet.extend(struct.pack('B', 0x00))  # Physical
        packet.extend(struct.pack('<H', universe))  # Universe (Low byte first)
        length = len(data)
        packet.extend(struct.pack('>H', length))  # Length (High byte first)
        packet.extend(data)
        return packet