Refactor code to be shorter

Client/Extract Pixiv/user.py

@@ -11,10 +11,6 @@ else:
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 
 inp = input("Enter User IDs (separated by spaces or commas): ")
-inps='1'
-while inps != "":
-    inps = input('')
-    inp += inps
 user_ids = re.findall(r"\d+", inp)
 
 user_ids = [int(uid) for uid in user_ids]

Client/actual_size.py

@@ -1,8 +0,0 @@
-import lmdb
-import os
-
-os.chdir(os.path.dirname(os.path.abspath(__file__)))
-db = lmdb.open("db", subdir=True, map_size=524288)
-
-print(db.stat())
-print(db.info())

Client/hunt.py

@@ -5,39 +5,32 @@ from concurrent.futures import ThreadPoolExecutor
 from tqdm import tqdm
 
 local = 0
-if local:
-    endpoint = "http://127.0.0.1:7860"
-else:
-    endpoint = "https://q6-p.hf.space"
-
+endpoint = "http://127.0.0.1:7860" if local else "https://q6-p.hf.space"
 img_base = 'https://i.pximg.net/img-original/img/'
 
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 os.makedirs("images/Stash", exist_ok=True)
 
-images_cache = os.listdir("images/Stash")
+images_cache = set(os.listdir("images/Stash"))
 
-db = lmdb.open("db", subdir=True, map_size=1048576)
-valid = [f for f in os.listdir() if f.endswith(".txt")]
+db = lmdb.open("db", subdir=True, map_size=1048576 * 2)
 
+valid = [f for f in os.listdir() if f.endswith(".txt")]
 for idx, file in enumerate(valid):
     print(f"{idx + 1}: {file}")
 
-inputs = input("Enter the index of the file: ")
-inputs = inputs.split()
+inputs = input("Enter the index of the file: ").split()
 indexs = []
-
 for inp in inputs:
-    if inp.isdigit():
-        indexs.append(int(inp) - 1)
-    elif "-" in inp:
+    if "-" in inp:
         start, end = map(int, inp.split("-"))
         indexs.extend(range(start - 1, end))
+    elif inp.isdigit():
+        indexs.append(int(inp) - 1)
 
 def download_image(args):
     post_id, url = args
-    full_url = img_base + url
-    response = requests.get(full_url, headers={
+    response = requests.get(img_base + url, headers={
         "Referer": "https://www.pixiv.net/",
         "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:126.0) Gecko/20100101 Firefox/126.0"
     })
@@ -45,54 +38,48 @@ def download_image(args):
         img_file.write(response.content)
     return post_id, True
 
+def decode_if_binary(val):
+    if type(val) == bytes:
+        return val.decode()
+    return val
+
 for index in indexs:
     group_name = valid[index].rsplit(".", 1)[0]
     os.makedirs(f"images/{group_name}", exist_ok=True)
     with open(valid[index], "r") as f:
-        post_ids = [x for x in f.read().split("\n") if x]
+        post_ids = f.read().split()
 
     with db.begin(write=True) as txn:
         post_ids_dict = {post_id: txn.get(post_id.encode()) for post_id in post_ids}
-        post_ids_dict = {post_id: url.decode() if url != None else None for post_id, url in post_ids_dict.items()}
+        filtered = [post_id for post_id in post_ids if post_ids_dict[post_id] == None and f"{post_id}.png" not in images_cache]
+        print(f"Group: {group_name}\nFiltered: {len(filtered)}/{len(post_ids)}")
 
-        filtered = [post_id for post_id, url in post_ids_dict.items() if url == None and f"{post_id}.png" not in images_cache]
-        print(f"Group: {group_name}\nFiltered: {len(filtered)}/{len(post_ids_dict)}")
         if filtered:
             data = requests.post(f'{endpoint}/pixif', json={"post_ids": filtered}).json()
-
             for post_id, url in data.items():
                 txn.put(post_id.encode(), url.encode())
                 post_ids_dict[post_id] = url
-
             no_exif = set(filtered) - set(data.keys())
             for post_id in no_exif:
                 txn.put(post_id.encode(), b'')
 
-    to_download = {post_id: url for post_id, url in post_ids_dict.items() 
-                   if url and f"{post_id}.png" not in images_cache}
-
-    print(f"Total images to download: {len(to_download)}")
+    to_download = {post_id: decode_if_binary(url) for post_id, url in post_ids_dict.items() if url and f"{post_id}.png" not in images_cache}
+    
+    if to_download:
+        print(f"Total images to download: {len(to_download)}")
+        with ThreadPoolExecutor(max_workers=20) as executor:
+            list(tqdm(executor.map(download_image, to_download.items()), total=len(to_download), desc="Downloading"))
 
-    max_workers = 30
-    to_download_items = list(to_download.items())
-
-    if to_download_items:
-        print("Starting download...")
-        with ThreadPoolExecutor(max_workers=max_workers) as executor:
-            for result in tqdm(executor.map(download_image, to_download_items), total=len(to_download_items), desc="Downloading"):
-                pass
-
-    images_cache = os.listdir("images/Stash")
+    images_cache.update(os.listdir("images/Stash"))
 
     print("Linking images to the group directory...")
-    for i, post_id in enumerate(post_ids_dict.keys()):
-        if f"{post_id}.png" in images_cache:
-            stash_path = f"images/Stash/{post_id}.png"
-            dest_path = f"images/{group_name}/{i}_{post_id}.png"
-            if not os.path.exists(dest_path):
-                os.link(stash_path, dest_path)
-
-    if len(os.listdir(f'images/{group_name}')) == 0:
+    for i, post_id in enumerate(post_ids):
+        stash_path = f"images/Stash/{post_id}.png"
+        dest_path = f"images/{group_name}/{i}_{post_id}.png"
+        if os.path.exists(stash_path) and not os.path.exists(dest_path):
+            os.link(stash_path, dest_path)
+
+    if not os.listdir(f'images/{group_name}'):
         os.rmdir(f"images/{group_name}")
 
 db.close()

Client/show_all.py

@@ -3,8 +3,7 @@ import lmdb
 
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 
-db = lmdb.open("db", subdir=True, map_size=1048576)
-
+db = lmdb.open("db", subdir=True, map_size=1048576 * 2)
 
 items = []
 
@@ -12,18 +11,14 @@ with db.begin() as txn:
     cursor = txn.cursor()
     
     for key, value in cursor:
-        if value != b'\x00':
-            try:
-                key_str = key.decode('utf-8')
-                value_str = value.decode('utf-8')
-                
-                int_part_str = value_str.split('p')[1].split('.')[0]
-                int_part = int(int_part_str)
-                items.append((key_str, int_part))
+        if value != b'':
+            key_str = key.decode('utf-8')
+            value_str = value.decode('utf-8')
+            
+            int_part_str = value_str.split('p')[1].split('.')[0]
+            int_part = int(int_part_str)
+            items.append((key_str, int_part))
             
-            except (IndexError, ValueError) as e:
-                print(f"Skipping key {key} due to parsing error: {e}")
-                continue
 
 sorted_items = sorted(items, key=lambda item: item[1])