1# Needed on case-insensitive filesystems
2from __future__ import absolute_import
3
4# Try to import PIL in either of the two ways it can be installed.
5try:
6    from PIL import Image, ImageDraw
7except ImportError:  # pragma: no cover
8    import Image
9    import ImageDraw
10
11import math
12
13
14class QRColorMask:
15    """
16    QRColorMask is used to color in the QRCode.
17    By the time apply_mask is called, the QRModuleDrawer of the StyledPilImage will have drawn all of the modules on the canvas
18    (the color of these modules will be mostly black, although antialiasing may result in gradiants)
19    In the base class, apply_mask is implemented such that the background color will remain, but the foreground pixels will be
20    replaced by a color determined by a call to get_fg_pixel. There is additional calculation done to preserve the gradiant artifacts
21    of antialiasing
22    All QRColorMask objects should be careful about RGB vs RGBA color spaces
23
24    For examples of what these look like, see doc/color_masks.png
25    """
26
27    back_color = (255,255,255)
28    has_transparency = False
29    paint_color = back_color
30
31    def initialize(self, styledPilImage, image):
32        self.paint_color = styledPilImage.paint_color
33
34    def apply_mask(self, image):
35        width, height = image.size
36        for x in range(width):
37            for y in range(height):
38                norm = self.extrap_color(self.back_color, self.paint_color, image.getpixel((x,y)))
39                if norm is not None:
40                    image.putpixel((x,y), self.interp_color(self.get_bg_pixel(image, x,y), self.get_fg_pixel(image, x,y), norm))
41                else:
42                    image.putpixel((x,y), self.get_bg_pixel(image, x,y))
43
44    def get_fg_pixel(self, image, x, y):
45        raise NotImplementedError("QRModuleDrawer.paint_fg_pixel")
46
47    def get_bg_pixel(self, image, x, y):
48        return self.back_color
49
50    # The following functions are helpful for color calculation:
51
52    # interpolate a number between two numbers
53    def interp_num(self, n1, n2, norm):
54        return int(n2 * norm + n1 * (1-norm))
55
56    # interpolate a color between two colorrs
57    def interp_color(self, col1, col2, norm):
58        return tuple(self.interp_num(col1[i], col2[i], norm) for i in range(len(col1)))
59
60    # find the interpolation coefficient between two numbers
61    def extrap_num(self, n1, n2, interped_num):
62        if n2 == n1:
63            return None
64        else:
65            return (interped_num - n1) / (n2 - n1)
66
67    # find the interpolation coefficient between two numbers
68    def extrap_color(self, col1, col2, interped_color):
69        normed = list(filter(lambda i: i is not None, [self.extrap_num(col1[i], col2[i], interped_color[i]) for i in range(len(col1))]))
70        if not normed:
71            return None
72        else:
73            return sum(normed) / len(normed)
74
75class SolidFillColorMask(QRColorMask):
76    """
77    Just fills in the background with one color and the foreground with another
78    """
79    def __init__(self, back_color = (255,255,255), front_color = (0,0,0)):
80        self.back_color = back_color
81        self.front_color = front_color
82        self.has_transparency = len(self.back_color) == 4
83
84    def apply_mask(self, image):
85        if self.back_color == (255,255,255) and self.front_color == (0,0,0):
86            # Optimization: the image is already drawn by QRModuleDrawer in black and white,
87            # so if these are also our mask colors we don't need to do anything.
88            # This is much faster than actually applying a mask.
89            pass
90        else:
91            # TODO there's probably a way to use PIL.ImageMath instead of doing the individual pixel comparisons
92            # that the base class uses, which would be a lot faster. (In fact doing this would probably remove
93            # the need for the B&W optimization above.)
94            QRColorMask.apply_mask(self, image)
95
96    def get_fg_pixel(self, image, x, y):
97        return self.front_color
98
99
100class RadialGradiantColorMask(QRColorMask):
101    """
102    Fills in the foreground with a radial gradiant from the center to the edge
103    """
104    def __init__(self, back_color = (255,255,255), center_color = (0,0,0), edge_color = (0,0,255)):
105        self.back_color = back_color
106        self.center_color = center_color
107        self.edge_color = edge_color
108        self.has_transparency = len(self.back_color) == 4
109
110    def get_fg_pixel(self, image, x, y):
111        width, _ = image.size
112        normedDistanceToCenter = math.sqrt((x - width/2) ** 2 + (y - width/2) ** 2) / (math.sqrt(2) * width/2)
113        return self.interp_color(self.center_color, self.edge_color, normedDistanceToCenter)
114
115class SquareGradiantColorMask(QRColorMask):
116    """
117    Fills in the foreground with a square gradiant from the center to the edge
118    """
119    def __init__(self, back_color = (255,255,255), center_color = (0,0,0), edge_color = (0,0,255)):
120        self.back_color = back_color
121        self.center_color = center_color
122        self.edge_color = edge_color
123        self.has_transparency = len(self.back_color) == 4
124
125    def get_fg_pixel(self, image, x, y):
126        width,_ = image.size
127        normedDistanceToCenter = max(abs(x - width/2), abs(y - width/2)) / (width/2)
128        return self.interp_color(self.center_color, self.edge_color, normedDistanceToCenter)
129
130
131class HorizontalGradiantColorMask(QRColorMask):
132    """
133    Fills in the foreground with a gradiant sweeping from the left to the right
134    """
135    def __init__(self, back_color = (255,255,255), left_color = (0,0,0), right_color = (0,0,255)):
136        self.back_color = back_color
137        self.left_color = left_color
138        self.right_color = right_color
139        self.has_transparency = len(self.back_color) == 4
140
141    def get_fg_pixel(self, image, x, y):
142        width,_ = image.size
143        return self.interp_color(self.left_color, self.right_color, x / width)
144
145class VerticalGradiantColorMask(QRColorMask):
146    """
147    Fills in the forefround with a gradiant sweeping from the top to the bottom
148    """
149    def __init__(self, back_color = (255,255,255), top_color = (0,0,0), bottom_color = (0,0,255)):
150        self.back_color = back_color
151        self.top_color = top_color
152        self.bottom_color = bottom_color
153        self.has_transparency = len(self.back_color) == 4
154
155    def get_fg_pixel(self, image, x, y):
156        width,_ = image.size
157        return self.interp_color(self.top_color, self.bottom_color, y / width)
158
159class ImageColorMask(QRColorMask):
160    """
161    Fills in the foreground with pixels from another image, either passed by path or passed by image object
162    """
163    def __init__(self, back_color = (255,255,255), color_mask_path=None, color_mask_image=None):
164        self.back_color = back_color
165        if color_mask_image:
166            self.color_img = color_mask_image
167        else:
168            self.color_img = Image.open(color_mask_path)
169
170        self.has_transparency = len(self.back_color) == 4
171
172    def initialize(self, styledPilImage, image):
173        self.paint_color = styledPilImage.paint_color
174        self.color_img = self.color_img.resize(image.size)
175
176    def get_fg_pixel(self, image, x, y):
177        width,_ = image.size
178        return self.color_img.getpixel((x,y))
179