-
Notifications
You must be signed in to change notification settings - Fork 0
/
transcribe.py
151 lines (123 loc) · 4.91 KB
/
transcribe.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
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import argparse
import csv
import os
import shutil
import warnings
import scenedetect
import transformers
from moviepy.editor import VideoFileClip
from scenedetect import ContentDetector, SceneManager
from scenedetect.scene_manager import save_images
from scenedetect.video_splitter import split_video_ffmpeg
transformers.logging.set_verbosity_error()
warnings.filterwarnings("ignore", category=FutureWarning)
def video_to_audio(input, output):
video_clip = VideoFileClip(input)
audio_clip = video_clip.audio
audio_clip.write_audiofile(output, verbose=False, logger=None)
audio_clip.close()
video_clip.close()
def init_folder(output_folder):
file = os.path.join(output_folder, "scenes.csv")
if not os.path.exists(output_folder):
os.makedirs(output_folder)
os.makedirs(os.path.join(output_folder, "videos"))
with open(file, "w") as f:
csv_writer = csv.writer(f)
csv_writer.writerow(["video", "screen_nr", "scene", "image", "start_time", "end_time", "transcription"])
return []
with open(file, "r") as f:
csv_reader = csv.reader(f)
files = [row[0] for row in csv_reader]
return files
def scene_generator(input_folder, input_file, output_folder):
input_path = os.path.join(input_folder, input_file)
scene_manager = SceneManager()
scene_manager.add_detector(ContentDetector())
video = scenedetect.open_video(input_path)
scene_manager.detect_scenes(video)
scene_list = scene_manager.get_scene_list(start_in_scene=True)
file_name, file_ext = os.path.splitext(input_file)
videos_dir = os.path.join(output_folder, "videos")
video_dir = os.path.join(videos_dir, file_name)
scenes_dir = os.path.join(video_dir, "scenes")
audio_dir = os.path.join(video_dir, "audio")
images_dir = os.path.join(video_dir, "images")
if os.path.exists(video_dir):
shutil.rmtree(video_dir)
for create_dir in [video_dir, scenes_dir, images_dir, audio_dir]:
os.makedirs(create_dir)
save_images(
scene_list=scene_list,
video=video,
num_images=1,
output_dir=images_dir,
image_name_template="$SCENE_NUMBER",
)
split_video_ffmpeg(
input_video_path=input_path,
scene_list=scene_list,
output_dir=scenes_dir,
output_file_template=f"$SCENE_NUMBER{file_ext}",
)
for i, (start_time, end_time) in enumerate(scene_list):
scene_nr = str(i + 1).zfill(3)
scene_file = f"{scene_nr}{file_ext}"
image_file = f"{scene_nr}.jpg"
audio_file = f"{scene_nr}.mp3"
scene_path = os.path.join(scenes_dir, scene_file)
image_path = os.path.join(images_dir, image_file)
audio_path = os.path.join(audio_dir, audio_file)
video_to_audio(scene_path, audio_path)
scene = {
"video": input_file,
"scene_nr": scene_nr,
"scene": scene_path,
"image": image_path,
"audio": audio_path,
"start_time": start_time,
"end_time": end_time,
}
yield scene
def main(output_folder, input_folder, input_files):
n = len(input_files)
print(f"Transcribing {n} files")
transcriber = transformers.pipeline("automatic-speech-recognition", "distil-whisper/distil-large-v3")
output_csv = os.path.join(output_folder, "scenes.csv")
with open(output_csv, "a") as f:
csv_writer = csv.writer(f)
for i, input_file in enumerate(input_files):
print(f"- {input_file}")
# separating the write, because if one scene has been written to CSV,
# we skip the whole video if the script is run again
rows = []
for scene in scene_generator(input_folder, input_file, output_folder):
res = transcriber(scene["audio"])
scene["transcription"] = res.get("text", "NO TRANSCRIPTION FOUND")
rows.append(
[
scene["video"],
scene["scene_nr"],
scene["scene"],
scene["image"],
scene["start_time"],
scene["end_time"],
scene["transcription"],
]
)
for row in rows:
csv_writer.writerow(row)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Convert videos to audio")
parser.add_argument("folder", type=str, help="Folder containing video files")
parser.add_argument("--output", type=str, help="Output folder", default="transcribed_scenes")
args = parser.parse_args()
done = init_folder(args.output)
todo = []
for file in os.listdir(args.folder):
if not file.endswith((".mp4", ".avi", ".mov", ".mkv")):
continue
if file in done:
continue
todo.append(file)
main(args.output, args.folder, todo)