Commit 7f425400 authored by 9wei's avatar 9wei
Browse files

Merge branch 'refactor' into 'main'

Refactor

See merge request !2
parents 8cae70e8 7b0e8af2
/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__
# 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
{
"python.linting.pylintEnabled": false,
"python.linting.flake8Enabled": true,
"python.linting.enabled": true
}
%% 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:
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
......
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,
)
# 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