Skip to content

DefocusLens

Thin-lens / circle-of-confusion model for fast depth-of-field and bokeh simulation, without full ray tracing. Useful when you need an inexpensive, differentiable defocus approximation.

deeplens.DefocusLens

DefocusLens(
    foclen,
    fnum,
    sensor_size=(8.0, 8.0),
    sensor_res=(2000, 2000),
    device=None,
    dtype=torch.float32,
)

Bases: Lens

Defocus lens that pre-computes the circle-of-confusion (CoC) PSF.

Rather than ray transfer (ABCD) matrices or thin-lens ray tracing, this model derives the circle of confusion from the focal length, F-number and focus distance, builds the corresponding PSF, and applies it directly. It simulates defocus blur (depth of field) but not higher-order optical aberrations. Useful as a fast baseline renderer, as commonly used in Blender and similar tools.

Attributes:

Name Type Description
foclen float

Focal length [mm].

fnum float

F-number.

foc_dist float

Current focus distance [mm], set by refocus (conventionally negative).

sensor_size tuple

Physical sensor size (W, H) [mm].

sensor_res tuple

Pixel resolution (W, H).

pixel_size float

Pixel pitch [mm].

Initialize a defocus lens.

A defocus lens models geometric defocus via the circle of confusion, which is wavelength-independent, so it takes no wavelength or default object-depth arguments (unlike the other lens classes).

Parameters:

Name Type Description Default
foclen float

Focal length in [mm].

required
fnum float

F-number.

required
sensor_size tuple

Physical sensor size as (W, H) in [mm]. Defaults to (8.0, 8.0).

(8.0, 8.0)
sensor_res tuple

Sensor resolution as (W, H) in pixels. Defaults to (2000, 2000).

(2000, 2000)
device str

Computation device. Defaults to None (auto-select GPU if available, else CPU).

None
dtype dtype

Data type for computations. Defaults to torch.float32.

float32
Source code in deeplens-src/deeplens/defocuslens.py
def __init__(
    self,
    foclen,
    fnum,
    sensor_size=(8.0, 8.0),
    sensor_res=(2000, 2000),
    device=None,
    dtype=torch.float32,
):
    """Initialize a defocus lens.

    A defocus lens models geometric defocus via the circle of confusion,
    which is wavelength-independent, so it takes no wavelength or default
    object-depth arguments (unlike the other lens classes).

    Args:
        foclen (float): Focal length in [mm].
        fnum (float): F-number.
        sensor_size (tuple, optional): Physical sensor size as (W, H) in [mm]. Defaults to (8.0, 8.0).
        sensor_res (tuple, optional): Sensor resolution as (W, H) in pixels. Defaults to (2000, 2000).
        device (str, optional): Computation device. Defaults to None
            (auto-select GPU if available, else CPU).
        dtype (torch.dtype, optional): Data type for computations. Defaults to torch.float32.
    """
    super(DefocusLens, self).__init__(
        device=device,
        dtype=dtype,
    )

    # Lens parameters
    self.foclen = foclen  # Focal length [mm]
    self.fnum = fnum

    # Configure sensor (sets sensor_size, sensor_res, pixel_size, r_sensor).
    self.set_sensor(sensor_size, sensor_res)
    self.astype(self.dtype)

    self.d_far = -20000.0
    self.d_close = -200.0
    self.refocus(foc_dist=-20000)

refocus

refocus(foc_dist)

Refocus the lens to a given object distance.

Parameters:

Name Type Description Default
foc_dist float

Focus distance in [mm], conventionally negative (object in front of the lens). Must be less than the focal length.

required

Raises:

Type Description
AssertionError

If foc_dist is not less than self.foclen.

Source code in deeplens-src/deeplens/defocuslens.py
def refocus(self, foc_dist):
    """Refocus the lens to a given object distance.

    Args:
        foc_dist (float): Focus distance in [mm], conventionally negative
            (object in front of the lens). Must be less than the focal length.

    Raises:
        AssertionError: If `foc_dist` is not less than `self.foclen`.
    """
    assert foc_dist < self.foclen, "Focus distance is too close."
    self.foc_dist = foc_dist

psf

psf(points, wvln=None, ks=PSF_KS, **kwargs)

Compute the defocus PSF as a circular disk of diameter CoC.

The PSF is a 2D blur disk whose diameter is the circle of confusion at each object depth, masked to a circle and normalized to sum to 1. With psf_type="gaussian" the disk is filled with a Gaussian falloff; with psf_type="pillbox" it is a flat-top (uniform) disk. The CoC model is wavelength-independent, so wvln is accepted for API uniformity with the other lens types but ignored.

Parameters:

Name Type Description Default
points Tensor

Object point positions in [mm], shape [N, 3] or [3]; the depth is taken from the z (third) coordinate.

required
wvln float or None

Wavelength in [µm]. Ignored (the CoC model is achromatic). Defaults to None.

None
ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS
**kwargs

Model-specific options. psf_type (str): "gaussian" (default) or "pillbox".

{}

Returns:

Name Type Description
psf Tensor

Normalized PSF kernel(s), shape [ks, ks] for a single point or [N, ks, ks] for N points.

Source code in deeplens-src/deeplens/defocuslens.py
def psf(self, points, wvln=None, ks=PSF_KS, **kwargs):
    """Compute the defocus PSF as a circular disk of diameter CoC.

    The PSF is a 2D blur disk whose diameter is the circle of confusion at
    each object depth, masked to a circle and normalized to sum to 1. With
    `psf_type="gaussian"` the disk is filled with a Gaussian falloff; with
    `psf_type="pillbox"` it is a flat-top (uniform) disk. The CoC model is
    wavelength-independent, so `wvln` is accepted for API uniformity with
    the other lens types but ignored.

    Args:
        points (torch.Tensor): Object point positions in [mm], shape [N, 3]
            or [3]; the depth is taken from the z (third) coordinate.
        wvln (float or None, optional): Wavelength in [µm]. Ignored (the CoC
            model is achromatic). Defaults to None.
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        **kwargs: Model-specific options. `psf_type` (str): "gaussian"
            (default) or "pillbox".

    Returns:
        psf (torch.Tensor): Normalized PSF kernel(s), shape [ks, ks] for a
            single point or [N, ks, ks] for N points.
    """
    psf_type = kwargs.get("psf_type", "gaussian")
    points = points.to(self.device)

    # Handle single point vs multiple points
    if len(points.shape) == 1:
        points = points.unsqueeze(0)
        single_point = True
    else:
        single_point = False

    # Calculate circle of confusion for each point
    depths = points[:, 2]  # Shape [N]
    coc_values = self.coc(depths)  # Shape [N]

    # Convert CoC from mm to pixels and add minimum value for numerical stability
    coc_pixel = torch.clamp(
        coc_values / self.pixel_size, min=0.5
    )  # Shape [N], minimum 0.5 pixels
    coc_pixel = coc_pixel.unsqueeze(-1).unsqueeze(
        -1
    )  # Shape [N, 1, 1], broadcasts with [ks, ks]
    coc_pixel_radius = coc_pixel / 2

    # Create an integer-centered coordinate grid. In particular, an even
    # kernel (the project default is 64) must still contain an explicit
    # origin sample. The previous half-pixel grid left an exact-focus CoC
    # mask empty, yielding a zero PSF instead of a peak-preserving delta.
    coords = torch.arange(ks, device=self.device, dtype=coc_pixel.dtype) - ks // 2
    x, y = torch.meshgrid(
        coords,
        coords,
        indexing="xy",
    )
    distance_sq = x**2 + y**2

    # Create PSF
    if psf_type == "gaussian":
        # Gaussian PSF
        psf = torch.exp(-distance_sq / (2 * coc_pixel_radius**2)) / (
            2 * np.pi * coc_pixel_radius**2
        )
    elif psf_type == "pillbox":
        # Pillbox PSF
        psf = torch.ones_like(x)
    else:
        raise ValueError(f"Invalid PSF type: {psf_type}")

    # Apply circular mask
    psf_mask = distance_sq < coc_pixel_radius**2
    psf = psf * psf_mask

    # Normalize PSF to sum to 1
    psf = psf / (psf.sum(dim=(-1, -2), keepdim=True) + EPSILON)

    if single_point:
        psf = psf.squeeze(0)

    return psf

coc

coc(depth)

Compute the circle-of-confusion (CoC) diameter from the paraxial defocus relation.

Depth is clamped to [self.d_far, self.d_close] and taken as an absolute distance before evaluating the CoC.

Parameters:

Name Type Description Default
depth Tensor

Object depth in [mm], shape [B] (or scalar). Conventionally negative (object in front of the lens).

required

Returns:

Name Type Description
coc Tensor

Circle-of-confusion diameter in [mm], same shape as depth.

Reference

[1] https://en.wikipedia.org/wiki/Circle_of_confusion

Source code in deeplens-src/deeplens/defocuslens.py
def coc(self, depth):
    """Compute the circle-of-confusion (CoC) diameter from the paraxial defocus relation.

    Depth is clamped to `[self.d_far, self.d_close]` and taken as an absolute
    distance before evaluating the CoC.

    Args:
        depth (torch.Tensor): Object depth in [mm], shape [B] (or scalar).
            Conventionally negative (object in front of the lens).

    Returns:
        coc (torch.Tensor): Circle-of-confusion diameter in [mm], same shape
            as `depth`.

    Reference:
        [1] https://en.wikipedia.org/wiki/Circle_of_confusion
    """
    depth = torch.as_tensor(depth, device=self.device)
    foc_dist = torch.tensor(
        self.foc_dist, device=self.device, dtype=depth.dtype
    ).abs()
    foclen = self.foclen
    fnum = self.fnum

    depth = torch.clamp(depth, self.d_far, self.d_close)
    depth = torch.abs(depth)

    # Calculate circle of confusion diameter, [mm]
    part1 = torch.abs(depth - foc_dist) / depth
    part2 = foclen**2 / (fnum * (foc_dist - foclen))
    coc = part1 * part2

    return coc

dof

dof(depth)

Compute the depth of field (DoF) at a given object depth.

Parameters:

Name Type Description Default
depth Tensor

Object depth in [mm], shape [B] (or scalar). Conventionally negative (object in front of the lens).

required

Returns:

Name Type Description
dof Tensor

Depth of field in [mm], same shape as depth.

Reference

[1] https://en.wikipedia.org/wiki/Depth_of_field

Source code in deeplens-src/deeplens/defocuslens.py
def dof(self, depth):
    """Compute the depth of field (DoF) at a given object depth.

    Args:
        depth (torch.Tensor): Object depth in [mm], shape [B] (or scalar).
            Conventionally negative (object in front of the lens).

    Returns:
        dof (torch.Tensor): Depth of field in [mm], same shape as `depth`.

    Reference:
        [1] https://en.wikipedia.org/wiki/Depth_of_field
    """
    depth = torch.as_tensor(depth, device=self.device)
    depth = torch.clamp(depth, self.d_far, self.d_close)
    depth_abs = torch.abs(depth)

    foclen = self.foclen
    fnum = self.fnum

    # Magnification factor
    m = foclen / (depth_abs - foclen)

    # CoC, [mm]
    coc = self.coc(depth)

    # Depth of field, [mm]
    part1 = 2 * fnum * coc * (m + 1)
    part2 = m**2 - (fnum * coc / foclen) ** 2
    dof = part1 / part2

    return dof

psf_rgb

psf_rgb(points, ks=PSF_KS, **kwargs)

Compute RGB PSF by replicating the monochrome PSF across three channels.

The defocus model is achromatic, so all channels share the same PSF.

Parameters:

Name Type Description Default
points Tensor

Object point positions in [mm], shape [N, 3].

required
ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS
**kwargs

Forwarded to psf.

{}

Returns:

Name Type Description
psf_rgb Tensor

RGB PSFs, shape [N, 3, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py
def psf_rgb(self, points, ks=PSF_KS, **kwargs):
    """Compute RGB PSF by replicating the monochrome PSF across three channels.

    The defocus model is achromatic, so all channels share the same PSF.

    Args:
        points (torch.Tensor): Object point positions in [mm], shape [N, 3].
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        **kwargs: Forwarded to `psf`.

    Returns:
        psf_rgb (torch.Tensor): RGB PSFs, shape [N, 3, ks, ks].
    """
    psf = self.psf(points, ks=ks, psf_type="gaussian", **kwargs)
    return psf.unsqueeze(1).repeat(1, 3, 1, 1)

psf_map

psf_map(grid=(5, 5), ks=PSF_KS, depth=None, **kwargs)

Compute a spatially-uniform monochrome PSF map.

Because the defocus model has no spatially-varying aberrations, every grid position receives the same on-axis PSF.

Parameters:

Name Type Description Default
grid tuple

Grid dimensions (rows, cols). Defaults to (5, 5).

(5, 5)
ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS
depth float or None

Object depth in [mm]. When None (default), falls back to self.obj_depth.

None
**kwargs

Forwarded to psf.

{}

Returns:

Name Type Description
psf_map Tensor

PSF map, shape [rows, cols, 1, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py
def psf_map(self, grid=(5, 5), ks=PSF_KS, depth=None, **kwargs):
    """Compute a spatially-uniform monochrome PSF map.

    Because the defocus model has no spatially-varying aberrations, every
    grid position receives the same on-axis PSF.

    Args:
        grid (tuple, optional): Grid dimensions (rows, cols).
            Defaults to (5, 5).
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        depth (float or None, optional): Object depth in [mm]. When None
            (default), falls back to `self.obj_depth`.
        **kwargs: Forwarded to `psf`.

    Returns:
        psf_map (torch.Tensor): PSF map, shape [rows, cols, 1, ks, ks].
    """
    depth = self.obj_depth if depth is None else depth
    points = torch.tensor([[0, 0, depth]], device=self.device)
    psf = self.psf(points=points, ks=ks, psf_type="gaussian", **kwargs)
    psf_map = psf.unsqueeze(0).unsqueeze(0).repeat(grid[0], grid[1], 1, 1, 1)
    return psf_map

psf_dp

psf_dp(points, ks=PSF_KS)

Generate left/right dual-pixel PSFs by masking the base PSF.

Takes the base defocus PSF and splits the aperture vertically into left and right halves to mimic a dual-pixel sensor. The half assigned to each sub-aperture is swapped depending on whether the object is nearer or farther than the focus distance, reproducing the depth-dependent left/right disparity that enables dual-pixel depth estimation and autofocus.

Parameters:

Name Type Description Default
points Tensor

Object point positions in [mm], shape [N, 3] with columns [x, y, z]; depth is taken from z.

required
ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS

Returns:

Name Type Description
psf_l Tensor

Left sub-aperture PSFs, shape [N, ks, ks].

psf_r Tensor

Right sub-aperture PSFs, shape [N, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py
def psf_dp(self, points, ks=PSF_KS):
    """Generate left/right dual-pixel PSFs by masking the base PSF.

    Takes the base defocus PSF and splits the aperture vertically into left
    and right halves to mimic a dual-pixel sensor. The half assigned to each
    sub-aperture is swapped depending on whether the object is nearer or
    farther than the focus distance, reproducing the depth-dependent left/right
    disparity that enables dual-pixel depth estimation and autofocus.

    Args:
        points (torch.Tensor): Object point positions in [mm], shape [N, 3]
            with columns [x, y, z]; depth is taken from z.
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.

    Returns:
        psf_l (torch.Tensor): Left sub-aperture PSFs, shape [N, ks, ks].
        psf_r (torch.Tensor): Right sub-aperture PSFs, shape [N, ks, ks].
    """
    depth = points[:, 2]

    # Get the base PSF
    psf_base = self.psf(points, ks=ks, psf_type="gaussian")
    device = psf_base.device

    # Create left and right masks for dual pixel simulation
    l_mask = torch.ones((ks, ks), device=device)
    r_mask = torch.ones((ks, ks), device=device)

    # Split aperture vertically (left half and right half)
    l_pixel, r_pixel = ks // 2, ks // 2 + 1
    l_mask[:, 0:l_pixel] = 0  # Block right side for left PSF
    r_mask[:, r_pixel:] = 0  # Block left side for right PSF

    # Determine focus positions
    depth = depth.to(device)
    foc_dist = torch.tensor(self.foc_dist, device=device, dtype=depth.dtype)
    near_focus_pos = depth > foc_dist  # Shape [N]

    # Apply masks based on focus position (vectorized)
    # For near focus: left PSF gets left mask, right PSF gets right mask
    # For far focus: masks are swapped to create opposite asymmetry
    nfp = near_focus_pos.unsqueeze(-1).unsqueeze(-1)  # [N, 1, 1]
    mask_l = torch.where(nfp, l_mask, r_mask)  # [N, ks, ks]
    mask_r = torch.where(nfp, r_mask, l_mask)  # [N, ks, ks]
    psf_l = psf_base * mask_l
    psf_r = psf_base * mask_r

    # Normalize PSFs
    psf_l = psf_l / (psf_l.sum(dim=(-1, -2), keepdim=True) + EPSILON)
    psf_r = psf_r / (psf_r.sum(dim=(-1, -2), keepdim=True) + EPSILON)

    return psf_l, psf_r

psf_rgb_dp

psf_rgb_dp(points, ks=PSF_KS)

Compute RGB dual-pixel PSFs for left and right sub-apertures.

Replicates the monochrome dual-pixel PSFs across three colour channels.

Parameters:

Name Type Description Default
points Tensor

Object point positions in [mm], shape [N, 3].

required
ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS

Returns:

Name Type Description
psf_l Tensor

Left sub-aperture RGB PSFs, shape [N, 3, ks, ks].

psf_r Tensor

Right sub-aperture RGB PSFs, shape [N, 3, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py
def psf_rgb_dp(self, points, ks=PSF_KS):
    """Compute RGB dual-pixel PSFs for left and right sub-apertures.

    Replicates the monochrome dual-pixel PSFs across three colour channels.

    Args:
        points (torch.Tensor): Object point positions in [mm], shape [N, 3].
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.

    Returns:
        psf_l (torch.Tensor): Left sub-aperture RGB PSFs, shape [N, 3, ks, ks].
        psf_r (torch.Tensor): Right sub-aperture RGB PSFs, shape [N, 3, ks, ks].
    """
    psf_l, psf_r = self.psf_dp(points, ks=ks)
    psf_l = psf_l.unsqueeze(1).repeat(1, 3, 1, 1)
    psf_r = psf_r.unsqueeze(1).repeat(1, 3, 1, 1)
    return psf_l, psf_r

psf_map_dp

psf_map_dp(grid=(5, 5), ks=PSF_KS, depth=None, **kwargs)

Compute spatially-uniform dual-pixel PSF maps.

Parameters:

Name Type Description Default
grid tuple

Grid dimensions (rows, cols). Defaults to (5, 5).

(5, 5)
ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS
depth float or None

Object depth in [mm]. When None (default), falls back to self.obj_depth.

None
**kwargs

Forwarded to psf_dp.

{}

Returns:

Name Type Description
psf_map_l Tensor

Left sub-aperture PSF map, shape [rows, cols, 1, ks, ks].

psf_map_r Tensor

Right sub-aperture PSF map, shape [rows, cols, 1, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py
def psf_map_dp(self, grid=(5, 5), ks=PSF_KS, depth=None, **kwargs):
    """Compute spatially-uniform dual-pixel PSF maps.

    Args:
        grid (tuple, optional): Grid dimensions (rows, cols).
            Defaults to (5, 5).
        ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        depth (float or None, optional): Object depth in [mm]. When None
            (default), falls back to `self.obj_depth`.
        **kwargs: Forwarded to `psf_dp`.

    Returns:
        psf_map_l (torch.Tensor): Left sub-aperture PSF map, shape
            [rows, cols, 1, ks, ks].
        psf_map_r (torch.Tensor): Right sub-aperture PSF map, shape
            [rows, cols, 1, ks, ks].
    """
    depth = self.obj_depth if depth is None else depth
    points = torch.tensor([[0, 0, depth]], device=self.device)
    psf_l, psf_r = self.psf_dp(points, ks=ks, **kwargs)
    psf_map_l = psf_l.unsqueeze(0).unsqueeze(0).repeat(grid[0], grid[1], 1, 1, 1)
    psf_map_r = psf_r.unsqueeze(0).unsqueeze(0).repeat(grid[0], grid[1], 1, 1, 1)
    return psf_map_l, psf_map_r

render_rgbd

render_rgbd(img_obj, depth_map, psf_ks=PSF_KS, num_layers=16)

Occlusion-aware RGBD rendering for defocus lens.

Uses back-to-front layered compositing to prevent color bleeding at depth discontinuities. Since defocus lenses have no spatially varying aberrations, rendering uses a spatially invariant PSF sampled across depth layers.

Parameters:

Name Type Description Default
img_obj Tensor

Object image, shape [B, C, H, W].

required
depth_map Tensor

Depth map in [mm], shape [B, 1, H, W] (or [B, H, W]). Values must be positive.

required
psf_ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS
num_layers int

Number of depth layers. Defaults to 16.

16

Returns:

Name Type Description
img_render Tensor

Rendered image, shape [B, C, H, W].

Reference

[1] "Dr.Bokeh: DiffeRentiable Occlusion-aware Bokeh Rendering", CVPR 2024.

Source code in deeplens-src/deeplens/defocuslens.py
def render_rgbd(
    self,
    img_obj,
    depth_map,
    psf_ks=PSF_KS,
    num_layers=16,
):
    """Occlusion-aware RGBD rendering for defocus lens.

    Uses back-to-front layered compositing to prevent color bleeding at depth
    discontinuities. Since defocus lenses have no spatially varying
    aberrations, rendering uses a spatially invariant PSF sampled across
    depth layers.

    Args:
        img_obj (torch.Tensor): Object image, shape [B, C, H, W].
        depth_map (torch.Tensor): Depth map in [mm], shape [B, 1, H, W]
            (or [B, H, W]). Values must be positive.
        psf_ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        num_layers (int, optional): Number of depth layers. Defaults to 16.

    Returns:
        img_render (torch.Tensor): Rendered image, shape [B, C, H, W].

    Reference:
        [1] "Dr.Bokeh: DiffeRentiable Occlusion-aware Bokeh Rendering", CVPR 2024.
    """
    if depth_map.min() < 0:
        raise ValueError("Depth map should be positive.")

    if len(depth_map.shape) == 3:
        depth_map = depth_map.unsqueeze(1)  # [B, H, W] -> [B, 1, H, W]

    depth_min = depth_map.min()
    depth_max = depth_map.max()

    # Sample depth layers
    disp_ref, depths_ref = self._sample_depth_layers(
        depth_min, depth_max, num_layers
    )

    # Compute PSF at each depth layer (spatially invariant, so patch_center=(0,0))
    points = torch.stack(
        [
            torch.zeros_like(depths_ref),
            torch.zeros_like(depths_ref),
            depths_ref,
        ],
        dim=-1,
    )
    psfs = self.psf_rgb(points=points, ks=psf_ks)  # [num_layers, 3, ks, ks]

    # Occlusion-aware rendering
    img_render = conv_psf_occlusion(img_obj, -depth_map, psfs, depths_ref)
    return img_render

render_rgbd_dp

render_rgbd_dp(rgb_img, depth, psf_ks=PSF_KS, num_layers=16)

Render left/right dual-pixel images from an RGBD input.

Computes dual-pixel PSFs at uniformly sampled reference depths and convolves the image with depth interpolation to produce the two sub-aperture views. Positive depths are negated internally so the object lies in front of the lens.

Parameters:

Name Type Description Default
rgb_img Tensor

RGB object image, shape [B, 3, H, W].

required
depth Tensor

Depth map in [mm], shape [B, 1, H, W].

required
psf_ks int

PSF kernel size in pixels. Defaults to PSF_KS.

PSF_KS
num_layers int

Number of depth layers. Defaults to 16.

16

Returns:

Name Type Description
img_left Tensor

Left sub-aperture image, shape [B, 3, H, W].

img_right Tensor

Right sub-aperture image, shape [B, 3, H, W].

Source code in deeplens-src/deeplens/defocuslens.py
def render_rgbd_dp(
    self,
    rgb_img,
    depth,
    psf_ks=PSF_KS,
    num_layers=16,
):
    """Render left/right dual-pixel images from an RGBD input.

    Computes dual-pixel PSFs at uniformly sampled reference depths and convolves
    the image with depth interpolation to produce the two sub-aperture views.
    Positive depths are negated internally so the object lies in front of the lens.

    Args:
        rgb_img (torch.Tensor): RGB object image, shape [B, 3, H, W].
        depth (torch.Tensor): Depth map in [mm], shape [B, 1, H, W].
        psf_ks (int, optional): PSF kernel size in pixels. Defaults to PSF_KS.
        num_layers (int, optional): Number of depth layers. Defaults to 16.

    Returns:
        img_left (torch.Tensor): Left sub-aperture image, shape [B, 3, H, W].
        img_right (torch.Tensor): Right sub-aperture image, shape [B, 3, H, W].
    """
    # Convert depth to negative values
    if (depth > 0).any():
        depth = -depth

    depth_min = depth.min()
    depth_max = depth.max()
    patch_center = (0.0, 0.0)

    # Calculate dual-pixel PSF at reference depths
    depths_ref = torch.linspace(
        depth_min, depth_max, num_layers, device=self.device
    )
    points = torch.stack(
        [
            torch.full_like(depths_ref, patch_center[0]),
            torch.full_like(depths_ref, patch_center[1]),
            depths_ref,
        ],
        dim=-1,
    )
    psfs_left, psfs_right = self.psf_rgb_dp(
        points=points, ks=psf_ks
    )  # shape [num_layers, 3, ks, ks]

    # Render dual-pixel image with PSF convolution and depth interpolation
    img_left = conv_psf_depth_interp(rgb_img, depth, psfs_left, depths_ref)
    img_right = conv_psf_depth_interp(rgb_img, depth, psfs_right, depths_ref)
    return img_left, img_right