Generates movie thumbnail images automatically

autothumbs.py

#!/usr/bin/env python3

#
# Automatic Movie Thumbnail Generator
#
# Scans a folder, and generates a thumbnail image for every movie found.
# Uses a grid layout that attempts to maximize utilized space while
# fitting under a given maximum image size, and having at least as many
# frames as requested.
#

#
# This requires MoviePy and Pillow:
#   pip install MoviePy
#   pip install Pillow
#
# The default setup uses the Comic Mono font:
#   https://dtinth.github.io/comic-mono-font/
#   (TTF file can be in same folder or installed on system.)
#
# The variables below can be overridden from the command line.
#
#
# Examples:
#   Process all files in current directory:
#     python autothumbs.py
#   Process single file:
#     python autothumbs.py -I file.mp4
#   Remove the timestamps and header:
#     python autothumbs.py -VPAD 1 -HEADPAD 1 -NOHEADER -NOTIMESTAMP
#   Custom header text:
#     python autothumbs.py -HEAD "File: &n" "Resolution: &x x &y"
#   Timestamp in top right instead of in VPAD area:
#     python autothumbs.py -VPAD 1 -TSCORNER 1 -TSX -3 -TSY 1 -TSOUTLINE 1
#

W = 1024
H = 1024
DIR_IN = "."
DIR_OUT = "Thumbs"
FORMAT = ".jpg"
QUALITY = 75
EXT = (".mp4",".mkv",".avi",".flv",".mpeg",".mov",".m4v")

WIPEDIR = False # delete contents of DIR_OUT before starting
SKIPDONE = False # don't recreate a thumbnail that already exists
RECURSE = False # recursive search of DIR_IN

# for testing
DEBUG = False # print debug info
TESTCOUNT = None # number of thumbnails to make before stopping
TESTSKIP = 0 # Skip this many thumbnails before starting

FPS_SOURCE = "fps" # this seems to give more predictable results than the default "tbr"
BADFPS = 1000 # over this give a warning (is probably gonna do a bad)

# other parameters for command line arguments
I=None # single input filename
O=None # single output filename
BG=(0,0,0) # background colour
COUNT=10 # minimum thumbnail count
INTSCALE=False # integer thumb scaling only
NOTIMESTAMP=False # hide timestamps (may want VPAD=1 if False)
NOHEADER=False # hide info header
HPAD=1 # horizontal padding
VPAD=18 # vertical padding, space for timestamps
HEADPAD=34 # header padding
TSFONT="ComicMono.ttf" # timestamp font
HEADFONT="ComicMono.ttf" # header font
TSX=1 # timestamp location in VPAD gutter
TSY=1
TSS=16 # timestamp font point size
TSC=(255,255,0) # timestamp colour
TSCORNER=2 # bottom left corner
TSOUTLINE=0 # timestamp outline
HEADX=1 # header text location
HEADY=1
HEADS=16 # header font point size
HEADC=(255,255,255) # header text colour
HEADH=16 # vertical distance between header text rows
HEAD=["&n","&t | &f FPS | &x x &y | &z MB"] # list of rows to generate for the header (see header_info below)
S=[] # pick specific frames (seconds)
F=[] # pick specific frames (frame number)

#
# Code
#

import sys
import argparse
import os
import math
import enum
import moviepy.editor
import PIL.Image
import PIL.ImageFont
import PIL.ImageDraw

# choose a timestamp format that fits the clip
TIMESTAMP_MID = 0 # mm:ss
TIMESTAMP_LONG = 1 # h:mm:ss
TIMESTAMP_SHORT = 2 # ss.ttt
def timestamp_mode(duration):
    if duration < 60: return TIMESTAMP_SHORT
    if duration >= (60*60): return TIMESTAMP_LONG
    return TIMESTAMP_MID
        
def timestamp(seconds,mode=TIMESTAMP_MID,roundup=False):
    if mode == TIMESTAMP_SHORT:
        (ms,s) = math.modf(seconds)
        ms *= 1000
        ms = math.ceil(ms) if roundup else math.floor(ms)
        return "%02d.%03d" % (s,ms)
    seconds = math.ceil(seconds) if roundup else math.floor(seconds)
    s = seconds % 60
    m = seconds // 60
    if mode == TIMESTAMP_LONG:
        h = m // 60
        m %= 60
        return "%d:%02d:%02d" % (h,m,s)
    # TIMESTAMP_MID
    return "%02d:%02d" % (m,s)      

def timestamp_to_seconds(s):
    ts = s.split(":")
    if len(ts) > 3:
        raise Exception("Timestamp has too many : separators: "+s)
    v = 0
    for t in s.split(":"):
        v *= 60
        try:
            v += float(t)
        except Exception:
            raise Exception("Timestamp invalid number: "+t+" ("+s+")")
    return v

# header info string, can replace:
# %n - filename
# %t - duration (auto type)
# %h - duration long
# %m - duration mid
# %s - duration small
# %f - framerate
# %x - width
# %y - height
# %z - file size (MB)
def headline(clip,s):
    if DEBUG: os = s
    sd = [
        ("&&","&"),
        ("&n",clip._headline_filename_),
        ("&t",timestamp(clip.duration,clip._headline_tsmode_,True)),
        ("&h",timestamp(clip.duration,TIMESTAMP_LONG,True)),
        ("&m",timestamp(clip.duration,TIMESTAMP_MID,True)),
        ("&s",timestamp(clip.duration,TIMESTAMP_SHORT,True)),
        ("&f","%5.2f" % clip.fps),
        ("&x","%d" % clip.w),
        ("&y","%d" % clip.h),
        ("&z","%.1f" % (clip._headline_size_ / (1024*1024))),
        ] # would like codec info but MoviePy doesn't seem to have it
    for (src,dst) in sd:
        s = s.replace(src,dst)
    if DEBUG: print("headline: '%s' -> '%s'" % (os,s))
    return s

def thumb(path,filename=None,
          pw=1024,ph=1024,bg=(0,0,0),
          count=10,intscale=False,
          hpad=1,
          vpad=18,tsx=1,tsy=1,tsfont="ComicMono.ttf",tss=16,tsc=(255,255,0),tscorner=2,tsoutline=0,
          headpad=34,headx=1,heady=1,headh=16,headfont="ComicMono.ttf",heads=16,headc=(255,255,255),head=[r"&n",r"&t | &f FPS | &x x &y | &z MB"],
          picks=[],pickf=[]
          ):
    # pw,ph = maximum image size (finds a grid that fills this space best)
    # bg = background colour
    # count = minimum number of thumbnails
    # intscale = only scale thumbnails down by an integer factor
    # hpad = pixel gutter between thumbs horizontally
    # vpad = pixel gutter between thumbs vertically
    # tsx,tsy = timestamp location in vpad gutter
    # tsfont,tss,tsc = timestamp font, points size, colour
    # tscorner = 0,1,2,3 = tl,tr,bl,br (right side is right-aligned)
    # tsoutline = black outline for timestamp (pixels)
    # headpad = pixel gutter header at top
    # headx,heady = header text location
    # headh = vertical distance between header text rows
    # headfont,heads,headc = header font, points size, colour
    # head = list of rows to generate for the header (see header_info above)
    # picks = list of specific frames to force (seconds)
    # pickf = list of specific frames to force (frame number)
    if DEBUG: print("thumb: " + path)
    clip = moviepy.editor.VideoFileClip(path,audio=False,fps_source=FPS_SOURCE)
    if (clip.fps <= 0 or clip.fps >= BADFPS):
        print("BAD FPS: %f" % clip.fps) # might help
    if filename == None: filename = path
    cw = clip.w
    ch = clip.h
    frames = int(clip.duration * clip.fps)
    if DEBUG: print("Resolution: %d x %d (%d frames @ %f FPS)" % (cw,ch,frames,clip.fps))
    # prepare fonts for timestamp and header
    tsmode = timestamp_mode(clip.duration)
    if tsfont: tsfont = PIL.ImageFont.truetype(tsfont,tss)
    if headfont: headfont = PIL.ImageFont.truetype(headfont,heads)
    # store information for headline
    clip._headline_filename_ = os.path.basename(filename)
    clip._headline_tsmode_ = tsmode
    clip._headline_size_ = os.path.getsize(path)
    # find best layout that fits >= count images into the desired width
    gw = 1
    gh = 1
    tw = cw
    th = ch
    px = -1
    for gwi in range(1,count+1): # try each grid width
        ghi = (count + (gwi-1)) // gwi # grid height to match width
        # padding and usable dimensions
        hp = hpad * (1+gwi)
        vp = headpad + (vpad * ghi)
        uw = pw - hp
        uh = ph - vp
        # find biggest thumbnail scale that fits (both axes)
        sw = cw / min(int(uw/gwi),cw)
        sh = ch / min(int(uh/ghi),ch)
        if sw <= 0: sw = cw # if nothing can fit, assume 1px scale
        if sh <= 0: sh = ch
        if intscale:
            sw = math.ceil(sw)
            sh = math.ceil(sh)
        s = max(sw,sh)
        #if DEBUG: print("S: %f (%f,%f)" % (s,sw,sh))
        # calculate resulting frame size
        twi = int(cw/s)
        thi = int(ch/s)
        # calculate resulting pixel coverage, keep if most
        pxi = (twi*thi)*(gwi*ghi)
        #if DEBUG: print("GI: %d x %d (%d x %d)" % (gwi,ghi,twi,thi))
        if (pxi > px):
            px = pxi
            tw = twi
            th = thi
            gw = gwi
            gh = ghi
    # expand grid vertically if underused
    while (headpad+((vpad+th)*(gh+1))) <= ph: gh += 1
    if DEBUG: print("Grid: %d x %d (%d x %d)" % (gw,gh,tw,th))
    # generate images
    tc = gw * gh
    iw = (tw * gw) + (hpad * (1+gw))
    ih = (th * gh) + headpad + (vpad * gh)
    img = PIL.Image.new("RGB",(iw,ih),color=bg)
    draw = PIL.ImageDraw.Draw(img)
    # chosen frames
    chosen = [int(frames*(x+0.5)/tc) for x in range(tc)] # evenly spaced
    if (len(pickf) + len(picks)) > 0:
        pick = pickf + [int(x * clip.fps) for x in picks] # convert seconds to frames
        pick = [max(0,min(frames-1,x)) for x in pick] # clamp to valid frames
        pick = list(dict.fromkeys(pick)) # eliminate duplicates
        chosen = pick[:]
        while len(chosen) < tc: # choose rest by subdivision of largest remaining space
            chosen = sorted(chosen)
            spans = [frames]
            if len(chosen) > 0:
                spans = [chosen[0]] + [(chosen[i+1]-chosen[i]) for i in range(0,len(chosen)-1)] + [frames-chosen[-1]]
            sr = max(spans)
            si = spans.index(sr) # largest space
            ss = 0
            if si > 0: ss = chosen[si-1]
            chosen.append(int(ss+(sr/2))) # add to middle
        for relax in range(10): # relaxation
            chosen = sorted(chosen)
            for i in range(len(chosen)):
                if chosen[i] in pick: continue # don't move picks
                vl = 0 if (i == 0) else chosen[i-1]
                vr = frames if (i == (len(chosen)-1)) else chosen[i+1]
                chosen[i] = int((vl+vr)/2)
        chosen = sorted(chosen)
        if DEBUG: print(chosen)
    # extract and assemble frames
    for tyi in range(0,gh):
        for txi in range(0,gw):
            ti = (tyi * gw) + txi
            tx = hpad + (txi * (tw + hpad))
            ty = headpad + (tyi * (th + vpad))
            frame = chosen[ti]
            frame_time = (frame) / clip.fps
            if DEBUG: print("frame: %d = %s (%f)" % (frame,timestamp(frame_time,tsmode),frame_time))
            timg = PIL.Image.fromarray(clip.get_frame(frame_time))
            timg = timg.resize((tw,th),resample=PIL.Image.Resampling.LANCZOS)
            img.paste(timg,(tx,ty))
            # timestamp
            if tsfont:
                tstext = timestamp(frame_time,tsmode)
                # bottom-left default
                tscx = tx
                tscy = ty+th
                if tscorner < 2: tscy = ty # top corners
                if (tscorner & 1): tscx = tx+tw-math.ceil(draw.textlength(tstext,tsfont)) # right corners
                tscx += tsx
                tscy += tsy
                if tsoutline > 0:
                    for tsoy in range(-tsoutline,tsoutline+1):
                        for tsox in range(-tsoutline,tsoutline+1):
                            draw.text((tscx+tsox,tscy+tsoy),tstext,fill=(0,0,0),font=tsfont)
                draw.text((tscx,tscy),tstext,fill=tsc,font=tsfont)
    # header
    if headfont:
        ty = heady
        for mode in head:
            s = headline(clip,mode)
            draw.text((headx,ty),s,fill=headc,font=headfont)
            ty += headh
    # done
    clip.close()
    return img

#
# MAIN
#

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ag = ap.add_argument_group("Input/Output")
    ag.add_argument("DIR_IN",nargs="?",help="Movie source directory (optional)")
    ag.add_argument("DIR_OUT",nargs="?",help="Thumbnail output directory (optional)")
    ag.add_argument("-I",help="Single input movie file instead of DIR_IN scan",metavar="filename")
    ag.add_argument("-O",help="Single output image instead of same folder, use with -I")
    ag = ap.add_argument_group("General")
    ag.add_argument("-W",type=int,help="Maximum thumbnail width",metavar="px")
    ag.add_argument("-H",type=int,help="Maximum thumbnail height",metavar="px")
    ag.add_argument("-COUNT",type=int,help="Minimum number of frames in thumbnail",metavar="frames")
    ag.add_argument("-S",action="append",help="Force specific frame by timestamp",metavar="h:mm:ss.ttt")
    ag.add_argument("-F",type=int,action="append",help="Force specific frame by number",metavar="frame")
    ag.add_argument("-BG",type=int,nargs=3,help="Background colour",metavar=("r","g","b"))
    ag.add_argument("-FORMAT",help="Image format extension: .jpg / .png",metavar="extension")
    ag.add_argument("-QUALITY",type=int,help="Image quality 0-100",metavar="percent")
    ag.add_argument("-HPAD",type=int,help="Horizontal pixel padding",metavar="px")
    ag.add_argument("-VPAD",type=int,help="Vertical pixel padding",metavar="px")
    ag.add_argument("-INTSCALE",action="store_const",const=True,help="Only allow frames to be shrunk by a whole-number factor")
    ag = ap.add_argument_group("Timestamps")
    ag.add_argument("-NOTIMESTAMP",action="store_const",const=True,help="Hide timestamp for each frame (see also VPAD)")
    ag.add_argument("-TSFONT",help="Timestamp font file (TTF)",metavar="filename")
    ag.add_argument("-TSS",type=int,help="Timestamp font size",metavar="points")
    ag.add_argument("-TSC",type=int,nargs=3,help="Timestamp font colour",metavar=("r","g","b"))
    ag.add_argument("-TSX",type=int,help="Timestamp X position relative to corner",metavar="px")
    ag.add_argument("-TSY",type=int,help="Timestamp Y position relative to corner",metavar="px")
    ag.add_argument("-TSCORNER",type=int,help="Timestamp corner: 0,1,2,3 = top left, top right (right-align), bottom left (default), bottom right",metavar="corner")
    ag.add_argument("-TSOUTLINE",type=int,help="Timestamp black outline",metavar="px")
    ag = ap.add_argument_group("Header")
    ag.add_argument("-NOHEADER",action="store_const",const=True,help="Hide header text at top (see also HEADPAD)")
    ag.add_argument("-HEADPAD",type=int,help="Header pixel padding",metavar="px")
    ag.add_argument("-HEADFONT",help="Header font file (TTF)",metavar="filename")
    ag.add_argument("-HEADS",type=int,help="Header font size",metavar="points")
    ag.add_argument("-HEADC",type=int,nargs=3,help="Header font colour",metavar=("r","g","b"))
    ag.add_argument("-HEADX",type=int,help="Header X position in HEADPAD",metavar="px")
    ag.add_argument("-HEADY",type=int,help="Header Y position in HEADPAD",metavar="px")
    ag.add_argument("-HEADH",type=int,help="Vertical spacing between header text rows",metavar="px")
    ag.add_argument("-HEAD",nargs="+",help="Header row text, substitutions: &n filename, &t clip time (&h/&m/&s long/mid/short), &f FPS, &x &y resolution, &z size (MB), && for &",metavar="row")
    ag = ap.add_argument_group("Directory Management")
    ag.add_argument("-WIPEDIR",action="store_const",const=True,help="Delete DIR_OUT contents before starting")
    ag.add_argument("-SKIPDONE",action="store_const",const=True,help="Skip if thumbnail already exists in DIR_OUT")
    ag.add_argument("-RECURSE",action="store_const",const=True,help="Search subdirectories of DIR_IN")
    ag = ap.add_argument_group("Debug")
    ag.add_argument("-FPS_SOURCE",help="FFMPEG FPS source: should be 'fps' or 'tbr'",metavar="source")
    ag.add_argument("-BADFPS",type=int,help="FPS threshold for BAD FPS warning",metavar="fps")
    ag.add_argument("-DEBUG",action="store_const",const=True,help="Print debug information")
    ag.add_argument("-TESTCOUNT",type=int,help="Quit after this many thumbnails (for testing)",metavar="count")
    ag.add_argument("-TESTSKIP",type=int,help="Skip this many thumbnails (for testing)",metavar="start")
    args = ap.parse_args()
    if DEBUG or args.DEBUG: print("Arguments:")
    for (k,v) in args.__dict__.items():
        if v != None:
            if DEBUG or args.DEBUG: print("  "+k+": "+str(v))
            globals()[k] = v

# command line argument conversion
BG = tuple(BG)
TSC = tuple(TSC)
HEADC = tuple(HEADC)
if NOTIMESTAMP: TSFONT = None
if NOHEADER: HEADFONT = None
S = [timestamp_to_seconds(x) for x in S]

print("Input: " + (DIR_IN if (I==None) else I))
print("Output: " + DIR_OUT)

# make sure the output folder exists
os.makedirs(DIR_OUT,exist_ok=True)

# clean output folder
if WIPEDIR:
    for (root,dirs,files) in os.walk(DIR_OUT):
        for f in files:
            path = os.path.join(root,f)
            print("DELETE: "+path)
            os.remove(path)

# build all videos
count = 0
for (root,dirs,files) in os.walk(DIR_IN):
    for f in files:
        if TESTCOUNT and count >= TESTCOUNT: break
        if f.lower().endswith(EXT) or I != None:
            count += 1
            if I != None: # single file override
                f = I
                root = ""
            #of = os.path.join(DIR_OUT,os.path.splitext(f)[0]+FORMAT)
            of = os.path.join(DIR_OUT,f+FORMAT) # keep extension
            if O != None:
                of = O
            if count < (TESTSKIP+1): continue
            if SKIPDONE and os.path.exists(of):
                print("SKIP: "+f)
                continue
            print("%3d %s" % (count-1,f))
            img = thumb(os.path.join(root,f),f,W,H,
                        bg=BG,
                        count=COUNT,
                        intscale=INTSCALE,
                        hpad=HPAD,vpad=VPAD,headpad=HEADPAD,
                        tsfont=TSFONT,headfont=HEADFONT,
                        tsx=TSX,tsy=TSY,tss=TSS,tsc=TSC,
                        tscorner=TSCORNER,tsoutline=TSOUTLINE,
                        headx=HEADX,heady=HEADY,heads=HEADS,headc=HEADC,headh=HEADH,
                        head=HEAD,
                        picks=S,pickf=F
                        )
            if img:
                img.save(of,quality=QUALITY)
            else:
                print("ERROR: No thumbnail generated!")
        if I != None: break # single file
    if not RECURSE or I != None: break # don't recurse
print("%d thumbnails" % count)