From 93357f81cba271e0f9511495a2224c3b31b10907 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 6 Jan 2023 22:49:57 +0000 Subject: [PATCH] trying to get basic model working --- data/dataset.py | 13 ++++- models/convnext.py | 110 +++++++++++++++++++++++++++++++++++++++--- models/depth_estim.py | 13 ++++- test.py | 16 ++++-- utils/data_utils.py | 8 ++- 5 files changed, 143 insertions(+), 17 deletions(-) diff --git a/data/dataset.py b/data/dataset.py index 674ebbc..8f4b8ed 100644 --- a/data/dataset.py +++ b/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 diff --git a/models/convnext.py b/models/convnext.py index 9a539b9..4d7b408 100644 --- a/models/convnext.py +++ b/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) diff --git a/models/depth_estim.py b/models/depth_estim.py index bbd3c3b..92cdcd2 100644 --- a/models/depth_estim.py +++ b/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 diff --git a/test.py b/test.py index b62a290..8b08a9a 100644 --- a/test.py +++ b/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) diff --git a/utils/data_utils.py b/utils/data_utils.py index 91621e7..1ef5999 100644 --- a/utils/data_utils.py +++ b/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)