1#!/usr/bin/env python3
2"""
3A script that converts images to an sRGB colour profile.
4
5This is particularly useful for screenshots on macOS, which are taken
6with the display's colour profile (e.g. Display LCD or Display P3), but
7which I want to convert to sRGB for converting on the web.
8
9Note: this is a potentially destructive script.  Don't run this on images
10you care about if you don't have a backup!
11
12*   It overwrites the original image file.
13*   It strips out EXIF metadata.
14
15Based on https://github.com/python-pillow/Pillow/issues/1662
16"""
17
18import io
19import sys
20import typing
21
22from PIL import Image, ImageCms, ImageOps
23from PIL.ImageCms import PyCMSError
24
25
26def get_profile_description(icc_profile):
27    f = io.BytesIO(icc_profile)
28    prf = ImageCms.ImageCmsProfile(f)
29    return prf.profile.profile_description
30
31
32def convert_image_to_srgb(im: Image) -> typing.Union[Image, None]:
33    """
34    Convert an image to sRGB and return a new Image instance.
35    """
36    icc_profile = im.info.get("icc_profile")
37
38    # If this image doesn't have a colour profile, we're done.
39    if icc_profile is None:
40        return None
41
42    # If the image has an EXIF orientation tag, it will be stripped out
43    # upon saving (like all EXIF metadata).
44    #
45    # To avoid any weird rotation issues, bake the rotation into the image.
46    # See https://github.com/python-pillow/Pillow/issues/4703#issuecomment-645219973
47    # or the associated test.
48    #
49    # See https://github.com/alexwlchan/scripts/issues/21
50    im = ImageOps.exif_transpose(im)
51
52    # Otherwise, convert the image to an sRGB colour profile and return that.
53    try:
54        return ImageCms.profileToProfile(
55            im,
56            inputProfile=io.BytesIO(icc_profile),
57            outputProfile=ImageCms.createProfile("sRGB"),
58        )
59    except PyCMSError as err:
60        profile_name = get_profile_description(icc_profile)
61
62        is_gray_xyz = im.mode == "L" and b"GRAYXYZ" in icc_profile
63        is_cmyk = im.mode == "CMYK" and profile_name == "U.S. Web Coated (SWOP) v2"
64
65        if err.args[0].args == ("cannot build transform",) and (is_gray_xyz or is_cmyk):
66            return ImageCms.profileToProfile(
67                im,
68                inputProfile=io.BytesIO(icc_profile),
69                outputProfile=ImageCms.createProfile("sRGB"),
70                outputMode="RGB",
71            )
72        else:
73            raise
74
75    return out_im
76
77
78if __name__ == "__main__":
79    if len(sys.argv) == 1:
80        sys.exit(f"Usage: {__file__} <PATH...>")
81
82    for path in sys.argv[1:]:
83        im = Image.open(path)
84        out_im = convert_image_to_srgb(im)
85
86        if out_im is not None:
87            out_im.save(path)
88
89        print(path)