Ma/video-subtitle-remover
1
0
Fork 0
mirror of https://github.com/YaoFANGUK/video-subtitle-remover.git synced 2026-02-04 04:34:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

非gui运行则关闭预览，减少计算

Browse Source
This commit is contained in:
YaoFANGUK
2023-12-26 17:27:53 +08:00
parent 935c341c32
commit f7ce8013e0
2 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
backend/main.py
View File

@@ -448,12 +448,14 @@ class SubtitleDetect:
class SubtitleRemover:
def __init__(self, vd_path, sub_area=None):
def __init__(self, vd_path, sub_area=None, gui_mode=False):
importlib.reload(config)
# 线程锁
self.lock = threading.RLock()
# 用户指定的字幕区域位置
self.sub_area = sub_area
# 是否为gui运行，gui运行需要显示预览
self.gui_mode = gui_mode
# 判断是否为图片
self.is_picture = False
if str(vd_path).endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
@@ -626,7 +628,8 @@ class SubtitleRemover:
print(
f'write frame: {start_frame_no + inner_index} with mask {sub_list[index]}')
inner_index += 1
self.preview_frame = cv2.hconcat([batch[i], inpainted_frame])
if self.gui_mode:
self.preview_frame = cv2.hconcat([batch[i], inpainted_frame])
self.update_progress(tbar, increment=len(batch))
# *********************** 批推理方案 end ***********************
@@ -676,7 +679,8 @@ class SubtitleRemover:
self.video_writer.write(inpainted_frame)
print(f'write frame: {start_frame_index + inner_index} with mask {sub_list[start_frame_index]}')
inner_index += 1
self.preview_frame = cv2.hconcat([batch[i], inpainted_frame])
if self.gui_mode:
self.preview_frame = cv2.hconcat([batch[i], inpainted_frame])
self.update_progress(tbar, increment=len(batch))
def lama_mode(self, sub_list, tbar):
@@ -696,7 +700,8 @@ class SubtitleRemover:
frame = cv2.inpaint(frame, mask, 3, cv2.INPAINT_TELEA)
else:
frame = self.lama_inpaint(frame, mask)
self.preview_frame = cv2.hconcat([original_frame, frame])
if self.gui_mode:
self.preview_frame = cv2.hconcat([original_frame, frame])
if self.is_picture:
cv2.imencode(self.ext, frame)[1].tofile(self.video_out_name)
else:
@@ -725,7 +730,8 @@ class SubtitleRemover:
original_frame = cv2.imread(self.video_path)
mask = create_mask(original_frame.shape[0:2], sub_list[1])
inpainted_frame = self.lama_inpaint(original_frame, mask)
self.preview_frame = cv2.hconcat([original_frame, inpainted_frame])
if self.gui_mode:
self.preview_frame = cv2.hconcat([original_frame, inpainted_frame])
cv2.imencode(self.ext, inpainted_frame)[1].tofile(self.video_out_name)
tbar.update(1)
self.progress_total = 100

3
gui.py
View File

@@ -16,6 +16,7 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import backend.main
class SubtitleRemoverGUI:
def __init__(self):
@@ -272,7 +273,7 @@ class SubtitleRemoverGUI:
def task():
while self.video_paths:
video_path = self.video_paths.pop()
self.sr = backend.main.SubtitleRemover(video_path, subtitle_area)
self.sr = backend.main.SubtitleRemover(video_path, subtitle_area, True)
self.sr.run()
Thread(target=task, daemon=True).start()
self.video_cap.release()