ImageMath.py 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265
  1. #
  2. # The Python Imaging Library
  3. # $Id$
  4. #
  5. # a simple math add-on for the Python Imaging Library
  6. #
  7. # History:
  8. # 1999-02-15 fl Original PIL Plus release
  9. # 2005-05-05 fl Simplified and cleaned up for PIL 1.1.6
  10. # 2005-09-12 fl Fixed int() and float() for Python 2.4.1
  11. #
  12. # Copyright (c) 1999-2005 by Secret Labs AB
  13. # Copyright (c) 2005 by Fredrik Lundh
  14. #
  15. # See the README file for information on usage and redistribution.
  16. #
  17. from __future__ import annotations
  18. import builtins
  19. from . import Image, _imagingmath
  20. class _Operand:
  21. """Wraps an image operand, providing standard operators"""
  22. def __init__(self, im):
  23. self.im = im
  24. def __fixup(self, im1):
  25. # convert image to suitable mode
  26. if isinstance(im1, _Operand):
  27. # argument was an image.
  28. if im1.im.mode in ("1", "L"):
  29. return im1.im.convert("I")
  30. elif im1.im.mode in ("I", "F"):
  31. return im1.im
  32. else:
  33. msg = f"unsupported mode: {im1.im.mode}"
  34. raise ValueError(msg)
  35. else:
  36. # argument was a constant
  37. if isinstance(im1, (int, float)) and self.im.mode in ("1", "L", "I"):
  38. return Image.new("I", self.im.size, im1)
  39. else:
  40. return Image.new("F", self.im.size, im1)
  41. def apply(self, op, im1, im2=None, mode=None):
  42. im1 = self.__fixup(im1)
  43. if im2 is None:
  44. # unary operation
  45. out = Image.new(mode or im1.mode, im1.size, None)
  46. im1.load()
  47. try:
  48. op = getattr(_imagingmath, op + "_" + im1.mode)
  49. except AttributeError as e:
  50. msg = f"bad operand type for '{op}'"
  51. raise TypeError(msg) from e
  52. _imagingmath.unop(op, out.im.id, im1.im.id)
  53. else:
  54. # binary operation
  55. im2 = self.__fixup(im2)
  56. if im1.mode != im2.mode:
  57. # convert both arguments to floating point
  58. if im1.mode != "F":
  59. im1 = im1.convert("F")
  60. if im2.mode != "F":
  61. im2 = im2.convert("F")
  62. if im1.size != im2.size:
  63. # crop both arguments to a common size
  64. size = (min(im1.size[0], im2.size[0]), min(im1.size[1], im2.size[1]))
  65. if im1.size != size:
  66. im1 = im1.crop((0, 0) + size)
  67. if im2.size != size:
  68. im2 = im2.crop((0, 0) + size)
  69. out = Image.new(mode or im1.mode, im1.size, None)
  70. im1.load()
  71. im2.load()
  72. try:
  73. op = getattr(_imagingmath, op + "_" + im1.mode)
  74. except AttributeError as e:
  75. msg = f"bad operand type for '{op}'"
  76. raise TypeError(msg) from e
  77. _imagingmath.binop(op, out.im.id, im1.im.id, im2.im.id)
  78. return _Operand(out)
  79. # unary operators
  80. def __bool__(self):
  81. # an image is "true" if it contains at least one non-zero pixel
  82. return self.im.getbbox() is not None
  83. def __abs__(self):
  84. return self.apply("abs", self)
  85. def __pos__(self):
  86. return self
  87. def __neg__(self):
  88. return self.apply("neg", self)
  89. # binary operators
  90. def __add__(self, other):
  91. return self.apply("add", self, other)
  92. def __radd__(self, other):
  93. return self.apply("add", other, self)
  94. def __sub__(self, other):
  95. return self.apply("sub", self, other)
  96. def __rsub__(self, other):
  97. return self.apply("sub", other, self)
  98. def __mul__(self, other):
  99. return self.apply("mul", self, other)
  100. def __rmul__(self, other):
  101. return self.apply("mul", other, self)
  102. def __truediv__(self, other):
  103. return self.apply("div", self, other)
  104. def __rtruediv__(self, other):
  105. return self.apply("div", other, self)
  106. def __mod__(self, other):
  107. return self.apply("mod", self, other)
  108. def __rmod__(self, other):
  109. return self.apply("mod", other, self)
  110. def __pow__(self, other):
  111. return self.apply("pow", self, other)
  112. def __rpow__(self, other):
  113. return self.apply("pow", other, self)
  114. # bitwise
  115. def __invert__(self):
  116. return self.apply("invert", self)
  117. def __and__(self, other):
  118. return self.apply("and", self, other)
  119. def __rand__(self, other):
  120. return self.apply("and", other, self)
  121. def __or__(self, other):
  122. return self.apply("or", self, other)
  123. def __ror__(self, other):
  124. return self.apply("or", other, self)
  125. def __xor__(self, other):
  126. return self.apply("xor", self, other)
  127. def __rxor__(self, other):
  128. return self.apply("xor", other, self)
  129. def __lshift__(self, other):
  130. return self.apply("lshift", self, other)
  131. def __rshift__(self, other):
  132. return self.apply("rshift", self, other)
  133. # logical
  134. def __eq__(self, other):
  135. return self.apply("eq", self, other)
  136. def __ne__(self, other):
  137. return self.apply("ne", self, other)
  138. def __lt__(self, other):
  139. return self.apply("lt", self, other)
  140. def __le__(self, other):
  141. return self.apply("le", self, other)
  142. def __gt__(self, other):
  143. return self.apply("gt", self, other)
  144. def __ge__(self, other):
  145. return self.apply("ge", self, other)
  146. # conversions
  147. def imagemath_int(self):
  148. return _Operand(self.im.convert("I"))
  149. def imagemath_float(self):
  150. return _Operand(self.im.convert("F"))
  151. # logical
  152. def imagemath_equal(self, other):
  153. return self.apply("eq", self, other, mode="I")
  154. def imagemath_notequal(self, other):
  155. return self.apply("ne", self, other, mode="I")
  156. def imagemath_min(self, other):
  157. return self.apply("min", self, other)
  158. def imagemath_max(self, other):
  159. return self.apply("max", self, other)
  160. def imagemath_convert(self, mode):
  161. return _Operand(self.im.convert(mode))
  162. ops = {}
  163. for k, v in list(globals().items()):
  164. if k[:10] == "imagemath_":
  165. ops[k[10:]] = v
  166. def eval(expression, _dict={}, **kw):
  167. """
  168. Evaluates an image expression.
  169. :param expression: A string containing a Python-style expression.
  170. :param options: Values to add to the evaluation context. You
  171. can either use a dictionary, or one or more keyword
  172. arguments.
  173. :return: The evaluated expression. This is usually an image object, but can
  174. also be an integer, a floating point value, or a pixel tuple,
  175. depending on the expression.
  176. """
  177. # build execution namespace
  178. args = ops.copy()
  179. for k in list(_dict.keys()) + list(kw.keys()):
  180. if "__" in k or hasattr(builtins, k):
  181. msg = f"'{k}' not allowed"
  182. raise ValueError(msg)
  183. args.update(_dict)
  184. args.update(kw)
  185. for k, v in args.items():
  186. if hasattr(v, "im"):
  187. args[k] = _Operand(v)
  188. compiled_code = compile(expression, "<string>", "eval")
  189. def scan(code):
  190. for const in code.co_consts:
  191. if type(const) is type(compiled_code):
  192. scan(const)
  193. for name in code.co_names:
  194. if name not in args and name != "abs":
  195. msg = f"'{name}' not allowed"
  196. raise ValueError(msg)
  197. scan(compiled_code)
  198. out = builtins.eval(expression, {"__builtins": {"abs": abs}}, args)
  199. try:
  200. return out.im
  201. except AttributeError:
  202. return out