release code

Author: liuyulinn

Diff for: .DS_Store

@@ -0,0 +1,172 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+downloads/
+eggs/
+.eggs/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+exps_*
+target_fn/
+.vscode/
+log.txt
+datasets
+tv.txt
+*.png
+plot/
+data/
+target_fn/
+flow/mobiusflow_nref.py
+temp.md
+
+ip.py
+target_fn_bak/
+name/
+*.npy
+/output
+/run
+/nflows
+*.txt
+temp*
+result/
+Public_prob_orientaion_ePublic_prob_orientation_estimation_with_matrix_fisher_distributions/
+tmp/
+
+plot/
+*.pth
+
+data/
+Matrixfisher/storage
+.ssh/
+*.pkl
+baseline/
+cluster/

Diff for: .gitignore

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 雨霖霖
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Diff for: LICENSE

@@ -0,0 +1,132 @@
+import lmdb
+import os
+import io
+import pickle
+from PIL import Image
+from geometric_utils import numpy_euler_to_rotation, numpy_quaternion_to_rot_mat, numpy_aet_to_rot
+import numpy as np
+import matplotlib.pyplot as plt
+from enum import IntEnum
+
+dataset_dir = 'dataset' # change to where you store datasets
+
+test_name = 'test_20V.Rawjpg.lmdb'
+train_name = 'train_100V.Rawjpg.lmdb'
+
+class ModelNetSo3():
+    def __init__(self, path=dataset_dir):
+        self.modelnet_dir = os.path.join(path, 'ModelNet10-SO3')
+
+    def get_train(self):
+        full = ModelNetSo3Subset(os.path.join(self.modelnet_dir, train_name)) # split TODO
+        return full
+
+    def get_eval(self):
+        return ModelNetSo3Subset(os.path.join(self.modelnet_dir, test_name))
+
+    def get_cad(self, class_idx, cad_idx):
+        return np.array([[0.0,0,0], [1,0,0], [0,1,0], [0,0,1]]), None
+
+class ModelNetSo3Classes(IntEnum):
+    BATHTUB=0
+    BED=1
+    CHAIR=2
+    DESK=3
+    DRESSER=4
+    MONITOR=5
+    NIGHT_STAND=6
+    SOFA=7
+    TABLE=8
+    TOILET=9
+
+    def __str__(self):
+        return self.name.lower()
+
+
+modelnetso3_str_enum_map = {}
+for v in ModelNetSo3Classes:
+    modelnetso3_str_enum_map[str(v)] = v
+
+class ModelNetSo3Subset():
+    def __init__(self, path, imsize=224):
+        self.path = path
+        lmdb_db = lmdb.open(path)
+        quaternions_path = os.path.join(path, 'viewID2quat.pkl')
+        euler_path = os.path.join(path, 'viewID2euler.pkl')
+        with open(quaternions_path, 'rb') as f:
+            quaterions = pickle.load(f, encoding='latin1')
+        with open(euler_path, 'rb') as f:
+            euler = pickle.load(f, encoding='latin1')
+        if os.path.exists(self.key_path()):
+            with open(self.key_path(), 'r') as f:
+                keys = [x.strip() for x in f.readlines()]
+        else:
+            with lmdb_db.begin() as txn:
+                keys = [key.decode('utf-8') for key, _ in txn.cursor()]
+            with open(self.key_path(), 'w') as f:
+                f.write('\n'.join(keys))
+
+        self.imsize = imsize
+        self.lmdb = lmdb_db
+        self.keys = keys
+        self.quaterions = quaterions
+        self.euler = euler
+
+    def key_path(self):
+        return os.path.join(self.path, 'keys.txt')
+
+    def __len__(self):
+        return len(self.keys)
+
+    def __getitem__(self, idx):
+        key = self.keys[idx]
+        cad_id, view_id = key.split('.')
+        class_str = cad_id[:cad_id.rfind('_')]
+        
+        class_idx = modelnetso3_str_enum_map[class_str]
+        with self.lmdb.begin() as txn:
+            bits = txn.get(key.encode('utf8'))
+        f = io.BytesIO(bits)
+        image = Image.open(f)
+        image = image.resize((self.imsize, self.imsize), Image.ANTIALIAS)
+        image = np.array(image.getdata()).reshape(self.imsize, self.imsize, 1).repeat(3,2)
+        image = image.astype(np.float32) / 255
+        image = image.transpose(2,0,1)
+        q = self.quaterions[key]
+        qx = np.array([q[1], -q[2], q[3], q[0]])
+        R = numpy_quaternion_to_rot_mat(qx)
+        extrinsic = np.eye(4)
+        extrinsic[:3,:3] = R
+        extrinsic[:3,3] = np.array([0,0,5])
+        intrinsic = np.array([[self.imsize/2, 0, self.imsize/2],
+                              [0, self.imsize/2, self.imsize/2],
+                              [0,0,1]])
+        extrinsic=extrinsic.astype(np.float32)
+        return image, extrinsic, class_idx, 0, intrinsic, 0
+
+
+
+
+def main():
+    dataset = ModelNetSo3()
+    tr = dataset.get_train()
+    print(len(tr))
+    exit(0)
+    for i in range(50000):
+        image, extrinsic, class_idx, hard, intrinsic, cad_idx = tr[i]
+        nodes, _ = dataset.get_cad(class_idx, cad_idx)
+        nodes_homo = np.ones((4, len(nodes)))
+        nodes_homo[:3, :] = nodes.transpose()
+        model = np.matmul(extrinsic, nodes_homo)
+        model = model[:3, :]
+        model /= model[2].reshape(1, -1)
+        mod_proj = np.matmul(intrinsic, model)
+        plt.imshow(image.transpose(1,2,0))
+        for i,c in enumerate(['r','g','b']):
+            x = [mod_proj[0,0], mod_proj[0, i+1]]
+            y = [mod_proj[1,0], mod_proj[1, i+1]]
+            plt.plot(x, y, c)
+        plt.show()
+
+if __name__ == '__main__':
+    main()