port of EasyVFR for VapourSynth

easyvfr.py

#!/usr/bin/env python3

'''port of EasyVFR for VapourSynth'''

__author__ = 'Chikuzen <chikuzen.mo at gmail dot com>'
__version__ = '0.2.5'


import vapoursynth as vs


def reduce_fraction(a, b):
    def get_gcd(x, y):
        if y == 0:
            return x
        return get_gcd(y, x % y)

    gcd = get_gcd(a, b)
    return a // gcd, b // gcd


def get_timecode_and_chapter(ch_clips, base_num, base_den):
    def generate_chapter(ch_number, pts, name):
        chapter = ('CHAPTER{num}={hour}:{minute}:{sec}.{msec}\n'
                   'CHAPTER{num}NAME={name}\n')
        hour, minute = divmod(int(round(pts / 1000)), 3600000000)
        minute, sec = divmod(minute, 60000000)
        sec, msec = divmod(sec, 1000000)
        msec /= 1000
        return chapter.format(num='%02i' % ch_number, hour='%02i' % hour,
                              minute='%02i' % minute, sec='%02i' % sec,
                              msec='%03i' % msec, name=name)

    timestamps = [0.0]
    chapters = []
    keyframes = []
    pts = 0.0 # nano_second
    ch_number = 1
    fr_number = 0

    for cc in ch_clips:
        clip = cc.clip

        if cc.name is not None:
            chapters.append(generate_chapter(ch_number, pts, cc.name))
            keyframes.append(fr_number)
            ch_number += 1

        cycle_clip, cycle_base = reduce_fraction(clip.fps_num, base_num)
        frame_duration = base_den * 1000000000 / clip.fps_num

        for count in range(clip.num_frames):
            pts += frame_duration
            timestamps.append(round(pts / 1000000, 3))

        fr_number += clip.num_frames
        # clip end correction
        if cycle_clip < cycle_base:
            adjust = clip.num_frames % cycle_clip
            if adjust != 0:
                pts += (base_den * 1000000000 * (adjust + 1) / base_num -
                        frame_duration * adjust)
            timestamps[-1] = round(pts / 1000000, 3)

    del timestamps[-1]

    return timestamps, chapters, keyframes


class InvalidArgument(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


class ChapterClip(object):
    '''Class for set chapter point to clip'''
    def __init__(self, clip, name=None):
        '''
        clip: VideoNode(aka clip) instance
        name: chapter name
        '''
        if not isinstance(clip, vs.VideoNode):
            raise TypeError('clip is not VideoNode instance')
        self.clip = clip
        self.name = None if name is None else str(name)


class EasyVFR(object):
    '''
    Class for splice different framerate clips and outputting timecode and
    chapter file
    '''
    def __init__(self, clips, base_num=30000, base_den=1001, base=None):
        '''
        clips: list of CFR clips.
        base_num: framerate numerator used as a base.
        base_den: framerate denominator used as a base.
        base: clip instance.
              if this is set, base_num and base_den are overwritten by
              the values of this.
        '''
        self.std = vs.get_core().std

        if isinstance(base, vs.VideoNode):
            self.base_num = base.fps_num
            self.base_den = base.fps_den
        else:
            self.base_num = int(base_num)
            self.base_den = int(base_den)
        if self.base_num < 1 or self.base_den < 1:
            raise InvalidArgument(
                'base_num and base_den must be higher than zero')

        self.ch_clips = []
        for i in range(len(clips)):
            cc = clips[i]
            if not isinstance(cc, ChapterClip):
                try:
                    cc = ChapterClip(clips[i])
                except:
                    raise TypeError('clips[%i] is not a clip.' % i)

            if cc.clip.fps_num == 0:
                raise InvalidArgument('clips[%i] is not CFR clip.' % i)

            n, d = reduce_fraction(cc.clip.fps_num, cc.clip.fps_den)
            cc.clip = self.std.AssumeFPS(cc.clip, fpsnum=n, fpsden=d)
            if cc.clip.fps_den != self.base_den:
                raise InvalidArgument('clips[%i] has invalid fps_den.' % i)

            self.ch_clips.append(cc)

        self.timestamps, self.chapters, self.keyframes = \
            get_timecode_and_chapter(self.ch_clips, self.base_num,
                                     self.base_den)

    def splice_clips(self):
        '''Return a clip which spliced all clips to one.'''
        clips = [cc.clip for cc in self.ch_clips]
        ret = self.std.Splice(clips, mismatch=True)
        return self.std.AssumeFPS(ret, fpsnum=self.base_num,
                                  fpsden=self.base_den)

    def write_timecode(self, filename):
        '''
        Output a matroska timecode v2 format file.
        filename: The name of the file to output.
        '''
        with open(filename, 'w') as o:
            print('# timecode format v2', file=o)
            print('\n'.join([str(pts) for pts in self.timestamps]), file=o)

    def write_chapter(self, filename):
        '''
        Output a simple(ogm) format chapter file.
        filename: The name of the file to output.
        '''
        with open(filename, 'wb') as o:
            o.write("\n".join(self.chapters).encode('utf-8'))

    def write_qpfile(self, filename):
        '''
        Output a x264 format qpfile.
        filename: The name of the file to output.
        '''
        with open(filename, 'w') as o:
            print('\n'.join(['%i K' % number for number in self.keyframes]),
                  file=o)

    def write_tcq(self, filename, ext_tc='tc', ext_ch='ch', ext_qp='qp'):
        '''
        Output timecode, chapter and qpfile at once.
        filename: The name of the files to output.
        ext_tc: file extension of timecode.
        ext_ch: file extension of chapter.
        ext_qp: file extension of qpfile.
        '''
        self.write_timecode('%s.%s' % (filename, ext_tc))
        self.write_chapter('%s.%s' % (filename, ext_ch))
        self.write_qpfile('%s.%s' % (filename, ext_qp))

if __name__ == '__main__':
    print('port of EasyVFR for VapourSynth version %s' % __version__)