41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
import argparse
|
||
|
||
from subtitles.parser import parse_srt_to_df
|
||
from scoring.bert_ranker import BERTImportanceRanker
|
||
from video_editor.highlight import VideoHighlighter
|
||
import config
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Оценка важности фраз в субтитрах с помощью BERT.")
|
||
parser.add_argument("--srt_path", help="Путь к .srt файлу субтитров")
|
||
parser.add_argument("--video_path", help="Путь к видео файлу")
|
||
parser.add_argument(
|
||
"--output", help="Путь к выходному CSV", default=config.CSV_OUTPUT)
|
||
args = parser.parse_args()
|
||
|
||
df = parse_srt_to_df(args.srt_path)
|
||
ranker = BERTImportanceRanker(config.MODEL_NAME)
|
||
df_scored = ranker.apply_to_dataframe(df)
|
||
important_df = (
|
||
df[df["importance_score"] > 0.9]
|
||
.loc[:, ["start", "end"]]
|
||
.reset_index(drop=True)
|
||
)
|
||
vh = VideoHighlighter(
|
||
video=args.video_path,
|
||
segments_df=important_df,
|
||
pad=1.5,
|
||
join_gap=0.3,
|
||
out_dir="clips",
|
||
concat=True,
|
||
)
|
||
vh.cut()
|
||
df_scored.to_csv(args.output, index=False, encoding='utf-8')
|
||
print(f"✅ Сохранено в {args.output}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|