diff --git a/src/show.py b/src/show.py index 232be40..3ed8fb5 100644 --- a/src/show.py +++ b/src/show.py @@ -1,11 +1,15 @@ from concurrent.futures import ProcessPoolExecutor, as_completed from dataclasses import dataclass, field +import math +import os +from pathlib import Path import re import numpy as np import cv2 import logging +import imdbinfo as imdb -from models import PathInfo +from models import FileCategory, FileInfo, PathInfo import tmdb @@ -242,63 +246,57 @@ def match_episodes_to_references( def process_show(info: PathInfo) -> PathInfo: # Gather meta information for identifying episodes - episode_durations: set[int] = set() + episode_durations: set[int] = ( + set() + ) # TODO: Store duration per specific episode for unique duration identification + episode_count = 0 # TODO: Per season for ep in imdb.get_all_episodes(info.imdb.imdb_id): + episode_count += 1 if ep.duration: episode_durations.add(int(ep.duration / 60)) - logging.debug(episode_durations) - - # Go over all files - count = 0 - for root, dirs, files in os.walk(info.path): - for filename in files: - filepath = Path(os.path.join(root, filename)) - file = FileInfo(filepath, info) - - if file.path.suffix != ".mkv" or file.duration_in_seconds is None: - continue - if ( - math.floor(file.duration_in_seconds / 60) in episode_durations - or math.ceil(file.duration_in_seconds / 60) in episode_durations - ): - print(f"{filename} {file.duration_in_seconds / 60} EPISODE") - count += 1 - else: - print(f"{filename} {file.duration_in_seconds / 60}") - - logging.info(f"Identified [{count}] episodes.") - - - - - - - - - - - if info.episodes is None: - logging.error( - "Episodes could not be identified, no reference matching possible." - ) - return info + logging.debug(f"Episode durations: {episode_durations}") + # Get reference images for episode identification if info.tmdb_id is None: logging.error("TMDB entry not identified, cannot find reference images.") return info - - # Match episodes to references references = tmdb.download_episode_images(info.tmdb_id) - matches = match_episodes_to_references( - [ - str(f.path.absolute()) - for f in info.episodes - if f.episode is None or f.season is None - ], - references.flatten(), - ) - logging.info(matches) + + # Go over all files + for root, dirs, files in os.walk(info.path): + # TODO: Detect season in dirs + + for filename in files: + filepath = Path(os.path.join(root, filename)) + file = FileInfo(filepath, info) + info.files.append(file) + + # Is file an episode? + if file.path.suffix != ".mkv" or file.duration_in_seconds is None: + continue + elif ( + math.floor(file.duration_in_seconds / 60) not in episode_durations + and math.ceil(file.duration_in_seconds / 60) not in episode_durations + ): + continue + + # ==== Handle episode ==== + file.category = FileCategory.EPISODE + + # TODO: Match episode to references + matches = match_episodes_to_references( + [ + str(f.path.absolute()) + for f in info.episodes + if f.episode is None or f.season is None + ], + references.flatten(), + ) + + if info.episodes is None: + logging.error("Episodes could not be identified.") + return info # Set new episode names if not matches.perfect_match or matches.reference_by_episode is None: