From e5780de5e52347ff21de86c42ec98fbb98e0f6ba Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 28 Dec 2022 19:17:25 +0000 Subject: [PATCH] generate normal image --- data/dataset.py | 25 ++++++++++++++++++------- data/render.py | 7 ++++--- utils/data_utils.py | 14 ++++++++++++++ 3 files changed, 36 insertions(+), 10 deletions(-) diff --git a/data/dataset.py b/data/dataset.py index df95b67..0175c10 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -6,6 +6,7 @@ os.environ["CUDA_VISIBLE_DEVICES"] = '0' import math import time import torch +import kornia import torchvision as tv import numpy as np import scipy.io as sio @@ -17,7 +18,7 @@ from render import Render, DepthRenderer sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from superSMPL.IO import readPC2Frame, writePC2 -from utils.data_utils import Frame, read_frames_file, mask_faces +from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth from cloth3d.read import DataReader from cloth3d.util import quads2tris @@ -71,10 +72,16 @@ class Dataset(Dataset): t = time.perf_counter() - start print(t) start = time.perf_counter() - depth_image, v2i = self.render(combined, combined_faces, info['camLoc']) + depth_image, depth_mask, v2i, K = self.render(combined, combined_faces, info['camLoc']) t = time.perf_counter() - start print(t) + pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True) + pc = pc[0,:,depth_mask].cpu().numpy() + + normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K) + + depth_image = depth_image.cpu().numpy() dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]] depth_image[v2i[idx, 1], v2i[idx, 0]] = 100 @@ -84,12 +91,16 @@ class Dataset(Dataset): import plotly.express as ex import plotly.graph_objects as go - ex.imshow(deodr_depth[:,:,0]).show() - ex.imshow(depth_image).show() + # ex.imshow(deodr_depth[:,:,0]).show() + # ex.imshow(depth_image).show() + ex.imshow(normal_image[0].cpu().numpy()).show() - go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'), - go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), - ]).show() + # go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'), + # go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'), + # ]).show() + + # go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), + # ]).show() # v2i_coverage = np.zeros((480, 640), dtype=np.float32) # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1 diff --git a/data/render.py b/data/render.py index b561c6d..5711023 100644 --- a/data/render.py +++ b/data/render.py @@ -163,13 +163,14 @@ class DepthRenderer(nn.Module): ext = torch.tensor(extrinsic(cam).astype(np.float32), device='cuda') # ext = proj(cam) - fp = np.array([240 * (480 / 36), 320 * (640 / 36)], dtype=np.float32) + # fp = np.array([240 * (480 / 36), 320 * (640 / 36)], dtype=np.float32) img_size = np.array([480, 640], dtype=np.float32) # K = np.eye(4, dtype=np.float32) # K[:3,:3] = intrinsic() + K = torch.tensor(intrinsic()[None, :, :], device='cuda') # cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=1,)# principal_point=fp[None, :], image_size=img_size[None, :]) - cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=torch.tensor(intrinsic()[None, :, :], device='cuda'), image_size=torch.tensor(img_size[None, :], device='cuda')) + cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(img_size[None, :], device='cuda')) raster_settings = RasterizationSettings( image_size=(480, 640), @@ -197,4 +198,4 @@ class DepthRenderer(nn.Module): # import plotly.express as ex # ex.imshow(depth.cpu().numpy()).show() - return depth, v2i + return depth, mask, v2i, K diff --git a/utils/data_utils.py b/utils/data_utils.py index 3316c2e..29092f0 100644 --- a/utils/data_utils.py +++ b/utils/data_utils.py @@ -1,9 +1,12 @@ import os import sys +import kornia +import torch import numpy as np import scipy.io as sio from math import cos, sin + sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from superSMPL.IO import readPC2Frame, readOBJ @@ -34,3 +37,14 @@ def read_frames_file(file): id, frame_nr, garment = line.split() frames.append(Frame(id, int(frame_nr), garment)) return frames + + +def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic): + depth_dilated = kornia.morphology.dilation(depth_image[:,None,:,:], torch.ones((3,3), device=depth_image.get_device()), engine='convolution') + depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution') + depth_dilated = torch.permute(depth_dilated, [0,2,3,1]) + depth_dilated[depth_mask,:] = depth_image[:,:,:,None][depth_mask,:] + depth_dilated = torch.permute(depth_dilated, [0,3,1,2]) + normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True) + normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None] + return normals \ No newline at end of file