353 lines
15 KiB
Python
353 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from __future__ import unicode_literals
|
|
from io import open
|
|
|
|
class AutoMKVChapters:
|
|
class Template:
|
|
def __init__(self):
|
|
from random import randint
|
|
self.uid = randint(10**4,10**6)
|
|
self.num_editions = 1
|
|
self.lang = ['eng']
|
|
self.country = ['us']
|
|
self.fps = '30'
|
|
self.ofps = '24'
|
|
self.qpf = '0'
|
|
self.idr = False
|
|
self.trims = None
|
|
self.kframes = None
|
|
|
|
def toxml(self,chapfile):
|
|
|
|
chf = open(chapfile+'.xml','w',encoding='utf-8')
|
|
head = '<?xml version="1.0" encoding="UTF-8"?>\n<!-- <!DOCTYPE Tags SYSTEM "matroskatags.dtd"> -->\n'
|
|
chf.write(head+'<Chapters>\n')
|
|
|
|
if self.num_editions > 1:
|
|
tagf = open(chapfile+'tags.xml','w')
|
|
tagf.write(head+'<Tags>\n')
|
|
else:
|
|
tagf = False
|
|
|
|
for ed in self.editions:
|
|
chf.write('\t<EditionEntry>\n')
|
|
chf.write('\t\t<EditionFlagHidden>{0:d}</EditionFlagHidden>\n'.format(ed.hidden) if ed.hidden else '')
|
|
chf.write('\t\t<EditionFlagDefault>{0:d}</EditionFlagDefault>\n'.format(ed.default) if ed.default else '')
|
|
chf.write('\t\t<EditionFlagOrdered>{0:d}</EditionFlagOrdered>\n'.format(ed.ordered) if ed.ordered else '')
|
|
chf.write('\t\t<EditionUID>{0:d}</EditionUID>\n'.format(ed.uid))
|
|
|
|
if tagf:
|
|
tagf.write('\t<Tag>\n\t\t<Targets>\n')
|
|
tagf.write('\t\t\t<EditionUID>{0:d}</EditionUID>\n'.format(ed.uid))
|
|
tagf.write('\t\t\t<TargetTypeValue>50</TargetTypeValue>\n\t\t</Targets>\n')
|
|
num_names = len(ed.name) if len(ed.name) < len(self.lang) else len(self.lang)
|
|
for i in range(num_names):
|
|
tagf.write('\t\t<Simple>\n\t\t\t<Name>TITLE</Name>\n')
|
|
tagf.write('\t\t\t<String>{0}</String>\n'.format(ed.name[i] if ed.name[i] != '' else ed.name[i-1]))
|
|
tagf.write('\t\t\t<TagLanguage>{0}</TagLanguage>\n'.format(self.lang[i]))
|
|
tagf.write('\t\t\t<DefaultLanguage>{0:d}</DefaultLanguage>\n'.format(1 if i == 0 else 0))
|
|
tagf.write('\t\t</Simple>\n')
|
|
tagf.write('\t</Tag>\n')
|
|
|
|
for ch in ed.chapters:
|
|
chf.write('\t\t<ChapterAtom>\n')
|
|
num_names = len(ch.name) if len(ch.name) < len(self.lang) else len(self.lang)
|
|
for i in range(num_names):
|
|
chf.write('\t\t\t<ChapterDisplay>\n')
|
|
chf.write('\t\t\t\t<ChapterString>{0}</ChapterString>\n'.format(ch.name[i] if ch.name[i] != '' else ch.name[i-1]))
|
|
chf.write('\t\t\t\t<ChapterLanguage>{0}</ChapterLanguage>\n'.format(self.lang[i]) if self.lang[i] != 'eng' else '')
|
|
chf.write('\t\t\t\t<ChapterCountry>{0}</ChapterCountry>\n'.format(self.country[i]) if i < len(self.country) else '')
|
|
chf.write('\t\t\t</ChapterDisplay>\n')
|
|
chf.write('\t\t\t<ChapterUID>{0:d}</ChapterUID>\n'.format(ch.uid))
|
|
chf.write('\t\t\t<ChapterTimeStart>{0}</ChapterTimeStart>\n'.format(ch.start))
|
|
chf.write('\t\t\t<ChapterTimeEnd>{0}</ChapterTimeEnd>\n'.format(ch.end) if ch.end else '')
|
|
chf.write('\t\t\t<ChapterFlagHidden>{0:d}</ChapterFlagHidden>\n'.format(ch.hidden) if ch.hidden != 0 else '')
|
|
chf.write('\t\t\t<ChapterFlagEnabled>{0:d}</ChapterFlagEnabled>\n'.format(ch.enabled) if ch.enabled != 1 else '')
|
|
chf.write('\t\t\t<ChapterSegmentUID format="hex">{0}</ChapterSegmentUID>\n'.format(ch.suid) if ch.suid else '')
|
|
chf.write('\t\t</ChapterAtom>\n')
|
|
|
|
chf.write('\t</EditionEntry>\n')
|
|
|
|
chf.write('</Chapters>\n')
|
|
|
|
if tagf:
|
|
tagf.write('</Tags>\n')
|
|
tagf.close()
|
|
|
|
if self.qpf != '0' and self.kframes:
|
|
from vfr import write_qpfile
|
|
if self.qpf != '1':
|
|
qpfile = self.qpf
|
|
else:
|
|
qpfile = chapfile+'.qpfile'
|
|
write_qpfile(qpfile, self.kframes, self.idr)
|
|
|
|
def connect_with_vfr(self,avs,label=None,clip=None):
|
|
"""
|
|
Connects templates.py with vfr.py, enabling its use outside of vfr.py.
|
|
|
|
Uses the same quirks as AMkvC but only for 24 and 30 fps.
|
|
Ex: inputfps=30 is understood as being '30*1000/1001'
|
|
|
|
"""
|
|
|
|
from vfr import parse_trims, fmt_time
|
|
|
|
# compensate for amkvc's fps assumption
|
|
if self.fps in ('24','30'):
|
|
fps = self.fps + '/1.001'
|
|
else:
|
|
fps = str(self.fps)
|
|
if self.ofps and self.ofps in ('24','30'):
|
|
ofps = self.ofps + '/1.001'
|
|
else:
|
|
ofps = str(self.ofps)
|
|
|
|
Trims2, Trims2ts = parse_trims(avs, fps, ofps, label=label, clip=clip)[2:4]
|
|
Trims2ts = [(fmt_time(i[0]),fmt_time(i[1]) if i[1] != 0 else None) for i in Trims2ts]
|
|
|
|
self.trims = Trims2ts
|
|
self.kframes = Trims2
|
|
|
|
def parse_mkv(self, path):
|
|
"""Parse a Matroska file for SegmentUID and Duration"""
|
|
import binascii
|
|
import struct
|
|
|
|
def get_data_len(byte):
|
|
"""Get the length (bytes) of the element data"""
|
|
n = ord(byte)
|
|
mask = 0b10000000
|
|
while not n & mask:
|
|
mask >>= 1
|
|
return n & ~mask
|
|
|
|
suid = tcscale = duration = 0
|
|
with open(path, 'rb') as file:
|
|
if file.read(4) != b'\x1A\x45\xDF\xA3': # not a Matroska file
|
|
return suid, duration
|
|
chunk_size = 100000 # 100 kB
|
|
i = 0
|
|
while True:
|
|
if suid and tcscale and duration:
|
|
break
|
|
bin = file.read(chunk_size)
|
|
if not bin:
|
|
break
|
|
suid_pos = bin.find(b'\x73\xA4\x90') # \x90 -> 16 bytes
|
|
if suid_pos != -1:
|
|
suid_pos = 4 + i * chunk_size + suid_pos + 3
|
|
file.seek(suid_pos)
|
|
suid = binascii.hexlify(file.read(16)).decode()
|
|
tcscale_pos = bin.find(b'\x2A\xD7\xB1')
|
|
if tcscale_pos != -1:
|
|
tcscale_pos = 4 + i * chunk_size + tcscale_pos + 3
|
|
file.seek(tcscale_pos)
|
|
tcscale_len = get_data_len(file.read(1))
|
|
tcscale = int(binascii.hexlify(file.read(tcscale_len)), 16)
|
|
duration_pos = bin.find(b'\x44\x89\x84') # float (4 bytes)
|
|
if duration_pos != -1:
|
|
duration_pos = 4 + i * chunk_size + duration_pos + 3
|
|
file.seek(duration_pos)
|
|
duration = struct.unpack('>f', file.read(4))[0]
|
|
if not duration: # double (8 bytes)
|
|
duration_pos = bin.find(b'\x44\x89\x88')
|
|
if duration_pos != -1:
|
|
duration_pos = 4 + i * chunk_size + duration_pos + 3
|
|
file.seek(duration_pos)
|
|
duration = struct.unpack('>d', file.read(8))[0]
|
|
if bin.find(b'\x1F\x43\xB6\x75') != -1:
|
|
# segment info should be before the clusters
|
|
break
|
|
i += 1
|
|
duration = duration * tcscale / 1000000
|
|
return suid, duration
|
|
|
|
class Edition:
|
|
def __init__(self):
|
|
self.default = 0
|
|
self.name = ['Default']
|
|
self.hidden = 0
|
|
self.ordered = 0
|
|
self.num_chapters = 1
|
|
self.uid = 0
|
|
|
|
class Chapter:
|
|
def __init__(self):
|
|
self.name = ['Chapter']
|
|
self.chapter = False
|
|
self.start = False
|
|
self.end = False
|
|
self.suid = False
|
|
self.hidden = 0
|
|
self.uid = 0
|
|
self.enabled = 1
|
|
|
|
def __init__(self, templatefile, output=None, avs=None, trims=None,
|
|
kframes=None, uid=None, label=None, ifps=None, clip=None,
|
|
idr=False):
|
|
try:
|
|
import configparser
|
|
except ImportError:
|
|
import ConfigParser as configparser
|
|
from io import open
|
|
|
|
# Init config
|
|
config = configparser.ConfigParser()
|
|
template = open(templatefile, encoding='utf-8')
|
|
|
|
# Read template
|
|
config.readfp(template)
|
|
template.close()
|
|
|
|
# Template defaults
|
|
self = self.Template()
|
|
self.editions = []
|
|
self.uid = uid if uid else self.uid
|
|
|
|
# Set mkvinfo path
|
|
from vfr import mkvmerge, parse_with_mkvmerge, fmt_time
|
|
from os.path import dirname, join, isfile
|
|
|
|
# Set placeholder for mkvinfo output
|
|
mkv_globbed = False
|
|
mkvinfo = {}
|
|
|
|
for k, v in config.items('info'):
|
|
if k == 'lang':
|
|
self.lang = v.split(',')
|
|
elif k == 'country':
|
|
self.country = v.split(',')
|
|
elif k == 'inputfps':
|
|
self.fps = v
|
|
elif k == 'outputfps':
|
|
self.ofps = v
|
|
elif k == 'createqpfile':
|
|
self.qpf = v
|
|
elif k == 'uid':
|
|
self.uid = int(v)
|
|
elif k == 'editions':
|
|
self.num_editions = int(v)
|
|
|
|
if avs and not ifps:
|
|
self.connect_with_vfr(avs, label, clip)
|
|
elif trims:
|
|
self.trims = trims
|
|
self.kframes = kframes
|
|
else:
|
|
self.trims = False
|
|
self.idr = idr
|
|
|
|
for i in range(self.num_editions):
|
|
from re import compile
|
|
ed = self.Edition()
|
|
ed.uid = self.uid * 100
|
|
self.uid += 1
|
|
cuid = ed.uid
|
|
ed.num = i+1
|
|
ed.chapters = []
|
|
stuff = {}
|
|
|
|
for k, v in config.items('edition{0:d}'.format(ed.num)):
|
|
if k == 'default':
|
|
ed.default = int(v)
|
|
elif k == 'name':
|
|
ed.name = v.split(',')
|
|
elif k == 'ordered':
|
|
ed.ordered = int(v)
|
|
elif k == 'hidden':
|
|
ed.hidden = int(v)
|
|
elif k == 'chapters':
|
|
ed.num_chapters = int(v)
|
|
for i in range(ed.num_chapters):
|
|
stuff[i+1] = []
|
|
elif k == 'uid':
|
|
ed.uid = int(v)
|
|
else:
|
|
opt_re = compile('(\d+)(\w+)')
|
|
ret = opt_re.search(k)
|
|
if ret:
|
|
stuff[int(ret.group(1))].append((ret.group(2),v))
|
|
|
|
for j in range(ed.num_chapters):
|
|
ch = self.Chapter()
|
|
cuid += 1
|
|
ch.uid = cuid
|
|
ch.num = j+1
|
|
|
|
for k, v in stuff[j+1]:
|
|
if k == 'name':
|
|
ch.name = v.split(',')
|
|
elif k == 'chapter':
|
|
ch.chapter = int(v)
|
|
elif k == 'start':
|
|
ch.start = v
|
|
elif k == 'end':
|
|
ch.end = v
|
|
elif k == 'suid':
|
|
ch.suid = v.strip() if ret else 0
|
|
elif k == 'hidden':
|
|
ch.hidden = int(v)
|
|
elif k == 'enabled':
|
|
ch.enabled = int(v)
|
|
|
|
if ch.suid and not isfile(ch.suid):
|
|
ch.suid = ch.suid.replace('0x','').lower().replace(' ','')
|
|
|
|
if ch.chapter and not (ch.start and ch.end):
|
|
ch.start, ch.end = self.trims[ch.chapter-1] if self.trims else (ch.start, ch.end)
|
|
elif ch.suid:
|
|
mkvfiles = []
|
|
if isfile(ch.suid):
|
|
mkvfiles = [ch.suid]
|
|
elif not mkv_globbed:
|
|
from glob import glob
|
|
mkvfiles = glob('*.mkv') + glob(join(dirname(avs),'*.mkv'))
|
|
mkv_globbed = True
|
|
if mkvfiles:
|
|
if parse_with_mkvmerge:
|
|
from subprocess import check_output
|
|
import json
|
|
for file in mkvfiles:
|
|
info = check_output([mkvmerge, '-i', '-F', 'json',
|
|
'--output-charset', 'utf-8', file]).decode('utf-8')
|
|
try:
|
|
props = json.loads(info).get("container", {}).get("properties", {})
|
|
ch.suid = props.get("segment_uid", 0)
|
|
duration = props.get("duration", 0)
|
|
mkvinfo[ch.suid] = {'file': file,
|
|
'duration': fmt_time(duration * 10**6)
|
|
if duration else 0}
|
|
except Exception:
|
|
pass
|
|
else:
|
|
for file in mkvfiles:
|
|
ch.suid, duration = self.parse_mkv(file)
|
|
mkvinfo[ch.suid] = {'file': file,
|
|
'duration': fmt_time(duration * 10**6)
|
|
if duration else 0}
|
|
if not (ch.start or ch.end):
|
|
ch.start = fmt_time(0) if not ch.start else ch.start
|
|
ch.end = mkvinfo[ch.suid]['duration'] if not ch.end and (ch.suid in mkvinfo) else ch.end
|
|
|
|
ed.chapters.append(ch)
|
|
self.editions.append(ed)
|
|
if output:
|
|
self.toxml(output)
|
|
|
|
|
|
def main(args):
|
|
|
|
template = args[0]
|
|
output = args[1]
|
|
avs = args[2] if len(args) == 3 else None
|
|
|
|
chaps = AutoMKVChapters(template,output,avs)
|
|
|
|
if __name__ == '__main__':
|
|
from sys import argv, exit
|
|
if len(argv) > 1:
|
|
main(argv[1:])
|
|
else:
|
|
exit("templates.py <template file> <output filenames w/o extension> [<avisynth file>]")
|