selmee/depth-pro

1# Copyright (C) 2024 Apple Inc. All Rights Reserved.
2
3import logging
4from pathlib import Path
5from typing import Any, Dict, List, Tuple, Union
6
7import numpy as np
8import pillow_heif
9from PIL import ExifTags, Image, TiffTags
10from pillow_heif import register_heif_opener
11
12register_heif_opener()
13LOGGER = logging.getLogger(__name__)
14
15
16def extract_exif(img_pil: Image) -> Dict[str, Any]:
17    """Return exif information as a dictionary.
18
19    Args:
20    ----
21        img_pil: A Pillow image.
22
23    Returns:
24    -------
25        A dictionary with extracted EXIF information.
26
27    """
28    # Get full exif description from get_ifd(0x8769):
29    # cf https://pillow.readthedocs.io/en/stable/releasenotes/8.2.0.html#image-getexif-exif-and-gps-ifd
30    img_exif = img_pil.getexif().get_ifd(0x8769)
31    exif_dict = {ExifTags.TAGS[k]: v for k, v in img_exif.items() if k in ExifTags.TAGS}
32
33    tiff_tags = img_pil.getexif()
34    tiff_dict = {
35        TiffTags.TAGS_V2[k].name: v
36        for k, v in tiff_tags.items()
37        if k in TiffTags.TAGS_V2
38    }
39    return {**exif_dict, **tiff_dict}
40
41
42def fpx_from_f35(width: float, height: float, f_mm: float = 50) -> float:
43    """Convert a focal length given in mm (35mm film equivalent) to pixels."""
44    return f_mm * np.sqrt(width**2.0 + height**2.0) / np.sqrt(36**2 + 24**2)
45
46
47def load_rgb(
48    path: Union[Path, str], auto_rotate: bool = True, remove_alpha: bool = True
49) -> Tuple[np.ndarray, List[bytes], float]:
50    """Load an RGB image.
51
52    Args:
53    ----
54        path: The url to the image to load.
55        auto_rotate: Rotate the image based on the EXIF data, default is True.
56        remove_alpha: Remove the alpha channel, default is True.
57
58    Returns:
59    -------
60        img: The image loaded as a numpy array.
61        icc_profile: The color profile of the image.
62        f_px: The optional focal length in pixels, extracting from the exif data.
63
64    """
65    LOGGER.debug(f"Loading image {path} ...")
66
67    path = Path(path)
68    if path.suffix.lower() in [".heic"]:
69        heif_file = pillow_heif.open_heif(path, convert_hdr_to_8bit=True)
70        img_pil = heif_file.to_pillow()
71    else:
72        img_pil = Image.open(path)
73
74    img_exif = extract_exif(img_pil)
75    icc_profile = img_pil.info.get("icc_profile", None)
76
77    # Rotate the image.
78    if auto_rotate:
79        exif_orientation = img_exif.get("Orientation", 1)
80        if exif_orientation == 3:
81            img_pil = img_pil.transpose(Image.ROTATE_180)
82        elif exif_orientation == 6:
83            img_pil = img_pil.transpose(Image.ROTATE_270)
84        elif exif_orientation == 8:
85            img_pil = img_pil.transpose(Image.ROTATE_90)
86        elif exif_orientation != 1:
87            LOGGER.warning(f"Ignoring image orientation {exif_orientation}.")
88
89    img = np.array(img_pil)
90    # Convert to RGB if single channel.
91    if img.ndim < 3 or img.shape[2] == 1:
92        img = np.dstack((img, img, img))
93
94    if remove_alpha:
95        img = img[:, :, :3]
96
97    LOGGER.debug(f"\tHxW: {img.shape[0]}x{img.shape[1]}")
98
99    # Extract the focal length from exif data.
100    f_35mm = img_exif.get(
101        "FocalLengthIn35mmFilm",
102        img_exif.get(
103            "FocalLenIn35mmFilm", img_exif.get("FocalLengthIn35mmFormat", None)
104        ),
105    )
106    if f_35mm is not None and f_35mm > 0:
107        LOGGER.debug(f"\tfocal length @ 35mm film: {f_35mm}mm")
108        f_px = fpx_from_f35(img.shape[1], img.shape[0], f_35mm)
109    else:
110        f_px = None
111
112    return img, icc_profile, f_px
113