Takes screenshots including subtitles for animebytes, automatically determining most suitable subtitle track. Also has a comparison feature but don't use that lol

Animebytes Screenshots.py

# recommended to specify default mpv path and output path below.
# Requirements:
#   - pip install pymediainfo colorama
#   - trackselect.lua (https://github.com/po5/trackselect) is required in mpv path unless subtitle track id is specified.
#   - oxipng (https://github.com/shssoichiro/oxipng) in path is required to use -opt.

mpv_path = "!CHANGE THIS!"
out_path = "!CHANGE THIS!"

import os
import subprocess
from argparse import ArgumentParser
import sys
import random
from pymediainfo import MediaInfo
from colorama import Fore, init
init(convert=True)

def isVideo(vidFile):
    if len(MediaInfo.parse(vidFile).video_tracks) == 0:
        return False
    else:
        return True


def main():
    parser = ArgumentParser(description="AB screenshotter")
    parser.add_argument('input', nargs='+', metavar="[file]", help="""Video file. Can use * to detect and process all video files in folder, ** to detect and process first video file in each subfolder, or *** for both.""")
    parser.add_argument('--screens', '-s', type=int, default=10, metavar="[int]", help="""Specify number of screenshots.""")
    parser.add_argument('--compare', '-c', nargs='*', type=int, default=-1, metavar="[int]", help="""Use comparison mode. Will generate the same frames for every file. Also takes frame numbers.""")
    parser.add_argument('--outDir', default=[out_path], metavar="[path]", help="""Directory screenshots should be output to.""")
    parser.add_argument('--sameFolder', action='store_true', help="""Create screenshot folder in the same directory as the input file.""")
    parser.add_argument('--subIndex', type=int, metavar="[int]", help="""Specify index of subtitle track to be used in input file. Use -1 for no subs.""")
    parser.add_argument('--pngComp', type=int, default=7, choices=range(1, 10), metavar="[int]", help="""Specify png compression level (1-9).""")
    parser.add_argument('--opt', action='store_true', help="""Optimize pngs.""")
    parser.add_argument('--pause', action='store_true', help="""Require input to take screenshots after each file.""")
    parser.add_argument('--mpvPath', default=[mpv_path], metavar="[path]", help="""Specify path of mpv executable.""")

    args = parser.parse_args()

    inFileLs = []
    specialBatch = False

    if len(args.input) == 1:
        #adds all video files in current directory to inFileLs
        if args.input[0] == "*" or args.input[0] == "***":
            specialBatch = True
            for file in os.listdir():
                if os.path.isfile(file):
                    if file.lower().endswith(".mkv") or isVideo(file):
                        inFileLs.append(file)

        #adds first video file from every subdirectory of current dir to inFileLs
        if args.input[0] == "**" or args.input[0] == "***":
            specialBatch = True
            for dir in sorted(os.listdir(os.getcwd())):
                dir = os.path.join(os.getcwd(), dir)
                if os.path.isdir(dir):
                    for file in sorted(os.listdir(dir)):
                        file = os.path.join(dir, file)
                        if os.path.isfile(file):
                            if file.lower().endswith(".mkv") or isVideo(file):
                                inFileLs.append(file)
                                break

        #otherwise add provided files to inFileLs
        if args.input[0] != "*" and args.input[0] != "**" and args.input[0] != "***":
            for file in args.input:
                if file.lower().endswith(".mkv") or isVideo(file):
                    inFileLs.append(file)

    elif len(args.input) > 1:
        for file in args.input:
            if file.lower().endswith(".mkv") or isVideo(file):
                inFileLs.append(file)

    #if no input files found, exit
    if len(inFileLs) == 0:
        print(Fore.RED + "No video files found." + Fore.WHITE)
        exit()

    mpvDir = os.path.dirname(args.mpvPath)
    outDirectory = args.outDir

    #checks if comparison mode is enabled
    comp = False
    randComp = False
    compFrames = []

    if type(args.compare) is not int:
        comp = True

        #if frames are specified, set number of screens to match and make sure they're all integers
        if len(args.compare) > 0:
            args.screens = len(args.compare)
            compFrames = args.compare
            for frame in compFrames:
                if type(frame) is not int:
                    print("\nComparison frame input must only contain integers.")
                    exit()

        #if not provided, frames will be randomly selected
        else:
            randComp = True

    fileNum = 0

    for inFile in inFileLs:

        if inFile != inFileLs[0] and args.pause:
            input()

        inFilename = os.path.basename(inFile)
        inDirectory = os.path.dirname(inFile)

        #if samefolder flag set, outdir is set to input's dir
        if args.sameFolder:
            if specialBatch:
                outDirectory = os.path.join(os.getcwd(), "screens")
            else:
                outDirectory = os.path.join(inDirectory, "screens")
        
        #create output folder
        if not os.path.exists(outDirectory):
            os.mkdir(outDirectory)

        #get length of video in frames
        length = round((MediaInfo.parse(inFile).general_tracks[0].duration / 1000) * float(MediaInfo.parse(inFile).video_tracks[0].frame_rate))
        
        #select random frame from video (making sure it hasn't been used before) and convert to timestamp
        usedFrames = []
        for i in range(0, args.screens):

            if (randComp and fileNum == 0) or not comp:
                newFrame = False
                while not newFrame:
                    randomFrame = random.randrange(0, length)
                    if randomFrame not in usedFrames:
                        newFrame = True
                usedFrames.append(randomFrame)
                if randComp:
                    compFrames.append(randomFrame)

            else:
                randomFrame = compFrames[i]


            randomSecond = randomFrame / float(MediaInfo.parse(inFile).video_tracks[0].frame_rate)

            #name output file
            outFile = os.path.join(outDirectory, inFilename + "_" + str(randomFrame) + ".png")

            #output screenshot progress
            if i == 0:
                print("\n" + Fore.CYAN + inFile)
            print("\n" + Fore.GREEN + str(i+1) + "/" + str(args.screens) + "\t" + Fore.WHITE + outFile)

            #build screenshot command
            mpvOptions = []
            
            if args.subIndex is None:
                mpvOptions.append("--script=" + os.path.join(mpvDir, "trackselect.lua"))
            elif args.subIndex == -1:
                mpvOptions.append("--no-sub")
            else:
                mpvOptions.append("--sid=" + str(args.subIndex))

            subprocess.call([args.mpvPath] + mpvOptions + ["--vo=image"] + ["--vo-image-format=png"] + ["--vf=format:rgb32"] + ["--no-audio"] + ["--untimed"] + ["--vo-image-png-compression=" + str(args.pngComp)] + ["--start=" + str(randomSecond)] + ["--frames=1"] + [inFile] + ["-o"] + [outFile], stdout=subprocess.DEVNULL)

            #oxipng optimisation command
            if args.opt:
                subprocess.call(["oxipng"] + [outFile] + ["-o"] + ["max"] + ["-p"])
            
        fileNum += 1

        print("\n")

if __name__ == "__main__":
    sys.exit(main())