add convnext implementation

data/dataset.py

@@ -13,7 +13,7 @@ import plotly.express as ex
 
 from torch.utils.data import Dataset
 
-from render import Render
+from render import Render, DepthRenderer
 
 
 class Dataset(Dataset):
@@ -22,7 +22,7 @@ class Dataset(Dataset):
         self.transforms = tv.transforms.ToTensor()
         self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))
 
-        self.render = Render(max_depth=10)
+        self.render = DepthRenderer()
         self.body_faces = np.load('custom_data/smpl_faces.npy')
         
 
@@ -41,12 +41,15 @@ class Dataset(Dataset):
 
         info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat'
         info = sio.loadmat(info_folder)
-        combined = np.concatenate((body, mesh_posed[mask, :]))
-        combined_faces = np.concatenate((self.body_faces, faces + 14475))
+        combined = torch.Tensor(np.concatenate((body, mesh_posed[mask, :]))).to('cuda')
+        combined_faces = torch.Tensor(np.concatenate((self.body_faces, faces + 14475))).to('cuda')
 
         start = time.perf_counter()
-        depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
+        # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
+        # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
+        depth_image = self.render(combined, combined_faces, info['camLoc'][0])
         t = time.perf_counter() - start
+        print(t)
 
         # image = self.transforms(image)
         return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask }

data/render.py

@@ -81,6 +81,7 @@ def intrinsic():
 					 [0,     fy_px, v],
 					 [0,     0,     1]], np.float32)
 				
+
 def extrinsic(camLoc):
 	R_w2bc = np.array([[0, 1, 0],
 					   [0, 0, 1],
@@ -122,3 +123,56 @@ class Render:
         mask = mask.astype(np.float32)
 
         return depth
+
+
+from torch import nn
+from pytorch3d.renderer import (
+    look_at_view_transform,
+    FoVPerspectiveCameras, 
+    PerspectiveCameras,
+    PointLights, 
+    DirectionalLights, 
+    Materials, 
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesUV,
+    TexturesVertex
+)
+from pytorch3d.structures import Meshes
+
+class DepthRenderer(nn.Module):
+    def __init__(self) -> None:
+         super().__init__()
+
+
+
+    def forward(self, verts, faces, cam):
+        R, T = look_at_view_transform(2.7, 0, 0) 
+        cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T)
+        ext = extrinsic(cam)
+        cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=50)
+
+        raster_settings = RasterizationSettings(
+                            image_size=512, 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+
+        rasterizer = MeshRasterizer(
+                    cameras=cameras, 
+                    raster_settings=raster_settings
+                )
+        depth = rasterizer(Meshes(verts[None,:,:], faces[None,:,:])).zbuf[0,:,:,0]
+
+        mask = depth != -1
+        # dmin = depth.min()
+        dmean = depth[mask].mean()
+        # dmax = depth[mask].max()
+        depth[mask] = depth[mask] - dmean
+        depth[~mask] = 0.0
+        # mask = mask.astype(np.float32)
+
+        return depth

models/convnext.py

@@ -0,0 +1,204 @@
+# from torchvision.models.convnext
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+from torchvision.ops.stochastic_depth import StochasticDepth
+from torchvision.models._api import register_model, Weights, WeightsEnum
+
+
+
+__all__ = [
+    "ConvNeXt",
+    "ConvNeXt_Tiny_Weights",
+    "ConvNeXt_Small_Weights",
+    "ConvNeXt_Base_Weights",
+    "ConvNeXt_Large_Weights",
+    "convnext_tiny",
+    "convnext_small",
+    "convnext_base",
+    "convnext_large",
+]
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+class CNBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        layer_scale: float,
+        stochastic_depth_prob: float,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.block = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True),
+            Permute([0, 2, 3, 1]),
+            norm_layer(dim),
+            nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+            nn.GELU(),
+            nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+            Permute([0, 3, 1, 2]),
+        )
+        self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.layer_scale * self.block(input)
+        result = self.stochastic_depth(result)
+        result += input
+        return result
+
+
+class CNBlockConfig:
+    # Stores information listed at Section 3 of the ConvNeXt paper
+    def __init__(
+        self,
+        input_channels: int,
+        out_channels: Optional[int],
+        num_layers: int,
+    ) -> None:
+        self.input_channels = input_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "input_channels={input_channels}"
+        s += ", out_channels={out_channels}"
+        s += ", num_layers={num_layers}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+
+class ConvNeXtBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super().__init__()
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        layers: List[nn.Module] = []
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        layers.append(
+            Conv2dNormActivation(
+                3,
+                firstconv_output_channels,
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                norm_layer=norm_layer,
+                activation_layer=None,
+                bias=True,
+            )
+        )
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        for cnf in block_setting:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                stage_block_id += 1
+            layers.append(nn.Sequential(*stage))
+            if cnf.out_channels is not None:
+                # Downsampling
+                layers.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.features = nn.Sequential(*layers)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+def _convnext(
+    block_setting: List[CNBlockConfig],
+    stochastic_depth_prob: float,
+    weights: Optional[WeightsEnum],
+    **kwargs: Any,
+) -> ConvNeXtBase:
+
+    model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+
+    if weights is not None:
+        model.load_state_dict(weights.get_state_dict())
+
+    return model
+
+
+class ConvNeXt(nn.Module):
+    def __init__(self, stochastic_depth_prob=0.0, **kwargs) -> None:
+        super().__init__()
+
+
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, 768, 9),
+            CNBlockConfig(768, None, 3),
+        ]
+        self.model = _convnext(block_setting, stochastic_depth_prob, None, **kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.model(x)
+        return x
+        
+
+
+if __name__ == '__main__':
+    model = ConvNeXt().to('cuda')
+    test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
+
+    pass