ImageMath.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314
  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. TYPE_CHECKING = False
  21. if TYPE_CHECKING:
  22. from collections.abc import Callable
  23. from types import CodeType
  24. from typing import Any
  25. class _Operand:
  26. """Wraps an image operand, providing standard operators"""
  27. def __init__(self, im: Image.Image):
  28. self.im = im
  29. def __fixup(self, im1: _Operand | float) -> Image.Image:
  30. # convert image to suitable mode
  31. if isinstance(im1, _Operand):
  32. # argument was an image.
  33. if im1.im.mode in ("1", "L"):
  34. return im1.im.convert("I")
  35. elif im1.im.mode in ("I", "F"):
  36. return im1.im
  37. else:
  38. msg = f"unsupported mode: {im1.im.mode}"
  39. raise ValueError(msg)
  40. else:
  41. # argument was a constant
  42. if isinstance(im1, (int, float)) and self.im.mode in ("1", "L", "I"):
  43. return Image.new("I", self.im.size, im1)
  44. else:
  45. return Image.new("F", self.im.size, im1)
  46. def apply(
  47. self,
  48. op: str,
  49. im1: _Operand | float,
  50. im2: _Operand | float | None = None,
  51. mode: str | None = None,
  52. ) -> _Operand:
  53. im_1 = self.__fixup(im1)
  54. if im2 is None:
  55. # unary operation
  56. out = Image.new(mode or im_1.mode, im_1.size, None)
  57. try:
  58. op = getattr(_imagingmath, f"{op}_{im_1.mode}")
  59. except AttributeError as e:
  60. msg = f"bad operand type for '{op}'"
  61. raise TypeError(msg) from e
  62. _imagingmath.unop(op, out.getim(), im_1.getim())
  63. else:
  64. # binary operation
  65. im_2 = self.__fixup(im2)
  66. if im_1.mode != im_2.mode:
  67. # convert both arguments to floating point
  68. if im_1.mode != "F":
  69. im_1 = im_1.convert("F")
  70. if im_2.mode != "F":
  71. im_2 = im_2.convert("F")
  72. if im_1.size != im_2.size:
  73. # crop both arguments to a common size
  74. size = (
  75. min(im_1.size[0], im_2.size[0]),
  76. min(im_1.size[1], im_2.size[1]),
  77. )
  78. if im_1.size != size:
  79. im_1 = im_1.crop((0, 0) + size)
  80. if im_2.size != size:
  81. im_2 = im_2.crop((0, 0) + size)
  82. out = Image.new(mode or im_1.mode, im_1.size, None)
  83. try:
  84. op = getattr(_imagingmath, f"{op}_{im_1.mode}")
  85. except AttributeError as e:
  86. msg = f"bad operand type for '{op}'"
  87. raise TypeError(msg) from e
  88. _imagingmath.binop(op, out.getim(), im_1.getim(), im_2.getim())
  89. return _Operand(out)
  90. # unary operators
  91. def __bool__(self) -> bool:
  92. # an image is "true" if it contains at least one non-zero pixel
  93. return self.im.getbbox() is not None
  94. def __abs__(self) -> _Operand:
  95. return self.apply("abs", self)
  96. def __pos__(self) -> _Operand:
  97. return self
  98. def __neg__(self) -> _Operand:
  99. return self.apply("neg", self)
  100. # binary operators
  101. def __add__(self, other: _Operand | float) -> _Operand:
  102. return self.apply("add", self, other)
  103. def __radd__(self, other: _Operand | float) -> _Operand:
  104. return self.apply("add", other, self)
  105. def __sub__(self, other: _Operand | float) -> _Operand:
  106. return self.apply("sub", self, other)
  107. def __rsub__(self, other: _Operand | float) -> _Operand:
  108. return self.apply("sub", other, self)
  109. def __mul__(self, other: _Operand | float) -> _Operand:
  110. return self.apply("mul", self, other)
  111. def __rmul__(self, other: _Operand | float) -> _Operand:
  112. return self.apply("mul", other, self)
  113. def __truediv__(self, other: _Operand | float) -> _Operand:
  114. return self.apply("div", self, other)
  115. def __rtruediv__(self, other: _Operand | float) -> _Operand:
  116. return self.apply("div", other, self)
  117. def __mod__(self, other: _Operand | float) -> _Operand:
  118. return self.apply("mod", self, other)
  119. def __rmod__(self, other: _Operand | float) -> _Operand:
  120. return self.apply("mod", other, self)
  121. def __pow__(self, other: _Operand | float) -> _Operand:
  122. return self.apply("pow", self, other)
  123. def __rpow__(self, other: _Operand | float) -> _Operand:
  124. return self.apply("pow", other, self)
  125. # bitwise
  126. def __invert__(self) -> _Operand:
  127. return self.apply("invert", self)
  128. def __and__(self, other: _Operand | float) -> _Operand:
  129. return self.apply("and", self, other)
  130. def __rand__(self, other: _Operand | float) -> _Operand:
  131. return self.apply("and", other, self)
  132. def __or__(self, other: _Operand | float) -> _Operand:
  133. return self.apply("or", self, other)
  134. def __ror__(self, other: _Operand | float) -> _Operand:
  135. return self.apply("or", other, self)
  136. def __xor__(self, other: _Operand | float) -> _Operand:
  137. return self.apply("xor", self, other)
  138. def __rxor__(self, other: _Operand | float) -> _Operand:
  139. return self.apply("xor", other, self)
  140. def __lshift__(self, other: _Operand | float) -> _Operand:
  141. return self.apply("lshift", self, other)
  142. def __rshift__(self, other: _Operand | float) -> _Operand:
  143. return self.apply("rshift", self, other)
  144. # logical
  145. def __eq__(self, other: _Operand | float) -> _Operand: # type: ignore[override]
  146. return self.apply("eq", self, other)
  147. def __ne__(self, other: _Operand | float) -> _Operand: # type: ignore[override]
  148. return self.apply("ne", self, other)
  149. def __lt__(self, other: _Operand | float) -> _Operand:
  150. return self.apply("lt", self, other)
  151. def __le__(self, other: _Operand | float) -> _Operand:
  152. return self.apply("le", self, other)
  153. def __gt__(self, other: _Operand | float) -> _Operand:
  154. return self.apply("gt", self, other)
  155. def __ge__(self, other: _Operand | float) -> _Operand:
  156. return self.apply("ge", self, other)
  157. # conversions
  158. def imagemath_int(self: _Operand) -> _Operand:
  159. return _Operand(self.im.convert("I"))
  160. def imagemath_float(self: _Operand) -> _Operand:
  161. return _Operand(self.im.convert("F"))
  162. # logical
  163. def imagemath_equal(self: _Operand, other: _Operand | float | None) -> _Operand:
  164. return self.apply("eq", self, other, mode="I")
  165. def imagemath_notequal(self: _Operand, other: _Operand | float | None) -> _Operand:
  166. return self.apply("ne", self, other, mode="I")
  167. def imagemath_min(self: _Operand, other: _Operand | float | None) -> _Operand:
  168. return self.apply("min", self, other)
  169. def imagemath_max(self: _Operand, other: _Operand | float | None) -> _Operand:
  170. return self.apply("max", self, other)
  171. def imagemath_convert(self: _Operand, mode: str) -> _Operand:
  172. return _Operand(self.im.convert(mode))
  173. ops = {
  174. "int": imagemath_int,
  175. "float": imagemath_float,
  176. "equal": imagemath_equal,
  177. "notequal": imagemath_notequal,
  178. "min": imagemath_min,
  179. "max": imagemath_max,
  180. "convert": imagemath_convert,
  181. }
  182. def lambda_eval(expression: Callable[[dict[str, Any]], Any], **kw: Any) -> Any:
  183. """
  184. Returns the result of an image function.
  185. :py:mod:`~PIL.ImageMath` only supports single-layer images. To process multi-band
  186. images, use the :py:meth:`~PIL.Image.Image.split` method or
  187. :py:func:`~PIL.Image.merge` function.
  188. :param expression: A function that receives a dictionary.
  189. :param **kw: Values to add to the function's dictionary.
  190. :return: The expression result. This is usually an image object, but can
  191. also be an integer, a floating point value, or a pixel tuple,
  192. depending on the expression.
  193. """
  194. args: dict[str, Any] = ops.copy()
  195. args.update(kw)
  196. for k, v in args.items():
  197. if isinstance(v, Image.Image):
  198. args[k] = _Operand(v)
  199. out = expression(args)
  200. try:
  201. return out.im
  202. except AttributeError:
  203. return out
  204. def unsafe_eval(expression: str, **kw: Any) -> Any:
  205. """
  206. Evaluates an image expression. This uses Python's ``eval()`` function to process
  207. the expression string, and carries the security risks of doing so. It is not
  208. recommended to process expressions without considering this.
  209. :py:meth:`~lambda_eval` is a more secure alternative.
  210. :py:mod:`~PIL.ImageMath` only supports single-layer images. To process multi-band
  211. images, use the :py:meth:`~PIL.Image.Image.split` method or
  212. :py:func:`~PIL.Image.merge` function.
  213. :param expression: A string containing a Python-style expression.
  214. :param **kw: Values to add to the evaluation context.
  215. :return: The evaluated expression. This is usually an image object, but can
  216. also be an integer, a floating point value, or a pixel tuple,
  217. depending on the expression.
  218. """
  219. # build execution namespace
  220. args: dict[str, Any] = ops.copy()
  221. for k in kw:
  222. if "__" in k or hasattr(builtins, k):
  223. msg = f"'{k}' not allowed"
  224. raise ValueError(msg)
  225. args.update(kw)
  226. for k, v in args.items():
  227. if isinstance(v, Image.Image):
  228. args[k] = _Operand(v)
  229. compiled_code = compile(expression, "<string>", "eval")
  230. def scan(code: CodeType) -> None:
  231. for const in code.co_consts:
  232. if type(const) is type(compiled_code):
  233. scan(const)
  234. for name in code.co_names:
  235. if name not in args and name != "abs":
  236. msg = f"'{name}' not allowed"
  237. raise ValueError(msg)
  238. scan(compiled_code)
  239. out = builtins.eval(expression, {"__builtins": {"abs": abs}}, args)
  240. try:
  241. return out.im
  242. except AttributeError:
  243. return out