From 37efdf60d379e1afff4d0737bd0d23532316457b Mon Sep 17 00:00:00 2001 From: sHa Date: Thu, 25 Dec 2025 05:13:56 +0000 Subject: [PATCH] feat: add frame class extraction and formatting for media files --- renamer/constants.py | 38 ++++++++++ renamer/extractor.py | 22 ++++-- renamer/extractors/filename_extractor.py | 27 ++++---- renamer/extractors/mediainfo_extractor.py | 50 ++++++++++---- renamer/formatters/media_formatter.py | 80 ++++++++++++++++------ renamer/formatters/resolution_formatter.py | 44 +++++++++++- renamer/test/test_filename_extractor.py | 24 +++++-- 7 files changed, 224 insertions(+), 61 deletions(-) diff --git a/renamer/constants.py b/renamer/constants.py index 8eec79b..7759964 100644 --- a/renamer/constants.py +++ b/renamer/constants.py @@ -31,4 +31,42 @@ SOURCE_DICT = { 'DVDRip': ['DVDRip', 'DVD-Rip', 'DVDRIP'], 'HDTV': ['HDTV'], 'BluRay': ['BluRay', 'BLURAY', 'Blu-ray'], +} + +FRAME_CLASSES = { + '480p': { + 'nominal_height': 480, + 'typical_widths': [640, 704, 720], + 'description': 'Standard Definition (SD) - DVD quality' + }, + '576p': { + 'nominal_height': 576, + 'typical_widths': [720, 768], + 'description': 'PAL Standard Definition (SD) - European DVD quality' + }, + '720p': { + 'nominal_height': 720, + 'typical_widths': [1280], + 'description': 'High Definition (HD) - 720p HD' + }, + '1080p': { + 'nominal_height': 1080, + 'typical_widths': [1920], + 'description': 'Full High Definition (FHD) - 1080p HD' + }, + '1440p': { + 'nominal_height': 1440, + 'typical_widths': [2560], + 'description': 'Quad High Definition (QHD) - 1440p 2K' + }, + '2160p': { + 'nominal_height': 2160, + 'typical_widths': [3840], + 'description': 'Ultra High Definition (UHD) - 2160p 4K' + }, + '4320p': { + 'nominal_height': 4320, + 'typical_widths': [7680], + 'description': 'Ultra High Definition (UHD) - 4320p 8K' + } } \ No newline at end of file diff --git a/renamer/extractor.py b/renamer/extractor.py index 5e55450..6afce80 100644 --- a/renamer/extractor.py +++ b/renamer/extractor.py @@ -27,14 +27,22 @@ class MediaExtractor: """Extract video source from filename""" return FilenameExtractor.extract_source(file_path) - def extract_resolution(self, file_path: Path) -> str | None: - """Extract resolution from media info or filename""" + def extract_frame_class(self, file_path: Path) -> str | None: + """Extract frame class from media info or filename""" # Try media info first - resolution = self.mediainfo_extractor.extract_resolution(file_path) - if resolution: - return resolution + frame_class = self.mediainfo_extractor.extract_frame_class(file_path) + if frame_class: + return frame_class # Fallback to filename - return FilenameExtractor.extract_resolution(file_path) + return FilenameExtractor.extract_frame_class(file_path) + + def extract_resolution(self, file_path: Path) -> str | None: + """Extract actual video resolution (WIDTHxHEIGHT) from media info""" + return self.mediainfo_extractor.extract_resolution(file_path) + + def extract_aspect_ratio(self, file_path: Path) -> str | None: + """Extract video aspect ratio from media info""" + return self.mediainfo_extractor.extract_aspect_ratio(file_path) def extract_hdr(self, file_path: Path) -> str | None: """Extract HDR info from media info""" @@ -54,7 +62,9 @@ class MediaExtractor: 'title': self.extract_title(file_path), 'year': self.extract_year(file_path), 'source': self.extract_source(file_path), + 'frame_class': self.extract_frame_class(file_path), 'resolution': self.extract_resolution(file_path), + 'aspect_ratio': self.extract_aspect_ratio(file_path), 'hdr': self.extract_hdr(file_path), 'audio_langs': self.extract_audio_langs(file_path), 'metadata': self.extract_metadata(file_path) diff --git a/renamer/extractors/filename_extractor.py b/renamer/extractors/filename_extractor.py index fcec976..da575a2 100644 --- a/renamer/extractors/filename_extractor.py +++ b/renamer/extractors/filename_extractor.py @@ -1,11 +1,19 @@ import re from pathlib import Path -from ..constants import SOURCE_DICT +from ..constants import SOURCE_DICT, FRAME_CLASSES class FilenameExtractor: """Class to extract information from filename""" + @staticmethod + def _get_frame_class_from_height(height: int) -> str: + """Get frame class from video height using FRAME_CLASSES constant""" + for frame_class, info in FRAME_CLASSES.items(): + if height == info['nominal_height']: + return frame_class + return 'Unclassified' + @staticmethod def extract_title(file_path: Path) -> str | None: """Extract movie title from filename""" @@ -40,20 +48,11 @@ class FilenameExtractor: return None @staticmethod - def extract_resolution(file_path: Path) -> str | None: - """Extract resolution from filename (e.g., 2160p, 1080p, 720p)""" + def extract_frame_class(file_path: Path) -> str | None: + """Extract frame class from filename (480p, 720p, 1080p, 2160p, etc.)""" file_name = file_path.name match = re.search(r'(\d{3,4})[pi]', file_name, re.IGNORECASE) if match: height = int(match.group(1)) - if height >= 2160: - return '2160p' - elif height >= 1080: - return '1080p' - elif height >= 720: - return '720p' - elif height >= 480: - return '480p' - else: - return f'{height}p' - return None \ No newline at end of file + return FilenameExtractor._get_frame_class_from_height(height) + return 'Unclassified' \ No newline at end of file diff --git a/renamer/extractors/mediainfo_extractor.py b/renamer/extractors/mediainfo_extractor.py index ceb2c3d..ace74cc 100644 --- a/renamer/extractors/mediainfo_extractor.py +++ b/renamer/extractors/mediainfo_extractor.py @@ -1,6 +1,7 @@ from pathlib import Path from pymediainfo import MediaInfo from collections import Counter +from ..constants import FRAME_CLASSES class MediaInfoExtractor: @@ -13,24 +14,49 @@ class MediaInfoExtractor: 'zh': 'chi', 'und': 'und' } - def extract_resolution(self, file_path: Path) -> str | None: - """Extract resolution from media info""" + def _get_frame_class_from_height(self, height: int) -> str: + """Get frame class from video height using FRAME_CLASSES constant""" + for frame_class, info in FRAME_CLASSES.items(): + if height == info['nominal_height']: + return frame_class + return 'Unclassified' + + def extract_frame_class(self, file_path: Path) -> str | None: + """Extract frame class from media info (480p, 720p, 1080p, etc.)""" try: media_info = MediaInfo.parse(file_path) video_tracks = [t for t in media_info.tracks if t.track_type == 'Video'] if video_tracks: height = getattr(video_tracks[0], 'height', None) if height: - if height >= 2160: - return '2160p' - elif height >= 1080: - return '1080p' - elif height >= 720: - return '720p' - elif height >= 480: - return '480p' - else: - return f'{height}p' + return self._get_frame_class_from_height(height) + except: + pass + return 'Unclassified' + + def extract_resolution(self, file_path: Path) -> str | None: + """Extract actual video resolution (WIDTHxHEIGHT) from media info""" + try: + media_info = MediaInfo.parse(file_path) + video_tracks = [t for t in media_info.tracks if t.track_type == 'Video'] + if video_tracks: + width = getattr(video_tracks[0], 'width', None) + height = getattr(video_tracks[0], 'height', None) + if width and height: + return f"{width}x{height}" + except: + pass + return None + + def extract_aspect_ratio(self, file_path: Path) -> str | None: + """Extract video aspect ratio from media info""" + try: + media_info = MediaInfo.parse(file_path) + video_tracks = [t for t in media_info.tracks if t.track_type == 'Video'] + if video_tracks: + aspect_ratio = getattr(video_tracks[0], 'display_aspect_ratio', None) + if aspect_ratio: + return str(aspect_ratio) except: pass return None diff --git a/renamer/formatters/media_formatter.py b/renamer/formatters/media_formatter.py index 09dffd9..e373932 100644 --- a/renamer/formatters/media_formatter.py +++ b/renamer/formatters/media_formatter.py @@ -9,34 +9,72 @@ from ..utils import detect_file_type class MediaFormatter: """Class to format media data for display""" - def format_file_info(self, file_path: Path, rename_data: dict) -> str: - """Format complete file information for display""" + def format_file_info_panel(self, file_path: Path, rename_data: dict) -> str: + """Format file information for the file info panel""" + output = [] + + # Panel title + output.append("[bold blue]FILE INFO[/bold blue]") + output.append("") # Empty line for spacing + # Get file stats size_full = SizeFormatter.format_size_full(file_path.stat().st_size) date_formatted = DateFormatter.format_modification_date(file_path.stat().st_mtime) - + # Get extension info ext_name = ExtensionExtractor.get_extension_name(file_path) ext_desc = ExtensionExtractor.get_extension_description(ext_name) meta_type, meta_desc = detect_file_type(file_path) match = ExtensionFormatter.check_extension_match(ext_name, meta_type) ext_info = ExtensionFormatter.format_extension_info(ext_name, ext_desc, meta_type, meta_desc, match) - + file_name = file_path.name - - # Build basic info - full_info = f"[bold blue]Path:[/bold blue] {str(file_path)}\n\n" - full_info += f"[bold green]Size:[/bold green] {size_full}\n" - full_info += f"[bold cyan]File:[/bold cyan] {file_name}\n" - full_info += f"{ext_info}\n" - full_info += f"[bold magenta]Modified:[/bold magenta] {date_formatted}" - - # Extra metadata - extra_text = self._format_extra_metadata(rename_data['metadata']) - if extra_text: - full_info += f"\n\n{extra_text}" - - return full_info + + output.append(f"[bold blue]Path:[/bold blue] {str(file_path)}") + output.append(f"[bold green]Size:[/bold green] {size_full}") + output.append(f"[bold cyan]File:[/bold cyan] {file_name}") + output.append(f"[bold magenta]Modified:[/bold magenta] {date_formatted}") + output.append(f"{ext_info}") + + return "\n".join(output) + + def format_filename_extraction_panel(self, rename_data: dict) -> str: + """Format filename extraction data for the filename panel""" + output = [] + output.append("[bold yellow]FILENAME EXTRACTION[/bold yellow]") + output.append("") # Empty line for spacing + output.append(f"[yellow]Title:[/yellow] {rename_data.get('title', 'Not found')}") + output.append(f"[yellow]Year:[/yellow] {rename_data.get('year', 'Not found')}") + output.append(f"[yellow]Source:[/yellow] {rename_data.get('source', 'Not found')}") + output.append(f"[yellow]Frame Class:[/yellow] {rename_data.get('frame_class', 'Not found')}") + return "\n".join(output) + + def format_metadata_extraction_panel(self, rename_data: dict) -> str: + """Format metadata extraction data for the metadata panel""" + output = [] + output.append("[bold cyan]METADATA EXTRACTION[/bold cyan]") + output.append("") # Empty line for spacing + metadata = rename_data.get('metadata', {}) + if metadata.get('duration'): + output.append(f"[cyan]Duration:[/cyan] {metadata['duration']:.1f} seconds") + if metadata.get('title'): + output.append(f"[cyan]Title:[/cyan] {metadata['title']}") + if metadata.get('artist'): + output.append(f"[cyan]Artist:[/cyan] {metadata['artist']}") + if not any(key in metadata for key in ['duration', 'title', 'artist']): + output.append("[dim]No metadata found[/dim]") + return "\n".join(output) if output else "[dim]No metadata found[/dim]" + + def format_mediainfo_extraction_panel(self, rename_data: dict) -> str: + """Format media info extraction data for the mediainfo panel""" + output = [] + output.append("[bold green]MEDIA INFO EXTRACTION[/bold green]") + output.append("") # Empty line for spacing + output.append(f"[green]Resolution:[/green] {rename_data.get('resolution', 'Not found')}") + output.append(f"[green]Aspect Ratio:[/green] {rename_data.get('aspect_ratio', 'Not found')}") + output.append(f"[green]HDR:[/green] {rename_data.get('hdr', 'Not found')}") + output.append(f"[green]Audio Languages:[/green] {rename_data.get('audio_langs', 'Not found')}") + return "\n".join(output) def format_proposed_name(self, rename_data: dict, ext_name: str) -> str: """Format the proposed filename""" @@ -49,8 +87,8 @@ class MediaFormatter: proposed_parts.append(rename_data['source']) tags = [] - if rename_data['resolution']: - tags.append(rename_data['resolution']) + if rename_data['frame_class'] and rename_data['frame_class'] != 'Unclassified': + tags.append(rename_data['frame_class']) if rename_data['hdr']: tags.append(rename_data['hdr']) if rename_data['audio_langs']: @@ -66,7 +104,9 @@ class MediaFormatter: lines.append(f"Movie title: {rename_data['title'] or 'Unknown'}") lines.append(f"Year: {rename_data['year'] or 'Unknown'}") lines.append(f"Video source: {rename_data['source'] or 'Unknown'}") + lines.append(f"Frame class: {rename_data['frame_class'] or 'Unknown'}") lines.append(f"Resolution: {rename_data['resolution'] or 'Unknown'}") + lines.append(f"Aspect ratio: {rename_data['aspect_ratio'] or 'Unknown'}") lines.append(f"HDR: {rename_data['hdr'] or 'No'}") lines.append(f"Audio langs: {rename_data['audio_langs'] or 'None'}") lines.append(f"Proposed filename: {proposed_name}") diff --git a/renamer/formatters/resolution_formatter.py b/renamer/formatters/resolution_formatter.py index 2861d43..adbc7b0 100644 --- a/renamer/formatters/resolution_formatter.py +++ b/renamer/formatters/resolution_formatter.py @@ -1,6 +1,44 @@ class ResolutionFormatter: - """Class for formatting video resolutions""" - + """Class for formatting video resolutions and frame classes""" + + @staticmethod + def get_frame_class_from_resolution(resolution: str) -> str: + """Convert resolution string (WIDTHxHEIGHT) to frame class (480p, 720p, etc.)""" + if not resolution: + return 'Unclassified' + + try: + # Extract height from WIDTHxHEIGHT format + if 'x' in resolution: + height = int(resolution.split('x')[1]) + else: + # Try to extract number directly + import re + match = re.search(r'(\d{3,4})', resolution) + if match: + height = int(match.group(1)) + else: + return 'Unclassified' + + if height == 4320: + return '4320p' + elif height >= 2160: + return '2160p' + elif height >= 1440: + return '1440p' + elif height >= 1080: + return '1080p' + elif height >= 720: + return '720p' + elif height >= 576: + return '576p' + elif height >= 480: + return '480p' + else: + return 'Unclassified' + except (ValueError, IndexError): + return 'Unclassified' + @staticmethod def format_resolution_p(height: int) -> str: """Format resolution as 2160p, 1080p, etc.""" @@ -14,7 +52,7 @@ class ResolutionFormatter: return '480p' else: return f'{height}p' - + @staticmethod def format_resolution_dimensions(width: int, height: int) -> str: """Format resolution as WIDTHxHEIGHT""" diff --git a/renamer/test/test_filename_extractor.py b/renamer/test/test_filename_extractor.py index ca9ff48..49e5be7 100644 --- a/renamer/test/test_filename_extractor.py +++ b/renamer/test/test_filename_extractor.py @@ -1,6 +1,7 @@ import pytest from pathlib import Path from ..extractors.filename_extractor import FilenameExtractor +from ..constants import FRAME_CLASSES def load_test_filenames(): @@ -17,6 +18,9 @@ def test_extract_title(filename): """Test title extraction from filename""" file_path = Path(filename) title = FilenameExtractor.extract_title(file_path) + # Print filename and extracted title clearly + print(f"\nFilename: \033[1;36m{filename}\033[0m") + print(f"Extracted title: \033[1;32m{title}\033[0m") # For now, just check it's not None and is string assert isinstance(title, str) or title is None @@ -39,15 +43,23 @@ def test_extract_source(filename): """Test source extraction from filename""" file_path = Path(filename) source = FilenameExtractor.extract_source(file_path) + # Print filename and extracted source clearly + print(f"\nFilename: \033[1;36m{filename}\033[0m") + print(f"Extracted source: \033[1;32m{source}\033[0m") # Source should be None or string assert isinstance(source, str) or source is None @pytest.mark.parametrize("filename", load_test_filenames()) -def test_extract_resolution(filename): - """Test resolution extraction from filename""" +def test_extract_frame_class(filename): + """Test frame class extraction from filename""" file_path = Path(filename) - resolution = FilenameExtractor.extract_resolution(file_path) - # Resolution should be None or string like '2160p' - if resolution: - assert 'p' in resolution or 'i' in resolution \ No newline at end of file + frame_class = FilenameExtractor.extract_frame_class(file_path) + # Print filename and extracted frame class clearly + print(f"\nFilename: \033[1;36m{filename}\033[0m") + print(f"Extracted frame_class: \033[1;32m{frame_class}\033[0m") + # Frame class should be a string + assert isinstance(frame_class, str) + # Should be one of the valid frame classes or 'Unclassified' + valid_classes = set(FRAME_CLASSES.keys()) | {'Unclassified'} + assert frame_class in valid_classes \ No newline at end of file