def startup(args: Args, chunk_queue: List[Chunk], cb: Callbacks): """ If resuming, open done file and get file properties from there else get file properties and """ # TODO: move this out and pass in total frames and initial frames done_path = args.temp / 'done.json' if args.resume and done_path.exists(): cb.run_callback("log", 'Resuming...\n') with open(done_path) as done_file: data = json.load(done_file) total = data['total'] done = len(data['done']) initial = sum(data['done'].values()) cb.run_callback("log", f'Resumed with {done} encoded clips done\n\n') else: initial = 0 total = frame_probe_fast(args.input, args.is_vs) d = {'total': total, 'done': {}} with open(done_path, 'w') as done_file: json.dump(d, done_file) clips = len(chunk_queue) args.workers = min(args.workers, clips) cb.run_callback( "log", f'\rQueue: {clips} Workers: {args.workers} Passes: {args.passes}\n' f'Params: {" ".join(args.video_params)}') cb.run_callback("newtask", "Encoding video", total) cb.run_callback("startencode", total, initial)
def _concatenate_mkvmerge(files, output, file_limit, cmd_limit, cb: Callbacks, flip=False): tmp_out = "{}.tmp{}.mkv".format(output, int(flip)) cmd = ["mkvmerge", "-o", tmp_out, files[0]] remaining = [] for i, file in enumerate(files[1:]): new_cmd = cmd + ['+{}'.format(file)] if sum(len(s) for s in new_cmd) < cmd_limit \ and (file_limit == -1 or i < max(1, file_limit - 10)): cmd = new_cmd else: remaining = files[i + 1:] break concat = subprocess.Popen(cmd, stdout=PIPE, universal_newlines=True) message, _ = concat.communicate() concat.wait() if concat.returncode != 0: cb.run_callback("log", message) print(message) raise Exception if len(remaining) > 0: return _concatenate_mkvmerge( [tmp_out] + remaining, output, file_limit, cmd_limit, cb, not flip) else: return tmp_out
def create_video_queue_segment(args: Args, split_locations: List[int], cb: Callbacks) -> List[Chunk]: """ Create a list of chunks using segmented files :param args: Args :param split_locations: a list of frames to split on :param cb: Reference to callback object to exit on failure :return: A list of chunks """ # segment into separate files segment(args.input, args.temp, split_locations, cb) # get the names of all the split files source_path = args.temp / 'split' queue_files = [x for x in source_path.iterdir() if x.suffix == '.mkv'] queue_files.sort(key=lambda p: p.stem) if len(queue_files) == 0: er = 'Error: No files found in temp/split, probably splitting not working' print(er) cb.run_callback("log", er) cb.run_callback("terminate", 1) chunk_queue = [ create_chunk_from_segment(args, index, file) for index, file in enumerate(queue_files) ] return chunk_queue
def process_encoding_pipe(pipes, encoder, cb: Callbacks): encoder_history = deque(maxlen=20) frame = 0 pipe = pipes[2] while True: line = pipe.stdout.readline().strip() if len(line) == 0 and pipe.poll() is not None: break if len(line) == 0: continue match = encoder.match_line(line, cb) if match: new = int(match.group(1)) if new > frame: cb.run_callback("newframes", new - frame) # counter.update(new - frame) frame = new if line: encoder_history.append(line) if pipe.returncode != 0 and pipe.returncode != -2: # -2 is Ctrl+C for aom print(f"\nEncoder encountered an error: {pipe.returncode}") print('\n'.join(encoder_history)) if pipe.returncode != 0 or pipes[0].returncode != 0 or pipes[ 1].returncode != 0: return 1 return 0
def plot_vmaf(source: Path, encoded: Path, args, model, vmaf_res, cb: Callbacks): """ Making VMAF plot after encode is done """ print('Calculating Vmaf...\r', end='') fl_path = encoded.with_name(f'{encoded.stem}_vmaflog').with_suffix(".json") # call_vmaf takes a chunk, so make a chunk of the entire source ffmpeg_gen_cmd = [ 'ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-i', source.as_posix(), "-pix_fmt", args.pix_format, '-f', 'yuv4mpegpipe', '-' ] input_chunk = Chunk(args.temp, 0, ffmpeg_gen_cmd, '', 0, 0) scores = call_vmaf(input_chunk, encoded, 0, model, vmaf_res, fl_path=fl_path) if not scores.exists(): print( f'Vmaf calculation failed for chunks:\n {source.name} {encoded.stem}' ) sys.exit() file_path = encoded.with_name(f'{encoded.stem}_plot').with_suffix('.png') cb.run_callback("plotvmaffile", scores, file_path)
def concatenate_mkvmerge(temp: Path, output, cb: Callbacks): """ Uses mkvmerge to concatenate encoded segments into the final file :param temp: the temp directory :param output: the final output file :param cb: the callbacks :return: None """ cb.run_callback("log", 'Concatenating\n') output = shlex.quote(output.as_posix()) encode_files = sorted((temp / 'encode').iterdir(), key=lambda x: int(x.stem) if x.stem.isdigit() else x.stem) encode_files = [shlex.quote(f.as_posix()) for f in encode_files] if platform.system() == "Linux": import resource file_limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE) cmd_limit = os.sysconf(os.sysconf_names['SC_ARG_MAX']) else: file_limit = -1 cmd_limit = 32767 audio_file = temp / "audio.mkv" audio = audio_file.as_posix() if audio_file.exists() else '' if len(encode_files) > 1: encode_files = [ _concatenate_mkvmerge(encode_files, output, file_limit, cmd_limit, cb) ] cmd = ['mkvmerge', '-o', output, encode_files[0]] if audio: cmd.append(audio) concat = subprocess.Popen(cmd, stdout=PIPE, universal_newlines=True) message, _ = concat.communicate() concat.wait() if concat.returncode != 0: cb.run_callback("log", message) print(message) raise Exception # remove temporary files used by recursive concat if os.path.exists("{}.tmp0.mkv".format(output)): os.remove("{}.tmp0.mkv".format(output)) if os.path.exists("{}.tmp1.mkv".format(output)): os.remove("{}.tmp1.mkv".format(output))
def match_line(self, line: str, cb: Callbacks): """Extract number of encoded frames from line. :param line: one line of text output from the encoder :param cb: Callbacks reference in case error :return: match object from re.search matching the number of encoded frames""" if 'error' in line.lower(): print('\n\nERROR IN ENCODING PROCESS\n\n', line) cb.run_callback("terminate", 1) return re.search(r"encoded.*? ([^ ]+?) ", line)
def time(video, split_frames, cb: Callbacks): """ Running Time splitting on source video. """ frames = frame_probe_fast(video) to_add = frames // split_frames splits = None splits = list(linspace(split_frames, frames - split_frames - (frames % split_frames), to_add, dtype=int)) newsplits = [] for i in splits: newsplits.append(i.item()) cb.run_callback("log", f'Found scenes: {len(splits)}\n') return newsplits
def extract_audio(input_vid: Path, temp, audio_params, cb: Callbacks): """Extracting audio from source, transcoding if needed.""" cb.run_callback("log", f'Audio processing\nParams: {" ".join(audio_params)}\n') audio_file = temp / 'audio.mkv' # Checking is source have audio track check = ['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-ss', '0', '-i', input_vid.as_posix(), '-t', '0', '-vn', '-c:a', 'copy', '-f', 'null', '-'] is_audio_here = len(subprocess.run(check, stdout=PIPE, stderr=STDOUT).stdout) == 0 # If source have audio track - process it if is_audio_here: cmd = ('ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-i', input_vid.as_posix(), '-vn', '-sn', '-dn', *audio_params, audio_file.as_posix()) subprocess.run(cmd)
def calc_split_locations(args: Args, cb: Callbacks) -> List[int]: """ Determines a list of frame numbers to split on with pyscenedetect or aom keyframes :param args: the Args :param cb: Callback reference for log/failure :return: A list of frame numbers """ # inherit video params from aom encode unless we are using a different encoder, then use defaults if args.split_method == 'none': cb.run_callback("log", 'Skipping scene detection\n') return [] sc = [] # Split from file if args.split_method == 'file': if args.scenes.exists(): # Read stats from CSV file opened in read mode: cb.run_callback("log", 'Using Saved Scenes\n') return read_scenes_from_file(args.scenes) # Splitting using PySceneDetect elif args.split_method == 'pyscene': cb.run_callback("log", f'Starting scene detection Threshold: {args.threshold}\n') try: sc = pyscene(args.input, args.threshold, args.is_vs, args.temp, cb) except Exception as e: cb.run_callback("log", f'Error in PySceneDetect: {e}\n') print(f'Error in PySceneDetect{e}\n') # Splitting based on aom keyframe placement elif args.split_method == 'ffmpeg': stat_file = args.temp / 'keyframes.log' sc = ffmpeg(args.input, args.threshold, args.is_vs, args.temp, cb) elif args.split_method == "time": sc = time(args.input, args.time_split_interval, cb) else: print(f'No valid split option: {args.split_method}\nValid options: "pyscene", "aom_keyframes"') cb.run_callback("terminate", 1) # Write scenes to file if args.scenes: write_scenes_to_file(sc, args.scenes) return sc
def encoding_loop(args: Args, cb: Callbacks, chunk_queue: List[Chunk]): """Creating process pool for encoders, creating progress bar.""" if args.workers != 0: with concurrent.futures.ThreadPoolExecutor( max_workers=args.workers) as executor: future_cmd = { executor.submit(encode, cmd, args, cb): cmd for cmd in chunk_queue } for future in concurrent.futures.as_completed(future_cmd): try: future.result() except Exception as exc: _, _, exc_tb = sys.exc_info() print(f'Encoding error {exc}\nAt line {exc_tb.tb_lineno}') cb.run_callback("terminate", 1)
def concat_routine(args: Args, cb: Callbacks): """ Runs the concatenation routine with args :param args: the Args :param cb: the callbacks :return: None """ try: if args.mkvmerge: concatenate_mkvmerge(args.temp, args.output_file, cb) else: concatenate_ffmpeg(args.temp, args.output_file, args.encoder, cb) except Exception as e: _, _, exc_tb = sys.exc_info() print(f'Concatenation failed, error\nAt line: {exc_tb.tb_lineno}\nError:{str(e)}') cb.run_callback("log", f'Concatenation failed, aborting, error: {e}\n') cb.run_callback("terminate", 1)
def concatenate_ffmpeg(temp: Path, output: Path, encoder: str, cb: Callbacks): """ Uses ffmpeg to concatenate encoded segments into the final file :param temp: the temp directory :param output: the final output file :param encoder: the encoder :param cb: the callbacks :return: None """ """With FFMPEG concatenate encoded segments into final file.""" cb.run_callback("log", 'Concatenating\n') with open(temp / "concat", 'w') as f: encode_files = sorted((temp / 'encode').iterdir()) # Replace all the ' with '/'' so ffmpeg can read the path correctly f.writelines(f'file {shlex.quote(str(file.absolute()))}\n' for file in encode_files) # Add the audio file if one was extracted from the input audio_file = temp / "audio.mkv" if audio_file.exists(): audio = ('-i', audio_file.as_posix(), '-c:a', 'copy', '-map', '1') else: audio = () if encoder == 'x265': cmd = ['ffmpeg', '-y', '-fflags', '+genpts', '-hide_banner', '-loglevel', 'error', '-f', 'concat', '-safe', '0', '-i', (temp / "concat").as_posix(), *audio, '-c', 'copy', '-movflags', 'frag_keyframe+empty_moov', '-map', '0', '-f', 'mp4', output.as_posix()] concat = subprocess.run(cmd, stdout=PIPE, stderr=STDOUT).stdout else: cmd = ['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-f', 'concat', '-safe', '0', '-i', (temp / "concat").as_posix(), *audio, '-c', 'copy', '-map', '0', output.as_posix()] concat = subprocess.run(cmd, stdout=PIPE, stderr=STDOUT).stdout if len(concat) > 0: cb.run_callback("log", concat.decode()) print(concat.decode()) raise Exception
def process_inputs(inputs, cb: Callbacks): # Check input file for being valid if not inputs: print('No input file') cb.run_callback("terminate", 1) if inputs[0].is_dir(): inputs = [ x for x in inputs[0].iterdir() if x.suffix in (".mkv", ".mp4", ".mov", ".avi", ".flv", ".m2ts") ] valid = np.array([i.exists() for i in inputs]) if not all(valid): print( f'File(s) do not exist: {", ".join([str(inputs[i]) for i in np.where(not valid)[0]])}' ) cb.run_callback("terminate", 1) return inputs
def segment(video: Path, temp: Path, frames: List[int], cb: Callbacks): """ Uses ffmpeg to segment the video into separate files. Splits the video by frame numbers or copies the video if no splits are needed :param video: the source video :param temp: the temp directory :param frames: the split locations :param cb: Callback reference :return: None """ cb.run_callback("log", 'Split Video\n') cmd = [ "ffmpeg", "-hide_banner", "-y", "-i", video.absolute().as_posix(), "-map", "0:v:0", "-an", "-c", "copy", "-avoid_negative_ts", "1", "-vsync", "0" ] if len(frames) > 0: cmd.extend([ "-f", "segment", "-segment_frames", ','.join([str(x) for x in frames]) ]) cmd.append(os.path.join(temp, "split", "%05d.mkv")) else: cmd.append(os.path.join(temp, "split", "0.mkv")) pipe = subprocess.Popen(cmd, stdout=PIPE, stderr=STDOUT) while True: line = pipe.stdout.readline().strip() if len(line) == 0 and pipe.poll() is not None: break cb.run_callback("log", 'Split Done\n')
def ffmpeg(video, threshold, is_vs, temp, cb: Callbacks): """ Running PySceneDetect detection on source video for segmenting. Optimal threshold settings 15-50 """ cb.run_callback( "log", f'Starting FFMPEG detection:\nThreshold: {threshold}, Is Vapoursynth input: {is_vs}\n' ) if is_vs: # Handling vapoursynth. Outputs vs to a file so ffmpeg can handle it. if sys.platform == "linux": vspipe_fifo = temp / 'vspipe.y4m' mkfifo(vspipe_fifo) else: vspipe_fifo = None vspipe_cmd = compose_vapoursynth_pipe(video, vspipe_fifo) vspipe_process = Popen(vspipe_cmd) finfo = "showinfo,select=gt(scene\\," + str( threshold) + "),select=eq(key\\,1),showinfo" ffmpeg_cmd = [ "ffmpeg", "-i", str(vspipe_fifo if is_vs else video.as_posix()), "-hide_banner", "-loglevel", "32", "-filter_complex", finfo, "-an", "-f", "null", "-" ] pipe = subprocess.Popen(ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) last_frame = -1 scenes = [] while True: line = pipe.stderr.readline().strip() if len(line) == 0 and pipe.poll() is not None: print(pipe.poll()) break if len(line) == 0: continue if line: cur_frame = re.search("n:\\ *[0-9]+", str(line)) if cur_frame is not None: frame_num = re.search("[0-9]+", cur_frame.group(0)) if frame_num is not None: frame_num = int(frame_num.group(0)) if frame_num < last_frame: scenes += [last_frame] else: last_frame = frame_num # If fed using a vspipe process, ensure that vspipe has finished. if is_vs: vspipe_process.wait() # Remove 0 from list if len(scenes) > 0 and scenes[0] == 0: scenes.remove(0) cb.run_callback("log", f'Found split points: {len(scenes)}\n') cb.run_callback("log", f'Splits: {scenes}\n') return scenes
def target_vmaf(chunk: Chunk, args: Args, cb: Callbacks): vmaf_cq = [] frames = chunk.frames q_list = [] score = 0 # Make middle probe middle_point = (args.min_q + args.max_q) // 2 q_list.append(middle_point) last_q = middle_point score = vmaf_probe(chunk, last_q, args) vmaf_cq.append((score, last_q)) if args.vmaf_steps < 3: #Use Euler's method with known relation between cq and vmaf vmaf_cq_deriv = -0.18 ## Formula -ln(1-score/100) = vmaf_cq_deriv*last_q + constant #constant = -ln(1-score/100) - vmaf_cq_deriv*last_q ## Formula -ln(1-args.vmaf_target/100) = vmaf_cq_deriv*cq + constant #cq = (-ln(1-args.vmaf_target/100) - constant)/vmaf_cq_deriv next_q = int(round(last_q + (transform_vmaf(args.vmaf_target)-transform_vmaf(score))/vmaf_cq_deriv)) #Clamp if next_q < args.min_q: next_q = args.min_q if args.max_q < next_q: next_q = args.max_q #Single probe cq guess or exit to avoid divide by zero if args.vmaf_steps == 1 or next_q == last_q: return next_q #Second probe at guessed value score_2 = vmaf_probe(chunk, next_q, args) #Calculate slope vmaf_cq_deriv = (transform_vmaf(score_2)-transform_vmaf(score))/(next_q-last_q) #Same deal different slope next_q = int(round(next_q+(transform_vmaf(args.vmaf_target)-transform_vmaf(score_2))/vmaf_cq_deriv)) #Clamp if next_q < args.min_q: next_q = args.min_q if args.max_q < next_q: next_q = args.max_q return next_q # Initialize search boundary vmaf_lower = score vmaf_upper = score vmaf_cq_lower = last_q vmaf_cq_upper = last_q # Branch if score < args.vmaf_target: next_q = args.min_q q_list.append(args.min_q) else: next_q = args.max_q q_list.append(args.max_q) # Edge case check score = vmaf_probe(chunk, next_q, args) vmaf_cq.append((score, next_q)) if next_q == args.min_q and score < args.vmaf_target: cb.run_callback("log", f"Chunk: {chunk.name}, Fr: {frames}\n" f"Q: {sorted([x[1] for x in vmaf_cq])}, Early Skip Low CQ\n" f"Vmaf: {sorted([x[0] for x in vmaf_cq], reverse=True)}\n" f"Target Q: {vmaf_cq[-1][1]} Vmaf: {vmaf_cq[-1][0]}\n\n") return next_q elif next_q == args.max_q and score > args.vmaf_target: cb.run_callback("log", f"Chunk: {chunk.name}, Fr: {frames}\n" f"Q: {sorted([x[1] for x in vmaf_cq])}, Early Skip High CQ\n" f"Vmaf: {sorted([x[0] for x in vmaf_cq], reverse=True)}\n" f"Target Q: {vmaf_cq[-1][1]} Vmaf: {vmaf_cq[-1][0]}\n\n") return next_q # Set boundary if score < args.vmaf_target: vmaf_lower = score vmaf_cq_lower = next_q else: vmaf_upper = score vmaf_cq_upper = next_q # VMAF search for _ in range(args.vmaf_steps - 2): new_point = weighted_search(vmaf_cq_lower, vmaf_lower, vmaf_cq_upper, vmaf_upper, args.vmaf_target) if new_point in [x[1] for x in vmaf_cq]: break q_list.append(new_point) score = vmaf_probe(chunk, new_point, args) vmaf_cq.append((score, new_point)) # Update boundary if score < args.vmaf_target: vmaf_lower = score vmaf_cq_lower = new_point else: vmaf_upper = score vmaf_cq_upper = new_point q, q_vmaf = get_target_q(vmaf_cq, args.vmaf_target) cb.run_callback("log", f'Chunk: {chunk.name}, Fr: {frames}\n' f'Q: {sorted([x[1] for x in vmaf_cq])}\n' f'Vmaf: {sorted([x[0] for x in vmaf_cq], reverse=True)}\n' f'Target Q: {q} Vmaf: {q_vmaf}\n\n') # Plot Probes if len(vmaf_cq) > 3: cb.run_callback("plotvmaf", args.vmaf_target, args.min_q, args.max_q, args.temp, vmaf_cq, chunk.name, frames) return q
def encode(chunk: Chunk, args: Args, cb: Callbacks): """ Encodes a chunk. :param chunk: The chunk to encode :param args: The cli args :param cb: The callback object :return: None """ try: st_time = time.time() chunk_frames = chunk.frames cb.run_callback("log", f'Enc: {chunk.name}, {chunk_frames} fr\n\n') # Target Vmaf Mode if args.vmaf_target: target_vmaf_routine(args, chunk, cb) ENCODERS[args.encoder].on_before_chunk(args, chunk) # Run all passes for this chunk for current_pass in range(1, args.passes + 1): try: enc = ENCODERS[args.encoder] err = 1 while err != 0: pipe = enc.make_encode_pipes(args, chunk, args.passes, current_pass, chunk.output) if not args.is_debug: err = process_encoding_pipe(pipe, enc, cb) else: err = process_enc_debug_pipes(pipe, enc, cb) if err != 0: print("Error running encode on chunk: " + str(chunk) + "... Retrying...") except Exception as e: _, _, exc_tb = sys.exc_info() traceback.print_exc() print(f'Error at encode {e}. At line {exc_tb.tb_lineno}') ENCODERS[args.encoder].on_after_chunk(args, chunk) # get the number of encoded frames, if no check assume it worked and encoded same number of frames encoded_frames = chunk_frames if args.no_check else frame_check_output( chunk, chunk_frames) # write this chunk as done if it encoded correctly if encoded_frames == chunk_frames: write_progress_file(Path(args.temp / 'done.json'), chunk, encoded_frames) enc_time = round(time.time() - st_time, 2) cb.run_callback( "log", f'Done: {chunk.name} Fr: {encoded_frames}\n' f'Fps: {round(encoded_frames / enc_time, 4)} Time: {enc_time} sec.\n\n' ) except Exception as e: _, _, exc_tb = sys.exc_info() print(f'Error in encoding loop {e}\nAt line {exc_tb.tb_lineno}')
def pyscene(video, threshold, is_vs, temp, cb: Callbacks): """ Running PySceneDetect detection on source video for segmenting. Optimal threshold settings 15-50 """ min_scene_len = 15 cb.run_callback( "log", f'Starting PySceneDetect:\nThreshold: {threshold}, Is Vapoursynth input: {is_vs}\n' ) if is_vs: if sys.platform == "linux": vspipe_fifo = temp / 'vspipe.y4m' mkfifo(vspipe_fifo) else: vspipe_fifo = None vspipe_cmd = compose_vapoursynth_pipe(video, vspipe_fifo) vspipe_process = Popen(vspipe_cmd) # Get number of frames from Vapoursynth script to pass as duration to VideoManager. # We need to pass the number of frames to the manager, otherwise it won't close the # receiving end of the pipe, and will simply sit waiting after vspipe has finished sending # the last frame. frames = frame_probe_fast(video) video_manager = VideoManager([str(vspipe_fifo if is_vs else video)]) scene_manager = SceneManager() scene_manager.add_detector( ContentDetector(threshold=threshold, min_scene_len=min_scene_len)) base_timecode = video_manager.get_base_timecode() video_manager.set_duration(duration=FrameTimecode( frames, video_manager.get_framerate()) if is_vs else None) cb.run_callback("newtask", "Pyscenedetect", frames) # Set downscale factor to improve processing speed. video_manager.set_downscale_factor() # Start video_manager. video_manager.start() scene_manager.detect_scenes(frame_source=video_manager, show_progress=True) # If fed using a vspipe process, ensure that vspipe has finished. if is_vs: vspipe_process.wait() # Obtain list of detected scenes. scene_list = scene_manager.get_scene_list(base_timecode) scenes = [int(scene[0].get_frames()) for scene in scene_list] # Remove 0 from list if scenes[0] == 0: scenes.remove(0) cb.run_callback("log", f'Found scenes: {len(scenes)}\n') return scenes