exapunks-hackmatch-bot/analyze.py

#!/usr/bin/env python3
import os
import numpy as np
from PIL import Image
from parser import BLOCK_SIZE, COLUMN_VSHIFT, MIN_COLUMN_VAL, MIN_COLUMN_SAT, \
                   detect_blocks, detect_exa, detect_held, print_board, \
                   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 = sys.argv[2] == 'y' if len(sys.argv) > 2 else False

    board = Image.open('screens/board%d.png' % ident).convert('HSV')
    #board.crop((0, 0, board.width, detect_columns(board))).show()
    state = State.detect(board)
    blocks = list(detect_blocks(board))
    exa = detect_exa(board)
    held = detect_held(board, exa)
    print_board(blocks, exa, held)

    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)