107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import secrets
|
|
from werkzeug.security import generate_password_hash
|
|
from duckduckgo_search import DDGS
|
|
import requests
|
|
import shutil
|
|
|
|
def set_login_info(username, password):
|
|
salt = secrets.token_hex(16)
|
|
password_hash = generate_password_hash(password + salt)
|
|
key = secrets.token_hex(32)
|
|
|
|
login_info = {
|
|
"login_info": {
|
|
"username": username,
|
|
"password_hash": password_hash,
|
|
"salt": salt
|
|
},
|
|
"key": key
|
|
}
|
|
|
|
os.makedirs('db', exist_ok=True)
|
|
with open('db/keys.json', 'w') as f:
|
|
json.dump(login_info, f, indent=4)
|
|
|
|
def get_next_batch_number():
|
|
existing_files = os.listdir('images')
|
|
batch_numbers = [int(f.split('_')[0]) for f in existing_files if '_' in f]
|
|
return max(batch_numbers, default=0) + 1
|
|
|
|
def download_images(query, total_images=100, batch_size=10):
|
|
os.makedirs('images', exist_ok=True)
|
|
ddgs = DDGS()
|
|
results = ddgs.images(query, max_results=total_images)
|
|
batch_number = get_next_batch_number()
|
|
downloaded_images = 0
|
|
valid_extensions = {'jpg', 'jpeg', 'png', 'gif'}
|
|
|
|
for i, result in enumerate(results):
|
|
image_url = result['image']
|
|
try:
|
|
response = requests.get(image_url)
|
|
if response.status_code == 200:
|
|
ext = image_url.split('.')[-1].lower()
|
|
if ext not in valid_extensions:
|
|
ext = 'jpg'
|
|
filename = f'{batch_number:04d}_{(downloaded_images % batch_size) + 1:05d}.{ext}'
|
|
with open(f'images/{filename}', 'wb') as f:
|
|
f.write(response.content)
|
|
downloaded_images += 1
|
|
else:
|
|
print(f"Failed to download image from {image_url}. Status code: {response.status_code}")
|
|
except requests.RequestException as e:
|
|
print(f"An error occurred while downloading image from {image_url}: {e}")
|
|
except Exception as e:
|
|
print(f"An unexpected error occurred: {e}")
|
|
|
|
if downloaded_images > 0 and downloaded_images % batch_size == 0:
|
|
batch_number += 1
|
|
|
|
print(f"Downloaded {downloaded_images} images for query '{query}'.")
|
|
|
|
def clear_cache():
|
|
cache_folder = 'cache'
|
|
if os.path.exists(cache_folder):
|
|
shutil.rmtree(cache_folder)
|
|
print(f"Cache directory '{cache_folder}' removed successfully.")
|
|
else:
|
|
print(f"Cache directory '{cache_folder}' does not exist.")
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='Set username and password for the application and optionally download images using DuckDuckGo.',
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
)
|
|
|
|
# Add an argument to enter user setting mode
|
|
parser.add_argument('--user-setting', action='store_true', help='Enter user setting mode to set username and password.')
|
|
|
|
# Create a userset argument group
|
|
userset_group = parser.add_argument_group('userset', 'Set username and password')
|
|
userset_group.add_argument('-u', '--username', type=str, help='The username to set.')
|
|
userset_group.add_argument('-p', '--password', type=str, help='The password to set.')
|
|
|
|
parser.add_argument('-q', '--query', type=str, help='The search query to download images for testing.')
|
|
parser.add_argument('-t', '--total_images', type=int, default=10, help='The total number of images to download.')
|
|
parser.add_argument('-b', '--batch_size', type=int, default=5, help='The number of images to download in each batch.')
|
|
parser.add_argument('-c', '--clear-cache', action='store_true', help='Remove the cache directory.')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.user_setting:
|
|
if not args.username or not args.password:
|
|
parser.error('--user-setting requires --username and --password.')
|
|
set_login_info(args.username, args.password)
|
|
print(f"Username and password set successfully. Key generated and saved in db/keys.json.")
|
|
elif args.clear_cache:
|
|
clear_cache()
|
|
elif args.query:
|
|
download_images(args.query, args.total_images, args.batch_size)
|
|
else:
|
|
parser.print_help()
|
|
|
|
if __name__ == '__main__':
|
|
main() |