ksnctf C92 E1 Mysterious Light

qr2txt.py

# -*- coding: utf-8 -*-

from PIL import Image

dth, lth = 50, 175
marks = {'dark': 'X', 'light': '_', 'unknown': '?'}


def get_pix(filename):
    with Image.open(filename) as img:
        width, height = img.size
        imgpix = [[] for i in range(height)]
        for y in range(height):
            for x in range(width):
                imgpix[y].append(img.getpixel((x, y)))
    return imgpix


def pix2txt(imgpix, unknown_patterns=[], skip=1):
    imgtxt = []
    for y in range(0, len(imgpix), skip):
        imgtxt.append('')
        for x in range(0, len(imgpix[0]), skip):
            pix = imgpix[y][x]
            for uptn in unknown_patterns:
                if uptn(pix):
                    mark = marks['unknown']
                    break
            else:
                avgpix = sum(pix)/3
                if avgpix < dth:
                    mark = marks['dark']
                elif avgpix > lth:
                    mark = marks['light']
                else:
                    mark = marks['unknown']
            imgtxt[int(y/skip)] += mark
    return imgtxt


def is_the_color(rgb, color):
    l = color
    ds = [i for i in range(3) if i != 1]
    return (rgb[l] > lth) and (rgb[ds[0]] < dth) and (rgb[ds[1]] < dth)

def is_red(rgb): return is_the_color(rgb, 0)
def is_blue(rgb): return is_the_color(rgb, 1)
def is_green(rgb): return is_the_color(rgb, 2)


def trim(imgtxt, udlr):
    trimtxt = []
    for y in range(udlr[0], len(imgtxt)-udlr[1]):
        trimtxt.append(imgtxt[y][udlr[2]:-udlr[3]])
    return trimtxt


def main(infilename, skip=1, trim_udlr=[0, 0, 0, 0]):
    imgpix = get_pix(infilename)
    imgtxt = pix2txt(imgpix, [is_red, is_blue, is_green], skip)
    imgtxt = trim(imgtxt, trim_udlr)
    outfilename = infilename.split('.')[0] + '.txt'
    with open(outfilename, 'w') as f:
        for line in imgtxt:
            f.write(line + '\n')


if __name__ == '__main__':
    main('flag2.png', 8, [4, 4, 4, 4])