-
Notifications
You must be signed in to change notification settings - Fork 5
/
gensrt.py
62 lines (44 loc) · 1.82 KB
/
gensrt.py
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
#!/usr/bin/env python3
import datetime
import json
import sys
from pathlib import Path
import srt
def generate_srt(json_input_file=None):
# get the list of image files in the img folder, sorted by their names
with open(json_input_file, "r") as f:
ocr_dict: dict = json.load(f)
subtitles = []
start_time = datetime.timedelta()
sorted_int_keys = sorted([int(k) for k in ocr_dict.keys()])
current_sub: srt.Subtitle = None
for frame_number in sorted(ocr_dict.keys()):
start_time: datetime.timedelta = datetime.timedelta(milliseconds=int(1000 * (float(frame_number) / 4 )))
end_time = start_time + datetime.timedelta(milliseconds=500)
body: str = ocr_dict.get(str(frame_number)).strip()
if body:
start_time: datetime.timedelta = datetime.timedelta(seconds=frame_number)
end_time = start_time + datetime.timedelta(milliseconds=1000)
sub = srt.Subtitle(None, start_time, end_time, body)
if not current_sub:
current_sub = sub
continue
# if it's duplicate content then add 1 second to current sub
if current_sub.content == body:
current_sub.end = current_sub.end + datetime.timedelta(milliseconds=1000)
else:
subtitles.append(current_sub)
print(current_sub.to_srt())
current_sub = sub
else:
if current_sub:
subtitles.append(current_sub)
print(current_sub.to_srt())
current_sub = None
return subtitles
json_input = sys.argv[1]
srt_output=sys.argv[2]
subtitles = generate_srt(json_input_file=json_input)
print('JSON input:', json_input)
print('SRT output:', srt_output)
Path(srt_output).write_text(srt.compose(subtitles), encoding='utf-8')