diff --git a/data/dataset.py b/data/dataset.py index 2592e18..0c1f8b0 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -192,6 +192,8 @@ class Dataset(Dataset): rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32) + view = torch.tensor(np.concatenate((np.array([info['zrot']], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device) + # image = self.transforms(image) data = { 'mesh_posed': mesh_posed, 'depth_image': depth_image, @@ -201,6 +203,7 @@ class Dataset(Dataset): 'depth_mask_garment': depth_mask_garment, 'depth_mask_body': depth_mask_body, 'rgb_image': rgb_image, + 'view': view, } self.cache[frame] = data diff --git a/models/depth_estim.py b/models/depth_estim.py index beee526..0ee0335 100644 --- a/models/depth_estim.py +++ b/models/depth_estim.py @@ -8,9 +8,26 @@ from torch.nn import functional as F sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from models.convnext import ConvNeXt, CNBlockConfig +class ViewEncoder(nn.Module): + def __init__(self, features=128, device='cuda'): + super().__init__() + self.features: int = features + + self.lin1 = nn.LazyLinear(features).to(device) + self.act1 = nn.GELU().to(device) + self.lin2 = nn.Linear(features, features * 15 * 20).to(device) + + + def forward(self, input: torch.Tensor) -> torch.Tensor: + x: torch.Tensor = self.lin1(input) + x = self.act1(x) + x = self.lin2(x) + x = x.view(x.shape[0],self.features,15,20) + + return x class DepthEstimation(nn.Module): - def __init__(self, device='cuda'): + def __init__(self, view_features=128, device='cuda'): super().__init__() @@ -22,7 +39,7 @@ class DepthEstimation(nn.Module): ] block_setting_reverse = [ - CNBlockConfig(768, 384, 3), + CNBlockConfig(768 + view_features, 384, 3), CNBlockConfig(384, 192, 9), CNBlockConfig(192, 96, 3), CNBlockConfig(96, 27, 3), @@ -34,10 +51,14 @@ class DepthEstimation(nn.Module): self.flatten = nn.Flatten() self.decoder = ConvNeXt(block_setting_reverse, reversed=True, device=device) + self.view_encoder = ViewEncoder(view_features) - def forward(self, x: torch.Tensor) -> torch.Tensor: - x, features = self.convnext(x) + + def forward(self, input) -> torch.Tensor: + image, view = input + x, features = self.convnext(image) + view_features = self.view_encoder(view) # x = self.flatten(x) - x = self.decoder([x, features]) + x = self.decoder([torch.concat((x, view_features), axis=1),features]) return x diff --git a/test.py b/test.py index 7336d95..76c0b7f 100644 --- a/test.py +++ b/test.py @@ -23,7 +23,7 @@ def train_step(batch): model.train() # train mode is on by default after construction # 1. Forward pass - y_pred = model(batch['rgb_image']) + y_pred = model([batch['rgb_image'], batch['view']]) # 2. Calculate loss loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) @@ -45,7 +45,7 @@ def valid_step(batch): model.eval() # put the model in evaluation mode for testing (inference) # 1. Forward pass with torch.inference_mode(): - pred = model(batch['rgb_image']) + pred = model([batch['rgb_image'], batch['view']]) # 2. Calculate the loss loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])