luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

fixed normalization, can overfit on single image

Browse Source
This commit is contained in:
luni
2023-01-24 22:45:14 +00:00
parent a0f3e95255
commit d062c2068d
6 changed files with 147 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+46 -30
View File
@@ -1,7 +1,7 @@
import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
import math
import time
@@ -19,7 +19,7 @@ from torch.utils.data import Dataset
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.render import Render, DepthRenderer
from superSMPL.IO import readPC2Frame, writePC2
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info
from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
from cloth3d.read import DataReader
from cloth3d.util import quads2tris, zRotMatrix
@@ -65,10 +65,11 @@ class Dataset(Dataset):
mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, None, zrot)
mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=self.device)
mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
# raw_mask = sio.loadmat(os.path.join(path, frame.garment + '_mask.mat'))['mask'].reshape((-1))
# mesh_mask = raw_mask > 0 #np.load(os.path.join(path, frame.garment + '_mask.npy'))
# TODO: what's the difference between masking and reading .obj file. is the obj file the not registered version?
faces = mask_faces(self.body_faces, mesh_mask)
# faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
# faces = mask_faces(self.body_faces, mesh_mask)
faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
# body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=True)
body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
@@ -76,17 +77,22 @@ class Dataset(Dataset):
body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=self.device)
body_faces = quads2tris(body_faces)
joints = torch.tensor(joints[0], device=self.device).to(torch.float32)
if self.random_view:
x_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
y_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
z_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi)
view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((x_angle, y_angle, z_angle)))
zRot = torch.tensor(zRotMatrix(-zrot), device=self.device)
mesh_posed = torch.matmul(zRot, mesh_posed.T).T
mesh_posed = mesh_posed - torch.tensor(trans, device=self.device)
body = torch.matmul(zRot, body.T).T
body = body - torch.tensor(trans, device=self.device)
# zRot = torch.tensor(zRotMatrix(-zrot), device=self.device)
# mesh_posed = torch.matmul(zRot, mesh_posed.T).T
# mesh_posed = mesh_posed - torch.tensor(trans, device=self.device)
# body = torch.matmul(zRot, body.T).T
# body = body - torch.tensor(trans, device=self.device)
mesh_posed = mesh_posed - joints[0,:][None,:]
body = body - joints[0,:][None,:]
transform = pytorch3d.transforms.Rotate(view)
mesh_posed = transform.transform_points(mesh_posed)
@@ -95,11 +101,13 @@ class Dataset(Dataset):
# body_no_trans = transform.transform_points(body_no_trans)
zRot = torch.tensor(zRotMatrix(zrot), device=self.device)
mesh_posed = mesh_posed + torch.tensor(trans, device=self.device)
mesh_posed = torch.matmul(zRot, mesh_posed.T).T
body = body + torch.tensor(trans, device=self.device)
body = torch.matmul(zRot, body.T).T
# zRot = torch.tensor(zRotMatrix(zrot), device=self.device)
# mesh_posed = mesh_posed + torch.tensor(trans, device=self.device)
# mesh_posed = torch.matmul(zRot, mesh_posed.T).T
# body = body + torch.tensor(trans, device=self.device)
# body = torch.matmul(zRot, body.T).T
mesh_posed = mesh_posed + joints[0,:][None,:]
body = body + joints[0,:][None,:]
else:
view = torch.eye(3, device=self.device)
@@ -141,27 +149,32 @@ class Dataset(Dataset):
bottom_right[:,1] += 1.25; bottom_right[:,2] -= 1.25
corners = torch.tensor(np.concatenate((top_left, bottom_right), axis=0), device=self.device)
corners_pixel = camera.transform_points_screen(corners)[:,:2].round().to(torch.int64)
root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=self.device))[0,:2].round()
root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=self.device))[0,:2].round().to(torch.int64)
cropper = lambda x: tv.transforms.functional.crop(x, corners_pixel[0,1], corners_pixel[0,0], corners_pixel[1,1] - corners_pixel[0,1], corners_pixel[1,0] - corners_pixel[0,0])
resizer = tv.transforms.Resize((256,256), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
resizer_mask = tv.transforms.Resize((256,256), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
resizer = lambda img, mask_croped, mask_resized: resize_image(img, mask_croped, mask_resized, (256, 256))
# bbox_before = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:]
depth_mask_garment = cropper(depth_mask_garment)
scale = depth_mask_garment.shape[0] / 256
depth_mask_garment = resizer(depth_mask_garment[None,None,:,:])[0,0,:,:]
depth_mask_garment = resizer_mask(depth_mask_garment[None,None,:,:])[0,0,:,:]
original = depth_image
original_mask = depth_mask
depth_mask = resizer(cropper(depth_mask)[None,None,:,:])[0,0,:,:]
scaled = resizer(cropper(depth_image)[None,None,:,:])[0,0,:,:] * depth_mask
normal_image = resizer(cropper(normal_image)[None,:,:,:])[0,:,:,:] * depth_mask[None,:,:]
depth_mask_original = depth_mask
depth_mask_croped = cropper(depth_mask)
depth_mask = resizer_mask(depth_mask_croped[None,None,:,:])[0,0,:,:]
scaled = resizer(cropper(depth_image)[None,None,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,0,:,:]
normal_image = resizer(cropper(normal_image)[None,:,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,:,:,:]
# bbox_after = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:]
# scale = (bbox_before[2] - bbox_before[0]) / (bbox_after[2] - bbox_after[0])
# root_dist = torch.tensor(np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2)), device=self.device)
root_dist = torch.tensor(np.abs(root_joint[0][0] - info['camLoc'][0]), device=self.device)
root_dist = torch.tensor(np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2)), device=self.device).to(torch.float32)
# root_dist = torch.tensor(np.abs(root_joint[0][0] - info['camLoc'][0]), device=self.device).to(torch.float32)
# scaled = root_dist - scaled
scaled = scaled - root_dist
@@ -268,7 +281,7 @@ class Dataset(Dataset):
# ex.imshow(v2i_coverage).show()
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(torch.float32).to(self.device)
rgb_image = resizer(cropper(rgb_image)[None,:,:,:])[0,:,:,:] * depth_mask[None,:,:]
rgb_image = resizer(cropper(rgb_image)[None,:,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,:,:,:]
# view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), trans.astype(np.float32)), axis=0), device=self.device)
@@ -278,8 +291,7 @@ class Dataset(Dataset):
data = { 'mesh_posed': mesh_posed,
'depth_image': scaled,
# 'depth_image': upscaled,
'mesh_mask': mesh_mask,
'depth_mask': depth_mask, #depth_mask, depth_mask_garment
'depth_mask': depth_mask_garment, #depth_mask, depth_mask_garment
'depth_mask_full': depth_mask, #depth_mask,
'normal_image': normal_image,
# 'depth_mask_garment': depth_mask_garment,
@@ -290,10 +302,14 @@ class Dataset(Dataset):
'root_dist': root_dist,
'scale': scale,
'root_joint_pixel': root_joint_pixel,
'joints': joints,
'original': original,
'original_mask': original_mask,
}
self.cache[frame] = data
return self.cache[frame]
if self.cache is not None:
self.cache[frame] = data
return data
if __name__ == '__main__':
@@ -302,8 +318,8 @@ if __name__ == '__main__':
device = "cuda" if torch.cuda.is_available() else "cpu"
# dataset = Dataset('/home/workspace/local-data/train')
dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_single.txt', random_view=False)
train_dataloader = DataLoader(dataset, batch_size=1, drop_last=True, shuffle=False)
dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True)
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
for batch in train_dataloader:
batch
data/preprocess.py
+2 -1
View File
@@ -9,6 +9,7 @@ from tqdm import tqdm
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from utils.data_utils import mask_faces, Frame
from superSMPL.IO import readPC2, writePC2
from cloth3d.util import loadInfo
def generate_faces(folder, frames):
@@ -71,7 +72,7 @@ def genereated_registered_and_mask(folder, is_skirt = False):
if __name__ == '__main__':
# folder = '/home/workspace/local-data/train'
# folder = '/home/workspace/local-data/valid'
folder = 'dataset/CLOTH3D/train'
folder = '/home/workspace/data/dataset/cloth3d/train'
# generate_faces(folder)
# genereated_registered_and_mask(folder)
data/render.py
+2 -2
View File
@@ -146,13 +146,13 @@ class DepthRenderer(nn.Module):
raster_settings=raster_settings
)
output = rasterizer(Meshes(verts[None,:,:], faces[None,:,:]))
v2i = cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
v2i = None #cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
depth = output.zbuf[0,:,:,0]
pix_to_face = output.pix_to_face[0,:,:,0]
mask = depth != -1
# dmin = depth.min()
dmean = depth[mask].mean()
# dmean = depth[mask].mean()
# dmax = depth[mask].max()
# depth[mask] = depth[mask] - dmean
# depth[~mask] = 0.0