split-video

#!/usr/bin/env python3
# -*- python -*-
#
# 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 argparse
import os
import datetime
import re
import math
import sys
from viddin import viddin

def build_argparser():
  parser = argparse.ArgumentParser()
  parser.add_argument("filename", help="Video to split")
  parser.add_argument("ranges", nargs="+", help="Split at offset(s) or ranges. Use negative number for offset from end. If argument is a filename, load offsets from it.")
  parser.add_argument("--transcode", action="store_true", help="transcode video for perfect split points")
  parser.add_argument("--transcode-flags", help="extra flags to pass to transcoder")
  parser.add_argument("--iframes", action="store_true", help="compare split positions to iframes from video")
  parser.add_argument("--output", help="name to use when creating splits")
  parser.add_argument("--outtype", help="force split to be this type")
  parser.add_argument("--quality", help="use this quality level when transcoding. Higher numbers are worse quality, default is 20")
  parser.add_argument("--trimend", help="Frames to remove from end")
  parser.add_argument("--chapters", action="store_true", help="Split a video by chapters")
  parser.add_argument("--debug", action="store_true", help="Turn on debug output")
  return parser

def findIframes(filename, output):
  if not os.path.exists(output) or os.path.getctime(output) < os.path.getctime(filename):
    if args.debug:
      print("Finding I frames " + filename)
    os.system("find-iframes \"%s\" \"%s\" > /dev/null 2>&1" % (filename, output))

def findNearest(pos, before):
  prev = 0
  print(pos, before)
  for frame in iframes:
    if pos < frame:
      if before:
        return prev
      else:
        return frame
    prev = frame
    
def splitVideo(segrange, input, output=None, trim=None, quality=None,
               transcodeFlag=False, transcodeOptions=None, iframeFlag=False):
  video, ext = os.path.splitext(input)

  vlen = viddin.getLength(input)
  if not segrange[1]:
    segrange[1] = vlen

  if not output:
    output = "%sc%s" % (video, ext)

  if trim:
    cmd = "vidinf \"%s\" | grep ID_VIDEO_FPS | sed -e 's/ID_VIDEO_FPS=//'" % input
    process = os.popen(cmd)
    fps = float(process.read())
    process.close()
    trim = float(trim) * 1.0 / fps

  if os.path.exists(output):
    os.remove(output)

  start = segrange[0]
  end = segrange[1]
  if trim:
    end -= trim

  if transcodeFlag: # and start > 0.0:
    # FIXME - match codec of source?
    qstr = ""
    if quality:
      qstr = "-crf %s" % (quality)
    # cmd = "ffmpeg -y -i \"%s\" -ss %f -to %f -c copy -vcodec libx264 -map 0 %s \"%s\"" \
    #     % (input, start, end, qstr, output)
    cmd = "rip-video --start-at %f --stop-at %f \"%s\" \"%s\"" % (start, end, input, output)
    if transcodeOptions:
      cmd += " " + transcodeOptions
    print(cmd)
    viddin.runCommand(cmd)
  else:
    if iframeFlag:
      start = findNearest(start, True)
      iend = findNearest(end, False)
      if iend:
        end = iend

    if start > 0.0:
      if end < vlen:
        temp = "split-%i%s" % (os.getpgid(0), ext)
        if os.path.exists(temp):
          os.remove(temp)
      else:
        temp = output
      cmd = "mkvmerge --split parts:%s-%s -o \"%s\" \"%s\"" % \
          (datetime.timedelta(seconds = start), datetime.timedelta(seconds = end),
           temp, input)
      print(cmd)
      viddin.runCommand(cmd)
    else:
      temp = input

    if end < vlen:
      tlen = end - start
      cmd = "ffmpeg -y -i \"%s\" -t %f -codec copy -map 0 \"%s\"" \
          % (temp, tlen, output)
      print(cmd)
      viddin.runCommand(cmd)
      
    if temp != input and temp != output:
      os.remove(temp)
  return output

def main():
  args = build_argparser().parse_args()
  segments = []
  if len(args.ranges) == 1 and os.path.isfile(args.ranges[0]):
    ranges = []
    with open(args.ranges[0]) as f:
      for line in f:
        line = line.strip()
        if line[0] >= '0' and line[0] <= '9':
          ranges.append(line)
    args.ranges = ranges

  for offset in args.ranges:
    if re.match("^[-+]?[0-9]+(\.[0-9]+)?$", offset) \
          or re.match("^[-+]?[0-9]+(:[0-9]+)+(\.[0-9]+)?$", offset):
      tc = viddin.decodeTimecode(offset)
      if offset[0] != '+':
        if len(segments) == 0:
          segments.append([0, tc])
        else:
          if not segments[-1][1]:
            segments[-1][1] = tc
      if tc > 0:
        segments.append([tc, None])
    elif re.match("^[0-9]+(.[0-9]+)?-[0-9]+(\.[0-9]+)?$", offset) \
          or re.match("^[0-9]+(:[0-9]+)+(\.[0-9]+)?-[0-9]+(:[0-9]+)+(\.[0-9]+)?$", offset):
      times = offset.split("-")
      segments.append([viddin.decodeTimecode(times[0]), viddin.decodeTimecode(times[1])])
    elif re.match("^[0-9]+(.[0-9]+)?-$", offset) \
          or re.match("^[0-9]+(:[0-9]+)+(\.[0-9]+)?-?$", offset):
      times = offset.split("-")
      segments.append([viddin.decodeTimecode(times[0]), None])
    else:
      print("Unrecognized timecode " + offset)
      sys.exit(1)

  if args.chapters:
    vfile = viddin.VideoSpec(args.filename)
    vfile.normalizeChapters()
    chapters = vfile.chapters
    segments = []
    for beg, end in zip(chapters, chapters[1:]):
      segments.append([beg.position, end.position])

  if len(segments) < 1:
    print("Nothing to do")
    return
  
  append = ""

  digits = int(math.floor(math.log(len(segments), 10)) + 1)
  append = "_%%0%ii" % digits

  video, ext = os.path.splitext(args.filename)
  if args.output:
    ovid, oext = os.path.splitext(args.output)
    append = ""
  else:
    ovid = video
    oext = ext

  ifile = video + ".iframes"
  #if not args.iframes and os.path.exists(ifile) \
  #      and os.path.getctime(ifile) >= os.path.getctime(args.filename):
  #  args.iframes = True
  if args.iframes:
    findIframes(args.filename, ifile)
    iframes = []
    with open(ifile) as f:
      for line in f:
        iframes.append(float(line))

  # FIXME - if transcoding and more than one segment do the entire thing
  #         and force frames? If only one segment then skip mkvmerge and
  #         let ffmpeg pull the segment

  outtype = args.outtype
  if not outtype:
    outtype = oext
    if not args.transcode:
      outtype = ext
  if outtype[0] != '.':
    outtype = "." + outtype
  if outtype == ".avi":
    outtype = ".mkv"

  for segrange in segments:
    output = ovid + append + outtype
    if len(append):
      output = output % (segments.index(segrange) + 1)
    splitVideo(segrange, args.filename, output, args.trimend, quality=args.quality,
               transcodeFlag=args.transcode, transcodeOptions=args.transcode_flags,
               iframeFlag=args.iframes)

  # FIXME - if transcoding force keyframes into the right spot, but
  #         re-read because ffmpeg isn't exact and puts them "close"

  return

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