diff --git a/buganime/buganime.py b/buganime/buganime.py index 1d0557e..dd465f3 100644 --- a/buganime/buganime.py +++ b/buganime/buganime.py @@ -110,7 +110,7 @@ def parse_filename(input_path: str) -> TVShow | Movie: return TVShow(name=match.group('name'), season=int(match.group('season')), episode=int(match.group('episode'))) # Structured TV Shows - dir_re = r'(?P[^\\]+?)[ -]+S(?:eason ?)?\d{1,2}(?:[ -][^\\]*)?' + dir_re = r'(?P[^\\]+?)[ -]+S(?:eason ?)?\d{1,2}(?:P\d{1,2})?(?:[ -][^\\]*)?' file_re = r'[^\\]*S(?P\d{1,2})E(?P\d{1,3})(?:[ -][^\\]*)?' if match := re.match(fr'^.*\\{dir_re}(?:\\.*)?\\{file_re}$', input_path): return TVShow(name=match.group('name'), season=int(match.group('season')), episode=int(match.group('episode'))) 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 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)), ]