evaluate-chroma-decoder

#!/usr/bin/python3
# Optimise filter configurations for ld-chroma-decoder using a genetic
# algorithm. Performance is evaluated by encoding test material with
# ld-chroma-encoder and measuring the similarity of the decoded result using
# SSIM.

# XXX Try generating an NTSC 3D FFT filter

import PIL.Image
import PIL.ImageDraw
import hashlib
import logging
import os
import random
import statistics
import sys
import time

from testvideos import *

logging.basicConfig(level=logging.INFO)
testcases = get_testcases()

"""
We represent threshold values 0.0-1.0 as integers from 0-100.

Grid sizes:
transform3d: 8 * 32 * ((16 / 8) + 1) = 768
transform2d: 16 * ((32 / 8) + 1) = 80
"""

# transform2d:
THRESHOLDS_Z = 1
THRESHOLDS_Y = 16
TILE_X = 32
# transform3d:
THRESHOLDS_Z = 8
THRESHOLDS_Y = 32
TILE_X = 16

THRESHOLDS_X = (TILE_X // 8) + 1

def cell(x, y, z):
    return (((z * THRESHOLDS_Y) + y) * THRESHOLDS_X) + x
def clamp_cell(v):
    return max(0, min(100, v))

THRESHOLDS_SIZE = THRESHOLDS_Z * THRESHOLDS_Y * THRESHOLDS_X
QUANTUM = 5
POPULATION_SIZE = 50
# Normally:
NUM_PREVIOUS_RANDOM = 0
# XXX This is really the same as is_resurrection (just without requiring existing scores)...
# After changing the set of tests:
#NUM_PREVIOUS_RANDOM = POPULATION_SIZE
NUM_CHILDREN = 1
# For experimentation:
#USE_TESTCASES = ["vqeg-mobilecalendar"]
# Small set for initial exploration:
#USE_TESTCASES = ["lavfi-pal75bars", "vqeg-mobilecalendar", "vqeg-harp", "ldv-shields", "ldv-stockholm"]
# Medium 625-line set:
#USE_TESTCASES = [
#    "lavfi-magenta", "lavfi-testsrc", "lavfi-pal75bars",
#    "vqeg-tree", "vqeg-barcelona", "vqeg-harp", "vqeg-movinggraphic", "vqeg-canoavalsesia",
#    "vqeg-f1car", "vqeg-fries", "vqeg-horizontalscrolling", "vqeg-rugby", "vqeg-mobilecalendar",
#    "ldv-parkrun", "ldv-shields", "ldv-stockholm",
#    ]
# Complete set:
USE_TESTCASES = sorted(testcases.keys())

# RNGs for creating random individuals, and making evolutionary choices.
# (The idea being that we cache the results of evaluating individuals, so
# generating the same random individuals saves some effort.)
create_random = random.Random(42)
choose_random = random.Random(73)

# Directory containing a cache of the individuals we've tried so far
hof_dir = os.path.join(cache_dir, "hof3d")
os.makedirs(hof_dir, exist_ok=True)

class Individual:
    def __init__(self, source, arg=None):
        # thresholds are immutable once an Individual is created.

        if source == "constant":
            self.thresholds = THRESHOLDS_SIZE * [arg]
        elif source == "random":
            self.thresholds = [create_random.randint(0, 100 / QUANTUM) * QUANTUM for i in range(THRESHOLDS_SIZE)]
        elif source == "copy":
            self.thresholds = arg
        else:
            raise ValueError("bad source: " + source)

        # Hash the thresholds to generate a unique filename
        self.hash = hashlib.sha256(" ".join(map(str, self.thresholds)).encode("UTF-8")).hexdigest()
        self.thresholds_name = os.path.join(hof_dir, self.hash + ".thresholds")
        self.scores_name = os.path.join(hof_dir, self.hash + ".scores")

        self.scores = None
        self.total_score = None

    def __repr__(self):
        return "[%s]" % (",".join(map(str, self.thresholds)))

    def write_thresholds(self):
        with open(self.thresholds_name + ".new", "w") as f:
            for z in range(THRESHOLDS_Z):
                for y in range(THRESHOLDS_Y):
                    for x in range(THRESHOLDS_X):
                        f.write("%.02f " % (self.thresholds[cell(x, y, z)] * 0.01))
                    f.write("\n")
                f.write("\n")
        os.rename(self.thresholds_name + ".new", self.thresholds_name)

    def read_scores(self):
        self.scores = {}
        try:
            with open(self.scores_name) as f:
                for line in f.readlines():
                    name, score = line.rstrip().split(",")
                    self.scores[name] = float(score)
        except IOError:
            pass

    def update_total(self):
        """Compute the total score as the product of all the scores we're currently using.
        Note this will give different results as the set of testcases changes,
        and it assumes there are scores for all current testcases."""
        self.total_score = 1.0
        for testcase_name in USE_TESTCASES:
            self.total_score *= self.scores[testcase_name]

    def write_scores(self):
        with open(self.scores_name + ".new", "w") as f:
            for name, score in sorted(self.scores.items()):
                f.write("%s,%f\n" % (name, score))
        os.rename(self.scores_name + ".new", self.scores_name)

def load_individual(filename, same_testcases=False):
    if not filename.endswith(".thresholds"):
        return None

    # Load the thresholds, checking they're the right size
    with open(os.path.join(hof_dir, filename)) as f:
        thresholds = [int(float(s) * 100) for s in f.read().rstrip().split()]
    if len(thresholds) != THRESHOLDS_SIZE:
        return None

    # Create the individual
    ind = Individual("copy", thresholds)
    if ind.hash != filename[:-len(".thresholds")]:
        # This shouldn't happen, but if it does (probably due to a leftover
        # file from earlier development), ignore it.
        return None
    ind.read_scores()

    if same_testcases:
        # Check we have scores for all the testcases we're using.
        for testcase_name in USE_TESTCASES:
            if testcase_name not in ind.scores:
                return None

    return ind

def show_stats():
    births = []
    mutations = {}

    for filename in sorted(os.listdir(hof_dir)):
        # We only want individuals with the current set of testcases
        ind = load_individual(filename, same_testcases=True)
        if ind is None:
            continue
        logging.info("reading %s", filename)

        birth = ind.scores.get("_birth")
        firstscore = ind.scores.get("_firstscore")
        parentscore = ind.scores.get("_parentscore")
        mutation = ind.scores.get("_mutation")
        delta = ind.scores.get("_delta", 0)
        radius = ind.scores.get("_radius", 0)

        if birth is not None:
            births.append((birth, ind))

        if (mutation is not None) and (firstscore is not None) and (parentscore is not None):
            mutation = int(mutation)
            delta = int(delta)
            if mutation < 6:
                key = "plane" + ("+%d" % delta if delta > 0 else "%d" % delta)
            elif mutation < 9:
                key = ("radius%d" % radius) + ("+%d" % delta if delta > 0 else "%d" % delta)
            elif mutation == 9:
                key = "crossover"
            elif mutation == 10:
                key = "crossover-radius%d" % radius
            else:
                key = "?"

            counts = mutations.setdefault(key, [0, 0])
            if firstscore > parentscore:
                counts[0] += 1
            counts[1] += 1

    births.sort()

    # Show the series of best results found over time.
    logging.info("writing result_history.png")
    best_inds = [births[0][1]]
    last_time = 0.0
    best_score = 0.0
    for birth, ind in births:
        ind.update_total()

        # It needs to be the best score seen so far...
        if ind.total_score < best_score:
            continue
        best_score = ind.total_score

        # And it needs to be at least 30m after the last one reported.
        if (birth - last_time) < (30 * 60):
            continue
        last_time = birth

        best_inds.append(ind)

    logging.info("%d inds", len(best_inds))
    COMPLEX_X = TILE_X // 2
    img_w = (len(best_inds) * (COMPLEX_X + 1)) + 1
    img_h = ((THRESHOLDS_Y + 1) * THRESHOLDS_Z) + 1
    img = PIL.Image.new("RGB", (img_w, img_h))
    draw = PIL.ImageDraw.Draw(img)
    # Plot to match ld-chroma-decoder's FFT visualisation in layout; the thresholds array
    # covers XTILE/8 <= x <= XTILE/4 (and reflects around XTILE/4)
    off_x = TILE_X // 8
    for i, ind in enumerate(best_inds):
        tl_x = (COMPLEX_X + 1) * i
        for z in range(THRESHOLDS_Z):
            tl_y = (THRESHOLDS_Y + 1) * z
            draw.rectangle([(tl_x, tl_y), (tl_x + COMPLEX_X + 1, tl_y + THRESHOLDS_Y + 1)], outline=0x004000, width=0)
            for y in range(THRESHOLDS_Y):
                for x in range(THRESHOLDS_X):
                    # Invert the value so black = always chroma, white = always luma
                    v = 255 - int((ind.thresholds[cell(x, y, z)] / 100.0) * 255)
                    img.putpixel((tl_x + 1 + x + off_x, tl_y + y + 1), (v << 16) | (v << 8) | v)
                    # And the reflection since the same threshold is used for both:
                    img.putpixel((tl_x + 1 + ((TILE_X // 2) - (x + off_x)), tl_y + y + 1), (v << 16) | (v << 8) | v)
    img.save("result_history.png")

    logging.info("writing births_scores.csv")
    with open("births_scores.csv", "w") as f:
        for birth, ind in births:
            f.write("%f,%f\n" % (birth, ind.total_score))

    logging.info("writing mutations.csv")
    with open("mutations.csv", "w") as f:
        for mutation, counts in sorted(mutations.items()):
            f.write("%s,%d,%d\n" % (mutation, counts[0], counts[1]))

if sys.argv[1:] == ["--stats"]:
    show_stats()
    sys.exit(0)

# Start with a known-fairly-good configuration.
population = [Individual("constant", 45)]

def insert_individual(new_population, ind):
    """Returns True if inserted, False if it was already there."""
    for other in new_population:
        if other.hash == ind.hash:
            return False
    new_population.append(ind)
    return True

# And some that were successfully evolved in the past...
# transform2d:
WINNERS2D = [
    "26ed6bbcbde9cbd3bdc0a1ec486ed577fabf9bf30ba20c2fb6f73af238fc172a.thresholds",
    "cbe73a8e39276df08e30f7a2f5b1cba6ad1851047d4cbadd738e2a959aaa5565.thresholds",
    ]
# transform3d:
WINNERS = [
    # trained with small set
    "0d34daf137fba85442f13fd0aec0e0443e95936f65dad6205ae085346a8f282f.thresholds",
    "e899cc440d4a31b2c050d087f6327fbdf18086e7e3e9623855d827ffa40dac8b.thresholds",
    "666702d506a3be1ca90945f8a6df29917f27aebbca67b656d852629c6a4592a8.thresholds",
    "8e2571c0a52ad1b677de895200c1d686a3c5ad5f484729d20dadb6183f090be9.thresholds",
    "6e26c6260a86e661a833ff973c6224c56001ceed15c02d20fd67cd64c90dfbc7.thresholds",
    "f9b046c8d30ce9fcdb320d9abd1057114ad56aa9ae56c67bea48ea9df972fd9d.thresholds",
    # trained with medium set
    "0c8bff696dd57e425435254a0a986dd87d54299e1c37e23a121a4ebf9936058b.thresholds",
    "330072406ae98ef88ad716916f01380a736f2fdfe932fe5f3d70e9e7554a8c44.thresholds",
    "6980dacead6876acb187e5795eaeda88536748e9f9c3b96c1d062e9376144e70.thresholds",
    "2c38d1c2f93df0df8c50408982b49de2e3c72a51d7ece6ff2f426b1bfd7baacf.thresholds",
    "5081649f3b1c9ad6192f49799d9b70f187fddd698d42daa274154ca91887fe6e.thresholds",
    "66a2a936fb27ed2c76a942f70c46fa7865ef57c71e6d1cc4c00e533d56e7c23f.thresholds",
    "eb10fd53e3edc6e821859d95c8c42731530f03bfc2e3050f66e257b99844178c.thresholds",
    "05490bfd61cedec6f97f1b451680bfe440521e439cd7301fedbc946e3f25bc53.thresholds",
]
for filename in WINNERS:
    ind = load_individual(filename)
    if ind is None:
        continue

    logging.info("Using previous winner HoF individual %s", ind.hash)
    insert_individual(population, ind)

# Find previously-saved configurations that have the set of testcases we're
# using, so we implicitly continue from a previous run. The cost of including
# these is minimal because we don't have to run any new tests, so we include
# all of them in the first generation.
# We have to sort here so that the shuffle below is deterministic.
hof_filenames = sorted(os.listdir(hof_dir))
for filename in hof_filenames:
    ind = load_individual(filename, same_testcases=True)
    if ind is None:
        continue

    logging.info("Using complete HoF individual %s", ind.hash)
    insert_individual(population, ind)

# Randomly select at most POPULATION_SIZE previous configurations (that don't
# have a complete set of tests already), for variety. This means you can
# continue with a different set of tests.
choose_random.shuffle(hof_filenames, choose_random.random)
count = 0
for filename in hof_filenames:
    if count >= NUM_PREVIOUS_RANDOM:
        break

    ind = load_individual(filename)
    if ind is None:
        continue

    if insert_individual(population, ind):
        logging.info("Using random HoF individual %s", ind.hash)
        count += 1

generation = 0
while True:
    logging.info("-" * 70)
    logging.info("Generation %d, population %d", generation, len(population))

    # Read in existing scores
    for ind in population:
        ind.read_scores()

    # Evaluate all the individuals against all the testcases.
    # Since the testcase data is large (several gigabytes), evaluate all
    # individuals against each testcase before moving on to the next testcase.
    for testcase_name in USE_TESTCASES:
        logging.info("Evaluating with %s", testcase_name)
        testcase = testcases[testcase_name]

        for ind in population:
            if testcase_name in ind.scores:
                #logging.info("Already done individual %s - score %f", ind.hash, ind.scores[testcase_name])
                continue

            ind.write_thresholds()
            decoder_args = ["-f", "transform3d", "--transform-thresholds", ind.thresholds_name]
            psnr, ssim = evaluate(testcase, decoder_args)
            logging.info("Testcase %s individual %s PSNR %f SSIM %f", testcase_name, ind.hash, psnr, ssim)

            ind.scores[testcase_name] = ssim
            ind.write_scores()

    # Periodically resurrect a set of random older individuals for variety
    is_resurrection = (generation % 50) == 0
    resurrected = []
    if is_resurrection:
        hof_filenames = os.listdir(hof_dir)
        choose_random.shuffle(hof_filenames, choose_random.random)
        for filename in hof_filenames:
            ind = load_individual(filename, same_testcases=True)
            if ind is None:
                continue

            insert_individual(resurrected, ind)
            if len(resurrected) >= POPULATION_SIZE:
                break

    # Update total scores
    for ind in population + resurrected:
        ind.update_total()

        # Record total score if this is the first time we've done it
        if "_firstscore" not in ind.scores:
            # XXX Hack because I introduced this partway through a run
            if "_birth" not in ind.scores:
                ind.scores["_birth"] = os.stat(ind.scores_name).st_mtime

            ind.scores["_firstscore"] = ind.total_score
            ind.write_scores()

    # Sort the best individuals first and trim to max size
    population.sort(key=lambda ind: -ind.total_score)
    population = population[:POPULATION_SIZE]

    # Add in the resurrected individuals
    if is_resurrection:
        for ind in resurrected:
            if insert_individual(population, ind):
                logging.info("resurrecting %s", ind.hash)
        population.sort(key=lambda ind: -ind.total_score)

    # Show stats
    logging.info("Generation %d leaderboard:", generation)
    for ind in population:
        logging.info("- %s - total_score %f", ind.hash, ind.total_score)
    total_scores = [ind.total_score for ind in population]
    logging.info("Generation %d: median score %f, best score %f", generation,
                 statistics.median(total_scores), max(total_scores))

    # Generate new children
    new_population = population[:]
    want_children = POPULATION_SIZE if is_resurrection else NUM_CHILDREN
    while len(new_population) < len(population) + want_children:

        # Choose a parent -- usually the best one
        # XXX This is more an SA than GA approach...
        if choose_random.random() < 0.8:
            parent = population[0]
        else:
            parent = choose_random.choice(population)
        logging.info("child mutating from %s", parent.hash)

        # Mutate
        # XXX These are all conservative changes -- it may be better to have a "set
        # to random" mutation to avoid getting stuck in a local maximum. (But that
        # may not be a problem depending on what the search space looks like...)
        thresholds = parent.thresholds[:]
        mutation = choose_random.choice([0, 6, 6, 6, 9, 10])
        if is_resurrection:
            mutation = 10
        delta = 0
        radius = 0
        if mutation == 0:
            # Raise/lower plane

            want_x, want_y, want_z = -1, -1, -1
            delta = QUANTUM * choose_random.choice([-1, 1])
            axis = choose_random.randrange(3)
            if axis == 0:
                want_x = choose_random.randrange(THRESHOLDS_X)
            elif axis == 1:
                want_y = choose_random.randrange(THRESHOLDS_Y)
            elif axis == 2:
                want_z = choose_random.randrange(THRESHOLDS_Z)
            logging.info("raise plane %d,%d,%d by %d", want_x, want_y, want_z, delta)

            for z in range(THRESHOLDS_Z):
                if want_z != -1 and want_z != z:
                    continue
                for y in range(THRESHOLDS_Y):
                    if want_y != -1 and want_y != y:
                        continue
                    for x in range(THRESHOLDS_X):
                        if want_x != -1 and want_x != x:
                            continue
                        idx = cell(x, y, z)
                        thresholds[idx] = clamp_cell(thresholds[idx] + delta)

        elif mutation == 6:
            # Raise/lower radius

            delta = QUANTUM * choose_random.choice([-3, -2, -2, -1, -1, -1, 1, 1, 1, 2, 2, 3])
            cx = choose_random.randrange(THRESHOLDS_X)
            cy = choose_random.randrange(THRESHOLDS_Y)
            cz = choose_random.randrange(THRESHOLDS_Z)
            radius = choose_random.randrange(min(THRESHOLDS_X, THRESHOLDS_Y, THRESHOLDS_Z))
            logging.info("raise radius %d around %d,%d,%d by %d", radius, cx, cy, cz, delta)

            r2 = radius * radius
            for z in range(THRESHOLDS_Z):
                for y in range(THRESHOLDS_Y):
                    for x in range(THRESHOLDS_X):
                        if ((cy - y) * (cy - y)) + ((cx - x) * (cx - x)) + ((cz - z) * (cz - z)) < r2:
                            idx = cell(x, y, z)
                            thresholds[idx] = clamp_cell(thresholds[idx] + delta)

        elif mutation == 9:
            # Plane crossover

            # Select a second parent
            parent2 = parent
            if len(population) > 1:
                while parent.hash == parent2.hash:
                    parent2 = choose_random.choice(population)
            if choose_random.random() > 0.5:
                parent, parent2 = parent2, parent

            # Select a random plane
            x_split, y_split, z_split = 0, 0, 0
            axis = choose_random.randrange(3)
            if axis == 0:
                x_split = choose_random.randint(1, THRESHOLDS_X - 1)
            elif axis == 1:
                y_split = choose_random.randint(1, THRESHOLDS_Y - 1)
            elif axis == 2:
                z_split = choose_random.randint(1, THRESHOLDS_Z - 1)

            logging.info("crossover between %s and %s at %d,%d,%d", parent.hash, parent2.hash, x_split, y_split, z_split)

            # Join the two halves together
            for z in range(THRESHOLDS_Z):
                for y in range(THRESHOLDS_Y):
                    for x in range(THRESHOLDS_X):
                        idx = cell(x, y, z)
                        if x >= x_split and y >= y_split and z >= z_split:
                            thresholds[idx] = parent.thresholds[idx]
                        else:
                            thresholds[idx] = parent2.thresholds[idx]

        elif mutation == 10:
            # Radius crossover

            # Select a second parent
            parent2 = parent
            if len(population) > 1:
                while parent.hash == parent2.hash:
                    parent2 = choose_random.choice(population)

            # Select a radius
            cx = choose_random.randrange(THRESHOLDS_X)
            cy = choose_random.randrange(THRESHOLDS_Y)
            cz = choose_random.randrange(THRESHOLDS_Z)
            radius = choose_random.randrange((THRESHOLDS_X + THRESHOLDS_Y + THRESHOLDS_Z) // 3)

            logging.info("crossover-radius between %s and %s around %d,%d,%d radius %d",
                         parent.hash, parent2.hash, cx, cy, cz, radius)

            # Insert the patch from parent2
            r2 = radius * radius
            for z in range(THRESHOLDS_Z):
                for y in range(THRESHOLDS_Y):
                    for x in range(THRESHOLDS_X):
                        idx = cell(x, y, z)
                        if ((cy - y) * (cy - y)) + ((cx - x) * (cx - x)) + ((cz - z) * (cz - z)) < r2:
                            thresholds[idx] = parent2.thresholds[idx]

        else:
            raise ValueError("unknown mutation %d" % mutation)

        # Insert the new child, if it doesn't duplicate one we already have
        child = Individual("copy", thresholds)
        insert_individual(new_population, child)

        # Record information about the child's creation, for later stats
        child.read_scores()
        if "_birth" not in child.scores:
            child.scores["_birth"] = time.time()
            child.scores["_parentscore"] = parent.total_score
            child.scores["_bestscore"] = population[0].total_score
            child.scores["_mutation"] = float(mutation)
            child.scores["_delta"] = delta
            child.scores["_radius"] = radius
            child.write_scores()

    population = new_population
    generation += 1