修复bug

backend/inpaint/lama_inpaint.py

@@ -1,6 +1,7 @@
 import os
 from typing import Union
 
+import cv2
 import torch
 import numpy as np
 from PIL import Image
@@ -18,10 +19,15 @@ class LamaInpaint:
         self.device = device
 
     def __call__(self, image: Union[Image.Image, np.ndarray], mask: Union[Image.Image, np.ndarray]):
+        if isinstance(image, np.ndarray):
+            orig_height, orig_width = image.shape[:2]
+        else:
+            orig_height, orig_width = np.array(image).shape[:2]
         image, mask = prepare_img_and_mask(image, mask, self.device)
         with torch.inference_mode():
             inpainted = self.model(image, mask)
             cur_res = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
-            cur_res = np.clip(cur_res * 255, 0, 255).astype(np.uint8)
+            cur_res = np.clip(cur_res * 255, 0, 255).astype('uint8')
+            cur_res = cur_res[:orig_height, :orig_width]
             return cur_res
 

backend/main.py

@@ -525,15 +525,12 @@ class SubtitleRemover:
         self.progress_total = 50 + self.progress_remover
 
     def run(self):
-
         # 记录开始时间
         start_time = time.time()
-        # 寻找字幕帧
+        # 重置进度条
         self.progress_total = 0
+        # 寻找字幕帧
         sub_list = self.sub_detector.find_subtitle_frame_no(sub_remover=self)
-        # 测试代码
-        # from test1_dict_raw import test1_raw
-        # sub_list = self.sub_detector.unify_regions(test1_raw)
         continuous_frame_no_list = self.sub_detector.find_continuous_ranges_with_same_mask(sub_list)
         # 获取场景分割的帧号
         scene_div_points = self.sub_detector.get_scene_div_frame_no(self.video_path)
@@ -546,9 +543,11 @@ class SubtitleRemover:
             self.lama_inpaint = LamaInpaint()
             original_frame = cv2.imread(self.video_path)
             mask = create_mask(original_frame.shape[0:2], sub_list[1])
-            frame = self.lama_inpaint(original_frame, mask)
-            cv2.imencode(self.ext, frame)[1].tofile(self.video_out_name)
-            self.preview_frame = cv2.hconcat([original_frame, frame])
+            inpainted_frame = self.lama_inpaint(original_frame, mask)
+            print(original_frame.shape)
+            print(inpainted_frame.shape)
+            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
         else:

backend/tools/merge_video.py

@@ -31,3 +31,5 @@ if __name__ == '__main__':
     merge_video(v0_path, v1_path, v2_path, video_out_path)
     # ffmpeg 命令 mp4转gif
     # ffmpeg -i demo3.mp4 -vf "scale=w=720:h=-1,fps=15,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse" -loop 0 -r 15 -f gif output.gif
+    # 宽度固定400，高度成比例：
+    # ffmpeg - i input.avi -vf scale=400:-2