image mosaic script v3

# This script takes a directory as its only argument and creates a mosaic image out of all the images in that folder, subfolders _not_ included.
# PIL (Python Image Library) is required to produce an output image.
# 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.
# Either give the folder path as command line argument or just change it below after the imports.

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 and run without arguments
if (len(sys.argv) > 1) and os.path.isdir(sys.argv[1]):
    path = sys.argv[1]
else:
    path = "E:/Pictures/anime/Sansha Sanyou/endcards"

#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 "mosaic"

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
padding_top = 2      #top, bottom and between rows

row_height = 540 - padding_top
row_width = 2*1920 - padding_left - padding_right

weight_accuracy = 1000  #determines with how much precision the aspect ratios of images will be handled
output_quality = 100    #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 os.path.splitext(f)[1].lower() in [".jpeg", ".jpg", ".png", ".gif", ".tiff", ".webp"]:
        if f.lower() == "mosaic.jpg": continue
        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

    # print(f, str(width), str(height))
    asp = round(weight_accuracy * float(width + padding_left) / 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)

# print(scaled_sum)
# print(k)
# print(Sol)

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_center
        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.BICUBIC)

            #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)

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