Skip to content

Commit

Permalink
rm preprocess
Browse files Browse the repository at this point in the history
  • Loading branch information
JudyYe committed Jun 16, 2022
1 parent bb8398d commit dcfdd56
Show file tree
Hide file tree
Showing 9 changed files with 112 additions and 276 deletions.
13 changes: 12 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -77,4 +77,15 @@ If you use find this code helpful, please consider citing:
+ [ ] test time refinement
+ [ ] predicted hand eval
+ [ ] add more demo images
+ [ ] models
+ [ ] models
### Acknowledgement
The repository is inspired from, The authors would like to thank the authors inspired from
- [FrankMocap](https://github.com/facebookresearch/frankmocap/blob/main/docs/INSTALL.md).
- [Phosa]()
- [DeepSDF]()
- [GraspingField]()
- [Mano]()
and many others...
2 changes: 1 addition & 1 deletion datasets/base_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def preload_anno(self, load_keys=['index']):
raise NotImplementedError

def get_sdf_files(self, cad_idx):
sdf_dir = osp.join(self.cfg.DB.DIR, 'sdf/SdfSamples/', self.dataset, 'all')
sdf_dir = osp.join(self.cfg.DB.DIR, 'mesh_sdf/SdfSamples/', self.dataset, 'all')
filename = osp.join(sdf_dir, cad_idx + '.npz')
assert osp.exists(filename), 'Not exists %s' % filename
return filename
Expand Down
8 changes: 5 additions & 3 deletions datasets/ho3d.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,14 +36,16 @@ def __init__(self, cfg, dataset: str, split='val', is_train=True,

if self.use_gt == 'none':
if not is_train:
# hand pose from frankmocap prediction
meta_folder = 'meta_plus'
else:
# hand pose from GT
meta_folder = 'meta_gt'
else:
meta_folder = 'meta_%s' % self.use_gt

self.cache_file = osp.join(self.data_dir, 'Cache', '%s_%s_%s.pkl' % (dataset, self.split, self.use_gt))
self.cache_mesh = osp.join(self.data_dir, 'Cache', '%s_%s_mesh.pkl' % (dataset, self.split))
self.cache_file = osp.join(osp.dirname(self.data_dir), 'cache', '%s_%s_%s.pkl' % (dataset, self.split, self.use_gt))
self.cache_mesh = osp.join(osp.dirname(self.data_dir), 'cache', '%s_%s_mesh.pkl' % (dataset, self.split))
self.mask_dir = ''
self.meta_dir = os.path.join(self.data_dir, '{}', '{}', meta_folder, '{}.pkl')
self.image_dir = osp.join(self.data_dir, '{}', '{}', 'rgb', '{}.jpg')
Expand All @@ -55,7 +57,6 @@ def preload_anno(self, load_keys=[]):
print('!! Load from cache !!', self.cache_file)
self.anno = pickle.load(open(self.cache_file, 'rb'))
else:
print(self.meta_dir)
print('creating cahce', self.meta_dir)
# filter 1e-3
df = pd.read_csv(osp.join(self.data_dir, '%s%s.csv' % (self.split, self.suf)))
Expand Down Expand Up @@ -140,6 +141,7 @@ def get_cam(self, idx):
return self.anno['cam'][idx]

def get_obj_mask(self, index):
"""fake mask"""
image = np.array(Image.open(self.image_dir.format(*index)))
H, W, _= image.shape
mask = Image.fromarray(np.ones([H, W]).astype(np.uint8) * 255 )
Expand Down
4 changes: 2 additions & 2 deletions datasets/mow.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,8 @@ def __init__(self, cfg, dataset: str, split='val', is_train=True,
if split == 'val':
self.split = 'test'
self.set_dir = osp.join(self.data_dir, '{}.lst')
self.cache_file = osp.join(self.data_dir, 'Cache', '%s_%s.pkl' % (dataset, self.split))
self.cache_mesh = osp.join(self.data_dir, 'Cache', '%s_%s_mesh.pkl' % (dataset, self.split))
self.cache_file = osp.join(osp.dirname(self.data_dir), 'cache', '%s_%s.pkl' % (dataset, self.split))
self.cache_mesh = osp.join(osp.dirname(self.data_dir), 'cache', '%s_%s_mesh.pkl' % (dataset, self.split))
self.mask_dir = osp.join(self.data_dir, 'results/{0}/{0}_mask.png')
self.image_dir = osp.join(self.data_dir, 'results/{0}/{0}.jpg')
self.shape_dir = osp.join(self.data_dir, 'results/{0}/{0}_norm.obj')
Expand Down
26 changes: 16 additions & 10 deletions datasets/obman.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import tqdm
from PIL import Image
from nnutils import mesh_utils, geom_utils

from nnutils.hand_utils import extract_rt_pose

class Obman(BaseData):
def __init__(self, cfg, dataset: str, split='val', is_train=True,
Expand All @@ -32,8 +32,8 @@ def __init__(self, cfg, dataset: str, split='val', is_train=True,
else:
self.cls = ''

self.cache_file = osp.join(self.data_dir, 'Cache', '%s_%s.pkl' % (dataset, self.split))
self.cache_mesh = osp.join(self.data_dir, 'Cache', '%s_%s_mesh.pkl' % (dataset, self.split))
self.cache_file = osp.join(osp.dirname(self.data_dir), 'cache', '%s_%s.pkl' % (dataset, self.split))
self.cache_mesh = osp.join(osp.dirname(self.data_dir), 'cache', '%s_%s_mesh.pkl' % (dataset, self.split))

self.shape_dir = os.path.join(self.cfg.DB.DIR, 'obmanboj', '{}', '{}', 'models', 'model_normalized.obj')
self.image_dir = osp.join(self.data_dir, split, 'rgb', '{}.jpg')
Expand All @@ -45,10 +45,10 @@ def preload_anno(self, load_keys=[]):
self.anno[key] = []

if self.cache and osp.exists(self.cache_file):
print('!! Load from cache !!')
print('!! Load from cache !!', self.cache_file)
self.anno = pickle.load(open(self.cache_file, 'rb'))
else:

print('making cache in %s' % self.cache_file)
index_list = [line.strip() for line in open(osp.join(self.data_dir, '%s.txt' % self.split))]
for i, index in enumerate(tqdm.tqdm(index_list)):
dset = self.cfg.DB.NAME
Expand All @@ -59,16 +59,21 @@ def preload_anno(self, load_keys=[]):
with open(meta_path, "rb") as meta_f:
meta_info = pickle.load(meta_f)

global_rt, art_pose, obj_pose = extract_rt_pose(meta_info, self.pose_wrapper)
meta_info['cTh'] = global_rt
meta_info['hA'] = art_pose
meta_info['cTo'] = obj_pose

self.anno['index'].append(index)
self.anno['cad_index'].append(osp.join(meta_info["class_id"], meta_info["sample_id"]))
cTo = torch.FloatTensor([meta_info['cTo']])
cTh = torch.FloatTensor([meta_info['cTh']])
cTo = meta_info['cTo']
cTh = meta_info['cTh']
hTc = geom_utils.inverse_rt(mat=cTh, return_mat=True)
hTo = torch.matmul(hTc, cTo)

self.anno['cTh'].append(cTh[0])
self.anno['hTo'].append(hTo[0])
self.anno['hA'].append(meta_info['hA'])
self.anno['hA'].append(meta_info['hA'].detach().numpy()[0])

os.makedirs(osp.dirname(self.cache_file), exist_ok=True)
print('save cache')
Expand All @@ -78,11 +83,11 @@ def preload_anno(self, load_keys=[]):

def preload_mesh(self):
if self.cache and osp.exists(self.cache_mesh):
print('!! Load from cache !!')
print('!! Load from cache !!', self.cache_mesh)
self.obj2mesh = pickle.load(open(self.cache_mesh, 'rb'))
else:
self.obj2mesh = {}
print('load mesh')
print('load mesh', self.cache_mesh)
for i, cls_id in tqdm.tqdm(enumerate(self.anno['cad_index']), total=len(self.anno['cad_index'])):
key = cls_id
cls, id = key.split('/')
Expand Down Expand Up @@ -139,3 +144,4 @@ def __getitem__(self, idx):
sample = {key: self.anno[key][idx] for key in self.anno}
sample['mesh'] = self.obj2mesh[sample['cad_index']]
return sample

19 changes: 13 additions & 6 deletions docs/preprocess.md
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
# Preprocess Data

Data preprocessing consists of two parts, sdf preprocessing and camera-hand-object transformation preprcoessing.

## Folder structure
```
data/
sdf/
cache/
mesh_sdf/
SdfSamples/
obman/all/
04074963/4e73215ae0f33d23a5e3ac6ff4952f3.npz
Expand All @@ -12,15 +15,19 @@ data/
mow/all/
gardening_v_d79ApGisN_Y_frame000046.npzw
obman/
obmanobj/
(obmanobj/)
04530566/72f4c3c433492d585001cb19c4a0eee4/models/model_normalized.obj
ho3d/
ho3dobj/models/
(ho3dobj/models/)
mow/
```

## use our cached data
(Coming soon)
## Use our preprocessed data
We provide our preprocessed `cache` and `mesh_sdf`. These can be downloaded from [here](). You need to unzip it and put it under `data/` folder.


## create cached data
## Create your own cache data
```
python preprocess/generate_sdf.py
```
(Coming soon)
61 changes: 60 additions & 1 deletion nnutils/hand_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# Written by Yufei Ye (https://github.com/JudyYe)
# --------------------------------------------------------
from __future__ import print_function
import numpy as np
import pickle
import os.path as osp
from typing import Tuple
Expand Down Expand Up @@ -264,4 +265,62 @@ def get_nTh(hand_wrapper, hA, r, inverse=False, center=None):

if inverse:
mat = geom_utils.inverse_rt(mat=mat, return_mat=True)
return mat
return mat



def extract_rt_pose(meta_info, pose_wrapper):
device = 'cpu'
cam_extr = np.array(
[
[1.0, 0.0, 0.0, 0.0],
[0.0, -1.0, 0.0, 0.0],
[0.0, 0.0, -1.0, 0.0],
[0, 0, 0, 1]
]
).astype(np.float32)
cTw = Transform3d(device=device, matrix=torch.FloatTensor([cam_extr]).transpose(1, 2).to(device))
wTo = Transform3d(device=device, matrix=torch.FloatTensor([meta_info['affine_transform']]).transpose(1, 2).to(device))
cTo = wTo.compose(cTw)
cTo_mat = cTo.get_matrix().transpose(-1, -2).cpu()

# src mesh
zeros = torch.zeros([1, 3], device=device, dtype=torch.float32)
art_pose = torch.FloatTensor([meta_info['pca_pose']]).to(device)
art_pose = pose_wrapper.pca_to_pose(art_pose) - pose_wrapper.hand_mean
hHand, _ = pose_wrapper(None, art_pose, zeros, zeros, mode='inner')
# dst mesh
wVerts = torch.FloatTensor([meta_info['verts_3d']]).to(device)
textures = TexturesVertex(torch.ones_like(wVerts)).to(device)
wHand = Meshes(wVerts, pose_wrapper.hand_faces, textures)

wTh = solve_rt(hHand.verts_padded(), wHand.verts_padded())
cTh = wTh.compose(cTw)
cTh_mat = cTh.get_matrix().transpose(-1, -2).cpu()

return cTh_mat, art_pose.cpu(), cTo_mat


def solve_rt(src_mesh, dst_mesh):
"""
(N, P, 3), (N, P, 3)
"""
device = src_mesh.device
src_centroid = torch.mean(src_mesh, -2, keepdim=True)
dst_centroid = torch.mean(dst_mesh, -2, keepdim=True)
src_bar = (src_mesh - src_centroid)
dst_bar = (dst_mesh - dst_centroid)
cov = torch.bmm(src_bar.transpose(-1, -2), dst_bar)
u, s, v = torch.svd(cov)
vh = v.transpose(-1, -2)
rot_t = torch.matmul(u, vh)
rot = rot_t.transpose(-1, -2) # v, uh

trans = dst_centroid - torch.matmul(src_centroid, rot_t) # (N, 1, 3)?

rot = Rotate(R=rot_t, device=device)
trans = Translate(trans.squeeze(1), device=device)

rt = rot.compose(trans)

return rt
4 changes: 3 additions & 1 deletion nnutils/mesh_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
import numpy as np
from scipy.spatial.distance import cdist
import re
from torch._six import container_abcs, string_classes, int_classes
import collections.abc as container_abcs
from torch._six import string_classes
int_classes = int

import trimesh
import torch
Expand Down
Loading

0 comments on commit dcfdd56

Please sign in to comment.