dailies.py

import os
import re
import sys
import yaml
import json
import glob
import subprocess as sp
from fileseq import FileSequence
import logging

log = logging.getLogger(__name__)


def set_logger(logger):
    global log
    log = logger

debug = os.environ.get('DEBUG')
if debug:
    log.setLevel(logging.DEBUG)
else:
    log.setLevel(logging.INFO)


class Dailies(object):
    """
    This class combines set of utilities for
    generating media for visual effects review sessions
    This currently include custom mov and slate generation
    """

    def __init__(self):

        script_dir = os.path.dirname(os.path.realpath(__file__))

        # Read module config from yaml file
        config_file = os.path.join(script_dir, 'config.yml')
        with open(config_file, 'r') as f:
            self.config = yaml.load(f)

        # Determine ffmpeg and ffprobe paths
        ffmpeg_dir = os.environ.get('FFMPEG_DIR')

        if os.environ.get('FFMPEG_DIR') is not None:
            # From the environmental variable
            ffmpeg_dir = os.environ.get('FFMPEG_DIR')
        elif config['ffmpeg_dir'][platform] is not None:
            # From configuration file
            raw_path = self.config['ffmpeg_dir'][platform]
            ffmpeg_dir = self._resolve_path(raw_path)

        else:
            exit_code = os.system('ffmpeg -version')
            if exit_code == 0:
                log.debug('Using system global ffmpeg')
            else:
                log.warning(
                    'Can not determine ffmpeg path. Make sure that ffmpeg '
                    'installed and in your PATH. You can also set FFMPEG_DIR '
                    'environmental variable or specify path in the config.yml'
                )
            ffmpeg_dir = ''

        self._ffmpeg = os.path.join(ffmpeg_dir, 'ffmpeg')
        self._ffprobe = os.path.join(ffmpeg_dir, 'ffprobe')

        log.debug('FFMPEG path: %s' % self._ffmpeg)

        # NOTE(Kirill): Probably should move all of this slate related variables
        # to the make_slate function

        # Slate resource files
        os.environ['RESOURCES'] = os.path.join(script_dir, 'resources')
        self.bars = self._resolve_path(self.config['bars'])
        self.color_bars = self._resolve_path(self.config['color_bar'])
        self.logo = self._resolve_path(self.config['company_logo'])
        # This paths used inside ffmpeg complex filter thus need special treatment
        self.logo_font_file = self._resolve_path(self.config['company_font'])
        self.font_file = self._resolve_path(self.config['body_font'])
        self.default_lut = self._resolve_path(self.config['default_lut'])

        # Global slate alignment properties
        self.left_text_margin = '(w)/2+150'
        self.top_text_margin = '380'
        self.font_size = 40
        self.line_spacing = 40
        self.font_color = 'White'

        self.new_x = '1920'
        self.new_y = '1080'

        # Contains data for all of the dynamic field used on slate
        self.fields_data = {
            'company_name': None,
            'project_name': None,
            'lut': None,
            'shot_name': None,
            'file_name': None,
            'fps': None,
            'frame_range': None,
            'frame_total': None,
            'handles': None,
            'comp_res': None,
            'date': None,
            'user': None,
            'description': None,
        }

        # This is to store temp files generated by this class
        self.tmp_files = []

    def _resolve_path(self, path):
        # Replace all of the environmental variable in path
        template_values = re.findall('\{(.*?)\}', path)
        # This will try to set all of the template values to
        # their corresponding environmental variable
        for v in template_values:
            try:
                v_val = os.environ[v]
            except KeyError:
                v_val = ''

            path = path.replace('{%s}' % v, v_val)

        return path

    def _get_tmp_dir(self):
        """
        :returns: (str) Path to temporary file directory for specific platform
        """
        tmp = None
        if sys.platform == 'darwin' or 'linux' in sys.platform:
            tmp = os.path.abspath(os.environ.get('TMPDIR'))
        elif sys.platform == 'win32' or sys.platform == 'cygwin':
            tmp = os.path.abspath(os.environ.get('TEMP')).replace(os.sep, '/')
            # On Windows, some usernames are shortened with a tilde.
            # Example: EVILEY~1 instead of evileye_52
            import getpass
            tmp = tmp.split('/')
            for item in tmp:
                if '~1' in item:
                    tmp[tmp.index(item)] = getpass.getuser()
            tmp = '/'.join(tmp)

        # Make sure that directory exists before returning it
        if not os.path.exists(tmp):
            os.makedirs(tmp)

        return tmp

    def _get_tmp_file(self, name):
        """
        This function does not create the actual file it only creates
        a temp folder for it and return the full path to use by the application

        :param name: (str) Name of desired temp file
        :returns: (str) Path to temporary file
        """
        tmp_dir = self._get_tmp_dir()
        if not os.path.exists(tmp_dir):
            os.mkdirs(tmp_dir)
        temp_file_path = os.path.join(tmp_dir, name)
        self.tmp_files.append(temp_file_path)

        return temp_file_path

    def _remove_tmp_files(self):
        """
        Remove temp files generated by this class
        """
        for f in self.tmp_files:
            if os.path.exists(f):
                os.remove(f)

    def _get_seq(self, path):
        """
        Crate file sequence object by given a file sequence path such as
        /path/image.%04d.dpx

        :param path: (str) File sequence path

        :returns: File Sequence object
        """
        s = re.sub(r'%[0-9]+d', '#', path)
        seq = FileSequence.findSequencesOnDisk(os.path.dirname(s))[0]

        return seq

    def _prep_for_filter(self, path):
        """
        Prepare file path to be passed to ffmpeg complex filter
        """
        # Use forward slashes
        path = path.replace('\\', '/')
        # Escape windows drive separator
        # Since complex filter use ':' to separate options we need to escape it.
        # Note that that we use double escape because first escaped will be
        # consumed by the shell whereas the second one by the ffmpeg itself
        path = path.replace(':', '\\\\:')

        return path

    def _check_exit_status(self, status, err_msg):
        """
        Check exit status of subprocess
        :param status: (bool) 0 - success, 1 - failure
        """
        # Check for errors and clean exit
        if status == 1:
            log.error(err_msg)
            log.info(
                'Please set DAILIES_DEBUG environmental variable '
                'to 3 to see debug info'
            )
            raise Exception(err_msg)

    def fields_from_dict(self, fields_dict):
        """
        Populate slate fields. Run before creating a slate
        """
        # TODO(Kirill): Made this a part of the make_slate function?
        for k, v in self.fields_data.items():
            self.fields_data[k] = fields_dict[k]

    def get_media_info(self, path):
        """
        (Kirill) This function is currently not used!

        Get information about movie file by using ffprobe

        :returns: (dict) Dictionary with media info about video stream
                         See ffprobe docs for more information
        """
        cmd = [
            str(self._ffprobe), '-v', 'quiet', '-select_streams', 'v',
            '-show_streams', '-print_format', 'json', str(path)
        ]

        output = None
        stream = {}

        try:
            result = sp.check_output(cmd)
            output = json.loads(result)
        except sp.CalledProcessError as e:
            log.error('ffprobe failed to extract information about the asset. %s' % e)
            log.debug('Test this command: %s' % ' '.join(cmd))
            raise
        except Exception as e:
            log.error('Error happened while executing CMD command. ' + str(e))
            log.debug('Test this command: %s' % ' '.join(cmd))
            raise

        if not output:
            log.warning('No media streams are found in %s' % path)
            return stream

        if len(output.get('streams')) == 1:
            stream = output['streams'][0]
        elif len(output.get('streams')) > 1:
            log.warning(
                'Media file %s contains more then one streams. '
                'Using the first one. '
            )
            stream = output['streams'][0]

        return stream

    def make_slate(self, src_seq, lut=''):
        """
        Generate an image slate out of given image sequence
        Slate information get set from self.fields_data member variable

        :param src_seq: (str) Path to source image sequence
                        in the following format /path/name_%04d.ext
        :param lut: (str) Path to optional cube LUT file

        :returns: (str) Path to generated slate image
        """

        output = self._get_tmp_file('tmp_slate.png')

        seq = self._get_seq(src_seq)

        # Alignment properties for slate field
        p = "x={left_text_margin}-text_w:y={top_text_margin}+{line_spacing}".format(
            left_text_margin=self.left_text_margin,
            top_text_margin=self.top_text_margin,
            line_spacing=self.line_spacing
        )

        # Alignment properties for slate field value
        pv = "x={left_text_margin}+10:y={top_text_margin}+{line_spacing}".format(
            left_text_margin=self.left_text_margin,
            top_text_margin=self.top_text_margin,
            line_spacing=self.line_spacing
        )

        # Slate main body text style
        text = "drawtext=fontsize={font_size}:fontcolor={font_color}:fontfile={font_file}:text".format(
            font_size=self.font_size,
            font_color=self.font_color,
            font_file=self._prep_for_filter(self.font_file)
        )

        if lut:
            lut_filter = (
                "[thumbnail] lut3d={lut_path} [thumbnail];"
            ).format(lut_path=self._prep_for_filter(lut))
        else:
            lut_filter = ''

        # ffmpeg complex filter to generate a slate. For more information see
        # section four (Filtergraph description) of the following docs
        # https://ffmpeg.org/ffmpeg-filters.html
        filters = (
            "[1:v] scale={new_x}:{new_y}, setsar=1:1 [base]; "
            "[0:v] scale={new_x}:{new_y} [thumbnail]; "
            "[thumbnail][3:v] overlay [thumbnail]; "
            "[thumbnail][3:v] overlay=x=(main_w-overlay_w):y=(main_h-overlay_h) [thumbnail]; "
            "[thumbnail] scale=(iw/4):(ih/4) [thumbnail]; {lut_filter}"
            "[base][thumbnail] overlay=((main_w-overlay_w)/2)-500:(main_h-overlay_h)/2 [base]; "
            "[2:v] scale=-1:-1 [self.bars]; "
            "[base][self.bars] overlay=x=(main_w-overlay_w):y=(main_h-overlay_h-50) [base]; "
            "[4:v] scale=(iw*0.2):(ih*0.2) [self.logo]; "
            "[base][self.logo] overlay=x=500:y=100 [base]; "
            "[base] "
              "drawtext=fontsize=80:fontcolor={font_color}:fontfile={logo_font_file}:text={company_name}:x=690:y=130, "
              "drawtext=fontsize=50:fontcolor={font_color}:fontfile={font_file}:text={project_name}:x=(w)/2:y=250, "
                   "{text}='LUT\: ':{p}*0, "
                   "{text}={lut}:{pv}*0, "
                   "{text}='Shot name\: ':{p}*1, "
                   "{text}={shot_name}:{pv}*1, "
                   "{text}='File name\: ':{p}*2, "
                   "{text}={file_name}:{pv}*2, "
                   "{text}='FPS\: ':{p}*3, "
                   "{text}={fps}:{pv}*3, "
                   "{text}='Frame range\: ':{p}*4, "
                   "{text}={frame_range}:{pv}*4, "
                   "{text}='Frame total\: ':{p}*5, "
                   "{text}={frame_total}:{pv}*5, "
                   "{text}='Handles\: ':{p}*6, "
                   "{text}={handles}:{pv}*6, "
                   "{text}='Comp resolution\: ':{p}*7, "
                   "{text}={comp_res}:{pv}*7, "
                   "{text}='Date\: ':{p}*8, "
                   "{text}={date}:{pv}*8, "
                   "{text}='User\: ':{p}*9, "
                   "{text}={user}:{pv}*9, "
                   "{text}='Description\: ':{p}*10, "
                   "{text}={description}:{pv}*10 "
        ).format(
            #
            # Global formatting and render values
            font_color=self.font_color, logo_font_file=self._prep_for_filter(self.logo_font_file),
            font_file=self._prep_for_filter(self.font_file), line_spacing=self.line_spacing,
            left_text_margin=self.left_text_margin,
            top_text_margin=self.top_text_margin, new_x=self.new_x,
            new_y=self.new_y, text=text, p=p, pv=pv, lut_filter=lut_filter,
            #
            # User defined fields slate values
            project_name=self.fields_data['project_name'],
            company_name=self.fields_data['company_name'],
            lut=self.fields_data['lut'],
            shot_name=self.fields_data['shot_name'],
            file_name=self.fields_data['file_name'],
            fps=self.fields_data['fps'],
            frame_range=self.fields_data['frame_range'],
            frame_total=self.fields_data['frame_total'],
            handles=self.fields_data['handles'],
            comp_res=self.fields_data['comp_res'],
            date=self.fields_data['date'],
            user=self.fields_data['user'],
            description=self.fields_data['description']
        )

        cmd = [
            self._ffmpeg, '-v', 'quiet', '-y', '-start_number', str(seq.start()), '-i', str(src_seq),
            '-f', 'lavfi', '-i', 'color=c=black', '-i', self.bars, '-i', self.color_bars,
            '-i', self.logo, '-vframes', '1', '-filter_complex', filters, str(output)
        ]

        if debug >= '2':
            log.debug('SLATE_COMMAND: %s' % cmd)
        if debug >= '3':
            cmd.pop(1)
            cmd.pop(1)

        exit_status = sp.call(cmd)

        msg = 'Error occurred during slate generation.'
        self._check_exit_status(exit_status, msg)

        return output

    def make_mov(self, src_seq, out_mov, preset='', burnin=True, slate=True, lut=''):
        """
        Generate movie out of image sequence with and optional slate and burnins

        :param src_seq: (str) Path to a file sequence /path/image.%04d.png
        :param out_mov: (str) Output path for generated movie file
        :param preset: (str) Name of video preset from 'config.yml'
        :param burnin: (bool) Generate burn in text with file name and frame
                       number on every frame of the video
        :param slate: (bool) Generate slate image and append it as the first
                      frame of the final video
        :param lut: (str) Path to optional cube LUT file

        :returns: (str) Path to generated movie file
        """

        # Get video settings for ffmpeg from the configuration file
        video_presets = self.config['video_pressets']
        video_settings = video_presets.get(preset)

        if video_settings is not None:
            video_settings = video_settings.split(' ')
        else:
            # Use default settings
            video_settings = [
                '-crf', '18', '-vcodec', 'mjpeg', '-pix_fmt', 'yuvj444p',
                '-qmin', '1', '-qmax', '1', '-r', '24'
            ]

        seq = self._get_seq(src_seq)

        parent, file_name = os.path.split(src_seq)

        filters = (
            "[1:v] scale={new_x}:{new_y}, setsar=1:1 [base]; "
            "[base] null "
        ).format(new_x=self.new_x, new_y=self.new_y)

        if lut == 'default':
            # Use default LUT specified in config.yml
            lut = self.default_lut

        if lut:
            # Apply cube LUT to the video
            lut_filter = (
                "[base]; [base] lut3d={lut_path}"
            ).format(lut_path=self._prep_for_filter(lut))

            filters += lut_filter

        if slate:
            # Concatenate our slate and the rest of the video
            slate_filter = (
                "[base]; "
                "[0:v] trim=start_frame=0:end_frame=1 [slate]; "
                "[slate][base] concat "
            )
            filters += slate_filter

        if burnin:
            # Add burin text to the filter
            # This filter generate text with original sequence file name
            # and frame number on the every frame of the video
            burnin_filter = (
                "[base]; [base] "
                  # Left button sequence name
                  "drawtext=fontsize=30: "
                  "fontcolor={font_color}: "
                  "fontfile={font_file}: "
                  "expansion=none: "
                  "text={file_name}: "
                  "x=10:y=(h-(text_h+10)): "
                  "enable='between(n,1,99999)', "
                  # Right button frame counter
                  "drawtext=fontsize=30: "
                  "fontcolor={font_color}: "
                  "fontfile={font_file}: "
                  "start_number=1000: "
                  "text=@frame_number \[{start}-{end}\]: "
                  "x=(w-(text_w+10)):y=(h-(text_h+10)): "
                  "enable='between(n,1,99999)'"
            ).format(
                font_color=self.font_color, font_file=self._prep_for_filter(self.font_file),
                file_name=file_name, start=seq.start(), end=seq.end()
            )
            filters += burnin_filter

        # Replace this frame number variable separately
        # since it conflict with the string format function tokens
        filters = filters.replace('@frame_number', '%{n}')

        # Generate an mov with the attached slate
        cmd = [self._ffmpeg, '-v', 'quiet', '-y']

        if slate:
            # Generate slate image
            slate = self.make_slate(src_seq, lut=lut)

            # Append slate stream as stream 0 to the final command
            cmd += ['-i', slate]
        else:
            # Append null source to not mess up indexes if the input for the filters
            cmd += ['-f', 'lavfi', '-i', 'nullsrc=s=256x256:d=5']

        cmd += ['-start_number', str(seq.start()), '-i', src_seq]
        cmd += video_settings
        cmd += ['-filter_complex', filters, out_mov]

        if debug >= '2':
            log.debug('MOV_COMMAND: %s' % cmd)
        if debug >= '3':
            # Remove silent flags from the final command
            cmd.pop(1)
            cmd.pop(1)

        exit_status = sp.call(cmd)

        msg = 'Error occurred during mov generation.'
        self._check_exit_status(exit_status, msg)

        self._remove_tmp_files()

        return out_mov