implement install.py

Author: mxh1999

embodiedscan/datasets/transforms/loading.py

@@ -93,6 +93,7 @@ def __repr__(self):
         return repr_str
 
 
+# TODO : refine
 @TRANSFORMS.register_module()
 class LoadAnnotations3D(LoadAnnotations):
     """Load Annotations3D.

embodiedscan/explorer.py

@@ -1,7 +1,7 @@
 import os
+import pickle
 from typing import List, Union
 
-import mmengine
 import numpy as np
 import open3d as o3d
 
@@ -68,7 +68,8 @@ def __init__(self,
         self.metainfo = None
         data_list = []
         for file in self.ann_files:
-            data = mmengine.load(file)
+            with open(file, 'rb') as f:
+                data = pickle.load(f)
             if self.metainfo is None:
                 self.metainfo = data['metainfo']
             else:

embodiedscan/visualization/continuous_drawer.py

@@ -1,7 +1,7 @@
 import os
+import pickle
 
 import cv2
-import mmengine
 import numpy as np
 import open3d as o3d
 
@@ -223,7 +223,8 @@ def begin(self):
         print('You can also press Esc to close window immediately,',
               'which may result in a segmentation fault.')
         self.gt = np.load(self.occ_path)
-        self.mask = mmengine.load(self.mask_path)
+        with open(self.mask_path, 'rb') as f:
+            self.mask = pickle.load(f)
 
         point_cloud_range = [-3.2, -3.2, -1.28 + 0.5, 3.2, 3.2, 1.28 + 0.5]
         occ_size = [40, 40, 16]

install.py

@@ -0,0 +1,110 @@
+import argparse
+import re
+import subprocess
+import sys
+
+
+def run_subprocess(command):
+    try:
+        process = subprocess.Popen(command,
+                                   stdout=subprocess.PIPE,
+                                   stderr=subprocess.PIPE,
+                                   universal_newlines=True)
+
+        # Read output and error in real-time
+        for line in process.stdout:
+            print(line.strip())
+        for line in process.stderr:
+            print(line.strip())
+
+        # Wait for the subprocess to finish
+        process.wait()
+
+        # Get the return code
+        return_code = process.returncode
+
+        if return_code != 0:
+            print(f'Command failed with return code {return_code}')
+
+    except subprocess.CalledProcessError as e:
+        print(f'Command failed with return code {e.returncode}')
+        print('Error output:')
+        print(e.output.decode())
+
+
+def pytorch3d_links():
+    try:
+        import torch
+    except ImportError as e:
+        print('Pytorch is not installed.')
+        raise e
+    cuda_version = torch.version.cuda
+    if cuda_version is None:
+        print('Pytorch is cpu only.')
+        raise NotImplementedError
+
+    pyt_version_str = torch.__version__.split('+')[0].replace('.', '')
+    cuda_version_str = torch.version.cuda.replace('.', '')
+    version_str = ''.join([
+        f'py3{sys.version_info.minor}_cu', cuda_version_str,
+        f'_pyt{pyt_version_str}'
+    ])
+    pytorch3d_links = f'https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html'  # noqa: E501
+    return pytorch3d_links
+
+
+def mmcv_links():
+    try:
+        import torch
+    except ImportError as e:
+        print('Pytorch is not installed.')
+        raise e
+    cuda_version = torch.version.cuda
+    if cuda_version is None:
+        print('Pytorch is cpu only.')
+        raise NotImplementedError
+
+    cuda_version_str = torch.version.cuda.replace('.', '')
+    pyt_version = torch.__version__.split('+')[0].split('.')
+    pyt_version_mmcv = pyt_version[0] + '.' + pyt_version[1]
+    mmcv_links = f'https://download.openmmlab.com/mmcv/dist/cu{cuda_version_str}/torch{pyt_version_mmcv}/index.html'  # noqa: E501
+    return mmcv_links
+
+
+def install_package(line):
+    pat = '(' + '|'.join(['>=', '==', '>', '<', '<=', '@']) + ')'
+    parts = re.split(pat, line, maxsplit=1)
+    package_name = parts[0].strip()
+    print('installing', package_name)
+    if package_name == 'pytorch3d':
+        links = pytorch3d_links()
+        run_subprocess(
+            [sys.executable, '-m', 'pip', 'install', 'pytorch3d', '-f', links])
+    elif package_name == 'mmcv':
+        links = mmcv_links()
+        run_subprocess(
+            [sys.executable, '-m', 'pip', 'install', line, '-f', links])
+    else:
+        run_subprocess([sys.executable, '-m', 'pip', 'install', line])
+
+
+def install_requires(fname):
+    with open(fname, 'r') as f:
+        for line in f.readlines():
+            line = line.strip()
+            if line:
+                install_package(line)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Install Embodiedscan from pre-built package.')
+    parser.add_argument('mode', default=None)
+    args = parser.parse_args()
+
+    install_requires('requirements/base.txt')
+    if args.mode == 'visual' or args.mode == 'all':
+        install_requires('requirements/visual.txt')
+
+    if args.mode == 'run' or args.mode == 'all':
+        install_requires('requirements/run.txt')

requirements/run.txt

@@ -0,0 +1,6 @@
+MinkowskiEngine @ git+https://github.com/NVIDIA/MinkowskiEngine.git
+mmcv==2.0.0rc4
+mmdet
+mmengine
+ninja
+pytorch3d