FakeVD/code_test/VGGish_Feature_Extractor/my_vggish_folder.py

import os
import torch
import numpy as np
from moviepy.editor import VideoFileClip
import pickle

from vggish_modified import VGGish as VGGish_modified  # 修改后的VGGish代码

# 定义输入视频文件夹和输出特征文件路径
video_folder_path = '/mnt/data10t/dazuoye/GROUP2024-GEN6/FakeSV/dataset/videos'
feature_path = '/mnt/data10t/dazuoye/GROUP2024-GEN6/FakeSV/code/VGGish_Feature_Extractor/my_dict_vid_audioconvfea.pkl'

# 加载模型
urls = {
    'vggish': 'https://github.com/harritaylor/torchvggish/releases/download/v0.1/vggish-10086976.pth',
    'pca': 'https://github.com/harritaylor/torchvggish/releases/download/v0.1/vggish_pca_params-970ea276.pth'
}
vggish_model = VGGish_modified(urls, pretrained=True)

# 初始化保存特征的字典
features_dict = {}

audio_file_path = '/mnt/data10t/dazuoye/GROUP2024-GEN6/FakeSV/dataset/1.wav'

# 遍历视频文件夹中的所有视频文件
for video_file_name in os.listdir(video_folder_path):
    video_file_path = os.path.join(video_folder_path, video_file_name)
    
    # 提取视频文件名（不包括扩展名）作为视频ID
    video_id = os.path.splitext(video_file_name)[0]
    # video_id = video_file_name.split('_')[1].split('.')[0]

    try:
        # 从视频中提取音频
        video = VideoFileClip(video_file_path)
        video.audio.write_audiofile(audio_file_path)

        # 提取特征
        features = vggish_model(audio_file_path)

        # 保存特征到字典中
        features_dict[video_id] = features.cpu().detach().numpy()

    except Exception as e:
        print(f"Error processing {video_file_name}: {e}")

# 保存特征字典到文件
with open(feature_path, 'wb') as f:
    pickle.dump(features_dict, f)

print(f"Audio features have been saved to {feature_path}")