def test_throw_exception_on_shifting_unknown_subtitle(self):
try:
unknown_file_path = os.path.join(self.__resource_tmp, "subtitle_test.unknown")
Path(unknown_file_path).touch()
Undertest.load(unknown_file_path)
except Exception as e:
self.assertTrue(isinstance(e, UnsupportedFormatException))
else:
self.fail("Should have thrown exception")
def test_throw_exception_on_missing_subtitle(self):
try:
unknown_file_path = os.path.join(self.__resource_tmp,
"subtitle_test.unknown")
Undertest.export_subtitle(unknown_file_path, None, "")
except Exception as e:
self.assertTrue(isinstance(e, UnsupportedFormatException))
else:
self.fail("Should have thrown exception")
def test_shift_stl_subtitle(self):
shifted_stl_file_path = os.path.join(self.__resource_tmp, "subtitle_test.srt")
Undertest.shift_subtitle(
self.__stl_file_path, 2, shifted_stl_file_path, suffix="_test"
)
with open(shifted_stl_file_path) as shifted:
for j, ls in enumerate(shifted):
pass
shifted_line_num = j + 1
self.assertEqual(32, shifted_line_num)
def test_save_subs_as_stl(self):
target_file_path = os.path.join(self.__resource_tmp,
"subtitle_converted.stl")
Undertest.save_subs_as_target_format(
Undertest.load(self.__stl_file_path).subs, self.__srt_file_path,
target_file_path)
with open(target_file_path) as target:
for j, lt in enumerate(target):
pass
target_line_num = j + 1
self.assertEqual(32, target_line_num)
def test_throw_exception_on_converting_to_unknown_subtitle(self):
try:
unknown_file_path = os.path.join(self.__resource_tmp,
"subtitle_test.unknown")
Path(unknown_file_path).touch()
Undertest.save_subs_as_target_format(
Undertest.load(self.__stl_file_path).subs,
self.__srt_file_path, unknown_file_path)
except Exception as e:
self.assertTrue(isinstance(e, UnsupportedFormatException))
else:
self.fail("Should have thrown exception")
def test_export_sami_subtitle(self):
target_file_path = os.path.join(self.__resource_tmp, "subtitle_test.srt")
Undertest.export_subtitle(
self.__stl_file_path,
Undertest.load(self.__stl_file_path).subs,
target_file_path,
)
with open(target_file_path) as target:
for j, lt in enumerate(target):
pass
target_line_num = j + 1
self.assertEqual(32, target_line_num)
def test_throw_exception_on_exporting_unknown_subtitle(self):
try:
unknown_file_path = os.path.join(self.__resource_tmp,
"subtitle_test.unknown")
Undertest.export_subtitle(
unknown_file_path,
Undertest.load(self.__ttml_file_path).subs,
"target",
)
except Exception as e:
self.assertTrue(isinstance(e, UnsupportedFormatException))
else:
self.fail("Should have thrown exception")
def test_shift_sami_subtitle(self):
shifted_sami_file_path = os.path.join(self.__resource_tmp, "subtitle_test.sami")
Undertest.shift_subtitle(
self.__sami_file_path, 2, shifted_sami_file_path, suffix="_test"
)
with open(self.__sami_file_path) as original:
for i, lo in enumerate(original):
pass
with open(shifted_sami_file_path) as shifted:
for j, ls in enumerate(shifted):
pass
original_line_num = i + 1
shifted_line_num = j + 1
self.assertEqual(original_line_num, shifted_line_num)
def test_export_sami_subtitle(self):
target_file_path = os.path.join(self.__resource_tmp, "subtitle_test.sami")
Undertest.export_subtitle(
self.__sami_file_path,
Undertest.load(self.__sami_file_path).subs,
target_file_path,
)
with open(self.__sami_file_path) as original:
for i, lo in enumerate(original):
pass
with open(target_file_path) as target:
for j, lt in enumerate(target):
pass
original_line_num = i + 1
target_line_num = j + 1
self.assertEqual(original_line_num, target_line_num)
def test_load(self):
srt_subtitle = Undertest.load(self.__srt_file_path)
ttml_subtitle = Undertest.load(self.__ttml_file_path)
vtt_subtitle = Undertest.load(self.__vtt_file_path)
ass_subtitle = Undertest.load(self.__ass_file_path)
ssa_subtitle = Undertest.load(self.__ssa_file_path)
microdvd_subtitle = Undertest.load(self.__microdvd_file_path)
mp2_subtitle = Undertest.load(self.__mpl2_file_path)
tmp_subtitle = Undertest.load(self.__tmp_file_path)
sami_subtitle = Undertest.load(self.__sami_file_path)
stl_subtitle = Undertest.load(self.__stl_file_path)
self.assertEqual(len(srt_subtitle.subs), len(ttml_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(vtt_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(ass_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(ssa_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(microdvd_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(mp2_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(tmp_subtitle.subs))
self.assertEqual(len(srt_subtitle.subs), len(sami_subtitle.subs))
self.assertEqual(7, len(stl_subtitle.subs))
def test_shift_srt_subtitle_without_destination_file_path(self):
shifted_srt_file_path = os.path.join(self.__resource_tmp,
"subtitle_test.srt")
another_shifted_srt_file_path = os.path.join(
self.__resource_tmp, "subtitle_test_aligned.srt")
Undertest.shift_subtitle(self.__srt_file_path,
2,
shifted_srt_file_path,
suffix="_test")
Undertest.shift_subtitle(shifted_srt_file_path,
2,
None,
suffix="_aligned")
with open(shifted_srt_file_path) as original:
for i, lo in enumerate(original):
pass
with open(another_shifted_srt_file_path) as shifted:
for j, ls in enumerate(shifted):
pass
original_line_num = i + 1
shifted_line_num = j + 1
self.assertEqual(original_line_num, shifted_line_num)
def test_get_vtt_file_path(self):
subtitle = Undertest.load_webvtt(self.__srt_file_path)
self.assertEqual(self.__srt_file_path, subtitle.subtitle_file_path)
def test_get_ttml_file_path_ttml(self):
subtitle = Undertest.load_ttml(self.__another_ttml_file_path)
self.assertEqual(self.__another_ttml_file_path, subtitle.subtitle_file_path)
def test_subtitle_extentions(self):
self.assertEqual({".srt", ".xml", ".ttml", ".dfxp", ".vtt", ".ssa", ".ass", ".sub", ".txt", ".tmp", ".smi", ".sami", ".stl"},
Undertest.subtitle_extensions())
def test_extract_text_from_stl(self):
text = Undertest.extract_text(self.__stl_file_path)
self.assertEqual(194, len(text))
def test_load_stl_subs(self):
subtitle = Undertest.load_stl(self.__stl_file_path)
self.assertGreater(len(subtitle.subs), 0)
def main():
if sys.version_info.major != 3:
print("Cannot find Python 3")
sys.exit(20)
try:
import subaligner
except ModuleNotFoundError:
print("Subaligner is not installed")
sys.exit(20)
from subaligner._version import __version__
parser = argparse.ArgumentParser(
description=
"Convert a subtitle from input format to output format (v%s)" %
__version__,
formatter_class=argparse.RawTextHelpFormatter)
required_args = parser.add_argument_group("required arguments")
required_args.add_argument(
"-i",
"--input_subtitle_path",
type=str,
default="",
help="File path or URL to the input subtitle file",
required=True,
)
required_args.add_argument(
"-o",
"--output_subtitle_path",
type=str,
default="",
help="File path to the output subtitle file",
required=True,
)
parser.add_argument("-d",
"--debug",
action="store_true",
help="Print out debugging information")
parser.add_argument("-q",
"--quiet",
action="store_true",
help="Switch off logging information")
parser.add_argument("-ver",
"--version",
action="version",
version=__version__)
FLAGS, unparsed = parser.parse_known_args()
if FLAGS.input_subtitle_path == "":
print("--input_subtitle_path was not passed in")
sys.exit(21)
if FLAGS.output_subtitle_path == "":
print("--output_subtitle_path was not passed in")
sys.exit(21)
local_subtitle_path = FLAGS.input_subtitle_path
from subaligner.logger import Logger
Logger.VERBOSE = FLAGS.debug
Logger.QUIET = FLAGS.quiet
from subaligner.subtitle import Subtitle
from subaligner.exception import UnsupportedFormatException, TerminalException
from subaligner.utils import Utils
try:
if FLAGS.input_subtitle_path.lower().startswith("http"):
_, local_subtitle_path = tempfile.mkstemp()
_, subtitle_file_extension = os.path.splitext(
FLAGS.input_subtitle_path.lower())
local_subtitle_path = "{}{}".format(local_subtitle_path,
subtitle_file_extension)
Utils.download_file(FLAGS.input_subtitle_path, local_subtitle_path)
subtitle = Subtitle.load(local_subtitle_path)
Subtitle.save_subs_as_target_format(subtitle.subs, local_subtitle_path,
FLAGS.output_subtitle_path)
print("Subtitle converted and saved to: {}".format(
FLAGS.output_subtitle_path))
except UnsupportedFormatException as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
_remove_tmp_files(FLAGS, local_subtitle_path)
sys.exit(23)
except TerminalException as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
_remove_tmp_files(FLAGS, local_subtitle_path)
sys.exit(24)
except Exception as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
_remove_tmp_files(FLAGS, local_subtitle_path)
sys.exit(1)
else:
_remove_tmp_files(FLAGS, local_subtitle_path)
sys.exit(0)
from subaligner.predictor import Predictor
from subaligner.subtitle import Subtitle
if __name__ == "__main__":
examples_dir = os.path.dirname(os.path.abspath(__file__))
output_dir = os.path.join(examples_dir, "tmp")
os.makedirs(output_dir, exist_ok=True)
video_file_path = os.path.join(examples_dir, "..",
"tests/subaligner/resource/test.mp4")
srt_file_path = os.path.join(examples_dir, "..",
"tests/subaligner/resource/test.srt")
predictor = Predictor()
subs, audio_file_path, voice_probabilities, frame_rate = predictor.predict_single_pass(
video_file_path, srt_file_path)
aligned_subtitle_path = os.path.join(output_dir, "test_aligned_1.srt")
Subtitle.export_subtitle(srt_file_path, subs, aligned_subtitle_path)
print("Aligned subtitle saved to: {}".format(aligned_subtitle_path))
log_loss = predictor.get_log_loss(voice_probabilities, subs)
print("Alignment finished with overall loss: {}".format(log_loss))
subs_list, subs, voice_probabilities, frame_rate = predictor.predict_dual_pass(
video_file_path, srt_file_path, stretch=False)
aligned_subtitle_path = os.path.join(output_dir, "test_aligned_2.srt")
Subtitle.export_subtitle(srt_file_path, subs_list, aligned_subtitle_path)
print("Aligned subtitle saved to: {}".format(aligned_subtitle_path))
log_loss = predictor.get_log_loss(voice_probabilities, subs)
print("Alignment finished with overall loss: {}".format(log_loss))
def test_load_vtt_subs(self):
subtitle = Undertest.load_webvtt(self.__vtt_file_path)
self.assertGreater(len(subtitle.subs), 0)
def test_get_microdvd_file_path(self):
subtitle = Undertest.load_microdvd(self.__microdvd_file_path)
self.assertEqual(self.__microdvd_file_path, subtitle.subtitle_file_path)
def test_remove_sound_effects_with_uppercase(self):
subtitle = Undertest.load(self.__srt_file_path)
new_subs = Undertest.remove_sound_effects_by_case(
subtitle.subs, se_uppercase=True
)
self.assertEqual(len(subtitle.subs) - 2, len(new_subs))
def test_get_stl_file_path(self):
subtitle = Undertest.load_stl(self.__stl_file_path)
self.assertEqual(self.__stl_file_path, subtitle.subtitle_file_path)
def test_extract_text_from_sami(self):
text = Undertest.extract_text(self.__sami_file_path)
with open(self.__subtxt_file_path) as target:
expected_text = target.read()
self.assertEqual(expected_text, text)
def main():
if sys.version_info.major != 3:
print("Cannot find Python 3")
sys.exit(20)
try:
import subaligner
except ModuleNotFoundError:
print("Subaligner is not installed")
sys.exit(20)
from subaligner._version import __version__
parser = argparse.ArgumentParser(
description="Run single-stage alignment (v%s)" % __version__,
formatter_class=argparse.RawTextHelpFormatter)
required_args = parser.add_argument_group("required arguments")
required_args.add_argument(
"-v",
"--video_path",
type=str,
default="",
help="File path or URL to the video file",
required=True,
)
from subaligner.subtitle import Subtitle
required_args.add_argument(
"-s",
"--subtitle_path",
type=str,
default="",
help=
"File path or URL to the subtitle file (Extensions of supported subtitles: {}) or selector for the embedded subtitle (e.g., embedded:page_num=888 or embedded:stream_index=0)"
.format(", ".join(Subtitle.subtitle_extensions())),
required=True,
)
parser.add_argument(
"-l",
"--max_logloss",
type=float,
default=float("inf"),
help="Max global log loss for alignment",
)
parser.add_argument(
"-tod",
"--training_output_directory",
type=str,
default=os.path.abspath(os.path.dirname(subaligner.__file__)),
help="Path to the output directory containing training results",
)
parser.add_argument(
"-o",
"--output",
type=str,
default="",
help="Path to the output subtitle file",
)
parser.add_argument("-d",
"--debug",
action="store_true",
help="Print out debugging information")
parser.add_argument("-q",
"--quiet",
action="store_true",
help="Switch off logging information")
parser.add_argument("-ver",
"--version",
action="version",
version=__version__)
FLAGS, unparsed = parser.parse_known_args()
if FLAGS.video_path == "":
print("--video_path was not passed in")
sys.exit(21)
if FLAGS.subtitle_path == "":
print("--subtitle_path was not passed in")
sys.exit(21)
if FLAGS.subtitle_path.lower().startswith("http") and FLAGS.output == "":
print(
"--output was not passed in for alignment on a remote subtitle file"
)
sys.exit(21)
if FLAGS.subtitle_path.lower().startswith(
"teletext:") and FLAGS.output == "":
print("--output was not passed in for alignment on embedded subtitles")
sys.exit(21)
local_video_path = FLAGS.video_path
local_subtitle_path = FLAGS.subtitle_path
from subaligner.logger import Logger
Logger.VERBOSE = FLAGS.debug
Logger.QUIET = FLAGS.quiet
from subaligner.predictor import Predictor
from subaligner.exception import UnsupportedFormatException
from subaligner.exception import TerminalException
from subaligner.utils import Utils
try:
if FLAGS.video_path.lower().startswith("http"):
_, local_video_path = tempfile.mkstemp()
_, video_file_extension = os.path.splitext(
FLAGS.video_path.lower())
local_video_path = "{}{}".format(local_video_path,
video_file_extension)
Utils.download_file(FLAGS.video_path, local_video_path)
if FLAGS.subtitle_path.lower().startswith("http"):
_, local_subtitle_path = tempfile.mkstemp()
_, subtitle_file_extension = os.path.splitext(
FLAGS.subtitle_path.lower())
local_subtitle_path = "{}{}".format(local_subtitle_path,
subtitle_file_extension)
Utils.download_file(FLAGS.subtitle_path, local_subtitle_path)
if FLAGS.subtitle_path.lower().startswith("embedded:"):
_, local_subtitle_path = tempfile.mkstemp()
_, subtitle_file_extension = os.path.splitext(FLAGS.output)
local_subtitle_path = "{}{}".format(local_subtitle_path,
subtitle_file_extension)
params = FLAGS.subtitle_path.lower().split(":")[1].split(",")
if params and "=" in params[0]:
params = {
param.split("=")[0]: param.split("=")[1]
for param in params
}
if "page_num" in params:
Utils.extract_teletext_as_subtitle(local_video_path,
int(params["page_num"]),
local_subtitle_path)
elif "stream_index" in params:
Utils.extract_matroska_subtitle(
local_video_path, int(params["stream_index"]),
local_subtitle_path)
else:
print("Embedded subtitle selector cannot be empty")
sys.exit(21)
predictor = Predictor()
subs, audio_file_path, voice_probabilities, frame_rate = predictor.predict_single_pass(
video_file_path=local_video_path,
subtitle_file_path=local_subtitle_path,
weights_dir=os.path.join(FLAGS.training_output_directory,
"models/training/weights"))
aligned_subtitle_path = "_aligned.".join(
FLAGS.subtitle_path.rsplit(".", 1)).replace(
".stl", ".srt") if FLAGS.output == "" else FLAGS.output
Subtitle.export_subtitle(local_subtitle_path, subs,
aligned_subtitle_path, frame_rate)
log_loss = predictor.get_log_loss(voice_probabilities, subs)
if log_loss is None or log_loss > FLAGS.max_logloss:
print("Alignment failed with a too high loss value: {}".format(
log_loss))
_remove_tmp_files(FLAGS, local_video_path, local_subtitle_path)
sys.exit(22)
print("Aligned subtitle saved to: {}".format(aligned_subtitle_path))
except UnsupportedFormatException as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
_remove_tmp_files(FLAGS, local_video_path, local_subtitle_path)
sys.exit(23)
except TerminalException as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
_remove_tmp_files(FLAGS, local_video_path, local_subtitle_path)
sys.exit(24)
except Exception as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
_remove_tmp_files(FLAGS, local_video_path, local_subtitle_path)
sys.exit(1)
else:
_remove_tmp_files(FLAGS, local_video_path, local_subtitle_path)
sys.exit(0)
def test_load_microdvd_subs(self):
subtitle = Undertest.load_microdvd(self.__microdvd_file_path)
self.assertGreater(len(subtitle.subs), 0)
def main():
if sys.version_info.major != 3:
print("Cannot find Python 3")
sys.exit(20)
try:
import subaligner
except ModuleNotFoundError:
print("Subaligner is not installed")
sys.exit(20)
parser = argparse.ArgumentParser(
description="Subaligner command line interface")
required_args = parser.add_argument_group("required arguments")
required_args.add_argument(
"-m",
"--mode",
type=str,
default="",
choices=["single", "dual"],
help="Alignment mode: either single or dual",
required=True,
)
required_args.add_argument(
"-v",
"--video_path",
type=str,
default="",
help="Path to the video file",
required=True,
)
required_args.add_argument(
"-s",
"--subtitle_path",
type=str,
default="",
help="Path to the subtitle file",
required=True,
)
parser.add_argument(
"-l",
"--max_logloss",
type=float,
default=float("inf"),
help="Max global log loss for alignment",
)
parser.add_argument(
"-so",
"--stretch_off",
action="store_true",
help="Switch off stretch on non-English speech and subtitles)",
)
parser.add_argument(
"-fos",
"--exit_segfail",
action="store_true",
help="Exit on any segment alignment failures",
)
parser.add_argument(
"-tod",
"--training_output_directory",
type=str,
default=os.path.abspath(os.path.dirname(subaligner.__file__)),
help="Path to the output directory containing training results",
)
parser.add_argument(
"-o",
"--output",
type=str,
default="",
help="Path to the output subtitle file",
)
parser.add_argument("-d",
"--debug",
action="store_true",
help="Print out debugging information")
parser.add_argument("-q",
"--quiet",
action="store_true",
help="Switch off logging information")
FLAGS, unparsed = parser.parse_known_args()
if FLAGS.mode == "":
print("--mode was not passed in")
sys.exit(21)
if FLAGS.video_path == "":
print("--video_path was not passed in")
sys.exit(21)
if FLAGS.subtitle_path == "":
print("--subtitle_path was not passed in")
sys.exit(21)
exit_segfail = FLAGS.exit_segfail
stretch = not FLAGS.stretch_off
from subaligner.logger import Logger
Logger.VERBOSE = FLAGS.debug
Logger.QUIET = FLAGS.quiet
from subaligner.predictor import Predictor
from subaligner.subtitle import Subtitle
from subaligner.exception import UnsupportedFormatException
from subaligner.exception import TerminalException
try:
predictor = Predictor()
if FLAGS.mode == "single":
aligned_subs, audio_file_path, voice_probabilities, frame_rate = predictor.predict_single_pass(
video_file_path=FLAGS.video_path,
subtitle_file_path=FLAGS.subtitle_path,
weights_dir=os.path.join(FLAGS.training_output_directory,
"models/training/weights"))
else:
aligned_subs, subs, voice_probabilities, frame_rate = predictor.predict_dual_pass(
video_file_path=FLAGS.video_path,
subtitle_file_path=FLAGS.subtitle_path,
weights_dir=os.path.join(FLAGS.training_output_directory,
"models/training/weights"),
stretch=stretch,
exit_segfail=exit_segfail,
)
aligned_subtitle_path = "_aligned.".join(
FLAGS.subtitle_path.rsplit(
".", 1)) if FLAGS.output == "" else FLAGS.output
Subtitle.export_subtitle(FLAGS.subtitle_path, aligned_subs,
aligned_subtitle_path, frame_rate)
print("Aligned subtitle saved to: {}".format(aligned_subtitle_path))
log_loss = predictor.get_log_loss(voice_probabilities, aligned_subs)
if log_loss is None or log_loss > FLAGS.max_logloss:
print("Alignment failed with a too high loss value: {}".format(
log_loss))
exit(22)
except UnsupportedFormatException as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
sys.exit(23)
except TerminalException as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
sys.exit(24)
except Exception as e:
print("{}\n{}".format(
str(e), "".join(traceback.format_stack()) if FLAGS.debug else ""))
sys.exit(1)
else:
exit(0)
def test_remove_sound_effects_with_out_suffix(self):
subtitle = Undertest.load(self.__srt_file_path)
new_subs = Undertest.remove_sound_effects_by_affixes(
subtitle.subs, se_prefix="("
)
self.assertEqual(len(subtitle.subs) - 2, len(new_subs))
