Bläddra i källkod

Add screenshot analysis utility

Taddeus Kroes 5 år sedan
förälder
incheckning
6cadb61cba
1 ändrade filer med 52 tillägg och 0 borttagningar
  1. 52 0
      analyze.py

+ 52 - 0
analyze.py

@@ -0,0 +1,52 @@
+#!/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)