377 lines
14 KiB
Python
377 lines
14 KiB
Python
import sys
|
|
import os
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import argparse
|
|
import glob
|
|
#import requests_async as requests
|
|
import asyncio
|
|
import aiohttp
|
|
from typing import IO
|
|
import aiofiles
|
|
import re
|
|
from colorama import Fore, Style
|
|
from PIL import Image
|
|
import io
|
|
|
|
# can tweak these you feel like it, but shouldn't be needed
|
|
unsafe_threshhold = 0.1 # higher values is more likely to be nsfw and will be skipped
|
|
aesthetic_threshhold = 5 # higher is more aesthetic, note laion2B-aesthetic already >7?
|
|
http_timeout = 10
|
|
|
|
# dont touch
|
|
downloaded_count = 0
|
|
current_parquet_file_downloaded_count = 0
|
|
logger_sp = None
|
|
|
|
def get_base_prefix_compat():
|
|
"""Get base/real prefix, or sys.prefix if there is none."""
|
|
return getattr(sys, "base_prefix", None) or getattr(sys, "real_prefix", None) or sys.prefix
|
|
|
|
def in_virtualenv():
|
|
return get_base_prefix_compat() != sys.prefix
|
|
|
|
def get_parser(**parser_kwargs):
|
|
parser = argparse.ArgumentParser(**parser_kwargs)
|
|
do_not_download = True
|
|
parser.add_argument(
|
|
"--laion_dir",
|
|
type=str,
|
|
nargs="?",
|
|
const=True,
|
|
default="./laion",
|
|
help="directory with laion parquet files, default is ./laion",
|
|
)
|
|
parser.add_argument(
|
|
"--search_text",
|
|
type=str,
|
|
const=True,
|
|
nargs="?",
|
|
default=None,
|
|
help="csv of words with AND logic, ex \"photo,man,dog\"",
|
|
),
|
|
parser.add_argument(
|
|
"--out_dir",
|
|
type=str,
|
|
nargs="?",
|
|
const=True,
|
|
default="./output",
|
|
help="directory to download files to, defaults is ./output",
|
|
),
|
|
parser.add_argument(
|
|
"--log_dir",
|
|
type=str,
|
|
nargs="?",
|
|
const=True,
|
|
default=None,
|
|
help="directory for logs, if ommitted will not log, logs may be large!",
|
|
),
|
|
parser.add_argument(
|
|
"--column",
|
|
type=str,
|
|
nargs="?",
|
|
const=True,
|
|
default="TEXT",
|
|
help="column to search for matches, defaults is 'TEXT', but you could use 'URL' if you wanted",
|
|
),
|
|
parser.add_argument(
|
|
"--limit",
|
|
type=int,
|
|
nargs="?",
|
|
const=True,
|
|
default=100,
|
|
help="max number of matching images to download, warning: may be slightly imprecise due to concurrency and http errors, defaults is 100",
|
|
),
|
|
parser.add_argument(
|
|
"--min_hw",
|
|
type=int,
|
|
nargs="?",
|
|
const=True,
|
|
default=512,
|
|
help="min height AND width of image to download, default is 512",
|
|
),
|
|
parser.add_argument(
|
|
"--force",
|
|
type=bool,
|
|
nargs="?",
|
|
const=True,
|
|
default=False,
|
|
help="forces a full download of all images, even if no search is provided, USE CAUTION!",
|
|
),
|
|
parser.add_argument(
|
|
"--parquet_skip",
|
|
type=int,
|
|
nargs="?",
|
|
const=True,
|
|
default=0,
|
|
help="skips the first n parquet files on disk, useful to resume",
|
|
),
|
|
parser.add_argument(
|
|
"--verbose",
|
|
type=bool,
|
|
nargs="?",
|
|
const=True,
|
|
default=False,
|
|
help="additional logging of URL and TEXT prefiltering",
|
|
),
|
|
parser.add_argument(
|
|
"--test",
|
|
action='store_const',
|
|
const=do_not_download,
|
|
default=not(do_not_download),
|
|
help="skips downloading, for checking filters, use with --verbose",
|
|
)
|
|
|
|
return parser
|
|
|
|
def cleanup_text(file_name: str):
|
|
# TODO: can be improved
|
|
|
|
file_name = re.sub("<div.*<\/div>", "", file_name)
|
|
file_name = re.sub("<span.*<\/span>", "", file_name)
|
|
file_name = re.sub("<a.*<\/a>", "", file_name)
|
|
file_name = file_name.replace('<p>', '').replace("</p>", "")
|
|
file_name = file_name.replace('<strong>', '').replace("</strong>", "")
|
|
file_name = file_name.replace('<em>', '').replace("</em>", "")
|
|
|
|
file_name = re.sub(r'[^\x00-\x7F]+', '', file_name) # remove non-ascii
|
|
|
|
file_name = file_name.replace(' & ', ' and ').replace(' &', ' and').replace('& ', 'and ') \
|
|
.replace(" + ", " and ").replace(" +", " and").replace("+ ", "and ")
|
|
|
|
file_name = file_name.replace('\t', ' ').replace('\n', ' ').replace('\r', ' ')
|
|
|
|
file_name = file_name.replace('\"t"', ' ')
|
|
|
|
file_name = file_name.replace(" ♥ ","love").replace("♥ ","love ").replace(" ♥"," love") \
|
|
.replace("♥"," love ")
|
|
|
|
# remove bad chars
|
|
file_name = file_name.replace('\"', '').replace('?', '') \
|
|
.replace('<', '').replace('>', '').replace('/', '').replace('*', '') \
|
|
.replace('!', '').replace('#', '').replace('$', '').replace('%', '') \
|
|
.replace('^', '').replace('(', '').replace(')', '')
|
|
|
|
# replace with space
|
|
file_name = file_name.replace(':',' ').replace('|',' ').replace('@', '') \
|
|
.replace("/", " ").replace("\\'", "\'").replace("\\", " ").replace('\\', ' ') \
|
|
.replace('_', ' ').replace("=", " ")
|
|
|
|
# replace foreign chars
|
|
file_name = file_name.replace('é', 'e').replace('è', 'e').replace('ê', 'e') \
|
|
.replace('ë', 'e').replace('à', 'a').replace('â', 'a').replace('ä', 'a') \
|
|
.replace('ç', 'c').replace('ù', 'u').replace('û', 'u').replace('ü', 'u') \
|
|
.replace('ô', 'o').replace('ö', 'o').replace('ï', 'i').replace('î', 'i') \
|
|
.replace('í', 'i').replace('ì', 'i').replace('ñ', 'n').replace('ß', 'ss') \
|
|
.replace('á', 'a').replace('ã', 'a').replace('å', 'a').replace('æ', 'ae') \
|
|
.replace('œ', 'oe').replace('ø', 'o').replace('ð', 'd').replace('þ', 'th') \
|
|
.replace('ý', 'y').replace('ÿ', 'y').replace('ž', 'z').replace('ž', 'z') \
|
|
.replace('š', 's').replace('đ', 'd').replace('ď', 'd').replace('č', 'c') \
|
|
.replace('ć', 'c').replace('ř', 'r').replace('ŕ', 'r').replace('ľ', 'l') \
|
|
.replace('ĺ', 'l').replace('ť', 't').replace('ň', 'n').replace('ņ', 'n') \
|
|
.replace('ď', 'd').replace('Ď', 'D').replace('Ť', 'T').replace('Ň', 'N')
|
|
|
|
_MAX_LENGTH = 240
|
|
if (len(file_name) > _MAX_LENGTH):
|
|
file_name = file_name[:_MAX_LENGTH]
|
|
|
|
return file_name
|
|
|
|
async def call_http(image_url: str, session: aiohttp.ClientSession):
|
|
#print(f"calling http and save to: {out_file_name}")
|
|
global downloaded_count
|
|
global http_timeout
|
|
global current_parquet_file_downloaded_count
|
|
try:
|
|
res = await session.request(method="GET", url=image_url, timeout=http_timeout)
|
|
|
|
if (res.status == 200):
|
|
return await res.content.read()
|
|
else:
|
|
print(f"{Fore.YELLOW}Failed to download image, HTTP response code: {res.status} for {Fore.LIGHTWHITE_EX}{image_url}{Style.RESET_ALL}")
|
|
downloaded_count -= 1
|
|
except Exception as e:
|
|
print(f"{Fore.YELLOW} *** Error downloading image: {Fore.LIGHTWHITE_EX}{image_url}{Fore.YELLOW}, ex: {str(e)}{Style.RESET_ALL}")
|
|
downloaded_count -= 1
|
|
pass
|
|
return None
|
|
|
|
async def save_img(buffer: io.BytesIO, full_outpath: str):
|
|
try:
|
|
async with aiofiles.open(full_outpath, "wb") as f:
|
|
await f.write(buffer.getbuffer())
|
|
except Exception as e:
|
|
print(f"{Fore.RED} *** Unable to write to disk: {Fore.LIGHTWHITE_EX}{full_outpath}{Style.RESET_ALL}")
|
|
print(f"{Fore.RED} *** ex: {Fore.LIGHTWHITE_EX}{str(e)}{Style.RESET_ALL}")
|
|
pass
|
|
|
|
def get_outpath_filename(data: any, full_outpath_noext: str, clean_text: str):
|
|
ext = "jpg"
|
|
full_outpath = None
|
|
buffer = None
|
|
try:
|
|
buffer = io.BytesIO(data)
|
|
image = Image.open(buffer)
|
|
ext = image.format.lower()
|
|
|
|
if (ext == "jpeg"):
|
|
ext = "jpg"
|
|
|
|
full_outpath = f"{full_outpath_noext}.{ext}"
|
|
except Exception as e:
|
|
print(f"{Fore.YELLOW} *** Possible corrupt image for text: {Fore.LIGHTWHITE_EX}{clean_text}{Style.RESET_ALL}")
|
|
print(f"{Fore.YELLOW} *** ex: {Fore.LIGHTWHITE_EX}{str(e)}{Style.RESET_ALL}")
|
|
pass
|
|
return full_outpath, buffer
|
|
|
|
async def download_image(image_url: str, clean_text: str, full_outpath_noext: IO, session: aiohttp.ClientSession):
|
|
http_content = await call_http(image_url=image_url, session=session)
|
|
|
|
buffer = None
|
|
|
|
if (http_content is not None):
|
|
full_outpath, buffer = get_outpath_filename(data=http_content, full_outpath_noext=full_outpath_noext, clean_text=clean_text)
|
|
|
|
if buffer is not None:
|
|
global downloaded_count
|
|
downloaded_count += 1
|
|
await save_img(buffer, full_outpath)
|
|
|
|
async def download_set_dict(opt, matches_dict: dict):
|
|
async with aiohttp.ClientSession() as session:
|
|
global downloaded_count
|
|
current_parquet_file_downloaded_count = 0
|
|
tasks = []
|
|
for row in matches_dict:
|
|
if downloaded_count < opt.limit:
|
|
current_parquet_file_downloaded_count += 1
|
|
pre_text=row["TEXT"]
|
|
image_url=row["URL"]
|
|
|
|
clean_text = cleanup_text(pre_text)
|
|
|
|
full_outpath_noext = os.path.join(opt.out_dir, clean_text)
|
|
|
|
if (opt.verbose):
|
|
print(f"{Fore.LIGHTGREEN_EX}***** Verbose log: ***** {Style.RESET_ALL}")
|
|
print(f"{Fore.LIGHTGREEN_EX} url: {image_url}{Style.RESET_ALL}")
|
|
print(f"{Fore.LIGHTGREEN_EX} text: {pre_text}{Style.RESET_ALL}")
|
|
print(f"{Fore.LIGHTGREEN_EX} captn: {clean_text}{Style.RESET_ALL}")
|
|
|
|
if any(glob.glob(full_outpath_noext + ".*")):
|
|
print(f"{Fore.YELLOW} already exists: {Fore.LIGHTWHITE_EX}{full_outpath_noext}{Fore.YELLOW}, skipping{Style.RESET_ALL}")
|
|
return
|
|
|
|
if not opt.test:
|
|
tasks.append(
|
|
download_image(image_url=image_url, clean_text=clean_text, full_outpath_noext=full_outpath_noext, session=session)
|
|
)
|
|
else:
|
|
current_parquet_file_downloaded_count += 1
|
|
downloaded_count += 1
|
|
if len(tasks) > 63:
|
|
await asyncio.gather(*tasks)
|
|
tasks = []
|
|
else:
|
|
print(f"{Fore.YELLOW} Limit reached: {opt.limit}, exiting...{Style.RESET_ALL}")
|
|
break
|
|
if not opt.test & len(tasks) > 0:
|
|
await asyncio.gather(*tasks)
|
|
print(f"{Fore.LIGHTBLUE_EX} Downloaded chunk of {current_parquet_file_downloaded_count} images{Style.RESET_ALL}")
|
|
|
|
def query_parquet(df: pd.DataFrame, opt):
|
|
# TODO: efficiency, expression tree?
|
|
matches = df
|
|
|
|
matches = matches[(matches.HEIGHT > opt.min_hw) & (matches.WIDTH > opt.min_hw)]
|
|
|
|
if 'punsafe' in matches.columns:
|
|
matches = matches[(matches.punsafe > unsafe_threshhold)]
|
|
|
|
if ('aesthetic' in matches):
|
|
matches = matches[(matches.aesthetic > aesthetic_threshhold)]
|
|
|
|
if opt.search_text:
|
|
for word in opt.search_text.split(","):
|
|
matches = matches[matches[opt.column].str.contains(word, case=False)]
|
|
|
|
matches = matches[~matches["URL"].str.contains("dreamstime.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("alamy.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("123rf.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("colourbox.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("envato.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("stockfresh.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("depositphotos.com", case=False)] # watermarks
|
|
matches = matches[~matches["URL"].str.contains("istockphoto.com", case=False)] # watermarks
|
|
|
|
return matches
|
|
|
|
async def download_laion_matches(opt):
|
|
print(f"{Fore.LIGHTBLUE_EX} Searching for {opt.search_text} in column: {opt.column} in {opt.laion_dir}/*.parquet{Style.RESET_ALL}")
|
|
|
|
for idx, file in enumerate(glob.iglob(os.path.join(opt.laion_dir, "*.parquet"))):
|
|
if idx < opt.parquet_skip:
|
|
print(f"{Fore.YELLOW} Skipping file {idx+1}/{opt.parquet_skip}: {file}{Style.RESET_ALL}")
|
|
continue
|
|
|
|
global downloaded_count
|
|
if downloaded_count < opt.limit:
|
|
print(f"{Fore.CYAN} reading file: {file}{Style.RESET_ALL}")
|
|
|
|
df = pd.read_parquet(file, engine="auto")
|
|
matches = query_parquet(df, opt)
|
|
# print(f"{Fore.CYAN} matches in current parquet file:{ Style.RESET_ALL}")
|
|
# print(matches)
|
|
|
|
match_dict = matches.to_dict('records') # TODO: pandas problems later in script... needs revisiting
|
|
|
|
await download_set_dict(opt, match_dict)
|
|
else:
|
|
print(f"{Fore.YELLOW}limit reached before reading next parquet file. idx: {idx}, filename: {file}{Style.RESET_ALL}")
|
|
break
|
|
|
|
def isWindows():
|
|
return sys.platform.startswith('win')
|
|
|
|
def ensure_path_exists(path: str):
|
|
if not os.path.exists(path):
|
|
print(f"{Fore.LIGHTBLUE_EX}creating path: {path}{Style.RESET_ALL}")
|
|
os.makedirs(path)
|
|
|
|
if __name__ == '__main__':
|
|
print(f"{Fore.CYAN}Launching...{Style.RESET_ALL}")
|
|
inVenv = in_virtualenv()
|
|
print(f"is running in venv: {inVenv}")
|
|
#assert inVenv, "Error loading venv. Please run 'source everydream-venv/bin/activate', or in windows 'everydream-venv/bin/activate.bat'"
|
|
|
|
parser = get_parser()
|
|
opt = parser.parse_args()
|
|
|
|
print(f"Test only mode: {opt.test}")
|
|
|
|
if(opt.search_text is None and opt.force is False):
|
|
print(f"{Fore.YELLOW}** No search terms provided, exiting...")
|
|
print(f"** Use --force to bypass safety to dump entire DB{Style.RESET_ALL}")
|
|
sys.exit(2)
|
|
|
|
ensure_path_exists(opt.out_dir)
|
|
|
|
if (opt.laion_dir[-1] != "/" or opt.laion_dir[-1] != "\\"):
|
|
opt.laion_dir += "/"
|
|
|
|
if (isWindows()):
|
|
print("{Fore.CYAN}Windows detected, using asyncio.WindowsSelectorEventLoopPolicy{Style.RESET_ALL}")
|
|
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
|
else:
|
|
print("{Fore.CYAN}Unix detected, using default asyncio event loop policy{Style.RESET_ALL}")
|
|
|
|
import time
|
|
s = time.perf_counter()
|
|
|
|
result = asyncio.run(download_laion_matches(opt))
|
|
|
|
elapsed = time.perf_counter() - s
|
|
print(f"{Fore.CYAN} **** Job Complete ****")
|
|
print(f" search_text: \"{opt.search_text}\", force: {opt.force}")
|
|
print(f" Total downloaded {downloaded_count} images")
|
|
print(f"{__file__} executed in {elapsed:0.2f} seconds.{Style.RESET_ALL}{Style.RESET_ALL}") |