nodes/spacepxl-ComfyUI-Depth-Pro
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

initial code from apple/ml-depth-pro

Browse Source
This commit is contained in:
spacepxl
2024-10-03 20:57:02 -04:00
parent 01b522070c
commit 36c3be8b80
14 changed files with 1973 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
depth_pro/__init__.py
+5
View File
@@ -0,0 +1,5 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
"""Depth Pro package."""
from .depth_pro import create_model_and_transforms # noqa
from .utils import load_rgb # noqa
depth_pro/cli/__init__.py
+4
View File
@@ -0,0 +1,4 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
"""Depth Pro CLI and tools."""
from .run import main as run_main # noqa
depth_pro/cli/run.py
+149
View File
@@ -0,0 +1,149 @@
#!/usr/bin/env python3
"""Sample script to run DepthPro.
Copyright (C) 2024 Apple Inc. All Rights Reserved.
"""
import argparse
import logging
from pathlib import Path
import numpy as np
import PIL.Image
import torch
from matplotlib import pyplot as plt
from tqdm import tqdm
from depth_pro import create_model_and_transforms, load_rgb
LOGGER = logging.getLogger(__name__)
def get_torch_device() -> torch.device:
"""Get the Torch device."""
device = torch.device("cpu")
if torch.cuda.is_available():
device = torch.device("cuda:0")
elif torch.backends.mps.is_available():
device = torch.device("mps")
return device
def run(args):
"""Run Depth Pro on a sample image."""
if args.verbose:
logging.basicConfig(level=logging.INFO)
# Load model.
model, transform = create_model_and_transforms(
device=get_torch_device(),
precision=torch.half,
)
model.eval()
image_paths = [args.image_path]
if args.image_path.is_dir():
image_paths = args.image_path.glob("**/*")
relative_path = args.image_path
else:
relative_path = args.image_path.parent
if not args.skip_display:
plt.ion()
fig = plt.figure()
ax_rgb = fig.add_subplot(121)
ax_disp = fig.add_subplot(122)
for image_path in tqdm(image_paths):
# Load image and focal length from exif info (if found.).
try:
LOGGER.info(f"Loading image {image_path} ...")
image, _, f_px = load_rgb(image_path)
except Exception as e:
LOGGER.error(str(e))
continue
# Run prediction. If `f_px` is provided, it is used to estimate the final metric depth,
# otherwise the model estimates `f_px` to compute the depth metricness.
prediction = model.infer(transform(image), f_px=f_px)
# Extract the depth and focal length.
depth = prediction["depth"].detach().cpu().numpy().squeeze()
if f_px is not None:
LOGGER.debug(f"Focal length (from exif): {f_px:0.2f}")
elif prediction["focallength_px"] is not None:
focallength_px = prediction["focallength_px"].detach().cpu().item()
LOGGER.info(f"Estimated focal length: {focallength_px}")
# Save Depth as npz file.
if args.output_path is not None:
output_file = (
args.output_path
/ image_path.relative_to(relative_path).parent
/ image_path.stem
)
LOGGER.info(f"Saving depth map to: {str(output_file)}")
output_file.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output_file, depth=depth)
# Save as color-mapped "turbo" jpg image.
cmap = plt.get_cmap("turbo_r")
normalized_depth = (depth - depth.min()) / (
depth.max() - depth.min()
)
color_depth = (cmap(normalized_depth)[..., :3] * 255).astype(
np.uint8
)
color_map_output_file = str(output_file) + ".jpg"
LOGGER.info(f"Saving color-mapped depth to: : {color_map_output_file}")
PIL.Image.fromarray(color_depth).save(
color_map_output_file, format="JPEG", quality=90
)
# Display the image and estimated depth map.
if not args.skip_display:
ax_rgb.imshow(image)
ax_disp.imshow(depth, cmap="turbo_r")
fig.canvas.draw()
fig.canvas.flush_events()
LOGGER.info("Done predicting depth!")
if not args.skip_display:
plt.show(block=True)
def main():
"""Run DepthPro inference example."""
parser = argparse.ArgumentParser(
description="Inference scripts of DepthPro with PyTorch models."
)
parser.add_argument(
"-i",
"--image-path",
type=Path,
default="./data/example.jpg",
help="Path to input image.",
)
parser.add_argument(
"-o",
"--output-path",
type=Path,
help="Path to store output files.",
)
parser.add_argument(
"--skip-display",
action="store_true",
help="Skip matplotlib display.",
)
parser.add_argument(
"-v",
"--verbose",
action="store_true",
help="Show verbose output."
)
run(parser.parse_args())
if __name__ == "__main__":
main()
depth_pro/depth_pro.py
+298
View File
@@ -0,0 +1,298 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
# Depth Pro: Sharp Monocular Metric Depth in Less Than a Second
from __future__ import annotations
from dataclasses import dataclass
from typing import Mapping, Optional, Tuple, Union
import torch
from torch import nn
from torchvision.transforms import (
Compose,
ConvertImageDtype,
Lambda,
Normalize,
ToTensor,
)
from .network.decoder import MultiresConvDecoder
from .network.encoder import DepthProEncoder
from .network.fov import FOVNetwork
from .network.vit_factory import VIT_CONFIG_DICT, ViTPreset, create_vit
@dataclass
class DepthProConfig:
"""Configuration for DepthPro."""
patch_encoder_preset: ViTPreset
image_encoder_preset: ViTPreset
decoder_features: int
checkpoint_uri: Optional[str] = None
fov_encoder_preset: Optional[ViTPreset] = None
use_fov_head: bool = True
DEFAULT_MONODEPTH_CONFIG_DICT = DepthProConfig(
patch_encoder_preset="dinov2l16_384",
image_encoder_preset="dinov2l16_384",
checkpoint_uri="./checkpoints/depth_pro.pt",
decoder_features=256,
use_fov_head=True,
fov_encoder_preset="dinov2l16_384",
)
def create_backbone_model(
preset: ViTPreset
) -> Tuple[nn.Module, ViTPreset]:
"""Create and load a backbone model given a config.
Args:
----
preset: A backbone preset to load pre-defind configs.
Returns:
-------
A Torch module and the associated config.
"""
if preset in VIT_CONFIG_DICT:
config = VIT_CONFIG_DICT[preset]
model = create_vit(preset=preset, use_pretrained=False)
else:
raise KeyError(f"Preset {preset} not found.")
return model, config
def create_model_and_transforms(
config: DepthProConfig = DEFAULT_MONODEPTH_CONFIG_DICT,
device: torch.device = torch.device("cpu"),
precision: torch.dtype = torch.float32,
) -> Tuple[DepthPro, Compose]:
"""Create a DepthPro model and load weights from `config.checkpoint_uri`.
Args:
----
config: The configuration for the DPT model architecture.
device: The optional Torch device to load the model onto, default runs on "cpu".
precision: The optional precision used for the model, default is FP32.
Returns:
-------
The Torch DepthPro model and associated Transform.
"""
patch_encoder, patch_encoder_config = create_backbone_model(
preset=config.patch_encoder_preset
)
image_encoder, _ = create_backbone_model(
preset=config.image_encoder_preset
)
fov_encoder = None
if config.use_fov_head and config.fov_encoder_preset is not None:
fov_encoder, _ = create_backbone_model(preset=config.fov_encoder_preset)
dims_encoder = patch_encoder_config.encoder_feature_dims
hook_block_ids = patch_encoder_config.encoder_feature_layer_ids
encoder = DepthProEncoder(
dims_encoder=dims_encoder,
patch_encoder=patch_encoder,
image_encoder=image_encoder,
hook_block_ids=hook_block_ids,
decoder_features=config.decoder_features,
)
decoder = MultiresConvDecoder(
dims_encoder=[config.decoder_features] + list(encoder.dims_encoder),
dim_decoder=config.decoder_features,
)
model = DepthPro(
encoder=encoder,
decoder=decoder,
last_dims=(32, 1),
use_fov_head=config.use_fov_head,
fov_encoder=fov_encoder,
).to(device)
if precision == torch.half:
model.half()
transform = Compose(
[
ToTensor(),
Lambda(lambda x: x.to(device)),
Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
ConvertImageDtype(precision),
]
)
if config.checkpoint_uri is not None:
state_dict = torch.load(config.checkpoint_uri, map_location="cpu")
missing_keys, unexpected_keys = model.load_state_dict(
state_dict=state_dict, strict=True
)
if len(unexpected_keys) != 0:
raise KeyError(
f"Found unexpected keys when loading monodepth: {unexpected_keys}"
)
# fc_norm is only for the classification head,
# which we would not use. We only use the encoding.
missing_keys = [key for key in missing_keys if "fc_norm" not in key]
if len(missing_keys) != 0:
raise KeyError(f"Keys are missing when loading monodepth: {missing_keys}")
return model, transform
class DepthPro(nn.Module):
"""DepthPro network."""
def __init__(
self,
encoder: DepthProEncoder,
decoder: MultiresConvDecoder,
last_dims: tuple[int, int],
use_fov_head: bool = True,
fov_encoder: Optional[nn.Module] = None,
):
"""Initialize DepthPro.
Args:
----
encoder: The DepthProEncoder backbone.
decoder: The MultiresConvDecoder decoder.
last_dims: The dimension for the last convolution layers.
use_fov_head: Whether to use the field-of-view head.
fov_encoder: A separate encoder for the field of view.
"""
super().__init__()
self.encoder = encoder
self.decoder = decoder
dim_decoder = decoder.dim_decoder
self.head = nn.Sequential(
nn.Conv2d(
dim_decoder, dim_decoder // 2, kernel_size=3, stride=1, padding=1
),
nn.ConvTranspose2d(
in_channels=dim_decoder // 2,
out_channels=dim_decoder // 2,
kernel_size=2,
stride=2,
padding=0,
bias=True,
),
nn.Conv2d(
dim_decoder // 2,
last_dims[0],
kernel_size=3,
stride=1,
padding=1,
),
nn.ReLU(True),
nn.Conv2d(last_dims[0], last_dims[1], kernel_size=1, stride=1, padding=0),
nn.ReLU(),
)
# Set the final convoultion layer's bias to be 0.
self.head[4].bias.data.fill_(0)
# Set the FOV estimation head.
if use_fov_head:
self.fov = FOVNetwork(num_features=dim_decoder, fov_encoder=fov_encoder)
@property
def img_size(self) -> int:
"""Return the internal image size of the network."""
return self.encoder.img_size
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
"""Decode by projection and fusion of multi-resolution encodings.
Args:
----
x (torch.Tensor): Input image.
Returns:
-------
The canonical inverse depth map [m] and the optional estimated field of view [deg].
"""
_, _, H, W = x.shape
assert H == self.img_size and W == self.img_size
encodings = self.encoder(x)
features, features_0 = self.decoder(encodings)
canonical_inverse_depth = self.head(features)
fov_deg = None
if hasattr(self, "fov"):
fov_deg = self.fov.forward(x, features_0.detach())
return canonical_inverse_depth, fov_deg
@torch.no_grad()
def infer(
self,
x: torch.Tensor,
f_px: Optional[Union[float, torch.Tensor]] = None,
interpolation_mode="bilinear",
) -> Mapping[str, torch.Tensor]:
"""Infer depth and fov for a given image.
If the image is not at network resolution, it is resized to 1536x1536 and
the estimated depth is resized to the original image resolution.
Note: if the focal length is given, the estimated value is ignored and the provided
focal length is use to generate the metric depth values.
Args:
----
x (torch.Tensor): Input image
f_px (torch.Tensor): Optional focal length in pixels corresponding to `x`.
interpolation_mode (str): Interpolation function for downsampling/upsampling.
Returns:
-------
Tensor dictionary (torch.Tensor): depth [m], focallength [pixels].
"""
if len(x.shape) == 3:
x = x.unsqueeze(0)
_, _, H, W = x.shape
resize = H != self.img_size or W != self.img_size
if resize:
x = nn.functional.interpolate(
x,
size=(self.img_size, self.img_size),
mode=interpolation_mode,
align_corners=False,
)
canonical_inverse_depth, fov_deg = self.forward(x)
if f_px is None:
f_px = 0.5 * W / torch.tan(0.5 * torch.deg2rad(fov_deg.to(torch.float)))
inverse_depth = canonical_inverse_depth * (W / f_px)
f_px = f_px.squeeze()
if resize:
inverse_depth = nn.functional.interpolate(
inverse_depth, size=(H, W), mode=interpolation_mode, align_corners=False
)
depth = 1.0 / torch.clamp(inverse_depth, min=1e-4, max=1e4)
return {
"depth": depth.squeeze(),
"focallength_px": f_px,
}
depth_pro/eval/boundary_metrics.py
+332
View File
@@ -0,0 +1,332 @@
from typing import List, Tuple
import numpy as np
def connected_component(r: np.ndarray, c: np.ndarray) -> List[List[int]]:
"""Find connected components in the given row and column indices.
Args:
----
r (np.ndarray): Row indices.
c (np.ndarray): Column indices.
Yields:
------
List[int]: Indices of connected components.
"""
indices = [0]
for i in range(1, r.size):
if r[i] == r[indices[-1]] and c[i] == c[indices[-1]] + 1:
indices.append(i)
else:
yield indices
indices = [i]
yield indices
def nms_horizontal(ratio: np.ndarray, threshold: float) -> np.ndarray:
"""Apply Non-Maximum Suppression (NMS) horizontally on the given ratio matrix.
Args:
----
ratio (np.ndarray): Input ratio matrix.
threshold (float): Threshold for NMS.
Returns:
-------
np.ndarray: Binary mask after applying NMS.
"""
mask = np.zeros_like(ratio, dtype=bool)
r, c = np.nonzero(ratio > threshold)
if len(r) == 0:
return mask
for ids in connected_component(r, c):
values = [ratio[r[i], c[i]] for i in ids]
mi = np.argmax(values)
mask[r[ids[mi]], c[ids[mi]]] = True
return mask
def nms_vertical(ratio: np.ndarray, threshold: float) -> np.ndarray:
"""Apply Non-Maximum Suppression (NMS) vertically on the given ratio matrix.
Args:
----
ratio (np.ndarray): Input ratio matrix.
threshold (float): Threshold for NMS.
Returns:
-------
np.ndarray: Binary mask after applying NMS.
"""
return np.transpose(nms_horizontal(np.transpose(ratio), threshold))
def fgbg_depth(
d: np.ndarray, t: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Find foreground-background relations between neighboring pixels.
Args:
----
d (np.ndarray): Depth matrix.
t (float): Threshold for comparison.
Returns:
-------
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
left, top, right, and bottom foreground-background relations.
"""
right_is_big_enough = (d[..., :, 1:] / d[..., :, :-1]) > t
left_is_big_enough = (d[..., :, :-1] / d[..., :, 1:]) > t
bottom_is_big_enough = (d[..., 1:, :] / d[..., :-1, :]) > t
top_is_big_enough = (d[..., :-1, :] / d[..., 1:, :]) > t
return (
left_is_big_enough,
top_is_big_enough,
right_is_big_enough,
bottom_is_big_enough,
)
def fgbg_depth_thinned(
d: np.ndarray, t: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Find foreground-background relations between neighboring pixels with Non-Maximum Suppression.
Args:
----
d (np.ndarray): Depth matrix.
t (float): Threshold for NMS.
Returns:
-------
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
left, top, right, and bottom foreground-background relations with NMS applied.
"""
right_is_big_enough = nms_horizontal(d[..., :, 1:] / d[..., :, :-1], t)
left_is_big_enough = nms_horizontal(d[..., :, :-1] / d[..., :, 1:], t)
bottom_is_big_enough = nms_vertical(d[..., 1:, :] / d[..., :-1, :], t)
top_is_big_enough = nms_vertical(d[..., :-1, :] / d[..., 1:, :], t)
return (
left_is_big_enough,
top_is_big_enough,
right_is_big_enough,
bottom_is_big_enough,
)
def fgbg_binary_mask(
d: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Find foreground-background relations between neighboring pixels in binary masks.
Args:
----
d (np.ndarray): Binary depth matrix.
Returns:
-------
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Four matrices indicating
left, top, right, and bottom foreground-background relations in binary masks.
"""
assert d.dtype == bool
right_is_big_enough = d[..., :, 1:] & ~d[..., :, :-1]
left_is_big_enough = d[..., :, :-1] & ~d[..., :, 1:]
bottom_is_big_enough = d[..., 1:, :] & ~d[..., :-1, :]
top_is_big_enough = d[..., :-1, :] & ~d[..., 1:, :]
return (
left_is_big_enough,
top_is_big_enough,
right_is_big_enough,
bottom_is_big_enough,
)
def edge_recall_matting(pr: np.ndarray, gt: np.ndarray, t: float) -> float:
"""Calculate edge recall for image matting.
Args:
----
pr (np.ndarray): Predicted depth matrix.
gt (np.ndarray): Ground truth binary mask.
t (float): Threshold for NMS.
Returns:
-------
float: Edge recall value.
"""
assert gt.dtype == bool
ap, bp, cp, dp = fgbg_depth_thinned(pr, t)
ag, bg, cg, dg = fgbg_binary_mask(gt)
return 0.25 * (
np.count_nonzero(ap & ag) / max(np.count_nonzero(ag), 1)
+ np.count_nonzero(bp & bg) / max(np.count_nonzero(bg), 1)
+ np.count_nonzero(cp & cg) / max(np.count_nonzero(cg), 1)
+ np.count_nonzero(dp & dg) / max(np.count_nonzero(dg), 1)
)
def boundary_f1(
pr: np.ndarray,
gt: np.ndarray,
t: float,
return_p: bool = False,
return_r: bool = False,
) -> float:
"""Calculate Boundary F1 score.
Args:
----
pr (np.ndarray): Predicted depth matrix.
gt (np.ndarray): Ground truth depth matrix.
t (float): Threshold for comparison.
return_p (bool, optional): If True, return precision. Defaults to False.
return_r (bool, optional): If True, return recall. Defaults to False.
Returns:
-------
float: Boundary F1 score, or precision, or recall depending on the flags.
"""
ap, bp, cp, dp = fgbg_depth(pr, t)
ag, bg, cg, dg = fgbg_depth(gt, t)
r = 0.25 * (
np.count_nonzero(ap & ag) / max(np.count_nonzero(ag), 1)
+ np.count_nonzero(bp & bg) / max(np.count_nonzero(bg), 1)
+ np.count_nonzero(cp & cg) / max(np.count_nonzero(cg), 1)
+ np.count_nonzero(dp & dg) / max(np.count_nonzero(dg), 1)
)
p = 0.25 * (
np.count_nonzero(ap & ag) / max(np.count_nonzero(ap), 1)
+ np.count_nonzero(bp & bg) / max(np.count_nonzero(bp), 1)
+ np.count_nonzero(cp & cg) / max(np.count_nonzero(cp), 1)
+ np.count_nonzero(dp & dg) / max(np.count_nonzero(dp), 1)
)
if r + p == 0:
return 0.0
if return_p:
return p
if return_r:
return r
return 2 * (r * p) / (r + p)
def get_thresholds_and_weights(
t_min: float, t_max: float, N: int
) -> Tuple[np.ndarray, np.ndarray]:
"""Generate thresholds and weights for the given range.
Args:
----
t_min (float): Minimum threshold.
t_max (float): Maximum threshold.
N (int): Number of thresholds.
Returns:
-------
Tuple[np.ndarray, np.ndarray]: Array of thresholds and corresponding weights.
"""
thresholds = np.linspace(t_min, t_max, N)
weights = thresholds / thresholds.sum()
return thresholds, weights
def invert_depth(depth: np.ndarray, eps: float = 1e-6) -> np.ndarray:
"""Inverts a depth map with numerical stability.
Args:
----
depth (np.ndarray): Depth map to be inverted.
eps (float): Minimum value to avoid division by zero (default is 1e-6).
Returns:
-------
np.ndarray: Inverted depth map.
"""
inverse_depth = 1.0 / depth.clip(min=eps)
return inverse_depth
def SI_boundary_F1(
predicted_depth: np.ndarray,
target_depth: np.ndarray,
t_min: float = 1.05,
t_max: float = 1.25,
N: int = 10,
) -> float:
"""Calculate Scale-Invariant Boundary F1 Score for depth-based ground-truth.
Args:
----
predicted_depth (np.ndarray): Predicted depth matrix.
target_depth (np.ndarray): Ground truth depth matrix.
t_min (float, optional): Minimum threshold. Defaults to 1.05.
t_max (float, optional): Maximum threshold. Defaults to 1.25.
N (int, optional): Number of thresholds. Defaults to 10.
Returns:
-------
float: Scale-Invariant Boundary F1 Score.
"""
assert predicted_depth.ndim == target_depth.ndim == 2
thresholds, weights = get_thresholds_and_weights(t_min, t_max, N)
f1_scores = np.array(
[
boundary_f1(invert_depth(predicted_depth), invert_depth(target_depth), t)
for t in thresholds
]
)
return np.sum(f1_scores * weights)
def SI_boundary_Recall(
predicted_depth: np.ndarray,
target_mask: np.ndarray,
t_min: float = 1.05,
t_max: float = 1.25,
N: int = 10,
alpha_threshold: float = 0.1,
) -> float:
"""Calculate Scale-Invariant Boundary Recall Score for mask-based ground-truth.
Args:
----
predicted_depth (np.ndarray): Predicted depth matrix.
target_mask (np.ndarray): Ground truth binary mask.
t_min (float, optional): Minimum threshold. Defaults to 1.05.
t_max (float, optional): Maximum threshold. Defaults to 1.25.
N (int, optional): Number of thresholds. Defaults to 10.
alpha_threshold (float, optional): Threshold for alpha masking. Defaults to 0.1.
Returns:
-------
float: Scale-Invariant Boundary Recall Score.
"""
assert predicted_depth.ndim == target_mask.ndim == 2
thresholds, weights = get_thresholds_and_weights(t_min, t_max, N)
thresholded_target = target_mask > alpha_threshold
recall_scores = np.array(
[
edge_recall_matting(
invert_depth(predicted_depth), thresholded_target, t=float(t)
)
for t in thresholds
]
)
weighted_recall = np.sum(recall_scores * weights)
return weighted_recall
depth_pro/eval/dis5k_sample_list.txt
+200
View File
@@ -0,0 +1,200 @@
DIS5K/DIS-TE1/im/12#Graphics#4#TrafficSign#8245751856_821be14f86_o.jpg
DIS5K/DIS-TE1/im/13#Insect#4#Butterfly#16023994688_7ff8cdccb1_o.jpg
DIS5K/DIS-TE1/im/14#Kitchenware#4#Kitchenware#IMG_20210520_205538.jpg
DIS5K/DIS-TE1/im/14#Kitchenware#8#SweetStand#4848284981_fc90f54b50_o.jpg
DIS5K/DIS-TE1/im/17#Non-motor Vehicle#4#Cart#15012855035_d10b57014f_o.jpg
DIS5K/DIS-TE1/im/2#Aircraft#5#Kite#13104545564_5afceec9bd_o.jpg
DIS5K/DIS-TE1/im/20#Sports#10#Skateboarding#8472763540_bb2390e928_o.jpg
DIS5K/DIS-TE1/im/21#Tool#14#Sword#32473146960_dcc6b77848_o.jpg
DIS5K/DIS-TE1/im/21#Tool#15#Tapeline#9680492386_2d2020f282_o.jpg
DIS5K/DIS-TE1/im/21#Tool#4#Flag#507752845_ef852100f0_o.jpg
DIS5K/DIS-TE1/im/21#Tool#6#Key#11966089533_3becd78b44_o.jpg
DIS5K/DIS-TE1/im/21#Tool#8#Scale#31946428472_d28def471b_o.jpg
DIS5K/DIS-TE1/im/22#Weapon#4#Rifle#8472656430_3eb908b211_o.jpg
DIS5K/DIS-TE1/im/8#Electronics#3#Earphone#1177468301_641df8c267_o.jpg
DIS5K/DIS-TE1/im/8#Electronics#9#MusicPlayer#2235782872_7d47847bb4_o.jpg
DIS5K/DIS-TE2/im/11#Furniture#13#Ladder#3878434417_2ed740586e_o.jpg
DIS5K/DIS-TE2/im/13#Insect#1#Ant#27047700955_3b3a1271f8_o.jpg
DIS5K/DIS-TE2/im/13#Insect#11#Spider#5567179191_38d1f65589_o.jpg
DIS5K/DIS-TE2/im/13#Insect#8#Locust#5237933769_e6687c05e4_o.jpg
DIS5K/DIS-TE2/im/14#Kitchenware#2#DishRack#70838854_40cf689da7_o.jpg
DIS5K/DIS-TE2/im/14#Kitchenware#8#SweetStand#8467929412_fef7f4275d_o.jpg
DIS5K/DIS-TE2/im/16#Music Instrument#2#Harp#28058219806_28e05ff24a_o.jpg
DIS5K/DIS-TE2/im/17#Non-motor Vehicle#1#BabyCarriage#29794777180_2e1695a0cf_o.jpg
DIS5K/DIS-TE2/im/19#Ship#3#Sailboat#22442908623_5977e3becf_o.jpg
DIS5K/DIS-TE2/im/2#Aircraft#5#Kite#44654358051_1400e71cc4_o.jpg
DIS5K/DIS-TE2/im/21#Tool#11#Stand#IMG_20210520_205442.jpg
DIS5K/DIS-TE2/im/21#Tool#17#Tripod#9318977876_34615ec9a0_o.jpg
DIS5K/DIS-TE2/im/5#Artifact#3#Handcraft#50860882577_8482143b1b_o.jpg
DIS5K/DIS-TE2/im/8#Electronics#10#Robot#3093360210_fee54dc5c5_o.jpg
DIS5K/DIS-TE2/im/8#Electronics#6#Microphone#47411477652_6da66cbc10_o.jpg
DIS5K/DIS-TE3/im/14#Kitchenware#4#Kitchenware#2451122898_ef883175dd_o.jpg
DIS5K/DIS-TE3/im/15#Machine#4#SewingMachine#9311164128_97ba1d3947_o.jpg
DIS5K/DIS-TE3/im/16#Music Instrument#2#Harp#7670920550_59e992fd7b_o.jpg
DIS5K/DIS-TE3/im/17#Non-motor Vehicle#1#BabyCarriage#8389984877_1fddf8715c_o.jpg
DIS5K/DIS-TE3/im/17#Non-motor Vehicle#3#Carriage#5947122724_98e0fc3d1f_o.jpg
DIS5K/DIS-TE3/im/2#Aircraft#2#Balloon#2487168092_641505883f_o.jpg
DIS5K/DIS-TE3/im/2#Aircraft#4#Helicopter#8401177591_06c71c8df2_o.jpg
DIS5K/DIS-TE3/im/20#Sports#1#Archery#12520003103_faa43ea3e0_o.jpg
DIS5K/DIS-TE3/im/21#Tool#11#Stand#IMG_20210709_221507.jpg
DIS5K/DIS-TE3/im/21#Tool#2#Clip#5656649687_63d0c6696d_o.jpg
DIS5K/DIS-TE3/im/21#Tool#6#Key#12878459244_6387a140ea_o.jpg
DIS5K/DIS-TE3/im/3#Aquatic#1#Lobster#109214461_f52b4b6093_o.jpg
DIS5K/DIS-TE3/im/4#Architecture#19#Windmill#20195851863_2627117e0e_o.jpg
DIS5K/DIS-TE3/im/5#Artifact#2#Cage#5821476369_ea23927487_o.jpg
DIS5K/DIS-TE3/im/8#Electronics#7#MobileHolder#49732997896_7f53c290b5_o.jpg
DIS5K/DIS-TE4/im/13#Insect#6#Centipede#15302179708_a267850881_o.jpg
DIS5K/DIS-TE4/im/17#Non-motor Vehicle#11#Tricycle#5771069105_a3aef6f665_o.jpg
DIS5K/DIS-TE4/im/17#Non-motor Vehicle#2#Bicycle#4245936196_fdf812dcb7_o.jpg
DIS5K/DIS-TE4/im/17#Non-motor Vehicle#9#ShoppingCart#4674052920_a5b7a2b236_o.jpg
DIS5K/DIS-TE4/im/18#Plant#1#Bonsai#3539420884_ca8973e2c0_o.jpg
DIS5K/DIS-TE4/im/2#Aircraft#6#Parachute#33590416634_9d6f2325e7_o.jpg
DIS5K/DIS-TE4/im/20#Sports#1#Archery#46924476515_0be1caa684_o.jpg
DIS5K/DIS-TE4/im/20#Sports#8#Racket#19337607166_dd1985fb59_o.jpg
DIS5K/DIS-TE4/im/21#Tool#6#Key#3193329588_839b0c74ce_o.jpg
DIS5K/DIS-TE4/im/5#Artifact#2#Cage#5821886526_0573ba2d0d_o.jpg
DIS5K/DIS-TE4/im/5#Artifact#3#Handcraft#50105138282_3c1d02c968_o.jpg
DIS5K/DIS-TE4/im/8#Electronics#1#Antenna#4305034305_874f21a701_o.jpg
DIS5K/DIS-TR/im/1#Accessories#1#Bag#15554964549_3105e51b6f_o.jpg
DIS5K/DIS-TR/im/1#Accessories#1#Bag#41104261980_098a6c4a56_o.jpg
DIS5K/DIS-TR/im/1#Accessories#2#Clothes#2284764037_871b2e8ca4_o.jpg
DIS5K/DIS-TR/im/1#Accessories#3#Eyeglasses#1824643784_70d0134156_o.jpg
DIS5K/DIS-TR/im/1#Accessories#3#Eyeglasses#3590020230_37b09a29b3_o.jpg
DIS5K/DIS-TR/im/1#Accessories#3#Eyeglasses#4809652879_4da8a69f3b_o.jpg
DIS5K/DIS-TR/im/1#Accessories#3#Eyeglasses#792204934_f9b28f99b4_o.jpg
DIS5K/DIS-TR/im/1#Accessories#5#Jewelry#13909132974_c4750c5fb7_o.jpg
DIS5K/DIS-TR/im/1#Accessories#7#Shoe#2483391615_9199ece8d6_o.jpg
DIS5K/DIS-TR/im/1#Accessories#8#Watch#4343266960_f6633b029b_o.jpg
DIS5K/DIS-TR/im/10#Frame#2#BicycleFrame#17897573_42964dd104_o.jpg
DIS5K/DIS-TR/im/10#Frame#5#Rack#15898634812_64807069ff_o.jpg
DIS5K/DIS-TR/im/10#Frame#5#Rack#23928546819_c184cb0b60_o.jpg
DIS5K/DIS-TR/im/11#Furniture#19#Shower#6189119596_77bcfe80ee_o.jpg
DIS5K/DIS-TR/im/11#Furniture#2#Bench#3263647075_9306e280b5_o.jpg
DIS5K/DIS-TR/im/11#Furniture#5#CoatHanger#12774091054_cd5ff520ef_o.jpg
DIS5K/DIS-TR/im/11#Furniture#6#DentalChair#13878156865_d0439dcb32_o.jpg
DIS5K/DIS-TR/im/11#Furniture#9#Easel#5861024714_2070cd480c_o.jpg
DIS5K/DIS-TR/im/12#Graphics#4#TrafficSign#40621867334_f3c32ec189_o.jpg
DIS5K/DIS-TR/im/13#Insect#1#Ant#3295038190_db5dd0d4f4_o.jpg
DIS5K/DIS-TR/im/13#Insect#10#Mosquito#24341339_a88a1dad4c_o.jpg
DIS5K/DIS-TR/im/13#Insect#11#Spider#27171518270_63b78069ff_o.jpg
DIS5K/DIS-TR/im/13#Insect#11#Spider#49925050281_fa727c154e_o.jpg
DIS5K/DIS-TR/im/13#Insect#2#Beatle#279616486_2f1e64f591_o.jpg
DIS5K/DIS-TR/im/13#Insect#3#Bee#43892067695_82cf3e536b_o.jpg
DIS5K/DIS-TR/im/13#Insect#6#Centipede#20874281788_3e15c90a1c_o.jpg
DIS5K/DIS-TR/im/13#Insect#7#Dragonfly#14106671120_1b824d77e4_o.jpg
DIS5K/DIS-TR/im/13#Insect#8#Locust#21637491048_676ef7c9f7_o.jpg
DIS5K/DIS-TR/im/13#Insect#9#Mantis#1381120202_9dff6987b2_o.jpg
DIS5K/DIS-TR/im/14#Kitchenware#1#Cup#12812517473_327d6474b8_o.jpg
DIS5K/DIS-TR/im/14#Kitchenware#10#WineGlass#6402491641_389275d4d1_o.jpg
DIS5K/DIS-TR/im/14#Kitchenware#3#Hydrovalve#3129932040_8c05825004_o.jpg
DIS5K/DIS-TR/im/14#Kitchenware#4#Kitchenware#2881934780_87d5218ebb_o.jpg
DIS5K/DIS-TR/im/14#Kitchenware#4#Kitchenware#IMG_20210520_205527.jpg
DIS5K/DIS-TR/im/14#Kitchenware#6#Spoon#32989113501_b69eccf0df_o.jpg
DIS5K/DIS-TR/im/14#Kitchenware#8#SweetStand#2867322189_c56d1e0b87_o.jpg
DIS5K/DIS-TR/im/15#Machine#1#Gear#19217846720_f5f2807475_o.jpg
DIS5K/DIS-TR/im/15#Machine#2#Machine#1620160659_9571b7a7ab_o.jpg
DIS5K/DIS-TR/im/16#Music Instrument#2#Harp#6012801603_1a6e2c16a6_o.jpg
DIS5K/DIS-TR/im/16#Music Instrument#5#Trombone#8683292118_d223c17ccb_o.jpg
DIS5K/DIS-TR/im/16#Music Instrument#6#Trumpet#8393262740_b8c216142c_o.jpg
DIS5K/DIS-TR/im/16#Music Instrument#8#Violin#1511267391_40e4949d68_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#1#BabyCarriage#6989512997_38b3dbc88b_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#12#Wheel#14627183228_b2d68cf501_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#12#Wheel#2932226475_1b2403e549_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#12#Wheel#5420155648_86459905b8_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#2#Bicycle#IMG_20210513_134904.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#3#Carriage#3311962551_6f211b7bd6_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#4#Cart#2609732026_baf7fff3a1_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#5#Handcart#5821282211_201cefeaf2_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#7#Mower#5779003232_3bb3ae531a_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#9#ShoppingCart#10051622843_ace07e32b8_o.jpg
DIS5K/DIS-TR/im/17#Non-motor Vehicle#9#ShoppingCart#8075259294_f23e243849_o.jpg
DIS5K/DIS-TR/im/18#Plant#2#Tree#44800999741_e377e16dbb_o.jpg
DIS5K/DIS-TR/im/2#Aircraft#1#Airplane#2631761913_3ac67d0223_o.jpg
DIS5K/DIS-TR/im/2#Aircraft#1#Airplane#37707911566_e908a261b6_o.jpg
DIS5K/DIS-TR/im/2#Aircraft#3#HangGlider#2557220131_b8506920c5_o.jpg
DIS5K/DIS-TR/im/2#Aircraft#4#Helicopter#6215659280_5dbd9b4546_o.jpg
DIS5K/DIS-TR/im/2#Aircraft#6#Parachute#20185790493_e56fcaf8c6_o.jpg
DIS5K/DIS-TR/im/20#Sports#1#Archery#3871269982_ae4c59a7eb_o.jpg
DIS5K/DIS-TR/im/20#Sports#9#RockClimbing#9662433268_51299bc50e_o.jpg
DIS5K/DIS-TR/im/21#Tool#14#Sword#26258479365_2950d7fa37_o.jpg
DIS5K/DIS-TR/im/21#Tool#15#Tapeline#15505703447_e0fdeaa5a6_o.jpg
DIS5K/DIS-TR/im/21#Tool#4#Flag#26678602024_9b665742de_o.jpg
DIS5K/DIS-TR/im/21#Tool#4#Flag#5774823110_d603ce3cc8_o.jpg
DIS5K/DIS-TR/im/21#Tool#5#Hook#6867989814_dba18d673c_o.jpg
DIS5K/DIS-TR/im/22#Weapon#4#Rifle#4451713125_cd91719189_o.jpg
DIS5K/DIS-TR/im/3#Aquatic#2#Seadragon#4910944581_913139b238_o.jpg
DIS5K/DIS-TR/im/4#Architecture#12#Scaffold#3661448960_8aff24cc4d_o.jpg
DIS5K/DIS-TR/im/4#Architecture#13#Sculpture#6385318715_9a88d4eba7_o.jpg
DIS5K/DIS-TR/im/4#Architecture#17#Well#5011603479_75cf42808a_o.jpg
DIS5K/DIS-TR/im/5#Artifact#2#Cage#4892828841_7f1bc05682_o.jpg
DIS5K/DIS-TR/im/5#Artifact#3#Handcraft#15404211628_9e9ff2ce2e_o.jpg
DIS5K/DIS-TR/im/5#Artifact#3#Handcraft#3200169865_7c84cfcccf_o.jpg
DIS5K/DIS-TR/im/5#Artifact#3#Handcraft#5859295071_c217e7c22f_o.jpg
DIS5K/DIS-TR/im/6#Automobile#10#SteeringWheel#17200338026_f1e2122d8e_o.jpg
DIS5K/DIS-TR/im/6#Automobile#3#Car#3780893425_1a7d275e09_o.jpg
DIS5K/DIS-TR/im/6#Automobile#5#Crane#15282506502_1b1132a7c3_o.jpg
DIS5K/DIS-TR/im/7#Electrical#1#Cable#16767791875_8e6df41752_o.jpg
DIS5K/DIS-TR/im/7#Electrical#1#Cable#3291433361_38747324c4_o.jpg
DIS5K/DIS-TR/im/7#Electrical#1#Cable#4195104238_12a754c61a_o.jpg
DIS5K/DIS-TR/im/7#Electrical#1#Cable#49645415132_61e5664ecf_o.jpg
DIS5K/DIS-TR/im/7#Electrical#1#Cable#IMG_20210521_232406.jpg
DIS5K/DIS-TR/im/7#Electrical#10#UtilityPole#3298312021_92f431e3e9_o.jpg
DIS5K/DIS-TR/im/7#Electrical#10#UtilityPole#47950134773_fbfff63f4e_o.jpg
DIS5K/DIS-TR/im/7#Electrical#11#VacuumCleaner#5448403677_6a29e21881_o.jpg
DIS5K/DIS-TR/im/7#Electrical#2#CeilingLamp#611568868_680ed5d39f_o.jpg
DIS5K/DIS-TR/im/7#Electrical#3#Fan#3391683115_990525a693_o.jpg
DIS5K/DIS-TR/im/7#Electrical#6#StreetLamp#150049122_0692266618_o.jpg
DIS5K/DIS-TR/im/7#Electrical#9#TransmissionTower#31433908671_7e7e277dfe_o.jpg
DIS5K/DIS-TR/im/8#Electronics#1#Antenna#8727884873_e0622ee5c4_o.jpg
DIS5K/DIS-TR/im/8#Electronics#2#Camcorder#4172690390_7e5f280ace_o.jpg
DIS5K/DIS-TR/im/8#Electronics#3#Earphone#413984555_f290febdf5_o.jpg
DIS5K/DIS-TR/im/8#Electronics#5#Headset#30574225373_3717ed9fa4_o.jpg
DIS5K/DIS-TR/im/8#Electronics#6#Microphone#538006482_4aae4f5bd6_o.jpg
DIS5K/DIS-TR/im/8#Electronics#9#MusicPlayer#1306012480_2ea80d2afd_o.jpg
DIS5K/DIS-TR/im/9#Entertainment#1#GymEquipment#33071754135_8f3195cbd1_o.jpg
DIS5K/DIS-TR/im/9#Entertainment#2#KidsPlayground#2305807849_be53d724ea_o.jpg
DIS5K/DIS-TR/im/9#Entertainment#2#KidsPlayground#3862040422_5bbf903204_o.jpg
DIS5K/DIS-TR/im/9#Entertainment#3#OutdoorFitnessEquipment#10814507005_3dacaa28b3_o.jpg
DIS5K/DIS-TR/im/9#Entertainment#4#FerrisWheel#81640293_4b0ee62040_o.jpg
DIS5K/DIS-TR/im/9#Entertainment#5#Swing#49867339188_08073f4b76_o.jpg
DIS5K/DIS-VD/im/1#Accessories#1#Bag#6815402415_e01c1a41e6_o.jpg
DIS5K/DIS-VD/im/1#Accessories#5#Jewelry#2744070193_1486582e8d_o.jpg
DIS5K/DIS-VD/im/10#Frame#1#BasketballHoop#IMG_20210521_232650.jpg
DIS5K/DIS-VD/im/10#Frame#5#Rack#6156611713_49ebf12b1e_o.jpg
DIS5K/DIS-VD/im/11#Furniture#11#Handrail#3276641240_1b84b5af85_o.jpg
DIS5K/DIS-VD/im/11#Furniture#13#Ladder#33423266_5391cf47e9_o.jpg
DIS5K/DIS-VD/im/11#Furniture#17#Table#3725111755_4fc101e7ab_o.jpg
DIS5K/DIS-VD/im/11#Furniture#2#Bench#35556410400_7235b58070_o.jpg
DIS5K/DIS-VD/im/11#Furniture#4#Chair#3301769985_e49de6739f_o.jpg
DIS5K/DIS-VD/im/11#Furniture#6#DentalChair#23811071619_2a95c3a688_o.jpg
DIS5K/DIS-VD/im/11#Furniture#9#Easel#8322807354_df6d56542e_o.jpg
DIS5K/DIS-VD/im/13#Insect#10#Mosquito#12391674863_0cdf430d3f_o.jpg
DIS5K/DIS-VD/im/13#Insect#7#Dragonfly#14693028899_344ea118f2_o.jpg
DIS5K/DIS-VD/im/14#Kitchenware#10#WineGlass#4450148455_8f460f541a_o.jpg
DIS5K/DIS-VD/im/14#Kitchenware#3#Hydrovalve#IMG_20210520_203410.jpg
DIS5K/DIS-VD/im/15#Machine#3#PlowHarrow#34521712846_df4babb024_o.jpg
DIS5K/DIS-VD/im/16#Music Instrument#5#Trombone#6222242743_e7189405cd_o.jpg
DIS5K/DIS-VD/im/17#Non-motor Vehicle#12#Wheel#25677578797_ea47e1d9e8_o.jpg
DIS5K/DIS-VD/im/17#Non-motor Vehicle#2#Bicycle#5153474856_21560b081b_o.jpg
DIS5K/DIS-VD/im/17#Non-motor Vehicle#7#Mower#16992510572_8a6ff27398_o.jpg
DIS5K/DIS-VD/im/19#Ship#2#Canoe#40571458163_7faf8b73d9_o.jpg
DIS5K/DIS-VD/im/2#Aircraft#1#Airplane#4270588164_66a619e834_o.jpg
DIS5K/DIS-VD/im/2#Aircraft#4#Helicopter#86789665_650b94b2ee_o.jpg
DIS5K/DIS-VD/im/20#Sports#14#Wakesurfing#5589577652_5061c168d2_o.jpg
DIS5K/DIS-VD/im/21#Tool#10#Spade#37018312543_63b21b0784_o.jpg
DIS5K/DIS-VD/im/21#Tool#14#Sword#24789047250_42df9bf422_o.jpg
DIS5K/DIS-VD/im/21#Tool#18#Umbrella#IMG_20210513_140445.jpg
DIS5K/DIS-VD/im/21#Tool#6#Key#43939732715_5a6e28b518_o.jpg
DIS5K/DIS-VD/im/22#Weapon#1#Cannon#12758066705_90b54295e7_o.jpg
DIS5K/DIS-VD/im/22#Weapon#4#Rifle#8019368790_fb6dc469a7_o.jpg
DIS5K/DIS-VD/im/3#Aquatic#5#Shrimp#2582833427_7a99e7356e_o.jpg
DIS5K/DIS-VD/im/4#Architecture#12#Scaffold#1013402687_590750354e_o.jpg
DIS5K/DIS-VD/im/4#Architecture#13#Sculpture#17176841759_272a3ed6e3_o.jpg
DIS5K/DIS-VD/im/4#Architecture#14#Stair#15079108505_0d11281624_o.jpg
DIS5K/DIS-VD/im/4#Architecture#19#Windmill#2928111082_ceb3051c04_o.jpg
DIS5K/DIS-VD/im/4#Architecture#3#Crack#3551574032_17dd106d31_o.jpg
DIS5K/DIS-VD/im/4#Architecture#5#GasStation#4564307581_c3069bdc62_o.jpg
DIS5K/DIS-VD/im/4#Architecture#8#ObservationTower#2704526950_d4f0ddc807_o.jpg
DIS5K/DIS-VD/im/5#Artifact#3#Handcraft#10873642323_1bafce3aa5_o.jpg
DIS5K/DIS-VD/im/6#Automobile#11#Tractor#8594504006_0c2c557d85_o.jpg
DIS5K/DIS-VD/im/8#Electronics#3#Earphone#8106454803_1178d867cc_o.jpg
depth_pro/network/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
"""Depth Pro network blocks."""
depth_pro/network/decoder.py
+206
View File
@@ -0,0 +1,206 @@
"""Copyright (C) 2024 Apple Inc. All Rights Reserved.
Dense Prediction Transformer Decoder architecture.
Implements a variant of Vision Transformers for Dense Prediction, https://arxiv.org/abs/2103.13413
"""
from __future__ import annotations
from typing import Iterable
import torch
from torch import nn
class MultiresConvDecoder(nn.Module):
"""Decoder for multi-resolution encodings."""
def __init__(
self,
dims_encoder: Iterable[int],
dim_decoder: int,
):
"""Initialize multiresolution convolutional decoder.
Args:
----
dims_encoder: Expected dims at each level from the encoder.
dim_decoder: Dim of decoder features.
"""
super().__init__()
self.dims_encoder = list(dims_encoder)
self.dim_decoder = dim_decoder
self.dim_out = dim_decoder
num_encoders = len(self.dims_encoder)
# At the highest resolution, i.e. level 0, we apply projection w/ 1x1 convolution
# when the dimensions mismatch. Otherwise we do not do anything, which is
# the default behavior of monodepth.
conv0 = (
nn.Conv2d(self.dims_encoder[0], dim_decoder, kernel_size=1, bias=False)
if self.dims_encoder[0] != dim_decoder
else nn.Identity()
)
convs = [conv0]
for i in range(1, num_encoders):
convs.append(
nn.Conv2d(
self.dims_encoder[i],
dim_decoder,
kernel_size=3,
stride=1,
padding=1,
bias=False,
)
)
self.convs = nn.ModuleList(convs)
fusions = []
for i in range(num_encoders):
fusions.append(
FeatureFusionBlock2d(
num_features=dim_decoder,
deconv=(i != 0),
batch_norm=False,
)
)
self.fusions = nn.ModuleList(fusions)
def forward(self, encodings: torch.Tensor) -> torch.Tensor:
"""Decode the multi-resolution encodings."""
num_levels = len(encodings)
num_encoders = len(self.dims_encoder)
if num_levels != num_encoders:
raise ValueError(
f"Got encoder output levels={num_levels}, expected levels={num_encoders+1}."
)
# Project features of different encoder dims to the same decoder dim.
# Fuse features from the lowest resolution (num_levels-1)
# to the highest (0).
features = self.convs[-1](encodings[-1])
lowres_features = features
features = self.fusions[-1](features)
for i in range(num_levels - 2, -1, -1):
features_i = self.convs[i](encodings[i])
features = self.fusions[i](features, features_i)
return features, lowres_features
class ResidualBlock(nn.Module):
"""Generic implementation of residual blocks.
This implements a generic residual block from
He et al. - Identity Mappings in Deep Residual Networks (2016),
https://arxiv.org/abs/1603.05027
which can be further customized via factory functions.
"""
def __init__(self, residual: nn.Module, shortcut: nn.Module | None = None) -> None:
"""Initialize ResidualBlock."""
super().__init__()
self.residual = residual
self.shortcut = shortcut
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Apply residual block."""
delta_x = self.residual(x)
if self.shortcut is not None:
x = self.shortcut(x)
return x + delta_x
class FeatureFusionBlock2d(nn.Module):
"""Feature fusion for DPT."""
def __init__(
self,
num_features: int,
deconv: bool = False,
batch_norm: bool = False,
):
"""Initialize feature fusion block.
Args:
----
num_features: Input and output dimensions.
deconv: Whether to use deconv before the final output conv.
batch_norm: Whether to use batch normalization in resnet blocks.
"""
super().__init__()
self.resnet1 = self._residual_block(num_features, batch_norm)
self.resnet2 = self._residual_block(num_features, batch_norm)
self.use_deconv = deconv
if deconv:
self.deconv = nn.ConvTranspose2d(
in_channels=num_features,
out_channels=num_features,
kernel_size=2,
stride=2,
padding=0,
bias=False,
)
self.out_conv = nn.Conv2d(
num_features,
num_features,
kernel_size=1,
stride=1,
padding=0,
bias=True,
)
self.skip_add = nn.quantized.FloatFunctional()
def forward(self, x0: torch.Tensor, x1: torch.Tensor | None = None) -> torch.Tensor:
"""Process and fuse input features."""
x = x0
if x1 is not None:
res = self.resnet1(x1)
x = self.skip_add.add(x, res)
x = self.resnet2(x)
if self.use_deconv:
x = self.deconv(x)
x = self.out_conv(x)
return x
@staticmethod
def _residual_block(num_features: int, batch_norm: bool):
"""Create a residual block."""
def _create_block(dim: int, batch_norm: bool) -> list[nn.Module]:
layers = [
nn.ReLU(False),
nn.Conv2d(
num_features,
num_features,
kernel_size=3,
stride=1,
padding=1,
bias=not batch_norm,
),
]
if batch_norm:
layers.append(nn.BatchNorm2d(dim))
return layers
residual = nn.Sequential(
*_create_block(dim=num_features, batch_norm=batch_norm),
*_create_block(dim=num_features, batch_norm=batch_norm),
)
return ResidualBlock(residual)
depth_pro/network/encoder.py
+332
View File
@@ -0,0 +1,332 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
# DepthProEncoder combining patch and image encoders.
from __future__ import annotations
import math
from typing import Iterable, Optional
import torch
import torch.nn as nn
import torch.nn.functional as F
class DepthProEncoder(nn.Module):
"""DepthPro Encoder.
An encoder aimed at creating multi-resolution encodings from Vision Transformers.
"""
def __init__(
self,
dims_encoder: Iterable[int],
patch_encoder: nn.Module,
image_encoder: nn.Module,
hook_block_ids: Iterable[int],
decoder_features: int,
):
"""Initialize DepthProEncoder.
The framework
1. creates an image pyramid,
2. generates overlapping patches with a sliding window at each pyramid level,
3. creates batched encodings via vision transformer backbones,
4. produces multi-resolution encodings.
Args:
----
img_size: Backbone image resolution.
dims_encoder: Dimensions of the encoder at different layers.
patch_encoder: Backbone used for patches.
image_encoder: Backbone used for global image encoder.
hook_block_ids: Hooks to obtain intermediate features for the patch encoder model.
decoder_features: Number of feature output in the decoder.
"""
super().__init__()
self.dims_encoder = list(dims_encoder)
self.patch_encoder = patch_encoder
self.image_encoder = image_encoder
self.hook_block_ids = list(hook_block_ids)
patch_encoder_embed_dim = patch_encoder.embed_dim
image_encoder_embed_dim = image_encoder.embed_dim
self.out_size = int(
patch_encoder.patch_embed.img_size[0] // patch_encoder.patch_embed.patch_size[0]
)
def _create_project_upsample_block(
dim_in: int,
dim_out: int,
upsample_layers: int,
dim_int: Optional[int] = None,
) -> nn.Module:
if dim_int is None:
dim_int = dim_out
# Projection.
blocks = [
nn.Conv2d(
in_channels=dim_in,
out_channels=dim_int,
kernel_size=1,
stride=1,
padding=0,
bias=False,
)
]
# Upsampling.
blocks += [
nn.ConvTranspose2d(
in_channels=dim_int if i == 0 else dim_out,
out_channels=dim_out,
kernel_size=2,
stride=2,
padding=0,
bias=False,
)
for i in range(upsample_layers)
]
return nn.Sequential(*blocks)
self.upsample_latent0 = _create_project_upsample_block(
dim_in=patch_encoder_embed_dim,
dim_int=self.dims_encoder[0],
dim_out=decoder_features,
upsample_layers=3,
)
self.upsample_latent1 = _create_project_upsample_block(
dim_in=patch_encoder_embed_dim, dim_out=self.dims_encoder[0], upsample_layers=2
)
self.upsample0 = _create_project_upsample_block(
dim_in=patch_encoder_embed_dim, dim_out=self.dims_encoder[1], upsample_layers=1
)
self.upsample1 = _create_project_upsample_block(
dim_in=patch_encoder_embed_dim, dim_out=self.dims_encoder[2], upsample_layers=1
)
self.upsample2 = _create_project_upsample_block(
dim_in=patch_encoder_embed_dim, dim_out=self.dims_encoder[3], upsample_layers=1
)
self.upsample_lowres = nn.ConvTranspose2d(
in_channels=image_encoder_embed_dim,
out_channels=self.dims_encoder[3],
kernel_size=2,
stride=2,
padding=0,
bias=True,
)
self.fuse_lowres = nn.Conv2d(
in_channels=(self.dims_encoder[3] + self.dims_encoder[3]),
out_channels=self.dims_encoder[3],
kernel_size=1,
stride=1,
padding=0,
bias=True,
)
# Obtain intermediate outputs of the blocks.
self.patch_encoder.blocks[self.hook_block_ids[0]].register_forward_hook(
self._hook0
)
self.patch_encoder.blocks[self.hook_block_ids[1]].register_forward_hook(
self._hook1
)
def _hook0(self, model, input, output):
self.backbone_highres_hook0 = output
def _hook1(self, model, input, output):
self.backbone_highres_hook1 = output
@property
def img_size(self) -> int:
"""Return the full image size of the SPN network."""
return self.patch_encoder.patch_embed.img_size[0] * 4
def _create_pyramid(
self, x: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Create a 3-level image pyramid."""
# Original resolution: 1536 by default.
x0 = x
# Middle resolution: 768 by default.
x1 = F.interpolate(
x, size=None, scale_factor=0.5, mode="bilinear", align_corners=False
)
# Low resolution: 384 by default, corresponding to the backbone resolution.
x2 = F.interpolate(
x, size=None, scale_factor=0.25, mode="bilinear", align_corners=False
)
return x0, x1, x2
def split(self, x: torch.Tensor, overlap_ratio: float = 0.25) -> torch.Tensor:
"""Split the input into small patches with sliding window."""
patch_size = 384
patch_stride = int(patch_size * (1 - overlap_ratio))
image_size = x.shape[-1]
steps = int(math.ceil((image_size - patch_size) / patch_stride)) + 1
x_patch_list = []
for j in range(steps):
j0 = j * patch_stride
j1 = j0 + patch_size
for i in range(steps):
i0 = i * patch_stride
i1 = i0 + patch_size
x_patch_list.append(x[..., j0:j1, i0:i1])
return torch.cat(x_patch_list, dim=0)
def merge(self, x: torch.Tensor, batch_size: int, padding: int = 3) -> torch.Tensor:
"""Merge the patched input into a image with sliding window."""
steps = int(math.sqrt(x.shape[0] // batch_size))
idx = 0
output_list = []
for j in range(steps):
output_row_list = []
for i in range(steps):
output = x[batch_size * idx : batch_size * (idx + 1)]
if j != 0:
output = output[..., padding:, :]
if i != 0:
output = output[..., :, padding:]
if j != steps - 1:
output = output[..., :-padding, :]
if i != steps - 1:
output = output[..., :, :-padding]
output_row_list.append(output)
idx += 1
output_row = torch.cat(output_row_list, dim=-1)
output_list.append(output_row)
output = torch.cat(output_list, dim=-2)
return output
def reshape_feature(
self, embeddings: torch.Tensor, width, height, cls_token_offset=1
):
"""Discard class token and reshape 1D feature map to a 2D grid."""
b, hw, c = embeddings.shape
# Remove class token.
if cls_token_offset > 0:
embeddings = embeddings[:, cls_token_offset:, :]
# Shape: (batch, height, width, dim) -> (batch, dim, height, width)
embeddings = embeddings.reshape(b, height, width, c).permute(0, 3, 1, 2)
return embeddings
def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
"""Encode input at multiple resolutions.
Args:
----
x (torch.Tensor): Input image.
Returns:
-------
Multi resolution encoded features.
"""
batch_size = x.shape[0]
# Step 0: create a 3-level image pyramid.
x0, x1, x2 = self._create_pyramid(x)
# Step 1: split to create batched overlapped mini-images at the backbone (BeiT/ViT/Dino)
# resolution.
# 5x5 @ 384x384 at the highest resolution (1536x1536).
x0_patches = self.split(x0, overlap_ratio=0.25)
# 3x3 @ 384x384 at the middle resolution (768x768).
x1_patches = self.split(x1, overlap_ratio=0.5)
# 1x1 # 384x384 at the lowest resolution (384x384).
x2_patches = x2
# Concatenate all the sliding window patches and form a batch of size (35=5x5+3x3+1x1).
x_pyramid_patches = torch.cat(
(x0_patches, x1_patches, x2_patches),
dim=0,
)
# Step 2: Run the backbone (BeiT) model and get the result of large batch size.
x_pyramid_encodings = self.patch_encoder(x_pyramid_patches)
x_pyramid_encodings = self.reshape_feature(
x_pyramid_encodings, self.out_size, self.out_size
)
# Step 3: merging.
# Merge highres latent encoding.
x_latent0_encodings = self.reshape_feature(
self.backbone_highres_hook0,
self.out_size,
self.out_size,
)
x_latent0_features = self.merge(
x_latent0_encodings[: batch_size * 5 * 5], batch_size=batch_size, padding=3
)
x_latent1_encodings = self.reshape_feature(
self.backbone_highres_hook1,
self.out_size,
self.out_size,
)
x_latent1_features = self.merge(
x_latent1_encodings[: batch_size * 5 * 5], batch_size=batch_size, padding=3
)
# Split the 35 batch size from pyramid encoding back into 5x5+3x3+1x1.
x0_encodings, x1_encodings, x2_encodings = torch.split(
x_pyramid_encodings,
[len(x0_patches), len(x1_patches), len(x2_patches)],
dim=0,
)
# 96x96 feature maps by merging 5x5 @ 24x24 patches with overlaps.
x0_features = self.merge(x0_encodings, batch_size=batch_size, padding=3)
# 48x84 feature maps by merging 3x3 @ 24x24 patches with overlaps.
x1_features = self.merge(x1_encodings, batch_size=batch_size, padding=6)
# 24x24 feature maps.
x2_features = x2_encodings
# Apply the image encoder model.
x_global_features = self.image_encoder(x2_patches)
x_global_features = self.reshape_feature(
x_global_features, self.out_size, self.out_size
)
# Upsample feature maps.
x_latent0_features = self.upsample_latent0(x_latent0_features)
x_latent1_features = self.upsample_latent1(x_latent1_features)
x0_features = self.upsample0(x0_features)
x1_features = self.upsample1(x1_features)
x2_features = self.upsample2(x2_features)
x_global_features = self.upsample_lowres(x_global_features)
x_global_features = self.fuse_lowres(
torch.cat((x2_features, x_global_features), dim=1)
)
return [
x_latent0_features,
x_latent1_features,
x0_features,
x1_features,
x_global_features,
]
depth_pro/network/fov.py
+82
View File
@@ -0,0 +1,82 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
# Field of View network architecture.
from typing import Optional
import torch
from torch import nn
from torch.nn import functional as F
class FOVNetwork(nn.Module):
"""Field of View estimation network."""
def __init__(
self,
num_features: int,
fov_encoder: Optional[nn.Module] = None,
):
"""Initialize the Field of View estimation block.
Args:
----
num_features: Number of features used.
fov_encoder: Optional encoder to bring additional network capacity.
"""
super().__init__()
# Create FOV head.
fov_head0 = [
nn.Conv2d(
num_features, num_features // 2, kernel_size=3, stride=2, padding=1
), # 128 x 24 x 24
nn.ReLU(True),
]
fov_head = [
nn.Conv2d(
num_features // 2, num_features // 4, kernel_size=3, stride=2, padding=1
), # 64 x 12 x 12
nn.ReLU(True),
nn.Conv2d(
num_features // 4, num_features // 8, kernel_size=3, stride=2, padding=1
), # 32 x 6 x 6
nn.ReLU(True),
nn.Conv2d(num_features // 8, 1, kernel_size=6, stride=1, padding=0),
]
if fov_encoder is not None:
self.encoder = nn.Sequential(
fov_encoder, nn.Linear(fov_encoder.embed_dim, num_features // 2)
)
self.downsample = nn.Sequential(*fov_head0)
else:
fov_head = fov_head0 + fov_head
self.head = nn.Sequential(*fov_head)
def forward(self, x: torch.Tensor, lowres_feature: torch.Tensor) -> torch.Tensor:
"""Forward the fov network.
Args:
----
x (torch.Tensor): Input image.
lowres_feature (torch.Tensor): Low resolution feature.
Returns:
-------
The field of view tensor.
"""
if hasattr(self, "encoder"):
x = F.interpolate(
x,
size=None,
scale_factor=0.25,
mode="bilinear",
align_corners=False,
)
x = self.encoder(x)[:, 1:].permute(0, 2, 1)
lowres_feature = self.downsample(lowres_feature)
x = x.reshape_as(lowres_feature) + lowres_feature
else:
x = lowres_feature
return self.head(x)
depth_pro/network/vit.py
+123
View File
@@ -0,0 +1,123 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
try:
from timm.layers import resample_abs_pos_embed
except ImportError as err:
print("ImportError: {0}".format(err))
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
def make_vit_b16_backbone(
model,
encoder_feature_dims,
encoder_feature_layer_ids,
vit_features,
start_index=1,
use_grad_checkpointing=False,
) -> nn.Module:
"""Make a ViTb16 backbone for the DPT model."""
if use_grad_checkpointing:
model.set_grad_checkpointing()
vit_model = nn.Module()
vit_model.hooks = encoder_feature_layer_ids
vit_model.model = model
vit_model.features = encoder_feature_dims
vit_model.vit_features = vit_features
vit_model.model.start_index = start_index
vit_model.model.patch_size = vit_model.model.patch_embed.patch_size
vit_model.model.is_vit = True
vit_model.model.forward = vit_model.model.forward_features
return vit_model
def forward_features_eva_fixed(self, x):
"""Encode features."""
x = self.patch_embed(x)
x, rot_pos_embed = self._pos_embed(x)
for blk in self.blocks:
if self.grad_checkpointing:
x = checkpoint(blk, x, rot_pos_embed)
else:
x = blk(x, rot_pos_embed)
x = self.norm(x)
return x
def resize_vit(model: nn.Module, img_size) -> nn.Module:
"""Resample the ViT module to the given size."""
patch_size = model.patch_embed.patch_size
model.patch_embed.img_size = img_size
grid_size = tuple([s // p for s, p in zip(img_size, patch_size)])
model.patch_embed.grid_size = grid_size
pos_embed = resample_abs_pos_embed(
model.pos_embed,
grid_size, # img_size
num_prefix_tokens=(
0 if getattr(model, "no_embed_class", False) else model.num_prefix_tokens
),
)
model.pos_embed = torch.nn.Parameter(pos_embed)
return model
def resize_patch_embed(model: nn.Module, new_patch_size=(16, 16)) -> nn.Module:
"""Resample the ViT patch size to the given one."""
# interpolate patch embedding
if hasattr(model, "patch_embed"):
old_patch_size = model.patch_embed.patch_size
if (
new_patch_size[0] != old_patch_size[0]
or new_patch_size[1] != old_patch_size[1]
):
patch_embed_proj = model.patch_embed.proj.weight
patch_embed_proj_bias = model.patch_embed.proj.bias
use_bias = True if patch_embed_proj_bias is not None else False
_, _, h, w = patch_embed_proj.shape
new_patch_embed_proj = torch.nn.functional.interpolate(
patch_embed_proj,
size=[new_patch_size[0], new_patch_size[1]],
mode="bicubic",
align_corners=False,
)
new_patch_embed_proj = (
new_patch_embed_proj * (h / new_patch_size[0]) * (w / new_patch_size[1])
)
model.patch_embed.proj = nn.Conv2d(
in_channels=model.patch_embed.proj.in_channels,
out_channels=model.patch_embed.proj.out_channels,
kernel_size=new_patch_size,
stride=new_patch_size,
bias=use_bias,
)
if use_bias:
model.patch_embed.proj.bias = patch_embed_proj_bias
model.patch_embed.proj.weight = torch.nn.Parameter(new_patch_embed_proj)
model.patch_size = new_patch_size
model.patch_embed.patch_size = new_patch_size
model.patch_embed.img_size = (
int(
model.patch_embed.img_size[0]
* new_patch_size[0]
/ old_patch_size[0]
),
int(
model.patch_embed.img_size[1]
* new_patch_size[1]
/ old_patch_size[1]
),
)
return model
depth_pro/network/vit_factory.py
+124
View File
@@ -0,0 +1,124 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
# Factory functions to build and load ViT models.
from __future__ import annotations
import logging
import types
from dataclasses import dataclass
from typing import Dict, List, Literal, Optional
import timm
import torch
import torch.nn as nn
from .vit import (
forward_features_eva_fixed,
make_vit_b16_backbone,
resize_patch_embed,
resize_vit,
)
LOGGER = logging.getLogger(__name__)
ViTPreset = Literal[
"dinov2l16_384",
]
@dataclass
class ViTConfig:
"""Configuration for ViT."""
in_chans: int
embed_dim: int
img_size: int = 384
patch_size: int = 16
# In case we need to rescale the backbone when loading from timm.
timm_preset: Optional[str] = None
timm_img_size: int = 384
timm_patch_size: int = 16
# The following 2 parameters are only used by DPT. See dpt_factory.py.
encoder_feature_layer_ids: List[int] = None
"""The layers in the Beit/ViT used to constructs encoder features for DPT."""
encoder_feature_dims: List[int] = None
"""The dimension of features of encoder layers from Beit/ViT features for DPT."""
VIT_CONFIG_DICT: Dict[ViTPreset, ViTConfig] = {
"dinov2l16_384": ViTConfig(
in_chans=3,
embed_dim=1024,
encoder_feature_layer_ids=[5, 11, 17, 23],
encoder_feature_dims=[256, 512, 1024, 1024],
img_size=384,
patch_size=16,
timm_preset="vit_large_patch14_dinov2",
timm_img_size=518,
timm_patch_size=14,
),
}
def create_vit(
preset: ViTPreset,
use_pretrained: bool = False,
checkpoint_uri: str | None = None,
use_grad_checkpointing: bool = False,
) -> nn.Module:
"""Create and load a VIT backbone module.
Args:
----
preset: The VIT preset to load the pre-defined config.
use_pretrained: Load pretrained weights if True, default is False.
checkpoint_uri: Checkpoint to load the wights from.
use_grad_checkpointing: Use grandient checkpointing.
Returns:
-------
A Torch ViT backbone module.
"""
config = VIT_CONFIG_DICT[preset]
img_size = (config.img_size, config.img_size)
patch_size = (config.patch_size, config.patch_size)
if "eva02" in preset:
model = timm.create_model(config.timm_preset, pretrained=use_pretrained)
model.forward_features = types.MethodType(forward_features_eva_fixed, model)
else:
model = timm.create_model(
config.timm_preset, pretrained=use_pretrained, dynamic_img_size=True
)
model = make_vit_b16_backbone(
model,
encoder_feature_dims=config.encoder_feature_dims,
encoder_feature_layer_ids=config.encoder_feature_layer_ids,
vit_features=config.embed_dim,
use_grad_checkpointing=use_grad_checkpointing,
)
if config.patch_size != config.timm_patch_size:
model.model = resize_patch_embed(model.model, new_patch_size=patch_size)
if config.img_size != config.timm_img_size:
model.model = resize_vit(model.model, img_size=img_size)
if checkpoint_uri is not None:
state_dict = torch.load(checkpoint_uri, map_location="cpu")
missing_keys, unexpected_keys = model.load_state_dict(
state_dict=state_dict, strict=False
)
if len(unexpected_keys) != 0:
raise KeyError(f"Found unexpected keys when loading vit: {unexpected_keys}")
if len(missing_keys) != 0:
raise KeyError(f"Keys are missing when loading vit: {missing_keys}")
LOGGER.info(model)
return model.model
depth_pro/utils.py
+112
View File
@@ -0,0 +1,112 @@
# Copyright (C) 2024 Apple Inc. All Rights Reserved.
import logging
from pathlib import Path
from typing import Any, Dict, List, Tuple, Union
import numpy as np
import pillow_heif
from PIL import ExifTags, Image, TiffTags
from pillow_heif import register_heif_opener
register_heif_opener()
LOGGER = logging.getLogger(__name__)
def extract_exif(img_pil: Image) -> Dict[str, Any]:
"""Return exif information as a dictionary.
Args:
----
img_pil: A Pillow image.
Returns:
-------
A dictionary with extracted EXIF information.
"""
# Get full exif description from get_ifd(0x8769):
# cf https://pillow.readthedocs.io/en/stable/releasenotes/8.2.0.html#image-getexif-exif-and-gps-ifd
img_exif = img_pil.getexif().get_ifd(0x8769)
exif_dict = {ExifTags.TAGS[k]: v for k, v in img_exif.items() if k in ExifTags.TAGS}
tiff_tags = img_pil.getexif()
tiff_dict = {
TiffTags.TAGS_V2[k].name: v
for k, v in tiff_tags.items()
if k in TiffTags.TAGS_V2
}
return {**exif_dict, **tiff_dict}
def fpx_from_f35(width: float, height: float, f_mm: float = 50) -> float:
"""Convert a focal length given in mm (35mm film equivalent) to pixels."""
return f_mm * np.sqrt(width**2.0 + height**2.0) / np.sqrt(36**2 + 24**2)
def load_rgb(
path: Union[Path, str], auto_rotate: bool = True, remove_alpha: bool = True
) -> Tuple[np.ndarray, List[bytes], float]:
"""Load an RGB image.
Args:
----
path: The url to the image to load.
auto_rotate: Rotate the image based on the EXIF data, default is True.
remove_alpha: Remove the alpha channel, default is True.
Returns:
-------
img: The image loaded as a numpy array.
icc_profile: The color profile of the image.
f_px: The optional focal length in pixels, extracting from the exif data.
"""
LOGGER.debug(f"Loading image {path} ...")
path = Path(path)
if path.suffix.lower() in [".heic"]:
heif_file = pillow_heif.open_heif(path, convert_hdr_to_8bit=True)
img_pil = heif_file.to_pillow()
else:
img_pil = Image.open(path)
img_exif = extract_exif(img_pil)
icc_profile = img_pil.info.get("icc_profile", None)
# Rotate the image.
if auto_rotate:
exif_orientation = img_exif.get("Orientation", 1)
if exif_orientation == 3:
img_pil = img_pil.transpose(Image.ROTATE_180)
elif exif_orientation == 6:
img_pil = img_pil.transpose(Image.ROTATE_270)
elif exif_orientation == 8:
img_pil = img_pil.transpose(Image.ROTATE_90)
elif exif_orientation != 1:
LOGGER.warning(f"Ignoring image orientation {exif_orientation}.")
img = np.array(img_pil)
# Convert to RGB if single channel.
if img.ndim < 3 or img.shape[2] == 1:
img = np.dstack((img, img, img))
if remove_alpha:
img = img[:, :, :3]
LOGGER.debug(f"\tHxW: {img.shape[0]}x{img.shape[1]}")
# Extract the focal length from exif data.
f_35mm = img_exif.get(
"FocalLengthIn35mmFilm",
img_exif.get(
"FocalLenIn35mmFilm", img_exif.get("FocalLengthIn35mmFormat", None)
),
)
if f_35mm is not None and f_35mm > 0:
LOGGER.debug(f"\tfocal length @ 35mm film: {f_35mm}mm")
f_px = fpx_from_f35(img.shape[1], img.shape[0], f_35mm)
else:
f_px = None
return img, icc_profile, f_px
requirements.txt
+4
View File
@@ -0,0 +1,4 @@
timm
numpy<2
pillow_heif
matplotlib