From 4c45e6a74cd5982a6f226165cc01046ec58680c0 Mon Sep 17 00:00:00 2001
From: Evilmass <evil0mass@gmail.com>
Date: Sun, 5 Mar 2023 14:55:32 +0800
Subject: [PATCH] multithreading event

---
 denoise_audio.py  | 21 ++++++++++++++++----
 download_video.py | 50 ++++++++++++++++++++++++++++++-----------------
 video2audio.py    | 27 ++++++++++++++++++++-----
 3 files changed, 71 insertions(+), 27 deletions(-)

diff --git a/denoise_audio.py b/denoise_audio.py
index 757e926..2f42382 100644
--- a/denoise_audio.py
+++ b/denoise_audio.py
@@ -1,18 +1,31 @@
 import os
+from concurrent.futures import ProcessPoolExecutor
+
 import torchaudio
+
+WORKERS = 3  # 6G显存+16G内存，请根据机器配置调整
 raw_audio_dir = "./raw_audio/"
 denoise_audio_dir = "./denoised_audio/"
 filelist = list(os.walk(raw_audio_dir))[0][2]
 
-for file in filelist:
+
+def denoise_audio(file):
     if file.endswith(".wav"):
         os.system(f"demucs --two-stems=vocals {raw_audio_dir}{file}")
-for file in filelist:
+
     file = file.replace(".wav", "")
-    wav, sr = torchaudio.load(f"./separated/htdemucs/{file}/vocals.wav", frame_offset=0, num_frames=-1, normalize=True,
+    wav, sr = torchaudio.load(f"./separated/htdemucs/{file}/vocals.wav",
+                              frame_offset=0,
+                              num_frames=-1,
+                              normalize=True,
                               channels_first=True)
     # merge two channels into one
     wav = wav.mean(dim=0).unsqueeze(0)
     if sr != 22050:
         wav = torchaudio.transforms.Resample(orig_freq=sr, new_freq=22050)(wav)
-    torchaudio.save(denoise_audio_dir + file + ".wav", wav, 22050, channels_first=True)
\ No newline at end of file
+    torchaudio.save(denoise_audio_dir + file + ".wav", wav, 22050, channels_first=True)
+
+
+if __name__ == "__main__":
+    with ProcessPoolExecutor(max_workers=WORKERS) as executor:
+        executor.map(denoise_audio, filelist)
diff --git a/download_video.py b/download_video.py
index 7d99b95..9d0b7ac 100644
--- a/download_video.py
+++ b/download_video.py
@@ -1,23 +1,37 @@
-from google.colab import files
-import shutil
 import os
+import random
+import shutil
+from concurrent.futures import ThreadPoolExecutor
+from google.colab import files
+
 basepath = os.getcwd()
-uploaded = files.upload() # 上传文件
+uploaded = files.upload()  # 上传文件
 for filename in uploaded.keys():
-    assert(filename.endswith(".txt")), "speaker-videolink info could only be .txt file!"
+    assert (filename.endswith(".txt")), "speaker-videolink info could only be .txt file!"
     shutil.move(os.path.join(basepath, filename), os.path.join("./speaker_links.txt"))
 
-with open("./speaker_links.txt", 'r', encoding='utf-8') as f:
-    lines = f.readlines()
-speakers = []
-for line in lines:
-    line = line.replace("\n", "").replace(" ", "")
-    if line == "":
-        continue
-    speaker, link = line.split("|")
-    if speaker not in speakers:
-        speakers.append(speaker)
-    # download link
-    import random
-    filename = speaker + "_" + str(random.randint(0, 1000000))
-    os.system(f"youtube-dl -f 0 {link} -o ./video_data/{filename}.mp4")
\ No newline at end of file
+
+def generate_infos():
+    infos = []
+    with open("./speaker_links.txt", 'r', encoding='utf-8') as f:
+        lines = f.readlines()
+    for line in lines:
+        line = line.replace("\n", "").replace(" ", "")
+        if line == "":
+            continue
+        speaker, link = line.split("|")
+        filename = speaker + "_" + str(random.randint(0, 1000000))
+        infos.append({"link": link, "filename": filename})
+    return infos
+
+
+def download_video(info):
+    link = info["link"]
+    filename = info["filename"]
+    os.system(f"youtube-dl -f 0 {link} -o ./video_data/{filename}.mp4")
+
+
+if __name__ == "__main__":
+    infos = generate_infos()
+    with ThreadPoolExecutor(max_workers=os.cpu_count() - 2) as executor:
+        executor.map(download_video, infos)
\ No newline at end of file
diff --git a/video2audio.py b/video2audio.py
index 411c861..e847929 100644
--- a/video2audio.py
+++ b/video2audio.py
@@ -1,10 +1,27 @@
-from moviepy.editor import AudioFileClip
 import os
+from concurrent.futures import ThreadPoolExecutor
+
+from moviepy.editor import AudioFileClip
+
 video_dir = "./video_data/"
 audio_dir = "./raw_audio/"
 filelist = list(os.walk(video_dir))[0][2]
-for file in filelist:
-    if file.endswith(".mp4"):
-        my_audio_clip = AudioFileClip(video_dir + file)
-        my_audio_clip.write_audiofile(audio_dir + file.rstrip(".mp4") + ".wav")
 
+
+def generate_infos():
+    videos = []
+    for file in filelist:
+        if file.endswith(".mp4"):
+            videos.append(file)
+    return videos
+
+
+def clip_file(file):
+    my_audio_clip = AudioFileClip(video_dir + file)
+    my_audio_clip.write_audiofile(audio_dir + file.rstrip(".mp4") + ".wav")
+
+
+if __name__ == "__main__":
+    infos = generate_infos()
+    with ThreadPoolExecutor(max_workers=os.cpu_count() - 2) as executor:
+        executor.map(clip_file, infos)