telegram-thumbnail-uploader/utils.py

212 lines
4.8 KiB
Python

import os
import re
import time
import logging
from pathlib import Path
from PIL import Image
logger = logging.getLogger(__name__)
BASE_DIR = os.path.dirname(__file__)
SINGLE_DIR = os.path.join(BASE_DIR, "single")
BATCH_DIR = os.path.join(BASE_DIR, "batch")
TEMP_DIR = os.path.join(BASE_DIR, "temp")
os.makedirs(TEMP_DIR, exist_ok=True)
VIDEO_FORMATS = (
".mp4", ".mkv", ".webm", ".avi", ".mov",
".m4v", ".flv", ".wmv", ".ts",
".mpg", ".mpeg", ".3gp", ".mp4v", ".vob"
)
THUMB_FORMATS = (".jpg", ".jpeg", ".png")
def find_matching_thumbnail(video_path):
video_name = Path(video_path).stem.lower()
# Exact match
for ext in THUMB_FORMATS:
exact_match = os.path.join(SINGLE_DIR, f"{video_name}{ext}")
if os.path.exists(exact_match):
logger.info(f"Exact thumbnail found: {exact_match}")
return exact_match
video_keywords = video_name.split(".")
for file in os.listdir(SINGLE_DIR):
if file.lower().endswith(THUMB_FORMATS):
thumb_name = Path(file).stem.lower()
if (
thumb_name in video_name
or video_name in thumb_name
or any(keyword in thumb_name for keyword in video_keywords)
):
return os.path.join(SINGLE_DIR, file)
return None
def get_local_videos():
video_files = []
logger.info(f"Scanning single folder: {SINGLE_DIR}")
videos = [
f for f in os.listdir(SINGLE_DIR)
if f.lower().endswith(VIDEO_FORMATS)
]
videos.sort()
for video in videos:
video_path = os.path.join(SINGLE_DIR, video)
thumb_path = find_matching_thumbnail(video_path)
video_files.append({
"video_path": video_path,
"thumb_path": thumb_path,
"filename": video,
"size": os.path.getsize(video_path),
"has_thumb": thumb_path is not None
})
logger.info(
f"Found video: {video} {'with thumbnail' if thumb_path else 'no thumbnail'}"
)
return video_files, {}
def process_local_thumbnail(thumb_path):
try:
img = Image.open(thumb_path).convert("RGB")
width, height = img.size
if width > height:
new_width = 320
new_height = int(320 * height / width)
else:
new_height = 320
new_width = int(320 * width / height)
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
output_path = os.path.join(
TEMP_DIR,
f"thumb_{int(time.time())}.jpg"
)
img.save(output_path, "JPEG", quality=95)
return output_path
except Exception as e:
logger.error(f"Thumbnail processing error: {e}")
return None
def find_batch_thumbnail(batch_dir):
for file in os.listdir(batch_dir):
if file.lower().endswith(THUMB_FORMATS):
return os.path.join(batch_dir, file)
return None
def extract_episode_number(filename):
patterns = [
r'[Ss](\d+)[Ee](\d+)',
r'[Ee][Pp]?\.?\s*(\d+)',
r'E(\d+)',
r'(\d+)'
]
filename = filename.lower()
for pattern in patterns:
match = re.search(pattern, filename)
if match:
try:
if len(match.groups()) == 2:
return int(match.group(2))
return int(match.group(1))
except:
continue
return 0
def get_batch_videos(batch_dir):
video_files = []
thumb_path = find_batch_thumbnail(batch_dir)
videos = []
for f in os.listdir(batch_dir):
if f.lower().endswith(VIDEO_FORMATS):
ep = extract_episode_number(f)
videos.append((f, ep))
videos.sort(key=lambda x: x[1])
for video, ep_num in videos:
video_path = os.path.join(batch_dir, video)
video_files.append({
"video_path": video_path,
"thumb_path": thumb_path,
"filename": video,
"episode": ep_num,
"size": os.path.getsize(video_path)
})
logger.info(f"Found video {video} episode {ep_num}")
return video_files
def natural_sort_key(s):
return [
int(text) if text.isdigit() else text.lower()
for text in re.split(r'([0-9]+)', s)
]
def get_batch_directories():
batch_dirs = []
if not os.path.exists(BATCH_DIR):
return batch_dirs
has_root_videos = False
for item in sorted(os.listdir(BATCH_DIR), key=natural_sort_key):
full = os.path.join(BATCH_DIR, item)
if os.path.isdir(full):
batch_dirs.append(full)
elif item.lower().endswith(VIDEO_FORMATS):
has_root_videos = True
if has_root_videos:
batch_dirs.insert(0, BATCH_DIR)
return batch_dirs