1
0

QoiImagePlugin.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. #
  2. # The Python Imaging Library.
  3. #
  4. # QOI support for PIL
  5. #
  6. # See the README file for information on usage and redistribution.
  7. #
  8. from __future__ import annotations
  9. import os
  10. from . import Image, ImageFile
  11. from ._binary import i32be as i32
  12. def _accept(prefix: bytes) -> bool:
  13. return prefix[:4] == b"qoif"
  14. class QoiImageFile(ImageFile.ImageFile):
  15. format = "QOI"
  16. format_description = "Quite OK Image"
  17. def _open(self) -> None:
  18. if not _accept(self.fp.read(4)):
  19. msg = "not a QOI file"
  20. raise SyntaxError(msg)
  21. self._size = i32(self.fp.read(4)), i32(self.fp.read(4))
  22. channels = self.fp.read(1)[0]
  23. self._mode = "RGB" if channels == 3 else "RGBA"
  24. self.fp.seek(1, os.SEEK_CUR) # colorspace
  25. self.tile = [ImageFile._Tile("qoi", (0, 0) + self._size, self.fp.tell())]
  26. class QoiDecoder(ImageFile.PyDecoder):
  27. _pulls_fd = True
  28. _previous_pixel: bytes | bytearray | None = None
  29. _previously_seen_pixels: dict[int, bytes | bytearray] = {}
  30. def _add_to_previous_pixels(self, value: bytes | bytearray) -> None:
  31. self._previous_pixel = value
  32. r, g, b, a = value
  33. hash_value = (r * 3 + g * 5 + b * 7 + a * 11) % 64
  34. self._previously_seen_pixels[hash_value] = value
  35. def decode(self, buffer: bytes | Image.SupportsArrayInterface) -> tuple[int, int]:
  36. assert self.fd is not None
  37. self._previously_seen_pixels = {}
  38. self._add_to_previous_pixels(bytearray((0, 0, 0, 255)))
  39. data = bytearray()
  40. bands = Image.getmodebands(self.mode)
  41. dest_length = self.state.xsize * self.state.ysize * bands
  42. while len(data) < dest_length:
  43. byte = self.fd.read(1)[0]
  44. value: bytes | bytearray
  45. if byte == 0b11111110 and self._previous_pixel: # QOI_OP_RGB
  46. value = bytearray(self.fd.read(3)) + self._previous_pixel[3:]
  47. elif byte == 0b11111111: # QOI_OP_RGBA
  48. value = self.fd.read(4)
  49. else:
  50. op = byte >> 6
  51. if op == 0: # QOI_OP_INDEX
  52. op_index = byte & 0b00111111
  53. value = self._previously_seen_pixels.get(
  54. op_index, bytearray((0, 0, 0, 0))
  55. )
  56. elif op == 1 and self._previous_pixel: # QOI_OP_DIFF
  57. value = bytearray(
  58. (
  59. (self._previous_pixel[0] + ((byte & 0b00110000) >> 4) - 2)
  60. % 256,
  61. (self._previous_pixel[1] + ((byte & 0b00001100) >> 2) - 2)
  62. % 256,
  63. (self._previous_pixel[2] + (byte & 0b00000011) - 2) % 256,
  64. self._previous_pixel[3],
  65. )
  66. )
  67. elif op == 2 and self._previous_pixel: # QOI_OP_LUMA
  68. second_byte = self.fd.read(1)[0]
  69. diff_green = (byte & 0b00111111) - 32
  70. diff_red = ((second_byte & 0b11110000) >> 4) - 8
  71. diff_blue = (second_byte & 0b00001111) - 8
  72. value = bytearray(
  73. tuple(
  74. (self._previous_pixel[i] + diff_green + diff) % 256
  75. for i, diff in enumerate((diff_red, 0, diff_blue))
  76. )
  77. )
  78. value += self._previous_pixel[3:]
  79. elif op == 3 and self._previous_pixel: # QOI_OP_RUN
  80. run_length = (byte & 0b00111111) + 1
  81. value = self._previous_pixel
  82. if bands == 3:
  83. value = value[:3]
  84. data += value * run_length
  85. continue
  86. self._add_to_previous_pixels(value)
  87. if bands == 3:
  88. value = value[:3]
  89. data += value
  90. self.set_as_raw(data)
  91. return -1, 0
  92. Image.register_open(QoiImageFile.format, QoiImageFile, _accept)
  93. Image.register_decoder("qoi", QoiDecoder)
  94. Image.register_extension(QoiImageFile.format, ".qoi")