From 29c5317a6967c213bd4894ea43b2d779d61ccb0a Mon Sep 17 00:00:00 2001
From: YaoFANGUK <flavioy@live.com>
Date: Wed, 13 Dec 2023 19:50:15 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8Dbug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/inpaint/lama_inpaint.py |  8 +++++++-
 backend/main.py                 | 15 +++++++--------
 backend/tools/merge_video.py    |  2 ++
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/backend/inpaint/lama_inpaint.py b/backend/inpaint/lama_inpaint.py
index e39a3bd..2be5325 100644
--- a/backend/inpaint/lama_inpaint.py
+++ b/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
 
diff --git a/backend/main.py b/backend/main.py
index bf89fcd..f914a56 100644
--- a/backend/main.py
+++ b/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:
diff --git a/backend/tools/merge_video.py b/backend/tools/merge_video.py
index 955e631..45a588c 100644
--- a/backend/tools/merge_video.py
+++ b/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