-
Notifications
You must be signed in to change notification settings - Fork 0
/
reconstruct.py
78 lines (75 loc) · 3.31 KB
/
reconstruct.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import numpy as np
from os import path
from sys import argv
from PIL import Image, ImageDraw
from compress_pickle import load
import functions as fun
import settings as sel
if fun.isNotebook():
(fPath, fName) = (
'/home/chipdelmal/Documents/Sync/LegoOptimizer/',
'marioTanooki.png'
)
else:
(fPath, fName) = (argv[1], argv[2])
BLOCKS_ALPHA = 1
(RB_COL, RN_COL, BK_COL) = ((0, 0, 0, 255), (255, 0, 0, 255), (255, 255, 255, 0))
(LW, LG) = (2, 1)
###############################################################################
# Get user selections from file
###############################################################################
SCALER = sel.USER_SEL['scaler']
###############################################################################
# Load image and decoded data
###############################################################################
img = Image.open(path.join(fPath, fName.split(".png")[0]+'_DWN.png'))
img = img.convert('RGBA')
img = img.resize((np.array(img.size)*SCALER).astype(int), resample=0)
# Read decoded data -----------------------------------------------------------
dFName = path.join(fPath, fName.split('.png')[0])+'_Decoded.pkl'
decoded = load(dFName)
###############################################################################
# Add block grid
###############################################################################
draw = ImageDraw.Draw(img)
# for x in range(0, img.width, SCALER):
# line = ((x, 0), (x, img.height))
# draw.line(line, fill=BK_COL, width=LG)
# for y in range(0, img.height, SCALER):
# line = ((0, y), (img.width, y))
# draw.line(line, fill=BK_COL, width=LG)
###############################################################################
# Annotate rectangles
###############################################################################
# Add frame all around image --------------------------------------------------
draw.rectangle(((0, 0), (img.width-LW/2, img.height-LW/2)), outline=RB_COL, width=LW)
# Iterate through the decoded array (rix: row index)
for rix in range(len(decoded)):
# Get row information (dRow: decoded row)
dRow = decoded[rix]
# Get the iterators started
(row, col) = (rix, 0)
# Load info from the block index
for bix in range(len(dRow)):
(bColor, bLensVct) = dRow[bix]
bColsLen = len(bLensVct)
for bCols in range(bColsLen):
# Setup the drawer to the left-top corner
tlCrnr = (SCALER*col+LW/2, SCALER*row+LW/2)
# The length of the block is defined by the array (height is constant for all)
(w, h) = (bLensVct[bCols]*SCALER, 1*SCALER)
rectCol = RB_COL if (w > 0) else RN_COL
w = abs(w)
blocks = (tlCrnr, (tlCrnr[0]+w-LW/2, tlCrnr[1]+h-LW/2))
# Draw the resulting block
draw = ImageDraw.Draw(img)
# draw.rectangle(blocks, fill=(*bColor, int(255*BLOCKS_ALPHA)))
draw.rectangle(blocks, outline=rectCol, width=LW)
# Shift column iterator
col = col + abs(bLensVct[bCols])
###############################################################################
# Export Resulting Image
###############################################################################
dFName = path.join(fPath, fName.split('.png')[0])+'_Lego.png'
img.save(dFName)
img.close()