ted-talks-annotation/downloader.py at master · nlpcl-lab/ted-talks-annotation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import os
from tqdm import tqdm
import subprocess
import pysrt
import csv

from mongoengine import connect

from config import Config
from annotation.models import Doc, Sent


def download_video_from_youtube(title, source):
    print('[download_video_from_youtube] start:', source)
    filename = title.replace(' ', '_')

    if Doc.objects.filter(title=title).count():
        doc = Doc.objects.get(title=title, source=source)
    else:
        doc = Doc(title=title, source=source)
        doc.save()

    subprocess.call(['youtube-dl',
                     # for sub
                     '--sub-lang', 'en',
                     '--write-sub',
                     '--convert-subs', 'srt',
                     # for audio
                     '--extract-audio',
                     '--audio-format', 'mp3',
                     # for filename
                     '--output', './data/videos/{}.%(ext)s'.format(filename),
                     source,
                     ])

    prefix_path = os.path.abspath(os.path.dirname(__file__) + './data/videos/{}'.format(filename))
    subs = align_subtitles(prefix_path)
    for index, sub in enumerate(subs):
        if Sent.objects.filter(doc=doc, index=index + 1).count():
            continue
        sent = Sent(
            doc=doc,
            index=index + 1,
            text=sub.eng_text,
            start_ts=sub.start.ordinal,
            end_ts=sub.end.ordinal)
        sent.save()

    print('[download_video_from_youtube] finish')


class Subtitle:
    def __init__(self):
        self.start = 0
        self.end = 0
        self.eng_text = ''
        self.kor_text = ''


def clean_text(text):
    return text.strip().replace('\n', ' ')


def align_subtitles(srt_path):
    eng_subs = pysrt.open(srt_path + '.en.srt')

    subs = []
    start = 0
    text = ''
    for eng_sub in eng_subs:
        text += ' ' + clean_text(eng_sub.text)
        if start == 0:
            start = eng_sub.start
        if text[-1] in ['.', ':', '-', '"', ')', '?', '!', ';']:
            sub = Subtitle()
            sub.eng_text = clean_text(text)
            sub.start = start
            sub.end = eng_sub.end
            subs.append(sub)
            text = ''
            start = 0

    return subs


if __name__ == '__main__':
    connect(**Config.MONGODB_SETTINGS)

    if not os.path.exists('./data/videos'):
        os.mkdir('./data/videos')

    with open('./data/video_list.csv') as f:
        reader = csv.DictReader(f, quotechar='"', delimiter=',', quoting=csv.QUOTE_ALL, skipinitialspace=True)
        for line in tqdm(reader):
            title, url = line['title'].strip(), line['url'].strip()
            download_video_from_youtube(title, url)