diff --git a/config.py b/config.py index 446d27d..122389e 100644 --- a/config.py +++ b/config.py @@ -82,6 +82,15 @@ def parse_args(): default=False, help="Extract bitmap (default: false)", ) + + parser.add_argument( + "--extractor-config-extract-disable-sdh", + dest="extractor_config_extract_sdh", + action="store_false", + default=True, + help="Disable sdh subtitles extraction", + ) + parser.add_argument( "--extractor-config-overwrite", action="store_true", @@ -162,6 +171,7 @@ def parse_args(): EXTRACTOR_EXCLUDE_FILE = config.extractor_exclude_file EXTRACTOR_EXCLUDE_APPEND = config.extractor_exclude_append EXTRACTOR_EXTRACT_BITMAP = config.extractor_extract_bitmap +EXTRACTOR_CONFIG_EXTRACT_SDH = config.extractor_config_extract_sdh EXTRACTOR_CONFIG_OVERWRITE = config.extractor_config_overwrite EXTRACTOR_CONFIG_DESIRED_FORMATS = config.extractor_config_desired_formats EXTRACTOR_CONFIG_LANGUAGES = config.extractor_config_languages diff --git a/extract/config.py b/extract/config.py index a565b90..61a17cb 100644 --- a/extract/config.py +++ b/extract/config.py @@ -2,8 +2,16 @@ Configuration management for subtitle extraction. """ +import logging from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from extract.prober import StreamInfo + +logger = logging.getLogger(__name__) + @dataclass class ExtractorConfig: @@ -15,6 +23,19 @@ class ExtractorConfig: # target languages languages: list[str] | tuple[str] = ("all",) unknown_language_as: str = "unknown" + extract_sdh: bool = True + + def is_stream_wanted(self, stream: "StreamInfo") -> bool: + + if self.extract_sdh == False and stream.is_sdh(): + logger.debug(f"Skipping unwanted SDH stream ({stream.index})") + return False + + is_wanted_lang = "all" in self.languages or stream.language in self.languages + + if not is_wanted_lang: + logger.debug( + f"Skipping unwanted language '{stream.language}' for stream {stream.index}" + ) - def is_language_wanted(self, language: str) -> bool: - return "all" in self.languages or language in self.languages + return is_wanted_lang diff --git a/extract/extractors/base.py b/extract/extractors/base.py index f57c956..ce48a69 100644 --- a/extract/extractors/base.py +++ b/extract/extractors/base.py @@ -63,10 +63,8 @@ def should_extract_stream( return False logger.debug(f"Overwriting existing file: {output_path}") - if not self.config.is_language_wanted(stream.language): - logger.debug( - f"Skipping unwanted language '{stream.language}' for stream {stream.index}" - ) + if not self.config.is_stream_wanted(stream): + logger.debug(f"Skipping unwanted stream {stream.index}") return False # Check codec support diff --git a/extract/prober.py b/extract/prober.py index ce8b0e5..99f8292 100644 --- a/extract/prober.py +++ b/extract/prober.py @@ -14,6 +14,25 @@ logger = logging.getLogger(__name__) +# { +# "streams": [ +# { +# "index": 2, +# "codec_name": "subrip", +# "codec_type": "subtitle", +# "disposition": { +# "hearing_impaired": 1, +# "default": 0, +# "forced": 0 +# }, +# "tags": { +# "language": "eng" +# } +# } +# ] +# } + + class StreamInfo: """Represents information about a subtitle stream.""" @@ -22,7 +41,9 @@ def __init__(self, stream_data: dict[str, Any]): self.index = stream_data.get("index") self.codec_name = stream_data.get("codec_name") self.codec_type = stream_data.get("codec_type") - self.data.setdefault("tags", {}) + + self.data["tags"] = stream_data.get("tags") or {} + self.data["disposition"] = stream_data.get("disposition") or {} @property def language(self) -> str: @@ -41,6 +62,9 @@ def title(self) -> str: def disposition(self) -> dict[str, int]: return self.data.get("disposition", {}) + def is_sdh(self) -> bool: + return bool(self.disposition.get("hearing_impaired", 0)) + def is_forced(self) -> bool: return bool(self.disposition.get("forced", 0)) diff --git a/main.py b/main.py index 1228e1b..5d1f0f7 100755 --- a/main.py +++ b/main.py @@ -43,6 +43,7 @@ def main(mainpath: str): "desired_formats": config.EXTRACTOR_CONFIG_DESIRED_FORMATS, "languages": config.EXTRACTOR_CONFIG_LANGUAGES, "unknown_language_as": config.EXTRACTOR_CONFIG_UNKNOWN_LANGUAGE_AS, + "extract_sdh": config.EXTRACTOR_CONFIG_EXTRACT_SDH, }, } )