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.
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]. Must be less than the focal length (i.e. beyond the focal point).	required

Raises:

Type	Description
AssertionError	If foc_dist >= 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].  Must be less than the
            focal length (i.e. beyond the focal point).

    Raises:
        AssertionError: If *foc_dist* >= ``self.foclen``.
    """
    assert foc_dist < self.foclen, "Focus distance is too close."
    self.foc_dist = foc_dist

psf

psf(points, ks=PSF_KS, psf_type='gaussian', **kwargs)

PSF is modeled as a 2D uniform circular disk with diameter CoC.

Parameters:

Name	Type	Description	Default
points	Tensor	Points of the object. Shape [N, 3] or [3].	required
ks	int	Kernel size.	PSF_KS
psf_type	str	PSF type. "gaussian" or "pillbox".	'gaussian'
**kwargs		Additional arguments for psf(). Currently not used.	{}

Returns:

Name	Type	Description
psf	Tensor	PSF kernels. Shape [ks, ks] or [N, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py

def psf(self, points, ks=PSF_KS, psf_type="gaussian", **kwargs):
    """PSF is modeled as a 2D uniform circular disk with diameter CoC.

    Args:
        points (torch.Tensor): Points of the object. Shape [N, 3] or [3].
        ks (int): Kernel size.
        psf_type (str): PSF type. "gaussian" or "pillbox".
        **kwargs: Additional arguments for psf(). Currently not used.

    Returns:
        psf (torch.Tensor): PSF kernels. Shape [ks, ks] or [N, ks, ks].
    """
    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 coordinate meshgrid
    x, y = torch.meshgrid(
        torch.linspace(-ks / 2 + 1 / 2, ks / 2 - 1 / 2, ks, device=self.device),
        torch.linspace(-ks / 2 + 1 / 2, ks / 2 - 1 / 2, ks, device=self.device),
        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)

Calculate circle of confusion (CoC) diameter [mm].

Parameters:

Name	Type	Description	Default
depth	Tensor	Depth of the object. Shape [B].	required

Returns:

Name	Type	Description
coc	Tensor	Circle of confusion diameter [mm]. Shape [B].

Reference

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

Source code in deeplens-src/deeplens/defocuslens.py

def coc(self, depth):
    """Calculate circle of confusion (CoC) diameter [mm].

    Args:
        depth (torch.Tensor): Depth of the object. Shape [B].

    Returns:
        coc (torch.Tensor): Circle of confusion diameter [mm]. Shape [B].

    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)

Calculate depth of field [mm].

Parameters:

Name	Type	Description	Default
depth	Tensor	Depth of the object. Shape [B].	required

Returns:

Name	Type	Description
dof	Tensor	Depth of field. Shape [B].

Reference

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

Source code in deeplens-src/deeplens/defocuslens.py

def dof(self, depth):
    """Calculate depth of field [mm].

    Args:
        depth (torch.Tensor): Depth of the object. Shape [B].

    Returns:
        dof (torch.Tensor): Depth of field. Shape [B].

    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	Point source positions, shape [N, 3].	required
ks	int	Kernel size. Defaults to PSF_KS.	PSF_KS
**kwargs		Forwarded to psf.	{}

Returns:

Type	Description
	torch.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): Point source positions, shape ``[N, 3]``.
        ks (int, optional): Kernel size. Defaults to ``PSF_KS``.
        **kwargs: Forwarded to `psf`.

    Returns:
        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	Kernel size. Defaults to PSF_KS.	PSF_KS
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
**kwargs		Forwarded to psf.	{}

Returns:

Type	Description
	torch.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): Kernel size. Defaults to ``PSF_KS``.
        depth (float, optional): Object depth [mm]. When ``None`` (default),
            falls back to ``self.obj_depth``.
        **kwargs: Forwarded to `psf`.

    Returns:
        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 dual-pixel PSF for left and right sub-apertures.

This function generates separate PSFs for left and right sub-apertures of a dual pixel sensor, which enables depth estimation and improved autofocus capabilities.

Parameters:

Name	Type	Description	Default
points	Tensor	Input tensor with shape [N, 3], where columns are [x, y, z] coordinates.	required
ks	int	Kernel size for PSF generation.	PSF_KS

Returns:

Name	Type	Description
tuple		(left_psf, right_psf) where each PSF tensor has shape [N, ks, ks].

Source code in deeplens-src/deeplens/defocuslens.py

def psf_dp(self, points, ks=PSF_KS):
    """Generate dual-pixel PSF for left and right sub-apertures.

    This function generates separate PSFs for left and right sub-apertures of a dual pixel sensor,
    which enables depth estimation and improved autofocus capabilities.

    Args:
        points (torch.Tensor): Input tensor with shape [N, 3], where columns are [x, y, z] coordinates.
        ks (int): Kernel size for PSF generation.

    Returns:
        tuple: (left_psf, right_psf) where each PSF tensor has 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	Point source positions, shape [N, 3].	required
ks	int	Kernel size. Defaults to PSF_KS.	PSF_KS

Returns:

Name	Type	Description
tuple		(psf_left, psf_right) each of 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): Point source positions, shape ``[N, 3]``.
        ks (int, optional): Kernel size. Defaults to ``PSF_KS``.

    Returns:
        tuple: ``(psf_left, psf_right)`` each of 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	Kernel size. Defaults to PSF_KS.	PSF_KS
depth	float	Object depth [mm]. When None (default), falls back to self.obj_depth.	None
**kwargs		Forwarded to psf_dp.	{}

Returns:

Name	Type	Description
tuple		(psf_map_left, psf_map_right) each of 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): Kernel size. Defaults to ``PSF_KS``.
        depth (float, optional): Object depth [mm]. When ``None`` (default),
            falls back to ``self.obj_depth``.
        **kwargs: Forwarded to `psf_dp`.

    Returns:
        tuple: ``(psf_map_left, psf_map_right)`` each of 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 [mm]. Shape [B, 1, H, W]. Values should be positive.	required
psf_ks	int	PSF kernel size. 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 (tensor): Object image. Shape [B, C, H, W].
        depth_map (tensor): Depth map [mm]. Shape [B, 1, H, W]. Values should be positive.
        psf_ks (int, optional): PSF kernel size. Defaults to PSF_KS.
        num_layers (int, optional): Number of depth layers. Defaults to 16.

    Returns:
        img_render (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 RGBD image with dual-pixel PSF.

Parameters:

Name	Type	Description	Default
rgb_img	tensor	[B, 3, H, W]	required
depth	tensor	[B, 1, H, W]	required
psf_ks	int	PSF kernel size. Defaults to PSF_KS.	PSF_KS
num_layers	int	Number of depth layers. Defaults to 16.	16

Returns:

Name	Type	Description
img_left	tensor	[B, 3, H, W]
img_right	tensor	[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 RGBD image with dual-pixel PSF.

    Args:
        rgb_img (tensor): [B, 3, H, W]
        depth (tensor): [B, 1, H, W]
        psf_ks (int, optional): PSF kernel size. Defaults to PSF_KS.
        num_layers (int, optional): Number of depth layers. Defaults to 16.

    Returns:
        img_left (tensor): [B, 3, H, W]
        img_right (tensor): [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