Spaces:
Running
Running
# setup.py | |
#!/usr/bin/env python | |
import os | |
import sys | |
import subprocess | |
import argparse | |
from pathlib import Path | |
def run_command(command): | |
print(f"Running: {command}") | |
process = subprocess.Popen(command, shell=True) | |
process.wait() | |
if process.returncode != 0: | |
print(f"Command failed: {command}") | |
sys.exit(1) | |
def main(): | |
parser = argparse.ArgumentParser(description="Setup NeuroNest project") | |
parser.add_argument("--download", action="store_true", help="Download checkpoint") | |
args = parser.parse_args() | |
# Create directory structure | |
dirs = ["app", "utils", "utils/contrast", "configs", "data"] | |
for dir_path in dirs: | |
os.makedirs(dir_path, exist_ok=True) | |
if args.download: | |
# Download the OneFormer checkpoint | |
checkpoint_dir = "data/oneformer_ade20k_swin_large" | |
if not os.path.exists(checkpoint_dir): | |
print(f"Creating checkpoint directory: {checkpoint_dir}") | |
os.makedirs(checkpoint_dir, exist_ok=True) | |
if not os.path.exists(f"{checkpoint_dir}/model_final.pth"): | |
print("Downloading OneFormer ADE20K checkpoint...") | |
# Use git-lfs if available, otherwise use curl/wget | |
if subprocess.call("which git-lfs", shell=True, stdout=subprocess.DEVNULL) == 0: | |
run_command("git lfs install") | |
run_command(f"git clone https://huggingface.co/shi-labs/oneformer_ade20k_swin_large {checkpoint_dir}") | |
else: | |
# Alternative download method using huggingface_hub | |
run_command("pip install huggingface_hub") | |
run_command(f"python -c \"from huggingface_hub import snapshot_download; snapshot_download('shi-labs/oneformer_ade20k_swin_large', local_dir='{checkpoint_dir}')\"") | |
print("✅ Setup complete!") | |
if __name__ == "__main__": | |
main() | |