Surface Normal Prediction and 3D Shape Understanding

Module: Computer Vision | Difficulty: Premium

Normal from Depth Map

Normal-Depth Consistency

Angular Error

Mean Angular Distance

import torch
import torch.nn as nn

class NormalEstimationNet(nn.Module):
    def __init__(self):
        super().__init__()
        import torchvision.models as models
        encoder = models.resnet50(pretrained=True)
        self.encoder = nn.Sequential(
            encoder.conv1, encoder.bn1, encoder.relu,
            encoder.maxpool, encoder.layer1, encoder.layer2,
            encoder.layer3, encoder.layer4,
        )
        self.decoder = nn.Sequential(
            nn.Conv2d(2048, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512), nn.ReLU(inplace=True),
            nn.Upsample(scale_factor=2, mode='bilinear',
                        align_corners=True),
            nn.Conv2d(512, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256), nn.ReLU(inplace=True),
            nn.Upsample(scale_factor=2, mode='bilinear',
                        align_corners=True),
            nn.Conv2d(256, 128, 3, padding=1, bias=False),
            nn.BatchNorm2d(128), nn.ReLU(inplace=True),
            nn.Upsample(scale_factor=2, mode='bilinear',
                        align_corners=True),
            nn.Conv2d(128, 3, 3, padding=1),
            nn.Tanh(),
        )

    def forward(self, x):
        features = self.encoder(x)
        normals = self.decoder(features)
        return nn.functional.normalize(normals, dim=1)

def depth_to_normals(depth, focal_length):
    B, _, H, W = depth.shape
    dx = (depth[:, :, :, 2:] - depth[:, :, :, :-2]) / 2
    dy = (depth[:, :, 2:, :] - depth[:, :, :-2, :]) / 2
    dx = nn.functional.pad(dx, (1, 1, 0, 0))
    dy = nn.functional.pad(dy, (0, 0, 1, 1))
    normals = torch.stack([
        -dx * focal_length,
        -dy * focal_length,
        torch.ones_like(depth),
    ], dim=1)
    return nn.functional.normalize(normals, dim=1)

Research Insight: Surface normal estimation provides complementary information to depth: normals capture local surface orientation, which is crucial for understanding geometry, material properties, and lighting. UNINormal demonstrated that normal estimation can be self-supervised using the consistency between predicted normals and depth gradients. The key insight is that normal estimation benefits from multi-task learning with depth and segmentation, as geometric and semantic information are complementary. Recent methods predict normals directly from point clouds, enabling real-time 3D reconstruction for robotics.