Creates an photo mosaic out of a base image and a search term. For fun only, the downloaded images may not be licensed for re-use

mosaic_builder.py

"""
Before running:
$ pip install pillow
$ pip install requests

Example usage:
$ python mosaic_builder.py puppies my_template_image.jpg

For optional arguments:
$ python mosaic_builder.py -h
"""

import math
import os
import random
import re
import shutil

from PIL import Image
import requests

AVAILABLE_COLORS = ['RED','YELLOW','GREEN','TEAL','BLUE','PURPLE','BLACK','WHITE','GRAY','ORANGE','PINK','BROWN', '']
MAX_TO_DOWNLOAD = 3
BING = False
SAVE_BLOCKED_IMAGE = False

def download_images(search_term, color, num_images, image_folder='images'):
    if color:
        color_dir = os.path.join(image_folder, color);
    else:
        color_dir = os.path.join(image_folder, 'uncolored')

    if os.path.exists(color_dir):
        num_already_downloaded = sum([1 for name in os.listdir(color_dir) if os.path.isfile(os.path.join(color_dir,name))])
    else:
        os.makedirs(color_dir)
        num_already_downloaded = 0

    num_to_download = num_images - num_already_downloaded

    if num_to_download > 0:

        user_agent_header = {'User-Agent' : 'Mozilla/5.0 (X11; Fedora; Linux x86_64; rv:42.0) Gecko/20100101 Firefox/42.0'}
        if color:
            search_term = color + '+' + search_term

        page_num = 1
        while(num_to_download > 0):
            print("DOWNLOADING: need {} of {} {}".format(num_to_download, color, search_term))
            
            if BING:
                request_url='https://www.bing.com/images/async?q=' + search_term + '&async=content&first=' + str(page_num) + '&adlt=off' + '&qft=+filterui:color2-FGcls_' + color
            else:
                request_url = 'https://www.google.com/search?tbm=isch&q=' + search_term +'&sout=1&start='+ str(page_num) #+ '&tbs=ic:specific,isc:' + color

            response = requests.get(request_url, headers=user_agent_header)
            if response.status_code != 200:
                print(response.text, response.status)
                return
            if BING:
                links = re.findall('imgurl:"(.*?)"',response.text)
            else:
                # links = re.findall('"ou":"(.*?)",', response.text)
                links = re.findall('src="(.*?gstatic.*?)"', response.text)

            # print(links)
            results_per_page = len(links)

            for link in links:
                if BING:
                    filename = link.split('/')[-1]
                else:
                    # This script assumes images are jpg, they will get deleted later if they are not
                    filename = link.split(':')[-1] + '.jpg'
                if os.path.exists(os.path.join(color_dir, filename)):
                    #we've already got this image, skip
                    continue
                
                try:
                    new_image = requests.get(link).content
                except:
                    print("Problem downloading {}".format(link))
                    continue
                with open(os.path.join(color_dir, filename),'wb') as imagefile:
                    imagefile.write(new_image)
                num_to_download -=1

            page_num += results_per_page

def get_region_color(region):
    histo = region.histogram()
    if len(histo) == 256:
        return None
    # split into red, green, blue
    r = histo[0:256]
    g = histo[256:256*2]
    b = histo[256*2: 256*3]

    r_avg = 0 if sum(r) == 0 else sum([count * i for i, count in enumerate(r)])/sum(r)
    g_avg = 0 if sum(g) == 0 else sum([count * i for i, count in enumerate(g)])/sum(g)
    b_avg = 0 if sum(b) == 0 else sum([count * i for i, count in enumerate(b)])/sum(b)
    avg_color = (int(r_avg), int(g_avg), int(b_avg)) 
    return avg_color

def nearest_color(original_rgb, available_colors):
    min_distance = float("inf")
    for rgb in available_colors:
        # euclidean distance between the available colors
        color_distance = math.sqrt((original_rgb[0] - rgb[0])**2 + (original_rgb[1] - rgb[1])**2 + (original_rgb[2] - rgb[2])**2)
        if color_distance < min_distance:
            color_result = rgb
            min_distance = color_distance

    return color_result

def get_color_order(image, region_size):
    grid_width = image.size[0]//region_size
    grid_height = image.size[1]//region_size

    color_order = []
    new_image = Image.new('RGB', image.size, (0, 255, 0))

    for ih in range(grid_height):
        for iw in range(grid_width):
            region = image.crop((iw*region_size, ih*region_size, iw*region_size+region_size, ih*region_size+region_size))
            color = get_region_color(region)
            patch = Image.new('RGB', (region_size, region_size), color)
            new_image.paste(patch, box=(iw*region_size, ih*region_size))
            color_order.append(color)

    if SAVE_BLOCKED_IMAGE:
        new_image.save('blocked_result.jpg')

    return color_order

def get_imagenames(folder='images'):
    color_dict = {}
    color_dirs = [name for name in os.listdir(folder) if os.path.isdir(os.path.join(folder, name))]
    for directory in color_dirs:
        names = [name for name in os.listdir(os.path.join(folder, directory)) if name.endswith('jpg')]
        for name in names:
            img = Image.open(os.path.join(folder, directory, name))
            calculated_color = get_region_color(img)
            if calculated_color:
                if calculated_color in color_dict:
                    color_dict[calculated_color].append(os.path.join(folder, directory, name))
                else:
                    color_dict[calculated_color] = [os.path.join(folder, directory, name)]

    return color_dict

def prune_images(source='images'):
    # clear out any images that can't be opened with PIL
    images = get_imagenames(source)
    for color, filenames in images.items():
        for filename in filenames:
            try:
                Image.open(filename)
            except OSError:
                print('Deleting: ', filename)
                os.remove(os.path.join(source, color, filename))

def assemble_image(color_order, image_size, region_size, image_dir, outfile):
    new_image = Image.new('RGB', image_size, (0, 255, 0))

    # create a pixelated image of the base image, for tuning regions size
    new_image_blocked = Image.new('RGB', image_size, (0, 255, 0))

    # Get a dict of image names
    image_files = get_imagenames(image_dir)
    available_colors = image_files.keys()

    for i, color in enumerate(color_order):
        nearest = nearest_color(color, available_colors)
        pupper_path = image_files[nearest][random.randint(0, len(image_files[nearest])-1)]
        patch = Image.open(pupper_path)
        patch = patch.resize((region_size, region_size))
        bbox = ((i*region_size)%image_size[0], (i // (image_size[0]//region_size))*region_size)
        new_image.paste(patch, box=bbox)
        patch.close()

        color_patch = Image.new('RGB', (region_size, region_size), nearest)
        new_image_blocked.paste(color_patch, box=bbox)
        

    new_image.save(outfile)

    if SAVE_BLOCKED_IMAGE:
        new_image_blocked.save('color_calc.jpg')

def main(search_term, block_size, base_image_filename, output_filename):

    img = Image.open(base_image_filename)

    # snap base image to multiple of region size
    grid_width, grid_height = img.size[0]//block_size, img.size[1]//block_size
    new_width = grid_width*block_size
    new_height = grid_height*block_size
    img = img.resize((new_width,new_height), Image.ANTIALIAS)

    color_order = get_color_order(img, block_size)
    # Download 100 of each color available on search engines, to try to get a wide enough variety of colors
    for color in AVAILABLE_COLORS:
        download_images(search_term, color, MAX_TO_DOWNLOAD, 'images')

    print('pruning bad images')
    prune_images('images')

    print("assembling mosaic")
    assemble_image(color_order, img.size, block_size, 'images', output_filename)
    print(img.size)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='Create a photo mosaic from images downloaded from Google')
    parser.add_argument(dest='search_term', help='Term to use for google image search')
    parser.add_argument(dest='base_image', help='Filename of image to use as a template to arrange mosaic on top of; Must be JPEG')
    parser.add_argument('-o', dest='destination', default='mosaic.jpg', help='Output filename; extension should be .jpg or .jpeg')
    parser.add_argument('-b', dest='block_size', default=64, type=int, help='Number of pixels per mosaic tile')

    args = parser.parse_args()
    main(args.search_term, args.block_size, args.base_image, args.destination)