From 245ceffcdceab0032ef400594d3e8be68ec15625 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 | 16 ++++++++-------- tests/test_buganime.py | 3 +++ 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/buganime/transcode.py b/buganime/transcode.py index fa9d758..79ad5f1 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,7 @@ 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 +136,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 diff --git a/tests/test_buganime.py b/tests/test_buganime.py index 49b8082..923311d 100644 --- a/tests/test_buganime.py +++ b/tests/test_buganime.py @@ -74,6 +74,9 @@ (r'C:\Temp\Torrents\Mushoku Tensei S01+SP 1080p Dual Audio BDRip 10 bits DDP x265-EMBER\Mushoku Tensei S01P01 1080p Dual Audio BDRip 10 bits DD ' r'x265-EMBER\S01E08-Turning Point 1 V2 [87C2150F].mkv', buganime.TVShow(name='Mushoku Tensei', season=1, episode=8)), + + (r'C:\Temp\Torrents\Mushoku Tensei S02P01+SP 1080p Dual Audio BDRip 10 bits DD+ x265-EMBER\S02E01-The Brokenhearted Mage [AFBB9792].mkv', + buganime.TVShow(name='Mushoku Tensei', season=2, episode=1)), ]