first commit

hoho.py

@@ -1,261 +0,0 @@
-import os 
-import json
-import shutil
-from pathlib import Path
-from typing import Dict
-
-from PIL import ImageFile
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-
-LOCAL_DATADIR = None
-
-def setup(local_dir='./data/usm-training-data/data'):
-    
-    # If we are in the test environment, we need to link the data directory to the correct location
-    tmp_datadir = Path('/tmp/data/data')
-    local_test_datadir = Path('./data/usm-test-data-x/data')
-    local_val_datadir = Path(local_dir)
-    
-    os.system('pwd')
-    os.system('ls -lahtr .')
-    
-    if tmp_datadir.exists() and not local_test_datadir.exists():
-        global LOCAL_DATADIR
-        LOCAL_DATADIR = local_test_datadir
-        # shutil.move(datadir, './usm-test-data-x/data')
-        print(f"Linking {tmp_datadir} to {LOCAL_DATADIR} (we are in the test environment)")
-        LOCAL_DATADIR.parent.mkdir(parents=True, exist_ok=True)
-        LOCAL_DATADIR.symlink_to(tmp_datadir)
-    else:
-        LOCAL_DATADIR = local_val_datadir
-        print(f"Using {LOCAL_DATADIR} as the data directory (we are running locally)")
-        
-    # os.system("ls -lahtr")
-    
-    assert LOCAL_DATADIR.exists(), f"Data directory {LOCAL_DATADIR} does not exist"
-    return LOCAL_DATADIR
-    
-    
-    
-    
-import importlib
-from pathlib import Path
-import subprocess
-
-def download_package(package_name, path_to_save='packages'):
-    """
-    Downloads a package using pip and saves it to a specified directory.
-
-    Parameters:
-    package_name (str): The name of the package to download.
-    path_to_save (str): The path to the directory where the package will be saved.
-    """
-    try:
-        # pip download webdataset -d packages/webdataset --platform manylinux1_x86_64 --python-version 38 --only-binary=:all:
-        subprocess.check_call([subprocess.sys.executable, "-m", "pip", "download", package_name, 
-                               "-d", str(Path(path_to_save)/package_name),  # Download the package to the specified directory
-                               "--platform", "manylinux1_x86_64",  # Specify the platform
-                               "--python-version", "38",  # Specify the Python version
-                               "--only-binary=:all:"])  # Download only binary packages
-        print(f'Package "{package_name}" downloaded successfully')
-    except subprocess.CalledProcessError as e:
-        print(f'Failed to downloaded package "{package_name}". Error: {e}')
-              
-              
-def install_package_from_local_file(package_name, folder='packages'):
-    """
-    Installs a package from a local .whl file or a directory containing .whl files using pip.
-
-    Parameters:
-    path_to_file_or_directory (str): The path to the .whl file or the directory containing .whl files.
-    """
-    try:
-        pth = str(Path(folder) / package_name)
-        subprocess.check_call([subprocess.sys.executable, "-m", "pip", "install", 
-                               "--no-index",  # Do not use package index
-                               "--find-links", pth,  # Look for packages in the specified directory or at the file
-                               package_name])  # Specify the package to install
-        print(f"Package installed successfully from {pth}")
-    except subprocess.CalledProcessError as e:
-        print(f"Failed to install package from {pth}. Error: {e}")
-        
-        
-def importt(module_name, as_name=None):
-    """
-    Imports a module and returns it.
-
-    Parameters:
-    module_name (str): The name of the module to import.
-    as_name (str): The name to use for the imported module. If None, the original module name will be used.
-
-    Returns:
-    The imported module.
-    """
-    for _ in range(2):
-        try:
-            if as_name is None:
-                print(f'imported {module_name}')
-                return importlib.import_module(module_name)
-            else:
-                print(f'imported {module_name} as {as_name}')
-                return importlib.import_module(module_name, as_name)
-        except ModuleNotFoundError as e:
-            install_package_from_local_file(module_name)
-            print(f"Failed to import module {module_name}. Error: {e}")
-            
-    
-def prepare_submission():
-    # Download packages from requirements.txt 
-    if Path('requirements.txt').exists():
-        print('downloading packages from requirements.txt')
-        Path('packages').mkdir(exist_ok=True)
-        with open('requirements.txt') as f:
-            packages = f.readlines()
-            for p in packages:
-                download_package(p.strip())
-                
-                
-    print('all packages downloaded. Don\'t foget to include the packages in the submission by adding them with git lfs.')
-        
-
-def Rt_to_eye_target(im, K, R, t):
-    height = im.height
-    focal_length = K[0,0]
-    fov = 2.0 * np.arctan2((0.5 * height), focal_length) / (np.pi / 180.0)
-
-    x_axis, y_axis, z_axis = R
-
-    eye = -(R.T @ t).squeeze()
-    z_axis = z_axis.squeeze()
-    target = eye + z_axis
-    up = -y_axis
-    
-    return eye, target, up, fov        
-
-
-########## general utilities ##########
-import contextlib
-import tempfile 
-from pathlib import Path
-
-@contextlib.contextmanager
-def working_directory(path):
-    """Changes working directory and returns to previous on exit."""
-    prev_cwd = Path.cwd()
-    os.chdir(path)
-    try:
-        yield
-    finally:
-        os.chdir(prev_cwd)
-        
-@contextlib.contextmanager
-def temp_working_directory():
-    with tempfile.TemporaryDirectory(dir='.') as D:
-        with working_directory(D):
-            yield
-
-
-############# Dataset ############# 
-def proc(row, split='train'):
-    # column_names_train = ['ade20k', 'depthcm', 'gestalt', 'colmap', 'KRt', 'mesh', 'wireframe']
-    # column_names_test = ['ade20k', 'depthcm', 'gestalt', 'colmap', 'KRt', 'wireframe']
-    # cols = column_names_train if split == 'train' else column_names_test
-    out = {}
-    for k, v in row.items():
-        colname = k.split('.')[0]
-        if colname in {'ade20k', 'depthcm', 'gestalt'}:
-            if colname in out:
-                out[colname].append(v)
-            else:
-                out[colname] = [v]
-        elif colname in {'wireframe', 'mesh'}:
-            # out.update({a: b.tolist() for a,b in v.items()})
-            out.update({a: b for a,b in v.items()})
-        elif colname in 'kr':
-            out[colname.upper()] = v
-        else:
-            out[colname] = v
-            
-    return Sample(out)
-
-
-class Sample(Dict):
-    def __repr__(self):
-        return str({k: v.shape if hasattr(v, 'shape') else [type(v[0])] if isinstance(v, list) else type(v) for k,v in self.items()})
-
-    
-        
-def get_params():
-    exmaple_param_dict = {
-        "competition_id": "usm3d/S23DR",
-        "competition_type": "script",
-        "metric": "custom",
-        "token": "hf_**********************************",
-        "team_id": "local-test-team_id",
-        "submission_id": "local-test-submission_id",
-        "submission_id_col": "__key__",
-        "submission_cols": [
-            "__key__",
-            "wf_edges",
-            "wf_vertices",
-            "edge_semantics"
-        ],
-        "submission_rows": 180,
-        "output_path": ".",
-        "submission_repo": "<THE HF MODEL ID of THIS REPO",
-        "time_limit": 7200,
-        "dataset": "usm3d/usm-test-data-x",
-        "submission_filenames": [
-            "submission.parquet"
-        ]
-    }
-    
-    param_path = Path('params.json')
-    
-    if not param_path.exists():
-        print('params.json not found (this means we probably aren\'t in the test env). Using example params.')
-        params = exmaple_param_dict
-    else:
-        print('found params.json (this means we are probably in the test env). Using params from file.')
-        with param_path.open() as f:
-            params = json.load(f)
-    print(params)
-    return params
-
-
-
-import webdataset as wds 
-import numpy as np
-
-def get_dataset(decode='pil', proc=proc, split='train', dataset_type='webdataset'):
-    if LOCAL_DATADIR is None:
-        raise ValueError('LOCAL_DATADIR is not set. Please run setup() first.')
-        
-    local_dir = Path(LOCAL_DATADIR)
-    if split != 'all':
-        local_dir = local_dir / split
-    
-    paths = [str(p) for p in local_dir.rglob('*.tar.gz')]
-    
-    dataset = wds.WebDataset(paths)
-    if decode is not None:
-        dataset = dataset.decode(decode)
-    else:
-        dataset = dataset.decode()
-    
-    dataset = dataset.map(proc)
-    
-    if dataset_type == 'webdataset':
-        return dataset
-    
-    if dataset_type == 'hf':
-        import datasets
-        from datasets import Features, Value, Sequence, Image, Array2D
-       
-        if split == 'train':
-            return datasets.IterableDataset.from_generator(lambda: dataset.iterator())
-        elif split == 'val':
-            return datasets.IterableDataset.from_generator(lambda: dataset.iterator())
-
-    
-