feat: add frame class extraction and formatting for media files
This commit is contained in:
@@ -32,3 +32,41 @@ SOURCE_DICT = {
|
|||||||
'HDTV': ['HDTV'],
|
'HDTV': ['HDTV'],
|
||||||
'BluRay': ['BluRay', 'BLURAY', 'Blu-ray'],
|
'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'
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -27,14 +27,22 @@ class MediaExtractor:
|
|||||||
"""Extract video source from filename"""
|
"""Extract video source from filename"""
|
||||||
return FilenameExtractor.extract_source(file_path)
|
return FilenameExtractor.extract_source(file_path)
|
||||||
|
|
||||||
def extract_resolution(self, file_path: Path) -> str | None:
|
def extract_frame_class(self, file_path: Path) -> str | None:
|
||||||
"""Extract resolution from media info or filename"""
|
"""Extract frame class from media info or filename"""
|
||||||
# Try media info first
|
# Try media info first
|
||||||
resolution = self.mediainfo_extractor.extract_resolution(file_path)
|
frame_class = self.mediainfo_extractor.extract_frame_class(file_path)
|
||||||
if resolution:
|
if frame_class:
|
||||||
return resolution
|
return frame_class
|
||||||
# Fallback to filename
|
# 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:
|
def extract_hdr(self, file_path: Path) -> str | None:
|
||||||
"""Extract HDR info from media info"""
|
"""Extract HDR info from media info"""
|
||||||
@@ -54,7 +62,9 @@ class MediaExtractor:
|
|||||||
'title': self.extract_title(file_path),
|
'title': self.extract_title(file_path),
|
||||||
'year': self.extract_year(file_path),
|
'year': self.extract_year(file_path),
|
||||||
'source': self.extract_source(file_path),
|
'source': self.extract_source(file_path),
|
||||||
|
'frame_class': self.extract_frame_class(file_path),
|
||||||
'resolution': self.extract_resolution(file_path),
|
'resolution': self.extract_resolution(file_path),
|
||||||
|
'aspect_ratio': self.extract_aspect_ratio(file_path),
|
||||||
'hdr': self.extract_hdr(file_path),
|
'hdr': self.extract_hdr(file_path),
|
||||||
'audio_langs': self.extract_audio_langs(file_path),
|
'audio_langs': self.extract_audio_langs(file_path),
|
||||||
'metadata': self.extract_metadata(file_path)
|
'metadata': self.extract_metadata(file_path)
|
||||||
|
|||||||
@@ -1,11 +1,19 @@
|
|||||||
import re
|
import re
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from ..constants import SOURCE_DICT
|
from ..constants import SOURCE_DICT, FRAME_CLASSES
|
||||||
|
|
||||||
|
|
||||||
class FilenameExtractor:
|
class FilenameExtractor:
|
||||||
"""Class to extract information from filename"""
|
"""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
|
@staticmethod
|
||||||
def extract_title(file_path: Path) -> str | None:
|
def extract_title(file_path: Path) -> str | None:
|
||||||
"""Extract movie title from filename"""
|
"""Extract movie title from filename"""
|
||||||
@@ -40,20 +48,11 @@ class FilenameExtractor:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def extract_resolution(file_path: Path) -> str | None:
|
def extract_frame_class(file_path: Path) -> str | None:
|
||||||
"""Extract resolution from filename (e.g., 2160p, 1080p, 720p)"""
|
"""Extract frame class from filename (480p, 720p, 1080p, 2160p, etc.)"""
|
||||||
file_name = file_path.name
|
file_name = file_path.name
|
||||||
match = re.search(r'(\d{3,4})[pi]', file_name, re.IGNORECASE)
|
match = re.search(r'(\d{3,4})[pi]', file_name, re.IGNORECASE)
|
||||||
if match:
|
if match:
|
||||||
height = int(match.group(1))
|
height = int(match.group(1))
|
||||||
if height >= 2160:
|
return FilenameExtractor._get_frame_class_from_height(height)
|
||||||
return '2160p'
|
return 'Unclassified'
|
||||||
elif height >= 1080:
|
|
||||||
return '1080p'
|
|
||||||
elif height >= 720:
|
|
||||||
return '720p'
|
|
||||||
elif height >= 480:
|
|
||||||
return '480p'
|
|
||||||
else:
|
|
||||||
return f'{height}p'
|
|
||||||
return None
|
|
||||||
@@ -1,6 +1,7 @@
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from pymediainfo import MediaInfo
|
from pymediainfo import MediaInfo
|
||||||
from collections import Counter
|
from collections import Counter
|
||||||
|
from ..constants import FRAME_CLASSES
|
||||||
|
|
||||||
|
|
||||||
class MediaInfoExtractor:
|
class MediaInfoExtractor:
|
||||||
@@ -13,24 +14,49 @@ class MediaInfoExtractor:
|
|||||||
'zh': 'chi', 'und': 'und'
|
'zh': 'chi', 'und': 'und'
|
||||||
}
|
}
|
||||||
|
|
||||||
def extract_resolution(self, file_path: Path) -> str | None:
|
def _get_frame_class_from_height(self, height: int) -> str:
|
||||||
"""Extract resolution from media info"""
|
"""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:
|
try:
|
||||||
media_info = MediaInfo.parse(file_path)
|
media_info = MediaInfo.parse(file_path)
|
||||||
video_tracks = [t for t in media_info.tracks if t.track_type == 'Video']
|
video_tracks = [t for t in media_info.tracks if t.track_type == 'Video']
|
||||||
if video_tracks:
|
if video_tracks:
|
||||||
height = getattr(video_tracks[0], 'height', None)
|
height = getattr(video_tracks[0], 'height', None)
|
||||||
if height:
|
if height:
|
||||||
if height >= 2160:
|
return self._get_frame_class_from_height(height)
|
||||||
return '2160p'
|
except:
|
||||||
elif height >= 1080:
|
pass
|
||||||
return '1080p'
|
return 'Unclassified'
|
||||||
elif height >= 720:
|
|
||||||
return '720p'
|
def extract_resolution(self, file_path: Path) -> str | None:
|
||||||
elif height >= 480:
|
"""Extract actual video resolution (WIDTHxHEIGHT) from media info"""
|
||||||
return '480p'
|
try:
|
||||||
else:
|
media_info = MediaInfo.parse(file_path)
|
||||||
return f'{height}p'
|
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:
|
except:
|
||||||
pass
|
pass
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -9,8 +9,14 @@ from ..utils import detect_file_type
|
|||||||
class MediaFormatter:
|
class MediaFormatter:
|
||||||
"""Class to format media data for display"""
|
"""Class to format media data for display"""
|
||||||
|
|
||||||
def format_file_info(self, file_path: Path, rename_data: dict) -> str:
|
def format_file_info_panel(self, file_path: Path, rename_data: dict) -> str:
|
||||||
"""Format complete file information for display"""
|
"""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
|
# Get file stats
|
||||||
size_full = SizeFormatter.format_size_full(file_path.stat().st_size)
|
size_full = SizeFormatter.format_size_full(file_path.stat().st_size)
|
||||||
date_formatted = DateFormatter.format_modification_date(file_path.stat().st_mtime)
|
date_formatted = DateFormatter.format_modification_date(file_path.stat().st_mtime)
|
||||||
@@ -24,19 +30,51 @@ class MediaFormatter:
|
|||||||
|
|
||||||
file_name = file_path.name
|
file_name = file_path.name
|
||||||
|
|
||||||
# Build basic info
|
output.append(f"[bold blue]Path:[/bold blue] {str(file_path)}")
|
||||||
full_info = f"[bold blue]Path:[/bold blue] {str(file_path)}\n\n"
|
output.append(f"[bold green]Size:[/bold green] {size_full}")
|
||||||
full_info += f"[bold green]Size:[/bold green] {size_full}\n"
|
output.append(f"[bold cyan]File:[/bold cyan] {file_name}")
|
||||||
full_info += f"[bold cyan]File:[/bold cyan] {file_name}\n"
|
output.append(f"[bold magenta]Modified:[/bold magenta] {date_formatted}")
|
||||||
full_info += f"{ext_info}\n"
|
output.append(f"{ext_info}")
|
||||||
full_info += f"[bold magenta]Modified:[/bold magenta] {date_formatted}"
|
|
||||||
|
|
||||||
# Extra metadata
|
return "\n".join(output)
|
||||||
extra_text = self._format_extra_metadata(rename_data['metadata'])
|
|
||||||
if extra_text:
|
|
||||||
full_info += f"\n\n{extra_text}"
|
|
||||||
|
|
||||||
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:
|
def format_proposed_name(self, rename_data: dict, ext_name: str) -> str:
|
||||||
"""Format the proposed filename"""
|
"""Format the proposed filename"""
|
||||||
@@ -49,8 +87,8 @@ class MediaFormatter:
|
|||||||
proposed_parts.append(rename_data['source'])
|
proposed_parts.append(rename_data['source'])
|
||||||
|
|
||||||
tags = []
|
tags = []
|
||||||
if rename_data['resolution']:
|
if rename_data['frame_class'] and rename_data['frame_class'] != 'Unclassified':
|
||||||
tags.append(rename_data['resolution'])
|
tags.append(rename_data['frame_class'])
|
||||||
if rename_data['hdr']:
|
if rename_data['hdr']:
|
||||||
tags.append(rename_data['hdr'])
|
tags.append(rename_data['hdr'])
|
||||||
if rename_data['audio_langs']:
|
if rename_data['audio_langs']:
|
||||||
@@ -66,7 +104,9 @@ class MediaFormatter:
|
|||||||
lines.append(f"Movie title: {rename_data['title'] or 'Unknown'}")
|
lines.append(f"Movie title: {rename_data['title'] or 'Unknown'}")
|
||||||
lines.append(f"Year: {rename_data['year'] or 'Unknown'}")
|
lines.append(f"Year: {rename_data['year'] or 'Unknown'}")
|
||||||
lines.append(f"Video source: {rename_data['source'] 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"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"HDR: {rename_data['hdr'] or 'No'}")
|
||||||
lines.append(f"Audio langs: {rename_data['audio_langs'] or 'None'}")
|
lines.append(f"Audio langs: {rename_data['audio_langs'] or 'None'}")
|
||||||
lines.append(f"Proposed filename: {proposed_name}")
|
lines.append(f"Proposed filename: {proposed_name}")
|
||||||
|
|||||||
@@ -1,5 +1,43 @@
|
|||||||
class ResolutionFormatter:
|
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
|
@staticmethod
|
||||||
def format_resolution_p(height: int) -> str:
|
def format_resolution_p(height: int) -> str:
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
import pytest
|
import pytest
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from ..extractors.filename_extractor import FilenameExtractor
|
from ..extractors.filename_extractor import FilenameExtractor
|
||||||
|
from ..constants import FRAME_CLASSES
|
||||||
|
|
||||||
|
|
||||||
def load_test_filenames():
|
def load_test_filenames():
|
||||||
@@ -17,6 +18,9 @@ def test_extract_title(filename):
|
|||||||
"""Test title extraction from filename"""
|
"""Test title extraction from filename"""
|
||||||
file_path = Path(filename)
|
file_path = Path(filename)
|
||||||
title = FilenameExtractor.extract_title(file_path)
|
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
|
# For now, just check it's not None and is string
|
||||||
assert isinstance(title, str) or title is None
|
assert isinstance(title, str) or title is None
|
||||||
|
|
||||||
@@ -39,15 +43,23 @@ def test_extract_source(filename):
|
|||||||
"""Test source extraction from filename"""
|
"""Test source extraction from filename"""
|
||||||
file_path = Path(filename)
|
file_path = Path(filename)
|
||||||
source = FilenameExtractor.extract_source(file_path)
|
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
|
# Source should be None or string
|
||||||
assert isinstance(source, str) or source is None
|
assert isinstance(source, str) or source is None
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("filename", load_test_filenames())
|
@pytest.mark.parametrize("filename", load_test_filenames())
|
||||||
def test_extract_resolution(filename):
|
def test_extract_frame_class(filename):
|
||||||
"""Test resolution extraction from filename"""
|
"""Test frame class extraction from filename"""
|
||||||
file_path = Path(filename)
|
file_path = Path(filename)
|
||||||
resolution = FilenameExtractor.extract_resolution(file_path)
|
frame_class = FilenameExtractor.extract_frame_class(file_path)
|
||||||
# Resolution should be None or string like '2160p'
|
# Print filename and extracted frame class clearly
|
||||||
if resolution:
|
print(f"\nFilename: \033[1;36m{filename}\033[0m")
|
||||||
assert 'p' in resolution or 'i' in resolution
|
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
|
||||||
Reference in New Issue
Block a user