wallpaper-finder.py

import praw
import re
import requests
import random
from io import BytesIO
from functools import cmp_to_key
from collections import namedtuple
from math import sqrt, ceil
from datetime import datetime
import PIL

try:
    import Image
except ImportError:
    from PIL import Image

#prints colors only on terminals that support True Color

CLIENT_ID = ""
CLIENT_SECRET = ""
USER_AGENT = "MyComputer:wallpaper-finder " + CLIENT_ID + ":v0.0.1"
SEARCH_TERMS = {"cityporn" : [["night", [["#1f3b68", 0.2], ["#487b91", 0.5], ["#b7682c", 0.5]]]]}
#colors must be in #rrggbb hex format
#[color, max_color_difference (max=2)]
RATIO = 1920/1080
MAX_RATIO_DIFFERENCE = 0.2
LIMIT = 100

Result = namedtuple('Result', ('url', 'colors'), verbose=False)
ColorInfo = namedtuple('ColorInfo', ('color', 'req_color', 'dist', 'density'), verbose=False)

def main():

    start_time = datetime.now()
    post_counter = 0
    results = []
    for sub_name in SEARCH_TERMS:
        for query in SEARCH_TERMS[sub_name]:
            for req_color in query[1]:
                req_color[0] = req_color[0][1:]
                req_color[0] = Point([int(req_color[0][:2], 16), int(req_color[0][2:4], 16), int(req_color[0][4:], 16)], 3, 1)

    reddit = praw.Reddit(client_id=CLIENT_ID,client_secret=CLIENT_SECRET, user_agent=USER_AGENT)

    for sub_name in SEARCH_TERMS.keys():
        sub = reddit.subreddit(sub_name)
        print("Searching in {}...".format(sub_name))
        for query in SEARCH_TERMS[sub_name]:
            print("Searching for {}...".format(query[0]))
            for post in sub.search(query[0], limit=ceil(LIMIT/len(SEARCH_TERMS.keys()))):
                post_counter += 1
                url = post.url
                title = post.title.replace(" ","")
                start = re.search(r"(?=\[\d+\s?(×|X|x)\s?\d+\])", title)
                end = re.search(r"(?=\]\d+\s?(×|X|x)\s?\d+\[)", title[::-1]) #Weak re engine implementantion :/

                if (start is not None and end is not None):
                    start = start.start()
                    end = len(title)-end.start()
                    dimensions = re.split(r"×|X|x", title[start:end][1:-1])
                    ratio = int(dimensions[0])/int(dimensions[1])

                    if (ratio - ratio * MAX_RATIO_DIFFERENCE < RATIO) and (RATIO < ratio + ratio * MAX_RATIO_DIFFERENCE):
                        if (re.match(r"^.*\.(jpg|jpeg|png)$", url)):
                            response = requests.get(url)
                            img = Image.open(BytesIO(response.content))
                            colors = colorz(img, n=len(query[1])) # [(center, num_points), ...]
                            size = img.size[0] * img.size[1]
                            counter = 0
                            for i in range(0, len(colors)):
                                color = colors[i][0]
                                percent = colors[i][1]
                                req_color = query[1][i][0]

                                if euclidean(color, req_color) <= (255/2) * percent:
                                    counter += 1
                                else:
                                    break

                            if counter == len(colors):
                                get_ci = lambda c, rc : ColorInfo(c[0], rc[0], euclidean(c[0], rc[0]), c[1]/size)
                                results.append(Result(url, [get_ci(c, rc) for c, rc in zip(colors, query[1])]))

    s = lambda r : sum([ci.density * ci.dist for ci in r.colors])
    #r = (url, [colorinfo1, ...])
    def cmp(r1, r2):
        x = s(r2)
        y = s(r1)

        if x < y:
            return 1
        if x > y:
            return -1
        return 0

    results.sort(key=cmp_to_key(cmp))

    p = lambda color: print("\x1b[38;2;{};{};{}m████\x1b[0m ".format(round(color.coords[0]), round(color.coords[1]), round(color.coords[2])), end="")

    for result in results:
        print("Requested : ", end="")
        for color in result.colors:
            color = color[1]#1f3b68
            p(color)
        print(result.url)

        print("Found     : ", end="")
        for color in result.colors:
            color = color[0]
            p(color)

        print("{}\n".format(round(s(result), 2)))

    end_time = datetime.now()
    print("Found {} results in {} .".format(len(results), end_time-start_time))

#n = dimension, ct = counter
Point = namedtuple('Point', ('coords', 'n', 'ct'))
#point = list_of_points, n = dimension, plen = number of total pixels that belong to this cluster
Cluster = namedtuple('Cluster', ('points', 'center', 'n', 'plen'))

def get_points(img):
    points = []
    w, h = img.size
    for count, color in img.getcolors(w * h):
        points.append(Point(color, 3, count))
    return points

rtoh = lambda rgb: '%s' % ''.join(('%02x' % p for p in rgb))

def colorz(img, n=3):
    img.thumbnail((200, 200), resample=PIL.Image.LANCZOS)
    w, h = img.size

    points = get_points(img)
    clusters = kmeans(points, n, 1)

    def cmp(c1, c2):
        if c1.plen < c2.plen:
            return 1
        elif c1.plen > c2.plen:
            return -1
        return 0

    clusters.sort(key=cmp_to_key(cmp))

    return [(c.center, c.plen) for c in clusters]

def euclidean(p1, p2):
    return sqrt(sum([(p1.coords[i] - p2.coords[i]) ** 2 for i in range(p1.n)]))

def calculate_center(points, n):
    vals = [0.0 for i in range(n)]
    plen = 0
    for p in points:
        plen += p.ct
        for i in range(n):
            vals[i] += (p.coords[i] * p.ct)
    return (Point([(v / plen) for v in vals], n, 1), plen)

def kmeans(points, k, min_diff):
    clusters = [Cluster([p], p, p.n, 0) for p in random.sample(points, k)]
    while 1:
        plists = [[] for i in range(k)]

        for p in points:
            smallest_distance = float('Inf')
            for i in range(k):
                distance = euclidean(p, clusters[i].center)
                if distance < smallest_distance:
                    smallest_distance = distance
                    idx = i
            plists[idx].append(p)

        diff = 0
        for i in range(k):
            old = clusters[i]
            center, plen = calculate_center(plists[i], old.n)
            new = Cluster(plists[i], center, old.n, plen)
            clusters[i] = new
            diff = max(diff, euclidean(old.center, new.center))

        if diff < min_diff:
            break

    return clusters

if __name__ == "__main__":
    main()