trying to get basic model working

data/dataset.py

@@ -24,7 +24,7 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda'):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True):
         self.folder = folder
         self.folder_rgb = folder_rgb
         self.device = device
@@ -37,6 +37,8 @@ class Dataset(Dataset):
         self.reader = DataReader(folder, 'superSMPL')
         
 
+        self.cache = dict() if use_cache else None
+
     def __len__(self):
         return len(self.frames)
 
@@ -44,6 +46,10 @@ class Dataset(Dataset):
     # TODO: make it work with two part outfits 
     def __getitem__(self, idx):
         frame : Frame = self.frames[idx]
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                return self.cache[frame]
+
         path = os.path.join(self.folder, frame.id)
         info = self.reader.read_info(frame.id)
         # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
@@ -189,7 +195,7 @@ class Dataset(Dataset):
 
 
         # image = self.transforms(image)
-        return { 'mesh_posed': mesh_posed,
+        data = { 'mesh_posed': mesh_posed,
                  'depth_image': depth_image, 
                  'mesh_mask': mesh_mask, 
                  'depth_mask': depth_mask, 
@@ -199,6 +205,9 @@ class Dataset(Dataset):
                  'rgb_image': rgb_image,
                 }
 
+        self.cache[frame] = data
+        return self.cache[frame]
+
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader

models/convnext.py

@@ -1,6 +1,4 @@
-# from torchvision.models.convnext
-import os
-os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+# taken from torchvision.models.convnext
 
 from functools import partial
 from typing import Any, Callable, List, Optional, Sequence
@@ -11,7 +9,7 @@ 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
+from torchvision.models._api import Weights, WeightsEnum
 
 
 
@@ -163,14 +161,112 @@ class ConvNeXtBase(nn.Module):
         return self._forward_impl(x)
 
 
+
+class ConvNeXtReverseBase(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] = []
+
+
+        # add bottleneck which will be executed at the end, to process after the final larger upsampling
+        stage: List[nn.Module] = []
+        for _ in range(3):
+            stage.append(block(1, layer_scale, 0.0))
+        layers.append(nn.Sequential(*stage))
+
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        layers.append(
+            # Convtr NormActivation(
+            #     firstconv_output_channels,
+            #     3,
+            #     kernel_size=4,
+            #     stride=4,
+            #     padding=0,
+            #     norm_layer=None,
+            #     activation_layer=None,
+            #     bias=True,
+            # )
+            nn.Sequential(
+                nn.ConvTranspose2d(
+                    firstconv_output_channels,
+                    1,
+                    kernel_size=4,
+                    stride=4,
+                    padding=0,
+                    ),
+            )
+        )
+
+        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:
+                # Upsampling
+                layers.append(
+                    nn.Sequential(
+                        norm_layer(cnf.out_channels),
+                        nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.features = nn.Sequential(*reversed(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)
+        x = torch.tanh(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],
+    reversed: bool,
     **kwargs: Any,
 ) -> ConvNeXtBase:
 
-    model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+    if reversed:
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+    else:
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
 
     if weights is not None:
         model.load_state_dict(weights.get_state_dict())
@@ -179,7 +275,7 @@ def _convnext(
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, stochastic_depth_prob=0.0, **kwargs) -> None:
+    def __init__(self, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
         super().__init__()
 
 
@@ -189,7 +285,7 @@ class ConvNeXt(nn.Module):
             CNBlockConfig(384, 768, 9),
             CNBlockConfig(768, None, 3),
         ]
-        self.model = _convnext(block_setting, stochastic_depth_prob, None, **kwargs)
+        self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.model(x)

models/depth_estim.py

@@ -3,15 +3,24 @@ import sys
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from models.convnext import ConvNeXt
 
 
 class DepthEstimation(nn.Module):
     def __init__(self, device='cuda'):
         super().__init__()
 
-        self.layer_1 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, device=device, padding='same')  
+        self.convnext = ConvNeXt()
+        self.flatten = nn.Flatten()
+        self.decoder = ConvNeXt(reversed=True)
+
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.layer_1(x)
+        x = self.convnext(x)
+        # x = self.flatten(x)
+        x = self.decoder(x)
         
         return x

test.py

@@ -1,4 +1,6 @@
 import os
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
 import sys
 
 import torch
@@ -7,9 +9,13 @@ import torch.optim as optim
 import torch.nn as nn
 
 from torch.utils.data import DataLoader
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
 from data.dataset import Dataset
 from utils.utils import Accumulator
 from models.depth_estim import DepthEstimation
+from models.convnext import ConvNeXt
 
 
 def train_step(batch):
@@ -20,7 +26,7 @@ def train_step(batch):
     y_pred = model(batch['rgb_image'])
 
     # 2. Calculate loss
-    loss = loss_fn(y_pred, batch['depth_image'][:,None,:,:])
+    loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
 
     # 3. Zero grad optimizer
     optimizer.zero_grad()
@@ -42,7 +48,7 @@ def valid_step(batch):
         pred = model(batch['rgb_image'])
     
         # 2. Calculate the loss
-        loss = loss_fn(pred, batch['depth_image'][:,None,:,:])
+        loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
 
@@ -57,14 +63,14 @@ def batch_to_device(batch, device):
 if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
-    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
+    train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
+    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
     train_dataloader = DataLoader(train_dataset, batch_size=2, drop_last=True, num_workers=0)
     valid_dataloader = DataLoader(valid_dataset, batch_size=2, drop_last=True, num_workers=0)
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
-    model = DepthEstimation()
+    model = DepthEstimation().to('cuda')
 
     loss_fn = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)

utils/data_utils.py

@@ -8,7 +8,7 @@ from math import cos, sin
 
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from superSMPL.IO import readPC2Frame, readOBJ
+# from superSMPL.IO import readPC2Frame, readOBJ
 
 
 class Frame:
@@ -17,6 +17,12 @@ class Frame:
         self.frame_nr = frame_nr
         self.garment = garment
 
+    def __hash__(self):
+        return hash((self.id, self.frame_nr))
+
+    def __eq__(self, other):
+        return (self.id, self.frame_nr) == (other.id, other.frame_nr)
+
 
 def mask_faces(F, mask):
     F_mask = np.take(mask, F.reshape(-1), axis=0)