利用python+whisper生成视频字幕文件

文章目录

      • 前言
      • 1.本地环境
      • 2.安装所需要的库
      • 3.导入相关库
      • 4.获取指定路径下的所有视频文件
      • 5.导入模型进行音频识别
      • 6.将识别结果转换为srt字幕文件
      • 7.完成代码

        前言

        最近在看一些网课,由于没有字幕看着非常费劲,需要全神贯注的去听。网上很多生成字幕的网站都需要收费,想用某映但是它的智能字幕不允许上传大于两小时的视频。

        基于这个问题就想着用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()