-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathbrushstroke_test.py
100 lines (83 loc) · 3.74 KB
/
brushstroke_test.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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import numpy as np
from DDP import *
import cv2
import pdb
def test_reward_model():
SAMPLE_IMAGE_PATH = 'brushstroke_sample_001.jpg'
image_desired = cv2.imread(SAMPLE_IMAGE_PATH)
image_desired = cv2.cvtColor(image_desired, cv2.COLOR_BGR2GRAY)
ret_info, image_desired = cv2.threshold(image_desired, 127, 255, cv2.THRESH_BINARY)
cv2.imshow('image_desired', image_desired)
n_white = cv2.countNonZero(image_desired)
n_black = image_desired.shape[0] * image_desired.shape[1] - n_white
print('#white: {}, #black: {}'.format(n_white, n_black))
cv2.waitKey()
def test_layers_extraction():
SAMPLE_IMAGE_PATH = 'brushstroke_sample_006.jpg'
DIVISION_LEVEL = 1
SAVE_DIR = 'save'
SAVE_IDENTIFIER = 'brushstroke_test'
SHOULD_SAVE_LAYERS = True
SHOULD_SAVE_RECOMPOSED_IMAGE = True
image_desired = cv2.imread(SAMPLE_IMAGE_PATH)
cv2.imshow('image_desired', image_desired)
print('divides color space into segments...', end='\r', flush=True)
portion = int(256 / (DIVISION_LEVEL + 1))
assert(portion * (DIVISION_LEVEL + 1) == 256)
segments = []
for i in range(DIVISION_LEVEL + 1):
segments.append(portion * i)
segments.append(255)
print('\033[Kcolor space segments ready (segments per channel: {})'.format(len(segments)))
print('prepares descriptions of layers...', end='\r', flush=True)
layers = []
for r in segments:
for g in segments:
for b in segments:
layer = {}
layer['color'] = np.array([r, g, b])
layer['luminance'] = np.array([0.2126, 0.7152, 0.0722]).dot(layer['color'])
layer['image'] = np.zeros((image_desired.shape[0], image_desired.shape[1]), np.uint8)
layers.append(layer)
print('\033[Klayer descriptions ready (layers: {})'.format(len(layers)))
print('generates layer images...', end='\r', flush=True)
for row in range(len(image_desired)):
for col in range(len(image_desired[row])):
pixel = image_desired[row][col]
minNorm = 256
layerIndex = -1
for l in range(len(layers)):
norm = np.linalg.norm(pixel - layers[l]['color'])
if norm < minNorm:
layerIndex = l
minNorm = norm
layers[layerIndex]['image'][row][col] = 255
print('\033[Kgenerates layer images... (completion: {}%)'.format(np.round(100. * float(row) / float(len(image_desired)), 1)), end='\r', flush=True)
print('\033[Klayer images ready (images: {})'.format(len(layers)))
print('recomposes image from layers...', end='\r', flush=True)
image_recomposed = np.zeros(image_desired.shape, np.uint8)
image_layers = []
for l in range(len(layers)):
layer = layers[l]
image_layer = np.zeros(image_desired.shape, np.uint8)
for row in range(len(layer['image'])):
for col in range(len(layer['image'][row])):
if layer['image'][row][col] == 255:
image_recomposed[row][col] = layer['color']
image_layer[row][col] = layer['color']
print('\033[Krecomposes image from layers... (layer: {}/{}, completion: {}%)'.format(l, len(layers), np.round(100. * float(row) / float(len(layer['image']))), 1), end='\r', flush=True)
image_layer[0][0] = np.array([255, 255, 255])
image_layers.append(image_layer)
if SHOULD_SAVE_LAYERS:
save_layer_path = SAVE_DIR + '/' + SAVE_IDENTIFIER + '_layer_' + str(l) + '.jpg'
cv2.imwrite(save_layer_path, image_layer)
cv2.imshow('image_layer_' + str(l), image_layer)
print('\033[Kimage recomposed from layers (layers: {})'.format(len(layers)))
if SHOULD_SAVE_RECOMPOSED_IMAGE:
save_recomposed_image_path = SAVE_DIR + '/' + SAVE_IDENTIFIER + '_recomposed.jpg'
cv2.imwrite(save_recomposed_image_path, image_recomposed)
cv2.imshow('image_recomposed', image_recomposed)
cv2.waitKey()
if __name__ == '__main__':
#test_reward_model()
test_layers_extraction()