codekingpro/portable-devtools

1from __future__ import annotations
2
3import os
4from io import BytesIO
5from typing import IO
6
7from . import ExifTags, Image, ImageFile
8
9try:
10    from . import _avif
11
12    SUPPORTED = True
13except ImportError:
14    SUPPORTED = False
15
16# Decoder options as module globals, until there is a way to pass parameters
17# to Image.open (see https://github.com/python-pillow/Pillow/issues/569)
18DECODE_CODEC_CHOICE = "auto"
19DEFAULT_MAX_THREADS = 0
20
21
22def get_codec_version(codec_name: str) -> str | None:
23    versions = _avif.codec_versions()
24    for version in versions.split(", "):
25        if version.split(" [")[0] == codec_name:
26            return version.split(":")[-1].split(" ")[0]
27    return None
28
29
30def _accept(prefix: bytes) -> bool | str:
31    if prefix[4:8] != b"ftyp":
32        return False
33    major_brand = prefix[8:12]
34    if major_brand in (
35        # coding brands
36        b"avif",
37        b"avis",
38        # We accept files with AVIF container brands; we can't yet know if
39        # the ftyp box has the correct compatible brands, but if it doesn't
40        # then the plugin will raise a SyntaxError which Pillow will catch
41        # before moving on to the next plugin that accepts the file.
42        #
43        # Also, because this file might not actually be an AVIF file, we
44        # don't raise an error if AVIF support isn't properly compiled.
45        b"mif1",
46        b"msf1",
47    ):
48        if not SUPPORTED:
49            return (
50                "image file could not be identified because AVIF support not installed"
51            )
52        return True
53    return False
54
55
56def _get_default_max_threads() -> int:
57    if DEFAULT_MAX_THREADS:
58        return DEFAULT_MAX_THREADS
59    if hasattr(os, "sched_getaffinity"):
60        return len(os.sched_getaffinity(0))
61    else:
62        return os.cpu_count() or 1
63
64
65class AvifImageFile(ImageFile.ImageFile):
66    format = "AVIF"
67    format_description = "AVIF image"
68    __frame = -1
69
70    def _open(self) -> None:
71        if not SUPPORTED:
72            msg = "image file could not be opened because AVIF support not installed"
73            raise SyntaxError(msg)
74
75        if DECODE_CODEC_CHOICE != "auto" and not _avif.decoder_codec_available(
76            DECODE_CODEC_CHOICE
77        ):
78            msg = "Invalid opening codec"
79            raise ValueError(msg)
80
81        assert self.fp is not None
82        self._decoder = _avif.AvifDecoder(
83            self.fp.read(),
84            DECODE_CODEC_CHOICE,
85            _get_default_max_threads(),
86        )
87
88        # Get info from decoder
89        self._size, self.n_frames, self._mode, icc, exif, exif_orientation, xmp = (
90            self._decoder.get_info()
91        )
92        self.is_animated = self.n_frames > 1
93
94        if icc:
95            self.info["icc_profile"] = icc
96        if xmp:
97            self.info["xmp"] = xmp
98
99        if exif_orientation != 1 or exif:
100            exif_data = Image.Exif()
101            if exif:
102                exif_data.load(exif)
103                original_orientation = exif_data.get(ExifTags.Base.Orientation, 1)
104            else:
105                original_orientation = 1
106            if exif_orientation != original_orientation:
107                exif_data[ExifTags.Base.Orientation] = exif_orientation
108                exif = exif_data.tobytes()
109        if exif:
110            self.info["exif"] = exif
111        self.seek(0)
112
113    def seek(self, frame: int) -> None:
114        if not self._seek_check(frame):
115            return
116
117        # Set tile
118        self.__frame = frame
119        self.tile = [ImageFile._Tile("raw", (0, 0) + self.size, 0, self.mode)]
120
121    def load(self) -> Image.core.PixelAccess | None:
122        if self.tile:
123            # We need to load the image data for this frame
124            data, timescale, pts_in_timescales, duration_in_timescales = (
125                self._decoder.get_frame(self.__frame)
126            )
127            self.info["timestamp"] = round(1000 * (pts_in_timescales / timescale))
128            self.info["duration"] = round(1000 * (duration_in_timescales / timescale))
129
130            if self.fp and self._exclusive_fp:
131                self.fp.close()
132            self.fp = BytesIO(data)
133
134        return super().load()
135
136    def load_seek(self, pos: int) -> None:
137        pass
138
139    def tell(self) -> int:
140        return self.__frame
141
142
143def _save_all(im: Image.Image, fp: IO[bytes], filename: str | bytes) -> None:
144    _save(im, fp, filename, save_all=True)
145
146
147def _save(
148    im: Image.Image, fp: IO[bytes], filename: str | bytes, save_all: bool = False
149) -> None:
150    info = im.encoderinfo.copy()
151    if save_all:
152        append_images = list(info.get("append_images", []))
153    else:
154        append_images = []
155
156    total = 0
157    for ims in [im] + append_images:
158        total += getattr(ims, "n_frames", 1)
159
160    quality = info.get("quality", 75)
161    if not isinstance(quality, int) or quality < 0 or quality > 100:
162        msg = "Invalid quality setting"
163        raise ValueError(msg)
164
165    duration = info.get("duration", 0)
166    subsampling = info.get("subsampling", "4:2:0")
167    speed = info.get("speed", 6)
168    max_threads = info.get("max_threads", _get_default_max_threads())
169    codec = info.get("codec", "auto")
170    if codec != "auto" and not _avif.encoder_codec_available(codec):
171        msg = "Invalid saving codec"
172        raise ValueError(msg)
173    range_ = info.get("range", "full")
174    tile_rows_log2 = info.get("tile_rows", 0)
175    tile_cols_log2 = info.get("tile_cols", 0)
176    alpha_premultiplied = bool(info.get("alpha_premultiplied", False))
177    autotiling = bool(info.get("autotiling", tile_rows_log2 == tile_cols_log2 == 0))
178
179    icc_profile = info.get("icc_profile", im.info.get("icc_profile"))
180    exif_orientation = 1
181    if exif := info.get("exif"):
182        if isinstance(exif, Image.Exif):
183            exif_data = exif
184        else:
185            exif_data = Image.Exif()
186            exif_data.load(exif)
187        if ExifTags.Base.Orientation in exif_data:
188            exif_orientation = exif_data.pop(ExifTags.Base.Orientation)
189            exif = exif_data.tobytes() if exif_data else b""
190        elif isinstance(exif, Image.Exif):
191            exif = exif_data.tobytes()
192
193    xmp = info.get("xmp")
194
195    if isinstance(xmp, str):
196        xmp = xmp.encode("utf-8")
197
198    advanced = info.get("advanced")
199    if advanced is not None:
200        if isinstance(advanced, dict):
201            advanced = advanced.items()
202        try:
203            advanced = tuple(advanced)
204        except TypeError:
205            invalid = True
206        else:
207            invalid = any(not isinstance(v, tuple) or len(v) != 2 for v in advanced)
208        if invalid:
209            msg = (
210                "advanced codec options must be a dict of key-value string "
211                "pairs or a series of key-value two-tuples"
212            )
213            raise ValueError(msg)
214
215    # Setup the AVIF encoder
216    enc = _avif.AvifEncoder(
217        im.size,
218        subsampling,
219        quality,
220        speed,
221        max_threads,
222        codec,
223        range_,
224        tile_rows_log2,
225        tile_cols_log2,
226        alpha_premultiplied,
227        autotiling,
228        icc_profile or b"",
229        exif or b"",
230        exif_orientation,
231        xmp or b"",
232        advanced,
233    )
234
235    # Add each frame
236    frame_idx = 0
237    frame_duration = 0
238    cur_idx = im.tell()
239    is_single_frame = total == 1
240    try:
241        for ims in [im] + append_images:
242            # Get number of frames in this image
243            nfr = getattr(ims, "n_frames", 1)
244
245            for idx in range(nfr):
246                ims.seek(idx)
247
248                # Make sure image mode is supported
249                frame = ims
250                rawmode = ims.mode
251                if ims.mode not in {"RGB", "RGBA"}:
252                    rawmode = "RGBA" if ims.has_transparency_data else "RGB"
253                    frame = ims.convert(rawmode)
254
255                # Update frame duration
256                if isinstance(duration, (list, tuple)):
257                    frame_duration = duration[frame_idx]
258                else:
259                    frame_duration = duration
260
261                # Append the frame to the animation encoder
262                enc.add(
263                    frame.tobytes("raw", rawmode),
264                    frame_duration,
265                    frame.size,
266                    rawmode,
267                    is_single_frame,
268                )
269
270                # Update frame index
271                frame_idx += 1
272
273                if not save_all:
274                    break
275
276    finally:
277        im.seek(cur_idx)
278
279    # Get the final output from the encoder
280    data = enc.finish()
281    if data is None:
282        msg = "cannot write file as AVIF (encoder returned None)"
283        raise OSError(msg)
284
285    fp.write(data)
286
287
288Image.register_open(AvifImageFile.format, AvifImageFile, _accept)
289if SUPPORTED:
290    Image.register_save(AvifImageFile.format, _save)
291    Image.register_save_all(AvifImageFile.format, _save_all)
292    Image.register_extensions(AvifImageFile.format, [".avif", ".avifs"])
293    Image.register_mime(AvifImageFile.format, "image/avif")
294