WatsonTang98 commited on
Commit
f4f33f0
1 Parent(s): 77424f3

Upload script.py

Browse files
Files changed (1) hide show
  1. script.py +314 -0
script.py ADDED
@@ -0,0 +1,314 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ### This is example of the script that will be run in the test environment.
2
+ ### Some parts of the code are compulsory and you should NOT CHANGE THEM.
3
+ ### They are between '''---compulsory---''' comments.
4
+ ### You can change the rest of the code to define and test your solution.
5
+ ### However, you should not change the signature of the provided function.
6
+ ### The script would save "submission.parquet" file in the current directory.
7
+ ### The actual logic of the solution is implemented in the `handcrafted_solution.py` file.
8
+ ### The `handcrafted_solution.py` file is a placeholder for your solution.
9
+ ### You should implement the logic of your solution in that file.
10
+ ### You can use any additional files and subdirectories to organize your code.
11
+
12
+ '''---compulsory---'''
13
+ # import subprocess
14
+ # from pathlib import Path
15
+ # def install_package_from_local_file(package_name, folder='packages'):
16
+ # """
17
+ # Installs a package from a local .whl file or a directory containing .whl files using pip.
18
+
19
+ # Parameters:
20
+ # path_to_file_or_directory (str): The path to the .whl file or the directory containing .whl files.
21
+ # """
22
+ # try:
23
+ # pth = str(Path(folder) / package_name)
24
+ # subprocess.check_call([subprocess.sys.executable, "-m", "pip", "install",
25
+ # "--no-index", # Do not use package index
26
+ # "--find-links", pth, # Look for packages in the specified directory or at the file
27
+ # package_name]) # Specify the package to install
28
+ # print(f"Package installed successfully from {pth}")
29
+ # except subprocess.CalledProcessError as e:
30
+ # print(f"Failed to install package from {pth}. Error: {e}")
31
+
32
+ # install_package_from_local_file('hoho')
33
+
34
+ import hoho; hoho.setup() # YOU MUST CALL hoho.setup() BEFORE ANYTHING ELSE
35
+ # import subprocess
36
+ # import importlib
37
+ # from pathlib import Path
38
+ # import subprocess
39
+
40
+
41
+ # ### The function below is useful for installing additional python wheels.
42
+ # def install_package_from_local_file(package_name, folder='packages'):
43
+ # """
44
+ # Installs a package from a local .whl file or a directory containing .whl files using pip.
45
+
46
+ # Parameters:
47
+ # path_to_file_or_directory (str): The path to the .whl file or the directory containing .whl files.
48
+ # """
49
+ # try:
50
+ # pth = str(Path(folder) / package_name)
51
+ # subprocess.check_call([subprocess.sys.executable, "-m", "pip", "install",
52
+ # "--no-index", # Do not use package index
53
+ # "--find-links", pth, # Look for packages in the specified directory or at the file
54
+ # package_name]) # Specify the package to install
55
+ # print(f"Package installed successfully from {pth}")
56
+ # except subprocess.CalledProcessError as e:
57
+ # print(f"Failed to install package from {pth}. Error: {e}")
58
+
59
+
60
+ # pip download webdataset -d packages/webdataset --platform manylinux1_x86_64 --python-version 38 --only-binary=:all:
61
+ # install_package_from_local_file('webdataset')
62
+ # install_package_from_local_file('tqdm')
63
+
64
+ ### Here you can import any library or module you want.
65
+ ### The code below is used to read and parse the input dataset.
66
+ ### Please, do not modify it.
67
+
68
+ import webdataset as wds
69
+ from tqdm import tqdm
70
+ from typing import Dict
71
+ import pandas as pd
72
+ from transformers import AutoTokenizer
73
+ import os
74
+ import time
75
+ import io
76
+ from PIL import Image as PImage
77
+ import numpy as np
78
+
79
+ from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
80
+ from hoho import proc, Sample
81
+
82
+ def convert_entry_to_human_readable(entry):
83
+ out = {}
84
+ already_good = ['__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces', 'face_semantics', 'K', 'R', 't']
85
+ for k, v in entry.items():
86
+ if k in already_good:
87
+ out[k] = v
88
+ continue
89
+ if k == 'points3d':
90
+ out[k] = read_points3D_binary(fid=io.BytesIO(v))
91
+ if k == 'cameras':
92
+ out[k] = read_cameras_binary(fid=io.BytesIO(v))
93
+ if k == 'images':
94
+ out[k] = read_images_binary(fid=io.BytesIO(v))
95
+ if k in ['ade20k', 'gestalt']:
96
+ out[k] = [PImage.open(io.BytesIO(x)).convert('RGB') for x in v]
97
+ if k == 'depthcm':
98
+ out[k] = [PImage.open(io.BytesIO(x)) for x in entry['depthcm']]
99
+ return out
100
+
101
+ '''---end of compulsory---'''
102
+
103
+ ### The part below is used to define and test your solution.
104
+ import subprocess
105
+ import sys
106
+ import os
107
+
108
+ import numpy as np
109
+ os.environ['MKL_THREADING_LAYER'] = 'GNU'
110
+ os.environ['MKL_SERVICE_FORCE_INTEL'] = '1'
111
+
112
+ def uninstall_package(package_name):
113
+ """
114
+ Uninstalls a package using pip.
115
+
116
+ Parameters:
117
+ package_name (str): The name of the package to uninstall.
118
+ """
119
+ try:
120
+ subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", package_name])
121
+ print(f"Package {package_name} uninstalled successfully")
122
+ except subprocess.CalledProcessError as e:
123
+ print(f"Failed to uninstall package {package_name}. Error: {e}")
124
+
125
+ # def download_packages(packages, folder='packages/torch'):
126
+ # """
127
+ # Downloads packages as .whl files into the specified folder using pip.
128
+
129
+ # Parameters:
130
+ # packages (list): List of packages to download with versions.
131
+ # folder (str): The folder where the .whl files will be saved.
132
+ # """
133
+ # Path(folder).mkdir(parents=True, exist_ok=True)
134
+ # try:
135
+ # subprocess.check_call([sys.executable, "-m", "pip", "download",
136
+ # "--platform", "manylinux1_x86_64",
137
+ # "--python-version", "38",
138
+ # "--only-binary=:all:",
139
+ # "-d", folder] + packages)
140
+ # print(f"Packages downloaded successfully into {folder}")
141
+ # except subprocess.CalledProcessError as e:
142
+ # print(f"Failed to download packages. Error: {e}")
143
+
144
+ def download_packages(packages, folder):
145
+ # Create the directory if it doesn't exist
146
+ if not os.path.exists(folder):
147
+ os.makedirs(folder)
148
+
149
+ try:
150
+ subprocess.check_call([
151
+ 'pip', 'download',
152
+ '--dest', folder,
153
+ '-f', 'https://download.pytorch.org/whl/cu121'
154
+ ] + packages)
155
+ print(f"Packages downloaded successfully to {folder}")
156
+ except subprocess.CalledProcessError as e:
157
+ print(f"Failed to download packages. Error: {e}")
158
+
159
+ # Set CUDA environment variables
160
+ os.environ['CUDA_HOME'] = '/usr/local/cuda-12.1'
161
+ os.environ['PATH'] = os.environ['CUDA_HOME'] + '/bin:' + os.environ['PATH']
162
+ os.environ['LD_LIBRARY_PATH'] = os.environ['CUDA_HOME'] + '/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
163
+
164
+ def install_package_from_local_file(package_name, folder='packages'):
165
+ """
166
+ Installs a package from a local .whl file or a directory containing .whl files using pip.
167
+
168
+ Parameters:
169
+ package_name (str): The name of the package to install.
170
+ folder (str): The folder where the .whl files are located.
171
+ """
172
+ try:
173
+ pth = str(Path(folder) / package_name)
174
+ subprocess.check_call([sys.executable, "-m", "pip", "install",
175
+ "--no-index", # Do not use package index
176
+ "--find-links", pth, # Look for packages in the specified directory or at the file
177
+ package_name]) # Specify the package to install
178
+ print(f"Package installed successfully from {pth}")
179
+ except subprocess.CalledProcessError as e:
180
+ print(f"Failed to install package from {pth}. Error: {e}")
181
+
182
+ def install_which():
183
+ try:
184
+ # Attempt to install which if it's not available
185
+ subprocess.check_call(['sudo', 'apt-get', 'install', '-y', 'which'])
186
+ print("Which installed successfully.")
187
+ except subprocess.CalledProcessError as e:
188
+ print(f"An error occurred while installing which: {e}")
189
+ sys.exit(1)
190
+
191
+ def setup_environment():
192
+ # Uninstall torch if it is already installed
193
+ # packages_to_uninstall = ['torch', 'torchvision', 'torchaudio']
194
+ # for package in packages_to_uninstall:
195
+ # uninstall_package(package)
196
+ # Download required packages
197
+ # pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
198
+ # pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://download.pytorch.org/whl/cu121
199
+ # pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121
200
+ # packages_to_download = ['torch==1.13.1', 'torchvision==0.14.1', 'torchaudio==0.13.1']
201
+ # packages_to_download = ['torch==2.1.0', 'torchvision==0.16.0', 'torchaudio==2.1.0']
202
+ # download_packages(packages_to_download, folder='packages/torch')
203
+
204
+ # Install ninja
205
+ # install_package_from_local_file('ninja', folder='packages')
206
+
207
+ # packages_to_download = ['torch==2.1.0', 'torchvision==0.16.0', 'torchaudio==2.1.0']
208
+ # download_folder = 'packages/torch'
209
+
210
+ # Download the packages
211
+ # download_packages(packages_to_download, download_folder)
212
+
213
+ # Install packages from local files
214
+ # install_package_from_local_file('torch', folder='packages')
215
+ # install_package_from_local_file('packages/torch/torchvision-0.16.0-cp38-cp38-manylinux1_x86_64.whl', folder='packages/torch')
216
+ # install_package_from_local_file('packages/torch/torchaudio-2.1.0-cp38-cp38-manylinux1_x86_64.whl', folder='packages/torch')
217
+ # install_package_from_local_file('scikit-learn', folder='packages')
218
+ # install_package_from_local_file('open3d', folder='packages')
219
+ # install_package_from_local_file('easydict', folder='packages')
220
+ # install_package_from_local_file('setuptools', folder='packages')
221
+ # install_package_from_local_file('ninja', folder='packages')
222
+ # download_packages(['scikit-learn'], folder='packages/scikit-learn')
223
+ # download_packages(['open3d'], folder='packages/open3d')
224
+ # download_packages(['easydict'], folder='packages/easydict')
225
+
226
+ # try:
227
+ # subprocess.check_call(['which', 'which'])
228
+ # except subprocess.CalledProcessError:
229
+ # install_which()
230
+
231
+ # Set environment variables for CUDA
232
+ os.environ['CUDA_HOME'] = '/usr/local/cuda'
233
+ os.environ['PATH'] = os.environ['CUDA_HOME'] + '/bin:' + os.environ['PATH']
234
+ os.environ['LD_LIBRARY_PATH'] = os.environ['CUDA_HOME'] + '/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
235
+ os.environ['LIBRARY_PATH'] = os.environ['CUDA_HOME'] + '/lib64'
236
+
237
+ # Print CUDA environment variables to verify
238
+ # print("CUDA_HOME:", os.environ['CUDA_HOME'])
239
+ # print("PATH:", os.environ['PATH'])
240
+ # print("LD_LIBRARY_PATH:", os.environ['LD_LIBRARY_PATH'])
241
+ # print("LIBRARY_PATH:", os.environ['LIBRARY_PATH'])
242
+
243
+ # Verify CUDA headers are accessible
244
+ cuda_include_path = os.path.join(os.environ['CUDA_HOME'], 'include')
245
+ if not os.path.exists(os.path.join(cuda_include_path, 'cuda.h')):
246
+ raise EnvironmentError(f"CUDA headers not found in {cuda_include_path}. Please check your CUDA installation.")
247
+
248
+ pc_util_path = os.path.join(os.getcwd(), 'pc_util')
249
+ if os.path.isdir(pc_util_path):
250
+ os.chdir(pc_util_path)
251
+ subprocess.check_call([sys.executable, "setup.py", "install"], cwd=pc_util_path)
252
+ os.chdir("..")
253
+
254
+ def setup_cuda_environment():
255
+ # cuda_home = '/usr/local/cuda'
256
+ # if not os.path.exists(cuda_home):
257
+ # raise EnvironmentError(f"CUDA_HOME directory {cuda_home} does not exist. Please install CUDA and set CUDA_HOME environment variable.")
258
+ # os.environ['CUDA_HOME'] = cuda_home
259
+ # os.environ['PATH'] = f"{cuda_home}/bin:{os.environ['PATH']}"
260
+ # os.environ['LD_LIBRARY_PATH'] = f"{cuda_home}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
261
+
262
+ os.environ['PATH'] = '/usr/local/cuda/bin'
263
+ os.environ['LD_LIBRARY_PATH'] = '/usr/local/cuda/lib64'
264
+ os.environ['LIBRARY_PATH'] = '/usr/local/cuda/lib64'
265
+
266
+ # usr_local_contents = os.listdir('/usr/local')
267
+ # # print("Items under /usr/local:")
268
+ # for item in usr_local_contents:
269
+ # print(item)
270
+
271
+ from pathlib import Path
272
+ def save_submission(submission, path):
273
+ """
274
+ Saves the submission to a specified path.
275
+
276
+ Parameters:
277
+ submission (List[Dict[]]): The submission to save.
278
+ path (str): The path to save the submission to.
279
+ """
280
+ sub = pd.DataFrame(submission, columns=["__key__", "wf_vertices", "wf_edges"])
281
+ sub.to_parquet(path)
282
+ print(f"Submission saved to {path}")
283
+
284
+ if __name__ == "__main__":
285
+ # setup_cuda_environment()
286
+ setup_environment()
287
+
288
+ from handcrafted_solution import predict
289
+ print ("------------ Loading dataset------------ ")
290
+ params = hoho.get_params()
291
+ dataset = hoho.get_dataset(decode=None, split='all', dataset_type='webdataset')
292
+
293
+ print('------------ Now you can do your solution ---------------')
294
+ solution = []
295
+ from concurrent.futures import ProcessPoolExecutor
296
+ with ProcessPoolExecutor(max_workers=1) as pool:
297
+ results = []
298
+ for i, sample in enumerate(tqdm(dataset)):
299
+ results.append(pool.submit(predict, sample, visualize=False))
300
+
301
+ for i, result in enumerate(tqdm(results)):
302
+ key, pred_vertices, pred_edges = result.result()
303
+ solution.append({
304
+ '__key__': key,
305
+ 'wf_vertices': pred_vertices.tolist(),
306
+ 'wf_edges': pred_edges
307
+ })
308
+ if i % 100 == 0:
309
+ # incrementally save the results in case we run out of time
310
+ print(f"Processed {i} samples")
311
+ # save_submission(solution, Path(params['output_path']) / "submission.parquet")
312
+ print('------------ Saving results ---------------')
313
+ save_submission(solution, Path(params['output_path']) / "submission.parquet")
314
+ print("------------ Done ------------ ")