2626from backend .config import *
2727from backend .tools .hardware_accelerator import HardwareAccelerator
2828from backend .tools import reformat
29+
2930from backend .tools .ocr import OcrRecogniser , get_coordinates
3031from backend .tools import subtitle_ocr
3132from backend .tools .paddle_model_config import PaddleModelConfig
3839import time
3940import pysrt
4041
42+
4143class SubtitleExtractor :
4244 """
4345 视频字幕提取类
@@ -110,16 +112,18 @@ def run(self):
110112 self .update_progress (ocr = 0 , frame_extract = 0 )
111113 self .append_output ('-----------------------------' )
112114 # 打印识别语言与识别模式
113- self .append_output (f" { tr ['Main' ]['RecSubLang' ]} :{ config .language .value } | { tr ['Main' ]['RecMode' ]} :{ config .mode .value } " )
115+ self .append_output (
116+ f" { tr ['Main' ]['RecSubLang' ]} :{ config .language .value } | { tr ['Main' ]['RecMode' ]} :{ config .mode .value } " )
114117 # 如果使用GPU加速,则打印GPU加速提示
115118 if self .hardware_accelerator .has_accelerator ():
116119 self .append_output (f" { tr ['Main' ]['AcceleratorON' ].format (self .hardware_accelerator .accelerator_name )} " )
117120
118121 # 打印视频帧数与帧率
119122 self .append_output (f" { tr ['Main' ]['FrameCount' ]} :{ self .frame_count } "
120- f" | { tr ['Main' ]['FrameRate' ]} :{ self .fps } " )
123+ f" | { tr ['Main' ]['FrameRate' ]} :{ self .fps } " )
121124 # 打印加载模型信息
122- self .append_output (f" DET: { os .path .basename (self .model_config .DET_MODEL_PATH )} | REC: { os .path .basename (self .model_config .REC_MODEL_PATH )} " )
125+ self .append_output (
126+ f" DET: { os .path .basename (self .model_config .DET_MODEL_PATH )} | REC: { os .path .basename (self .model_config .REC_MODEL_PATH )} " )
123127 self .append_output ('-----------------------------' )
124128 # 打印视频帧提取开始提示
125129 self .append_output (tr ['Main' ]['StartProcessFrame' ])
@@ -197,31 +201,31 @@ def capture_frame_with_subtitle_area(self):
197201 # 确保输出目录存在
198202 if not os .path .exists (self .temp_output_dir ):
199203 os .makedirs (self .temp_output_dir )
200-
204+
201205 # 确保视频已打开
202206 if not self .video_cap .isOpened ():
203207 self .video_cap = cv2 .VideoCapture (self .video_path )
204-
208+
205209 # 将视频指针设置到第一帧
206210 # self.video_cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
207-
211+
208212 # 读取第一帧
209213 ret , frame = self .video_cap .read ()
210-
214+
211215 if ret :
212216 # 如果有字幕区域,绘制矩形
213217 sub_area = self .sub_area
214218 if sub_area is not None :
215219 # 绘制绿色矩形框
216220 cv2 .rectangle (frame , (sub_area .xmin , sub_area .ymin ), (sub_area .xmax , sub_area .ymax ), (0 , 255 , 0 ), 2 )
217221 # 添加文字标注
218- cv2 .putText (frame , "Subtitle Area" , (sub_area .xmin , sub_area .ymin - 10 ),
219- cv2 .FONT_HERSHEY_SIMPLEX , 0.9 , (0 , 255 , 0 ), 2 )
220-
222+ cv2 .putText (frame , "Subtitle Area" , (sub_area .xmin , sub_area .ymin - 10 ),
223+ cv2 .FONT_HERSHEY_SIMPLEX , 0.9 , (0 , 255 , 0 ), 2 )
224+
221225 # 保存图像
222226 output_path = os .path .join (self .temp_output_dir , 'sub_area.jpg' )
223227 cv2 .imwrite (output_path , frame )
224-
228+
225229 # 重置视频指针到第一帧
226230 self .video_cap .set (cv2 .CAP_PROP_POS_FRAMES , 0 )
227231
@@ -308,7 +312,8 @@ def extract_frame_by_det(self):
308312 dt_box , rec_res = self .ocr .predict (frame )
309313 area_text1 = "" .join (self .__get_area_text ((dt_box , rec_res )))
310314 if start_frame_no not in compare_ocr_result_cache .keys ():
311- compare_ocr_result_cache [current_frame_no ] = {'text' : area_text1 , 'dt_box' : dt_box , 'rec_res' : rec_res }
315+ compare_ocr_result_cache [current_frame_no ] = {'text' : area_text1 , 'dt_box' : dt_box ,
316+ 'rec_res' : rec_res }
312317 frame_lru_list .append ((frame , current_frame_no ))
313318 ocr_args_list .append ((self .frame_count , current_frame_no ))
314319 # 缓存头帧
@@ -327,7 +332,8 @@ def extract_frame_by_det(self):
327332 # 如果在找结束帧的时候
328333 if is_finding_end_frame_no :
329334 # 判断该帧与头帧ocr内容是否一致,若不一致则找到尾,尾巴为前一帧
330- if not self ._compare_ocr_result (compare_ocr_result_cache , None , start_frame_no , frame , current_frame_no ):
335+ if not self ._compare_ocr_result (compare_ocr_result_cache , None , start_frame_no , frame ,
336+ current_frame_no ):
331337 is_finding_end_frame_no = False
332338 is_finding_start_frame_no = True
333339 end_frame_no = current_frame_no - 1
@@ -349,7 +355,7 @@ def extract_frame_by_det(self):
349355 frame_lru_list .pop (0 )
350356
351357 # if len(start_end_frame_no) > 0:
352- # self.append_output(start_end_frame_no)
358+ # self.append_output(start_end_frame_no)
353359
354360 while len (ocr_args_list ) > 1 :
355361 total_frame_count , ocr_info_frame_no = ocr_args_list .pop (0 )
@@ -384,6 +390,7 @@ def extract_frame_by_vsf(self):
384390 if self .video_cap :
385391 self .video_cap .release ()
386392 self .video_cap = None
393+
387394 def count_process ():
388395 duration_ms = (self .frame_count / self .fps ) * 1000
389396 last_total_ms = 0
@@ -481,10 +488,11 @@ def vsf_output(out, ):
481488 # 计算进度
482489 try :
483490 self .vsf_running = True
484- Thread (target = count_process , daemon = True ).start ()
491+ Thread (target = count_process , daemon = True ).start ()
485492 # 已知BUG: test_chinese_cht.flv在net drive上会导致无法停止, 但在本地不会, 可能是vsf的原因
486493 p = subprocess .Popen (cmd , stdout = subprocess .PIPE , stderr = subprocess .PIPE , bufsize = 1 ,
487- close_fds = 'posix' in sys .builtin_module_names , shell = False , creationflags = subprocess .CREATE_NEW_PROCESS_GROUP )
494+ close_fds = 'posix' in sys .builtin_module_names , shell = False ,
495+ creationflags = subprocess .CREATE_NEW_PROCESS_GROUP )
488496 ProcessManager .instance ().add_process (p )
489497 self .manage_process (p .pid )
490498 p .wait ()
@@ -500,14 +508,15 @@ def vsf_output(out, ):
500508 self .vsf_running = True
501509 try :
502510 p = subprocess .Popen (cmd , stdout = subprocess .PIPE , stderr = subprocess .PIPE , bufsize = 1 ,
503- close_fds = 'posix' in sys .builtin_module_names , shell = True ,
504- start_new_session = True )
511+ close_fds = 'posix' in sys .builtin_module_names , shell = True ,
512+ start_new_session = True )
505513 Thread (target = vsf_output , daemon = True , args = (p .stderr ,)).start ()
506514 ProcessManager .instance ().add_process (p )
507515 self .manage_process (p .pid )
508516 p .wait ()
509517 finally :
510518 self .vsf_running = False
519+
511520 def filter_watermark (self ):
512521 """
513522 去除原始字幕文本中的水印区域的文本
@@ -773,7 +782,9 @@ def _remove_duplicate_subtitle(self):
773782 while idx_j < content_list_len :
774783 # 计算当前行与下一行的Levenshtein距离
775784 # 判决idx_j的下一帧是否与idx_i不同,若不同(或者是最后一帧)则找到结束帧
776- if idx_j + 1 == content_list_len or ratio (i .content .replace (' ' , '' ), content_list [idx_j + 1 ].content .replace (' ' , '' )) < (config .thresholdTextSimilarity .value / 100.0 ):
785+ if idx_j + 1 == content_list_len or ratio (i .content .replace (' ' , '' ),
786+ content_list [idx_j + 1 ].content .replace (' ' , '' )) < (
787+ config .thresholdTextSimilarity .value / 100.0 ):
777788 # 若找到终点帧,定义字幕结束帧帧号
778789 end_frame = content_list [idx_j ].no
779790 if not self .use_vsf :
@@ -836,18 +847,21 @@ def _unite_coordinates(self, coordinates_list):
836847 indexed = sorted (enumerate (coordinates_list ), key = lambda x : x [1 ][0 ])
837848 # parent数组用于并查集
838849 parent = list (range (n ))
850+
839851 def find (i ):
840852 while parent [i ] != i :
841853 parent [i ] = parent [parent [i ]]
842854 i = parent [i ]
843855 return i
856+
844857 def union (i , j ):
845858 ri , rj = find (i ), find (j )
846859 if ri != rj :
847860 # 保留较小索引的坐标作为代表
848861 if ri > rj :
849862 ri , rj = rj , ri
850863 parent [rj ] = ri
864+
851865 # 滑动窗口:xmin已排序,只要xmin差值超过容忍度就移动左边界
852866 left = 0
853867 for right in range (n ):
@@ -1001,16 +1015,18 @@ def get_ocr_progress():
10011015 # self.append_output(f'recv total_ms:{total_ms}')
10021016 if current_frame_no == - 1 :
10031017 return
1018+
10041019 options = {
10051020 'REC_CHAR_TYPE' : config .language .value ,
10061021 'DROP_SCORE' : config .dropScore .value / 100.0 ,
10071022 'SUB_AREA_DEVIATION_RATE' : config .subtitleAreaDeviationRate .value / 100.0 ,
10081023 'DEBUG_OCR_LOSS' : config .debugOcrLoss .value ,
10091024 'HARDWARD_ACCELERATOR' : self .hardware_accelerator ,
10101025 }
1011- process , task_queue , progress_queue = subtitle_ocr .async_start (self .video_path , self .raw_subtitle_path , self .sub_area , options )
1026+ process , task_queue , progress_queue = subtitle_ocr .async_start (self .video_path , self .raw_subtitle_path ,
1027+ self .sub_area , options )
10121028 ProcessManager .instance ().add_process (process )
1013- self .manage_process (process . pid )
1029+ self .manage_process (getattr ( process , ' pid' , None ) )
10141030 self .subtitle_ocr_task_queue = task_queue
10151031 self .subtitle_ocr_progress_queue = progress_queue
10161032 # 开启线程负责更新OCR进度
@@ -1035,23 +1051,23 @@ def append_output(self, *args):
10351051 def add_progress_listener (self , listener ):
10361052 """
10371053 添加进度监听器
1038-
1054+
10391055 Args:
10401056 listener: 一个回调函数,接收参数 (progress_ocr, progress_frame_extract, progress_total, isFinished)
10411057 """
10421058 if listener not in self .progress_listeners :
10431059 self .progress_listeners .append (listener )
1044-
1060+
10451061 def remove_progress_listener (self , listener ):
10461062 """
10471063 移除进度监听器
1048-
1064+
10491065 Args:
10501066 listener: 要移除的监听器函数
10511067 """
10521068 if listener in self .progress_listeners :
10531069 self .progress_listeners .remove (listener )
1054-
1070+
10551071 def notify_progress_listeners (self ):
10561072 """
10571073 通知所有进度监听器当前进度
@@ -1065,6 +1081,7 @@ def notify_progress_listeners(self):
10651081 def manage_process (pid ):
10661082 pass
10671083
1084+
10681085if __name__ == '__main__' :
10691086 multiprocessing .set_start_method ("spawn" )
10701087 # 提示用户输入视频路径
0 commit comments