Files
Server-scripts/media_editing/sonar_translation.py
T

385 lines
10 KiB
Python
Executable File

#!/home/lucas/scripts/.venv/bin/python
import re
import requests
import argparse
from pathlib import Path
import logging
weight_jaccard = 0.6
weight_fuzzy = 1 - weight_jaccard
from serverlibrary import *
logging.setLoggerClass(CustomLogger)
logger: CustomLogger = logging.getLogger(__name__) # type: ignore[assignment]
def find_tmdbid_by_title(title, media_type="tv", year=None, language="de-DE"):
"""
Search TMDB by title.
media_type:
tv | movie
"""
params = {
"query": title,
"language": language,
}
if year:
if media_type == "tv":
params["first_air_date_year"] = year
else:
params["year"] = year
return tmdb_get(
f"/search/{media_type}",
params,
)
def find_by_tmdb_id(tmdb_id, language="de-DE"):
"""
Try to determine whether a TMDB ID is a TV show or movie.
Returns:
("tv", data)
("movie", data)
(None, None)
"""
try:
data = get_tv_details(tmdb_id, language)
return "tv", data
except requests.HTTPError:
pass
try:
data = get_movie_details(tmdb_id, language)
return "movie", data
except requests.HTTPError:
pass
return None, None
def get_tv_details(tmdb_id, language="de-DE"):
"""
Retrieve TV show details in German.
"""
return tmdb_get(
f"/tv/{tmdb_id}",
{
"language": language,
},
)
def get_movie_details(tmdb_id, language="de-DE"):
"""
Retrieve movie details in German.
"""
return tmdb_get(
f"/movie/{tmdb_id}",
{
"language": language,
},
)
def get_season(tmdb_id, season_number):
return tmdb_get(
f"/tv/{tmdb_id}/season/{season_number}",
{
"language": "de-DE",
},
)
def get_episode_list(show, series_path: Path, get_all: bool = False):
# _, show = find_by_tmdb_id(tmdb_id)
file_seasons = sorted(
{
int(m.group(1))
for d in series_path.iterdir()
if d.is_dir() and (m := re.match(r"^Season\s+0*(\d+)$", d.name, re.I))
}
)
episodes = []
for season in show.get("seasons", []):
season_number = season["season_number"]
# skip specials if desired
if not get_all and (season_number == 0 or season_number not in file_seasons):
continue
season_data = get_season(show.get("id"), season_number)
for episode in season_data.get("episodes", []):
episodes.append(
{
"season": season_number,
"episode": episode["episode_number"],
"title": episode["name"],
"air_date": episode["air_date"],
"runtime": episode.get("runtime"),
"tmdb_id": episode["id"],
}
)
return episodes
def print_results(result):
media_type, data = result
logger.info(f"Type : {media_type}")
logger.info(f"TMDB ID : {data.get('id')}")
logger.info(f"Title : {data.get('name') or data.get('title')}")
logger.info(
f"Original : {data.get('original_name') or data.get('original_title')}"
)
logger.info(
f"Year : {(data.get('first_air_date') or data.get('release_date') or '')[:4]}"
)
logger.info(f"Rating : {data.get('vote_average')}")
logger.info(f"Votes : {data.get('vote_count')}")
logger.info(f"Status : {data.get('status')}")
logger.info(f"Tagline : {data.get('tagline')}")
logger.info(f"Overview : {data.get('overview')}")
genres = ", ".join(g["name"] for g in data.get("genres", []))
logger.info(f"Genres : {genres}")
logger.info(
f"Poster : https://image.tmdb.org/t/p/original{data.get('poster_path', '')}"
)
logger.info(
f"Backdrop : https://image.tmdb.org/t/p/original{data.get('backdrop_path', '')}"
)
if media_type == "tv":
logger.info(f"Seasons : {data.get('number_of_seasons')}")
logger.info(f"Episodes : {data.get('number_of_episodes')}")
def normalize(text, title, year):
text = text.lower()
# remove series title
text = re.sub(re.escape(title.lower()), " ", text)
# remove year
text = re.sub(rf"\b{year}\b", " ", text)
# remove anything in [] including the brackets
text = re.sub(r"\[[^\]]*\]", " ", text)
# remove punctuation
text = re.sub(r"[^\w\s]", " ", text)
# remove episode numbers like S01E01
text = re.sub(r"\bs\d+e\d+\b", " ", text, flags=re.I)
# collapse spaces
text = re.sub(r"\s+", " ", text).strip()
return text
def best_episode_match(filename, episodes, title, year):
filename_norm = normalize(Path(filename).stem, title, year)
best = None
best_score = 0
for episode in episodes:
score_fuzzy = score_match_fuzzy(
filename_norm, normalize(episode["title"], title, year)
)
score_jaccard = score_match_jaccard(
filename_norm, normalize(episode["title"], title, year)
)
score = combined_score(score_fuzzy, score_jaccard, weight_fuzzy, weight_jaccard)
if score > best_score:
best_score = score
best = episode
return best, best_score
def rename_media(episodes, series_path, file_title, tmdb_title):
series = series_path.name
episode_map = {(ep["season"], ep["episode"]): ep for ep in episodes}
file_list = []
new_file_list = []
for file in series_path.rglob("*"):
match_SE_name = re.search(r"S(\d{2})E(\d{2})", file.name, re.I)
if not match_SE_name:
match_season_folder = re.search(r"Season (\d{2})", file.name, re.I)
if match_season_folder:
continue # witchout bothering anyone
logger.error(f"No season/episode number found: {file.name}")
continue
season = int(match_SE_name.group(1))
episode = int(match_SE_name.group(2))
episode_number = episode_map.get((season, episode))
if not episode_number:
logger.warning(f"Episode not found: S{season:02d}E{episode:02d}")
continue
episode_season_number_match = re.match(
r"^(.*?S\d{2}E\d{2})\s*-\s*(.*?)\s*(\[[^\]]+\].*)?$",
file.stem,
re.I,
)
if not episode_season_number_match:
logger.warning(f"S{season:02d}E{episode:02d} not found")
continue
prefix = episode_season_number_match.group(1).replace(
file_title, tmdb_title
) # Andor (2022) - S01E07
quality = (
episode_season_number_match.group(3) or ""
) # [WEBRip-1080p][AAC 5.1][h265]
new_name = (
f"{prefix} - {episode_number['title']}"
f"{(' ' + quality) if quality else ''}"
f"{file.suffix}"
)
logger.info(f"{file.name} -> {new_name}")
file_list.append(file)
new_file_list.append(new_name)
logger.info(f"{series} -> {tmdb_title}")
file_list.append(series_path)
new_file_list.append(tmdb_title)
has_changes = False
for file, new_file in zip(file_list, new_file_list):
if file.name != new_file:
has_changes = True
break
if not has_changes:
logger.info("No changes necessary.")
return
answer = input("Rename files? [y/N]: ").strip().lower()
if answer in ("y", "yes"):
logger.info("Proceeding...")
for file, new_file in zip(file_list, new_file_list):
file.rename(file.with_name(new_file))
logger.info("All files renamed")
return
logger.warning("Cancelled.")
return
def process_series(series_path: Path, skip_rename: bool = False):
series = series_path.name
m = re.match(r"^(.*?)\s*\((\d{4})\)\s*\[(\d+)\]$", series)
if not m:
logger.warning(f"Skipping invalid folder: {series}")
return
file_title = m.group(1).strip()
# file_year = int(m.group(2))
file_tmdb_id = int(m.group(3))
if file_tmdb_id == 0:
logger.error("TMDB ID required")
return
try:
data = find_by_tmdb_id(file_tmdb_id)
if not data:
logger.error("TMDB lookup failed")
return
media_type, media_data = data
tmdb_title = f"{(media_data.get('name') or media_data.get('title'))}"
tmdb_year = f"{(media_data.get('first_air_date') or media_data.get('release_date') or '')[:4]}"
new_series_folder_name = f"{tmdb_title} ({tmdb_year}) [{media_data.get('id')}]"
# print_results(data)
get_all = False
extra_episodes_path = series_path / "Extra Episodes"
if extra_episodes_path.is_dir():
get_all = True
episodes = get_episode_list(media_data, series_path, get_all=get_all)
process_extra_episodes(episodes, extra_episodes_path, tmdb_title, tmdb_year)
if skip_rename:
return
rename_media(episodes, series_path, file_title, new_series_folder_name)
except Exception as e:
logger.error(f"Error processing {series}: {e}")
def process_extra_episodes(
episode_list: list, extra_episodes_path: Path, tmdb_title, tmdb_year
):
for file in extra_episodes_path.iterdir():
if not (file.suffix == ".mkv" or file.suffix == ".mp4"):
continue
best_matc, best_score = best_episode_match(
file, episode_list, tmdb_title, tmdb_year
)
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",
)
args = parser.parse_args()
setup_logging(args.log_level, "sonar_translation")
input_path = Path(args.path)
if not input_path.is_dir():
parser.error(f"Directory not found: {input_path}")
# Input is already a series folder
if re.match(r"^(.*?)\s*\((\d{4})\)\s*\[(\d+)\]$", input_path.name):
process_series(input_path)
# 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)
if __name__ == "__main__":
main()