Skip to content

Commit

Permalink
Merge branch 'development'
Browse files Browse the repository at this point in the history
  • Loading branch information
@Audionut
Audionut committed Dec 4, 2024
2 parents b31ccac + b06178c commit 09e870f
Show file tree
Hide file tree
Showing 6 changed files with 110 additions and 72 deletions.
11 changes: 11 additions & 0 deletions data/example-config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,16 @@
# Number of screenshots to capture
"screens": "6",

# Number of cutoff screenshots
# If there are at least this many screenshots already, perhaps pulled from existing
# description, skip creating and uploading any further screenshots.
"cutoff_screens": "3",

# multi processing task limit
# When capturing/optimizing images, limit to this many concurrent tasks
# defaults to 'os.cpu_count()'
"task_limit": "1",

# Providing the option to change the size of the screenshot thumbnails where supported.
# Default is 350, ie [img=350]
"thumbnail_size": "350",
Expand Down Expand Up @@ -152,6 +162,7 @@
"passkey": "HDB passkey",
"announce_url": "https://hdbits.org/announce/Custom_Announce_URL",
# "anon": False,
"img_rehost": True,
},
"HDT": {
"username": "username",
Expand Down
2 changes: 1 addition & 1 deletion src/args.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ def parse(self, args, meta):
parser.add_argument('-s', '--screens', nargs='*', required=False, help="Number of screenshots", default=int(self.config['DEFAULT']['screens']))
parser.add_argument('-mf', '--manual_frames', required=False, help="Comma-separated frame numbers to use as screenshots", type=str, default=None)
parser.add_argument('-c', '--category', nargs='*', required=False, help="Category [MOVIE, TV, FANRES]", choices=['movie', 'tv', 'fanres'])
parser.add_argument('-t', '--type', nargs='*', required=False, help="Type [DISC, REMUX, ENCODE, WEBDL, WEBRIP, HDTV, DVDRIP]", choices=['disc', 'remux', 'encode', 'webdl', 'web-dl', 'webrip', 'hdtv', 'dvdrip'])
parser.add_argument('-t', '--type', nargs='*', required=False, help="Type [DISC, REMUX, ENCODE, WEBDL, WEBRIP, HDTV, DVDRIP]", choices=['disc', 'remux', 'encode', 'webdl', 'web-dl', 'webrip', 'hdtv', 'dvdrip'], dest="manual_type")
parser.add_argument('--source', nargs='*', required=False, help="Source [Blu-ray, BluRay, DVD, HDDVD, WEB, HDTV, UHDTV, LaserDisc, DCP]", choices=['Blu-ray', 'BluRay', 'DVD', 'HDDVD', 'WEB', 'HDTV', 'UHDTV', 'LaserDisc', 'DCP'], dest="manual_source")
parser.add_argument('-res', '--resolution', nargs='*', required=False, help="Resolution [2160p, 1080p, 1080i, 720p, 576p, 576i, 480p, 480i, 8640p, 4320p, OTHER]", choices=['2160p', '1080p', '1080i', '720p', '576p', '576i', '480p', '480i', '8640p', '4320p', 'other'])
parser.add_argument('-tmdb', '--tmdb', nargs='*', required=False, help="TMDb ID", type=str, dest='tmdb_manual')
Expand Down
149 changes: 88 additions & 61 deletions src/prep.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -419,6 +419,10 @@ async def handle_image_list(self, meta, tracker_name):
console.print(f"[green]Images retained from {tracker_name}.")

async def gather_prep(self, meta, mode):
meta['cutoff'] = int(self.config['DEFAULT'].get('cutoff_screens', 3))
task_limit = self.config['DEFAULT'].get('task_limit', "0")
if int(task_limit) > 0:
meta['task_limit'] = task_limit
meta['mode'] = mode
base_dir = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
meta['isdir'] = os.path.isdir(meta['path'])
Expand Down Expand Up @@ -655,8 +659,7 @@ async def process_tracker(tracker_name, meta):
s.terminate()

meta['tmdb'] = meta.get('tmdb_manual', None)
if meta.get('type', None) is None:
meta['type'] = self.get_type(video, meta['scene'], meta['is_disc'])
meta['type'] = self.get_type(video, meta['scene'], meta['is_disc'], meta)
if meta.get('category', None) is None:
meta['category'] = self.get_cat(video)
else:
Expand Down Expand Up @@ -1164,6 +1167,7 @@ def is_scene(self, video, meta, imdb=None):
with open(nfo_file_path, 'wb') as f:
f.write(nfo_response.content)
meta['nfo'] = True
meta['auto_nfo'] = True
console.print(f"[green]NFO downloaded to {nfo_file_path}")
else:
console.print("[yellow]NFO file not available for download.")
Expand Down Expand Up @@ -1201,8 +1205,8 @@ def disc_screenshots(self, meta, filename, bdinfo, folder_id, base_dir, use_vs,
meta['image_list'] = []
existing_images = [img for img in meta['image_list'] if isinstance(img, dict) and img.get('img_url', '').startswith('http')]

if len(existing_images) >= 3 and not force_screenshots:
console.print("[yellow]There are already at least 3 images in the image list. Skipping additional screenshots.")
if len(existing_images) >= meta.get('cutoff') and not force_screenshots:
console.print("[yellow]There are already at least {} images in the image list. Skipping additional screenshots.".format(meta.get('cutoff')))
return

if num_screens is None:
Expand Down Expand Up @@ -1231,6 +1235,8 @@ def disc_screenshots(self, meta, filename, bdinfo, folder_id, base_dir, use_vs,
return

console.print("[bold yellow]Saving Screens...")
capture_results = []
task_limit = int(meta.get('task_limit', os.cpu_count()))

if use_vs:
from src.vs import vs_screengn
Expand All @@ -1253,24 +1259,24 @@ def disc_screenshots(self, meta, filename, bdinfo, folder_id, base_dir, use_vs,
for i in range(num_screens + 1)
]

with Pool(processes=min(len(capture_tasks), os.cpu_count())) as pool:
with Pool(processes=min(len(capture_tasks), task_limit)) as pool:
capture_results = list(
tqdm(
pool.imap_unordered(self.capture_disc_task, capture_tasks),
total=len(capture_tasks),
desc="Capturing Screenshots"
)
)

if len(capture_results) > num_screens:
smallest = min(capture_results, key=os.path.getsize)
if meta['debug']:
console.print(f"[yellow]Removing smallest image: {smallest} ({os.path.getsize(smallest)} bytes)[/yellow]")
os.remove(smallest)
capture_results.remove(smallest)
if capture_results:
if len(capture_results) > num_screens:
smallest = min(capture_results, key=os.path.getsize)
if meta['debug']:
console.print(f"[yellow]Removing smallest image: {smallest} ({os.path.getsize(smallest)} bytes)[/yellow]")
os.remove(smallest)
capture_results.remove(smallest)

optimize_tasks = [(result, self.config) for result in capture_results if result and os.path.exists(result)]
with Pool(processes=min(len(optimize_tasks), os.cpu_count())) as pool:
with Pool(processes=min(len(optimize_tasks), task_limit)) as pool:
optimized_results = list(
tqdm(
pool.imap_unordered(self.optimize_image_task, optimize_tasks),
Expand Down Expand Up @@ -1336,8 +1342,8 @@ def dvd_screenshots(self, meta, disc_num, num_screens=None):
meta['image_list'] = []
existing_images = [img for img in meta['image_list'] if isinstance(img, dict) and img.get('img_url', '').startswith('http')]

if len(existing_images) >= 3:
console.print("[yellow]There are already at least 3 images in the image list. Skipping additional screenshots.")
if len(existing_images) >= meta.get('cutoff'):
console.print("[yellow]There are already at least {} images in the image list. Skipping additional screenshots.".format(meta.get('cutoff')))
return

if num_screens is None:
Expand Down Expand Up @@ -1422,12 +1428,13 @@ def _is_vob_good(n, loops, num_screens):
voblength, n = _is_vob_good(0, 0, num_screens)
ss_times = self.valid_ss_time([], num_screens + 1, voblength)
tasks = []
task_limit = int(meta.get('task_limit', os.cpu_count()))
for i in range(num_screens + 1):
image = f"{meta['base_dir']}/tmp/{meta['uuid']}/{meta['discs'][disc_num]['name']}-{i}.png"
input_file = f"{meta['discs'][disc_num]['path']}/VTS_{main_set[i % len(main_set)]}"
tasks.append((input_file, image, ss_times[i], meta, width, height, w_sar, h_sar))

with Pool(processes=min(num_screens + 1, os.cpu_count())) as pool:
with Pool(processes=min(num_screens + 1, task_limit)) as pool:
results = list(tqdm(pool.imap_unordered(self.capture_dvd_screenshot, tasks), total=len(tasks), desc="Capturing Screenshots"))

if len(glob.glob1(f"{meta['base_dir']}/tmp/{meta['uuid']}/", f"{meta['discs'][disc_num]['name']}-*")) > num_screens:
Expand All @@ -1451,7 +1458,7 @@ def _is_vob_good(n, loops, num_screens):

optimize_tasks = [(image, self.config) for image in results if image and os.path.exists(image)]

with Pool(processes=min(len(optimize_tasks), os.cpu_count())) as pool:
with Pool(processes=min(len(optimize_tasks), task_limit)) as pool:
optimize_results = list( # noqa F841
tqdm(
pool.imap_unordered(self.optimize_image_task, optimize_tasks),
Expand Down Expand Up @@ -1486,8 +1493,8 @@ def screenshots(self, path, filename, folder_id, base_dir, meta, num_screens=Non

existing_images = [img for img in meta['image_list'] if isinstance(img, dict) and img.get('img_url', '').startswith('http')]

if len(existing_images) >= 3 and not force_screenshots:
console.print("[yellow]There are already at least 3 images in the image list. Skipping additional screenshots.")
if len(existing_images) >= meta.get('cutoff') and not force_screenshots:
console.print("[yellow]There are already at least {} images in the image list. Skipping additional screenshots.".format(meta.get('cutoff')))
return

if num_screens is None:
Expand Down Expand Up @@ -1534,35 +1541,42 @@ def screenshots(self, path, filename, folder_id, base_dir, meta, num_screens=Non
else:
ss_times = self.valid_ss_time([], num_screens + 1, length)

capture_tasks = []
capture_results = []
task_limit = int(meta.get('task_limit', os.cpu_count()))

capture_tasks = []
for i in range(num_screens + 1):
image_path = os.path.abspath(f"{base_dir}/tmp/{folder_id}/{filename}-{i}.png")
if not os.path.exists(image_path) or meta.get('retake', False):
capture_tasks.append((path, ss_times[i], image_path, width, height, w_sar, h_sar, loglevel))
else:
if meta['debug']:
console.print(f"[yellow]Skipping existing screenshot: {image_path}")
elif meta['debug']:
console.print(f"[yellow]Skipping existing screenshot: {image_path}")

if not capture_tasks:
console.print("[yellow]All screenshots already exist. Skipping capture process.")
else:
capture_results = []
with Pool(processes=len(capture_tasks)) as pool:
for result in tqdm(pool.imap_unordered(self.capture_screenshot, capture_tasks), total=len(capture_tasks), desc="Capturing Screenshots"):
with Pool(processes=min(len(capture_tasks), task_limit)) as pool:
for result in tqdm(pool.imap_unordered(self.capture_screenshot, capture_tasks),
total=len(capture_tasks),
desc="Capturing Screenshots"):
capture_results.append(result)

if len(capture_results) > num_screens:
smallest = min(capture_results, key=os.path.getsize)
if meta['debug']:
console.print(f"[yellow]Removing smallest image: {smallest} ({os.path.getsize(smallest)} bytes)[/yellow]")
os.remove(smallest)
capture_results.remove(smallest)
if capture_results and len(capture_results) > num_screens:
smallest = min(capture_results, key=os.path.getsize)
if meta['debug']:
console.print(f"[yellow]Removing smallest image: {smallest} ({os.path.getsize(smallest)} bytes)[/yellow]")
os.remove(smallest)
capture_results.remove(smallest)

optimize_tasks = [(result, self.config) for result in capture_results if "Error" not in result]
optimize_results = []
with Pool(processes=len(optimize_tasks)) as pool:
for result in tqdm(pool.imap_unordered(self.optimize_image_task, optimize_tasks), total=len(optimize_tasks), desc="Optimizing Images"):
optimize_results.append(result)
if optimize_tasks:
with Pool(processes=min(len(optimize_tasks), task_limit)) as pool:
for result in tqdm(pool.imap_unordered(self.optimize_image_task, optimize_tasks),
total=len(optimize_tasks),
desc="Optimizing Images"):
optimize_results.append(result)

valid_results = []
for image_path in optimize_results:
Expand Down Expand Up @@ -1681,25 +1695,28 @@ def optimize_image_task(self, args):
Get type and category
"""

def get_type(self, video, scene, is_disc):
filename = os.path.basename(video).lower()
if "remux" in filename:
type = "REMUX"
elif any(word in filename for word in [" web ", ".web.", "web-dl"]):
type = "WEBDL"
elif "webrip" in filename:
type = "WEBRIP"
# elif scene == True:
# type = "ENCODE"
elif "hdtv" in filename:
type = "HDTV"
elif is_disc is not None:
type = "DISC"
elif "dvdrip" in filename:
type = "DVDRIP"
# exit()
def get_type(self, video, scene, is_disc, meta):
if meta.get('manual_type'):
type = meta.get('manual_type')
else:
type = "ENCODE"
filename = os.path.basename(video).lower()
if "remux" in filename:
type = "REMUX"
elif any(word in filename for word in [" web ", ".web.", "web-dl", "webdl"]):
type = "WEBDL"
elif "webrip" in filename:
type = "WEBRIP"
# elif scene == True:
# type = "ENCODE"
elif "hdtv" in filename:
type = "HDTV"
elif is_disc is not None:
type = "DISC"
elif "dvdrip" in filename:
type = "DVDRIP"
# exit()
else:
type = "ENCODE"
return type

def get_cat(self, video):
Expand Down Expand Up @@ -2988,7 +3005,7 @@ def upload_screens(self, meta, screens, img_host_num, i, total_screens, custom_i

# Define host-specific limits
host_limits = {
"imgbox": 6,
"imgbb": 1,
# Other hosts can use the default pool size
}
default_pool_size = os.cpu_count()
Expand Down Expand Up @@ -3079,7 +3096,7 @@ async def imgbox_upload(self, chdir, image_glob, meta, return_dict):
return []

async def get_name(self, meta):
type = meta.get('type', "")
type = meta.get('type', "").upper()
title = meta.get('title', "")
alt_title = meta.get('aka', "")
year = meta.get('year', "")
Expand Down Expand Up @@ -3206,6 +3223,7 @@ async def get_name(self, meta):
console.print(f"--category [yellow]{meta['category']}")
console.print(f"--type [yellow]{meta['type']}")
console.print(f"--source [yellow]{meta['source']}")
console.print("[bold green]If you specified type, try also specifying source")

exit()
name_notag = name
Expand Down Expand Up @@ -3562,25 +3580,34 @@ def clean_text(text):
uuid = meta['uuid']
current_dir_path = "*.nfo"
specified_dir_path = os.path.join(base_dir, "tmp", uuid, "*.nfo")

if meta['debug']:
console.print(f"specified_dir_path: {specified_dir_path}")
if meta.get('nfo') and not content_written:
nfo_files = glob.glob(current_dir_path)
if not nfo_files:
if meta['auto_nfo'] is True:
nfo_files = glob.glob(specified_dir_path)
scene_nfo = True
else:
nfo_files = glob.glob(current_dir_path)
if meta['debug']:
console.print(f"Glob current_dir_path matches: {glob.glob(current_dir_path)}")
console.print(f"Glob specified_dir_path matches: {glob.glob(specified_dir_path)}")
if not nfo_files:
console.print("NFO was set but no nfo file was found")
description.write("\n")
return meta

if nfo_files:
console.print("We found nfo")
nfo = nfo_files[0]
try:
with open(nfo, 'r', encoding="utf-8") as nfo_file:
nfo_content = nfo_file.read()
console.print("NFO content read with utf-8 encoding.")
if meta['debug']:
console.print("NFO content read with utf-8 encoding.")
except UnicodeDecodeError:
console.print("utf-8 decoding failed, trying latin1.")
if meta['debug']:
console.print("utf-8 decoding failed, trying latin1.")
with open(nfo, 'r', encoding="latin1") as nfo_file:
nfo_content = nfo_file.read()
console.print("NFO content read with latin1 encoding.")

if scene_nfo is True:
description.write(f"[center][spoiler=Scene NFO:][code]{nfo_content}[/code][/spoiler][/center]\n")
Expand Down
2 changes: 1 addition & 1 deletion src/trackers/HDB.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ def __init__(self, config):
self.source_flag = 'HDBits'
self.username = config['TRACKERS']['HDB'].get('username', '').strip()
self.passkey = config['TRACKERS']['HDB'].get('passkey', '').strip()
self.rehost_images = config['TRACKERS']['HDB'].get('img_rehost', False)
self.rehost_images = config['TRACKERS']['HDB'].get('img_rehost', True)
self.signature = None
self.banned_groups = [""]

Expand Down
Loading

0 comments on commit 09e870f

Please sign in to comment.