image mosaic script v6

## This script requires Python 3.
## PIL (Python Imaging Library) is required to produce an output image. I use Pillow 3.4.0. Get it here: https://pypi.python.org/pypi/Pillow
## It takes a directory as its only argument and creates a mosaic image out of all the images in that folder, subfolders NOT included.
## Its output is a 'mosaic.jpg' in that same folder.
## If there already is a 'mosaic.jpg' in that folder, it will not be included in the mosaic, but the file will be overwritten with the new one without a confirmation prompt.
## Read the description of variables below.
## Usually only the adjustment of row_height and row_width is required.
## output_quality can be adjusted down from 100 to lower file sizes.

# pastebin link: http://pastebin.com/gP3UgLBC

import os
import sys
import subprocess as sub
import random
from PIL import Image
from math import floor
from copy import deepcopy

#give path as command line argument or just change it below
if (len(sys.argv) > 1) and os.path.isdir(sys.argv[1]):
    path = sys.argv[1]
else:
    path = "E:/Pictures/anime/Flip Flappers/Papika/n"

if not os.path.isdir(path):
    print('Path "'+path+'" not found.')
    sys.exit(0)

#force_aspect_ratio = 0: allow row heights to exceed row_height
#force_aspect_ratio > 0: add padding between images in a row or between rows until desired aspect ratio (e.g. force_aspect_ratio = 16/9) is reached
force_aspect_ratio = 0

random_pick = True      #the images in each row will likely be the same, but arranged in a different order if you run the script twice
outputs = 1             #if >1 and random_pick=True, then multiple different mosaics will be created in a subfolder "mosaics"

padding_left = 2        #only on the left end of each row
padding_right = 2       #only on the right end of each row
padding_center = 3      #between images in a row; should not be too different from padding_left
padding_top = 2         #top, bottom and between rows

row_height = 540 - padding_top
# row_height = 1080 - padding_top
row_width = 2*1920 + 960 - padding_left - padding_right   #final image width without padding

weight_accuracy = 1000  #determines with how much precision the aspect ratios of images will be handled
output_quality = 95     #jpg quality of the output image

############
#http://stackoverflow.com/questions/7938809/dynamic-programming-linear-partitioning-please-help-grok/7942946#7942946
############
from operator import itemgetter

def linear_partition(seq, k):
    if k <= 0:
        return []
    n = len(seq) - 1
    if k > n:
        return map(lambda x: [x], seq)
    table, solution = linear_partition_table(seq, k)
    k, ans = k-2, []
    while k >= 0:
        ans = [[seq[i] for i in range(solution[n-1][k]+1, n+1)]] + ans
        n, k = solution[n-1][k], k-1
    return [[seq[i] for i in range(0, n+1)]] + ans

def linear_partition_table(seq, k):
    n = len(seq)
    table = [[0] * k for x in range(n)]
    solution = [[0] * (k-1) for x in range(n-1)]
    for i in range(n):
        table[i][0] = seq[i] + (table[i-1][0] if i else 0)
    for j in range(k):
        table[0][j] = seq[0]
    for i in range(1, n):
        for j in range(1, k):
            table[i][j], solution[i-1][j-1] = min(((max(table[x][j-1], table[i][0]-table[x][0]), x) for x in range(i)), key=itemgetter(0))
    return (table, solution)
############


Sol = {}
scaled_sum = 0
random.seed()

#info[f] = [width, height, scale]
info = {}

#calculate image Sol
print('calculating image weights')
root, dirs, files = next(os.walk(path))
for f in files:
    if f.lower() == "mosaic.jpg": continue
    if os.path.splitext(f)[1].lower() in [".jpeg", ".jpg", ".png", ".gif", ".tiff", ".webp"]:
        try:
            image = Image.open(os.path.join(root, f))
            width, height = image.size
            info[f] = [width, height, 1]
        except:
            print('Error while opening "'+f+'"')
    else:
        # print("unknown: ", f)
        continue

    #padding_center is included in the aspect ratios here, because the linear partition solution has to consider the sum of all pixels in a row.
    #Otherwise the rows with more images would receive more padding, resulting in wider rows.
    #Since it's added for every image, it's added once too much for every row. In order to do it right, the final placement would have to be known beforehand.
    #I'm assuming here that padding_left and padding_right aren't too different from padding_center, so the error compared to the entire row_width should be hardly noticeable.
    asp = round(weight_accuracy * float(width + padding_center) / float(height))
    info[f][2] = row_height / (height + padding_top)
    scaled_sum += info[f][2] * width

    if asp in Sol:
        Sol[asp].append(f)
    else:
        Sol[asp] = [f]


#calculate number of rows
#images scaled to row height -> sum of scaled widths / row_height
k = round(scaled_sum / row_width)

seq = []
for w in Sol:
    for i in range(len(Sol[w])):
        seq.append(w)
print('fitting '+str(len(seq))+' images into '+str(k)+' rows')
base_solution = linear_partition(seq, k)


if not random_pick: outputs = 1
mosaic_solutions = []

while outputs > 0:
    outputs -= 1

    print('calculating final image size')
    # {row: [(file1, sizex, sizey, posx, posy), (file2, sizex, sizey, posx, posy), ...]}
    mosaic_solution = {}
    mosaic_width, mosaic_height = (row_width + padding_left + padding_right, row_height + padding_top)
    random.seed()
    S = deepcopy(Sol)
    solution = deepcopy(base_solution)

    #current "cursor" position
    pos = [padding_left, padding_top]
    for row in range(len(solution)):
        random.seed()
        random.shuffle(solution[row])

        mosaic_solution[row] = []

        #rescale the row to fit target row_width
        row_scale = 1
        solution_width = padding_left
        temp_S = {}
        for col in range(len(solution[row])):
            f = S[solution[row][col]].pop(0)
            if solution[row][col] in temp_S:
                temp_S[solution[row][col]].append(f)
            else:
                temp_S[solution[row][col]] = [f]

            width, height, scale = info[f][0:3]
            solution_width += floor(scale*width) + padding_center

        row_scale = row_width / solution_width

        for w in temp_S:
            for f in temp_S[w]:
                S[w].insert(0, f)

        for col in range(len(solution[row])):
            #desired weight = solution[row][col]
            if random_pick:
                random.seed()
                random.shuffle(S[solution[row][col]])
                f = S[solution[row][col]].pop(0)
            else:
                f = S[solution[row][col]].pop(0)
                if len(S[solution[row][col]]) == 0: S.pop(solution[row][col])

            width, height, scale = info[f][0:3]

            mosaic_solution[row].append([f, floor(row_scale*scale*width), floor(row_scale*scale*height), pos[0], pos[1]])
            pos[0] += floor(row_scale*scale*width) + padding_center
            if pos[0] > mosaic_width: mosaic_width = pos[0]

        pos[0] = padding_left
        pos[1] += floor(row_scale*row_height) + padding_top
        if pos[1] > mosaic_height: mosaic_height = pos[1]

    if force_aspect_ratio > 0:
        #mosaic height is not guaranteed to be the desired height
        #calculate new image width from desired aspect ratio
        print('refitting into aspect ratio: '+str(force_aspect_ratio))

        #if mosaic is too high -> more padding between images
        if force_aspect_ratio > float(mosaic_width / mosaic_height):
            new_mosaic_width = round(mosaic_height * force_aspect_ratio)
            new_padding = new_mosaic_width - mosaic_width - padding_left - padding_right
            mosaic_width = new_mosaic_width

            #reposition images, distribute padding equally *between* the pictures, leave outer padding
            print("new padding: "+str(new_padding))
            for row in range(len(mosaic_solution)):
                if len(mosaic_solution[row]) == 0: next
                new_row_padding = round(new_padding / (len(mosaic_solution[row]) - 1))
                print("new row padding: "+str(new_row_padding))
                for col in range(len(mosaic_solution[row])):
                    mosaic_solution[row][col][3] += new_row_padding * col

        #if mosaic is too wide -> more padding between rows
        else:
            new_mosaic_height = round(mosaic_width / force_aspect_ratio)
            new_padding = new_mosaic_height - mosaic_height - padding_top - padding_top
            mosaic_height = new_mosaic_height

            #reposition images, distribute padding equally between the rows, leave top and bottom padding
            print("new padding: "+str(new_padding))
            for row in range(len(mosaic_solution)):
                if len(mosaic_solution[row]) == 0: next
                new_row_padding = round(new_padding / (len(mosaic_solution) - 1))
                print("new row padding: "+str(new_row_padding))
                for col in range(len(mosaic_solution[row])):
                    mosaic_solution[row][col][4] += new_row_padding * row

    mosaic_solutions.append(mosaic_solution)

if len(mosaic_solutions) > 1 and not os.path.isdir(os.path.join(path, 'mosaics')):
    os.mkdir(os.path.join(path, 'mosaics'))

for s in range(len(mosaic_solutions)):
    if len(mosaic_solutions) > 1:
        print('\ncreating mosaics/mosaic'+str(s+1)+'.jpg')
        mosaic_fname = os.path.join(path, 'mosaics', 'mosaic'+str(s+1)+'.jpg')
    else:
        print('creating mosaic.jpg\n')
        mosaic_fname = os.path.join(path, 'mosaic.jpg')

    mosaic = Image.new("RGB", (mosaic_width, mosaic_height), color=0)
    mosaic_solution = mosaic_solutions[s]
    for row in range(len(mosaic_solution)):
        for col in range(len(mosaic_solution[row])):
            f, sizex, sizey, posx, posy = mosaic_solution[row][col]

            temp_image = Image.open(os.path.join(root, f))
            #rescale if necessary
            if sizex != temp_image.size[0]:
                temp_image = temp_image.resize((sizex, sizey), resample = Image.ANTIALIAS)

            #place the image into the mosaic
            mosaic.paste(temp_image, (posx, posy))

    print('mosaic dimensions:', str(mosaic.size))
    mosaic.save(mosaic_fname, quality=output_quality)

    print('mosaic file size:', "{:.2f}".format(os.path.getsize(mosaic_fname)/1024/1024) + "MB")

    if len(mosaic_solutions) == 1:
        # mosaic.show()
        sub.Popen('"'+mosaic_fname+'"', shell=True)