如何在 python 中栅格化一批多边形

How to rasterize a batch of polygons in python

我有一批多边形,它们可以是 NumPy 数组、torch 张量或任何其他形状为 [230_000,3,2] 的 nd 数组。 230_000是多边形的个数,3表示是三角形,2是它的x和y坐标。

我还有一批形状特征[230_000,3,3是颜色],最好在光栅化的时候插值特征的颜色。

我正在尝试找到一种方法来光栅化所有这些,使输出看起来像这样:

算法规则不多。最重要的是填充的多边形在 nd 数组上栅格化并且速度很快。批量大小将是数十万个多边形。 for 循环将无济于事。

如果您打算使用 python 创建光栅,我会查看 Rasterio。如果您的主要目标是使用 Python 创建上述可视化效果,则更简单的方法是使用 shapely 和 matplotlib。您似乎也可能遗漏了一块——您还需要一个生成形状的函数。你提到:

230_000 is the number of polygons, 3 indicates it's a triangle, 2 is its x and y coordinates

2是什么点的x和y坐标?需要更多信息才能了解您需要生成的形状类型。

据我所知,你需要做的是:

我确实最终想出了一种在 pytorch 中完成它的方法。速度不快,但有效。

import torch
import torch.nn as nn
import cv2

class TriRasterizer(nn.Module):
    def __init__(self,shape=(480,640),channels:int=3,_trainable=False):
        super(TriRasterizer, self).__init__()
        assert len(shape) == 2
        self._shape = shape
        self._backdrop = nn.Parameter(torch.zeros(*shape,channels,dtype=torch.float),requires_grad=_trainable)
        # Feel free to mess around with the quality
        self.quality = 0.1
        ...

    @staticmethod
    @torch.jit.script
    def _fill(poly,area,_type:torch.dtype=torch.float32):
        """
        The main part
        The algorithm takes the tensor points, repeats interleave them by their area
        Then for each tiangles in the batch:
            there should be the-area more triangles exactly the same
            the goal is for each of those extra triangles,
                interpolate each of those points such that: each points fills up the traingle

        """
        # calculate center of the base
        base_p = (poly[:,0] + poly[:,1])/2
        a_diff = (poly[:,2]-base_p) . detach()

        # then height
        heights = torch.sqrt_(torch.square_(a_diff[...,0]) + torch.square_(a_diff[...,1]))

        # comulative sum of the areas
        # This gives positional embeding for each points when subtracted from torch.arange
        cum   = torch.cumsum(area,dim=0).int()
        poly  = poly.type(_type)
        heights  = heights.type(torch.int16)
        poly    = poly.repeat_interleave(area,dim=0)
        cum_    = cum.repeat_interleave(area,dim=0)
        heights = heights.repeat_interleave(area,dim=0)

        # normalized makes the first point = 0 and last = 1 for the repeated points
        normalized = 1-(cum_-torch.arange(len(cum_)))/cum_

        # this is the exponential, it reduces pixels as it gets closer to the tip of the triangle
        exp = 1-torch.sqrt_(1-normalized)
        exp = exp.type(torch.float32)

        # This iterpolates more of the x as it gets to the tip of the triangle
        ij__interp = (exp % (torch.tensor(1,dtype=torch.float16)/heights))[...,None]
        # * heights[...,None] just brings the exponentials max back to 1
        # the random normal washes the points and makes it appear blurry
        ij_interp = ij__interp*heights[...,None] #+ torch.randn_like(ij__interp)/70

        # this makes the makes the c interpolation exponentialy faster
        k_interp = (exp[...,None] - ij__interp)  #+ torch.randn_like(ij__interp)/70

        # Finally interpolate them
        # weigh a and b first with respect to the calculated exponential
        ij_trunc = poly[:,0] * (1-ij_interp) + poly[:,1]*ij_interp
        # then weigh the output with c with respect to the modulus of the exponential
        filled   = poly[:,2] * k_interp  + ij_trunc * (1-k_interp)
        return filled

    def _compute_area(self,poly):
        # Formula to find area from a set of points
        area = torch.abs(0.5 * ((poly[..., 0, 0] * (poly[..., 1, 1] - poly[..., 2, 1])) +
                   (poly[...,1,0]*(poly[...,2,1]-poly[...,0,1])) +
                   (poly[...,2,0]*(poly[...,0,1]-poly[...,1,1]))))
        return area

    def forward(self,traingles,features):
        return self.rasterize(traingles,features)


    @torch.no_grad()
    def rasterize(self,traingles,features):
        raster = self._backdrop.clone()
        # first find the area
        # the lower the quality the more spaces in between pixels
        area = (self._compute_area(traingles)*self.quality).int()

        # fill in points based on computed area
        filled_cords = self._fill(traingles,area,_type=torch.int16)
        filled_cords = filled_cords.long()

        # clamp the points to prevent out of bounds exception
        filled_cords[...,0] = filled_cords[...,0].clamp(min=0,max=self._shape[1]-1)
        filled_cords[...,1] = filled_cords[...,1].clamp(min=0,max=self._shape[0]-1)

        # also fill in features based on computed area
        filled_feats = self._fill(features,area,_type=torch.float32)
        pixels = (filled_cords[...,1],filled_cords[...,0])

        # compute the times the same point has been placed on the raster
        sum_ = torch.zeros_like(raster).index_put(pixels,torch.ones_like(filled_feats),accumulate=True)

        # rasterize them using indexing
        raster = raster.index_put(pixels,filled_feats,accumulate=False)

        # average those sums
        raster = raster/sum_
        return raster
#

if __name__ == '__main__':
    tiangles = torch.tensor([
        [[100, 0], [0, 200], [20, 300]]
    ], dtype=torch.float) + 100

    print(tiangles.shape)
    tiangles = tiangles.expand(3,-1,-1).clone()

    features = torch.randn(*tiangles.shape[:-1], 3)
    rasterizer = TriRasterizer()
    j = torch.zeros(*tiangles.shape)

    while True:
        features = features*0.01 + torch.randn(*tiangles.shape[:-1], 3)*0.09
        j += torch.randn(*tiangles.shape)*0.1
        tiangles += j
        raster = rasterizer(tiangles,features+1)
        cv2.imshow('x',raster.numpy())
        cv2.waitKey(20)

我还发现了其他方法,例如绑定 C++ 或 PyOpengl