2019-07-14 13:12:16 +00:00
#!/usr/bin/env python
"""
2019-10-17 21:47:31 +00:00
Generic script to generate keyframes for all files of a given extension
using kagefunc ' s generate_keyframes function.
2019-08-01 09:46:35 +00:00
Dependencies :
2019-08-16 07:24:41 +00:00
* VapourSynth
* wwxd ( https : / / github . com / dubhater / vapoursynth - wwxd )
2020-05-17 02:27:11 +00:00
* vapoursynth - scxvid ( Optional : if set to use over wwxd ) ( https : / / github . com / dubhater / vapoursynth - scxvid )
2019-07-14 13:12:16 +00:00
"""
import argparse
2020-05-05 04:14:30 +00:00
import glob
import mimetypes
2019-07-14 13:12:16 +00:00
import os
2019-08-19 21:42:51 +00:00
from ast import literal_eval
2020-05-05 04:14:30 +00:00
import vapoursynth as vs
2019-07-14 13:12:16 +00:00
core = vs . core
2020-05-05 04:14:30 +00:00
__author__ = " LightArrowsEXE "
__license__ = ' MIT '
2020-05-17 02:44:46 +00:00
__version__ = ' 1.2 '
2020-05-05 04:14:30 +00:00
2019-08-01 09:46:35 +00:00
# Slightly modified from kagefunc to remove some dependencies
2019-08-16 07:24:41 +00:00
def generate_keyframes ( clip : vs . VideoNode , out_path = None , no_header = False ) - > None :
2019-08-01 09:46:35 +00:00
"""
2020-05-17 02:27:11 +00:00
The exact same version as found in kagefunc .
2019-08-01 09:46:35 +00:00
"""
2020-05-05 04:14:30 +00:00
clip = core . resize . Bilinear ( clip , 640 , 360 , format = vs . YUV420P8 )
2019-10-17 21:47:31 +00:00
2020-05-17 02:27:11 +00:00
if args . scxvid :
clip = core . scxvid . Scxvid ( clip )
else :
clip = core . wwxd . WWXD ( clip ) # speed up the analysis by resizing first
2020-05-17 02:44:46 +00:00
out_txt = ' ' if no_header else " # WWXD log file, using qpfile format \n # Please do not modify this file \n \n "
2019-10-17 21:47:31 +00:00
2019-08-01 09:46:35 +00:00
for i in range ( clip . num_frames ) :
2020-05-17 02:27:11 +00:00
if args . scxvid :
if clip . get_frame ( i ) . props [ " _SceneChangePrev " ] == 1 :
out_txt + = " %d I -1 \n " % i
else :
if clip . get_frame ( i ) . props [ " Scenechange " ] == 1 :
out_txt + = " %d I -1 \n " % i
2019-10-17 09:48:08 +00:00
if i % 1 == 0 :
print ( f " Progress: { i } / { clip . num_frames } frames " , end = " \r " )
2019-08-01 09:46:35 +00:00
text_file = open ( out_path , " w " )
text_file . write ( out_txt )
text_file . close ( )
def main ( ) :
2019-10-17 22:03:47 +00:00
if args . outfile and not args . file :
print ( " Warning: Please set --file (-F) when using --outfile (-O)! " )
return
2019-08-16 07:24:41 +00:00
if args . file :
2019-10-17 21:47:31 +00:00
files = [ args . file ]
2019-07-14 13:12:16 +00:00
else :
2019-10-17 22:03:47 +00:00
files = glob . glob ( ' **/* ' , recursive = True ) if args . recursive else glob . glob ( ' * ' )
2019-08-16 07:24:41 +00:00
2019-07-14 13:12:16 +00:00
for f in files :
2020-05-05 04:14:30 +00:00
mime = mimetypes . types_map . get ( os . path . splitext ( f ) [ - 1 ] , " " )
# Not entirely sure why mkv's fail, as they have a mimetype of "video/x-matroska"
if mime . startswith ( " video/ " ) or f . endswith ( ' .m2ts ' ) or f . endswith ( ' .mkv ' ) :
2019-11-21 09:37:19 +00:00
if args . check_exists :
2020-05-05 04:14:30 +00:00
if os . path . exists ( f " { os . path . splitext ( f ) [ 0 ] } _keyframes.txt " ) :
2019-11-21 09:37:19 +00:00
print ( f " \n Keyframes already exist for { f } . Skipping. " )
continue
2019-10-17 21:47:31 +00:00
2020-05-05 04:14:30 +00:00
src = core . lsmas . LWLibavSource ( f ) if f . endswith ( " .m2ts " ) else core . ffms2 . Source ( f )
2019-10-17 21:47:31 +00:00
2020-05-05 04:14:30 +00:00
mime = mimetypes . types_map . get ( os . path . splitext ( f ) [ - 1 ] ) or " Unknown "
print ( f " \n Video info \n Filename: { f } \n Dimensions: { src . width } x { src . height } \n Framerate: { src . fps } \n Format: { src . format . name } \n Mimetype: { mime } \n " )
2019-08-19 21:42:51 +00:00
if args . trims :
trims = literal_eval ( args . trims )
if type ( trims ) is not tuple :
trims = ( trims , )
2019-08-19 22:40:56 +00:00
try :
src = core . std . Splice ( [ src [ slice ( * trim ) ] for trim in trims ] )
except :
2019-10-17 21:47:31 +00:00
print ( " TypeError: Please make sure you’re using a list for this function. \n Example: -T \" [24,-24] \" , -T \" [None,30000],[None,-24] \" , -T \" [None,16000],[16100,16200],[16300,None] \" " )
2019-08-19 22:40:56 +00:00
return
2019-08-19 21:42:51 +00:00
2019-10-17 21:47:31 +00:00
if args . outfile and args . file :
generate_keyframes ( src , os . path . join ( os . path . dirname ( f ) , args . outfile ) , args . noheader )
2019-08-16 07:24:41 +00:00
else :
2020-05-05 04:14:30 +00:00
generate_keyframes ( src , os . path . abspath ( f " { os . path . splitext ( f ) [ 0 ] } _keyframes.txt " ) , args . noheader )
2019-10-17 21:47:31 +00:00
print ( f " Progress: { src . num_frames } / { src . num_frames } frames " )
2019-08-19 21:42:51 +00:00
2019-11-21 09:37:19 +00:00
try :
2020-05-05 04:14:30 +00:00
os . remove ( f " { f } .lwi " ) if f . endswith ( " .m2ts " ) else os . remove ( f " { f } .ffindex " )
2019-11-21 09:37:19 +00:00
except FileNotFoundError :
pass
2019-10-17 21:47:31 +00:00
2020-05-05 04:14:30 +00:00
print ( f " Output: { args . outfile } " ) if args . outfile and args . file else print ( f " Output: { os . path . splitext ( f ) [ 0 ] } _keyframes.txt " )
2019-10-17 21:47:31 +00:00
2019-07-14 13:12:16 +00:00
if __name__ == " __main__ " :
parser = argparse . ArgumentParser ( )
2019-08-16 07:24:41 +00:00
parser . add_argument ( " -F " , " --file " ,
2019-10-17 21:47:31 +00:00
help = " generate keyframes for a specific file " ,
action = " store " )
2020-05-17 02:27:11 +00:00
parser . add_argument ( " --scxvid " ,
action = " store_true " , default = False ,
help = " use scxvid for keyframe generating (default: %(default)s ) " )
2019-07-14 13:12:16 +00:00
parser . add_argument ( " -R " , " --recursive " ,
2020-05-05 04:14:30 +00:00
action = " store_true " , default = False ,
help = " search files recursively (default: %(default)s ) " )
2019-08-16 07:24:41 +00:00
parser . add_argument ( " -N " , " --noheader " ,
2020-05-05 04:14:30 +00:00
action = " store_true " , default = False ,
help = " do not include header line for aegisub (default: %(default)s ) " )
2019-08-16 07:24:41 +00:00
parser . add_argument ( " -O " , " --outfile " ,
2020-05-05 04:14:30 +00:00
action = " store " , default = None ,
help = " name for keyframes file output (Note: requires --file (-F) to be set) " )
2019-08-19 21:42:51 +00:00
parser . add_argument ( " -T " , " --trims " ,
2020-05-05 04:14:30 +00:00
action = " store " ,
2019-10-17 21:47:31 +00:00
help = " string of trims to source file. " \
2020-05-05 04:14:30 +00:00
" format: \" [inclusive,exclusive],[inclusive,exclusive],[None,exclusive],[inclusive,None] \" " )
2019-11-21 09:37:19 +00:00
parser . add_argument ( " -C " , " --check_exists " ,
2020-05-05 04:14:30 +00:00
action = " store_true " , default = False ,
help = " Check if keyframe file already exists (default: %(default)s ) " )
2020-05-17 02:27:11 +00:00
# TO-DO: Allow other formats than qpfile format
2019-07-14 13:12:16 +00:00
args = parser . parse_args ( )
2019-08-01 09:46:35 +00:00
main ( )
2019-11-21 09:37:19 +00:00
print ( f " \n Done generating keyframes. " )