gba_processing.py

from pygame.image import load as load_image, save as save_image
from pygame.surfarray import array3d, make_surface
from pygame import Surface
from colorsys import rgb_to_hls
from sys import argv
from math import sqrt
from numpy import array

formatted_colors = {}
def fmt_color(c):
    rgb = c if isinstance(c,tuple) else color_components(c)
    get_fmt = lambda c:"#{:02X}{:02X}{:02X}".format(*rgb)
    #get_fmt = lambda c:"rgb({},{},{})".format(*rgb)
    if c not in formatted_colors:
        formatted_colors[c] = get_fmt(c)
    return formatted_colors[c]

def gba_component(std_component):
    return int(round((((std_component / 255) * 31))) * 8)

def gba_filter(c):return tuple(map(gba_component,c))

def color_distance(a,b):
    hsl_a, hsl_b = rgb_to_hls(*map(int,a)),rgb_to_hls(*map(int,b))
    return sqrt(sum(map(lambda a,b:((a-b)**2), hsl_a,hsl_b)))

def n_to_t(n):return (n[0],n[1],n[2])
def t_to_n(t):return array(t)

def process_surface(surf):
    rows = array3d(surf)
    to_swap,color_heat = {},{}
    for row in rows:
        for pixel in row:
            pixel = n_to_t(pixel)
            if pixel in to_swap:
                color_heat[to_swap[pixel]]-=1
                continue
            nc = gba_filter(pixel)
            to_swap[pixel] = nc
            color_heat[nc] = -1
    colors = sorted(map(to_swap.get,to_swap.keys()),key=color_heat.get)
    used_colors = colors[:16]
    r_cs = [c for c in to_swap.keys() if to_swap[c] not in used_colors]
    for color in r_cs:
        dist_to = lambda cc:(cc,-1 * color_distance(color,cc))
        color_distances = dict(map(dist_to,used_colors))
        s = sorted(color_distances.keys(),key=color_distances.get)
        to_swap[color] = s[0]
    u_colors = []
    def color(c):
        r=to_swap.get(c,c)
        u_colors.append(r)
        return (r[1],r[0],r[2])
    def column(row):return array([color(n_to_t(c)) for c in row])
    new_things = array([column(row) for row in rows])
    u_colors = set(u_colors)
    print(*map(fmt_color,u_colors),sep='\n')
    return make_surface(new_things)

def append_fn(fn,ap):return fn[:fn.find(".")] + ap + fn[fn.find("."):]
def base_fn(fn):return fn[fn.rfind("\\")+1:]

for fn in argv[1:]:
    nfn = append_fn(fn,"-processed")
    print(base_fn(fn),":",base_fn(nfn))
    img = load_image(fn)
    new_img,dimensions = process_surface(img), img.get_size()
    new_image = Surface(tuple(map(lambda a,b:a*b,dimensions,[2,1])))
    new_image.blit(img,(0,0))
    new_image.blit(new_img,(img.get_width(),0))
    save_image(new_img, nfn)
    save_image(new_image, append_fn(fn,"-combined"))
    print("="*80)