diff options
Diffstat (limited to 'get_images.py')
| -rw-r--r-- | get_images.py | 73 |
1 files changed, 64 insertions, 9 deletions
diff --git a/get_images.py b/get_images.py index b4fef24..37becdd 100644 --- a/get_images.py +++ b/get_images.py @@ -1,17 +1,32 @@ from dataclasses import dataclass +from PIL import Image, ImageDraw from io import StringIO from lxml import etree import requests +import logging import urllib import random +import utils import time +import json import cv2 import os +with open("config.json", "r") as f: + CONFIG = json.load(f) + +logging.basicConfig( + format = "%(levelname)s\t[%(asctime)s]\t%(message)s", + level = logging.INFO, + handlers=[ + logging.FileHandler(CONFIG["logpath"]), + logging.StreamHandler() + ]) + # all of these tags are added to all queries. Preceded with '-' to blacklist -base_tags = ["yaoi", "-muscle", "-comic"] +base_tags = CONFIG["base_tags"] # one of these will be added -search_tags = ["looking_at_another", "kiss", "trap", "2boys", "promare"] +search_tags = CONFIG["search_tags"] def get_random_searchtag(): return [random.choice(search_tags)] @@ -30,7 +45,10 @@ def get_image(tags): search_url = "https://safebooru.org/index.php?page=post&s=list&tags=%s&pid=%i" % ("+".join(base_tags+tags), (random.randint(1, get_num_pages(tags))-1)*5*8) tree = etree.parse(StringIO(requests.get(search_url).text), etree.HTMLParser()) elements = [e for e in tree.xpath("/html/body/div[6]/div/div[2]/div[1]")[0].iter(tag = "a")] - element = random.choice(elements) + try: + element = random.choice(elements) + except IndexError: + raise ConnectionError("Couldn't find any images") simg = SafebooruImage( id = get_id_from_url(element.get("href")), tags = element.find("img").get("alt").split(), @@ -67,6 +85,16 @@ def get_num_pages(tags): else: return int(int(urllib.parse.parse_qs(page_element.get("href"))["pid"][0]) / (5*8)) +def append_blacklisted(id_): + with open(CONFIG["blacklist"], "a") as f: + f.write(str(id_) + "\n") + +def id_is_blacklisted(id_): + if not os.path.exists(CONFIG["blacklist"]): + return False + with open(CONFIG["blacklist"], "r") as f: + return str(id_) in f.read().splitlines() + @dataclass class DownloadedImage: imurl: str @@ -83,11 +111,20 @@ class DownloadedImage: def __exit__(self, type, value, traceback): os.remove(self.filename) -def main(): +def main(draw_faces = False): try: simg = get_image(get_random_searchtag()) except ConnectionError: + logging.warning("Retried since couldn't get source...") main() + return + + if id_is_blacklisted(simg.id): + logging.info("Retried, already posted image...") + main() + return + + append_blacklisted(simg.id) with DownloadedImage(simg.imurl) as impath: img = cv2.imread(impath) @@ -101,14 +138,32 @@ def main(): scaleFactor = 1.1, minNeighbors = 5, minSize = (24, 24)) - for (x, y, w, h) in faces: - cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2) + if draw_faces: + for (x, y, w, h) in faces: + cv2.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2) + + logging.info("Found image %i faces, id: %i" % (len(faces), simg.id)) + + pilimg = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) + text = utils.get_quote(CONFIG["texts"]) + logging.info(text) + font = utils.set_font(pilimg, text) + draw = ImageDraw.Draw(pilimg) + lines = utils.messages_multiline(text, font, pilimg) + colours = utils.get_colors(impath) + + (x, y, faces) = utils.randomize_location(pilimg, lines, font, faces) + for line in lines: + height = font.getsize(line[1])[1] + utils.draw_with_border(x, y, line, colours[0], colours[1], font, draw) + y = y + height + + pilimg.save("img.png") + return "img.png", simg.source, text - cv2.imshow("Press <q> to exit", img) - cv2.waitKey(0) if __name__ == "__main__": - main() + print(main()) |