guess-chapters

#!/usr/bin/env python3
# -*- python -*-
#
# Copyright 2017 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 argparse
import datetime
import tempfile
import sys
import re
import viddin

def build_argparser():
  parser = argparse.ArgumentParser()
  parser.add_argument("file", help="Video file to guess")
  parser.add_argument("offsets", nargs="*", help="Approximate location of expected chapters")
  parser.add_argument("-f", "--force", action="store_true", help="Force a chapter at locations even if there's no matching silence")
  parser.add_argument("-a", "--add", action="store_true", help="Add chapters to existing chapters")
  parser.add_argument("--debug", action="store_true", help="Turn on debug output")
  parser.add_argument("--distance", default=10, help="Maximum number of seconds away from specified offset")
  parser.add_argument("-c", "--cuts", action="store_true",
                      help="Use scene cuts if there's no black")
  return parser

def closestExpected(best):
  global expected

  begin = float(best[1])
  end = float(best[2])
  center = begin + (end - begin) / 2
  closest = None
  for offset in expected:
    diff = abs(offset - center)
    if not closest or diff < closest[0]:
      closest = [diff, center, begin, end]

  if closest:
    return [closest[2], closest[3]]
  return None

def findNearest(position, chapters):
  for chap in chapters:
    pos = chap
    if isinstance(chap, viddin.Chapter):
      pos = chap.position
    if pos > position + 2:
      break
    diff = abs(pos - position)
    if diff < 1:
      return chap

  return None

def main():
  args = build_argparser().parse_args()

  vfile = viddin.Media(args.file, None)
  doEdit = False

  viddin.findBlack(args.file)
  viddin.findSilence(args.file)
  if args.cuts:
    viddin.findCuts(args.file)

  expected = None
  if len(args.offsets):
    expected = []
    for offset in args.offsets:
      cidx = offset.find("=")
      cname = None
      if cidx > 0:
        cname = offset[cidx+1:]
        offset = offset[:cidx]
      expected.append(viddin.Chapter(viddin.decodeTimecode(offset), cname))

    changed = []
    found = []
    for pos in expected:
      idx, chap = vfile.chapterWithID(pos)
      if idx is not None:
        if pos.name is not None and pos.name != chap.name:
          changed.append(viddin.Chapter(chap.position, pos.name))
        found.append(pos)
    if len(changed):
      vfile.addChapters(changed)
      doEdit = True
    if len(found):
      expected = [x for x in expected if x not in found]

  if expected is None or len(expected):
    video, ext = os.path.splitext(args.file)
    black = viddin.loadSplits(video + ".blk")
    silence = viddin.loadSplits(video + ".sil")

    splits = []
    for row in black:
      match = viddin.bestSilence(row, silence)
      if match:
        splits.append((match[1] - match[0]) / 2 + match[0])
      else:
        splits.append(0 - ((row[1] - row[0]) / 2 + row[0]))

    if expected:
      changed = []
      for offset in expected:
        bdiff = -1
        boffset = None
        pos = offset.position
        if pos < 0:
          pos = vfile.length + pos
        for row in splits:
          if row >= 0 or args.force:
            diff = abs(pos - abs(row))
            if not boffset or diff < bdiff:
              bdiff = diff
              boffset = abs(row)
        idx, echap = vfile.chapterWithID(boffset)
        if boffset and bdiff < args.distance \
           and (idx is None or echap.name != offset.name):
          changed.append(viddin.Chapter(boffset, offset.name))
      if len(changed):
        vfile.addChapters(changed)
        doEdit = True

  if doEdit:
    vfile.writeChapters()

  if expected is None:
    chapters = vfile.chapters
    for chap in chapters:
      if findNearest(chap.position, splits) is None:
        splits.append(chap.position)
    splits = sorted(splits, key=abs)
    
    for row in splits:
      if row >= 0 or args.force:
        print("%0.3f\t%s" % (row, viddin.formatTimecode(abs(row))), end="")
        chap = findNearest(row, chapters)
        if chap is not None:
          print("\t%s\t%s" % (viddin.formatTimecode(chap.position), chap.name), end="")
        print()
        
  return

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