feat: add frame class extraction and formatting for media files

This commit is contained in:
sHa
2025-12-25 05:13:56 +00:00
parent 4e7cf94d58
commit 37efdf60d3
7 changed files with 224 additions and 61 deletions

View File

@@ -32,3 +32,41 @@ SOURCE_DICT = {
'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'
}
}

View File

@@ -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)

View File

@@ -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
return FilenameExtractor._get_frame_class_from_height(height)
return 'Unclassified'

View File

@@ -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

View File

@@ -9,8 +9,14 @@ 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)
@@ -24,19 +30,51 @@ class MediaFormatter:
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}"
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}")
# Extra metadata
extra_text = self._format_extra_metadata(rename_data['metadata'])
if extra_text:
full_info += f"\n\n{extra_text}"
return "\n".join(output)
return full_info
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}")

View File

@@ -1,5 +1,43 @@
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:

View File

@@ -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
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