Skip to content
GitLab
Commit 7f425400 authored by 9wei's avatar 9wei
Browse files

Merge branch 'refactor' into 'main' 

Refactor

See merge request !2
parents 8cae70e8 7b0e8af2
Hide whitespace changes
Inline Side-by-side
.gitignore
View file @ 7f425400
/mydataset/FPS1024
/mydataset/ycb_video_obj_ply
/mydataset/my_train
/mydataset/my_val
/mydataset/my_test
/.dvc/cache
/.dvc/tmp
/.dvc/config.local
/log
__pycache__
.pre-commit-config.yaml 0 → 100644
View file @ 7f425400
# See https://pre-commit.com for more information
# See https://pre-commit.com/hooks.html for more hooks
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.0.1
hooks:
- id: trailing-whitespace
- id: end-of-file-fixer
- id: check-yaml
- id: check-json
- id: check-added-large-files
- repo: https://github.com/adrienverge/yamllint
rev: v1.26.3
hooks:
- id: yamllint
name: Lint GitHub Action workflow files
files: ^\.github/workflows/
args:
- -d
- >
{
extends: default,
rules: {
document-start:{
present: true,
level: warning
},
line-length: {
max: 88,
level: warning
},
indentation: {
spaces: 2,
indent-sequences: false,
level: warning
}
}
}
- repo: https://github.com/zricethezav/gitleaks
rev: v7.6.1
hooks:
- id: gitleaks
name: 🔍 Check git leaks
- repo: https://github.com/jumanjihouse/pre-commit-hooks
rev: 2.1.5
hooks:
- id: shellcheck
- id: shfmt
- repo: https://github.com/psf/black
rev: 21.11b1
hooks:
- id: black
name: 🐍 Format Python code PEP8
args:
- --line-length=88
- repo: https://gitlab.com/pycqa/flake8
rev: 3.9.2
hooks:
- id: flake8
name: 🐍 Check Python code
args:
- --max-line-length=88
- --ignore=F841,W503
- --exclude=train.py
- repo: https://github.com/pycqa/isort
rev: 5.10.1
hooks:
- id: isort
name: 🐍 Sort Python imports
# - repo: https://github.com/asottile/pyupgrade
# rev: v2.29.1
# hooks:
# - id: pyupgrade
# name: 🐍 Upgrade Python code
# args: [--py36-plus]
# Black and flake8 for notebook
- repo: https://github.com/nbQA-dev/nbQA
rev: 1.1.1
hooks:
- id: nbqa-black
- id: nbqa-flake8
- id: nbqa-isort
# Lint Dockerfile
- repo: https://github.com/hadolint/hadolint
rev: v2.8.0
hooks:
- id: hadolint
.vscode/settings.json 0 → 100644
View file @ 7f425400
{
"python.linting.pylintEnabled": false,
"python.linting.flake8Enabled": true,
"python.linting.enabled": true
}
Notebooks/data-explore.ipynb 0 → 100644
View file @ 7f425400
%% Cell type:code id: tags:
``` python
import os
import numpy as np
import tensorflow as tf
import tensorflow_graphics.geometry.transformation as tfgt
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
```
%% Cell type:code id: tags:
``` python
train_filenames = [
"../mydataset/FPS1024/train_files_FPS1024_0_0.tfrecords",
"../mydataset/FPS1024/train_files_FPS1024_0_1.tfrecords",
]
tr_dataset = tf.data.TFRecordDataset(train_filenames).shuffle(1000000)
```
%% Cell type:code id: tags:
``` python
def decode(serialized_example, total_num_point):
features = tf.io.parse_example(
[serialized_example],
features={
"xyz": tf.io.FixedLenFeature([total_num_point, 3], tf.float32),
"rgb": tf.io.FixedLenFeature([total_num_point, 3], tf.float32),
"translation": tf.io.FixedLenFeature([3], tf.float32),
"quaternion": tf.io.FixedLenFeature([4], tf.float32),
"num_valid_points_in_segment": tf.io.FixedLenFeature([], tf.int64),
"seq_id": tf.io.FixedLenFeature([], tf.int64),
"frame_id": tf.io.FixedLenFeature([], tf.int64),
"class_id": tf.io.FixedLenFeature([], tf.int64),
},
)
return features
BATCH_SIZE = 1
TOTAL_NUM_POINT = 1024
tr_dataset = tr_dataset.map(lambda x: decode(x, TOTAL_NUM_POINT))
```
%% Cell type:code id: tags:
``` python
print(tr_dataset)
```
%% Cell type:code id: tags:
``` python
element = iter(tr_dataset).get_next()
```
%% Cell type:code id: tags:
``` python
# seg = tf.concat(values=[element.get("xyz"), element.get("rgb")], axis=0)
seg = tf.squeeze(element.get("xyz"))
seg.shape
```
%% Cell type:markdown id: tags:
%% Cell type:code id: tags:
``` python
rotation_matrix = tfgt.rotation_matrix_3d.from_quaternion(
element.get("quaternion"),
)
```
%% Cell type:code id: tags:
``` python
# Flatten tensor to 1D array
pos = tf.concat(
[
tf.reshape(rotation_matrix, -1),
tf.reshape(element.get("translation"), -1),
tf.constant([0, 0, 0, 1], dtype=tf.float32),
tf.reshape(tf.cast(element.get("class_id"), tf.float32), -1),
],
axis=0,
)
```
%% Cell type:code id: tags:
``` python
np.save("../mydataset/my_train/0000_pos.npy", pos.numpy())
```
%% Cell type:code id: tags:
``` python
np.save("../mydataset/my_train/0000_seg.npy", seg.numpy())
```
%% Cell type:markdown id: tags:
losses.py
View file @ 7f425400
import torch
CLASS_NAME = ["cl0", "cl1", "cl2", "cl3", "cl4"]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def skew_symmetric(axag_unit):
"""
Create the skew symmetric matrix for the input vector
......@@ -20,7 +22,7 @@ def skew_symmetric(axag_unit):
row1 = torch.cat(
[
torch.zeros((sh[0], 1, 1), dtype=torch.float64).cuda(),
torch.zeros((sh[0], 1, 1), dtype=torch.float64).to(device),
-axag_unit_exp[:, 2, :, :],
axag_unit_exp[:, 1, :, :],
],
......@@ -29,7 +31,7 @@ def skew_symmetric(axag_unit):
row2 = torch.cat(
[
axag_unit_exp[:, 2, :, :],
torch.zeros((sh[0], 1, 1), dtype=torch.float64).cuda(),
torch.zeros((sh[0], 1, 1), dtype=torch.float64).to(device),
-axag_unit_exp[:, 0, :, :],
],
dim=2,
......@@ -39,7 +41,7 @@ def skew_symmetric(axag_unit):
[
-axag_unit_exp[:, 1, :, :],
axag_unit_exp[:, 0, :, :],
torch.zeros((sh[0], 1, 1), dtype=torch.float64).cuda(),
torch.zeros((sh[0], 1, 1), dtype=torch.float64).to(device),
],
dim=2,
)
......@@ -93,7 +95,7 @@ def exponential_map(axag, EPS=1e-2):
torch.eye(3, dtype=torch.float64)
.unsqueeze(0)
.repeat(axag.shape[0], 1, 1)
.cuda()
.to(device)
)
axag_exp = (
batch_identity
......@@ -129,14 +131,14 @@ def logarithm(R, b_deal_with_sym=False, EPS=1e-2):
"""
B = R.shape[0]
trace_all = torch.zeros(B, dtype=torch.float64).cuda()
trace_all = torch.zeros(B, dtype=torch.float64).to(device)
for iii in range(B):
trace = torch.trace(R[iii, :, :].squeeze(0))
trace_all[iii] = trace
trace_temp = (trace_all - 1) / 2
# take the safe acos
o_n_e = torch.ones((B,), dtype=torch.float64).cuda()
o_n_e = torch.ones((B,), dtype=torch.float64).to(device)
trace_temp = clip_by_tensor(trace_temp, -o_n_e, o_n_e)
theta = torch.acos(trace_temp)
......@@ -216,8 +218,10 @@ def get_loss(end_points):
if __name__ == "__main__":
label = torch.tensor([[0.6977, 0.8248, 0.9367]], dtype=torch.float64).cuda()
pred = torch.tensor([[-2.100418, -2.167796, 0.2733]], dtype=torch.float64).cuda()
label = torch.tensor([[0.6977, 0.8248, 0.9367]], dtype=torch.float64).to(device)
pred = torch.tensor([[-2.100418, -2.167796, 0.2733]], dtype=torch.float64).to(
device
)
# print(torch.matmul(pred, label))
print(get_rotation_error(pred, label))
# import cv2
......
mydataset/dataset_converter.py 0 → 100644
View file @ 7f425400
import glob
import os
from typing import List
import numpy as np
import tensorflow as tf
import tensorflow_graphics.geometry.transformation as tfgt
class TFRecord2NumPy:
"""Convert TF record Dataset to NumPy array."""
# @tf.function(experimental_follow_type_hints=True)
def __init__(self, dataset: tf.data.TFRecordDataset):
self.dataset = dataset.map(lambda x: self._decode(x, 1024))
def _decode(self, serialized_example, total_num_point: int):
features = tf.io.parse_example(
[serialized_example],
features={
"xyz": tf.io.FixedLenFeature([total_num_point, 3], tf.float32),
"rgb": tf.io.FixedLenFeature([total_num_point, 3], tf.float32),
"translation": tf.io.FixedLenFeature([3], tf.float32),
"quaternion": tf.io.FixedLenFeature([4], tf.float32),
"num_valid_points_in_segment": tf.io.FixedLenFeature([], tf.int64),
"seq_id": tf.io.FixedLenFeature([], tf.int64),
"frame_id": tf.io.FixedLenFeature([], tf.int64),
"class_id": tf.io.FixedLenFeature([], tf.int64),
},
)
return features
def _set_rotation_matrix(self, element) -> np.ndarray:
rot = tfgt.rotation_matrix_3d.from_quaternion(
element["quaternion"],
)
assert rot.shape == (1, 3, 3)
return rot
def convert_to_seg_npy(self, element) -> np.ndarray:
seg = tf.squeeze(element["xyz"]).numpy()
assert seg.shape == (1024, 3)
return seg
def convert_to_pos_npy(self, element) -> np.array:
rotation_matrix = self._set_rotation_matrix(element)
pos = tf.concat(
[
tf.reshape(rotation_matrix, -1),
tf.reshape(element.get("translation"), -1),
tf.constant([0, 0, 0, 1], dtype=tf.float32),
tf.reshape(tf.cast(element.get("class_id"), tf.float32), -1),
],
axis=0,
).numpy()
assert pos.shape == (17,)
return pos
def save_npy(
self, seg: np.ndarray, pos: np.array, dir_path: str, file_basename: int
):
np.save(os.path.join(dir_path, f"{file_basename:05d}_seg"), seg)
np.save(os.path.join(dir_path, f"{file_basename:05d}_pos"), pos)
def prepare_file_paths(dir_name: str) -> List[str]:
"Return a list of *.tfrecords files"
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
DIR_NAME = "FPS1024"
FILE_PATHS = list(glob.iglob(os.path.join(ROOT_DIR, DIR_NAME, "*.tfrecords")))
return FILE_PATHS
if __name__ == "__main__":
# Env needs for CPU
if len(tf.config.list_physical_devices("GPU")) == 0:
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
TRAIN_FILES = prepare_file_paths("FPS1024")
tr_dataset = tf.data.TFRecordDataset(TRAIN_FILES).shuffle(100000)
data_converter = TFRecord2NumPy(tr_dataset)
ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
TRAIN_DIR = os.path.join(ROOT_DIR, "my_train")
for count, element in enumerate(data_converter.dataset):
if count > 11000:
break
pos = data_converter.convert_to_pos_npy(element)
seg = data_converter.convert_to_seg_npy(element)
data_converter.save_npy(
seg=seg,
pos=pos,
dir_path=TRAIN_DIR,
file_basename=count,
)
mydataset/model_util_my → mydataset/model_util_my.py
View file @ 7f425400
# Copyright (c) Facebook, Inc. and its affiliates.
#
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import numpy as np
import open3d as o3d
import sys
import os
import sys
import numpy as np
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
ROOT_DIR = os.path.dirname(BASE_DIR)
sys.path.append(os.path.join(ROOT_DIR, 'utils'))
import pc_util
sys.path.append(os.path.join(ROOT_DIR, "utils"))
class myDatasetConfig(object):
def __init__(self):
self.num_class = 5
self.num_class = 18
self.type2class = {'cl0':0, 'cl1':1, 'cl2':2, 'cl3':3, 'cl4':4}
self.type2class = {f"cl{i}": i for i in range(self.num_class)}
self.class2type = {self.type2class[t]: t for t in self.type2class}
self.type2onehotclass = {'cl0': np.array([1, 0, 0, 0, 0]),
'cl1': np.array([0, 1, 0, 0, 0]),
'cl2': np.array([0, 0, 1, 0, 0]),
'cl3': np.array([0, 0, 0, 1, 0]),
'cl4': np.array([0, 0, 0, 0, 1])}
# 2D array
self.onehot_encoding = np.eye(self.num_class)[
np.array([range(self.num_class)]).reshape(-1)
]
def sem2class(self, cls):
onehot = self.type2onehotclass[self.class2type[cls]]
# Select ith row of the 2D array
onehot = self.onehot_encoding[int(cls), :]
return onehot
def size2class(self, type_name):
''' Convert 3D box size (l,w,h) to size class and size residual '''
"""Convert 3D box size (l,w,h) to size class and size residual"""
size_class = self.type2class[type_name] # 0
# size_residual = size - self.type_mean_size[type_name] # 尺寸
return size_class
def class2size(self, pred_cls):
''' Inverse function to size2class '''
"""Inverse function to size2class"""
mean_size = self.type_mean_size[self.class2type[pred_cls]]
return mean_size
def class2sem(self, pred_cls):
sem = self.class2type[pred_cls]
return sem