101 lines
3.6 KiB
Python
101 lines
3.6 KiB
Python
import cv2
|
|
import numpy as np
|
|
import requests
|
|
import logging
|
|
|
|
from models import ReferenceShowImages
|
|
|
|
# ==== CONFIGURATION ====
|
|
TMDB_API_KEY = "b7006350eb3eeb4cf7d9cb6db44cdc0b" # <-- Replace with your TMDB API key
|
|
BASE_URL = "https://api.themoviedb.org/3"
|
|
IMG_BASE = "https://image.tmdb.org/t/p/original"
|
|
|
|
|
|
def tmdb_request(endpoint: str, params: dict = {}):
|
|
"""Helper to query TMDB API with authentication."""
|
|
params["api_key"] = TMDB_API_KEY
|
|
response = requests.get(f"{BASE_URL}{endpoint}", params=params)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
|
|
def search_show(imdb_id: str, title: str, year: int) -> str | None:
|
|
"""Find TMDB ID by IMDb ID first, fallback to title/year. Returns TMDB id if successful, otherwise None."""
|
|
# Try external source (IMDb ID)
|
|
try:
|
|
res = tmdb_request("/find/" + imdb_id, {"external_source": "imdb_id"})
|
|
if res.get("tv_results"):
|
|
return res["tv_results"][0]["id"]
|
|
except Exception as e:
|
|
logging.warning(
|
|
"TMDB lookup with IMDB ID failed for SHOW, falling back to search:", e
|
|
)
|
|
|
|
# Fallback to title/year search
|
|
res = tmdb_request("/search/tv", {"query": title, "first_air_date_year": year})
|
|
if res.get("results"):
|
|
return res["results"][0]["id"]
|
|
|
|
logging.error(f"Unable to find show for title [{title}] and year [{year}].")
|
|
return None
|
|
|
|
|
|
def download_image(img_path: str) -> np.ndarray | None:
|
|
"""Download a single image from TMDB and return as numpy array (BGR). On error returns None."""
|
|
url = IMG_BASE + img_path
|
|
try:
|
|
r = requests.get(url, stream=True, timeout=10)
|
|
r.raise_for_status()
|
|
img_array = np.frombuffer(r.content, np.uint8)
|
|
img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
logging.warning(f"cv2.imdecode failed for downloaded image [{url}].")
|
|
return None
|
|
return img
|
|
except Exception as e:
|
|
logging.warning(f"Failed to download image [{url}]: {e}")
|
|
return None
|
|
|
|
|
|
def download_episode_images(
|
|
tmdb_id: str, seasons: list[int] | None = None
|
|
) -> ReferenceShowImages:
|
|
"""Loop through all seasons and episodes, downloading images. Given an TMDB id."""
|
|
show_details = tmdb_request(f"/tv/{tmdb_id}")
|
|
|
|
# Download images for seasons
|
|
season_episode_images = {}
|
|
for season in show_details.get("seasons", []):
|
|
season_number = season["season_number"]
|
|
if seasons is not None and season_number not in seasons:
|
|
continue
|
|
|
|
logging.info(f"Fetching season [{season_number}] images.")
|
|
season_episode_images[season_number] = {}
|
|
season_details = tmdb_request(f"/tv/{tmdb_id}/season/{season_number}")
|
|
|
|
# Download images for episodes
|
|
for episode in season_details.get("episodes", []):
|
|
ep_num = episode["episode_number"]
|
|
season_episode_images[season_number][ep_num] = []
|
|
|
|
# Fetch episode images
|
|
images = tmdb_request(
|
|
f"/tv/{tmdb_id}/season/{season_number}/episode/{ep_num}/images"
|
|
)
|
|
|
|
for idx, still in enumerate(images.get("stills", [])):
|
|
image = download_image(still["file_path"])
|
|
if image is not None:
|
|
season_episode_images[season_number][ep_num].append(image)
|
|
|
|
return ReferenceShowImages(season_episode_images)
|
|
|
|
|
|
# if __name__ == "__main__":
|
|
# tv_id = search_show(IMDB_ID, TITLE, YEAR)
|
|
# if tv_id:
|
|
# print(f"Found TMDB TV ID: {tv_id}")
|
|
# download_episode_images(tv_id)
|
|
# else:
|
|
# print("Could not find show on TMDB.")
|