Softrasterizer for Mask from 2D points about softras HOT 2 OPEN

def create_silhouette(
        size, batch_size, edge_points, offset_points, pixel_points, n_rays,
        sigma=.02,
        ):
    """
    Create a differentiable silhouette as mentioned in the SoftRas paper.
    https://github.com/ShichenLiu/SoftRas
    """
    start_points = offset_points
    x = torch.linspace(0, 1, size)
    y = torch.linspace(0, 1, size)
    mx, my = torch.meshgrid(x, y)
    # in the 3rd channel we save, if the pixel belongs to the triangle
    Ps = torch.stack((mx, my))
    # extend grid
    # first we need a grid for each ray
    Ps = Ps.expand(n_rays, 2, size, size)
    # second we need it for each patch in the batch
    Ps = Ps.expand(batch_size, n_rays, 2, size, size)
    
    # this vector saves the distances and the property of belonging to a
    # triangle for each pixel for each ray for each patch in the batch
    # we have 4 channels: one for the belonging: -1 / +1
    # three for the smallest distances to each line segment
    # in the end we can take the min over the channel axis
    belonging = torch.empty(batch_size, n_rays, 1, size, size)
    distances = torch.empty(batch_size, n_rays, 3, size, size)
    
    # collapse batch_size and the two size axes
    # put the two sizes in front, because of expand statements in for loop
    Ps = Ps.permute(3, 4, 0, 1, 2)
    belonging = belonging.permute(3, 4, 0, 1, 2)
    distances = distances.permute(3, 4, 0, 1, 2)
    Ps = Ps.flatten(0, 2)
    belonging = belonging.flatten(0, 2)
    distances = distances.flatten(0, 2)
    # do the same with the points C
    C = start_points.expand(size, batch_size, 2)
    C = C.expand(size, size, batch_size, 2)
    C = C.flatten(0, 2)
    
    # I think we now have to use a for loop to iterate over the different
    # triangles in each patch
    for i in range(n_rays):
        # C is always in the center
        # since the rays are distributed counterclockwise, we always get a
        # triangle, where ABC are in clockwise order
        B = edge_points[:,i,:]
        A = edge_points[:,(i+1)%n_rays,:]
        # now expand these points of the triangles in each patch to image size
        B = B.expand(size, batch_size, 2)
        B = B.expand(size, size, batch_size, 2)
        A = A.expand(size, batch_size, 2)
        A = A.expand(size, size, batch_size, 2)
        # and permute axes
        B = B.flatten(0, 2)
        A = A.flatten(0, 2)
        
        # lets follow this website
        # https://blackpawn.com/texts/pointinpoly/default.html
        # P = grid without the belonging channel
        P = Ps[:, i, :]
        v0 = C - A
        v1 = B - A
        v2 = P - A
        
        dot00 = my_dot(v0, v0)
        dot01 = my_dot(v0, v1)
        dot02 = my_dot(v0, v2)
        dot11 = my_dot(v1, v1)
        dot12 = my_dot(v1, v2)

        # Compute barycentric coordinates
        invDenom = 1 / (dot00 * dot11 - dot01 * dot01)
        u = (dot11 * dot02 - dot01 * dot12) * invDenom
        v = (dot00 * dot12 - dot01 * dot02) * invDenom

        # Check if point is in triangle
        belonging[:, i, :] = ((u >= 0) * (v >= 0) * (u + v < 1))*2 - 1
        
        distances[:, i, :1] = distance_linesegment_no_sqrt(P, A, B)
        distances[:, i, 1:2] = distance_linesegment_no_sqrt(P, B, C)
        distances[:, i, 2:] = distance_linesegment_no_sqrt(P, C, A)
        
    # after for loop lets reshape all needed tensors to the correct shape
    belonging = belonging.reshape(size, size, batch_size, n_rays, 1)
    distances = distances.reshape(size, size, batch_size, n_rays, 3)
    # and permute
    belonging = belonging.permute(2, 3, 4, 0, 1)
    distances = distances.permute(2, 3, 4, 0, 1)
    # finally take the minimum of the distances to each line segment of triangle
    distances, _ = distances.min(dim=2, keepdim=True)
    # now calculate the probabilities for each triangle
    probabilities = torch.sigmoid(belonging*distances/sigma)
    # substract the probabilities from 1
    probabilities = 1 - probabilities
    # we now have to multiply over the rays axis and substract that from one
    prod = torch.prod(probabilities, dim=1)
    silhouette = 1 - prod
    
    return silhouette
        
def my_dot(v1, v2):
#    bs = v1.size[0]
    x = v1[:, 0] * v2[:, 0]
    y = v1[:, 1] * v2[:, 1]
    dot = x+y
    dot = dot.expand(1, -1)
    dot = dot.transpose(0, 1)
    return dot

def distance_linesegment_no_sqrt(P0, P1, P2):
    # we will use this post for the shortest distance of point to triangle
    # https://stackoverflow.com/a/10984080/10106730
    b = P0 - P1
    a = P2 - P1
    r = my_dot(a, b) / my_dot(a, a)
    
    output = torch.empty_like(r)
    output[r<=0] = my_dot(b, b)[r<=0]
    
    c = P0 - P2
    output[r>=1] = my_dot(c, c)[r>=1]
    
    output[(r>0)*(r<1)] = (my_dot(b, b) - r.pow(2) * my_dot(a, a))[(r>0)*(r<1)]
    return output