def run_test(base_path, plots_path, test_name, params): def safe_add_key(args, key, name): if name in params: args.extend((key, str(params[name]))) def safe_add_path(args, folder, key, name): if name in params: args.extend((key, os.path.join(folder, params[name]))) logging.info('Testing "{0}"'.format(test_name)) folder = os.path.join(base_path, params['folder']) cmd = ["sushi"] safe_add_path(cmd, folder, '--src', 'src') safe_add_path(cmd, folder, '--dst', 'dst') safe_add_path(cmd, folder, '--src-keyframes', 'src-keyframes') safe_add_path(cmd, folder, '--dst-keyframes', 'dst-keyframes') safe_add_path(cmd, folder, '--src-timecodes', 'src-timecodes') safe_add_path(cmd, folder, '--dst-timecodes', 'dst-timecodes') safe_add_path(cmd, folder, '--script', 'script') safe_add_path(cmd, folder, '--chapters', 'chapters') safe_add_path(cmd, folder, '--src-script', 'src-script') safe_add_path(cmd, folder, '--dst-script', 'dst-script') safe_add_key(cmd, '--max-kf-distance', 'max-kf-distance') safe_add_key(cmd, '--max-ts-distance', 'max-ts-distance') safe_add_key(cmd, '--max-ts-duration', 'max-ts-duration') output_path = os.path.join(folder, params['dst']) + '.sushi.test.ass' cmd.extend(('-o', output_path)) if plots_path: cmd.extend(('--test-shift-plot', os.path.join(plots_path, '{0}.png'.format(test_name)))) log_path = os.path.join(folder, 'sushi_test.log') with open(log_path, "w") as log_file: try: subprocess.call(cmd, stderr=log_file, stdout=log_file) except Exception as e: logging.critical('Sushi failed on test "{0}": {1}'.format(test_name, e.message)) return False with set_file_logger(log_path): ideal_path = os.path.join(folder, params['ideal']) try: timecodes = Timecodes.from_file(os.path.join(folder, params['dst-timecodes'])) except KeyError: timecodes = Timecodes.cfr(params['fps']) return compare_scripts(ideal_path, output_path, timecodes, test_name, params['expected_errors'])
def test_cfr_timecodes_v1_with_overrides(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,23.976000\n3000,5000,23.976000' parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0/23.976, parsed.get_frame_size(0)) self.assertAlmostEqual(1.0/23.976, parsed.get_frame_size(25)) self.assertAlmostEqual(1.0/23.976*100, parsed.get_frame_time(100)) self.assertEqual(0, parsed.get_frame_time(0))
def test_cfr_timecodes_v1(self): text = '# timecode format v1\nAssume 23.976024' parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0/23.976024, parsed.get_frame_size(0)) self.assertAlmostEqual(1.0/23.976024, parsed.get_frame_size(25)) self.assertAlmostEqual(1.0/23.976024*100, parsed.get_frame_time(100)) self.assertEqual(0, parsed.get_frame_time(0)) self.assertEqual(0, parsed.get_frame_number(0)) self.assertEqual(27461, parsed.get_frame_number(1145.353))
def test_cfr_timecodes_v2(self): text = '# timecode format v2\n' + '\n'.join(str(1000 * x / 23.976) for x in range(0, 30000)) parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0/23.976, parsed.get_frame_size(0)) self.assertAlmostEqual(1.0/23.976, parsed.get_frame_size(25)) self.assertAlmostEqual(1.0/23.976*100, parsed.get_frame_time(100)) self.assertEqual(0, parsed.get_frame_time(0)) self.assertEqual(0, parsed.get_frame_number(0)) self.assertEqual(27461, parsed.get_frame_number(1145.353))
def test_get_frame_number(self): tcs = Timecodes.cfr(24000.0 / 1001.0) self.assertEqual(tcs.get_frame_number(0), 0) self.assertEqual(tcs.get_frame_number(1145.353), 27461) self.assertEqual(tcs.get_frame_number(1001.0 / 24000.0 * 1234567), 1234567)
def test_get_frame_size(self): tcs = Timecodes.cfr(23.976) t1 = tcs.get_frame_size(0) t2 = tcs.get_frame_size(1000) self.assertAlmostEqual(1.0 / 23.976, t1) self.assertAlmostEqual(t1, t2)
def test_get_frame_time_insane(self): tcs = Timecodes.cfr(23.976) t = tcs.get_frame_time(100000) self.assertAlmostEqual(100000.0 / 23.976, t)
def test_get_frame_size(self): tcs = Timecodes.cfr(23.976) t1 = tcs.get_frame_size(0) t2 = tcs.get_frame_size(1000) self.assertAlmostEqual(1.0/23.976, t1) self.assertAlmostEqual(t1, t2)
def run(args): ignore_chapters = args.chapters_file is not None and args.chapters_file.lower( ) == 'none' write_plot = plot_enabled and args.plot_path if write_plot: plt.clf() plt.ylabel('Shift, seconds') plt.xlabel('Event index') # first part should do all possible validation and should NOT take significant amount of time check_file_exists(args.source, 'Source') check_file_exists(args.destination, 'Destination') check_file_exists(args.src_timecodes, 'Source timecodes') check_file_exists(args.dst_timecodes, 'Source timecodes') check_file_exists(args.script_file, 'Script') if not ignore_chapters: check_file_exists(args.chapters_file, 'Chapters') if args.src_keyframes not in ('auto', 'make'): check_file_exists(args.src_keyframes, 'Source keyframes') if args.dst_keyframes not in ('auto', 'make'): check_file_exists(args.dst_keyframes, 'Destination keyframes') if (args.src_timecodes and args.src_fps) or (args.dst_timecodes and args.dst_fps): raise SushiError( 'Both fps and timecodes file cannot be specified at the same time') src_demuxer = Demuxer(args.source) dst_demuxer = Demuxer(args.destination) if src_demuxer.is_wav and not args.script_file: raise SushiError("Script file isn't specified") if (args.src_keyframes and not args.dst_keyframes) or (args.dst_keyframes and not args.src_keyframes): raise SushiError( 'Either none or both of src and dst keyframes should be provided') create_directory_if_not_exists(args.temp_dir) # selecting source audio if src_demuxer.is_wav: src_audio_path = args.source else: src_audio_path = format_full_path(args.temp_dir, args.source, '.sushi.wav') src_demuxer.set_audio(stream_idx=args.src_audio_idx, output_path=src_audio_path, sample_rate=args.sample_rate) # selecting destination audio if dst_demuxer.is_wav: dst_audio_path = args.destination else: dst_audio_path = format_full_path(args.temp_dir, args.destination, '.sushi.wav') dst_demuxer.set_audio(stream_idx=args.dst_audio_idx, output_path=dst_audio_path, sample_rate=args.sample_rate) # selecting source subtitles if args.script_file: src_script_path = args.script_file else: stype = src_demuxer.get_subs_type(args.src_script_idx) src_script_path = format_full_path(args.temp_dir, args.source, '.sushi' + stype) src_demuxer.set_script(stream_idx=args.src_script_idx, output_path=src_script_path) script_extension = get_extension(src_script_path) if script_extension not in ('.ass', '.srt'): raise SushiError('Unknown script type') # selection destination subtitles if args.output_script: dst_script_path = args.output_script dst_script_extension = get_extension(args.output_script) if dst_script_extension != script_extension: raise SushiError( "Source and destination script file types don't match ({0} vs {1})" .format(script_extension, dst_script_extension)) else: dst_script_path = format_full_path(args.temp_dir, args.destination, '.sushi' + script_extension) # selecting chapters if args.grouping and not ignore_chapters: if args.chapters_file: if get_extension(args.chapters_file) == '.xml': chapter_times = chapters.get_xml_start_times( args.chapters_file) else: chapter_times = chapters.get_ogm_start_times( args.chapters_file) elif not src_demuxer.is_wav: chapter_times = src_demuxer.chapters output_path = format_full_path(args.temp_dir, src_demuxer.path, ".sushi.chapters.txt") src_demuxer.set_chapters(output_path) else: chapter_times = [] else: chapter_times = [] # selecting keyframes and timecodes if args.src_keyframes: def select_keyframes(file_arg, demuxer): auto_file = format_full_path(args.temp_dir, demuxer.path, '.sushi.keyframes.txt') if file_arg in ('auto', 'make'): if file_arg == 'make' or not os.path.exists(auto_file): if not demuxer.has_video: raise SushiError( "Cannot make keyframes for {0} because it doesn't have any video!" .format(demuxer.path)) demuxer.set_keyframes(output_path=auto_file) return auto_file else: return file_arg def select_timecodes(external_file, fps_arg, demuxer): if external_file: return external_file elif fps_arg: return None elif demuxer.has_video: path = format_full_path(args.temp_dir, demuxer.path, '.sushi.timecodes.txt') demuxer.set_timecodes(output_path=path) return path else: raise SushiError( 'Fps, timecodes or video files must be provided if keyframes are used' ) src_keyframes_file = select_keyframes(args.src_keyframes, src_demuxer) dst_keyframes_file = select_keyframes(args.dst_keyframes, dst_demuxer) src_timecodes_file = select_timecodes(args.src_timecodes, args.src_fps, src_demuxer) dst_timecodes_file = select_timecodes(args.dst_timecodes, args.dst_fps, dst_demuxer) # after this point nothing should fail so it's safe to start slow operations # like running the actual demuxing src_demuxer.demux() dst_demuxer.demux() try: if args.src_keyframes: src_timecodes = Timecodes.cfr( args.src_fps) if args.src_fps else Timecodes.from_file( src_timecodes_file) src_keytimes = [ src_timecodes.get_frame_time(f) for f in parse_keyframes(src_keyframes_file) ] dst_timecodes = Timecodes.cfr( args.dst_fps) if args.dst_fps else Timecodes.from_file( dst_timecodes_file) dst_keytimes = [ dst_timecodes.get_frame_time(f) for f in parse_keyframes(dst_keyframes_file) ] script = AssScript.from_file( src_script_path ) if script_extension == '.ass' else SrtScript.from_file( src_script_path) script.sort_by_time() src_stream = WavStream(src_audio_path, sample_rate=args.sample_rate, sample_type=args.sample_type) dst_stream = WavStream(dst_audio_path, sample_rate=args.sample_rate, sample_type=args.sample_type) calculate_shifts( src_stream, dst_stream, script.events, chapter_times=chapter_times, window=args.window, max_window=args.max_window, rewind_thresh=args.rewind_thresh if args.grouping else 0, max_ts_duration=args.max_ts_duration, max_ts_distance=args.max_ts_distance) events = script.events if write_plot: plt.plot([x.shift for x in events], label='From audio') if args.grouping: if not ignore_chapters and chapter_times: groups = groups_from_chapters(events, chapter_times) for g in groups: fix_near_borders(g) smooth_events([x for x in g if not x.linked], args.smooth_radius) groups = split_broken_groups(groups, args.min_group_size) else: fix_near_borders(events) smooth_events([x for x in events if not x.linked], args.smooth_radius) groups = detect_groups(events, args.min_group_size) if write_plot: plt.plot([x.shift for x in events], label='Borders fixed') for g in groups: start_shift = g[0].shift end_shift = g[-1].shift avg_shift = average_shifts(g) logging.info( u'Group (start: {0}, end: {1}, lines: {2}), ' u'shifts (start: {3}, end: {4}, average: {5})'.format( format_time(g[0].start), format_time(g[-1].end), len(g), start_shift, end_shift, avg_shift)) if args.src_keyframes: for e in (x for x in events if x.linked): e.resolve_link() for g in groups: snap_groups_to_keyframes( g, chapter_times, args.max_ts_duration, args.max_ts_distance, src_keytimes, dst_keytimes, src_timecodes, dst_timecodes, args.max_kf_distance, args.kf_mode) if args.write_avs: write_shift_avs(dst_script_path + '.avs', groups, src_audio_path, dst_audio_path) else: fix_near_borders(events) if write_plot: plt.plot([x.shift for x in events], label='Borders fixed') if args.src_keyframes: for e in (x for x in events if x.linked): e.resolve_link() snap_groups_to_keyframes(events, chapter_times, args.max_ts_duration, args.max_ts_distance, src_keytimes, dst_keytimes, src_timecodes, dst_timecodes, args.max_kf_distance, args.kf_mode) for event in events: event.apply_shift() script.save_to_file(dst_script_path) if write_plot: plt.plot([ x.shift + (x._start_shift + x._end_shift) / 2.0 for x in events ], label='After correction') plt.legend(fontsize=5, frameon=False, fancybox=False) plt.savefig(args.plot_path, dpi=300) finally: if args.cleanup: src_demuxer.cleanup() dst_demuxer.cleanup()
def test_vfr_timecodes_v1_frame_time_outside_of_defined_range(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(1000.968, parsed.get_frame_time(number=25000), places=3)
def test_vft_timecodes_v1_frame_size_between_override_blocks(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0 / 23.976, parsed.get_frame_size(timestamp=87.496))
def test_vft_timecodes_v1_frame_size_between_override_blocks(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0/23.976, parsed.get_frame_size(timestamp=87.496))
def test_vfr_timecodes_v1_frame_size_outside_of_defined_range(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0/23.976, parsed.get_frame_size(timestamp=5000.0))
def test_get_frame_time_zero(self): tcs = Timecodes.cfr(23.976) t = tcs.get_frame_time(0) self.assertEqual(t, 0)
def test_get_frame_number(self): tcs = Timecodes.cfr(24000.0/1001.0) self.assertEqual(tcs.get_frame_number(0), 0) self.assertEqual(tcs.get_frame_number(1145.353), 27461) self.assertEqual(tcs.get_frame_number(1001.0/24000.0 * 1234567), 1234567)
def test_vfr_timecodes_v1_frame_size_outside_of_defined_range(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(1.0 / 23.976, parsed.get_frame_size(timestamp=5000.0))
def test_vfr_timecodes_v1_frame_time_at_first_frame(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(0, parsed.get_frame_time(number=0))
def test_vft_timecodes_v1_frame_time_between_override_blocks(self): text = '# timecode format v1\nAssume 23.976000\n0,2000,29.970000\n3000,4000,59.940000' parsed = Timecodes.parse(text) self.assertAlmostEqual(87.579, parsed.get_frame_time(number=2500), places=3)
def run(args): ignore_chapters = args.chapters_file is not None and args.chapters_file.lower() == 'none' write_plot = plot_enabled and args.plot_path if write_plot: plt.clf() plt.ylabel('Shift, seconds') plt.xlabel('Event index') # first part should do all possible validation and should NOT take significant amount of time check_file_exists(args.source, 'Source') check_file_exists(args.destination, 'Destination') check_file_exists(args.src_timecodes, 'Source timecodes') check_file_exists(args.dst_timecodes, 'Source timecodes') check_file_exists(args.script_file, 'Script') if not ignore_chapters: check_file_exists(args.chapters_file, 'Chapters') if args.src_keyframes not in ('auto', 'make'): check_file_exists(args.src_keyframes, 'Source keyframes') if args.dst_keyframes not in ('auto', 'make'): check_file_exists(args.dst_keyframes, 'Destination keyframes') if (args.src_timecodes and args.src_fps) or (args.dst_timecodes and args.dst_fps): raise SushiError('Both fps and timecodes file cannot be specified at the same time') src_demuxer = Demuxer(args.source) dst_demuxer = Demuxer(args.destination) if src_demuxer.is_wav and not args.script_file: raise SushiError("Script file isn't specified") if (args.src_keyframes and not args.dst_keyframes) or (args.dst_keyframes and not args.src_keyframes): raise SushiError('Either none or both of src and dst keyframes should be provided') create_directory_if_not_exists(args.temp_dir) # selecting source audio if src_demuxer.is_wav: src_audio_path = args.source else: src_audio_path = format_full_path(args.temp_dir, args.source, '.sushi.wav') src_demuxer.set_audio(stream_idx=args.src_audio_idx, output_path=src_audio_path, sample_rate=args.sample_rate) # selecting destination audio if dst_demuxer.is_wav: dst_audio_path = args.destination else: dst_audio_path = format_full_path(args.temp_dir, args.destination, '.sushi.wav') dst_demuxer.set_audio(stream_idx=args.dst_audio_idx, output_path=dst_audio_path, sample_rate=args.sample_rate) # selecting source subtitles if args.script_file: src_script_path = args.script_file else: stype = src_demuxer.get_subs_type(args.src_script_idx) src_script_path = format_full_path(args.temp_dir, args.source, '.sushi'+ stype) src_demuxer.set_script(stream_idx=args.src_script_idx, output_path=src_script_path) script_extension = get_extension(src_script_path) if script_extension not in ('.ass', '.srt'): raise SushiError('Unknown script type') # selection destination subtitles if args.output_script: dst_script_path = args.output_script dst_script_extension = get_extension(args.output_script) if dst_script_extension != script_extension: raise SushiError("Source and destination script file types don't match ({0} vs {1})" .format(script_extension, dst_script_extension)) else: dst_script_path = format_full_path(args.temp_dir, args.destination, '.sushi' + script_extension) # selecting chapters if args.grouping and not ignore_chapters: if args.chapters_file: if get_extension(args.chapters_file) == '.xml': chapter_times = chapters.get_xml_start_times(args.chapters_file) else: chapter_times = chapters.get_ogm_start_times(args.chapters_file) elif not src_demuxer.is_wav: chapter_times = src_demuxer.chapters output_path = format_full_path(args.temp_dir, src_demuxer.path, ".sushi.chapters.txt") src_demuxer.set_chapters(output_path) else: chapter_times = [] else: chapter_times = [] # selecting keyframes and timecodes if args.src_keyframes: def select_keyframes(file_arg, demuxer): auto_file = format_full_path(args.temp_dir, demuxer.path, '.sushi.keyframes.txt') if file_arg in ('auto', 'make'): if file_arg == 'make' or not os.path.exists(auto_file): if not demuxer.has_video: raise SushiError("Cannot make keyframes for {0} because it doesn't have any video!" .format(demuxer.path)) demuxer.set_keyframes(output_path=auto_file) return auto_file else: return file_arg def select_timecodes(external_file, fps_arg, demuxer): if external_file: return external_file elif fps_arg: return None elif demuxer.has_video: path = format_full_path(args.temp_dir, demuxer.path, '.sushi.timecodes.txt') demuxer.set_timecodes(output_path=path) return path else: raise SushiError('Fps, timecodes or video files must be provided if keyframes are used') src_keyframes_file = select_keyframes(args.src_keyframes, src_demuxer) dst_keyframes_file = select_keyframes(args.dst_keyframes, dst_demuxer) src_timecodes_file = select_timecodes(args.src_timecodes, args.src_fps, src_demuxer) dst_timecodes_file = select_timecodes(args.dst_timecodes, args.dst_fps, dst_demuxer) # after this point nothing should fail so it's safe to start slow operations # like running the actual demuxing src_demuxer.demux() dst_demuxer.demux() try: if args.src_keyframes: src_timecodes = Timecodes.cfr(args.src_fps) if args.src_fps else Timecodes.from_file(src_timecodes_file) src_keytimes = [src_timecodes.get_frame_time(f) for f in keyframes.parse_keyframes(src_keyframes_file)] dst_timecodes = Timecodes.cfr(args.dst_fps) if args.dst_fps else Timecodes.from_file(dst_timecodes_file) dst_keytimes = [dst_timecodes.get_frame_time(f) for f in keyframes.parse_keyframes(dst_keyframes_file)] script = AssScript.from_file(src_script_path) if script_extension == '.ass' else SrtScript.from_file(src_script_path) script.sort_by_time() src_stream = WavStream(src_audio_path, sample_rate=args.sample_rate, sample_type=args.sample_type) dst_stream = WavStream(dst_audio_path, sample_rate=args.sample_rate, sample_type=args.sample_type) search_groups = prepare_search_groups(script.events, source_duration=src_stream.duration_seconds, chapter_times=chapter_times, max_ts_duration=args.max_ts_duration, max_ts_distance=args.max_ts_distance) calculate_shifts(src_stream, dst_stream, search_groups, normal_window=args.window, max_window=args.max_window, rewind_thresh=args.rewind_thresh if args.grouping else 0) events = script.events if write_plot: plt.plot([x.shift for x in events], label='From audio') if args.grouping: if not ignore_chapters and chapter_times: groups = groups_from_chapters(events, chapter_times) for g in groups: fix_near_borders(g) smooth_events([x for x in g if not x.linked], args.smooth_radius) groups = split_broken_groups(groups) else: fix_near_borders(events) smooth_events([x for x in events if not x.linked], args.smooth_radius) groups = detect_groups(events) if write_plot: plt.plot([x.shift for x in events], label='Borders fixed') for g in groups: start_shift = g[0].shift end_shift = g[-1].shift avg_shift = average_shifts(g) logging.info(u'Group (start: {0}, end: {1}, lines: {2}), ' u'shifts (start: {3}, end: {4}, average: {5})' .format(format_time(g[0].start), format_time(g[-1].end), len(g), start_shift, end_shift, avg_shift)) if args.src_keyframes: for e in (x for x in events if x.linked): e.resolve_link() for g in groups: snap_groups_to_keyframes(g, chapter_times, args.max_ts_duration, args.max_ts_distance, src_keytimes, dst_keytimes, src_timecodes, dst_timecodes, args.max_kf_distance, args.kf_mode) else: fix_near_borders(events) if write_plot: plt.plot([x.shift for x in events], label='Borders fixed') if args.src_keyframes: for e in (x for x in events if x.linked): e.resolve_link() snap_groups_to_keyframes(events, chapter_times, args.max_ts_duration, args.max_ts_distance, src_keytimes, dst_keytimes, src_timecodes, dst_timecodes, args.max_kf_distance, args.kf_mode) for event in events: event.apply_shift() script.save_to_file(dst_script_path) if write_plot: plt.plot([x.shift + (x._start_shift + x._end_shift)/2.0 for x in events], label='After correction') plt.legend(fontsize=5, frameon=False, fancybox=False) plt.savefig(args.plot_path, dpi=300) finally: if args.cleanup: src_demuxer.cleanup() dst_demuxer.cleanup()
def test_get_frame_time_insane(self): tcs = Timecodes.cfr(23.976) t = tcs.get_frame_time(100000) self.assertAlmostEqual(100000.0/23.976, t)