79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
from os import environ, makedirs, path, walk, listdir
|
|
from shutil import move
|
|
from dotenv import load_dotenv
|
|
from base64 import b64decode
|
|
import re
|
|
import requests
|
|
|
|
load_dotenv()
|
|
|
|
DOWNLOAD_PATH=environ.get("DOWNLOAD_PATH")
|
|
TESTING_PATH=environ.get("TESTING_PATH")
|
|
TRAINING_PATH=environ.get("TRAINING_PATH")
|
|
|
|
def prepare_dirs():
|
|
makedirs(DOWNLOAD_PATH, exist_ok=True)
|
|
makedirs(TESTING_PATH, exist_ok=True)
|
|
makedirs(TRAINING_PATH, exist_ok=True)
|
|
|
|
def fetch_captcha(id):
|
|
print(f"Fetching captcha with id {id}")
|
|
captcha = requests.get(f"{environ.get('CAPTCHA_AGGREGATOR_API')}/captcha/{id}").json()["captcha"]
|
|
|
|
with open(f"{DOWNLOAD_PATH}/{captcha['hash']}_{captcha['solution']}.jpeg", 'wb') as captcha_file:
|
|
captcha_file.write(b64decode(captcha['image']))
|
|
|
|
def search_saved_captcha(hash, path):
|
|
print(f"searching captcha with hash {hash} in {path}")
|
|
regex = re.compile(hash + '_\\w{6}\\.jpeg')
|
|
|
|
for _, _, files in walk(path):
|
|
for file in files:
|
|
if regex.match(file):
|
|
return True
|
|
return False
|
|
|
|
def search_and_download_new(captchas):
|
|
print(f"Searching and downloading new captchas")
|
|
for captcha in captchas:
|
|
id = captcha["id"]
|
|
hash = captcha["hash"]
|
|
training_exists = search_saved_captcha(hash, TRAINING_PATH)
|
|
testing_exists = search_saved_captcha(hash, TESTING_PATH)
|
|
new_exists = search_saved_captcha(hash, DOWNLOAD_PATH)
|
|
if not training_exists and not testing_exists and not new_exists:
|
|
fetch_captcha(id)
|
|
|
|
def sort_datasets():
|
|
print(f"Sorting datasets")
|
|
percent_of_testing = int(environ.get("PERCENT_OF_TESTING"))
|
|
amount_of_new_data = len([file for file in listdir(DOWNLOAD_PATH) if path.isfile(f'{DOWNLOAD_PATH}/{file}')])
|
|
print(amount_of_new_data)
|
|
amount_to_send_to_test = round(amount_of_new_data * (percent_of_testing / 100))
|
|
print(amount_to_send_to_test)
|
|
for _, _, files in walk(DOWNLOAD_PATH):
|
|
for index, file in enumerate(files):
|
|
if index < amount_to_send_to_test:
|
|
move(f"{DOWNLOAD_PATH}/{file}", TESTING_PATH)
|
|
else:
|
|
move(f"{DOWNLOAD_PATH}/{file}", TRAINING_PATH)
|
|
|
|
def download_dataset():
|
|
prepare_dirs()
|
|
|
|
captchas = requests.get(f"{environ.get('CAPTCHA_AGGREGATOR_API')}/captcha/all").json()["captchas"]
|
|
|
|
search_and_download_new(captchas)
|
|
sort_datasets()
|
|
|
|
|
|
|
|
def train_nn():
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
download_dataset()
|
|
train_nn()
|
|
|
|
|