Plex-Meta-Manager/modules/overlays.py
2022-05-07 11:46:57 -04:00

230 lines
12 KiB
Python

import os, time
from modules import util
from modules.builder import CollectionBuilder
from modules.util import Failed
from plexapi.exceptions import BadRequest
from plexapi.video import Show, Season, Episode
from PIL import Image
logger = util.logger
class Overlays:
def __init__(self, config, library):
self.config = config
self.library = library
self.overlays = []
def run_overlays(self):
logger.info("")
logger.separator(f"{self.library.name} Library Overlays")
logger.info("")
overlay_rating_keys = {}
item_keys = {}
os.makedirs(self.library.overlay_backup, exist_ok=True)
overlay_updated = {}
overlay_images = {}
item_overlays = {}
if self.library.remove_overlays:
logger.info("")
logger.separator(f"Removing Overlays for the {self.library.name} Library")
logger.info("")
else:
for overlay_file in self.library.overlay_files:
for k, v in overlay_file.overlays.items():
try:
builder = CollectionBuilder(self.config, overlay_file, k, v, library=self.library, overlay=True)
logger.info("")
logger.separator(f"Gathering Items for {k} Overlay", space=False, border=False)
if builder.overlay not in overlay_rating_keys:
overlay_rating_keys[builder.overlay] = []
if builder.filters or builder.tmdb_filters:
logger.info("")
for filter_key, filter_value in builder.filters:
logger.info(f"Collection Filter {filter_key}: {filter_value}")
for filter_key, filter_value in builder.tmdb_filters:
logger.info(f"Collection Filter {filter_key}: {filter_value}")
for method, value in builder.builders:
logger.debug("")
logger.debug(f"Builder: {method}: {value}")
logger.info("")
builder.filter_and_save_items(builder.gather_ids(method, value))
if builder.added_items:
for item in builder.added_items:
item_keys[item.ratingKey] = item
if item.ratingKey not in overlay_rating_keys[builder.overlay]:
overlay_rating_keys[builder.overlay].append(item.ratingKey)
if builder.remove_overlays:
for rk in overlay_rating_keys[builder.overlay]:
for remove_overlay in builder.remove_overlays:
if remove_overlay in overlay_rating_keys and rk in overlay_rating_keys[remove_overlay]:
overlay_rating_keys[remove_overlay].remove(rk)
except Failed as e:
logger.error(e)
for overlay_name, over_keys in overlay_rating_keys.items():
clean_name, _ = util.validate_filename(overlay_name)
image_compare = None
if self.config.Cache:
_, image_compare, _ = self.config.Cache.query_image_map(overlay_name, f"{self.library.image_table_name}_overlays")
overlay_file = os.path.join(self.library.overlay_folder, f"{clean_name}.png")
overlay_size = os.stat(overlay_file).st_size
overlay_updated[overlay_name] = not image_compare or str(overlay_size) != str(image_compare)
overlay_images[overlay_name] = Image.open(overlay_file).convert("RGBA")
for over_key in over_keys:
if over_key not in item_overlays:
item_overlays[over_key] = []
item_overlays[over_key].append(overlay_name)
if self.config.Cache:
self.config.Cache.update_image_map(overlay_name, f"{self.library.image_table_name}_overlays", overlay_name, overlay_size)
def find_poster_url(plex_item):
if self.library.is_movie:
if plex_item.ratingKey in self.library.movie_rating_key_map:
return self.config.TMDb.get_movie(self.library.movie_rating_key_map[plex_item.ratingKey]).poster_url
elif self.library.is_show:
check_key = plex_item.ratingKey if isinstance(plex_item, Show) else plex_item.show().ratingKey
tmdb_id = self.config.Convert.tvdb_to_tmdb(self.library.show_rating_key_map[check_key])
if isinstance(plex_item, Show) and plex_item.ratingKey in self.library.show_rating_key_map:
return self.config.TMDb.get_show(tmdb_id).poster_url
elif isinstance(plex_item, Season):
return self.config.TMDb.get_season(tmdb_id, plex_item.seasonNumber).poster_url
elif isinstance(plex_item, Episode):
return self.config.TMDb.get_episode(tmdb_id, plex_item.seasonNumber, plex_item.episodeNumber).still_url
def get_overlay_items(libtype=None):
return [o for o in self.library.search(label="Overlay", libtype=libtype) if o.ratingKey not in item_overlays]
remove_overlays = get_overlay_items()
if self.library.is_show:
remove_overlays.extend(get_overlay_items(libtype="episode"))
remove_overlays.extend(get_overlay_items(libtype="season"))
elif self.library.is_music:
remove_overlays.extend(get_overlay_items(libtype="album"))
for i, item in enumerate(remove_overlays, 1):
logger.ghost(f"Restoring: {i}/{len(remove_overlays)} {item.title}")
clean_name, _ = util.validate_filename(item.title)
poster, _, item_dir = self.library.find_assets(
name="poster" if self.library.asset_folders else clean_name,
folder_name=clean_name if self.library.asset_folders else None,
prefix=f"{item.title}'s "
)
is_url = False
original = None
if poster:
poster_location = poster.location
elif os.path.exists(os.path.join(self.library.overlay_backup, f"{item.ratingKey}.png")):
original = os.path.join(self.library.overlay_backup, f"{item.ratingKey}.png")
poster_location = original
elif os.path.exists(os.path.join(self.library.overlay_backup, f"{item.ratingKey}.jpg")):
original = os.path.join(self.library.overlay_backup, f"{item.ratingKey}.jpg")
poster_location = original
else:
is_url = True
poster_location = find_poster_url(item)
if poster_location:
self.library.upload_poster(item, poster_location, url=is_url)
self.library.edit_tags("label", item, remove_tags=["Overlay"])
if original:
os.remove(original)
else:
logger.error(f"No Poster found to restore for {item.title}")
logger.exorcise()
if item_overlays:
logger.info("")
logger.separator(f"Applying Overlays for the {self.library.name} Library")
logger.info("")
for i, (over_key, over_names) in enumerate(item_overlays.items(), 1):
try:
item = item_keys[over_key]
logger.ghost(f"Overlaying: {i}/{len(item_overlays)} {item.title}")
image_compare = None
overlay_compare = None
if self.config.Cache:
image, image_compare, _ = self.config.Cache.query_image_map(item.ratingKey, f"{self.library.image_table_name}_overlays")
overlay_compare = [] if overlay_compare is None else util.get_list(overlay_compare)
has_overlay = any([item_tag.tag.lower() == "overlay" for item_tag in item.labels])
overlay_change = False if has_overlay else True
if not overlay_change:
for oc in overlay_compare:
if oc not in over_names:
overlay_change = True
if not overlay_change:
for over_name in over_names:
if over_name not in overlay_compare or overlay_updated[over_name]:
overlay_change = True
clean_name, _ = util.validate_filename(item.title)
poster, _, item_dir = self.library.find_assets(
name="poster" if self.library.asset_folders else clean_name,
folder_name=clean_name if self.library.asset_folders else None,
prefix=f"{item.title}'s "
)
has_original = False
changed_image = False
new_backup = None
if poster:
if image_compare and str(poster.compare) != str(image_compare):
changed_image = True
elif has_overlay:
if os.path.exists(os.path.join(self.library.overlay_backup, f"{item.ratingKey}.png")):
has_original = os.path.join(self.library.overlay_backup, f"{item.ratingKey}.png")
elif os.path.exists(os.path.join(self.library.overlay_backup, f"{item.ratingKey}.jpg")):
has_original = os.path.join(self.library.overlay_backup, f"{item.ratingKey}.jpg")
else:
self.library.reload(item)
new_backup = find_poster_url(item)
if new_backup is None:
new_backup = item.posterUrl
else:
self.library.reload(item)
new_backup = item.posterUrl
if new_backup:
changed_image = True
image_response = self.config.get(new_backup)
if image_response.status_code >= 400:
raise Failed(f"Overlay Error: Poster Download Failed for {item.title}")
i_ext = "jpg" if image_response.headers["Content-Type"] == "image/jpeg" else "png"
backup_image_path = os.path.join(self.library.overlay_backup, f"{item.ratingKey}.{i_ext}")
with open(backup_image_path, "wb") as handler:
handler.write(image_response.content)
while util.is_locked(backup_image_path):
time.sleep(1)
has_original = backup_image_path
poster_uploaded = False
if changed_image or overlay_change:
new_poster = Image.open(poster.location if poster else has_original).convert("RGBA")
temp = os.path.join(self.library.overlay_folder, f"temp.png")
try:
for over_name in over_names:
new_poster = new_poster.resize(overlay_images[over_name].size, Image.ANTIALIAS)
new_poster.paste(overlay_images[over_name], (0, 0), overlay_images[over_name])
new_poster.save(temp, "PNG")
self.library.upload_poster(item, temp)
self.library.edit_tags("label", item, add_tags=["Overlay"])
self.library.reload(item)
poster_uploaded = True
logger.info(f"Detail: Overlays: {', '.join(over_names)} applied to {item.title}")
except (OSError, BadRequest) as e:
logger.stacktrace()
raise Failed(f"Overlay Error: {e}")
if self.config.Cache:
if poster_uploaded:
self.config.Cache.update_image_map(
item.ratingKey, self.library.image_table_name, item.thumb,
poster.compare if poster else item.thumb, overlay=','.join(over_names)
)
except Failed as e:
logger.error(e)
logger.exorcise()