flandre/src/nodes/ImageFFMPEG.py

import logging
import subprocess
import time

import numpy as np
import zmq

from BusClient import BusClient
from config import VIDEO_HEIGHT, VIDEO_WIDTH
from nodes.Node import Node
from utils.Msg import BMMsg, SetWindowVisibleMsg, Msg

logger = logging.getLogger(__name__)


class ImageFFMPEG(Node):
    topics = [BMMsg]

    def __init__(self, level=logging.INFO):
        super().__init__(level=level)
        self.buffer = np.zeros((VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=np.uint8) + 128
        self.buffer = self.buffer.tobytes()

    def loop(self):
        self.c = BusClient(BMMsg, pub=False, conflare=True, poller=True)
        p = subprocess.Popen(['ffmpeg',
                              '-f', 'rawvideo',
                              '-pixel_format', 'rgb24',
                              '-video_size', '1080x1920',
                              '-framerate', '24',
                              '-hwaccel', 'nvdec',
                              '-i', '-',
                              '-vcodec', 'h264_nvenc',
                              '-preset', 'fast',
                              '-gpu', '1',

                              # '-profile:v', 'high',
                              '-zerolatency', '1',
                              # '-tune', 'ull',
                              # '-level', '42',

                              '-pix_fmt', 'yuv420p',
                              # '-vcodec', 'libx264',
                              # '-b:v', '40M',
                              '-f', 'flv',
                              'rtmp://q1hyb.as/live/bscan'
                              # '-f', 'mpegts',
                              # 'srt://localhost:10080?streamid=#!::r=live/livestream,m=publish'
                              ],
                             stdin=subprocess.PIPE,
                             )
        lasttime = time.time()
        while True:
            # socks = dict(self.c.poller.poll(1 / 30))
            events = self.c.poller.poll(1000 / 24)
            if events:
                msg: BMMsg = Msg.decode_msg(events[0][0].recv())
                b = np.frombuffer(msg.data, dtype=np.uint8)
                b = np.reshape(b, (VIDEO_HEIGHT, VIDEO_WIDTH, 4))[:, :, :3]
                self.buffer = b.tobytes()
            p.stdin.write(self.buffer)
            # time.sleep(1 / 60)
            currenttime = time.time()
            logger.debug(f'{currenttime - lasttime}')
            lasttime = currenttime


if __name__ == '__main__':
    ImageFFMPEG()()