update documentation and remove some unnecessary lines

neural_renderer/perspective.py

@@ -3,15 +3,18 @@
 import torch
 
 def perspective(vertices, angle=30.):
-        assert(vertices.ndimension() == 3)
-        device = vertices.device
-        angle = torch.tensor(angle / 180 * math.pi).to(device)
-        angle = angle[None]
-
-        width = torch.tan(angle)
-        width = width[:, None] 
-        z = vertices[:, :, 2]
-        x = vertices[:, :, 0] / z / width
-        y = vertices[:, :, 1] / z / width
-        vertices = torch.stack((x,y,z), dim=2)
-        return vertices
+    '''
+    Compute perspective distortion from a given angle
+    '''
+    if (vertices.ndimension() != 3):
+        raise ValueError('vertices Tensor should have 3 dimensions')
+    device = vertices.device
+    angle = torch.tensor(angle / 180 * math.pi, dtype=torch.float32, device=device)
+    angle = angle[None]
+    width = torch.tan(angle)
+    width = width[:, None] 
+    z = vertices[:, :, 2]
+    x = vertices[:, :, 0] / z / width
+    y = vertices[:, :, 1] / z / width
+    vertices = torch.stack((x,y,z), dim=2)
+    return vertices

neural_renderer/projection.py

@@ -1,22 +1,27 @@
 import torch
-import torch.nn as nn
 
 
 def projection(vertices, P, dist_coeffs, orig_size):
+    '''
+    Calculate projective transformation of vertices given a projection matrix
+    P: 3x4 projection matrix
+    dist_coeffs: vector of distortion coefficients
+    orig_size: original size of image captured by the camera
+    '''
     vertices = torch.cat([vertices, torch.ones_like(vertices[:, :, None, 0])], dim=-1)
-    # vertices = torch.bmm(vertices, P.transpose(2,1))
     vertices = torch.bmm(vertices, P.transpose(2,1))
     x, y, z = vertices[:, :, 0], vertices[:, :, 1], vertices[:, :, 2]
     x_ = x / (z + 1e-5)
     y_ = y / (z + 1e-5)
-    # vertices[:, :, :-1] = (vertices[:, :, :-1] / (vertices[:, :, None, -1] + 1e-5))
+
+    # Get distortion coefficients from vector
     k1 = dist_coeffs[:, None, 0]
     k2 = dist_coeffs[:, None, 1]
     p1 = dist_coeffs[:, None, 2]
     p2 = dist_coeffs[:, None, 3]
     k3 = dist_coeffs[:, None, 4]
+
     # we use x_ for x' and x__ for x'' etc.
-    # x_, y_ = vertices[:, :, 0], vertices[:, :, 1]
     r = torch.sqrt(x_ ** 2 + y_ ** 2)
     x__ = x_*(1 + k1*(r**2) + k2*(r**4) + k3*(r**6)) + 2*p1*x_*y_ + p2*(r**2 + 2*x_**2)
     y__ = y_*(1 + k1*(r**2) + k2*(r**4) + k3 *(r**6)) + p1*(r**2 + 2*y_**2) + 2*p2*x_*y_