diff --git a/data/example-config.py b/data/example-config.py index b6342c58..989c2d54 100644 --- a/data/example-config.py +++ b/data/example-config.py @@ -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", @@ -152,6 +162,7 @@ "passkey": "HDB passkey", "announce_url": "https://hdbits.org/announce/Custom_Announce_URL", # "anon": False, + "img_rehost": True, }, "HDT": { "username": "username", diff --git a/src/args.py b/src/args.py index c4057e76..5c348cfc 100644 --- a/src/args.py +++ b/src/args.py @@ -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') diff --git a/src/prep.py b/src/prep.py index ae04ca52..43a7916a 100644 --- a/src/prep.py +++ b/src/prep.py @@ -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']) @@ -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: @@ -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.") @@ -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: @@ -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 @@ -1253,7 +1259,7 @@ 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), @@ -1261,16 +1267,16 @@ def disc_screenshots(self, meta, filename, bdinfo, folder_id, base_dir, use_vs, 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), @@ -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: @@ -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: @@ -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), @@ -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: @@ -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: @@ -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): @@ -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() @@ -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', "") @@ -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 @@ -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") diff --git a/src/trackers/HDB.py b/src/trackers/HDB.py index a59c42e8..e7f0b21f 100644 --- a/src/trackers/HDB.py +++ b/src/trackers/HDB.py @@ -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 = [""] diff --git a/src/trackers/MTV.py b/src/trackers/MTV.py index 9d8cbe84..84f402b4 100644 --- a/src/trackers/MTV.py +++ b/src/trackers/MTV.py @@ -62,7 +62,7 @@ async def upload_with_retry(self, meta, cookiefile, common, img_host_index=1): img_host_index += 1 continue - new_images_key = f'mtv_images_key_{img_host_index}' + new_images_key = 'mtv_images_key' if image_list is not None: image_list = meta[new_images_key] break @@ -117,7 +117,7 @@ async def upload_with_retry(self, meta, cookiefile, common, img_host_index=1): source_id = await self.get_source_id(meta) origin_id = await self.get_origin_id(meta) des_tags = await self.get_tags(meta) - await self.edit_desc(meta, images_reuploaded, image_list) + await self.edit_desc(meta) group_desc = await self.edit_group_desc(meta) mtv_name = await self.edit_name(meta) @@ -183,7 +183,7 @@ async def handle_image_upload(self, meta, img_host_index=1, approved_image_hosts retry_mode = False images_reuploaded = False - new_images_key = f'mtv_images_key_{img_host_index}' + new_images_key = 'mtv_images_key' discs = meta.get('discs', []) # noqa F841 filelist = meta.get('video', []) filename = meta['filename'] @@ -267,7 +267,7 @@ async def handle_image_upload(self, meta, img_host_index=1, approved_image_hosts return meta[new_images_key], False, images_reuploaded # Return retry_mode and images_reuploaded - async def edit_desc(self, meta, images_reuploaded, valid_images): + async def edit_desc(self, meta): base = open(f"{meta['base_dir']}/tmp/{meta['uuid']}/DESCRIPTION.txt", 'r', encoding='utf-8').read() with open(f"{meta['base_dir']}/tmp/{meta['uuid']}/[{self.tracker}]DESCRIPTION.txt", 'w', encoding='utf-8') as desc: @@ -283,8 +283,8 @@ async def edit_desc(self, meta, images_reuploaded, valid_images): elif mi_dump: desc.write("[mediainfo]" + mi_dump + "[/mediainfo]\n\n") - if valid_images: - images = valid_images + if meta['mtv_images_key']: + images = meta['mtv_images_key'] else: images = meta['image_list'] if len(images) > 0: diff --git a/upload.py b/upload.py index d34bc4d0..38a69c59 100644 --- a/upload.py +++ b/upload.py @@ -198,7 +198,7 @@ def merge_meta(meta, saved_meta, path): 'trackers', 'dupe', 'debug', 'anon', 'category', 'type', 'screens', 'nohash', 'manual_edition', 'imdb', 'tmdb_manual', 'mal', 'manual', 'hdb', 'ptp', 'blu', 'no_season', 'no_aka', 'no_year', 'no_dub', 'no_tag', 'no_seed', 'client', 'desclink', 'descfile', 'desc', 'draft', 'modq', 'region', 'freeleech', 'personalrelease', 'unattended', 'manual_season', 'manual_episode', 'torrent_creation', 'qbit_tag', 'qbit_cat', - 'skip_imghost_upload', 'imghost', 'manual_source', 'webdv', 'hardcoded-subs', 'dual_audio' + 'skip_imghost_upload', 'imghost', 'manual_source', 'webdv', 'hardcoded-subs', 'dual_audio', 'manual_type' ] sanitized_saved_meta = {} for key, value in saved_meta.items(): @@ -254,8 +254,8 @@ async def process_meta(meta, base_dir): with open(f"{meta['base_dir']}/tmp/{meta['uuid']}/meta.json", 'w') as f: json.dump(meta, f, indent=4) meta['name_notag'], meta['name'], meta['clean_name'], meta['potential_missing'] = await prep.get_name(meta) - - if len(meta.get('image_list', [])) < 3 and meta.get('skip_imghost_upload', False) is False: + meta['cutoff'] = int(config['DEFAULT'].get('cutoff_screens', 3)) + if len(meta.get('image_list', [])) < meta.get('cutoff') and meta.get('skip_imghost_upload', False) is False: if 'image_list' not in meta: meta['image_list'] = [] return_dict = {}