charlesCXK / TorchSSC

Implement some state-of-the-art methods of Semantic Scene Completion (SSC) task in PyTorch. [1] 3D Sketch-aware Semantic Scene Completion via Semi-supervised Structure Prior (CVPR 2020)
MIT License
62 stars 14 forks source link

Borders calculation code #5

Closed luigui2906 closed 4 years ago

luigui2906 commented 4 years ago

Hi,

First congrats for the good work. I was wondering if you could share the code to calculate the borders on the 3D labels grid. I wanted to test the network in other datasets and need to do this preprocessing step. I've implemented a function myself by using Sobel as explained in your article but there are some slight differences between my calculated borders and yours :

from scipy import ndimage

gt = gt.reshape(_W, _H, _D)  # reshape to 3D [60, 36, 60]
gt_sobel = np.copy(gt)
gt_sobel[gt_sobel == 255] = 0  # Unknown voxels considered as free
dx = ndimage.sobel(gt_sobel, axis=0)  # x derivative
dy = ndimage.sobel(gt_sobel, axis=1)  # y derivative
dz = ndimage.sobel(gt_sobel, axis=2)  # z derivative
dx[dx != 0] = 1 # Binarizing
dy[dy != 0] = 1 # Binarizing
dz[dz != 0] = 1 # Binarizing
borders = ((dx + dy + dz) > 1).astype(np.uint) # Borders

Is there any further step..? Thanks in advance.

charlesCXK commented 4 years ago

I hope the code below could help you.


class Sobel3D(nn.Module):
    def __init__(self):
        super(Sobel3D, self).__init__()
        self.kernel_x = [
            [[1, 0, -1],
             [1, 0, -1],
             [1, 0, -1], ],
            [[2, 0, -2],
             [2, 0, -2],
             [2, 0, -2], ],
            [[1, 0, -1],
             [1, 0, -1],
             [1, 0, -1], ]]

        self.kernel_y = [
            [[1, 2, 1],
             [1, 2, 1],
             [1, 2, 1], ],
            [[0, 0, 0],
             [0, 0, 0],
             [0, 0, 0], ],
            [[-1, -2, -1],
             [-1, -2, -1],
             [-1, -2, -1], ]]

        self.kernel_z = [
            [[1, 2, 1],
             [0, 0, 0],
             [-1, -2, -1]],
            [[1, 2, 1],
             [0, 0, 0],
             [-1, -2, -1]],
            [[1, 2, 1],
             [0, 0, 0],
             [-1, -2, -1], ]]

        self.kernel_x = torch.FloatTensor(self.kernel_x).unsqueeze(0).unsqueeze(0)
        self.kernel_y = torch.FloatTensor(self.kernel_y).unsqueeze(0).unsqueeze(0)
        self.kernel_z = torch.FloatTensor(self.kernel_z).unsqueeze(0).unsqueeze(0)
        self.weight_x = nn.Parameter(data=self.kernel_x, requires_grad=False)#.cuda()
        self.weight_y = nn.Parameter(data=self.kernel_y, requires_grad=False)#.cuda()
        self.weight_z = nn.Parameter(data=self.kernel_z, requires_grad=False)#.cuda()

        self.conv_x = nn.Conv3d(1, 1, kernel_size=3, padding=1, bias=False)
        self.conv_x.weight = self.weight_x
        self.conv_y = nn.Conv3d(1, 1, kernel_size=3, padding=1, bias=False)
        self.conv_y.weight = self.weight_y
        self.conv_z = nn.Conv3d(1, 1, kernel_size=3, padding=1, bias=False)
        self.conv_z.weight = self.weight_z

    def forward(self, data):
        data[data == 255] = 0
        # print(self.conv_x.weight)

        data = data.float()

        dx = self.conv_x(data)
        dy = self.conv_y(data)
        dz = self.conv_z(data)
        dx[abs(dx) > 0] = 1
        dy[abs(dy) > 0] = 1
        dz[abs(dz) > 0] = 1
        cat = torch.cat([dx, dy, dz], dim=1)
        norm = torch.norm(cat, dim=1, keepdim=True)
        thresh = 1
        norm[norm <= thresh] = 0
        norm[norm > thresh] = 1
        return norm

for i in range(1, 1449+1):
    print(i)
    label = np.load(os.path.join(root, '%04d'%i + '.npz'))['arr_0'].reshape(60, 36, 60)
    label = torch.from_numpy(label).view(1, 1, 60, 36, 60)
    sobel = Sobel3D()
    sketch = sobel(label)
    sketch = sketch.numpy().squeeze(0).squeeze(0)
    np.save(os.path.join(root.replace('Label', 'Sketch3D'), '%04d'%i + '.npy'), sketch)
luigui2906 commented 4 years ago

Thanks :)