338 lines
9.1 KiB
Python
Executable File
338 lines
9.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
weight_jaccard = 0.5
|
|
weight_fuzzy = 1 - weight_jaccard
|
|
|
|
from serverlibrary import (
|
|
combined_score,
|
|
score_match_fuzzy,
|
|
score_match_jaccard,
|
|
score_match_recall_jaccard,
|
|
tvdb_get,
|
|
words,
|
|
)
|
|
|
|
MATCH_THRESHOLD = 0.85
|
|
|
|
|
|
class CustomLogger(logging.Logger):
|
|
def trace(self, msg: str, *args, **kwargs) -> None:
|
|
if self.isEnabledFor(TRACE):
|
|
self._log(TRACE, msg, args, **kwargs)
|
|
|
|
|
|
TRACE = 5
|
|
|
|
logging.setLoggerClass(CustomLogger)
|
|
logger: CustomLogger = logging.getLogger(__name__) # type: ignore[assignment]
|
|
|
|
|
|
def normalize(text: str) -> str:
|
|
text = text.lower()
|
|
|
|
text = re.sub(r"\.[a-z0-9]{2,4}$", "", text)
|
|
text = re.sub(
|
|
r"\b(1080p|720p|2160p|x264|x265|h264|h265|bluray|web-dl|webrip|dvdrip)\b",
|
|
"",
|
|
text,
|
|
flags=re.IGNORECASE,
|
|
)
|
|
|
|
text = re.sub(r"[._-]", " ", text)
|
|
text = re.sub(r"ß", "ss", text)
|
|
text = re.sub(r"\s+", " ", text)
|
|
text = re.sub(r"\((\d+)\s+\d+\)", r"(\1)", text)
|
|
text = re.sub(r"\((\d+)_+\d+\)", r"(\1)", text)
|
|
|
|
return text.strip()
|
|
|
|
|
|
def get_all_episodes_by_tvdbid(tvdb_id: int):
|
|
series = tvdb_get(
|
|
f"/series/{tvdb_id}/extended",
|
|
{"language": "deu"},
|
|
)["data"]
|
|
|
|
episodes = []
|
|
|
|
for season in series.get("seasons", []):
|
|
logger.info(f"Requesting Season: {season.get('number')}")
|
|
season_data = {}
|
|
try:
|
|
season_data = tvdb_get(
|
|
f"/seasons/{season['id']}/extended",
|
|
{"language": "deu"},
|
|
)["data"]
|
|
except requests.exceptions.RequestException:
|
|
logger.error(
|
|
"Failed to fetch season %s",
|
|
season["id"],
|
|
)
|
|
for ep in season_data.get("episodes", []):
|
|
episodes.append(
|
|
{
|
|
"id": ep["id"],
|
|
"season": ep["seasonNumber"],
|
|
"episode": ep["number"],
|
|
"title": ep["name"],
|
|
}
|
|
)
|
|
|
|
return series, episodes
|
|
|
|
|
|
def parse_series_name(series_name: str):
|
|
"""
|
|
Universe (1987) [66757]
|
|
-> ("Universe", 1987, 66757)
|
|
Universe (1987)
|
|
-> ("Universe", 1987, None)
|
|
Universe
|
|
-> ("Universe", None, None)
|
|
"""
|
|
|
|
tmdb_id = None
|
|
year = None
|
|
|
|
id_match = re.search(r"\[(\d+)]", series_name)
|
|
if id_match:
|
|
tmdb_id = int(id_match.group(1))
|
|
|
|
year_match = re.search(r"\((\d{4})\)", series_name)
|
|
if year_match:
|
|
year = int(year_match.group(1))
|
|
|
|
title = re.sub(r"\[\d+]", "", series_name)
|
|
title = re.sub(r"\(\d{4}\)", "", title)
|
|
title = title.strip()
|
|
|
|
return title, year, tmdb_id
|
|
|
|
|
|
def get_tvdbid_by_name(series_name: str):
|
|
title, year, tmdb_id = parse_series_name(series_name)
|
|
|
|
params = {
|
|
"query": title,
|
|
"type": "series",
|
|
"language": "deu",
|
|
}
|
|
|
|
if year:
|
|
params["year"] = str(year)
|
|
|
|
results = tvdb_get(
|
|
"/search",
|
|
params,
|
|
)["data"]
|
|
|
|
if not results:
|
|
raise RuntimeError(f"Series not found: {series_name}")
|
|
|
|
# Filter by year if available
|
|
if tmdb_id is not None:
|
|
filtered = []
|
|
|
|
for result in results:
|
|
result_tmdb_id = result.get("remote_ids")
|
|
for id in result_tmdb_id:
|
|
if id.get("sourceName") == "TheMovieDB.com":
|
|
result_tmdb_id = id.get("id")
|
|
|
|
if result_tmdb_id and int(result_tmdb_id) == tmdb_id:
|
|
filtered.append(result)
|
|
|
|
if filtered:
|
|
results = filtered
|
|
|
|
return results[0]
|
|
|
|
|
|
def best_episode_match(filename, episodes):
|
|
search_words = words(normalize(filename))
|
|
logger.debug(f"Search Words: {search_words}")
|
|
best_score = 0.0
|
|
best_episode = None
|
|
|
|
for episode in episodes:
|
|
if not episode["title"] or episode["title"] == "Unkown":
|
|
continue
|
|
candidate_words = words(normalize(episode["title"]))
|
|
fuzzy = score_match_fuzzy(
|
|
search_words,
|
|
candidate_words,
|
|
)
|
|
jaccard = score_match_jaccard(
|
|
search_words,
|
|
candidate_words,
|
|
)
|
|
recall_jaccard1 = score_match_recall_jaccard(
|
|
search_words,
|
|
candidate_words,
|
|
)
|
|
recall_jaccard2 = score_match_recall_jaccard(
|
|
candidate_words,
|
|
search_words,
|
|
)
|
|
|
|
score = combined_score(
|
|
fuzzy,
|
|
jaccard,
|
|
weight_fuzzy=weight_fuzzy,
|
|
weight_jaccard=weight_jaccard,
|
|
)
|
|
score = max(score, recall_jaccard1, recall_jaccard2)
|
|
|
|
if "hyänen" in search_words and score > 0.6:
|
|
logger.warning(
|
|
f"SEARCH WORDS: {search_words}, CANDIDATE WORDS: {candidate_words}"
|
|
)
|
|
|
|
if score > best_score:
|
|
logger.trace(f"{search_words}:\t\t{candidate_words},\t\t{score}")
|
|
best_score = score
|
|
best_episode = episode
|
|
|
|
return best_episode, best_score
|
|
|
|
|
|
def process_series(series_path: Path, dry_run: bool = False):
|
|
series_name = series_path.name
|
|
logger.debug(f"Pricess series with name : {series_name}")
|
|
if Path(f"{series_name}_episodes.json").is_file():
|
|
logger.info("load json files")
|
|
with open(f"{series_name}_media_data.json", encoding="utf-8") as f:
|
|
media_data = json.load(f)
|
|
with open(f"{series_name}_episodes.json", encoding="utf-8") as f:
|
|
episodes = json.load(f)
|
|
tvdb_id = int(media_data["tvdb_id"])
|
|
for ep in episodes:
|
|
logger.debug(f"{ep.get('title')}")
|
|
else:
|
|
logger.info("QUery from TVDB")
|
|
media_data = get_tvdbid_by_name(series_name)
|
|
tvdb_id = int(media_data["tvdb_id"])
|
|
series_data, episodes = get_all_episodes_by_tvdbid(tvdb_id)
|
|
logger.info("Dump to json")
|
|
with open(f"{series_name}_media_data.json", "w", encoding="utf-8") as f:
|
|
json.dump(media_data, f, ensure_ascii=False, indent=2)
|
|
with open(f"{series_name}_series_data.json", "w", encoding="utf-8") as f:
|
|
json.dump(series_data, f, ensure_ascii=False, indent=2)
|
|
with open(f"{series_name}_episodes.json", "w", encoding="utf-8") as f:
|
|
json.dump(episodes, f, ensure_ascii=False, indent=2)
|
|
logger.info(
|
|
f"Found series name from TVDB: {media_data['name']}, ID: {tvdb_id}, Year: {media_data.get('year')}"
|
|
)
|
|
extra_episodes_path = series_path / "Extra Episodes"
|
|
|
|
for file in extra_episodes_path.iterdir():
|
|
if not file.is_file():
|
|
continue
|
|
|
|
filename = re.sub(media_data["name"], "", file.stem)
|
|
|
|
episode, score = best_episode_match(
|
|
filename,
|
|
episodes,
|
|
)
|
|
if episode is None:
|
|
continue
|
|
|
|
if score < MATCH_THRESHOLD:
|
|
logger.warning(
|
|
f"LOW CONFIDENCE: {filename} {score:.2f}: {normalize(episode.get('title'))}"
|
|
)
|
|
continue
|
|
|
|
safe_title = re.sub(
|
|
r'[<>:"/\\|?*]',
|
|
"",
|
|
episode["title"],
|
|
)
|
|
|
|
new_name = (
|
|
f"{media_data.get('name')}"
|
|
f" - S{episode['season']:02d}"
|
|
f"E{episode['episode']:02d}"
|
|
f" - {safe_title}"
|
|
f"{file.suffix}"
|
|
)
|
|
|
|
target = file.with_name(new_name)
|
|
if score < 1:
|
|
logger.debug(f"{file.name} -> {new_name}, score={score:.3f}")
|
|
elif score == 1:
|
|
logger.trace(f"{file.name} -> {new_name}, score={score:.3f}")
|
|
|
|
if not dry_run:
|
|
file.rename(target)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"path",
|
|
help='Either a series folder "Maya the Bee (1975) [10620]" or a root folder containing such folders',
|
|
)
|
|
|
|
parser.add_argument(
|
|
"-ll",
|
|
"--log-level",
|
|
type=str.upper,
|
|
default="INFO",
|
|
choices=["TRACE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
|
|
help="Set log level",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="Dry run only",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
dry_run = args.dry_run
|
|
LOG_LEVELS = {
|
|
"TRACE": TRACE,
|
|
"DEBUG": logging.DEBUG,
|
|
"INFO": logging.INFO,
|
|
"WARNING": logging.WARNING,
|
|
"ERROR": logging.ERROR,
|
|
"CRITICAL": logging.CRITICAL,
|
|
}
|
|
|
|
logging.basicConfig(
|
|
level=LOG_LEVELS[args.log_level.upper()],
|
|
format="%(asctime)s %(levelname)s %(message)s",
|
|
)
|
|
|
|
input_path = Path(args.path).expanduser().resolve()
|
|
|
|
if not input_path.is_dir():
|
|
parser.error(f"Directory not found: {input_path}")
|
|
|
|
logger.info(f"Pricess folder: {input_path}")
|
|
|
|
# Input is already a series folder
|
|
if re.match(r"^(.*?)\s*\((\d{4})\)\s*\[(\d+)]$", input_path.name):
|
|
logger.info(f"Process series: {input_path.name}")
|
|
process_series(input_path, dry_run)
|
|
|
|
# Input is a root folder containing series folders
|
|
else:
|
|
for series_path in sorted(input_path.iterdir()):
|
|
if series_path.is_dir():
|
|
process_series(series_path, dry_run)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|