exapunks-hackmatch-bot/analyze.py

#!/usr/bin/env python3
import os
import numpy as np
from PIL import Image
from detection import BLOCK_SIZE, COLUMN_VSHIFT, MIN_COLUMN_VAL, \
                      MIN_COLUMN_SAT, detect_columns
from strategy import State


def cut_board(board):
    y = detect_columns(board) - BLOCK_SIZE
    i = row = 0
    while y >= 0:
        for col, shift in enumerate(COLUMN_VSHIFT):
            x = col * BLOCK_SIZE
            block = board.crop((x, y + shift,
                                x + BLOCK_SIZE, y + shift + BLOCK_SIZE))
            yield col, row, block
            i += 1
        y -= BLOCK_SIZE
        row += 1


def make_bitmap(board):
    h, s, v = board.split()
    v = [int(v > MIN_COLUMN_VAL and s > MIN_COLUMN_SAT) * 255
         for h, s, v in zip(h.getdata(), s.getdata(), v.getdata())]
    v = Image.fromarray(np.uint8(v).reshape((board.height, board.width)))
    return Image.merge('RGB', (v, v, v))


if __name__ == '__main__':
    import sys
    ident = int(sys.argv[1])
    extensive = len(sys.argv) > 2 and sys.argv[2] == 'y'

    board = Image.open('screens/board%d.png' % ident).convert('HSV')
    #board.crop((0, 0, board.width, detect_columns(board))).show()
    state = State.detect(board)
    state.print()

    if extensive:
        make_bitmap(board).save('screens/bitmap%d.png' % ident)

        os.makedirs('blocks/board%d' % ident, exist_ok=True)
        for row, col, block in cut_board(board):
            f = 'blocks/board%d/row%d-col%d.png' % (ident, row, col)
            block.convert('RGB').save(f)