- Extract matching strategies into strategies.py module (#10) * DurationMatchStrategy: match files by runtime similarity * SequentialStrategy: fill remaining gaps in order * DiscMappingStrategy: explicit disc-to-episode mapping * ConstraintAwareStrategy: validate disc capacity + duration - Support multiple video formats: .mkv, .mp4, .avi, .mov, .wmv, .flv, .webm, .m4v (#12) - Make MediaInfo fallback ratio configurable via config.json duration_minutes_per_gb (#13) Closes #10, #12, #13
This commit is contained in:
Jarian
parent
b30c798280
commit
c818e69b0f
@ -1,5 +1,20 @@
|
||||
"""Episode Matcher module for classifying and renaming video files."""
|
||||
from .file_classifier import FileClassifier
|
||||
from .episode_renamer import EpisodeRenamer
|
||||
from .strategies import (
|
||||
DurationMatchStrategy,
|
||||
SequentialStrategy,
|
||||
DiscMappingStrategy,
|
||||
ConstraintAwareStrategy,
|
||||
match_episodes,
|
||||
)
|
||||
|
||||
__all__ = ['FileClassifier', 'EpisodeRenamer']
|
||||
__all__ = [
|
||||
'FileClassifier',
|
||||
'EpisodeRenamer',
|
||||
'DurationMatchStrategy',
|
||||
'SequentialStrategy',
|
||||
'DiscMappingStrategy',
|
||||
'ConstraintAwareStrategy',
|
||||
'match_episodes',
|
||||
]
|
||||
|
||||
@ -14,48 +14,67 @@ try:
|
||||
except ImportError:
|
||||
config_manager = None
|
||||
|
||||
# Supported video file extensions
|
||||
DEFAULT_VIDEO_EXTENSIONS = {'.mkv', '.mp4', '.avi', '.mov', '.wmv', '.flv', '.webm', '.m4v'}
|
||||
|
||||
|
||||
class FileClassifier:
|
||||
"""Classifies video files as episodes or extras based on file size and video duration."""
|
||||
|
||||
def __init__(self, folder_path: str):
|
||||
"""Initialize with folder path containing video files."""
|
||||
def __init__(self, folder_path: str, video_extensions: set = None):
|
||||
"""Initialize with folder path containing video files.
|
||||
|
||||
Args:
|
||||
folder_path: Path to folder containing video files.
|
||||
video_extensions: Set of file extensions to scan (e.g., {'.mkv', '.mp4'}).
|
||||
Defaults to DEFAULT_VIDEO_EXTENSIONS if not provided.
|
||||
"""
|
||||
self.folder_path = Path(folder_path)
|
||||
self.video_extensions = video_extensions or DEFAULT_VIDEO_EXTENSIONS
|
||||
self.video_files = self._get_video_files()
|
||||
self.file_info = self._analyze_files()
|
||||
|
||||
def _get_video_files(self) -> List[Path]:
|
||||
"""Get all MKV files in the folder, excluding system files."""
|
||||
"""Get all video files in the folder, excluding system files."""
|
||||
if not self.folder_path.exists():
|
||||
raise FileNotFoundError(f"Folder not found: {self.folder_path}")
|
||||
|
||||
all_mkv_files = list(self.folder_path.glob("*.mkv"))
|
||||
|
||||
video_files = []
|
||||
for ext in self.video_extensions:
|
||||
video_files.extend(self.folder_path.glob(f"*{ext}"))
|
||||
|
||||
# Filter out macOS resource fork files and other system files
|
||||
video_files = []
|
||||
for file in all_mkv_files:
|
||||
filtered = []
|
||||
for file in video_files:
|
||||
if file.name.startswith('._'):
|
||||
continue # Skip macOS resource fork files
|
||||
continue
|
||||
if file.name.startswith('.'):
|
||||
continue # Skip any hidden files
|
||||
video_files.append(file)
|
||||
continue
|
||||
filtered.append(file)
|
||||
|
||||
if not video_files:
|
||||
print(f"Warning: No valid MKV files found in {self.folder_path}")
|
||||
if not filtered:
|
||||
ext_list = ', '.join(sorted(self.video_extensions))
|
||||
print(f"Warning: No video files found ({ext_list}) in {self.folder_path}")
|
||||
|
||||
return video_files
|
||||
return filtered
|
||||
|
||||
def _get_file_size(self, file_path: Path) -> int:
|
||||
"""Get file size in bytes."""
|
||||
return file_path.stat().st_size
|
||||
|
||||
def _get_fallback_duration_ratio(self) -> float:
|
||||
"""Get minutes-per-GB ratio for fallback duration estimation."""
|
||||
if config_manager:
|
||||
return config_manager.get_duration_minutes_per_gb()
|
||||
return 45.0
|
||||
|
||||
def _get_video_duration(self, file_path: Path) -> float:
|
||||
"""Get video duration in minutes using MediaInfo."""
|
||||
fallback_ratio = self._get_fallback_duration_ratio()
|
||||
if MediaInfo is None:
|
||||
print("MediaInfo not available, using file size as proxy for duration")
|
||||
# Rough estimate: 1GB ≈ 45 minutes for typical video
|
||||
print(f"MediaInfo not available, using file size as proxy for duration ({fallback_ratio} min/GB)")
|
||||
size_gb = self._get_file_size(file_path) / (1024**3)
|
||||
return size_gb * 45
|
||||
return size_gb * fallback_ratio
|
||||
|
||||
try:
|
||||
media_info = MediaInfo.parse(str(file_path))
|
||||
@ -74,10 +93,10 @@ class FileClassifier:
|
||||
duration_minutes = float(duration_ms) / (1000 * 60)
|
||||
|
||||
# Sanity check: if duration seems unreasonable, fall back to size estimation
|
||||
if duration_minutes > 300: # More than 5 hours is likely wrong
|
||||
if duration_minutes > 300:
|
||||
print(f"Warning: Suspicious duration ({duration_minutes:.1f} min) for {file_path.name}, using size estimation")
|
||||
size_gb = self._get_file_size(file_path) / (1024**3)
|
||||
return size_gb * 45
|
||||
return size_gb * fallback_ratio
|
||||
|
||||
return duration_minutes
|
||||
except Exception as e:
|
||||
@ -85,7 +104,7 @@ class FileClassifier:
|
||||
|
||||
# Fallback to size-based estimation
|
||||
size_gb = self._get_file_size(file_path) / (1024**3)
|
||||
return size_gb * 45
|
||||
return size_gb * fallback_ratio
|
||||
|
||||
def _analyze_files(self) -> List[Dict]:
|
||||
"""Analyze all video files to get size and duration info."""
|
||||
|
||||
200
src/Matcher/strategies.py
Normal file
200
src/Matcher/strategies.py
Normal file
@ -0,0 +1,200 @@
|
||||
"""Matching strategies for episode-to-file mapping."""
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional
|
||||
|
||||
|
||||
class DurationMatchStrategy:
|
||||
"""Match files to TVDB episodes by duration similarity."""
|
||||
|
||||
def __init__(self, tolerance_minutes: float = 1.0):
|
||||
self.tolerance = tolerance_minutes
|
||||
|
||||
def match(self, files: List[Dict], tvdb_episodes: List[Dict],
|
||||
allowed_episodes: Optional[List[int]] = None) -> List[Dict]:
|
||||
matched = []
|
||||
used_files = set()
|
||||
used_episodes = set()
|
||||
|
||||
for file_info in files:
|
||||
if id(file_info) in used_files:
|
||||
continue
|
||||
|
||||
file_duration = file_info['duration']
|
||||
best_match = None
|
||||
best_score = float('inf')
|
||||
|
||||
for tvdb_ep in tvdb_episodes:
|
||||
if tvdb_ep['episode_number'] in used_episodes:
|
||||
continue
|
||||
if allowed_episodes and tvdb_ep['episode_number'] not in allowed_episodes:
|
||||
continue
|
||||
|
||||
tvdb_duration = tvdb_ep.get('runtime', 0)
|
||||
if tvdb_duration == 0:
|
||||
continue
|
||||
|
||||
duration_diff = abs(file_duration - tvdb_duration)
|
||||
if duration_diff < best_score:
|
||||
best_score = duration_diff
|
||||
best_match = tvdb_ep
|
||||
|
||||
if best_match and best_score <= self.tolerance:
|
||||
matched.append({
|
||||
'file_info': file_info['file_info'],
|
||||
'episode_number': best_match['episode_number'],
|
||||
'tvdb_info': best_match,
|
||||
'duration_diff': best_score,
|
||||
'strategy': 'duration_match'
|
||||
})
|
||||
used_files.add(id(file_info))
|
||||
used_episodes.add(best_match['episode_number'])
|
||||
|
||||
return matched, used_files, used_episodes
|
||||
|
||||
|
||||
class SequentialStrategy:
|
||||
"""Assign remaining files to episodes in sequential order."""
|
||||
|
||||
def match(self, remaining_files: List[Dict], tvdb_episodes: List[Dict],
|
||||
used_episodes: set, allowed_episodes: List[int]) -> List[Dict]:
|
||||
matched = []
|
||||
available = sorted(
|
||||
[e for e in allowed_episodes if e not in used_episodes]
|
||||
)
|
||||
|
||||
sorted_files = sorted(remaining_files, key=lambda x: x['filename'])
|
||||
|
||||
for i, file_info in enumerate(sorted_files):
|
||||
if i >= len(available):
|
||||
break
|
||||
|
||||
episode_number = available[i]
|
||||
tvdb_match = next(
|
||||
(ep for ep in tvdb_episodes if ep['episode_number'] == episode_number),
|
||||
None
|
||||
)
|
||||
|
||||
duration_diff = None
|
||||
if tvdb_match and tvdb_match.get('runtime', 0) > 0:
|
||||
duration_diff = abs(file_info['duration'] - tvdb_match['runtime'])
|
||||
|
||||
matched.append({
|
||||
'file_info': file_info['file_info'],
|
||||
'episode_number': episode_number,
|
||||
'tvdb_info': tvdb_match,
|
||||
'duration_diff': duration_diff,
|
||||
'strategy': 'sequential'
|
||||
})
|
||||
|
||||
return matched
|
||||
|
||||
|
||||
class DiscMappingStrategy:
|
||||
"""Match files to episodes using explicit disc-to-episode mapping."""
|
||||
|
||||
def __init__(self, disc_mapping: Dict[int, List[int]]):
|
||||
self.disc_mapping = disc_mapping
|
||||
|
||||
def match(self, episodes_info: List[Dict], tvdb_episodes: List[Dict]) -> List[Dict]:
|
||||
matched = []
|
||||
disc_groups = {}
|
||||
|
||||
for info in episodes_info:
|
||||
disc_num = info['disc_number']
|
||||
if disc_num is not None:
|
||||
disc_groups.setdefault(disc_num, []).append(info)
|
||||
|
||||
for disc_num in sorted(disc_groups.keys()):
|
||||
files = sorted(disc_groups[disc_num], key=lambda x: x['filename'])
|
||||
expected = self.disc_mapping.get(disc_num, [])
|
||||
|
||||
for i, file_info in enumerate(files):
|
||||
if i >= len(expected):
|
||||
break
|
||||
|
||||
episode_number = expected[i]
|
||||
tvdb_episode = next(
|
||||
(ep for ep in tvdb_episodes if ep['episode_number'] == episode_number),
|
||||
None
|
||||
)
|
||||
|
||||
matched.append({
|
||||
'file_info': file_info['file_info'],
|
||||
'episode_number': episode_number,
|
||||
'tvdb_info': tvdb_episode,
|
||||
'strategy': 'disc_mapping'
|
||||
})
|
||||
|
||||
return matched
|
||||
|
||||
|
||||
class ConstraintAwareStrategy:
|
||||
"""Match files respecting disc capacity and duration constraints."""
|
||||
|
||||
def __init__(self, max_duration_diff: float = 2.0,
|
||||
min_file_duration_buffer: float = 2.0):
|
||||
self.max_duration_diff = max_duration_diff
|
||||
self.min_file_buffer = min_file_duration_buffer
|
||||
|
||||
def validate(self, file_info: Dict, tvdb_episode: Dict,
|
||||
disc_assignments: Dict[int, int],
|
||||
disc_capacity: Dict[int, int]) -> bool:
|
||||
file_duration = file_info['duration']
|
||||
tvdb_duration = tvdb_episode.get('runtime', 0)
|
||||
file_disc = file_info['disc_number']
|
||||
|
||||
if tvdb_duration > 0:
|
||||
diff = abs(file_duration - tvdb_duration)
|
||||
if diff > self.max_duration_diff:
|
||||
return False
|
||||
if file_duration < (tvdb_duration - self.min_file_buffer):
|
||||
return False
|
||||
|
||||
if file_disc and disc_capacity:
|
||||
max_eps = disc_capacity.get(file_disc, 0)
|
||||
current = disc_assignments.get(file_disc, 0)
|
||||
if current >= max_eps:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def match_episodes(files: List[Dict], tvdb_episodes: List[Dict],
|
||||
disc_mapping: Optional[Dict[int, List[int]]] = None,
|
||||
disc_capacity: Optional[Dict[int, int]] = None
|
||||
) -> List[Dict]:
|
||||
"""
|
||||
Multi-strategy episode matcher.
|
||||
|
||||
Pipeline:
|
||||
1. If disc_mapping provided, use DiscMappingStrategy
|
||||
2. Otherwise, use DurationMatchStrategy for close matches
|
||||
3. Fill gaps with SequentialStrategy
|
||||
|
||||
Args:
|
||||
files: Episode info dicts with file_info, disc_number, duration, filename
|
||||
tvdb_episodes: TVDB episode data with runtime
|
||||
disc_mapping: Optional disc-to-episode mapping
|
||||
disc_capacity: Optional disc capacity constraints
|
||||
|
||||
Returns:
|
||||
List of match dicts sorted by episode number
|
||||
"""
|
||||
all_allowed = list(range(1, len(tvdb_episodes) + 1))
|
||||
|
||||
if disc_mapping:
|
||||
return DiscMappingStrategy(disc_mapping).match(files, tvdb_episodes)
|
||||
|
||||
# Phase 1: Duration matching
|
||||
duration_strategy = DurationMatchStrategy()
|
||||
matched, used_files, used_eps = duration_strategy.match(
|
||||
files, tvdb_episodes, all_allowed
|
||||
)
|
||||
|
||||
# Phase 2: Sequential fill for remaining
|
||||
remaining = [f for f in files if id(f) not in used_files]
|
||||
sequential = SequentialStrategy()
|
||||
matched.extend(sequential.match(remaining, tvdb_episodes, used_eps, all_allowed))
|
||||
|
||||
matched.sort(key=lambda x: x['episode_number'])
|
||||
return matched
|
||||
@ -23,6 +23,8 @@ class ConfigManager:
|
||||
default_config = {
|
||||
"tvdb_api_key": "",
|
||||
"default_episode_duration": 45,
|
||||
"duration_minutes_per_gb": 45,
|
||||
"video_extensions": [".mkv", ".mp4", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"],
|
||||
"classification_thresholds": {
|
||||
"size_threshold_ratio": 0.3,
|
||||
"duration_threshold_ratio": 0.4
|
||||
@ -65,6 +67,16 @@ class ConfigManager:
|
||||
"""Get default episode duration in minutes."""
|
||||
return self.config.get("default_episode_duration", 45)
|
||||
|
||||
def get_duration_minutes_per_gb(self) -> float:
|
||||
"""Get fallback duration estimation ratio (minutes per GB)."""
|
||||
return float(self.config.get("duration_minutes_per_gb", 45))
|
||||
|
||||
def get_video_extensions(self) -> list:
|
||||
"""Get list of supported video file extensions."""
|
||||
return self.config.get("video_extensions", [
|
||||
".mkv", ".mp4", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"
|
||||
])
|
||||
|
||||
def get_classification_thresholds(self) -> Dict[str, float]:
|
||||
"""Get classification threshold ratios."""
|
||||
return self.config.get("classification_thresholds", {
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user