From c0e2d9da263f0675eeb661fe664425a48c9b16c1 Mon Sep 17 00:00:00 2001 From: Bugale Bugalit Date: Mon, 2 Sep 2024 00:48:31 +0300 Subject: [PATCH] fix: handle irregular aspect ratios properly --- buganime/transcode.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/buganime/transcode.py b/buganime/transcode.py index fa9d758..3bc7f27 100644 --- a/buganime/transcode.py +++ b/buganime/transcode.py @@ -54,6 +54,12 @@ def __init__(self, input_path: str, output_path: str, height_out: int, width_out self.__video_info = video_info self.__height_out = height_out self.__width_out = width_out + self.__upscale_height_out = self.__height_out + self.__upscale_width_out = self.__width_out + if self.__video_info.width / self.__video_info.height > self.__width_out / self.__height_out: + self.__upscale_height_out = round(self.__video_info.height * self.__width_out / self.__video_info.width) + else: + self.__upscale_width_out = round(self.__video_info.width * self.__height_out / self.__video_info.height) model = Transcoder.Module(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4) if torch.cuda.is_available(): model.load_state_dict(torch.load(MODEL_PATH)['params'], strict=True) @@ -88,13 +94,8 @@ async def __write_output_frames(self, frames: AsyncIterator[bytes]) -> None: os.link(self.__input_path, os.path.join(temp_dir, 'input.mkv')) else: shutil.copy(self.__input_path, os.path.join(temp_dir, 'input.mkv')) - width_out = self.__width_out - height_out = self.__height_out - if self.__video_info.width / self.__video_info.height > self.__width_out / self.__height_out: - height_out = round(self.__video_info.height * self.__width_out / self.__video_info.width) - else: - width_out = round(self.__video_info.width * self.__height_out / self.__video_info.height) - args = ('-f', 'rawvideo', '-framerate', str(self.__video_info.fps), '-pix_fmt', 'rgb24', '-s', f'{width_out}x{height_out}', + args = ('-f', 'rawvideo', '-framerate', str(self.__video_info.fps), '-pix_fmt', 'rgb24', + '-s', f'{self.__upscale_width_out}x{self.__upscale_height_out}', '-i', 'pipe:', '-i', 'input.mkv', '-map', '0', '-map', f'1:{self.__video_info.audio_index}', '-vf', f'subtitles=input.mkv:si={self.__video_info.subtitle_index}, pad={self.__width_out}:{self.__height_out}:(ow-iw)/2:(oh-ih)/2:black', @@ -136,7 +137,7 @@ async def __upscale_frame(self, frame: bytes) -> bytes: async with self.__gpu_lock: frame_cpu = await asyncio.to_thread(self.__gpu_upscale, frame_arr) return cast(bytes, await asyncio.to_thread( - lambda: cv2.resize(frame_cpu.numpy(), (self.__width_out, self.__height_out), interpolation=cv2.INTER_LANCZOS4).tobytes())) + lambda: cv2.resize(frame_cpu.numpy(), (self.__upscale_width_out, self.__upscale_height_out), interpolation=cv2.INTER_LANCZOS4).tobytes())) async def __generate_upscaling_tasks(self) -> None: assert self.__frame_tasks_queue