利用python+whisper生成视频字幕文件 python利用openai开源模型whisper批量生成字幕文件

前言

最近在看一些网课，由于没有字幕看着非常费劲，需要全神贯注的去听。网上很多生成字幕的网站都需要收费，想用某映但是它的智能字幕不允许上传大于两小时的视频。
基于这个问题就想着用openai开源的whisper来试试，最终整体的效果还行，硬件不行识别的有点慢，准确率不算高，但是配合音频基本能理解是什么意思，主要看视频更加轻松了。
注：由于我有很多视频，所以才用python自己写脚本批量处理，如不需要或者觉得写脚本麻烦可以看看WhisperDesktop，它识别更快一点，而且资源占用更低

1.本地环境

GPU：GTX 1650 4G
Cuda：10.1
Python：3.8.0
Pytorch：1.7.1

2.安装所需要的库

在已有的python环境上安装openai-whisper、ffmpeg和zhconv，其中zhconv是用来进行简繁体的。

pip install openai-whisper
pip install ffmpeg-python
pip install zhconv

3.导入相关库

import whisper
import os
import datetime,time
from zhconv import convert # 简繁体转换
from tqdm import tqdm
import imageio # 用来获取视频时长

4.获取指定路径下的所有视频文件

# 获取mp4文件列表
def find_files(path,suffix):
	"""
 用来获取path下的所有suffix格式文件
 @params:
 path - Required : 目标路径 (str)
 suffix - Required : 视频文件格式 (str)
 """
 mp4_files = []
 for root, dirs, files in os.walk(path):
 for file in files:
 if file.endswith('.'+suffix):
 mp4_files.append(os.path.abspath(os.path.join(root, file)))
 return mp4_files
# 主文件夹
file_path = r'E:\视频'
mp4_files = find_files(file_path,suffix='mp4')

5.导入模型进行音频识别

whisper默认将模型权重下载到当前用户的.cache/whisper路径下，如果想使用其他路径可以通过参数download_root进行设置。

# 获取模型
model = whisper.load_model('small')

接下批量进行识别，生成srt格式的字幕文件

for file in tqdm(mp4_files):
 # 字幕文件保存路径
 # xxx.mp4 --> xxx. + srt
 # 如果是其他格式，如mpweg需要改一下，这里因为都是mp4就直接对字符串切片了
 save_file = file[:-3] + "srt"
 # 判断文件是否存在，存在则说明已经有字幕，跳出不识别
 if os.path.exists(save_file):
 time.sleep(0.01)
 continue
 # 获取当前视频识别开始时间
 start_time = datetime.datetime.now()
 print('正在识别：{} --{}'.format('\\'.join(file.split('\\')[2:]),start_time.strftime('%Y-%m-%d %H:%M:%S')))
 # 获取视频时长
 video = imageio.get_reader(file)
 duration = seconds_to_hmsm(video.get_meta_data()['duration'])
 video.close()
 print('视频时长：{}'.format(duration))
 # 文字识别
 res = model.transcribe(file,fp16=False,language='Chinese')

6.将识别结果转换为srt字幕文件

srt字幕文件内容格式如下，第一行是顺序，第二行是开始时间到结束时间，第三行是字幕文字。

1
00:00:00,000 --> 00:00:03,399
第一句
2
00:00:04,160 --> 00:00:06,560
第二句

由于whisper的识别结果中，其时间是秒数，不是srt的hh:mm:ss:mm时间格式，所有需要进行转换。

# 秒转时分秒毫秒
def seconds_to_hmsm(seconds):
	"""
 输入一个秒数，输出为H:M:S:M时间格式
 @params:
 seconds - Required : 秒 (float)
 """
 hours = str(int(seconds // 3600))
 minutes = str(int((seconds % 3600) // 60))
 seconds = seconds % 60
 milliseconds = str(int(int((seconds - int(seconds)) * 1000))) # 毫秒留三位
 seconds = str(int(seconds))
 # 补0
 if len(hours) < 2:
 hours = '0' + hours
 if len(minutes) < 2:
 minutes = '0' + minutes
 if len(seconds) < 2:
 seconds = '0' + seconds
 if len(milliseconds) < 3:
 milliseconds = '0'*(3-len(milliseconds)) + milliseconds
 return f"{hours}:{minutes}:{seconds},{milliseconds}"

将结果转为srt字幕文件：

# 写入字幕文件
 with open(save_file,'w',encoding='utf-8') as f:
 i = 1
 for r in res['segments']:
 f.write(str(i)+'\n')
 f.write(seconds_to_hmsm(float(r['start']))+' --> '+seconds_to_hmsm(float(r['end']))+'\n')
 i += 1
 f.write(convert(r['text'], 'zh-cn')+'\n') # 结果可能是繁体，转为简体zh-cn
 f.write('\n')

7.完成代码

import whisper
import os
import datetime,time
from zhconv import convert # 简繁体转换
from tqdm import tqdm
import imageio # 用来获取视频时长
# 获取mp4文件列表
def find_files(path,suffix):
	"""
 用来获取path下的所有suffix格式文件
 @params:
 path - Required : 目标路径 (str)
 suffix - Required : 视频文件格式 (str)
 """
 mp4_files = []
 for root, dirs, files in os.walk(path):
 for file in files:
 if file.endswith('.'+suffix):
 mp4_files.append(os.path.abspath(os.path.join(root, file)))
 return mp4_files
# 秒转时分秒毫秒
def seconds_to_hmsm(seconds):
	"""
 输入一个秒数，输出为H:M:S:M时间格式
 @params:
 seconds - Required : 秒 (float)
 """
 hours = str(int(seconds // 3600))
 minutes = str(int((seconds % 3600) // 60))
 seconds = seconds % 60
 milliseconds = str(int(int((seconds - int(seconds)) * 1000))) # 毫秒留三位
 seconds = str(int(seconds))
 # 补0
 if len(hours) < 2:
 hours = '0' + hours
 if len(minutes) < 2:
 minutes = '0' + minutes
 if len(seconds) < 2:
 seconds = '0' + seconds
 if len(milliseconds) < 3:
 milliseconds = '0'*(3-len(milliseconds)) + milliseconds
 return f"{hours}:{minutes}:{seconds},{milliseconds}"
def main():
	# 主文件夹
	file_path = r'E:\视频'
	mp4_files = find_files(file_path,suffix='mp4')
	
	# 获取模型
	model = whisper.load_model('small')
	for file in tqdm(mp4_files):
	 # 字幕文件保存路径
	 # xxx.mp4 --> xxx. + srt
	 # 如果是其他格式，如mpweg需要改一下，这里因为都是mp4就直接对字符串切片了
	 save_file = file[:-3] + "srt"
	 # 判断文件是否存在，存在则说明已经有字幕，跳出不识别
	 if os.path.exists(save_file):
	 time.sleep(0.01)
	 continue
	 # 获取当前视频识别开始时间
	 start_time = datetime.datetime.now()
	 print('正在识别：{} --{}'.format('\\'.join(file.split('\\')[2:]),start_time.strftime('%Y-%m-%d %H:%M:%S')))
	 # 获取视频时长
	 video = imageio.get_reader(file)
	 duration = seconds_to_hmsm(video.get_meta_data()['duration'])
	 video.close()
	 print('视频时长：{}'.format(duration))
	
	 # 文字识别
	 res = model.transcribe(file,fp16=False,language='Chinese')
	
	# 写入字幕文件
	 with open(save_file,'w',encoding='utf-8') as f:
	 i = 1
	 for r in res['segments']:
	 f.write(str(i)+'\n')
	 f.write(seconds_to_hmsm(float(r['start']))+' --> '+seconds_to_hmsm(float(r['end']))+'\n')
	 i += 1
	 f.write(convert(r['text'], 'zh-cn')+'\n') # 结果可能是繁体，转为简体zh-cn
	 f.write('\n')
	 # 获取当前视频识别结束时间
	 end_time = datetime.datetime.now()
	 print('完成识别：{} --{}'.format('\\'.join(file.split('\\')[2:]),end_time.strftime('%Y-%m-%d %H:%M:%S')))
	 print('花费时间:',end_time-start_time)
if __name__ == "__main__":
	main()