trim-shorts

#!/usr/bin/env python3
#
# Copyright 2016 by Chris Osborn <fozztexx@fozztexx.com>
#
# This file is part of viddin.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License at <http://www.gnu.org/licenses/> for
# more details.

import os
import datetime
import tvdb_api
import sys
from operator import itemgetter
import re
import argparse
import glob
import math
import csv
import mimetypes
from viddin import viddin

def build_argparser():
  parser = argparse.ArgumentParser()
  parser.add_argument("series", help="Name of series on thetvdb.com")
  parser.add_argument("offsets", nargs="+", help="Trim amount from each short. Negative numbers are offset from end. Pattern repeats through files.")
  parser.add_argument("--dvdmissing", action="store_true", help="use episode data when dvd info is missing")
  parser.add_argument("--dvdignore", action="store_true", help="ignore dvd episode data")
  parser.add_argument("--enumerate", action="store_true", help="Episodes are in order but may have gaps")
  parser.add_argument("--transcode", action="store_true", help="Transcode instead of cut")
  parser.add_argument("--transcode-flags", help="extra flags to pass to transcoder")
  parser.add_argument("--first", default=1, help="Start enumeration at this number")
  parser.add_argument("--minsize", help="Ignore videos smaller than this")
  parser.add_argument("--iframes", action="store_true", help="Split at nearest iframe")
  parser.add_argument("--skipexisting", action="store_true", help="Don't redo existing trims")
  parser.add_argument("--trimend", help="Frames to remove from end")
  parser.add_argument("--debug", action="store_true", help="Turn on debug output")
  return parser

def bestMatch(offset, filename, debugFlag=False):
  video, ext = os.path.splitext(filename)
  if debugFlag:
    print("Finding offset", offset)
  if offset < 0:
    offset = viddin.getLength(filename) + offset
    
  best = None
  with open(video + ".blk") as f:
    for line in f:
      info = line.split()
      begin = float(info[1])
      end = float(info[2])
      seconds = begin + (end - begin) / 2
      diff = abs(seconds - offset)
      if not best or diff < best[0]:
        best = [diff, seconds, begin, end]
        
  bestsil = None
  with open(video + ".sil") as f:
    for line in f:
      info = line.split()
      begin = float(info[1])
      end = float(info[2])
      seconds = begin + (end - begin) / 2
      diff = abs(seconds - best[1])
      if ((begin >= best[2] and begin <= best[3]) or (end >= best[2] and end <= best[3]) or \
            (best[2] >= begin and best[2] <= end) or (best[3] >= begin and best[3] <= end)) \
            and (not bestsil or diff < bestsil[0]):
        bestsil = [diff, seconds, begin, end]

  if bestsil:
    if debugFlag:
      print("Silence: ", bestsil[2], bestsil[3], best[2], best[3])
    begin = bestsil[2]
    if begin < best[2]:
      begin = best[2]
    end = bestsil[3]
    if end > best[3]:
      end = best[3]
    split = begin + (end - begin) / 2
    if debugFlag:
      print("Splitting at ", split)
    return split

  if args.debug:
    print("Black: ", best)
  return best[1]

def main():
  args = build_argparser().parse_args()
  # FIXME - if args.series is file then trim just that one file

  if os.path.isfile(args.series):
    mime = mimetypes.guess_type(args.series)
    ftype = mime[0].split('/')[0]
    if ftype != "video":
      series = viddin.loadEpisodeInfoCSV(args.series)
  else:
    series = viddin.loadEpisodeInfoTVDB(args.series, dvdIgnore=args.dvdignore,
                                        dvdMissing=args.dvdmissing)

  dvdorder = viddin.EpisodeList(series, "dvdSeason", "dvdEpisode")

  if len(series):
    videos = [file for file in glob.glob("[0-9]x[0-9][0-9].mkv")]
    videos.sort()
  else:
    videos = [args.series]

  minsize = 0
  if args.minsize:
    val = re.sub("^( *[0-9]+).*", "\\1", args.minsize)
    minsize = int(val)
    suffix = args.minsize[-1]
    try:
      magnitude = ["B", "K", "M", "G"].index(suffix.upper())
    except ValueError:
      magnitude = -1
    if magnitude >= 0:
      minsize *= math.pow(1024, magnitude)

  if viddin.isint(args.first):
    first = int(args.first)
  else:
    episode = dvdorder.findEpisode(args.first)
    first = dvdorder.indexOf(episode) + 1
    print(episode)

  # FIXME - do merging of multi chapter segments
  epnum = 0
  for i in range(0, len(videos)):
    oindex = epnum % len(args.offsets)
    filename = videos[i]
    if minsize:
      size = os.path.getsize(filename)
      if size < minsize:
        continue

    if args.debug:
      print(filename)
    video, ext = os.path.splitext(filename)
    epid = video
    if args.enumerate:
      epid = epnum + first - 1
    episode = dvdorder.findEpisode(epid)
    if args.debug:
      print(episode)

    if episode:
      eptitle = "%s %s%s" % (dvdorder.formatEpisodeID(episode), episode.title, ext)
    else:
      eptitle = video + "_1" + ext

    if not args.skipexisting or not os.path.exists(eptitle):
      olen = viddin.decodeTimecode(args.offsets[oindex])

      offset = None
      if olen:
        viddin.findBlack(filename)
        viddin.findSilence(filename)
        split = bestMatch(olen, filename)
        offset = "+"
        if olen < 0:
          offset = "0-"
        offset += str(split)

      output = video + "_1" + ext

      if offset is not None or args.transcode:
        if offset is not None:
          cmd = ["split-video"]
          if args.iframes:
            cmd.append("--iframes")
          if args.transcode:
            cmd.append("--transcode")
          if args.transcode_flags:
            cmd.append('--transcode-flags="' + args.transcode_flags + '"')
          if args.trimend:
            cmd.extend(["--trimend", args.trimend])
        else:
          cmd = ["rip-video"]
          if args.transcode_flags:
            cmd.extend(args.transcode_flags.split())
        cmd.append(filename)
        if offset is not None:
          cmd.append(offset)
        else:
          cmd.append(output)

        viddin.runCommand(cmd, args.debug)
        os.rename(output, eptitle)
        os.remove(video + "_1.log")
      else:
        if os.path.exists(eptitle):
          os.remove(eptitle)
        os.link(filename, eptitle)
      print(eptitle)
    else:
      if args.debug:
        print("Skipped " + eptitle)

    epnum += 1

  return

if __name__ == '__main__':
  exit(main() or 0)