1"""
2Functions for interacting with images/videos.
3
4Dependencies:
5* ffprobe
6* https://github.com/alexwlchan/create_thumbnail
7* https://github.com/alexwlchan/dominant_colours
8* https://github.com/alexwlchan/get_live_text
9
10References:
11* https://alexwlchan.net/2021/dominant-colours/
12* https://alexwlchan.net/2025/detecting-av1-videos/
13* https://stackoverflow.com/a/58567453
14
15"""
16
17from fractions import Fraction
18import json
19from pathlib import Path
20import subprocess
21from typing import Literal, NotRequired, TypedDict, TYPE_CHECKING
22
23if TYPE_CHECKING:
24    import PIL.Image
25
26
27__all__ = [
28    "convert_srt_to_vtt",
29    "create_image_entity",
30    "create_video_entity",
31    "get_media_paths",
32    "get_tint_colour",
33    "is_av1_video",
34    "ImageEntity",
35    "MediaEntity",
36    "VideoEntity",
37]
38
39
40def is_av1_video(path: str | Path) -> bool:
41    """
42    Return True if a video is encoded with AV1, False otherwise.
43    """
44    # fmt: off
45    cmd = [
46        "ffprobe",
47        #
48        # Set the logging level
49        "-loglevel", "error",
50        #
51        # Select the first video stream
52        "-select_streams", "v:0",
53        #
54        # Print the codec_name (e.g. av1)
55        "-show_entries", "stream=codec_name",
56        #
57        # Print just the value
58        "-output_format", "default=noprint_wrappers=1:nokey=1",
59        #
60        # Name of the video to check
61        str(path),
62    ]
63    # fmt: on
64
65    output = subprocess.check_output(cmd, text=True)
66
67    return output.strip() == "av1"
68
69
70class ImageEntity(TypedDict):
71    """
72    ImageEntity contains all the fields I need to render an image
73    in a web page.
74    """
75
76    type: Literal["image"]
77
78    # The path to the image on disk
79    path: str
80
81    # The path to a low-resolution thumbnail
82    thumbnail_path: NotRequired[str]
83
84    # The display resolution of the image
85    width: int
86    height: int
87
88    # A hex-encoded colour which is prominent in this image.
89    tint_colour: str
90
91    # Whether the image is animated (GIF and WebP only)
92    is_animated: NotRequired[Literal[True]]
93
94    # Whether the image has transparent pixels
95    has_transparency: NotRequired[Literal[True]]
96
97    # The alt text of the image, if available
98    alt_text: NotRequired[str]
99
100    # The source URL of the image, if available
101    source_url: NotRequired[str]
102
103
104class SubtitlesEntity(TypedDict):
105    """
106    SubtitlesEntity describes a set of subtitles.
107    """
108
109    path: str
110    label: str
111
112
113class VideoEntity(TypedDict):
114    """
115    VideoEntity contains all the fields I need to render a video
116    in a web page.
117    """
118
119    type: Literal["video"]
120
121    # The path to the video on disk
122    path: str
123
124    # The poster image for the video
125    poster: ImageEntity
126
127    # The display resolution of the video
128    width: int
129    height: int
130
131    # The duration of the video, as an HOURS:MM:SS.MICROSECONDS string
132    duration: str
133
134    # Path to the subtitles for the video, if available
135    subtitles: NotRequired[list[SubtitlesEntity]]
136
137    # The source URL of the image, if available
138    source_url: NotRequired[str]
139
140    # Whether the video should play automatically. This is used for
141    # videos that are substituting for animated GIFs.
142    autoplay: NotRequired[Literal[True]]
143
144
145MediaEntity = ImageEntity | VideoEntity
146
147
148def get_media_paths(e: MediaEntity) -> set[Path]:
149    """
150    Return a list of all media paths represented by this media entity.
151    """
152    result: set[str | Path] = set()
153
154    try:
155        e["type"]
156    except KeyError:
157        raise TypeError(f"Entity does not have a type: {e}")
158
159    if e["type"] == "video":
160        result.add(e["path"])
161
162        try:
163            for se in e["subtitles"]:
164                result.add(se["path"])
165        except KeyError:
166            pass
167        for p in get_media_paths(e["poster"]):
168            result.add(p)
169    elif e["type"] == "image":
170        result.add(e["path"])
171        try:
172            result.add(e["thumbnail_path"])
173        except KeyError:
174            pass
175    else:
176        raise TypeError(f"Unrecognised entity type: {e['type']}")
177
178    return {Path(p) for p in result}
179
180
181class ThumbnailConfig(TypedDict):
182    out_dir: Path | str
183    width: NotRequired[int]
184    height: NotRequired[int]
185
186
187def create_image_entity(
188    path: str | Path,
189    *,
190    background: str = "#ffffff",
191    alt_text: str | None = None,
192    source_url: str | None = None,
193    thumbnail_config: ThumbnailConfig | None = None,
194    generate_transcript: bool = False,
195) -> ImageEntity:
196    """
197    Create an ImageEntity for a saved image.
198    """
199    from PIL import Image, ImageOps
200
201    with Image.open(path) as im:
202        # Account for EXIF orientation in the dimensions.
203        # See https://alexwlchan.net/til/2024/photos-can-have-orientation-in-exif/
204        transposed_im = ImageOps.exif_transpose(im)
205
206        entity: ImageEntity = {
207            "type": "image",
208            "path": str(path),
209            "tint_colour": get_tint_colour(path, background=background),
210            "width": transposed_im.width,
211            "height": transposed_im.height,
212        }
213
214        if _is_animated(im):
215            entity["is_animated"] = True
216
217        if _has_transparency(im):
218            entity["has_transparency"] = True
219
220    if thumbnail_config is not None:
221        entity["thumbnail_path"] = _create_thumbnail(path, thumbnail_config)
222
223    if alt_text is not None and generate_transcript:
224        raise TypeError("You cannot set alt_text and generate_transcript=True!")
225
226    elif alt_text is not None:
227        entity["alt_text"] = alt_text
228    elif generate_transcript:
229        transcript = _get_transcript(path)
230        if transcript is not None:
231            entity["alt_text"] = transcript
232
233    if source_url is not None:
234        entity["source_url"] = source_url
235
236    return entity
237
238
239def create_video_entity(
240    video_path: str | Path,
241    *,
242    poster_path: str | Path,
243    subtitles_path: str | Path | None = None,
244    source_url: str | None = None,
245    autoplay: bool = False,
246    thumbnail_config: ThumbnailConfig | None = None,
247    background: str = "#ffffff",
248) -> VideoEntity:
249    """
250    Create a video entity for files on disk.
251    """
252    width, height, duration = _get_video_data(video_path)
253    poster = create_image_entity(
254        poster_path, thumbnail_config=thumbnail_config, background=background
255    )
256
257    entity: VideoEntity = {
258        "type": "video",
259        "path": str(video_path),
260        "width": width,
261        "height": height,
262        "duration": duration,
263        "poster": poster,
264    }
265
266    if subtitles_path:
267        entity["subtitles"] = [{"path": str(subtitles_path), "label": "English"}]
268
269    if source_url:
270        entity["source_url"] = source_url
271
272    if autoplay:
273        entity["autoplay"] = autoplay
274
275    return entity
276
277
278def _is_animated(im: "PIL.Image.Image") -> bool:
279    """
280    Return True if an image is animated, False otherwise.
281    """
282    return getattr(im, "is_animated", False)
283
284
285def _has_transparency(im: "PIL.Image.Image") -> bool:
286    """
287    Return True if an image has transparent pixels, False otherwise.
288
289    By Vinyl Da.i'gyu-Kazotetsu on Stack Overflow:
290    https://stackoverflow.com/a/58567453
291    """
292    if im.info.get("transparency", None) is not None:
293        return True
294    if im.mode == "P":
295        transparent = im.info.get("transparency", -1)
296        for _, index in im.getcolors():  # type: ignore
297            # TODO: Find an image that hits this branch, so I can
298            # include it in the test suite.
299            if index == transparent:  # pragma: no cover
300                return True
301    elif im.mode == "RGBA":
302        extrema = im.getextrema()
303        if extrema[3][0] < 255:  # type: ignore
304            return True
305    return False
306
307
308def get_tint_colour(path: str | Path, *, background: str) -> str:
309    """
310    Get the tint colour for an image.
311    """
312    if background == "white":
313        background = "#ffffff"
314    elif background == "black":
315        background = "#000000"
316
317    cmd = [
318        "dominant_colours",
319        str(path),
320        f"--best-against-bg={background}",
321        "--no-palette",
322    ]
323
324    result = subprocess.check_output(cmd, text=True)
325    return result.strip()
326
327
328def _get_transcript(path: str | Path) -> str | None:
329    """
330    Get the transcript for an image (if any).
331    """
332    result = subprocess.check_output(["get_live_text", str(path)], text=True)
333
334    return result.strip() or None
335
336
337def _create_thumbnail(path: str | Path, thumbnail_config: ThumbnailConfig) -> str:
338    """
339    Create a thumbnail for an image and return the path.
340    """
341    cmd = ["create_thumbnail", str(path), "--out-dir", thumbnail_config["out_dir"]]
342
343    if "width" in thumbnail_config:
344        cmd.extend(["--width", str(thumbnail_config["width"])])
345
346    if "height" in thumbnail_config:
347        cmd.extend(["--height", str(thumbnail_config["height"])])
348
349    return subprocess.check_output(cmd, text=True)
350
351
352def _get_video_data(video_path: str | Path) -> tuple[int, int, str]:
353    """
354    Return the dimensions and duration of a video, as a width/height fraction.
355    """
356    cmd = [
357        "ffprobe",
358        #
359        # verbosity level = error
360        "-v",
361        "error",
362        #
363        # only get information about the first video stream
364        "-select_streams",
365        "v:0",
366        #
367        # only gather the entries I'm interested in
368        "-show_entries",
369        "stream=width,height,sample_aspect_ratio,duration",
370        #
371        # print the duration in HH:MM:SS.microseconds format
372        "-sexagesimal",
373        #
374        # print output in JSON, which is easier to parse
375        "-print_format",
376        "json",
377        #
378        # input file
379        str(video_path),
380    ]
381
382    output = subprocess.check_output(cmd)
383    ffprobe_resp = json.loads(output)
384
385    # The output will be structured something like:
386    #
387    #   {
388    #       "streams": [
389    #           {
390    #               "width": 1920,
391    #               "height": 1080,
392    #               "sample_aspect_ratio": "45:64"
393    #           }
394    #       ],
395    #       …
396    #   }
397    #
398    # If the video doesn't specify a pixel aspect ratio, then it won't
399    # have a `sample_aspect_ratio` key.
400    video_stream = ffprobe_resp["streams"][0]
401
402    try:
403        pixel_aspect_ratio = Fraction(
404            video_stream["sample_aspect_ratio"].replace(":", "/")
405        )
406    except KeyError:
407        pixel_aspect_ratio = Fraction(1)
408
409    width = round(video_stream["width"] * pixel_aspect_ratio)
410    height = video_stream["height"]
411    duration = video_stream["duration"]
412
413    return width, height, duration
414
415
416def convert_srt_to_vtt(srt_path: Path) -> Path:
417    """
418    Convert an SRT subtitles file to a VTT file.
419
420    SRT subtitles are commonly used for video platform uploads (like Vimeo
421    or YouTube), but I want VTT files for use in the HTML5 <video> tag.
422
423    This deletes the SRT file and returns the path to new VTT file.
424    """
425    if srt_path.suffix != ".srt":
426        raise ValueError(
427            f"convert_srt_to_vtt can only take .srt files, got {srt_path!r}"
428        )
429
430    with open(srt_path, encoding="utf-8-sig") as srt_file:
431        srt_lines = srt_file.readlines()
432
433    if len(srt_lines) == 0:
434        raise ValueError("cannot convert empty SRT file")
435
436    vtt_path = srt_path.with_suffix(".vtt")
437    assert vtt_path != srt_path, vtt_path
438
439    with open(vtt_path, "w", encoding="utf-8") as vtt_file:
440        # Write the required WEBVTT header and blank line
441        vtt_file.write("WEBVTT\n\n")
442
443        for line in srt_lines:
444            # Remove caption numbers; they're not required in VTT files
445            if line.strip().isdigit():
446                continue
447
448            # VTT timestamps uses commas rather than periods in timestamps
449            if "-->" in line:
450                line = line.replace(",", ".")
451
452            vtt_file.write(line)
453
454    srt_path.unlink()
455    return vtt_path