find-intro

#!/usr/bin/env python3
# -*- python -*-
#
# Copyright 2018 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.

# FIXME - make this search for theme song:
# https://dev.to/hiisi13/find-an-audio-within-another-audio-in-10-lines-of-python-1866
# https://github.com/craigfrancis/audio-detect

import os
import argparse
import datetime
import sys
import re
import viddin

DISTANCE = 2
CHAPTER_INTRO = "Intro"
CHAPTER_SKIP = "Skip Intro"

def build_argparser():
  parser = argparse.ArgumentParser()
  parser.add_argument("files", nargs="+", help="Video file to guess")
  parser.add_argument("--audio", help="File with theme song to search for")
  parser.add_argument("--pattern", nargs="+", type=float, help="Pattern to look for")
  parser.add_argument("--add", action="store_true", help="add chapter if pattern found")
  parser.add_argument("--chapter-only", action="store_true",
                      help="one end of audio must align with existing chapter")
  parser.add_argument("--distance", default=10,
                      help="Maximum number of seconds away from specified offset")
  parser.add_argument("--margin", default=DISTANCE, type=float, help="allowed margin away")
  return parser

def matchesPattern(pos, splits, pattern):
  testpattern = []
  for idx in range(len(pattern), 0, -1):
    testpattern.append(abs(splits[pos]) - abs(splits[pos - idx]))
  for idx, (a, b) in enumerate(zip(testpattern, pattern)):
    if abs(a - b) > DISTANCE:
      return False
  return True

def findCredits(episode, pattern, addFlag):
  splits = episode.loadSplits()

  for idx in range(len(pattern), len(splits)):
    if matchesPattern(idx, splits, pattern):
      pos = abs(splits[idx])
      doEdit = episode.addChapters([viddin.Chapter(pos, CHAPTER_SKIP)])
      print("Setting opening for %s to %0.3f" % (filename, pos))
      if doEdit and addFlag:
        episode.writeChapters()
      return

  print("Couldn't find opening for", filename)
  return

def bestSplit(position, splits, margin=DISTANCE):
  bdiff = -1
  boffset = None
  for row in splits:
    diff = abs(position - abs(row))
    if not boffset or diff < bdiff:
      bdiff = diff
      boffset = abs(row)
  if boffset and bdiff < margin:
    return boffset, bdiff
  return None

def findAudio(episode, theme, addFlag, chapterFlag, margin):
  intro_idx, _ = episode.chapterWithID(CHAPTER_INTRO)
  skip_idx, _ = episode.chapterWithID(CHAPTER_SKIP)
  if intro_idx is not None and skip_idx is not None:
    return

  splits = episode.loadSplits()

  base, ext = os.path.splitext(episode.path)
  aud = base + ".aud"

  # FIXME - need to keep track of which audio file was used to create .aud, not just timestamp
  if os.path.exists(aud) and os.path.getmtime(aud) < os.path.getmtime(theme):
    os.unlink(aud)

  if not os.path.exists(aud):
    cmd = ["find-audio", theme, episode.path]
    viddin.runCommand(cmd)
  audio = viddin.loadSplits(base + ".aud")
  silence = viddin.loadSplits(base + ".sil")

  chapters = episode.chapters

  _, begin, end = audio[0]
  audio_chaps = [None, None]
  was_chaps = [None, None]

  offset = None
  for chap in chapters:
    b_diff = abs(chap.position - begin)
    e_diff = abs(chap.position - end)
    if e_diff < margin:
      offset = begin - end
      was_chaps[1] = chap
      break
    elif b_diff < margin:
      offset = end - begin
      was_chaps[0] = chap
      break

  if offset:
    audio_chaps[0] = chap.position
    pos = chap.position + offset
    best = bestSplit(pos, splits, margin)
    if not best:
      best = bestSplit(pos, [x[0] for x in silence], margin)
    if best:
      audio_chaps[1] = best[0]
    else:
      audio_chaps[1] = pos
  elif chapterFlag:
    print("Unable to find audio", episode.path)
    return
  else:
    # No matching chapter
    best = bestSplit(end, splits)
    if not best:
      best = bestSplit(end, [x[0] for x in silence])
    if best:
      audio_chaps[1] = best[0]
    else:
      audio_chaps[1] = end

    best = bestSplit(begin, splits)
    if not best:
      best = bestSplit(begin, [x[0] for x in silence])
    if best:
      audio_chaps[0] = best[0]
    else:
      audio_chaps[0] = audio_chaps[1] - (end - begin)

  audio_chaps.sort()
  chaps = [viddin.Chapter(audio_chaps[0], CHAPTER_INTRO),
           viddin.Chapter(audio_chaps[1], CHAPTER_SKIP)]
  identical = []
  for idx in range(len(chaps)):
    c = chaps[idx]
    for vc in chapters:
      if c.name == vc.name and abs(c.position - vc.position) < margin:
        identical.append(idx)
        break

  for idx in reversed(identical):
    del chaps[idx]
    del was_chaps[idx]

  if chaps:
    doEdit = episode.addChapters(chaps)
    if doEdit:
      for c, w in zip(chaps, was_chaps):
        print("  %s: %0.3f" % (c.name, c.position), end="")
        if w:
          print(" (was %s)" % (w.name), end="")
        print()

    if doEdit and addFlag:
      episode.writeChapters()

  return

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

  if args.audio:
    for path in args.files:
      if not os.path.exists(path):
        print("No such file", path)
        exit(1)
      print(path)
      episode = viddin.Media(path)
      findAudio(episode, args.audio, args.add, args.chapter_only, args.margin)
  else:
    pattern = []
    for idx in range(len(args.pattern)-1):
      pattern.append(abs(args.pattern[-1]) - abs(args.pattern[idx]))

    for path in args.files:
      if not os.path.exists(path):
        print("No such file", path)
        exit(1)
      episode = viddin.Media(path)
      findCredits(episode, pattern, args.add)
  return

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