platypush/platypush/plugins/camera/model/writer/ffmpeg.py

import logging
import subprocess
import threading
import time

from abc import ABC
from typing import Optional, List

from PIL.Image import Image

from platypush.plugins.camera.model.writer import VideoWriter, FileVideoWriter, StreamWriter


logger = logging.getLogger('ffmpeg-writer')


class FFmpegWriter(VideoWriter, ABC):
    """
    Generic FFmpeg encoder for camera frames.
    """

    def __init__(self, *args, input_file: str = '-', input_format: str = 'rawvideo', input_codec: Optional[str] = None,
                 output_file: str = '-', output_format: Optional[str] = None, output_codec: Optional[str] = None,
                 pix_fmt: Optional[str] = None, output_opts: Optional[List[str]] = None, **kwargs):
        super().__init__(*args, **kwargs)

        self.input_file = input_file
        self.input_format = input_format
        self.input_codec = input_codec
        self.output_file = output_file
        self.output_format = output_format
        self.output_codec = output_codec
        self.width, self.height = self.camera.effective_resolution()
        self.pix_fmt = pix_fmt
        self.output_opts = output_opts or []

        logger.info('Starting FFmpeg. Command: {}'.format(' '.join(self.ffmpeg_args)))
        self.ffmpeg = subprocess.Popen(self.ffmpeg_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

    @property
    def ffmpeg_args(self):
        return ['ffmpeg', '-y',
                '-f', self.input_format,
                *(('-vcodec', self.input_codec) if self.input_codec else ()),
                *(('-pix_fmt', self.pix_fmt) if self.pix_fmt else ()),
                '-s', '{}x{}'.format(self.width, self.height),
                '-r', str(self.camera.info.fps),
                '-i', self.input_file,
                *(('-f', self.output_format) if self.output_format else ()),
                *self.output_opts,
                *(('-vcodec', self.output_codec) if self.output_codec else ()),
                self.output_file]

    def is_closed(self):
        return self.closed or not self.ffmpeg or self.ffmpeg.poll() is not None

    def write(self, image: Image):
        if self.is_closed():
            return

        try:
            self.ffmpeg.stdin.write(image.convert('RGB').tobytes())
        except Exception as e:
            logger.warning('FFmpeg send error: {}'.format(str(e)))
            self.close()

    def close(self):
        if not self.is_closed():
            if self.ffmpeg and self.ffmpeg.stdin:
                try:
                    self.ffmpeg.stdin.close()
                except (IOError, OSError):
                    pass

            if self.ffmpeg:
                self.ffmpeg.terminate()
                try:
                    self.ffmpeg.wait(timeout=5.0)
                except subprocess.TimeoutExpired:
                    logger.warning('FFmpeg has not returned - killing it')
                    self.ffmpeg.kill()

            if self.ffmpeg and self.ffmpeg.stdout:
                try:
                    self.ffmpeg.stdout.close()
                except (IOError, OSError):
                    pass

        self.ffmpeg = None
        super().close()


class FFmpegFileWriter(FileVideoWriter, FFmpegWriter):
    """
    Write camera frames to a file using FFmpeg.
    """

    def __init__(self, *args, video_file: str, **kwargs):
        FileVideoWriter.__init__(self, *args, video_file=video_file, **kwargs)
        FFmpegWriter.__init__(self, *args, pix_fmt='rgb24', output_file=self.video_file, **kwargs)


class FFmpegStreamWriter(StreamWriter, FFmpegWriter, ABC):
    """
    Stream camera frames using FFmpeg.
    """

    def __init__(self, *args, output_format: str, **kwargs):
        StreamWriter.__init__(self, *args, **kwargs)
        FFmpegWriter.__init__(self, *args, pix_fmt='rgb24', output_format=output_format,
                              output_opts=[
                                  '-tune', '-zerolatency', '-preset', 'superfast', '-trellis', '0',
                                  '-fflags', 'nobuffer'], **kwargs)
        self._reader = threading.Thread(target=self._reader_thread)
        self._reader.start()

    def encode(self, image: Image) -> bytes:
        return image.convert('RGB').tobytes()

    def _reader_thread(self):
        start_time = time.time()

        while not self.is_closed():
            try:
                data = self.ffmpeg.stdout.read(1 << 15)
            except Exception as e:
                logger.warning('FFmpeg reader error: {}'.format(str(e)))
                break

            if not data:
                continue

            if self.frame is None:
                latency = time.time() - start_time
                logger.info('FFmpeg stream latency: {} secs'.format(latency))

            with self.ready:
                self.frame = data
                self.frame_time = time.time()
                self.ready.notify_all()

            self._sock_send(self.frame)

    def write(self, image: Image):
        if self.is_closed():
            return

        data = self.encode(image)
        try:
            self.ffmpeg.stdin.write(data)
        except Exception as e:
            logger.warning('FFmpeg send error: {}'.format(str(e)))
            self.close()

    def close(self):
        super().close()
        if self._reader and self._reader.is_alive() and threading.get_ident() != self._reader.ident:
            self._reader.join(timeout=5.0)
            self._reader = None


class MKVStreamWriter(FFmpegStreamWriter):
    mimetype = 'video/webm'

    def __init__(self, *args, **kwargs):
        super().__init__(*args, output_format='matroska', **kwargs)


class H264StreamWriter(FFmpegStreamWriter):
    mimetype = 'video/h264'

    def __init__(self, *args, **kwargs):
        super().__init__(*args, output_format='h264', **kwargs)


class H265StreamWriter(FFmpegStreamWriter):
    mimetype = 'video/h265'

    def __init__(self, *args, **kwargs):
        super().__init__(*args, output_format='h265', **kwargs)


# vim:sw=4:ts=4:et: