diff --git a/.github/workflows/check-code.yml b/.github/workflows/check-code.yml index fbbf420..42f8f67 100644 --- a/.github/workflows/check-code.yml +++ b/.github/workflows/check-code.yml @@ -50,4 +50,4 @@ jobs: - name: Install test dependencies run: python -m pip install -r dev-requirements.txt - name: Test - run: pytest tests + run: pytest tests --full-trace diff --git a/.vscode/settings.json b/.vscode/settings.json index e624a3a..e8d0021 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,13 +1,6 @@ { - "python.linting.enabled": true, - "python.linting.flake8Enabled": true, - "python.linting.flake8Args": [], - "python.linting.mypyEnabled": true, - "python.linting.mypyArgs": ["--follow-imports=silent"], - "python.linting.pylintEnabled": true, - "python.linting.pylintArgs": ["--load-plugins", "pylint_pytest"], - "python.testing.pytestEnabled": true, - "files.associations": { - "*.yaml": "home-assistant" - } + "flake8.args": [], + "mypy-type-checker.args": ["--follow-imports=silent"], + "pylint.args": ["--load-plugins", "pylint_pytest"], + "python.testing.pytestEnabled": true } diff --git a/buganime/buganime.py b/buganime/buganime.py index 8d5fea0..737b132 100644 --- a/buganime/buganime.py +++ b/buganime/buganime.py @@ -44,9 +44,12 @@ class Movie: def parse_streams(streams: Any) -> transcode.VideoInfo: def _get_video_stream() -> Any: - for stream in streams: + video_streams = [stream for stream in streams if stream['codec_type'] == 'video'] + if len(video_streams) == 1: + return video_streams[0] + for stream in video_streams: match stream: - case {'codec_type': 'video', 'disposition': {'default': 1}}: + case {'disposition': {'default': 1}}: return stream raise RuntimeError('No default video stream found') @@ -79,15 +82,16 @@ def _get_subtitle_stream_index() -> int: video = _get_video_stream() return transcode.VideoInfo(audio_index=_get_audio_stream()['index'], subtitle_index=_get_subtitle_stream_index(), width=video['width'], height=video['height'], fps=video['r_frame_rate'], - frames=int(video['tags'].get('NUMBER_OF_FRAMES') or video['tags'].get('NUMBER_OF_FRAMES-eng') or 0)) + frames=int(video.get('tags', {}).get('NUMBER_OF_FRAMES') or video.get('tags', {}).get('NUMBER_OF_FRAMES-eng') or 0)) def parse_filename(input_path: str) -> TVShow | Movie: # Remove metadata in brackets/parentheses and extension (e.g. hash, resolution, etc.) - input_path = input_path.replace('_', ' ') + input_path = re.sub(r'[_+]', ' ', input_path) input_path = re.sub(r'\[[^\]]*\]', '', input_path) input_path = re.sub(r'\([^\)]*\)', '', input_path) input_path = re.sub(r'\d{3,4}p[ -][^\\]*', '', input_path) + input_path = re.sub(r'[ -]*\\[ -]*', r'\\', input_path) input_path = os.path.splitext(input_path)[0].strip(' -') # Remove extension and directories @@ -103,13 +107,13 @@ def parse_filename(input_path: str) -> TVShow | Movie: return TVShow(name=match.group('name'), season=int(match.group('season')), episode=int(match.group('episode'))) # Other standalone TV Shows - if match := re.match(r'^(?P.+?)[ -]+(?:S(?:eason ?)?(?P\d{1,2})[ -]*)?E?(?P\d{1,3})(?:v\d+)?$', input_name): + if match := re.match(r'^(?P.+?)[ -]* (?:S(?:eason ?)?(?P\d{1,2})[ -]*)?E?(?P\d{1,3})(?:v\d+)?(?:[ -].*)?$', input_name): return TVShow(name=match.group('name'), season=int(match.group('season') or '1'), episode=int(match.group('episode'))) # Structured TV Shows - dir_re = r'(?P[^\\]+?)[ -]+S(?:eason ?)?(?P\d{1,2})[ -][^\\]*' - file_re = r'[^\\]*S\d{1,2}?E(?P\d{1,3})(?:[ -][^\\]*)?' - if match := re.match(fr'^.*\\{dir_re}\\{file_re}$', input_path): + dir_re = r'(?P[^\\]+?)[ -]+S(?:eason ?)?\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'))) return Movie(name=input_name) @@ -137,20 +141,20 @@ def process_file(input_path: str) -> None: logging.info('ffprobe %s wrote %s, %s', str(proc.args), proc.stderr, proc.stdout) video_info = parse_streams(json.loads(proc.stdout)['streams']) - try: - with lock_mutex(name=UPSCALE_MUTEX_NAME): - logging.info('Running Upscaler') - asyncio.run(transcode.Transcoder(input_path=input_path, output_path=output_path, height_out=2160, video_info=video_info).run()) - logging.info('Upscaler for %s finished', input_path) - except Exception: - logging.exception('Failed to convert %s', input_path) + with lock_mutex(name=UPSCALE_MUTEX_NAME): + logging.info('Running Upscaler') + asyncio.run(transcode.Transcoder(input_path=input_path, output_path=output_path, height_out=2160, width_out=3840, video_info=video_info).run()) + logging.info('Upscaler for %s finished', input_path) def process_path(input_path: str) -> None: if os.path.isdir(input_path): for root, _, files in os.walk(input_path): for file in files: - process_file(input_path=os.path.join(root, file)) + try: + process_file(input_path=os.path.join(root, file)) + except Exception: + logging.exception('Failed to convert %s', input_path) else: process_file(input_path=input_path) diff --git a/buganime/externals/Anime4KCPP_CLI/Anime4KCPPCore.dll b/buganime/externals/Anime4KCPP_CLI/Anime4KCPPCore.dll deleted file mode 100644 index 4cd0583..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/Anime4KCPPCore.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/Anime4KCPP_CLI.exe b/buganime/externals/Anime4KCPP_CLI/Anime4KCPP_CLI.exe deleted file mode 100644 index 9edc90a..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/Anime4KCPP_CLI.exe and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/concrt140.dll b/buganime/externals/Anime4KCPP_CLI/concrt140.dll deleted file mode 100644 index e1fad94..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/concrt140.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/msvcp140.dll b/buganime/externals/Anime4KCPP_CLI/msvcp140.dll deleted file mode 100644 index 94cbba7..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/msvcp140.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/opencv_videoio_ffmpeg450_64.dll b/buganime/externals/Anime4KCPP_CLI/opencv_videoio_ffmpeg450_64.dll deleted file mode 100644 index 95bae6f..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/opencv_videoio_ffmpeg450_64.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/opencv_videoio_msmf450_64.dll b/buganime/externals/Anime4KCPP_CLI/opencv_videoio_msmf450_64.dll deleted file mode 100644 index d02cbcc..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/opencv_videoio_msmf450_64.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/opencv_world450.dll b/buganime/externals/Anime4KCPP_CLI/opencv_world450.dll deleted file mode 100644 index 4643f8a..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/opencv_world450.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/openh264-1.8.0-win64.dll b/buganime/externals/Anime4KCPP_CLI/openh264-1.8.0-win64.dll deleted file mode 100644 index 99fb7b2..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/openh264-1.8.0-win64.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/vcruntime140.dll b/buganime/externals/Anime4KCPP_CLI/vcruntime140.dll deleted file mode 100644 index f59e67e..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/vcruntime140.dll and /dev/null differ diff --git a/buganime/externals/Anime4KCPP_CLI/vcruntime140_1.dll b/buganime/externals/Anime4KCPP_CLI/vcruntime140_1.dll deleted file mode 100644 index f564204..0000000 Binary files a/buganime/externals/Anime4KCPP_CLI/vcruntime140_1.dll and /dev/null differ diff --git a/buganime/transcode.py b/buganime/transcode.py index a50a925..fa9d758 100644 --- a/buganime/transcode.py +++ b/buganime/transcode.py @@ -3,6 +3,8 @@ import tempfile import asyncio import logging +import warnings +import shutil from dataclasses import dataclass from typing import AsyncIterator, cast, Optional @@ -44,17 +46,21 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor: tensor = body(tensor) return cast(torch.Tensor, self.__upsampler(tensor) + base) - def __init__(self, input_path: str, output_path: str, height_out: int, video_info: VideoInfo) -> None: + def __init__(self, input_path: str, output_path: str, height_out: int, width_out: int, video_info: VideoInfo) -> None: if not os.path.isfile(MODEL_PATH): with open(MODEL_PATH, 'wb') as file: file.write(requests.get(MODEL_URL, timeout=600).content) self.__input_path, self.__output_path = input_path, output_path self.__video_info = video_info self.__height_out = height_out - self.__width_out = round(self.__video_info.width * self.__height_out / self.__video_info.height) + self.__width_out = width_out model = Transcoder.Module(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4) - model.load_state_dict(torch.load(MODEL_PATH)['params'], strict=True) - self.__model = model.eval().cuda().half() + if torch.cuda.is_available(): + model.load_state_dict(torch.load(MODEL_PATH)['params'], strict=True) + self.__model = model.eval().cuda().half() + else: + model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu'))['params'], strict=True) + self.__model = model.eval() self.__gpu_lock: Optional[asyncio.Lock] = None self.__frame_tasks_queue: Optional[asyncio.Queue[Optional[asyncio.Task[bytes]]]] = None @@ -78,10 +84,20 @@ async def __read_input_frames(self) -> AsyncIterator[bytes]: async def __write_output_frames(self, frames: AsyncIterator[bytes]) -> None: with tempfile.TemporaryDirectory() as temp_dir: - os.link(self.__input_path, os.path.join(temp_dir, 'input.mkv')) - args = ('-f', 'rawvideo', '-framerate', str(self.__video_info.fps), '-pix_fmt', 'rgb24', '-s', f'{self.__width_out}x{self.__height_out}', + if os.path.splitdrive(self.__input_path)[0] == os.path.splitdrive(temp_dir)[0]: + 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}', '-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}', + '-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', *FFMPEG_OUTPUT_ARGS, self.__output_path, '-loglevel', 'warning', '-y') proc = await asyncio.subprocess.create_subprocess_exec('ffmpeg', *args, stdin=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, @@ -103,7 +119,10 @@ async def __write_output_frames(self, frames: AsyncIterator[bytes]) -> None: @retry.retry(RuntimeError, tries=10, delay=1) def __gpu_upscale(self, frame: torch.Tensor) -> torch.Tensor: with torch.no_grad(): - frame_float = frame.cuda().permute(2, 0, 1).half() / 255 + if torch.cuda.is_available(): + frame_float = frame.cuda().permute(2, 0, 1).half() / 255 + else: + frame_float = frame.permute(2, 0, 1) / 255 frame_upscaled_float = self.__model(frame_float.unsqueeze(0)).data.squeeze().clamp_(0, 1) return cast(torch.Tensor, (frame_upscaled_float * 255.0).round().byte().permute(1, 2, 0).cpu()) @@ -111,7 +130,8 @@ async def __upscale_frame(self, frame: bytes) -> bytes: if self.__video_info.height == self.__height_out: return frame with torch.no_grad(): - frame_arr = torch.frombuffer(frame, dtype=torch.uint8).reshape([self.__video_info.height, self.__video_info.width, 3]) + with warnings.catch_warnings(action='ignore'): + frame_arr = torch.frombuffer(frame, dtype=torch.uint8).reshape([self.__video_info.height, self.__video_info.width, 3]) assert self.__gpu_lock async with self.__gpu_lock: frame_cpu = await asyncio.to_thread(self.__gpu_upscale, frame_arr) diff --git a/dev-requirements.txt b/dev-requirements.txt index 0469855..0ca5412 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -7,4 +7,5 @@ mypy tqdm-stubs types-retry types-requests +torch<2.3.0 # Due to https://github.com/pytorch/pytorch/issues/124897 -r requirements.txt \ No newline at end of file diff --git a/tests/data/0.mkv b/tests/data/0.mkv index 70308e5..945fc29 100644 Binary files a/tests/data/0.mkv and b/tests/data/0.mkv differ diff --git a/tests/data/1.mkv b/tests/data/1.mkv new file mode 100644 index 0000000..3c2b161 Binary files /dev/null and b/tests/data/1.mkv differ diff --git a/tests/data/2.mkv b/tests/data/2.mkv new file mode 100644 index 0000000..9935d87 Binary files /dev/null and b/tests/data/2.mkv differ diff --git a/tests/data/7.json b/tests/data/7.json new file mode 100644 index 0000000..dd615ad --- /dev/null +++ b/tests/data/7.json @@ -0,0 +1,346 @@ +{ + "streams": [ + { + "index": 0, + "codec_name": "hevc", + "codec_long_name": "H.265 / HEVC (High Efficiency Video Coding)", + "profile": "Main", + "codec_type": "video", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "width": 1920, + "height": 1080, + "coded_width": 1920, + "coded_height": 1080, + "closed_captions": 0, + "film_grain": 0, + "has_b_frames": 4, + "sample_aspect_ratio": "1:1", + "display_aspect_ratio": "16:9", + "pix_fmt": "yuv420p", + "level": 150, + "color_range": "tv", + "color_space": "bt709", + "color_transfer": "bt709", + "color_primaries": "bt709", + "chroma_location": "left", + "refs": 1, + "r_frame_rate": "24000/1001", + "avg_frame_rate": "24000/1001", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "extradata_size": 116, + "disposition": { + "default": 1, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + } + }, + { + "index": 1, + "codec_name": "aac", + "codec_long_name": "AAC (Advanced Audio Coding)", + "profile": "LC", + "codec_type": "audio", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "sample_fmt": "fltp", + "sample_rate": "48000", + "channels": 2, + "channel_layout": "stereo", + "bits_per_sample": 0, + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "extradata_size": 2, + "disposition": { + "default": 1, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "language": "jpn", + "title": "Japanese" + } + }, + { + "index": 2, + "codec_name": "ass", + "codec_long_name": "ASS (Advanced SSA) subtitle", + "codec_type": "subtitle", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "duration_ts": 1453920, + "duration": "1453.920000", + "extradata_size": 610, + "disposition": { + "default": 1, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "language": "eng", + "title": "English" + } + }, + { + "index": 3, + "codec_name": "ass", + "codec_long_name": "ASS (Advanced SSA) subtitle", + "codec_type": "subtitle", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "duration_ts": 1453920, + "duration": "1453.920000", + "extradata_size": 610, + "disposition": { + "default": 0, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "language": "ind", + "title": "Indonesian" + } + }, + { + "index": 4, + "codec_name": "ass", + "codec_long_name": "ASS (Advanced SSA) subtitle", + "codec_type": "subtitle", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "duration_ts": 1453920, + "duration": "1453.920000", + "extradata_size": 610, + "disposition": { + "default": 0, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "language": "tha", + "title": "Thai" + } + }, + { + "index": 5, + "codec_name": "ass", + "codec_long_name": "ASS (Advanced SSA) subtitle", + "codec_type": "subtitle", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "duration_ts": 1453920, + "duration": "1453.920000", + "extradata_size": 610, + "disposition": { + "default": 0, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "language": "vie", + "title": "Vietnamese" + } + }, + { + "index": 6, + "codec_name": "ass", + "codec_long_name": "ASS (Advanced SSA) subtitle", + "codec_type": "subtitle", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/1000", + "start_pts": 0, + "start_time": "0.000000", + "duration_ts": 1453920, + "duration": "1453.920000", + "extradata_size": 610, + "disposition": { + "default": 0, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "language": "chi", + "title": "Chinese (Simplified)" + } + }, + { + "index": 7, + "codec_type": "attachment", + "codec_tag_string": "[0][0][0][0]", + "codec_tag": "0x0000", + "r_frame_rate": "0/0", + "avg_frame_rate": "0/0", + "time_base": "1/90000", + "start_pts": 0, + "start_time": "0.000000", + "duration_ts": 130852800, + "duration": "1453.920000", + "extradata_size": 313724, + "disposition": { + "default": 0, + "dub": 0, + "original": 0, + "comment": 0, + "lyrics": 0, + "karaoke": 0, + "forced": 0, + "hearing_impaired": 0, + "visual_impaired": 0, + "clean_effects": 0, + "attached_pic": 0, + "timed_thumbnails": 0, + "captions": 0, + "descriptions": 0, + "metadata": 0, + "dependent": 0, + "still_image": 0 + }, + "tags": { + "filename": "NotoSans-Medium.ttf", + "mimetype": "font/ttf" + } + } + ], + "format": { + "filename": "C:\\Temp\\Ooi! Tonbo - S01E01 - 1080p WEB HEVC -NanDesuKa (B-Global).mkv", + "nb_streams": 8, + "nb_programs": 0, + "format_name": "matroska,webm", + "format_long_name": "Matroska / WebM", + "start_time": "0.000000", + "duration": "1453.920000", + "size": "306731820", + "bit_rate": "1687750", + "probe_score": 100, + "tags": { + "encoder": "\"...\"", + "MOVIE/ENCODER": "Lavf58.76.100" + } + } +} \ No newline at end of file diff --git a/tests/test_buganime.py b/tests/test_buganime.py index 3eb22a3..4722af5 100644 --- a/tests/test_buganime.py +++ b/tests/test_buganime.py @@ -1,6 +1,13 @@ import os import tempfile import json +import subprocess +import functools +import typing + +import cv2 +import pytest + from buganime import buganime, transcode NAME_CONVERSIONS = [ @@ -40,8 +47,8 @@ (r'C:\Kaguya-sama - Love is War - S01E06.mkv', buganime.TVShow(name='Kaguya-sama - Love is War', season=1, episode=6)), - (r'C:\Kaguya-sama wa Kokurasetai S03 1080p Dual Audio WEBRip AAC x265-EMBER\S03E01-Miko Iino Wants to Be Soothed Kaguya ' - r'Doesn’t Realize Chika Fujiwara Wants to Battle [8933E8C9].mkv', + (r'C:\Kaguya-sama wa Kokurasetai S03 1080p Dual Audio WEBRip AAC x265-EMBER\S03E01-Miko Iino Wants to Be Soothed Kaguya Doesn’t Realize Chika Fujiwara ' + r'Wants to Battle [8933E8C9].mkv', buganime.TVShow(name='Kaguya-sama wa Kokurasetai', season=3, episode=1)), (r'C:\Kaguya-sama wa Kokurasetai S2 - OVA - 1080p WEB H.264 -NanDesuKa (B-Global).mkv', @@ -52,9 +59,25 @@ (r'C:\Watashi no Shiawase na Kekkon - S01E01 - MULTi.mkv', buganime.TVShow(name='Watashi no Shiawase na Kekkon', season=1, episode=1)), + + (r'C:\Monogatari Series\15. Zoku Owarimonogatari\Zoku Owarimonogatari 01 - Koyomi Reverse, Part 1.mkv', + buganime.TVShow(name='Zoku Owarimonogatari', season=1, episode=1)), + + (r'C:\SNAFU S01-S03+OVA 1080p Dual Audio BDRip 10 bits DD x265-EMBER\SNAFU S02+OVA 1080p Dual Audio BDRip 10 bits DD x265-EMBER\Series\S02E01-Nobody Knows' + r'Why They Came to the Service Club [7CE95AC0].mkv', + buganime.TVShow(name='SNAFU', season=2, episode=1)), + + (r'C:\Temp\Torrents\SNAFU S01-S03+OVA 1080p Dual Audio BDRip 10 bits DD x265-EMBER\SNAFU S02+OVA 1080p Dual Audio BDRip 10 bits DD x265-EMBER\OVA\S02E14 ' + r'[OVA]-Undoubtedly, Girls Are Made of Sugar, Spice, and Everything Nice [7E9E8A1F].mkv', + buganime.TVShow(name='SNAFU', season=2, episode=14)), ] +@pytest.mark.parametrize('path,result', NAME_CONVERSIONS) +def test_parse_filename(path: str, result: buganime.TVShow | buganime.Movie) -> None: + assert buganime.parse_filename(path) == result + + STREAM_CONVERSIONS = [ ('0.json', transcode.VideoInfo(audio_index=1, subtitle_index=1, width=1920, height=1080, fps='24000/1001', frames=34094)), ('1.json', transcode.VideoInfo(audio_index=1, subtitle_index=3, width=1920, height=1080, fps='24000/1001', frames=34095)), @@ -63,21 +86,58 @@ ('4.json', transcode.VideoInfo(audio_index=2, subtitle_index=1, width=1920, height=1080, fps='24000/1001', frames=34047)), ('5.json', transcode.VideoInfo(audio_index=2, subtitle_index=1, width=1920, height=1080, fps='24000/1001', frames=35638)), ('6.json', transcode.VideoInfo(audio_index=1, subtitle_index=0, width=1920, height=1080, fps='30/1', frames=7425)), + ('7.json', transcode.VideoInfo(audio_index=1, subtitle_index=0, width=1920, height=1080, fps='24000/1001', frames=0)), ] -def test_parse_filename() -> None: - for path, result in NAME_CONVERSIONS: - assert buganime.parse_filename(path) == result +@pytest.mark.parametrize('filename,result', STREAM_CONVERSIONS) +def test_parse_streams(filename: str, result: transcode.VideoInfo) -> None: + with open(os.path.join(os.path.dirname(__file__), 'data', filename), 'rb') as file: + assert buganime.parse_streams(json.loads(file.read())['streams']) == result -def test_parse_streams() -> None: - for filename, result in STREAM_CONVERSIONS: - with open(os.path.join(os.path.dirname(__file__), 'data', filename), 'rb') as file: - assert buganime.parse_streams(json.loads(file.read())['streams']) == result +def _check_side_bars(frame: typing.Any, bar_size: int) -> None: + assert max(cv2.mean(frame[0:, :bar_size])[:3]) < 1 + assert max(cv2.mean(frame[0:, -bar_size:])[:3]) < 1 + assert min(cv2.mean(frame[:1, bar_size:-bar_size])[:3]) > 254 + assert min(cv2.mean(frame[-1:, bar_size:-bar_size])[:3]) > 254 + + +def _check_top_bottom_bars(frame: typing.Any, bar_size: int) -> None: + assert max(cv2.mean(frame[:bar_size])[:3]) < 1 + assert max(cv2.mean(frame[-bar_size:])[:3]) < 1 + assert min(cv2.mean(frame[bar_size:-bar_size, :1])[:3]) > 254 + assert min(cv2.mean(frame[bar_size:-bar_size, -1:])[:3]) > 254 + + +VIDEO_TESTS = [ + ('0.mkv', '24000/1001', None), + + # 1900x1080 -> 3840x2160, validate black bars on left/right + ('1.mkv', '24000/1001', functools.partial(_check_side_bars, bar_size=20)), + + # 1940x1080 -> 3840x2160, validate black bars on top/bottom + ('2.mkv', '24000/1001', functools.partial(_check_top_bottom_bars, bar_size=11)), +] -def test_sanity() -> None: +@pytest.mark.parametrize('filename,fps,check_func', VIDEO_TESTS) +def test_transcode(filename: str, fps: str, check_func: typing.Callable[[typing.Any], None] | None) -> None: with tempfile.TemporaryDirectory() as tempdir: buganime.OUTPUT_DIR = tempdir - buganime.process_file(os.path.join(os.path.dirname(__file__), 'data', '0.mkv')) + output_path = os.path.join(tempdir, 'Movies', filename) + buganime.process_file(os.path.join(os.path.dirname(__file__), 'data', filename)) + assert os.path.isfile(output_path) + stream = json.loads(subprocess.run(['ffprobe', '-show_format', '-show_streams', '-of', 'json', output_path], text=True, capture_output=True, + check=True, encoding='utf-8').stdout)['streams'][0] + assert stream['codec_name'] == 'hevc' + assert stream['width'] == 3840 + assert stream['height'] == 2160 + assert stream['r_frame_rate'] == fps + if check_func is not None: + video = cv2.VideoCapture(output_path) + try: + frame = video.read()[1] + check_func(frame) + finally: + video.release()