forked from Moussa/3D-Models-automaton
-
Notifications
You must be signed in to change notification settings - Fork 2
/
imageprocessor.py
155 lines (140 loc) · 6.1 KB
/
imageprocessor.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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
"""
Deals with image processing. This file is the bottleneck of the process, and
is thus using numpy for efficiency. As a result, it is not very easy to read.
"""
from PIL.Image import ANTIALIAS, fromarray, new
from numpy import array, dstack, inner, uint8, where
class ImageProcessor(object):
"""
A class to handle all the imageprocessing done on the screenshots.
Deals with blending (finding transparency), cropping, and stitching.
"""
def __init__(self, y_rotations, x_rotations):
self.target_dimension = 280
self.target_size = 512 * 1024 # 512 KB
self.cropping = {'left': [], 'top': [], 'right': [], 'bottom': []}
self.images = []
self.y_rotations = y_rotations
self.x_rotations = 2*x_rotations+1 # Total vertical rotations
def find_minimum_bounds(self, white_image, black_image):
"""
Finds the extrema of the black pixels of the first black image,
in order to find the limits of the HLMV viewport.
We do this by sampling 3 lines in each direction (horizontal / vertical)
"""
black_arr = array(black_image, dtype=int)
horizontal_samples = [
len(black_arr[:, 0, 0]) * 1 // 4,
len(black_arr[:, 0, 0]) * 2 // 4,
len(black_arr[:, 0, 0]) * 3 // 4,
]
# This returns a 2d array of [[x coordinates], [y coordinates]] representing the zeros (all black pixels) in the image along 3 horizontal lines
h_black_pixels = where(black_arr[horizontal_samples, :, :].sum(axis=2) == 0)
vertical_samples = [
len(black_arr[0, :, 0]) * 1 // 4,
len(black_arr[0, :, 0]) * 2 // 4,
len(black_arr[0, :, 0]) * 3 // 4,
]
# This returns a 2d array of [[x coordinates], [y coordinates]] representing the zeros (all black pixels) in the image along 3 vertical lines
v_black_pixels = where(black_arr[:, vertical_samples, :].sum(axis=2) == 0)
return (
h_black_pixels[1].min(), # Left
v_black_pixels[0].min(), # Top
h_black_pixels[1].max(), # Right
v_black_pixels[0].max(), # Bottom
)
def blend(self, white_image, black_image):
"""
Blends the two images into an alpha image using percieved luminescence.
https://en.wikipedia.org/wiki/Luma_(video)#Use_of_relative_luminance
Then, finds the closest-cropped lines that are all white.
Uses numpy because traversing python arrays is very slow.
"""
# white_arr[:, :, 0] means:
# Treat the white image like a 3D array (x, y, RGB).
# Then, select all X coordinates, all Y coordinates, but only Z[0] (red).
# Similarly, white_arr[:, :, 1] is all of the green values.
# "inner" is the sum of the products of each pair of elements.
# So, we subtract the red values between white and black,
# multiply by .299, then add the results to green and blue.
# This needs to be dtype=int to prevent an overflow when adding
white_arr = array(white_image, dtype=int)
black_arr = array(black_image, dtype=int)
blended_arr = dstack((
(white_arr[:, :, 0] + black_arr[:, :, 0])/2,
(white_arr[:, :, 1] + black_arr[:, :, 1])/2,
(white_arr[:, :, 2] + black_arr[:, :, 2])/2,
255 - inner(white_arr - black_arr, [.299, .587, .114])
))
# Calculate crop lines by looking for all-white && all-black pixels, i.e. places where the luma is zero.
# np.any() will return 'True' for any rows which contain nonzero integers (because zero is Falsy).
# Then, we use nonzero() to get the only indices which are 'True', which are the rows with content.
# (nonzero returns a tuple for some reason, so we also have to [0] it.)
horizontal = blended_arr[:, :, 3].any(axis=0).nonzero()[0]
vertical = blended_arr[:, :, 3].any(axis=1).nonzero()[0]
self.cropping['left'].append(horizontal[0])
self.cropping['top'].append(vertical[0])
self.cropping['right'].append(horizontal[-1])
self.cropping['bottom'].append(vertical[-1])
# This needs to be a uint8 to render correctly.
blended_image = fromarray(blended_arr.astype(uint8), mode='RGBA')
blended_image = blended_image.crop((
horizontal[0],
vertical[0],
horizontal[-1],
vertical[-1]
))
self.images.append(blended_image)
def stitch(self):
"""
Crops the images to a shared size, then pastes them together.
Prompts for login and uploads to the wiki when done.
"""
# Determining crop bounds
min_cropping = (
min(self.cropping['left']),
min(self.cropping['top']),
max(self.cropping['right']),
max(self.cropping['bottom'])
)
print('Min cropping: ' + str(min_cropping))
max_frame_size = (
min_cropping[2] - min_cropping[0],
min_cropping[3] - min_cropping[1]
)
print('Max frame size: ' + str(max_frame_size))
target_ratio = self.target_dimension / max(max_frame_size)
print('Target scaling ratio: %f' % target_ratio)
max_frame_size = (
int(target_ratio * max_frame_size[0]),
int(target_ratio * max_frame_size[1])
)
print('Scaled max frame size: ' + str(max_frame_size))
# Pasting together
full_image = new(mode='RGBA', color=(255, 255, 255, 255), size=((
(max_frame_size[0]+1)*self.y_rotations*self.x_rotations,
max_frame_size[1]
)))
curr_offset = 0
offset_map = []
for i, image in enumerate(self.images):
image = image.resize((
int(image.width*target_ratio),
int(image.height*target_ratio),
), ANTIALIAS)
left_crop = int(target_ratio*(self.cropping['left'][i]-min_cropping[0]))
top_crop = int(target_ratio*(self.cropping['top'][i]-min_cropping[1]))
full_image.paste(image, (curr_offset, top_crop), image)
# Offset map adds 1 manually for some reason
offset_map += [curr_offset-i, image.height, left_crop]
# Increase by 1 each time to add a 1px gap
curr_offset += image.width+1
full_image = full_image.crop((
0,
0,
curr_offset,
max_frame_size[1],
))
full_offset_map = "%d,%d,%d,%d," % (curr_offset, max_frame_size[0], max_frame_size[1], self.x_rotations)
full_offset_map += ",".join([str(o) for o in offset_map])
return (full_image, full_offset_map)