├── .all-contributorsrc ├── .env.example ├── .github └── workflows │ ├── dockerhub.yml │ ├── mkdocs.yml │ ├── publish.yml │ └── tests.yml ├── .gitignore ├── Dockerfile ├── MANIFEST.in ├── README.md ├── data ├── camera_data.json ├── datasets.json ├── object_data.json └── stage_data.json ├── docs ├── assets │ ├── DeepAI.png │ ├── favicon.ico │ ├── logo.png │ └── simian_banner.png ├── backgrounds.md ├── batch_rendering.md ├── cameras.md ├── combiner.md ├── getting_started.md ├── index.md ├── object.md ├── postprocessing.md ├── rendering.md ├── scene.md ├── transform.md └── worker.md ├── mkdocs.yml ├── requirements.txt ├── scripts ├── data │ ├── get_data.sh │ ├── get_objaverse_data.py │ ├── get_polyhaven_background_data.py │ └── get_polyhaven_texture_data.py ├── demo │ └── demo.py ├── filter │ ├── caption_combination_pairs.py │ ├── combinations_add_placeholder.py │ ├── combinations_to_captions.py │ ├── combinations_to_motion_labels.py │ ├── get_captions.py │ ├── get_ontop_captions.py │ ├── rewrite_captions_gem.py │ ├── rewrite_captions_gpt.py │ └── test_combinations.py ├── setup_redis.sh └── start_x_server.py ├── setup.py └── simian ├── __init__.py ├── background.py ├── batch.py ├── camera.py ├── combiner.py ├── distributed.py ├── object.py ├── postprocessing.py ├── prompts.py ├── render.py ├── scene.py ├── server.py ├── tests ├── __run__.py ├── background_test.py ├── batch_test.py ├── camera_test.py ├── combiner_test.py ├── new_camera_test.py ├── object_test.py ├── postprocessing_test.py ├── transform_test.py └── worker_test.py ├── transform.py ├── vendor ├── __init__.py └── objaverse.py └── worker.py /.all-contributorsrc: -------------------------------------------------------------------------------- 1 | { 2 | "files": [ 3 | "README.md" 4 | ], 5 | "imageSize": 100, 6 | "commit": false, 7 | "commitType": "docs", 8 | "commitConvention": "angular", 9 | "contributors": [ 10 | { 11 | "login": "eric-prog", 12 | "name": "Eric S", 13 | "avatar_url": "https://avatars.githubusercontent.com/u/59460685?v=4", 14 | "profile": "https://ericsheen.tech/", 15 | "contributions": [ 16 | "infra", 17 | "code", 18 | "test", 19 | "doc" 20 | ] 21 | }, 22 | { 23 | "login": "lalalune", 24 | "name": "M̵̞̗̝̼̅̏̎͝Ȯ̴̝̻̊̃̋̀Õ̷̼͋N̸̩̿͜ ̶̜̠̹̼̩͒", 25 | "avatar_url": "https://avatars.githubusercontent.com/u/18633264?v=4", 26 | "profile": "https://github.com/lalalune", 27 | "contributions": [ 28 | "infra", 29 | "code", 30 | "test", 31 | "doc" 32 | ] 33 | } 34 | ], 35 | "contributorsPerLine": 7, 36 | "skipCi": true, 37 | "repoType": "github", 38 | "repoHost": "https://github.com", 39 | "projectName": "Simian", 40 | "projectOwner": "DeepAI Research" 41 | } 42 | -------------------------------------------------------------------------------- /.env.example: -------------------------------------------------------------------------------- 1 | REDIS_HOST=redis-xxx.xxx.us-west-1-2.ec2.redns.redis-cloud.com 2 | REDIS_PORT=1337 3 | REDIS_USER=default 4 | REDIS_PASSWORD=NZxxxxxxxxxxxxxxxxxxxxxxxxxxxxxBD 5 | HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxw 6 | HF_REPO_ID=org/repo 7 | HF_PATH=./ 8 | VAST_API_KEY=1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef 9 | OPENAI_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx 10 | GOOGLE_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx -------------------------------------------------------------------------------- /.github/workflows/dockerhub.yml: -------------------------------------------------------------------------------- 1 | name: Publish Docker image 2 | 3 | on: 4 | release: 5 | types: [published] 6 | 7 | env: 8 | REGISTRY: index.docker.io 9 | IMAGE_NAME: arfx/simian-worker 10 | 11 | jobs: 12 | push_to_registry: 13 | name: Push Docker image to Docker Hub 14 | runs-on: ubuntu-latest 15 | permissions: 16 | packages: write 17 | contents: read 18 | attestations: write 19 | id-token: write 20 | steps: 21 | - name: Check out the repo 22 | uses: actions/checkout@v4 23 | 24 | - name: Log in to Docker Hub 25 | uses: docker/login-action@f4ef78c080cd8ba55a85445d5b36e214a81df20a 26 | with: 27 | username: ${{ secrets.DOCKER_USERNAME }} 28 | password: ${{ secrets.DOCKER_PASSWORD }} 29 | 30 | - name: Extract metadata (tags, labels) for Docker 31 | id: meta 32 | uses: docker/metadata-action@9ec57ed1fcdbf14dcef7dfbe97b2010124a938b7 33 | with: 34 | images: ${{ env.IMAGE_NAME }} 35 | 36 | - name: Build and push Docker image 37 | id: push 38 | uses: docker/build-push-action@3b5e8027fcad23fda98b2e3ac259d8d67585f671 39 | with: 40 | context: . 41 | file: ./Dockerfile 42 | push: true 43 | tags: ${{ steps.meta.outputs.tags }} 44 | labels: ${{ steps.meta.outputs.labels }} 45 | 46 | 47 | - name: Generate artifact attestation 48 | uses: actions/attest-build-provenance@v1 49 | with: 50 | subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME}} 51 | subject-digest: ${{ steps.push.outputs.digest }} 52 | push-to-registry: true 53 | 54 | -------------------------------------------------------------------------------- /.github/workflows/mkdocs.yml: -------------------------------------------------------------------------------- 1 | name: mkdocs 2 | on: 3 | push: 4 | branches: 5 | - main 6 | permissions: 7 | contents: write 8 | jobs: 9 | deploy: 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@v4 13 | - name: Configure Git Credentials 14 | run: | 15 | git config user.name github-actions[bot] 16 | git config user.email 41898282+github-actions[bot]@users.noreply.github.com 17 | - uses: actions/setup-python@v5 18 | with: 19 | python-version: 3.x 20 | - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV # (3)! 21 | - uses: actions/cache@v4 22 | with: 23 | key: mkdocs-material-${{ env.cache_id }} 24 | path: .cache 25 | restore-keys: | 26 | mkdocs-material- 27 | - run: pip install mkdocs-material mkdocstrings mkdocstrings-python 28 | - run: mkdocs gh-deploy --force 29 | -------------------------------------------------------------------------------- /.github/workflows/publish.yml: -------------------------------------------------------------------------------- 1 | name: Upload Python Package 2 | 3 | on: 4 | release: 5 | types: [published] 6 | 7 | permissions: 8 | contents: read 9 | 10 | jobs: 11 | deploy: 12 | 13 | runs-on: ubuntu-latest 14 | 15 | steps: 16 | - uses: actions/checkout@v3 17 | - name: Set up Python 18 | uses: actions/setup-python@v3 19 | with: 20 | python-version: '3.x' 21 | - name: Install dependencies 22 | run: | 23 | python -m pip install --upgrade pip 24 | pip install build 25 | - name: Extract package version 26 | id: extract_version 27 | run: echo "package_version=$(echo $GITHUB_REF | cut -d / -f 3)" >> $GITHUB_ENV 28 | - name: Write package version to file 29 | run: echo "${{ env.package_version }}" > version.txt 30 | - name: Build package 31 | run: python -m build 32 | - name: Publish package 33 | uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 34 | with: 35 | user: ${{ secrets.PYPI_USERNAME }} 36 | password: ${{ secrets.PYPI_PASSWORD }} 37 | -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- 1 | name: Run Tests 2 | 3 | on: 4 | push: 5 | branches: [ main ] 6 | pull_request: 7 | branches: [ main ] 8 | env: 9 | HF_TOKEN: ${{ secrets.HF_TOKEN }} 10 | HF_REPO_ID: ${{ secrets.HF_REPO_ID }} 11 | HF_PATH: ${{ secrets.HF_PATH }} 12 | REDIS_HOST: ${{ secrets.REDIS_HOST }} 13 | REDIS_PORT: ${{ secrets.REDIS_PORT }} 14 | REDIS_USER: ${{ secrets.REDIS_USER }} 15 | REDIS_PASSWORD: ${{ secrets.REDIS_PASSWORD }} 16 | 17 | jobs: 18 | run-tests: 19 | runs-on: ubuntu-latest 20 | 21 | steps: 22 | - uses: actions/checkout@v2 23 | 24 | - name: Set up Python 25 | uses: actions/setup-python@v2 26 | with: 27 | python-version: 3.11.9 28 | 29 | - name: Install dependencies 30 | run: | 31 | python -m pip install --upgrade pip 32 | pip install -r requirements.txt 33 | 34 | - name: Get the data 35 | run: bash scripts/data/get_data.sh 36 | 37 | - name: Generate combinations 38 | run: python3.11 -m simian.combiner 39 | 40 | - name: Run tests 41 | run: python3.11 -m simian.tests.__run__ 42 | 43 | - name: Check test results 44 | if: failure() 45 | run: | 46 | echo "Some tests failed. Please check the test output for more details." 47 | exit 1 -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | __pycache__ 3 | simian/__pycache__ 4 | github 5 | hf-objaverse-v1 6 | sketchfab 7 | smithsonian 8 | thingiverse 9 | combinations.json 10 | annotations*.json 11 | captions.json 12 | datasets 13 | renders 14 | backgrounds 15 | examples 16 | materials 17 | sampled_df.json 18 | *.blend1 19 | *.obj 20 | venv/ 21 | .env 22 | version.txt 23 | dist 24 | config.json 25 | chroma_data 26 | chroma_db 27 | chroma.log 28 | get_captions_1.json 29 | get_captions_2.json 30 | get_captions_processed.json 31 | combinations_updated.json 32 | combinations_processed.json 33 | stationary 34 | combined_captions.py 35 | combinations 36 | combinations_0-26.json 37 | combinations_27-53.json -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM --platform=linux/x86_64 ubuntu:24.04 2 | 3 | ENV DEBIAN_FRONTEND=noninteractive 4 | ENV TZ=UTC 5 | 6 | RUN apt-get update && \ 7 | apt-get install -y \ 8 | python3-pip \ 9 | python3 \ 10 | xorg \ 11 | git \ 12 | && apt-get install -y software-properties-common && \ 13 | add-apt-repository ppa:deadsnakes/ppa && \ 14 | apt-get update && \ 15 | apt-get install -y python3.11 python3.11-distutils python3.11-dev && \ 16 | update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 && \ 17 | update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1 && \ 18 | apt-get clean \ 19 | && rm -rf /var/lib/apt/lists/* 20 | 21 | COPY scripts/ ./scripts/ 22 | RUN bash scripts/data/get_data.sh 23 | 24 | COPY requirements.txt . 25 | 26 | RUN python3.11 -m pip install --upgrade --ignore-installed setuptools wheel 27 | RUN python3.11 -m pip install omegaconf requests argparse numpy scipy rich chromadb bpy boto3 28 | RUN python3.11 -m pip install distributask==0.0.40 29 | 30 | COPY simian/ ./simian/ 31 | COPY data/ ./data/ 32 | 33 | CMD ["python3.11", "-m", "celery", "-A", "simian.worker", "worker", "--loglevel=info", "--concurrency=1"] -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include version.txt 2 | include README.md 3 | include requirements.txt 4 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Simian 2 | A synthetic data generator for video caption pairs. 3 | 4 | 5 | 6 | ![docs](https://github.com/RaccoonResearch/simian/actions/workflows/mkdocs.yml/badge.svg) 7 | ![tests](https://github.com/RaccoonResearch/simian/actions/workflows/tests.yml/badge.svg) 8 | ![docker](https://github.com/RaccoonResearch/simian/actions/workflows/dockerhub.yml/badge.svg) 9 | [![PyPI version](https://badge.fury.io/py/simian3d.svg)](https://badge.fury.io/py/simian3d) 10 | [![License](https://img.shields.io/badge/License-MIT-blue)](https://github.com/RaccoonResearch/simian/blob/main/LICENSE) 11 | [![stars](https://img.shields.io/github/stars/RaccoonResearch/simian?style=social)](https://github.com/RaccoonResearch/simian) 12 | 13 | Simian creates synthetic data that is usable for generative video and video captioning tasks. The data consists of videos and captions. The videos are generated using Blender, a 3D modeling software. 14 | 15 | ## 🖥️ Setup 16 | 17 | > **_NOTE:_** Simian requires Python 3.11. 18 | 19 | 1. Install dependences: 20 | ```bash 21 | pip install -r requirements.txt 22 | ``` 23 | 24 | 2. Download the datasets: 25 | ```bash 26 | ./scripts/data/get_data.sh 27 | ``` 28 | 29 | 3. [OPTIONAL] If you're on a headless Linux server, install Xorg and start it: 30 | 31 | ```bash 32 | sudo apt-get install xserver-xorg -y && \ 33 | sudo python3 scripts/start_x_server.py start 34 | ``` 35 | 36 | ## 📸 Usage 37 | 38 | ### Generating Combinations 39 | 40 | Generate scenes without movement (static videos): 41 | ```bash 42 | python3 -m simian.combiner --count 1000 --seed 42 43 | ``` 44 | 45 | Add movement to all or no objects (camera is stationary): 46 | ```bash 47 | python3 -m simian.combiner --count 1000 --seed 42 --movement 48 | ``` 49 | 50 | Allow objects to be on top of each other (static or movement): 51 | ```bash 52 | python3 -m simian.combiner --count 1000 --seed 42 --ontop 53 | ``` 54 | 55 | Make camera follow an object (camera follows object): 56 | ```bash 57 | python3 -m simian.combiner --count 1000 --seed 42 --camera_follow 58 | ``` 59 | 60 | Randomly apply movement, object stacking, and camera follow effects: 61 | ```bash 62 | python3 -m simian.combiner --count 1000 --seed 42 --random 63 | ``` 64 | 65 | ### Generating Videos or Images 66 | 67 | Configure the flags as needed: 68 | - `--width` and `--height` are the resolution of the video. 69 | - `--start_index` and `--end_index` are the number of videos in the combinations you want to run. 0-100 will compile all 100 videos. 70 | - `--combination_index` is the index of the combination to render. 71 | - `--output_dir` is the directory to save the rendered video. 72 | - `--hdri_path` is the directory containing the background images. 73 | - `--start_frame` and `--end_frame` are the start and end frames of the video. 74 | - `--images` adding this will output images instead of video at random frames. Creates multiple images per combination of varying sizes 75 | - `blend_file ` allows users to upload and use their own blend files as the terrain 76 | - `animation_length` is a percentage from 0-100 which describes how fast the animation should occur within the frames 77 | 78 | Or generate all or part of the combination set using the `batch.py` script: 79 | 80 | 81 | ### Generating Videos or Images 82 | 83 | Run: 84 | ``` 85 | python3 -m simian.batch 86 | ``` 87 | 88 | #### Batched Option 89 | 90 | To generate videos use the arrow keys and select enter. Then 91 | 92 | To generate a video(s): 93 | ```bash 94 | --start_index 0 --end_index 1000 --width 1024 --height 576 --start_frame 1 --end_frame 2 95 | ``` 96 | 97 | To generate an video(s) with your own blend file: 98 | ```bash 99 | --start_index 0 --end_index 1000 --width 1024 --height 576 --start_frame 1 --end_frame 3 ---blend 100 | ``` 101 | 102 | To generate an image(s): 103 | ```bash 104 | --start_index 0 --end_index 1000 --width 1024 --height 576 --start_frame 1 --end_frame 2 --images 105 | ``` 106 | 107 | #### Prompt Option 108 | 109 | Must first embedd all the data 110 | 111 | ``` 112 | python3 server/server.py 113 | ``` 114 | 115 | #### Set Up Redis 116 | You can make a free Redis account [here](https://redis.io/try-free/). 117 | 118 | For local testing and multiple local workers, you can use the following script to set up a local instance of Redis: 119 | ```bash 120 | scripts/setup_redis.sh 121 | ``` 122 | 123 | #### Huggingface API Key 124 | 125 | You can get a Huggingface API key [here](https://huggingface.co/settings/tokens). 126 | 127 | Now, start your workers 128 | ```bash 129 | export REDIS_HOST=.com 130 | export REDIS_PORT=1337 131 | export REDIS_USER=default 132 | export REDIS_PASSWORD= 133 | export HF_TOKEN= 134 | export HF_REPO_ID= 135 | celery -A simian.worker worker --loglevel=info 136 | ``` 137 | 138 | You can also build and run the worker with Docker 139 | ```bash 140 | # build the container 141 | docker build -t simian-worker . 142 | 143 | # run the container with .env 144 | docker run --env-file .env simian-worker 145 | 146 | # run the container with environment variables 147 | docker run -e REDIS_HOST={myhost} -e REDIS_PORT={port} -e REDIS_USER=default -e REDIS_PASSWORD={some password} -e HF_TOKEN={token} -e HF_REPO_ID={repo_id} simian-worker 148 | ``` 149 | 150 | Finally, issue work to your task queue 151 | 152 | ```bash 153 | python3 -m simian.distributed --width 1024 --height 576 154 | ``` 155 | 156 | If you want to use a custom or hosted Redis instance (recommended), you can add the redis details like this: 157 | ```bash 158 | export REDIS_HOST=.com 159 | export REDIS_PORT=1337 160 | export REDIS_USER=default 161 | export REDIS_PASSWORD= 162 | ``` 163 | 164 | To run all tests 165 | 166 | ``` 167 | python3 -m simian.tests.__run__ 168 | ``` 169 | 170 | To run tests look into the test folder and run whichever test file you want 171 | 172 | ```bash 173 | python3 -m simian.tests.object_test 174 | ``` 175 | 176 | ## 📁 Datasets 177 | 178 | We are currently using the following datasets: 179 | [Objaverse](https://huggingface.co/datasets/allenai/objaverse) 180 | 181 | Backgrounds are loaded from: 182 | [Poly Haven](https://polyhaven.com) 183 | 184 | ## 🦝 Contributing 185 | 186 | We welcome contributions! We're especially interested in help adding and refining datasets, improving generation quality, adding new features and dynamics and allowing the project to meet more use cases. 187 | 188 | ### How to contribute 189 | 190 | 1. Check out the issues here. 191 | 2. Join our Discord here. 192 | 3. Get in touch with us so we can coordinate on development. 193 | 4. Or, you know, just YOLO a pull request. We're pretty chill. 194 | 195 | ## 📜 License 196 | 197 | This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. 198 | 199 | If you use it, please cite us: 200 | 201 | ```bibtex 202 | @misc{Simverse, 203 | author = {Deep AI, Inc}, 204 | title = {Simverse: A Synthetic Data Generator for Video Caption Pairs}, 205 | year = {2024}, 206 | publisher = {GitHub}, 207 | howpublished = {\url{https://github.com/DeepAI-Research/Simverse}} 208 | } 209 | ``` 210 | 211 | ## Contributors ✨ 212 | 213 | 214 | [![All Contributors](https://img.shields.io/badge/all_contributors-2-orange.svg?style=flat-square)](#contributors-) 215 | 216 | 217 | This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome! 218 | 219 | 220 | 221 | 222 | 223 | 224 | 225 | 226 | 227 | 228 | 229 |
Eric S
Eric S
🚇💻 ⚠️ 📖		M̵̞̗̝̼̅̏̎͝Ȯ̴̝̻̊̃̋̀Õ̷̼͋N̸̩̿͜ ̶̜̠̹̼̩͒
M̵̞̗̝̼̅̏̎͝Ȯ̴̝̻̊̃̋̀Õ̷̼͋N̸̩̿͜ ̶̜̠̹̼̩͒
🚇 💻 ⚠️ 📖
230 | 231 | 232 | 233 | 234 | 235 | 236 | ## Sponsors 237 | 238 |

239 |

Deep AI Research is sponsored by the following organizations:
240 |

241 |

242 |

DeepAI
243 |

244 |

245 |

Interested in working with us? Join our Discord or post an issue to get in touch.
246 |

-------------------------------------------------------------------------------- /data/datasets.json: -------------------------------------------------------------------------------- 1 | { 2 | "models": [ 3 | "cap3d_captions" 4 | ], 5 | "backgrounds": ["hdri_data"] 6 | } -------------------------------------------------------------------------------- /data/object_data.json: -------------------------------------------------------------------------------- 1 | { 2 | "object_list_intro": [ 3 | "The scene features ", 4 | "Present in this scene are ", 5 | "Objects visible in the environment include ", 6 | "The composition contains ", 7 | "Arranged within the scene are ", 8 | "Key elements in this setting comprise ", 9 | "The visual space is occupied by ", 10 | "Notable items in the scene are ", 11 | "The frame encompasses ", 12 | "Captured in this view are ", 13 | "The scene is populated with ", 14 | "Central to this composition are ", 15 | "Prominently displayed are ", 16 | "The environment showcases ", 17 | "Focal points of the scene include ", 18 | "Among the objects present are ", 19 | "The visual narrative includes ", 20 | "Constituting the scene's elements are ", 21 | "Depicted in this setting are ", 22 | "The scene is characterized by the presence of " 23 | ], 24 | "descriptions": [ 25 | "The subject is ", 26 | "The focus is on ", 27 | "The view is focused on ", 28 | "The camera is pointed at ", 29 | "The frame is centered on ", 30 | "Featuring ", 31 | "There are \"\" in the scene", 32 | "The scene contains \"\"", 33 | "Put the objects \"\" into the scene", 34 | "Add the objects \"\" to the scene", 35 | "Focus on ", 36 | "" 37 | ], 38 | "movement_speed_description": { 39 | "0.25": [ 40 | "moderate", 41 | "average", 42 | "normal", 43 | "standard", 44 | "typical", 45 | "regular", 46 | "usual", 47 | "slowly", 48 | "gradually" 49 | ], 50 | "0.5": [ 51 | "fast", 52 | "quick", 53 | "rapid", 54 | "swift", 55 | "speedy", 56 | "hasty", 57 | "brisk", 58 | "hurried", 59 | "hurriedly", 60 | "hastily", 61 | "promptly", 62 | "quickly", 63 | "rapidly", 64 | "speedily", 65 | "swiftly", 66 | "fleetly", 67 | "briskly" 68 | ] 69 | }, 70 | "movement_description_relationship": [ 71 | "The is traveling at per second.", 72 | "Watch as the moves at units/sec.", 73 | "The shifts at a pace of each second.", 74 | "Observe the heading at a speed of per second.", 75 | "The drifts at units.", 76 | " movement of the is at .", 77 | "The advances at units/s.", 78 | "Moving at , the progresses.", 79 | "The proceeds to the at a speed of .", 80 | "The glides at a velocity of each second.", 81 | "The moves at units/sec.", 82 | "With a speed of , the moves .", 83 | "The travels at a rate of per second.", 84 | "The transitions at units/s.", 85 | "The heads at .", 86 | "At , the shifts .", 87 | "The moves at units/s towards the .", 88 | "The slides at .", 89 | "Watch the go at units.", 90 | "The navigates at a speed of per second.", 91 | "The maneuvers at .", 92 | "Moving to the at , the advances.", 93 | "The shifts to the at each second.", 94 | "The proceeds at units.", 95 | "The glides to the at a velocity of per second.", 96 | "The heads to the at .", 97 | "Traveling at , the progresses.", 98 | "The drifts to the at a pace of units/sec.", 99 | "Moving at , the travels.", 100 | "The advances at units.", 101 | "At a speed of per second, the moves .", 102 | "The transitions to the at units.", 103 | "The navigates at .", 104 | "The shifts at units.", 105 | "The maneuvers at a rate of per second.", 106 | "The heads at units.", 107 | "The slides to the at .", 108 | "The glides at each second.", 109 | "The proceeds at a speed of .", 110 | "The moves at a velocity of .", 111 | "Moving to the at , the proceeds.", 112 | "The drifts at units.", 113 | "At , the transitions .", 114 | "The travels to the at .", 115 | "The heads to the at a speed of units.", 116 | "Moving at a rate of , the advances.", 117 | "The maneuvers at units.", 118 | "At a pace of , the travels .", 119 | "The slides at .", 120 | "The moves at units." 121 | ], 122 | "ontop_description_relationship": [ 123 | " is on top of ", 124 | " is below ", 125 | " is stacked above ", 126 | " is underneath ", 127 | " is placed over ", 128 | " is positioned under ", 129 | " is resting on ", 130 | " is supporting ", 131 | " is situated atop ", 132 | " is lying beneath ", 133 | " is layered above ", 134 | " is seated under ", 135 | " is on the top of ", 136 | " is at the bottom of ", 137 | " is elevated over ", 138 | " is settled below ", 139 | " is placed on ", 140 | " is arranged below " 141 | ], 142 | "name_description_relationship": [ 143 | "There is a ", 144 | "The is present in the scene", 145 | " can be observed", 146 | "The is a notable element", 147 | " stands out in the image", 148 | "The scene features a ", 149 | "A is visible", 150 | "The ", 151 | " is ", 152 | " ", 153 | " of proportions", 154 | " is in scale", 155 | "A version of ", 156 | "The , in dimension", 157 | " measures ", 158 | "A tall ", 159 | "The , standing at ", 160 | "There is a high ", 161 | "The reaches in height", 162 | " ", 163 | " spans ", 164 | "A long ", 165 | "The , measuring ", 166 | "There is a wide ", 167 | "The extends ", 168 | " ", 169 | " occupies space" 170 | ], 171 | "relationships": { 172 | "to_the_left": [ 173 | "is to the left of", 174 | "is left of", 175 | "is beside", 176 | "is on the left side of", 177 | "is left-side of", 178 | "is on the left of" 179 | ], 180 | "to_the_right": [ 181 | "is to the right of", 182 | "is right of", 183 | "is beside", 184 | "is next to", 185 | "is adjacent to", 186 | "is right-side of", 187 | "is on the right of" 188 | ], 189 | "in_front_of": [ 190 | "is in front of", 191 | "is front of", 192 | "is before", 193 | "is ahead of" 194 | ], 195 | "behind": [ 196 | "is behind", 197 | "is to the back of", 198 | "is in back of", 199 | "is after", 200 | "is at the rear of", 201 | "is to the rear of" 202 | ] 203 | }, 204 | "scales": { 205 | "tiny": { 206 | "names": [ 207 | "tiny", 208 | "mini", 209 | "miniscule", 210 | "minature", 211 | "petite", 212 | "extra small" 213 | ], 214 | "factor": 0.25 215 | }, 216 | "small": { 217 | "names": [ 218 | "small", 219 | "little", 220 | "mini", 221 | "petite", 222 | "miniature", 223 | "short", 224 | "undersized", 225 | "smaller" 226 | ], 227 | "factor": 0.5 228 | }, 229 | "small-medium": { 230 | "names": [ 231 | "small-medium", 232 | "medium-small", 233 | "medium-sized", 234 | "average sized", 235 | "below average sized", 236 | "smaller than average", 237 | "smaller than usual", 238 | "shorter than normal", 239 | "small-ish" 240 | ], 241 | "factor": 0.75 242 | }, 243 | "medium": { 244 | "names": [ 245 | "medium", 246 | "average", 247 | "normal", 248 | "standard", 249 | "typical", 250 | "regular", 251 | "standard", 252 | "usual" 253 | ], 254 | "factor": 1.0 255 | }, 256 | "medium-large": { 257 | "names": [ 258 | "medium-large", 259 | "large-medium", 260 | "large-ish", 261 | "moderate-large", 262 | "mid-large", 263 | "considerable", 264 | "generous", 265 | "expansive", 266 | "sizable", 267 | "hefty", 268 | "bulky" 269 | ], 270 | "factor": 1.5 271 | }, 272 | "large": { 273 | "names": [ 274 | "large", 275 | "big", 276 | "huge", 277 | "massive", 278 | "giant", 279 | "tall" 280 | ], 281 | "factor": 2.25 282 | }, 283 | "huge": { 284 | "names": [ 285 | "huge", 286 | "towering", 287 | "massive", 288 | "giant", 289 | "gigantic", 290 | "enormous", 291 | "colossal", 292 | "really big", 293 | "really tall", 294 | "really large", 295 | "very big", 296 | "very tall", 297 | "very large", 298 | "extra large" 299 | ], 300 | "factor": 3.0 301 | } 302 | } 303 | } -------------------------------------------------------------------------------- /data/stage_data.json: -------------------------------------------------------------------------------- 1 | { 2 | "material_names": [ 3 | "floor material", 4 | "floor texture", 5 | "floor", 6 | "ground material", 7 | "ground texture", 8 | "ground", 9 | "stage material", 10 | "stage texture", 11 | "stage", 12 | "flooring", 13 | "flooring material", 14 | "flooring texture" 15 | ], 16 | "background_names": [ 17 | "background", 18 | "background scene", 19 | "background view", 20 | "background setting", 21 | "background environment", 22 | "backdrop", 23 | "scene", 24 | "setting", 25 | "envionrment", 26 | "view", 27 | "panorama", 28 | "landscape", 29 | "scenery" 30 | ] 31 | } -------------------------------------------------------------------------------- /docs/assets/DeepAI.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepAI-Research/Simverse/73a463c49ae8c6dbc6ff14ea1dcae1ba76b39ecf/docs/assets/DeepAI.png -------------------------------------------------------------------------------- /docs/assets/favicon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepAI-Research/Simverse/73a463c49ae8c6dbc6ff14ea1dcae1ba76b39ecf/docs/assets/favicon.ico -------------------------------------------------------------------------------- /docs/assets/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepAI-Research/Simverse/73a463c49ae8c6dbc6ff14ea1dcae1ba76b39ecf/docs/assets/logo.png -------------------------------------------------------------------------------- /docs/assets/simian_banner.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepAI-Research/Simverse/73a463c49ae8c6dbc6ff14ea1dcae1ba76b39ecf/docs/assets/simian_banner.png -------------------------------------------------------------------------------- /docs/backgrounds.md: -------------------------------------------------------------------------------- 1 | # Backgrounds 2 | 3 | Backgrounds are loaded from datasources, such as Poly Haven. The `background` module is responsible for managing the backgrounds that are loaded into the scene. It handles loading the backgrounds, setting their position, scale and orientation, as well as how materials and textures are handled. 4 | 5 | ::: simian.background 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/batch_rendering.md: -------------------------------------------------------------------------------- 1 | # Batch Rendering 2 | 3 | Local batch processing is handled by the `batch` module. This module is responsible for generating videos in bulk. It calls the `render` module to generate videos, iterating over all combinations from the supplied start index to the end index. 4 | 5 | ::: simian.batch 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/cameras.md: -------------------------------------------------------------------------------- 1 | # Cameras 2 | 3 | The `camera` module is responsible for managing the camera in Blender. It handles setting up the camera, setting the camera position, orientation, and field of view. 4 | 5 | ::: simian.camera 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/combiner.md: -------------------------------------------------------------------------------- 1 | # Combiner 2 | 3 | The `combiner` module is responible for creating the combinations.json file which is used to store the combinations of assets that will be used to render the final video. It handles reading the assets from the assets directory, creating the combinations, and writing the combinations to the combinations.json file. 4 | 5 | ::: simian.combiner 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/getting_started.md: -------------------------------------------------------------------------------- 1 | # Getting Started 2 | 3 | Below are some quick notes to get you up and running. Please read through the rest of the documentation for more detailed information. 4 | 5 | ## 🖥️ Setup 6 | 7 | > **_NOTE:_** Simian requires Python 3.11. 8 | 9 | 1. Install dependences: 10 | ```bash 11 | pip install -r requirements.txt 12 | ``` 13 | 14 | 2. Download the datasets: 15 | ```bash 16 | ./scripts/data/get_data.sh 17 | ``` 18 | 19 | 3. [OPTIONAL] If you're on a headless Linux server, install Xorg and start it: 20 | 21 | ```bash 22 | sudo apt-get install xserver-xorg -y && \ 23 | sudo python3 scripts/start_x_server.py start 24 | ``` 25 | 26 | ## 📸 Usage 27 | 28 | ### Generating Combinations 29 | 30 | ```bash 31 | python3 -m simian.combiner --count 1000 --seed 42 32 | ``` 33 | 34 | ### Generating Videos 35 | To generate a video(s): 36 | ```bash 37 | python3 -m simian.batch --start_index 0 --end_index 1000 --width 1024 --height 576 --start_frame 1 --end_frame 65 38 | ``` 39 | 40 | To generate an image(s): 41 | ```bash 42 | python3 -m simian.batch --start_index 0 --end_index 1000 --width 1024 --height 576 --start_frame 1 --end_frame 65 43 | ``` 44 | 45 | You can generate individually: 46 | ```bash 47 | # MacOS 48 | python -m simian.render 49 | 50 | # Linux 51 | python -m simian.render 52 | 53 | ## Kitchen sink 54 | python -m simian.render -- --width 1920 --height 1080 --combination_index 0 --output_dir ./renders --hdri_path ./backgrounds --start_frame 1 --end_frame 65 55 | ``` 56 | 57 | Configure the flags as needed: 58 | - `--width` and `--height` are the resolution of the video. 59 | - `--combination_index` is the index of the combination to render. 60 | - `--output_dir` is the directory to save the rendered video. 61 | - `--hdri_path` is the directory containing the background images. 62 | - `--start_frame` and `--end_frame` are the start and end frames of the video. 63 | 64 | Or generate all or part of the combination set using the `batch.py` script: 65 | 66 | ```bash 67 | python3 -m simian.batch --start_index 0 --end_index 1000 --width 1024 --height 576 --start_frame 1 --end_frame 65 68 | ``` 69 | 70 | ### Clean up Captions 71 | 72 | Make captions more prompt friendly. 73 | 74 | > **_NOTE:_** Create a .env file and add your OpenAI API key 75 | ```bash 76 | python3 scripts/rewrite_captions.py 77 | ``` 78 | 79 | ### Distributed rendering 80 | Rendering can be distributed across multiple machines using the "simian.py" and "worker.py" scripts. 81 | 82 | You will need to set up Redis and obtain Huggingface API key to use this feature. 83 | 84 | #### Set Up Redis 85 | You can make a free Redis account [here](https://redis.io/try-free/). 86 | 87 | For local testing and multiple local workers, you can use the following script to set up a local instance of Redis: 88 | ```bash 89 | scripts/setup_redis.sh 90 | ``` 91 | 92 | #### Huggingface API Key 93 | 94 | You can get a Huggingface API key [here](https://huggingface.co/settings/tokens). 95 | 96 | Now, start your workers 97 | ```bash 98 | export REDIS_HOST=.com 99 | export REDIS_PORT=1337 100 | export REDIS_USER=default 101 | export REDIS_PASSWORD= 102 | export HF_TOKEN= 103 | export HF_REPO_ID= 104 | celery -A simian.worker worker --loglevel=info 105 | ``` 106 | 107 | You can also build and run the worker with Docker 108 | ```bash 109 | # build the container 110 | docker build -t simian-worker . 111 | 112 | # run the container with .env 113 | docker run --env-file .env simian-worker 114 | 115 | # run the container with environment variables 116 | docker run -e REDIS_HOST={myhost} -e REDIS_PORT={port} -e REDIS_USER=default -e REDIS_PASSWORD={some password} -e HF_TOKEN={token} -e HF_REPO_ID={repo_id} simian-worker 117 | ``` 118 | 119 | Finally, issue work to your task queue 120 | 121 | ```bash 122 | python3 -m simian.distributed --width 1024 --height 576 123 | ``` 124 | 125 | If you want to use a custom or hosted Redis instance (recommended), you can add the redis details like this: 126 | ```bash 127 | export REDIS_HOST=.com 128 | export REDIS_PORT=1337 129 | export REDIS_USER=default 130 | export REDIS_PASSWORD= 131 | ``` 132 | 133 | To run tests look into the test folder and run whichever test file you want 134 | 135 | ```bash 136 | python object_test.py 137 | ``` -------------------------------------------------------------------------------- /docs/index.md: -------------------------------------------------------------------------------- 1 | # Simian 2 | 3 | 4 | 5 | Simian is a synthetic data generation tool for creating image caption or video caption pairs. What is novel about Simian is that it focuses on high-quality, specific captioning for features that are difficult to capture in current image and video description models-- for example camera movement, actor intent or physical interactions. 6 | 7 | ## Welcome 8 | 9 | This documentation is intended to help you understand the structure of the Simian codebase and how to use it to generate synthetic data for your projects. Please select a topic from the sidebar to get started. 10 | 11 | ## Core Use Cases 12 | Use Simian data to train image, video and 3D models. 13 | 14 | - **Image and Video Generation**: Generate synthetic data for training image and video captioning models. Use Simian to create synthetic data for training image and video captioning models on real data. Mix with real data to add high-level controllability to your model. 15 | 16 | - **3D Model Generation**: Generate random views or arranged spherical views for training 3D diffusion models. 17 | 18 | - **Image and Video Description**: Use synthetic data to improve the quality of your captioning, descriptive or contrastive model on real data. 19 | 20 | - **3D Object Detection**: Generate synthetic data for training 3D object detection models. Use Simian to create synthetic data for training 3D object detection models on real data. 21 | 22 | ## Getting Started 23 | 24 | Please visit the [Getting Started](/getting_started) page to learn how to set up your environment and generate synthetic data. -------------------------------------------------------------------------------- /docs/object.md: -------------------------------------------------------------------------------- 1 | # Object 2 | 3 | The `object` module contains classes and functions for managing objects in the scene. It handles loading objects from datasources, setting their position, scale, and orientation, as well as how materials and textures are handled. 4 | 5 | Objects are loaded from datasources, such as Objaverse. The `object` module is responsible for managing the objects that are loaded into the scene. It handles loading the objects, setting their position, scale and orientation, as well as how materials and textures are handled. 6 | 7 | ::: simian.object 8 | :docstring: 9 | :members: 10 | :undoc-members: 11 | :show-inheritance: -------------------------------------------------------------------------------- /docs/postprocessing.md: -------------------------------------------------------------------------------- 1 | # Postprocessing 2 | 3 | The `postprocessing` module is responsible for applying postprocessing effects to the rendered videos. It handles adding effects such as color correction, vignetting, and lens distortion to the videos. 4 | 5 | ::: simian.postprocessing 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/rendering.md: -------------------------------------------------------------------------------- 1 | # Rendering 2 | 3 | Scene assembly and generation from individual combinations is handled by the `render` module. This module is responsible for creating the scene, setting up the camera, and rendering the scene. Most of the other modules are imported into this module. 4 | 5 | ::: simian.render 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/scene.md: -------------------------------------------------------------------------------- 1 | # Scene 2 | 3 | The `scene` module is responsible for managing the scene in the Simian application. It handles creating the scene, setting up the camera, and rendering the scene. Most of the other modules are imported into this module. 4 | 5 | ::: simian.scene 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/transform.md: -------------------------------------------------------------------------------- 1 | # Transform 2 | 3 | The `transform` module is responsible for applying transformations to the assets in the scene. It handles scaling, rotating, and translating the assets to create the desired effect. 4 | 5 | ::: simian.transform 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /docs/worker.md: -------------------------------------------------------------------------------- 1 | # Worker 2 | 3 | The `worker` module is responsible for managing the worker threads in a distributed rendering setup. It handles creating and managing worker threads, sending tasks to the workers, and collecting the results. 4 | 5 | ::: simian.worker 6 | :docstring: 7 | :members: 8 | :undoc-members: 9 | :show-inheritance: -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: Simian 2 | theme: 3 | name: material 4 | logo: assets/logo.png 5 | favicon: assets/favicon.ico 6 | plugins: 7 | - search 8 | - autorefs 9 | - mkdocstrings: 10 | enabled: true 11 | default_handler: python -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | omegaconf 2 | requests 3 | argparse 4 | numpy 5 | scipy 6 | distributask 7 | bpy 8 | rich 9 | chromadb 10 | boto3 11 | sentence_transformers 12 | questionary 13 | google-generativeai 14 | datasets -------------------------------------------------------------------------------- /scripts/data/get_data.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | git clone https://github.com/RaccoonResearch/simdata 3 | mv simdata/datasets datasets 4 | mv simdata/examples examples 5 | rm -rf simdata -------------------------------------------------------------------------------- /scripts/data/get_objaverse_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import objaverse 4 | import objaverse.xl as oxl 5 | import json 6 | 7 | import requests 8 | 9 | annotations = oxl.get_annotations( 10 | download_dir="./", 11 | ) 12 | 13 | annotations = objaverse.load_annotations() 14 | 15 | # save annotations to json. annotations is a dict 16 | with open("annotations.json", "w") as file: 17 | json.dump(annotations, file, indent=4) 18 | 19 | annotations = json.load(open("annotations.json")) 20 | 21 | categories = set() 22 | for model in annotations.values(): 23 | for category in model["categories"]: 24 | categories.add(category["name"]) 25 | 26 | for category in categories: 27 | category_data = [ 28 | model 29 | for model in annotations.values() 30 | if category in [cat["name"] for cat in model["categories"]] 31 | ] 32 | with open(f"data/{category}.json", "w") as file: 33 | json.dump(category_data, file, indent=4) 34 | 35 | files = os.listdir("./datasets") 36 | 37 | model_count = 0 38 | 39 | for file in files: 40 | data = json.load(open(f"./datasets/{file}")) 41 | 42 | # Filter data entries ensuring all required fields are present and meet criteria 43 | filtered_data = [] 44 | for model in data: 45 | # Check for presence of all necessary keys and that their values meet criteria 46 | if ( 47 | "faceCount" in model 48 | and 50 < model["faceCount"] < 1000000 49 | and "vertexCount" in model 50 | and 50 < model["vertexCount"] < 1000000 51 | and "tags" in model 52 | and not any( 53 | tag["name"].lower().find("scan") != -1 54 | or tag["name"].lower().find("photogrammetry") != -1 55 | for tag in model["tags"] 56 | ) 57 | and "archives" in model 58 | and not any( 59 | archive_data.get("textureCount") == 0 60 | for archive_data in model["archives"].values() 61 | ) 62 | ): 63 | filtered_data.append(model) 64 | 65 | annotations_slim = [ 66 | { 67 | "uid": model["uid"], 68 | "name": model["name"], 69 | "categories": [cat["name"] for cat in model["categories"]], 70 | "description": model["description"], 71 | "tags": [tag["name"] for tag in model["tags"]], 72 | } 73 | for model in filtered_data 74 | ] 75 | model_count += len(annotations_slim) 76 | # Overwrite the file with filtered data 77 | with open(f"./datasets/{file}", "w") as file: 78 | json.dump(annotations_slim, file, indent=4) 79 | 80 | print(f"Extracted data for {model_count} models") 81 | 82 | files = os.listdir("./datasets") 83 | for file in files: 84 | data = json.load(open(f"./datasets/{file}")) 85 | 86 | # for each entry get https://api.sketchfab.com/v3/models/ 87 | # get the "description" field and save it to the existing entry 88 | for model in data: 89 | uid = model["uid"] 90 | request = requests.get(f"https://api.sketchfab.com/v3/models/{uid}") 91 | model_data = request.json() 92 | if "description" in model_data: 93 | model["description"] = model_data["description"] 94 | print(f"got description for {uid}: {model_data['description']}") 95 | else: 96 | print(f"no description for {uid}") 97 | model["description"] = "" 98 | time.sleep(0.1) 99 | 100 | with open(f"./datasets/{file}", "w") as file: 101 | json.dump(data, file, indent=4) 102 | -------------------------------------------------------------------------------- /scripts/data/get_polyhaven_background_data.py: -------------------------------------------------------------------------------- 1 | import requests 2 | import json 3 | import os 4 | 5 | # JSON file path 6 | json_file = "../hdri_data.json" 7 | 8 | # Check if the JSON file exists, otherwise create an empty dictionary 9 | if os.path.exists(json_file): 10 | with open(json_file, "r") as file: 11 | hdri_data = json.load(file) 12 | else: 13 | hdri_data = {} 14 | 15 | # Step 1: GET request to retrieve the list of HDRIs 16 | url = "https://api.polyhaven.com/assets?type=hdris" 17 | response = requests.get(url) 18 | data = response.json() 19 | 20 | # Step 2: Iterate over the keys (IDs) of the JSON object 21 | for id, asset_data in data.items(): 22 | # Check if the ID is already in the JSON file 23 | if id in hdri_data: 24 | print(f"Skipping ID: {id} (already exists in the JSON file)") 25 | continue 26 | 27 | # Step 3: Visit the URL for each ID 28 | file_url = f"https://api.polyhaven.com/files/{id}" 29 | file_response = requests.get(file_url) 30 | file_data = file_response.json() 31 | 32 | # Step 4: Extract the URL associated with the HDR file inside the 8k hdri key'd value 33 | hdr_url = file_data["hdri"]["8k"]["hdr"]["url"] 34 | 35 | # Step 5: Extract the name, categories, and tags from the asset data 36 | name = asset_data["name"] 37 | categories = asset_data["categories"] 38 | tags = asset_data["tags"] 39 | 40 | # Step 6: Add the ID, URL, name, categories, and tags to the dictionary 41 | hdri_data[id] = { 42 | "url": hdr_url, 43 | "name": name, 44 | "categories": categories, 45 | "tags": tags, 46 | } 47 | 48 | # Save the updated dictionary to the JSON file 49 | with open(json_file, "w") as file: 50 | json.dump(hdri_data, file, indent=4) 51 | 52 | print(f"Added ID: {id}, URL: {hdr_url}, Name: {name}") 53 | 54 | print("JSON file updated successfully.") 55 | -------------------------------------------------------------------------------- /scripts/data/get_polyhaven_texture_data.py: -------------------------------------------------------------------------------- 1 | import requests 2 | import json 3 | import os 4 | 5 | if not os.path.exists("./datasets"): 6 | os.makedirs("./datasets") 7 | 8 | json_file = "./datasets/texture_data.json" 9 | 10 | if os.path.exists(json_file): 11 | with open(json_file, "r") as file: 12 | texture_data = json.load(file) 13 | else: 14 | texture_data = {} 15 | 16 | url = "https://api.polyhaven.com/assets?type=textures" 17 | response = requests.get(url) 18 | data = response.json() 19 | 20 | for id, asset_data in data.items(): 21 | if id in texture_data: 22 | print(f"Skipping ID: {id} (already exists in the JSON file)") 23 | continue 24 | 25 | file_url = f"https://api.polyhaven.com/files/{id}" 26 | file_response = requests.get(file_url) 27 | file_data = file_response.json() 28 | 29 | texture_maps = {} 30 | 31 | map_types = file_data.keys() 32 | 33 | for map_type in map_types: 34 | preferred_resolutions = ["4k", "2k"] # Prioritize higher resolutions first 35 | file_formats = ["jpg", "png"] # Common file formats 36 | 37 | for resolution in preferred_resolutions: 38 | if resolution in file_data[map_type]: 39 | for file_format in file_formats: 40 | if file_format in file_data[map_type][resolution]: 41 | texture_maps[map_type] = file_data[map_type][resolution][ 42 | file_format 43 | ]["url"] 44 | break # Stop checking once the first applicable format is found 45 | if map_type in texture_maps: 46 | break # Break the outer loop if a URL has been found 47 | 48 | name = asset_data["name"] 49 | categories = asset_data["categories"] 50 | tags = asset_data["tags"] 51 | 52 | texture_data[id] = { 53 | "maps": texture_maps, 54 | "name": name, 55 | "categories": categories, 56 | "tags": tags, 57 | } 58 | 59 | with open(json_file, "w") as file: 60 | json.dump(texture_data, file, indent=4) 61 | 62 | print(f"Added ID: {id}, Name: {name}") 63 | 64 | print("JSON file updated successfully.") 65 | -------------------------------------------------------------------------------- /scripts/demo/demo.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import random 4 | import logging 5 | 6 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 7 | logger = logging.getLogger(__name__) 8 | 9 | def load_combinations(file_path: str): 10 | logger.debug(f"Loading combinations from file: {file_path}") 11 | try: 12 | with open(file_path, "r") as file: 13 | return json.load(file) 14 | except Exception as e: 15 | logger.error(f"Failed to load combinations from file: {file_path} with error: {str(e)}") 16 | return None 17 | 18 | def extract_random_indices(data, count=5): 19 | logger.debug(f"Extracting {count} random indices from data") 20 | combinations = data["combinations"] 21 | total_count = len(combinations) 22 | random_indices = random.sample(range(total_count), count) 23 | extracted_combinations = [combinations[i] for i in random_indices] 24 | return extracted_combinations 25 | 26 | def save_to_file(data, file_path: str): 27 | try: 28 | with open(file_path, "w") as file: 29 | json.dump(data, file, indent=4) 30 | logger.info(f"Saved extracted combinations to {file_path}") 31 | except Exception as e: 32 | logger.error(f"Failed to save combinations to file: {file_path} with error: {str(e)}") 33 | 34 | if __name__ == "__main__": 35 | # Path to the combinations file 36 | combinations_file_path = "combinations.json" 37 | 38 | # Load combinations from the file 39 | if os.path.exists(combinations_file_path): 40 | data = load_combinations(combinations_file_path) 41 | if data: 42 | extracted_combinations = extract_random_indices(data) 43 | output_path = "demo.json" 44 | save_to_file({"combinations": extracted_combinations}, output_path) 45 | else: 46 | logger.error(f"Failed to load data from {combinations_file_path}") 47 | else: 48 | logger.error(f"{combinations_file_path} file not found") 49 | -------------------------------------------------------------------------------- /scripts/filter/caption_combination_pairs.py: -------------------------------------------------------------------------------- 1 | import json 2 | import jsonlines 3 | 4 | def create_finetuning_data(input_file, output_file): 5 | # Read the input JSON file 6 | with open(input_file, 'r') as f: 7 | data = json.load(f) 8 | 9 | # Create input-output pairs for fine-tuning 10 | finetuning_data = [] 11 | if 'combinations' in data: 12 | for combination in data['combinations']: 13 | caption = combination.get('caption', '') 14 | if caption: 15 | finetuning_data.append({ 16 | "input": caption, 17 | "output": json.dumps(combination) # Convert the entire combination to a JSON string 18 | }) 19 | 20 | # Write the fine-tuning data to the output JSONL file 21 | with jsonlines.open(output_file, 'w') as writer: 22 | writer.write_all(finetuning_data) 23 | 24 | print(f"Fine-tuning data has been written to {output_file}") 25 | 26 | # Usage 27 | input_file = './combinations.json' 28 | output_file = 'finetune_data.jsonl' 29 | create_finetuning_data(input_file, output_file) 30 | -------------------------------------------------------------------------------- /scripts/filter/combinations_add_placeholder.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | def replace_background_values(background): 4 | background["name"] = "" 5 | background["url"] = "" 6 | background["id"] = "" 7 | background["from"] = "" 8 | 9 | def replace_object_ids(objects): 10 | for obj in objects: 11 | obj["uid"] = "" 12 | 13 | def replace_map_values(maps): 14 | for key in maps: 15 | maps[key] = f"<{key.lower()}_placeholder>" 16 | 17 | def process_combinations(data): 18 | for combination in data["combinations"]: 19 | if "objects" in combination: 20 | replace_object_ids(combination["objects"]) 21 | 22 | if "background" in combination: 23 | replace_background_values(combination["background"]) 24 | 25 | if "stage" in combination and "material" in combination["stage"]: 26 | material = combination["stage"]["material"] 27 | if "maps" in material: 28 | replace_map_values(material["maps"]) 29 | 30 | def main(): 31 | input_file = "./combinations.json" 32 | output_file = "combinations_processed.json" 33 | 34 | try: 35 | with open(input_file, 'r') as file: 36 | data = json.load(file) 37 | 38 | process_combinations(data) 39 | 40 | with open(output_file, 'w') as file: 41 | json.dump(data, file, indent=4) 42 | 43 | print(f"Processing complete. Output saved to {output_file}") 44 | 45 | except FileNotFoundError: 46 | print(f"Error: The file {input_file} was not found.") 47 | except json.JSONDecodeError: 48 | print(f"Error: The file {input_file} is not a valid JSON file.") 49 | except Exception as e: 50 | print(f"An unexpected error occurred: {str(e)}") 51 | 52 | if __name__ == "__main__": 53 | main() -------------------------------------------------------------------------------- /scripts/filter/combinations_to_captions.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | import argparse 4 | 5 | 6 | def main(json_file, video_folder=None): 7 | # Load the combinations.json file 8 | with open(json_file, "r") as f: 9 | data = json.load(f) 10 | 11 | combinations = data.get("combinations", []) 12 | 13 | captions = [] 14 | 15 | if video_folder: 16 | # List and sort all .mp4 and .png files in the video folder based on numerical index 17 | video_files = [f for f in os.listdir(video_folder) if f.endswith(".mp4")] 18 | image_files = [f for f in os.listdir(video_folder) if f.endswith(".png")] 19 | 20 | video_files.sort(key=lambda x: int(os.path.splitext(x)[0].split("_")[0])) 21 | image_files.sort(key=lambda x: int(os.path.splitext(x)[0].split("_")[0])) 22 | 23 | # Debugging: Print files count and sorted lists 24 | print(f"Found {len(video_files)} video files.") 25 | print(f"Sorted video files: {video_files}") 26 | print(f"Found {len(image_files)} image files.") 27 | print(f"Sorted image files: {image_files}") 28 | 29 | # Process each combination, ensuring we do not exceed the number of available videos or images 30 | for i, obj in enumerate(combinations): 31 | if i >= len(video_files) and i >= len(image_files): 32 | print( 33 | f"Warning: Only {len(video_files)} videos and {len(image_files)} images found. Stopping at index {i}." 34 | ) 35 | break 36 | 37 | # Access the 'caption' field 38 | caption = obj.get("caption") 39 | 40 | # Ensure the required field is present 41 | if caption: 42 | # Create a new object with the desired format 43 | caption_obj = { 44 | "cap": [caption], 45 | } 46 | if i < len(video_files): 47 | caption_obj["path"] = f"{video_folder}/{video_files[i]}" 48 | if i < len(image_files): 49 | caption_obj["path"] = f"{video_folder}/{image_files[i]}" 50 | 51 | captions.append(caption_obj) 52 | else: 53 | print(f"Missing caption in combination at index {i}: {obj}") 54 | 55 | else: 56 | # Process each combination without video or image paths 57 | for i, obj in enumerate(combinations): 58 | # Access the 'caption' field 59 | caption = obj.get("caption") 60 | 61 | # Ensure the required field is present 62 | if caption: 63 | # Create a new object with the desired format 64 | caption_obj = {"cap": [caption]} 65 | captions.append(caption_obj) 66 | else: 67 | print(f"Missing caption in combination at index {i}: {obj}") 68 | 69 | # Write the captions array to captions.json file 70 | with open("captions.json", "w") as out_f: 71 | json.dump(captions, out_f, indent=2) 72 | 73 | 74 | # Run the script with the required arguments 75 | # python scripts/combinations_to_captions.py --json_file ./combinations.json [--video_folder ./renders] 76 | if __name__ == "__main__": 77 | parser = argparse.ArgumentParser( 78 | description="Convert combinations to captions format." 79 | ) 80 | parser.add_argument( 81 | "--json_file", type=str, required=True, help="Path to combinations.json" 82 | ) 83 | parser.add_argument("--video_folder", type=str, help="Path to videos folder") 84 | args = parser.parse_args() 85 | 86 | main(args.json_file, args.video_folder) 87 | -------------------------------------------------------------------------------- /scripts/filter/combinations_to_motion_labels.py: -------------------------------------------------------------------------------- 1 | # combinations_to_labels.py 2 | import json 3 | import os 4 | import argparse 5 | 6 | 7 | def main(json_file): 8 | # Load the combinations.json file 9 | with open(json_file, "r") as f: 10 | data = json.load(f) 11 | 12 | combinations = data.get("combinations", []) 13 | 14 | # Create a list to store the labels 15 | labels = [] 16 | 17 | # Process each combination, ensuring we do not exceed the number of available videos 18 | for i, obj in enumerate(combinations): 19 | # Access the 'animation' field 20 | animation = obj.get("animation") 21 | 22 | # Use a default motion value of "NONE" if 'animation' field is missing 23 | if animation: 24 | # Extract the 'motion' value, using "NONE" as default if missing 25 | motion = animation.get("name", "NONE") 26 | else: 27 | motion = "NONE" 28 | 29 | # Convert the motion value to the appropriate enum label 30 | label = motion.upper() 31 | labels.append(label) 32 | 33 | # Write the labels to labels.txt file 34 | with open("labels.txt", "w") as out_f: 35 | for i, label in enumerate(labels): 36 | out_f.write(f"{i}.mp4 {label}\n") 37 | 38 | 39 | if __name__ == "__main__": 40 | parser = argparse.ArgumentParser( 41 | description="Convert combinations to labels format." 42 | ) 43 | parser.add_argument( 44 | "--json_file", 45 | type=str, 46 | default="combinations.json", 47 | required=False, 48 | help="Path to combinations.json", 49 | ) 50 | args = parser.parse_args() 51 | 52 | main(args.json_file) 53 | -------------------------------------------------------------------------------- /scripts/filter/get_captions.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | # Load the existing combinations.json file 4 | with open('./combinations.json', 'r') as file: 5 | data = json.load(file) 6 | 7 | # Initialize a list to store the extracted captions 8 | captions_list = [] 9 | 10 | # Extract captions for each combination 11 | for combination in data['combinations']: 12 | captions = { 13 | 'index': combination.get('index', ''), 14 | 'caption': combination.get('caption', ''), 15 | 'objects_caption': combination.get('objects_caption', ''), 16 | 'background_caption': combination.get('background_caption', ''), 17 | 'orientation_caption': combination.get('orientation_caption', ''), 18 | 'framing_caption': combination.get('framing_caption', ''), 19 | 'animation_caption': combination.get('animation_caption', ''), 20 | 'stage_caption': combination.get('stage_caption', ''), 21 | 'postprocessing_caption': combination.get('postprocessing_caption', '') 22 | } 23 | captions_list.append(captions) 24 | 25 | # Determine the number of files needed 26 | max_rows_per_file = 5000 27 | num_files = (len(captions_list) + max_rows_per_file - 1) // max_rows_per_file 28 | 29 | # Save the extracted captions to multiple JSON files 30 | for i in range(num_files): 31 | start_index = i * max_rows_per_file 32 | end_index = start_index + max_rows_per_file 33 | file_data = captions_list[start_index:end_index] 34 | with open(f'get_captions_{i + 1}.json', 'w') as file: 35 | json.dump(file_data, file, indent=4) 36 | 37 | print(f"Captions extracted and saved to {num_files} files.") 38 | -------------------------------------------------------------------------------- /scripts/filter/get_ontop_captions.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | def filter_combinations(file_path): 4 | with open(file_path, 'r') as file: 5 | data = json.load(file) 6 | 7 | original_count = len(data['combinations']) 8 | filtered_combinations = [] 9 | 10 | for combination in data['combinations']: 11 | placements = [obj['placement'] for obj in combination['objects']] 12 | if len(placements) != len(set(placements)): 13 | filtered_combinations.append(combination) 14 | 15 | data['combinations'] = filtered_combinations 16 | data['count'] = len(filtered_combinations) 17 | 18 | with open(file_path, 'w') as file: 19 | json.dump(data, file, indent=4) 20 | 21 | print(f"Original combinations: {original_count}") 22 | print(f"Filtered combinations: {len(filtered_combinations)}") 23 | print(f"Removed {original_count - len(filtered_combinations)} combinations") 24 | print(f"Updated JSON saved to {file_path}") 25 | 26 | # Usage 27 | file_path = './combinations.json' # Replace with your actual file path 28 | filter_combinations(file_path) -------------------------------------------------------------------------------- /scripts/filter/rewrite_captions_gpt.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import requests 4 | from dotenv import load_dotenv 5 | 6 | load_dotenv() 7 | MODEL = "gpt-4o" 8 | 9 | def rewrite_caption(caption_arr, context_string): 10 | split_captions = [caption.split(", ") for caption in caption_arr] 11 | caption_string = json.dumps(split_captions) 12 | 13 | content = f"{context_string}\n\n{caption_string}" 14 | 15 | print("Caption context:") 16 | print(content) 17 | 18 | headers = { 19 | "Content-Type": "application/json", 20 | "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}", 21 | } 22 | 23 | data = { 24 | "model": MODEL, 25 | "messages": [{"role": "user", "content": content}], 26 | "temperature": 0.5, 27 | "top_p": 0.8, 28 | "frequency_penalty": 0, 29 | "presence_penalty": 0, 30 | } 31 | 32 | response = requests.post( 33 | "https://api.openai.com/v1/chat/completions", headers=headers, json=data 34 | ) 35 | 36 | if response.status_code == 200: 37 | response_data = response.json() 38 | captions_content = response_data["choices"][0]["message"]["content"].strip() 39 | print("API Response:\n", captions_content) # Debug print 40 | 41 | try: 42 | # Try parsing the response as a JSON list 43 | rewritten_captions_content = json.loads(captions_content) 44 | if not isinstance(rewritten_captions_content, list): 45 | raise ValueError("Parsed JSON is not a list") 46 | except (json.JSONDecodeError, ValueError) as e: 47 | print("JSON Decode Error or Value Error:", e) 48 | # Fallback to splitting the response by newlines, assuming each line is a caption 49 | rewritten_captions_content = captions_content.split("\n") 50 | 51 | print("==== Captions rewritten by OpenAI ====") 52 | return rewritten_captions_content 53 | else: 54 | print(f"Error: {response.status_code} - {response.text}") 55 | return [] 56 | 57 | 58 | def read_combinations_and_get_array_of_just_captions(): 59 | # we want just the caption text and store in array 60 | with open("combinations.json") as file: 61 | data = json.load(file) 62 | combinations = data.get("combinations") 63 | captions = [] 64 | for obj in combinations: 65 | caption = obj.get("caption") 66 | captions.append(caption) 67 | print("==== File read, captions extracted ====") 68 | return captions 69 | 70 | 71 | def write_to_file(i, rewritten_captions): 72 | with open("combinations.json", "r+") as file: 73 | data = json.load(file) 74 | 75 | combinations = data.get("combinations") 76 | 77 | for j in range(len(rewritten_captions)): 78 | print("==> Rewriting caption: ", i + j) 79 | combinations[i + j]["caption"] = rewritten_captions[j] 80 | 81 | data["combinations"] = combinations 82 | 83 | file.seek(0) 84 | json.dump(data, file, indent=4) 85 | file.truncate() 86 | 87 | print("==== File captions rewritten, check file ====") 88 | 89 | 90 | def rewrite_captions_in_batches(combinations, context_string): 91 | batch_size = 10 92 | num_combinations = len(combinations) 93 | 94 | for i in range(0, num_combinations, batch_size): 95 | captions_batch = combinations[i : i + batch_size] 96 | rewritten_captions = rewrite_caption(captions_batch, context_string) 97 | write_to_file(i, rewritten_captions) 98 | 99 | 100 | if __name__ == "__main__": 101 | CONTEXT_STRING = """ 102 | examples: 103 | Before: "2ft Vehicle Mine GTA2 3d remake: A green and red spaceship with a circular design and a red button. 7 LOW POLY PROPS: a pink square with a small brown box and a stick on it, resembling a basketball court. Modified Colt 1911 Pistol is a gun LOW POLY PROPS is to the left of and behind Modified Colt 1911 Pistol. Modified Colt 1911 Pistol is to the right of and in front of LOW POLY PROPS. Vehicle Mine GTA2 3d remake is and in front of LOW POLY PROPS. The view is set sharply downward, looking 189\u00b0 to the right. Set the focal length of the camera to 75 mm. Semi-close perspective The panorama is Monbachtal Riverbank. The floor is Shell Floor. The scene transitions swiftly with enhanced animation speed." 104 | After: "Vehicle Mine GTA2 3d remake is a green and red spaceship with a circular design and a red button. LOW POLY PROPS is pink square with a small brown box and a stick on it, resembling a basketball court. Modified Colt 1911 Pistol is a gun LOW POLY PROPS is to the left of and behind Modified Colt 1911 Pistol. Modified Colt 1911 Pistol is to the right of and in front of LOW POLY PROPS. Vehicle Mine GTA2 3d remake is and in front of LOW POLY PROPS. The view is set sharply downward and looking 189 degrees to the right. Somewhat close perspective and the floor is a Shell Floor. Scene transitions swiftly with enhanced animation speed." 105 | 106 | Before: "Dulal Das Test File (height: 5feet) is a tan leather recliner chair and ottoman. Stylized Apple = a pink apple or peach on a plate. Stylized Apple is to the left of Dulal Das Test File. Dulal Das Test File is to the right of Stylized Apple. The lens is oriented direct right, with a 30 forward tilt. Set the fov of the camera to 32 degrees. (61.00 mm focal length) Standard medium. The scenery is Limehouse. The ground material is Gravel Floor. A standard animation speed is applied to the scene." 107 | After: "Dulal Das Test File is a tan leather recliner chair and ottoman. Stylized Apple is pink apple or peach on a plate. Dulal Das Test File is right of Stylized Apple. The lens is oriented direct right, with a slight forward tilt. Set the field of view of the camera to 32 degrees. The scenery is Limehouse with ground material of Gravel Floor. A normal animation speed for the scene." 108 | 109 | Before: "Best Japanese Curry is 1 and A bowl of stew with rice and meat. Apple is 7feet and an apple. Apple is and behind Best Japanese Curry. Best Japanese Curry is and in front of Apple. Direct the camera sharp right back, set tilt to steeply angled down. The focal length is 29 mm. Taking in the whole scene. The scene has a noticeable bloom effect. Motion blur is set to medium intensity in the scene. The backdrop is Pump House. The floor texture is Wood Plank Wall. The scene moves with a relaxed animation speed." 110 | After: "Best Japanese Curry is a bowl of stew with rice and meat. Apple is an apple placed behind Best Japanese Curry. Best Japanese Curry is in front of Apple. The camera is very angled right and very down, make sure to capture whole scene. Subtle glow effect, and medium blur when movement. Background is Pump House with a wood plank floor. Slow animation speed maybe" 111 | 112 | Before: "[Black Mercedes-Benz G-Class SUV - Black Mercedes-Benz G-Class SUV - 0.75meters] a damaged room with a kitchen and bathroom, including a refrigerator, sinks, and a toilet, in an upside-down house and destroyed building. (size: 1.5meters) a damaged room with a kitchen and bathroom, including a refrigerator, sinks, and a toilet, in an upside-down house and destroyed building. is and in front of Black Mercedes-Benz G-Class SUV. Black Mercedes-Benz G-Class SUV is and behind a damaged room with a kitchen and bathroom, including a refrigerator, sinks, and a toilet, in an upside-down house and destroyed building.. Rotate the camera to heavy left rear and tilt it 51 forward. The camera has a 17 mm focal length. Extremely wide coverage The landscape is Woods. The floor material is Blue Painted Planks. The scene has a faster animation speed of 156%." 113 | After: "there are 2 objects: Black Mercedes-Benz G-Class SUV, a damaged room with a kitchen and bathroom, including a refrigerator, sinks, and a toilet, in an upside-down house and destroyed building. the damaged room with a kitchen and bathroom is infront of Black Mercedes-Benz G-Class SUV. Black Mercedes-Benz G-Class SUV is and behind the damaged room. Rotate the camera a lot left and tilt it forward, wide angle. wood landscape. The floor is blue painted planks. fast animation speed." 114 | 115 | Above are caption example before and after. Go through array of captions and make it sound more human. Change complex director-like words like "bloom" or "pitch", change them to synonyms that are easier to understand and that any person would most likely use. 116 | Feel free to change/remove exact values like degrees. Instead of 32 degrees left you can say slightly to the left. Combine sentences maybe. 117 | Use synonyms/words to use. You can even remove some words/sentences but capture some of the holistic important details. 118 | 119 | Below are captions that NEED to be shortened/simplified, and more human-like. Return an array of rewritten captions. DO NOT wrap the strings in quotes in the array and return in format: ["caption1", "caption2", "caption3"] 120 | """ 121 | 122 | combinations = read_combinations_and_get_array_of_just_captions() 123 | rewrite_captions_in_batches(combinations, CONTEXT_STRING) 124 | -------------------------------------------------------------------------------- /scripts/filter/test_combinations.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | 4 | def find_duplicate_placements(combinations): 5 | """ 6 | Find combinations with duplicate placements. 7 | 8 | Args: 9 | combinations (list): List of combination data. 10 | 11 | Returns: 12 | list: List of index numbers with duplicate placements. 13 | """ 14 | indices_with_duplicates = [] 15 | 16 | for combination in combinations: 17 | index = combination["index"] 18 | placements = [obj["placement"] for obj in combination["objects"]] 19 | if len(placements) != len(set(placements)): 20 | indices_with_duplicates.append(index) 21 | 22 | return indices_with_duplicates 23 | 24 | 25 | # Load combinations.json file 26 | with open("combinations.json", "r") as file: # Adjust the path as needed 27 | data = json.load(file) 28 | 29 | # Find indices with duplicate placements 30 | duplicates = find_duplicate_placements(data["combinations"]) 31 | 32 | print("Indices with duplicate placements:", duplicates) 33 | -------------------------------------------------------------------------------- /scripts/setup_redis.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Function to check if Redis is installed 4 | check_redis_installed() { 5 | if command -v redis-server >/dev/null; then 6 | echo "Redis is already installed." 7 | return 0 8 | else 9 | echo "Redis is not installed." 10 | return 1 11 | fi 12 | } 13 | 14 | # Function to install Redis 15 | install_redis() { 16 | echo "Installing Redis..." 17 | if [[ "$OSTYPE" == "linux-gnu"* ]]; then 18 | # Linux installation (for Debian-based systems) 19 | sudo apt-get update 20 | sudo apt-get install redis-server -y 21 | elif [[ "$OSTYPE" == "darwin"* ]]; then 22 | # macOS installation 23 | brew install redis 24 | else 25 | echo "Unsupported OS" 26 | exit 1 27 | fi 28 | } 29 | 30 | # Function to start Redis with default configuration 31 | start_redis() { 32 | echo "Starting Redis..." 33 | redis-server --daemonize yes 34 | echo "Redis is running on port 6379 with default username and password." 35 | } 36 | 37 | # Main execution flow 38 | check_redis_installed || install_redis 39 | start_redis -------------------------------------------------------------------------------- /scripts/start_x_server.py: -------------------------------------------------------------------------------- 1 | # Taken from https://github.com/allenai/ai2thor/blob/main/scripts/ai2thor-xorg 2 | # Starts an x-server to support running Blender on a headless machine with 3 | # dedicated NVIDIA GPUs 4 | 5 | import argparse 6 | 7 | #!/usr/bin/env python3 8 | import os 9 | import platform 10 | import re 11 | import shlex 12 | import signal 13 | import subprocess 14 | import sys 15 | import time 16 | 17 | # Turning off automatic black formatting for this script as it breaks quotes. 18 | # fmt: off 19 | from typing import List 20 | 21 | PID_FILE = "/var/run/ai2thor-xorg.pid" 22 | CONFIG_FILE = "/tmp/ai2thor-xorg.conf" 23 | 24 | DEFAULT_HEIGHT = 768 25 | DEFAULT_WIDTH = 1024 26 | 27 | 28 | def process_alive(pid): 29 | """ 30 | Use kill(0) to determine if pid is alive 31 | :param pid: process id 32 | :rtype: bool 33 | """ 34 | try: 35 | os.kill(pid, 0) 36 | except OSError: 37 | return False 38 | 39 | return True 40 | 41 | 42 | def find_devices(excluded_device_ids): 43 | devices = [] 44 | id_counter = 0 45 | for r in pci_records(): 46 | if r.get("Vendor", "") == "NVIDIA Corporation" and r["Class"] in [ 47 | "VGA compatible controller", 48 | "3D controller", 49 | ]: 50 | bus_id = "PCI:" + ":".join( 51 | map(lambda x: str(int(x, 16)), re.split(r"[:\.]", r["Slot"])) 52 | ) 53 | 54 | if id_counter not in excluded_device_ids: 55 | devices.append(bus_id) 56 | 57 | id_counter += 1 58 | 59 | if not devices: 60 | print("Error: ai2thor-xorg requires at least one NVIDIA device") 61 | sys.exit(1) 62 | 63 | return devices 64 | 65 | def active_display_bus_ids(): 66 | # this determines whether a monitor is connected to the GPU 67 | # if one is, the following Option is added for the Screen "UseDisplayDevice" "None" 68 | command = "nvidia-smi --query-gpu=pci.bus_id,display_active --format=csv,noheader" 69 | active_bus_ids = set() 70 | result = subprocess.run(command, shell=True, stdout=subprocess.PIPE) 71 | if result.returncode == 0: 72 | for line in result.stdout.decode().strip().split("\n"): 73 | nvidia_bus_id, display_status = re.split(r",\s?", line.strip()) 74 | bus_id = "PCI:" + ":".join( 75 | map(lambda x: str(int(x, 16)), re.split(r"[:\.]", nvidia_bus_id)[1:]) 76 | ) 77 | if display_status.lower() == "enabled": 78 | active_bus_ids.add(bus_id) 79 | 80 | return active_bus_ids 81 | 82 | def pci_records(): 83 | records = [] 84 | command = shlex.split("lspci -vmm") 85 | output = subprocess.check_output(command).decode() 86 | 87 | for devices in output.strip().split("\n\n"): 88 | record = {} 89 | records.append(record) 90 | for row in devices.split("\n"): 91 | key, value = row.split("\t") 92 | record[key.split(":")[0]] = value 93 | 94 | return records 95 | 96 | 97 | def read_pid(): 98 | if os.path.isfile(PID_FILE): 99 | with open(PID_FILE) as f: 100 | return int(f.read()) 101 | else: 102 | return None 103 | 104 | 105 | def start(display: str, excluded_device_ids: List[int], width: int, height: int): 106 | pid = read_pid() 107 | 108 | if pid and process_alive(pid): 109 | print("Error: ai2thor-xorg is already running with pid: %s" % pid) 110 | sys.exit(1) 111 | 112 | with open(CONFIG_FILE, "w") as f: 113 | f.write(generate_xorg_conf(excluded_device_ids, width=width, height=height)) 114 | 115 | log_file = "/var/log/ai2thor-xorg.%s.log" % display 116 | error_log_file = "/var/log/ai2thor-xorg-error.%s.log" % display 117 | command = shlex.split( 118 | "Xorg -quiet -maxclients 1024 -noreset +extension GLX +extension RANDR +extension RENDER -logfile %s -config %s :%s" 119 | % (log_file, CONFIG_FILE, display) 120 | ) 121 | 122 | pid = None 123 | with open(error_log_file, "w") as error_log_f: 124 | proc = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=error_log_f) 125 | pid = proc.pid 126 | try: 127 | proc.wait(timeout=0.25) 128 | except subprocess.TimeoutExpired: 129 | pass 130 | 131 | if pid and process_alive(pid): 132 | with open(PID_FILE, "w") as f: 133 | f.write(str(proc.pid)) 134 | else: 135 | print("Error: error with command '%s'" % " ".join(command)) 136 | with open(error_log_file, "r") as f: 137 | print(f.read()) 138 | 139 | 140 | def print_config(excluded_device_ids: List[int], width: int, height: int): 141 | print(generate_xorg_conf(excluded_device_ids, width=width, height=height)) 142 | 143 | 144 | def stop(): 145 | pid = read_pid() 146 | if pid and process_alive(pid): 147 | os.kill(pid, signal.SIGTERM) 148 | 149 | for i in range(10): 150 | time.sleep(0.2) 151 | if not process_alive(pid): 152 | os.unlink(PID_FILE) 153 | break 154 | 155 | 156 | def generate_xorg_conf( 157 | excluded_device_ids: List[int], width: int, height: int 158 | ): 159 | devices = find_devices(excluded_device_ids) 160 | active_display_devices = active_display_bus_ids() 161 | 162 | xorg_conf = [] 163 | 164 | device_section = """ 165 | Section "Device" 166 | Identifier "Device{device_id}" 167 | Driver "nvidia" 168 | VendorName "NVIDIA Corporation" 169 | BusID "{bus_id}" 170 | EndSection 171 | """ 172 | server_layout_section = """ 173 | Section "ServerLayout" 174 | Identifier "Layout0" 175 | {screen_records} 176 | EndSection 177 | """ 178 | screen_section = """ 179 | Section "Screen" 180 | Identifier "Screen{screen_id}" 181 | Device "Device{device_id}" 182 | DefaultDepth 24 183 | Option "AllowEmptyInitialConfiguration" "True" 184 | Option "Interactive" "False" 185 | {extra_options} 186 | SubSection "Display" 187 | Depth 24 188 | Virtual {width} {height} 189 | EndSubSection 190 | EndSection 191 | """ 192 | screen_records = [] 193 | for i, bus_id in enumerate(devices): 194 | extra_options = "" 195 | if bus_id in active_display_devices: 196 | # See https://github.com/allenai/ai2thor/pull/990 197 | # when a monitor is connected, this option must be used otherwise 198 | # Xorg will fail to start 199 | extra_options = 'Option "UseDisplayDevice" "None"' 200 | xorg_conf.append(device_section.format(device_id=i, bus_id=bus_id)) 201 | xorg_conf.append(screen_section.format(device_id=i, screen_id=i, width=width, height=height, extra_options=extra_options)) 202 | screen_records.append( 203 | 'Screen {screen_id} "Screen{screen_id}" 0 0'.format(screen_id=i) 204 | ) 205 | 206 | xorg_conf.append( 207 | server_layout_section.format(screen_records="\n ".join(screen_records)) 208 | ) 209 | 210 | output = "\n".join(xorg_conf) 211 | return output 212 | 213 | 214 | # fmt: on 215 | 216 | if __name__ == "__main__": 217 | if os.geteuid() != 0: 218 | path = os.path.abspath(__file__) 219 | print("Executing ai2thor-xorg with sudo") 220 | args = ["--", path] + sys.argv[1:] 221 | os.execvp("sudo", args) 222 | 223 | if platform.system() != "Linux": 224 | print("Error: Can only run ai2thor-xorg on linux") 225 | sys.exit(1) 226 | 227 | parser = argparse.ArgumentParser() 228 | parser.add_argument( 229 | "--exclude-device", 230 | help="exclude a specific GPU device", 231 | action="append", 232 | type=int, 233 | default=[], 234 | ) 235 | parser.add_argument( 236 | "--width", 237 | help="width of the screen to start (should be greater than the maximum" 238 | f" width of any ai2thor instance you will start) [default: {DEFAULT_WIDTH}]", 239 | type=int, 240 | default=DEFAULT_WIDTH, 241 | ) 242 | parser.add_argument( 243 | "--height", 244 | help="height of the screen to start (should be greater than the maximum" 245 | f" height of any ai2thor instance you will start) [default: {DEFAULT_HEIGHT}]", 246 | type=int, 247 | default=DEFAULT_HEIGHT, 248 | ) 249 | parser.add_argument( 250 | "command", 251 | help="command to be executed", 252 | choices=["start", "stop", "print-config"], 253 | ) 254 | parser.add_argument( 255 | "display", help="display to be used", nargs="?", type=int, default=0 256 | ) 257 | args = parser.parse_args() 258 | if args.command == "start": 259 | start( 260 | display=args.display, 261 | excluded_device_ids=args.exclude_device, 262 | height=args.height, 263 | width=args.width, 264 | ) 265 | elif args.command == "stop": 266 | stop() 267 | elif args.command == "print-config": 268 | print_config( 269 | excluded_device_ids=args.exclude_device, 270 | width=args.width, 271 | height=args.height, 272 | ) 273 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | import sys 3 | import os 4 | 5 | # get the cwd where the setup.py file is located 6 | file_path = os.path.dirname(os.path.realpath(__file__)) 7 | 8 | long_description = "" 9 | with open(os.path.join(file_path, "README.md"), "r") as fh: 10 | long_description = fh.read() 11 | long_description = long_description.split("\n") 12 | long_description = [line for line in long_description if not "")[0].split("<")[0] for line in install_requires 25 | ] 26 | 27 | setup( 28 | name="simian3d", 29 | version=version, 30 | description="A synthetic data generator for video caption pairs.", 31 | long_description=long_description, 32 | long_description_content_type="text/markdown", 33 | url="https://github.com/RaccoonResearch/Simian", 34 | author="Raccoon Research", 35 | author_email="shawmakesmagic@gmail.com", 36 | license="MIT", 37 | packages=["simian"], 38 | install_requires=install_requires, 39 | classifiers=[ 40 | "Development Status :: 4 - Beta", 41 | "Intended Audience :: Science/Research", 42 | "License :: OSI Approved :: MIT License", 43 | "Operating System :: POSIX :: Linux", 44 | "Programming Language :: Python :: 3", 45 | "Operating System :: MacOS :: MacOS X", 46 | "Operating System :: Microsoft :: Windows", 47 | ], 48 | ) 49 | -------------------------------------------------------------------------------- /simian/__init__.py: -------------------------------------------------------------------------------- 1 | from .background import * 2 | from .render import * 3 | from .camera import * 4 | from .distributed import * 5 | from .combiner import * 6 | from .object import * 7 | from .postprocessing import * 8 | from .scene import * 9 | from .transform import * 10 | from .prompts import * 11 | from .server import * 12 | from .worker import * 13 | from .batch import * 14 | -------------------------------------------------------------------------------- /simian/background.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Dict 3 | import requests 4 | import bpy 5 | import logging 6 | 7 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 8 | logger = logging.getLogger(__name__) 9 | 10 | 11 | def get_hdri_path(hdri_path: str, combination: Dict) -> str: 12 | """ 13 | Get the local file path for the background HDR image. 14 | 15 | Args: 16 | hdri_path (str): The base directory for storing background images. 17 | combination (Dict): The combination dictionary containing background 18 | Returns: 19 | str: The local file path for the background HDR image. 20 | """ 21 | background = combination["background"] 22 | background_id = background["id"] 23 | background_from = background["from"] 24 | hdri_path = f"{hdri_path}/{background_from}/{background_id}.hdr" 25 | 26 | return hdri_path 27 | 28 | 29 | def get_background(hdri_path: str, combination: Dict) -> None: 30 | """ 31 | Download the background HDR image if it doesn't exist locally. 32 | 33 | This function checks if the background HDR image specified in the combination dictionary 34 | exists locally. If it doesn't exist, it downloads the image from the provided URL and 35 | saves it to the local file path. 36 | 37 | Args: 38 | hdri_path (str): The base directory for storing background images. 39 | combination (Dict): The combination dictionary containing background information. 40 | 41 | Returns: 42 | None 43 | """ 44 | hdri_path = get_hdri_path(hdri_path, combination) 45 | 46 | background = combination["background"] 47 | background_url = background["url"] 48 | 49 | # make sure each folder in the path exists 50 | os.makedirs(os.path.dirname(hdri_path), exist_ok=True) 51 | 52 | if not os.path.exists(hdri_path): 53 | # logger.info(f"Downloading {background_url} to {hdri_path}") 54 | response = requests.get(background_url) 55 | with open(hdri_path, "wb") as file: 56 | file.write(response.content) 57 | # else: 58 | # logger.info(f"Background {hdri_path} already exists") 59 | 60 | 61 | def set_background(hdri_path: str, combination: Dict) -> None: 62 | """ 63 | Set the background HDR image of the scene. 64 | 65 | This function sets the background HDR image of the scene using the provided combination 66 | dictionary. It ensures that the world nodes are used and creates the necessary nodes 67 | (Environment Texture, Background, and World Output) if they don't exist. It then loads 68 | the HDR image, connects the nodes, and enables the world background in the render settings. 69 | 70 | Args: 71 | hdri_path (str): The base directory for storing background images. 72 | combination (Dict): The combination dictionary containing background information. 73 | 74 | Returns: 75 | None 76 | """ 77 | get_background(hdri_path, combination) 78 | hdri_path = get_hdri_path(hdri_path, combination) 79 | 80 | # Check if the scene has a world, and create one if it doesn't 81 | if bpy.context.scene.world is None: 82 | bpy.context.scene.world = bpy.data.worlds.new("World") 83 | 84 | # Ensure world nodes are used 85 | bpy.context.scene.world.use_nodes = True 86 | tree = bpy.context.scene.world.node_tree 87 | 88 | # Clear existing nodes 89 | tree.nodes.clear() 90 | 91 | # Create the Environment Texture node 92 | env_tex_node = tree.nodes.new(type="ShaderNodeTexEnvironment") 93 | env_tex_node.location = (-300, 0) 94 | 95 | # Load the HDR image 96 | env_tex_node.image = bpy.data.images.load(hdri_path) 97 | 98 | # Create the Background node 99 | background_node = tree.nodes.new(type="ShaderNodeBackground") 100 | background_node.location = (0, 0) 101 | 102 | # Connect the Environment Texture node to the Background node 103 | tree.links.new(env_tex_node.outputs["Color"], background_node.inputs["Color"]) 104 | 105 | # Create the World Output node 106 | output_node = tree.nodes.new(type="ShaderNodeOutputWorld") 107 | output_node.location = (300, 0) 108 | 109 | # Connect the Background node to the World Output 110 | tree.links.new(background_node.outputs["Background"], output_node.inputs["Surface"]) 111 | 112 | # Enable the world background in the render settings 113 | bpy.context.scene.render.film_transparent = False 114 | 115 | # logger.info(f"Set background to {hdri_path}") 116 | 117 | 118 | def create_photosphere( 119 | hdri_path: str, combination: Dict, scale: float = 10 120 | ) -> bpy.types.Object: 121 | """ 122 | Create a photosphere object in the scene. 123 | 124 | This function creates a UV sphere object in the scene and positions it at (0, 0, 3). 125 | It smooths the sphere, inverts its normals, and renames it to "Photosphere". It then 126 | calls the `create_photosphere_material` function to create a material for the photosphere 127 | using the environment texture as emission. 128 | 129 | Args: 130 | hdri_path (str): The base directory for storing background images. 131 | combination (Dict): The combination dictionary containing background information. 132 | 133 | Returns: 134 | bpy.types.Object: The created photosphere object. 135 | """ 136 | bpy.ops.mesh.primitive_uv_sphere_add( 137 | segments=64, ring_count=32, radius=scale, location=(0, 0, 3) 138 | ) 139 | 140 | bpy.ops.object.shade_smooth() 141 | 142 | # invert the UV sphere normals 143 | bpy.ops.object.mode_set(mode="EDIT") 144 | bpy.ops.mesh.select_all(action="SELECT") 145 | bpy.ops.mesh.flip_normals() 146 | bpy.ops.object.mode_set(mode="OBJECT") 147 | 148 | sphere = bpy.context.object 149 | sphere.name = "Photosphere" 150 | sphere.data.name = "PhotosphereMesh" 151 | create_photosphere_material(hdri_path, combination, sphere) 152 | return sphere 153 | 154 | 155 | def create_photosphere_material( 156 | hdri_path: str, combination: Dict, sphere: bpy.types.Object 157 | ) -> None: 158 | """ 159 | Create a material for the photosphere object using the environment texture as emission. 160 | 161 | This function creates a new material for the provided photosphere object. It sets up 162 | the material nodes to use the environment texture as emission and assigns the material 163 | to the photosphere object. 164 | 165 | Args: 166 | hdri_path (str): The base directory for storing background images. 167 | combination (Dict): The combination dictionary containing background information. 168 | sphere (bpy.types.Object): The photosphere object to assign the material to. 169 | 170 | Returns: 171 | None 172 | """ 173 | # Create a new material 174 | mat = bpy.data.materials.new(name="PhotosphereMaterial") 175 | mat.use_nodes = True 176 | nodes = mat.node_tree.nodes 177 | nodes.clear() 178 | 179 | # Create and connect the nodes 180 | emission = nodes.new(type="ShaderNodeEmission") 181 | env_tex = nodes.new(type="ShaderNodeTexEnvironment") 182 | env_tex.image = bpy.data.images.load(get_hdri_path(hdri_path, combination)) 183 | mat.node_tree.links.new(env_tex.outputs["Color"], emission.inputs["Color"]) 184 | output = nodes.new(type="ShaderNodeOutputMaterial") 185 | mat.node_tree.links.new(emission.outputs["Emission"], output.inputs["Surface"]) 186 | 187 | # Assign material to the sphere 188 | if sphere.data.materials: 189 | sphere.data.materials[0] = mat 190 | else: 191 | sphere.data.materials.append(mat) 192 | 193 | # logger.info("Material created and applied to Photosphere") 194 | -------------------------------------------------------------------------------- /simian/camera.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import math 3 | 4 | import bpy 5 | from mathutils import Vector 6 | 7 | import numpy as np 8 | from scipy.spatial.transform import Rotation as R 9 | 10 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 11 | logger = logging.getLogger(__name__) 12 | 13 | 14 | def rotate_points(points, angles): 15 | """ 16 | Rotate points by given angles in degrees for (x, y, z) rotations. 17 | 18 | Args: 19 | points (np.ndarray): The points to rotate. 20 | angles (tuple): The angles to rotate by in degrees for (x, y, z) rotations. 21 | 22 | Returns: 23 | np.ndarray: The rotated points. 24 | """ 25 | 26 | rotation = R.from_euler("xyz", angles, degrees=True) 27 | return np.array([rotation.apply(point) for point in points]) 28 | 29 | 30 | def compute_camera_distance(points, fov_deg): 31 | """ 32 | Calculate the camera distance required to frame the bounding sphere of the points. 33 | 34 | Args: 35 | points (np.ndarray): The points to frame. 36 | fov_deg (float): The field of view in degrees. 37 | 38 | Returns: 39 | tuple: The camera distance, centroid of the points, and the radius of the bounding sphere. 40 | 41 | """ 42 | # Calculate the center of the bounding sphere (use the centroid for simplicity) 43 | centroid = np.mean(points, axis=0) 44 | # Calculate the radius as the max distance from the centroid to any point 45 | radius = np.max(np.linalg.norm(points - centroid, axis=1)) 46 | # Calculate the camera distance using the radius and the field of view 47 | fov_rad = math.radians(fov_deg) 48 | distance = radius / math.tan(fov_rad / 2) 49 | return distance, centroid, radius 50 | 51 | 52 | def perspective_project(points, camera_distance, fov_deg, aspect_ratio=1.0): 53 | """ 54 | Project points onto a 2D plane using a perspective projection considering the aspect ratio. 55 | 56 | Args: 57 | points (np.ndarray): The points to project. 58 | camera_distance (float): The distance of the camera from the origin. 59 | fov_deg (float): The field of view in degrees. 60 | aspect_ratio (float): The aspect ratio of the screen. Defaults to 1.0. 61 | 62 | Returns: 63 | np.ndarray: The screen space coordinates of the projected 64 | """ 65 | screen_points = [] 66 | fov_rad = math.radians(fov_deg) 67 | f = 1.0 / math.tan(fov_rad / 2) 68 | for point in points: 69 | # Translate point to camera's frame of reference (camera along the positive x-axis) 70 | p_cam = np.array([camera_distance - point[0], point[1], point[2]]) 71 | # Apply perspective projection if the point is in front of the camera 72 | if p_cam[0] > 0: 73 | x = (p_cam[1] * f) / (p_cam[0] * aspect_ratio) # Adjust x by aspect ratio 74 | y = p_cam[2] * f / p_cam[0] 75 | # Normalize to range [0, 1] for OpenGL screen space 76 | screen_x = (x + 1) / 2 77 | screen_y = (y + 1) / 2 78 | screen_points.append((screen_x, screen_y)) 79 | return np.array(screen_points) 80 | 81 | 82 | def create_camera_rig() -> bpy.types.Object: 83 | """ 84 | Creates a camera rig consisting of multiple objects in Blender. 85 | 86 | Returns: 87 | dict: A dictionary containing the created objects: 88 | - camera_animation_root: The root object of the camera animation hierarchy. 89 | - camera_orientation_pivot_yaw: The yaw pivot object for camera orientation. 90 | - camera_orientation_pivot_pitch: The pitch pivot object for camera orientation. 91 | - camera_framing_pivot: The pivot object for camera framing. 92 | - camera_animation_pivot: The pivot object for camera animation. 93 | - camera_object: The camera object. 94 | - camera: The camera data. 95 | """ 96 | camera_animation_root = bpy.data.objects.new("CameraAnimationRoot", None) 97 | bpy.context.scene.collection.objects.link(camera_animation_root) 98 | 99 | camera_orientation_pivot_yaw = bpy.data.objects.new( 100 | "CameraOrientationPivotYaw", None 101 | ) 102 | camera_orientation_pivot_yaw.parent = camera_animation_root 103 | bpy.context.scene.collection.objects.link(camera_orientation_pivot_yaw) 104 | 105 | camera_orientation_pivot_pitch = bpy.data.objects.new( 106 | "CameraOrientationPivotPitch", None 107 | ) 108 | camera_orientation_pivot_pitch.parent = camera_orientation_pivot_yaw 109 | bpy.context.scene.collection.objects.link(camera_orientation_pivot_pitch) 110 | 111 | camera_framing_pivot = bpy.data.objects.new("CameraFramingPivot", None) 112 | camera_framing_pivot.parent = camera_orientation_pivot_pitch 113 | bpy.context.scene.collection.objects.link(camera_framing_pivot) 114 | 115 | camera_animation_pivot = bpy.data.objects.new("CameraAnimationPivot", None) 116 | camera_animation_pivot.parent = camera_framing_pivot 117 | bpy.context.scene.collection.objects.link(camera_animation_pivot) 118 | 119 | camera = bpy.data.cameras.new("Camera") 120 | camera_object = bpy.data.objects.new("Camera", camera) 121 | 122 | # Rotate the Camera 90º 123 | camera_object.delta_rotation_euler = [1.5708, 0, 1.5708] 124 | camera_object.data.lens_unit = "FOV" 125 | 126 | camera_object.parent = camera_animation_pivot 127 | bpy.context.scene.collection.objects.link(camera_object) 128 | 129 | bpy.context.scene.camera = camera_object 130 | 131 | return { 132 | "camera_animation_root": camera_animation_root, 133 | "camera_orientation_pivot_yaw": camera_orientation_pivot_yaw, 134 | "camera_orientation_pivot_pitch": camera_orientation_pivot_pitch, 135 | "camera_framing_pivot": camera_framing_pivot, 136 | "camera_animation_pivot": camera_animation_pivot, 137 | "camera_object": camera_object, 138 | "camera": camera, 139 | } 140 | 141 | 142 | def set_camera_settings(combination: dict) -> None: 143 | """ 144 | Applies camera settings from a combination to the Blender scene. 145 | 146 | This function updates various camera settings including orientation, pivot adjustments, and 147 | framing based on the provided combination dictionary. 148 | 149 | Args: 150 | combination (dict): A dictionary containing camera settings including 'fov', 'animation', 151 | and orientation details. 152 | Returns: 153 | None 154 | """ 155 | camera = bpy.context.scene.objects["Camera"] 156 | camera_data = camera.data 157 | 158 | postprocessing = combination.get("postprocessing", {}) 159 | 160 | # Apply bloom settings 161 | bloom_settings = postprocessing.get("bloom", {}) 162 | threshold = bloom_settings.get("threshold", 0.8) 163 | intensity = bloom_settings.get("intensity", 0.5) 164 | radius = bloom_settings.get("radius", 5.0) 165 | bpy.context.scene.eevee.use_bloom = True 166 | bpy.context.scene.eevee.bloom_threshold = threshold 167 | bpy.context.scene.eevee.bloom_intensity = intensity 168 | bpy.context.scene.eevee.bloom_radius = radius 169 | 170 | # Apply SSAO settings 171 | ssao_settings = postprocessing.get("ssao", {}) 172 | distance = ssao_settings.get("distance", 0.2) 173 | factor = ssao_settings.get("factor", 0.5) 174 | bpy.context.scene.eevee.use_gtao = True 175 | bpy.context.scene.eevee.gtao_distance = distance 176 | bpy.context.scene.eevee.gtao_factor = factor 177 | 178 | # Apply SSRR settings 179 | ssrr_settings = postprocessing.get("ssrr", {}) 180 | max_roughness = ssrr_settings.get("max_roughness", 0.5) 181 | thickness = ssrr_settings.get("thickness", 0.1) 182 | bpy.context.scene.eevee.use_ssr = True 183 | bpy.context.scene.eevee.use_ssr_refraction = True 184 | bpy.context.scene.eevee.ssr_max_roughness = max_roughness 185 | bpy.context.scene.eevee.ssr_thickness = thickness 186 | 187 | # Apply motion blur settings 188 | motionblur_settings = postprocessing.get("motionblur", {}) 189 | shutter_speed = motionblur_settings.get("shutter_speed", 0.5) 190 | bpy.context.scene.eevee.use_motion_blur = True 191 | bpy.context.scene.eevee.motion_blur_shutter = shutter_speed 192 | 193 | # Get the initial lens value from the combination 194 | initial_lens = combination["framing"]["fov"] 195 | 196 | # Get the first keyframe's angle_offset value, if available 197 | animation = combination["animation"] 198 | keyframes = animation["keyframes"] 199 | if ( 200 | keyframes 201 | and "Camera" in keyframes[0] 202 | and "angle_offset" in keyframes[0]["Camera"] 203 | ): 204 | angle_offset = keyframes[0]["Camera"]["angle_offset"] 205 | camera_data.angle = math.radians(initial_lens + angle_offset) 206 | else: 207 | camera_data.angle = math.radians(initial_lens) 208 | 209 | orientation_data = combination["orientation"] 210 | orientation = {"pitch": orientation_data["pitch"], "yaw": orientation_data["yaw"]} 211 | 212 | # Rotate CameraOrientationPivotYaw by the Y 213 | camera_orientation_pivot_yaw = bpy.data.objects.get("CameraOrientationPivotYaw") 214 | camera_orientation_pivot_yaw.rotation_euler[2] = orientation["yaw"] * math.pi / 180 215 | 216 | # Rotate CameraOrientationPivotPitch by the X 217 | camera_orientation_pivot_pitch = bpy.data.objects.get("CameraOrientationPivotPitch") 218 | camera_orientation_pivot_pitch.rotation_euler[1] = ( 219 | orientation["pitch"] * -math.pi / 180 220 | ) 221 | 222 | # set the camera framerate to 30 223 | bpy.context.scene.render.fps = 30 224 | 225 | 226 | def set_camera_animation(combination: dict, frame_interval: int, animation_length: int) -> None: 227 | """ 228 | Applies the specified animation to the camera based on the keyframes from the camera_data.json file. 229 | The total animation frames are fixed to ensure consistent speed. 230 | 231 | Args: 232 | combination (dict): The combination dictionary containing animation data. 233 | 234 | Returns: 235 | None 236 | """ 237 | animation = combination["animation"] 238 | speed_factor = animation.get("speed_factor", 1)/1.5 239 | keyframes = animation["keyframes"] 240 | adjusted_frame_interval = frame_interval * (animation_length / 100) 241 | 242 | for i, keyframe in enumerate(keyframes): 243 | for obj_name, transforms in keyframe.items(): 244 | obj = bpy.data.objects.get(obj_name) 245 | if obj is None: 246 | raise ValueError(f"Object {obj_name} not found in the scene") 247 | frame = int(i * adjusted_frame_interval) 248 | for transform_name, value in transforms.items(): 249 | if transform_name == "position": 250 | obj.location = [coord * speed_factor for coord in value] 251 | obj.keyframe_insert(data_path="location", frame=frame) 252 | elif transform_name == "rotation": 253 | obj.rotation_euler = [ 254 | math.radians(angle * speed_factor) for angle in value 255 | ] 256 | obj.keyframe_insert(data_path="rotation_euler", frame=frame) 257 | elif transform_name == "scale": 258 | obj.scale = [coord * speed_factor for coord in value] 259 | obj.keyframe_insert(data_path="scale", frame=frame) 260 | elif transform_name == "angle_offset" and obj_name == "Camera": 261 | camera_data = bpy.data.objects["Camera"].data 262 | camera_data.angle = math.radians( 263 | combination["framing"]["fov"] + value 264 | ) 265 | camera_data.keyframe_insert(data_path="lens", frame=frame) 266 | 267 | bpy.context.scene.frame_set(0) 268 | 269 | 270 | def position_camera(combination: dict, focus_object: bpy.types.Object) -> None: 271 | """ 272 | Positions the camera based on the coverage factor and lens values. 273 | 274 | Args: 275 | combination (dict): The combination dictionary containing coverage factor and lens values. 276 | focus_object (bpy.types.Object): The object to focus the camera on. 277 | 278 | Returns: 279 | None 280 | """ 281 | camera = bpy.context.scene.objects["Camera"] 282 | 283 | # Get the bounding box of the focus object in world space 284 | bpy.context.view_layer.update() 285 | 286 | bbox = [ 287 | focus_object.matrix_world @ Vector(corner) for corner in focus_object.bound_box 288 | ] 289 | bbox_points = np.array([corner.to_tuple() for corner in bbox]) 290 | 291 | # Rotate points as per the desired view angle if any 292 | # Assuming we want to compute this based on some predefined rotation angles 293 | rotation_angles = (45, 45, 45) # Example rotation angles 294 | rotated_points = rotate_points(bbox_points, rotation_angles) 295 | 296 | # scale rotated_points by combination["framing"]["coverage_factor"] 297 | rotated_points *= combination["framing"]["coverage_factor"] 298 | 299 | # Calculate the camera distance to frame the rotated bounding box correctly 300 | fov_deg = combination["framing"][ 301 | "fov" 302 | ] # Get the FOV from combination or default to 45 303 | aspect_ratio = ( 304 | bpy.context.scene.render.resolution_x / bpy.context.scene.render.resolution_y 305 | ) 306 | 307 | camera_distance, centroid, radius = compute_camera_distance( 308 | rotated_points, fov_deg / aspect_ratio 309 | ) 310 | 311 | # Set the camera properties 312 | camera.data.angle = math.radians(fov_deg) # Set camera FOV 313 | if aspect_ratio >= 1: 314 | camera.data.sensor_fit = "HORIZONTAL" 315 | else: 316 | camera.data.sensor_fit = "VERTICAL" 317 | 318 | bbox = [ 319 | focus_object.matrix_world @ Vector(corner) for corner in focus_object.bound_box 320 | ] 321 | bbox_min = min(bbox, key=lambda v: v.z) 322 | bbox_max = max(bbox, key=lambda v: v.z) 323 | 324 | # Calculate the height of the bounding box 325 | bbox_height = bbox_max.z - bbox_min.z 326 | 327 | # Position the camera based on the computed distance 328 | camera.location = Vector((camera_distance, 0, 0)) # Adjust this as needed 329 | 330 | 331 | # Set the position of the CameraAnimationRoot object to slightly above the focus object center, quasi-rule of thirds 332 | # bbox_height / 2 is the center of the bounding box, bbox_height / 1.66 is more aesthetically pleasing 333 | bpy.data.objects["CameraAnimationRoot"].location = focus_object.location + Vector( 334 | (0, 0, bbox_height/2) 335 | ) 336 | -------------------------------------------------------------------------------- /simian/distributed.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import logging 4 | import os 5 | import time 6 | from tqdm import tqdm 7 | from typing import Dict 8 | 9 | from distributask.distributask import Distributask 10 | 11 | from .worker import run_job 12 | 13 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 14 | logger = logging.getLogger(__name__) 15 | 16 | if __name__ == "__main__": 17 | 18 | def get_env_vars(path: str = ".env") -> Dict[str, str]: 19 | """Get the environment variables from the specified file.""" 20 | env_vars = {} 21 | if not os.path.exists(path): 22 | return env_vars 23 | with open(path, "r") as f: 24 | for line in f: 25 | key, value = line.strip().split("=") 26 | env_vars[key] = value 27 | return env_vars 28 | 29 | def get_settings(args): 30 | env_vars = get_env_vars(".env") 31 | 32 | # Get settings from parsed argements or env, set defaults 33 | settings = { 34 | "start_index": args.start_index or int(env_vars.get("START_INDEX", 0)), 35 | "combinations_file": args.combinations_file 36 | or env_vars.get("COMBINATIONS_FILE", "combinations.json"), 37 | "end_index": args.end_index or int(env_vars.get("END_INDEX", 100)), 38 | "start_frame": args.start_frame or int(env_vars.get("START_FRAME", 0)), 39 | "end_frame": args.end_frame or int(env_vars.get("END_FRAME", 300)), 40 | "width": args.width or int(env_vars.get("WIDTH", 1280)), 41 | "height": args.height or int(env_vars.get("HEIGHT", 720)), 42 | "output_dir": args.output_dir or env_vars.get("OUTPUT_DIR", "./renders"), 43 | "hdri_path": args.hdri_path or env_vars.get("HDRI_PATH", "./backgrounds"), 44 | "max_price": args.max_price or float(env_vars.get("MAX_PRICE", 0.1)), 45 | "max_nodes": args.max_nodes or int(env_vars.get("MAX_NODES", 1)), 46 | "api_key": args.api_key or env_vars.get("VAST_API_KEY", ""), 47 | "redis_host": args.redis_host or env_vars.get("REDIS_HOST", "localhost"), 48 | "redis_port": args.redis_port or int(env_vars.get("REDIS_PORT", 6379)), 49 | "redis_user": args.redis_user or env_vars.get("REDIS_USER", ""), 50 | "redis_password": args.redis_password or env_vars.get("REDIS_PASSWORD", ""), 51 | "amazon_key_id": args.amazon_key_id or env_vars.get("AMAZON_KEY_ID", ""), 52 | "amazon_secret_key": args.amazon_secret_key 53 | or env_vars.get("AMAZON_SECRET_KEY", ""), 54 | "s3_bucket_name": args.s3_bucket_name or env_vars.get("S3_BUCKET_NAME", ""), 55 | "hf_token": args.hf_token or env_vars.get("HF_TOKEN", ""), 56 | "hf_repo_id": args.hf_repo_id or env_vars.get("HF_REPO_ID", ""), 57 | "broker_pool_limit": args.broker_pool_limit 58 | or int(env_vars.get("BROKER_POOL_LIMIT", 1)), 59 | "render_batch_size": args.render_batch_size 60 | or int(env_vars.get("RENDER_BATCH_SIZE", 1)), 61 | "inactivity_check_interval": args.inactivity_check or int(env_vars.get("INACTVITY_INTERVAL", 600)), 62 | "upload_destination": args.upload_dest or env_vars.get("UPLOAD_DEST", "s3") 63 | } 64 | 65 | # Load combinations from file 66 | with open(settings["combinations_file"], "r") as f: 67 | combinations = json.load(f) 68 | settings["combinations"] = combinations["combinations"] 69 | 70 | return settings 71 | 72 | def start_new_job(args): 73 | logger.info("Starting a new job...") 74 | settings = get_settings(args) 75 | 76 | # Override environment variables with provided arguments 77 | os.environ["REDIS_HOST"] = args.redis_host or settings["redis_host"] 78 | os.environ["REDIS_PORT"] = str(args.redis_port or settings["redis_port"]) 79 | os.environ["REDIS_USER"] = args.redis_user or settings["redis_user"] 80 | os.environ["REDIS_PASSWORD"] = args.redis_password or settings["redis_password"] 81 | os.environ["HF_TOKEN"] = args.hf_token or settings["hf_token"] 82 | os.environ["HF_REPO_ID"] = args.hf_repo_id or settings["hf_repo_id"] 83 | 84 | job_config = { 85 | "max_price": settings["max_price"], 86 | "max_nodes": settings["max_nodes"], 87 | "start_index": settings["start_index"], 88 | "end_index": settings["end_index"], 89 | "combinations": settings["combinations"], 90 | "width": settings["width"], 91 | "height": settings["height"], 92 | "output_dir": settings["output_dir"], 93 | "hdri_path": settings["hdri_path"], 94 | "start_frame": settings["start_frame"], 95 | "end_frame": settings["end_frame"], 96 | "api_key": settings["api_key"], 97 | "hf_token": settings["hf_token"], 98 | "hf_repo_id": settings["hf_repo_id"], 99 | "redis_host": settings["redis_host"], 100 | "redis_port": settings["redis_port"], 101 | "redis_user": settings["redis_user"], 102 | "redis_password": settings["redis_password"], 103 | "amazon_key_id": settings["amazon_key_id"], 104 | "s3_bucket_name": settings["s3_bucket_name"], 105 | "amazon_secret_key": settings["amazon_secret_key"], 106 | "broker_pool_limit": settings["broker_pool_limit"], 107 | "render_batch_size": settings["render_batch_size"], 108 | "inactivity_check_interval": settings["inactivity_check_interval"], 109 | "upload_destination": settings["upload_destination"] 110 | } 111 | 112 | instance_env = { 113 | "VAST_API_KEY": settings["api_key"], 114 | "HF_TOKEN": settings["hf_token"], 115 | "HF_REPO_ID": settings["hf_repo_id"], 116 | "REDIS_HOST": settings["redis_host"], 117 | "REDIS_PORT": settings["redis_port"], 118 | "REDIS_USER": settings["redis_user"], 119 | "REDIS_PASSWORD": settings["redis_password"], 120 | "AWS_ACCESS_KEY_ID": settings["amazon_key_id"], 121 | "AWS_SECRET_ACCESS_KEY": settings["amazon_secret_key"], 122 | "S3_BUCKET_NAME": settings["s3_bucket_name"], 123 | "BROKER_POOL_LIMIT": settings["broker_pool_limit"], 124 | } 125 | 126 | print("*** JOB CONFIG") 127 | for key, value in job_config.items(): 128 | if key != "combinations": 129 | print(f"{key}: {value}") 130 | 131 | distributask = Distributask( 132 | hf_repo_id=job_config["hf_repo_id"], 133 | hf_token=job_config["hf_token"], 134 | vast_api_key=job_config["api_key"], 135 | redis_host=job_config["redis_host"], 136 | redis_port=job_config["redis_port"], 137 | redis_username=job_config["redis_user"], 138 | redis_password=job_config["redis_password"], 139 | broker_pool_limit=job_config["broker_pool_limit"], 140 | ) 141 | 142 | max_price = job_config["max_price"] 143 | max_nodes = job_config["max_nodes"] 144 | docker_image = "antbaez/simian-worker:latest" 145 | module_name = "simian.worker" 146 | 147 | # Rent and set up vastai nodes with docker image 148 | print("Searching for nodes...") 149 | num_nodes_avail = len(distributask.search_offers(max_price)) 150 | print("Total nodes available: ", num_nodes_avail) 151 | 152 | rented_nodes = distributask.rent_nodes( 153 | max_price, max_nodes, docker_image, module_name, env_settings=instance_env 154 | ) 155 | 156 | print("Total nodes rented: ", len(rented_nodes)) 157 | 158 | distributask.register_function(run_job) 159 | 160 | while True: 161 | user_input = input("press r when workers are ready: ") 162 | if user_input == "r": 163 | break 164 | 165 | tasks = [] 166 | 167 | batch_size = job_config["render_batch_size"] 168 | # Submit tasks to queue in batches 169 | for combination_index in range( 170 | job_config["start_index"], 171 | job_config["end_index"], 172 | batch_size, 173 | ): 174 | task = distributask.execute_function( 175 | "run_job", 176 | { 177 | "combination_indeces": [ 178 | index 179 | for index in range( 180 | combination_index, 181 | min(combination_index + batch_size, settings["end_index"]), 182 | ) 183 | ], 184 | "combinations": [ 185 | job_config["combinations"][index] 186 | for index in range( 187 | combination_index, 188 | min(combination_index + batch_size, settings["end_index"]), 189 | ) 190 | ], 191 | "width": job_config["width"], 192 | "height": job_config["height"], 193 | "output_dir": job_config["output_dir"], 194 | "hdri_path": job_config["hdri_path"], 195 | "upload_dest": job_config["upload_destination"], 196 | "start_frame": job_config["start_frame"], 197 | "end_frame": job_config["end_frame"] 198 | }, 199 | ) 200 | tasks.append(task) 201 | 202 | # distributask.monitor_tasks(tasks, show_time_left=False) 203 | 204 | print("Tasks sent. Starting monitoring") 205 | inactivity_log = {node["instance_id"]: 0 for node in rented_nodes} 206 | 207 | start_time = time.time() 208 | with tqdm(total=len(tasks), unit="task") as pbar: 209 | while not all(task.ready() for task in tasks): 210 | 211 | current_tasks = sum([task.ready() for task in tasks]) 212 | pbar.update(current_tasks - pbar.n) 213 | 214 | time.sleep(1) 215 | current_time = time.time() 216 | # check if node is inactive at set interval 217 | if current_time - start_time > settings["inactivity_check_interval"]: 218 | start_time = time.time() 219 | for node in rented_nodes: 220 | # get log with api call 221 | log_response = distributask.get_node_log(node) 222 | if log_response: 223 | # if "Task completed" in two consecutive logs, terminate node 224 | try: 225 | last_msg = log_response.text.splitlines()[-3:] 226 | if ( 227 | any("Task completed" in msg for msg in last_msg) 228 | and inactivity_log[node["instance_id"]] == 0 229 | ): 230 | inactivity_log[node["instance_id"]] = 1 231 | elif ( 232 | any("Task completed" in msg for msg in last_msg) 233 | and inactivity_log[node["instance_id"]] == 1 234 | ): 235 | distributask.terminate_nodes([node]) 236 | print("node terminated") 237 | else: 238 | inactivity_log[node["instance_id"]] == 0 239 | except: 240 | pass 241 | 242 | print("All tasks have been completed!") 243 | 244 | parser = argparse.ArgumentParser(description="Simian CLI") 245 | parser.add_argument("--start_index", type=int, help="Starting index for rendering") 246 | parser.add_argument( 247 | "--combinations-file", help="Path to the combinations JSON file" 248 | ) 249 | parser.add_argument("--end_index", type=int, help="Ending index for rendering") 250 | parser.add_argument("--start_frame", type=int, help="Starting frame number") 251 | parser.add_argument("--end_frame", type=int, help="Ending frame number") 252 | parser.add_argument("--width", type=int, help="Rendering width in pixels") 253 | parser.add_argument("--height", type=int, help="Rendering height in pixels") 254 | parser.add_argument("--output_dir", help="Output directory") 255 | parser.add_argument("--hdri_path", help="HDRI path") 256 | parser.add_argument("--max_price", type=float, help="Maximum price per hour") 257 | parser.add_argument("--max_nodes", type=int, help="Maximum number of nodes") 258 | parser.add_argument("--api_key", help="Vast.ai API key") 259 | parser.add_argument("--redis_host", help="Redis host") 260 | parser.add_argument("--redis_port", type=int, help="Redis port") 261 | parser.add_argument("--redis_user", help="Redis user") 262 | parser.add_argument("--redis_password", help="Redis password") 263 | parser.add_argument("--hf_token", help="Hugging Face token") 264 | parser.add_argument("--hf_repo-id", help="Hugging Face repository ID") 265 | parser.add_argument("--amazon_key_id", help="amazon secret key id") 266 | parser.add_argument("--amazon_secret_key", help="amazon secret access key") 267 | parser.add_argument("--s3_bucket_name", help="amazon s3 bucket name") 268 | parser.add_argument( 269 | "--broker_pool_limit", type=int, help="Limit on redis pool size" 270 | ) 271 | parser.add_argument( 272 | "--render_batch_size", type=int, help="Batch size of simian rendering" 273 | ) 274 | parser.add_argument("--inactivity_check", type=int, help="The interval of checking and terminating nodes for inactivity in seconds" 275 | ) 276 | parser.add_argument("--upload_dest", type=int, help="The desired destination of uploads at task completion" 277 | ) 278 | args = parser.parse_args() 279 | 280 | start_new_job(args) 281 | -------------------------------------------------------------------------------- /simian/postprocessing.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | from typing import Any 3 | 4 | 5 | def setup_compositor_for_black_and_white(context: bpy.types.Context) -> None: 6 | """ 7 | Sets up the compositor for a black and white effect using Blender's Eevee engine. 8 | Assumes that the context provided is valid and that the scene uses Eevee. 9 | 10 | Args: 11 | context (bpy.types.Context): The Blender context. 12 | 13 | Returns: 14 | None 15 | """ 16 | # Ensure the use of nodes in the scene's compositing. 17 | scene = context.scene 18 | scene.use_nodes = True 19 | tree = scene.node_tree 20 | 21 | # Create necessary nodes 22 | render_layers = tree.nodes.new(type="CompositorNodeRLayers") 23 | hue_sat = tree.nodes.new(type="CompositorNodeHueSat") 24 | composite = tree.nodes.new( 25 | type="CompositorNodeComposite" 26 | ) # Ensure there's a composite node 27 | 28 | # Position nodes 29 | render_layers.location = (-300, 0) 30 | hue_sat.location = (100, 0) 31 | composite.location = (300, 0) 32 | 33 | # Configure Hue Sat node for desaturation 34 | hue_sat.inputs["Saturation"].default_value = ( 35 | 0 # Reduce saturation to zero to get grayscale 36 | ) 37 | # increase contrast 38 | hue_sat.inputs["Value"].default_value = 2.0 39 | 40 | # Link nodes 41 | links = tree.links 42 | links.new(render_layers.outputs["Image"], hue_sat.inputs["Image"]) 43 | links.new( 44 | hue_sat.outputs["Image"], composite.inputs["Image"] 45 | ) # Direct output to the composite node 46 | 47 | 48 | def setup_compositor_for_cel_shading(context: bpy.types.Context) -> None: 49 | """ 50 | Sets up the compositor for a cel shading effect using Blender's Eevee engine. 51 | The node setup is based on the theory of using normal and diffuse passes to create 52 | a stylized, non-photorealistic look with subtle transitions between shadows and highlights. 53 | Assumes the context provided is valid and that the scene uses Eevee. 54 | 55 | Args: 56 | context (bpy.types.Context): The Blender context. 57 | 58 | Returns: 59 | None 60 | """ 61 | # Ensure the use of nodes in the scene's compositing 62 | scene: bpy.types.Scene = context.scene 63 | scene.use_nodes = True 64 | tree: bpy.types.NodeTree = scene.node_tree 65 | 66 | # Create necessary nodes for compositing 67 | # Render Layers Node: Provides access to the render passes 68 | # Normal Node: Converts the normal pass into a usable format 69 | # Color Ramp Nodes: Control the shading and highlight areas 70 | # Mix RGB Nodes: Combine the shading and highlight with the diffuse color 71 | # Composite Node: The final output node 72 | render_layers: Any = tree.nodes.new(type="CompositorNodeRLayers") 73 | normal_node: Any = tree.nodes.new(type="CompositorNodeNormal") 74 | color_ramp_shadow: Any = tree.nodes.new(type="CompositorNodeValToRGB") 75 | color_ramp_highlight: Any = tree.nodes.new(type="CompositorNodeValToRGB") 76 | mix_rgb_shadow: Any = tree.nodes.new(type="CompositorNodeMixRGB") 77 | mix_rgb_highlight: Any = tree.nodes.new(type="CompositorNodeMixRGB") 78 | alpha_over: Any = tree.nodes.new( 79 | type="CompositorNodeAlphaOver" 80 | ) # Add an Alpha Over node 81 | composite: Any = tree.nodes.new(type="CompositorNodeComposite") 82 | 83 | # Configure Mix RGB nodes 84 | # Multiply blend mode for shadows to darken the diffuse color slightly 85 | # Overlay blend mode for highlights to create subtle bright highlights 86 | # Reference: https://docs.blender.org/manual/en/latest/compositing/types/color/mix.html 87 | mix_rgb_shadow.blend_type = "MULTIPLY" 88 | mix_rgb_highlight.blend_type = "OVERLAY" 89 | mix_rgb_shadow.use_clamp = True 90 | mix_rgb_highlight.use_clamp = True 91 | 92 | # Configure Shadow Color Ramp 93 | color_ramp_shadow.color_ramp.interpolation = "EASE" 94 | color_ramp_shadow.color_ramp.elements[0].position = 0.5 95 | color_ramp_shadow.color_ramp.elements[1].position = 0.8 96 | color_ramp_shadow.color_ramp.elements[0].color = (0.5, 0.5, 0.5, 1) # Mid Gray 97 | color_ramp_shadow.color_ramp.elements[1].color = (1, 1, 1, 1) # White 98 | 99 | # Configure Highlight Color Ramp 100 | color_ramp_highlight.color_ramp.interpolation = "EASE" 101 | color_ramp_highlight.color_ramp.elements[0].position = 0.8 102 | color_ramp_highlight.color_ramp.elements[1].position = 0.95 103 | color_ramp_highlight.color_ramp.elements[0].color = (0, 0, 0, 1) # Black 104 | color_ramp_highlight.color_ramp.elements[1].color = (1, 1, 1, 1) # White 105 | 106 | # Adjust the Mix RGB nodes 107 | mix_rgb_shadow.blend_type = "MULTIPLY" 108 | mix_rgb_shadow.inputs[0].default_value = 1.0 # Reduce the shadow intensity 109 | 110 | mix_rgb_highlight.blend_type = ( 111 | "SCREEN" # Change to 'SCREEN' for better highlight blending 112 | ) 113 | mix_rgb_highlight.inputs[0].default_value = 0.5 # Reduce the highlight intensity 114 | 115 | # Link nodes 116 | links: Any = tree.links 117 | # Connect Normal pass to Normal node for shading information 118 | links.new(render_layers.outputs["Normal"], normal_node.inputs["Normal"]) 119 | # Connect Normal node to Shadow Color Ramp for shadow intensity control 120 | links.new(normal_node.outputs["Dot"], color_ramp_shadow.inputs["Fac"]) 121 | # Connect Normal node to Highlight Color Ramp for highlight intensity control 122 | links.new(normal_node.outputs["Dot"], color_ramp_highlight.inputs["Fac"]) 123 | # Connect Shadow Color Ramp to Mix RGB Shadow node for shadow color blending 124 | links.new(color_ramp_shadow.outputs["Image"], mix_rgb_shadow.inputs[2]) 125 | # Connect Diffuse pass to Mix RGB Shadow node as the base color 126 | links.new(render_layers.outputs["Image"], mix_rgb_shadow.inputs[1]) 127 | # Connect Mix RGB Shadow to Mix RGB Highlight for combining shadows with the base color 128 | links.new(mix_rgb_shadow.outputs["Image"], mix_rgb_highlight.inputs[1]) 129 | # Connect Highlight Color Ramp to Mix RGB Highlight for adding highlights 130 | links.new(color_ramp_highlight.outputs["Image"], mix_rgb_highlight.inputs[2]) 131 | # Connect Mix RGB Highlight to Composite node for final output 132 | links.new( 133 | mix_rgb_highlight.outputs["Image"], alpha_over.inputs[2] 134 | ) # Plug the cel-shaded result into the foreground input of Alpha Over 135 | links.new( 136 | render_layers.outputs["Env"], alpha_over.inputs[1] 137 | ) # Plug the environment pass into the background input of Alpha Over 138 | links.new( 139 | alpha_over.outputs["Image"], composite.inputs["Image"] 140 | ) # Plug the Alpha Over output into the Composite node 141 | 142 | 143 | def setup_compositor_for_depth(context: bpy.types.Context) -> None: 144 | """ 145 | Sets up the compositor for rendering a depth map using Blender's Eevee engine. 146 | Assumes that the context provided is valid and that the scene uses Eevee. 147 | 148 | Args: 149 | context (bpy.types.Context): The Blender context. 150 | 151 | Returns: 152 | None 153 | """ 154 | # Ensure the use of nodes in the scene's compositing. 155 | scene: bpy.types.Scene = context.scene 156 | scene.use_nodes = True 157 | tree: bpy.types.NodeTree = scene.node_tree 158 | # Create necessary nodes 159 | render_layers: Any = tree.nodes.new(type="CompositorNodeRLayers") 160 | normalize: Any = tree.nodes.new(type="CompositorNodeNormalize") 161 | composite: Any = tree.nodes.new( 162 | type="CompositorNodeComposite" 163 | ) # Ensure there's a composite node 164 | 165 | # Position nodes 166 | render_layers.location = (-300, 0) 167 | normalize.location = (0, 0) 168 | composite.location = (300, 0) 169 | 170 | # Enable the Depth pass in the View Layer properties 171 | view_layer: bpy.types.ViewLayer = context.view_layer 172 | view_layer.use_pass_z = True 173 | 174 | # Link nodes 175 | links: Any = tree.links 176 | links.new(render_layers.outputs["Depth"], normalize.inputs["Value"]) 177 | links.new( 178 | normalize.outputs["Value"], composite.inputs["Image"] 179 | ) # Direct output to the composite node 180 | 181 | 182 | effects = { 183 | "black_and_white": setup_compositor_for_black_and_white, 184 | "cel_shading": setup_compositor_for_cel_shading, 185 | "depth": setup_compositor_for_depth, 186 | } 187 | 188 | 189 | def enable_effect(context: bpy.types.Context, effect_name: str) -> None: 190 | """ 191 | Enables the cel shading effect in the compositor. 192 | 193 | Args: 194 | context (bpy.types.Context): The Blender context. 195 | effect_name (str): The name of the effect to enable. 196 | 197 | Returns: 198 | None 199 | """ 200 | # Make sure compositor use is turned on 201 | context.scene.render.use_compositing = True 202 | context.scene.use_nodes = True 203 | tree: bpy.types.NodeTree = context.scene.node_tree 204 | 205 | # Clear existing nodes 206 | for node in tree.nodes: 207 | tree.nodes.remove(node) 208 | 209 | # Enable Normal Pass and Diffuse Pass 210 | # These passes provide the necessary information for shading and color 211 | # Reference: https://docs.blender.org/manual/en/latest/render/layers/passes.html 212 | view_layer: bpy.types.ViewLayer = context.view_layer 213 | view_layer.use_pass_normal = True 214 | view_layer.use_pass_diffuse_color = True 215 | view_layer.use_pass_environment = True # Enable the environment pass 216 | view_layer.use_pass_z = True 217 | 218 | # get the function for the effect 219 | effect_function: Any = effects.get(effect_name) 220 | 221 | # Set up the nodes for cel shading 222 | effect_function(context) 223 | -------------------------------------------------------------------------------- /simian/render.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import logging 4 | import os 5 | import ssl 6 | import sys 7 | import bpy 8 | import random 9 | from rich.console import Console 10 | 11 | console = Console() 12 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 13 | logger = logging.getLogger(__name__) 14 | 15 | ssl._create_default_https_context = ssl._create_unverified_context 16 | 17 | from .camera import ( 18 | create_camera_rig, 19 | position_camera, 20 | set_camera_animation, 21 | set_camera_settings, 22 | ) 23 | from .transform import( 24 | find_largest_length, 25 | place_objects_on_grid, 26 | apply_animation 27 | ) 28 | from .object import ( 29 | apply_all_modifiers, 30 | apply_and_remove_armatures, 31 | get_meshes_in_hierarchy, 32 | join_objects_in_hierarchy, 33 | load_object, 34 | lock_all_objects, 35 | normalize_object_scale, 36 | optimize_meshes_in_hierarchy, 37 | set_pivot_to_bottom, 38 | unlock_objects, 39 | unparent_keep_transform, 40 | ) 41 | from .background import create_photosphere, set_background 42 | from .scene import apply_stage_material, create_stage, initialize_scene 43 | from .vendor import objaverse 44 | 45 | 46 | def read_combination(combination_file: str, index: int = 0) -> dict: 47 | """ 48 | Reads a specified camera combination from a JSON file. 49 | 50 | Args: 51 | combination_file (str): Path to the JSON file containing camera combinations. 52 | index (int): Index of the camera combination to read from the JSON file. Defaults to 0. 53 | 54 | Returns: 55 | dict: The camera combination data. 56 | """ 57 | with open(combination_file, "r") as file: 58 | data = json.load(file) 59 | combinations_data = data["combinations"] 60 | return combinations_data[index] 61 | 62 | 63 | def load_user_blend_file(user_blend_file): 64 | """ 65 | Loads a user-specified Blender file as the base scene. 66 | 67 | Args: 68 | user_blend_file (str): Path to the user-specified Blender file. 69 | 70 | Returns: 71 | bool: True if the file was successfully loaded, False otherwise. 72 | """ 73 | if not os.path.exists(user_blend_file): 74 | logger.error(f"Blender file {user_blend_file} does not exist.") 75 | return False 76 | 77 | try: 78 | bpy.ops.wm.open_mainfile(filepath=user_blend_file) 79 | # logger.info(f"Opened user-specified Blender file {user_blend_file} as the base scene.") 80 | return True 81 | except Exception as e: 82 | logger.error(f"Failed to load Blender file {user_blend_file}: {e}") 83 | return False 84 | 85 | 86 | def select_focus_object(all_objects): 87 | for obj_dict in all_objects: 88 | obj = list(obj_dict.keys())[0] # This is the actual Blender object 89 | properties = obj_dict[obj] 90 | if isinstance(properties, dict) and properties.get("camera_follow", {}).get("follow", False): 91 | return obj 92 | 93 | # If no object with camera_follow, fall back to placement 4 or the first object 94 | for obj_dict in all_objects: 95 | obj = list(obj_dict.keys())[0] 96 | properties = obj_dict[obj] 97 | if isinstance(properties, dict) and properties.get("placement") == 4: 98 | return obj 99 | 100 | # If no object with placement 4 is found, return the first object 101 | return list(all_objects[0].keys())[0] if all_objects else None 102 | 103 | 104 | def render_scene( 105 | output_dir: str, 106 | context: bpy.types.Context, 107 | combination_file, 108 | start_frame: int = 1, 109 | end_frame: int = 65, 110 | combination_index=0, 111 | combination=None, 112 | render_images: bool =False, 113 | user_blend_file = None, 114 | animation_length: int = 100 115 | ) -> None: 116 | """ 117 | Renders a scene with specified parameters. 118 | 119 | Args: 120 | output_dir (str): Path to the directory where the rendered video will be saved. 121 | context (bpy.types.Context): Blender context. 122 | combination_file (str): Path to the JSON file containing camera combinations. 123 | start_frame (int): Start frame of the animation. Defaults to 1. 124 | end_frame (int): End frame of the animation. Defaults to 65. 125 | combination_index (int): Index of the camera combination to use from the JSON file. Defaults to 0. 126 | render_images (bool): Flag to indicate if images should be rendered instead of videos. 127 | user_blend_file (str): Path to the user-specified Blender file to use as the base scene 128 | animation_length (int): Percentage animation length. Defaults to 100. 129 | 130 | Returns: 131 | None 132 | """ 133 | 134 | console.print("Rendering scene with combination ", style="orange_red1", end="") 135 | console.print(f"{combination_index}", style="bold bright_green") 136 | 137 | os.makedirs(output_dir, exist_ok=True) 138 | 139 | initialize_scene() 140 | 141 | if user_blend_file: 142 | bpy.ops.wm.open_mainfile(filepath=user_blend_file) 143 | if not load_user_blend_file(user_blend_file): 144 | logger.error(f"Unable to load user-specified Blender file: {user_blend_file}") 145 | return # Exit the function if the file could not be loaded 146 | 147 | context.scene.render.engine = 'BLENDER_EEVEE' 148 | 149 | create_camera_rig() 150 | 151 | scene = context.scene 152 | 153 | scene.frame_start = start_frame 154 | scene.frame_end = end_frame 155 | 156 | # Lock and hide all scene objects before doing any object operations 157 | initial_objects = lock_all_objects() 158 | 159 | if combination is not None: 160 | combination = json.loads(combination) 161 | else: 162 | combination = read_combination(combination_file, combination_index) 163 | all_objects = [] 164 | 165 | focus_object = None 166 | 167 | for object_data in combination["objects"]: 168 | object_file = objaverse.load_objects([object_data["uid"]])[object_data["uid"]] 169 | 170 | load_object(object_file) 171 | obj = [obj for obj in context.view_layer.objects.selected][0] 172 | 173 | apply_and_remove_armatures() 174 | apply_all_modifiers(obj) 175 | join_objects_in_hierarchy(obj) 176 | optimize_meshes_in_hierarchy(obj) 177 | 178 | meshes = get_meshes_in_hierarchy(obj) 179 | obj = meshes[0] 180 | 181 | unparent_keep_transform(obj) 182 | set_pivot_to_bottom(obj) 183 | 184 | obj.scale = [object_data["scale"]["factor"] for _ in range(3)] 185 | normalize_object_scale(obj) 186 | obj.name = object_data["uid"] 187 | 188 | all_objects.append({obj: object_data}) 189 | 190 | largest_length = find_largest_length(all_objects) 191 | 192 | if not user_blend_file: 193 | set_background(args.hdri_path, combination) 194 | create_photosphere(args.hdri_path, combination).scale = (10, 10, 10) 195 | stage = create_stage(combination) 196 | apply_stage_material(stage, combination) 197 | 198 | unlock_objects(initial_objects) 199 | 200 | set_camera_settings(combination) 201 | set_camera_animation(combination, end_frame-start_frame, animation_length) 202 | 203 | place_objects_on_grid(all_objects, largest_length) 204 | yaw = combination["orientation"]["yaw"] 205 | 206 | focus_object = select_focus_object(all_objects) 207 | if focus_object is None or not isinstance(focus_object, bpy.types.Object): 208 | logger.error("No valid focus object found or focus object is not a Blender object. Cannot position camera.") 209 | return 210 | 211 | position_camera(combination, focus_object) 212 | 213 | if not combination.get("no_movement", False): 214 | check_camera_follow = any(obj.get("camera_follow", {}).get("follow", False) for obj in combination['objects']) 215 | apply_animation(all_objects, focus_object, yaw, scene.frame_start, end_frame, check_camera_follow) 216 | 217 | sizes = [ 218 | (1920, 1080), 219 | (1024, 1024), 220 | # able to add more options here 221 | ] 222 | 223 | if render_images: 224 | # Render a specific frame as an image with a random size 225 | middle_frame = (scene.frame_start + scene.frame_end) // 2 226 | size = random.choice(sizes) 227 | scene.frame_set(middle_frame) 228 | scene.render.resolution_x = size[0] 229 | scene.render.resolution_y = size[1] 230 | scene.render.resolution_percentage = 100 231 | render_path = os.path.join( 232 | output_dir, 233 | f"{combination_index}_frame_{middle_frame}_{size[0]}x{size[1]}.png", 234 | ) 235 | scene.render.filepath = render_path 236 | bpy.ops.render.render(write_still=True) 237 | logger.info(f"Rendered image saved to {render_path}") 238 | else: 239 | # Render the entire animation as a video 240 | scene.render.resolution_x = 1920 241 | scene.render.resolution_y = 1080 242 | scene.render.resolution_percentage = 100 243 | scene.render.image_settings.file_format = "FFMPEG" 244 | scene.render.ffmpeg.format = "MPEG4" 245 | scene.render.ffmpeg.codec = "H264" 246 | scene.render.ffmpeg.constant_rate_factor = "PERC_LOSSLESS" 247 | scene.render.ffmpeg.ffmpeg_preset = "BEST" 248 | render_path = os.path.join(output_dir, f"{combination_index}.mp4") 249 | scene.render.filepath = render_path 250 | bpy.ops.render.render(animation=True) 251 | 252 | # uncomment this to prevent generation of blend files 253 | bpy.ops.wm.save_as_mainfile( 254 | filepath=os.path.join(output_dir, f"{combination_index}.blend") 255 | ) 256 | 257 | logger.info(f"Rendered video saved to {render_path}") 258 | 259 | 260 | if __name__ == "__main__": 261 | parser = argparse.ArgumentParser() 262 | parser.add_argument( 263 | "--output_dir", 264 | type=str, 265 | default="renders", 266 | required=False, 267 | help="Path to the directory where the rendered video or images will be saved.", 268 | ) 269 | parser.add_argument( 270 | "--combination_file", 271 | type=str, 272 | default="combinations.json", 273 | help="Path to the JSON file containing camera combinations.", 274 | ) 275 | parser.add_argument( 276 | "--hdri_path", 277 | type=str, 278 | default="backgrounds", 279 | help="Path to the directory where the background HDRs will be saved.", 280 | ) 281 | parser.add_argument( 282 | "--combination_index", 283 | type=int, 284 | default=0, 285 | help="Index of the camera combination to use from the JSON file.", 286 | required=False, 287 | ) 288 | parser.add_argument( 289 | "--start_frame", 290 | type=int, 291 | default=1, 292 | help="Start frame of the animation.", 293 | required=False, 294 | ) 295 | parser.add_argument( 296 | "--end_frame", 297 | type=int, 298 | default=65, 299 | help="End frame of the animation.", 300 | required=False, 301 | ) 302 | parser.add_argument( 303 | "--width", type=int, default=1920, help="Render output width.", required=False 304 | ) 305 | parser.add_argument( 306 | "--height", type=int, default=1080, help="Render output height.", required=False 307 | ) 308 | parser.add_argument( 309 | "--combination", type=str, default=None, help="Combination dictionary." 310 | ) 311 | parser.add_argument( 312 | "--images", 313 | action="store_true", 314 | help="Generate images instead of videos.", 315 | ) 316 | parser.add_argument( 317 | "--blend", 318 | type=str, 319 | default=None, 320 | help="Path to the user-specified Blender file to use as the base scene.", 321 | required=False, 322 | ) 323 | parser.add_argument( 324 | "--animation_length", 325 | type=int, 326 | default=100, 327 | help="Percentage animation length. Defaults to 100%.", 328 | required=False 329 | ) 330 | 331 | if "--" in sys.argv: 332 | argv = sys.argv[sys.argv.index("--") + 1 :] 333 | else: 334 | argv = [] 335 | 336 | args = parser.parse_args(argv) 337 | 338 | context = bpy.context 339 | scene = context.scene 340 | render = scene.render 341 | 342 | if args.combination is not None: 343 | combination = json.loads(args.combination) 344 | else: 345 | combination = read_combination(args.combination_file, args.combination_index) 346 | 347 | # get the object uid from the 'object' column, which is a dictionary 348 | objects_column = combination["objects"] 349 | 350 | for object in objects_column: 351 | uid = object["uid"] 352 | 353 | downloaded = objaverse.load_objects([uid]) 354 | 355 | # Render the images 356 | render_scene( 357 | start_frame=args.start_frame, 358 | end_frame=args.end_frame, 359 | output_dir=args.output_dir, 360 | context=context, 361 | combination_file=args.combination_file, 362 | combination_index=args.combination_index, 363 | combination=args.combination, 364 | render_images=args.images, 365 | user_blend_file=args.blend, 366 | ) 367 | -------------------------------------------------------------------------------- /simian/scene.py: -------------------------------------------------------------------------------- 1 | from math import cos, sin 2 | import bpy 3 | from typing import Tuple 4 | import bmesh 5 | import os 6 | import requests 7 | 8 | 9 | def initialize_scene() -> None: 10 | # start bpy from scratch 11 | bpy.ops.wm.read_factory_settings(use_empty=True) 12 | 13 | # delete all objects 14 | bpy.ops.object.select_all(action="SELECT") 15 | 16 | bpy.ops.object.delete() 17 | 18 | # set render mode of blend file to eevee 19 | bpy.context.scene.render.engine = 'BLENDER_EEVEE' 20 | 21 | 22 | def download_texture(url: str, material_name: str, texture_name: str) -> str: 23 | """ 24 | Downloads the texture from the given URL and saves it in the materials/ folder. 25 | Returns the local file path of the downloaded texture. 26 | 27 | Args: 28 | url (str): The URL of the texture to download. 29 | material_name (str): The name of the material. 30 | texture_name (str): The name of the texture. 31 | 32 | Returns: 33 | str: The local file path of the downloaded texture. 34 | """ 35 | materials_dir = os.path.join("materials", material_name) 36 | os.makedirs(materials_dir, exist_ok=True) 37 | 38 | local_path = os.path.join(materials_dir, f"{texture_name}.jpg") 39 | 40 | if not os.path.exists(local_path): 41 | response = requests.get(url) 42 | with open(local_path, "wb") as file: 43 | file.write(response.content) 44 | 45 | return local_path 46 | 47 | 48 | def create_stage( 49 | combination: dict, 50 | stage_size: Tuple[int, int] = (100, 100), 51 | stage_height: float = 0.002, 52 | ) -> bpy.types.Object: 53 | """ 54 | Creates a simple stage object in the scene. 55 | 56 | Args: 57 | combination (dict): A dictionary containing the stage settings. 58 | stage_size (Tuple[int, int], optional): The size of the stage in Blender units (width, height). Defaults to (100, 100). 59 | stage_height (float, optional): The height of the stage above the ground plane. Defaults to 0.002. 60 | 61 | Returns: 62 | bpy.types.Object: The created stage object. 63 | """ 64 | # Create a new plane object 65 | bpy.ops.mesh.primitive_plane_add(size=1) 66 | stage = bpy.context.active_object 67 | 68 | stage_data = combination.get("stage", {}) 69 | stage_material = stage_data.get("material", {}) 70 | uv_scale = stage_data.get("uv_scale", [1.0, 1.0]) 71 | uv_rotation = stage_data.get("uv_rotation", 0.0) 72 | 73 | # Scale the stage to the desired size 74 | stage.scale = (stage_size[0], stage_size[1], 1) 75 | 76 | # Set the stage location to be at the bottom of the scene 77 | stage.location = (0, 0, stage_height) 78 | 79 | # Rename the stage object 80 | stage.name = "Stage" 81 | 82 | # Rescale the UVs based on the inverse of the scale multiplied by 2 83 | scale_x = stage_size[0] * uv_scale[0] 84 | scale_y = stage_size[1] * uv_scale[1] 85 | 86 | # Enter edit mode 87 | bpy.ops.object.mode_set(mode="EDIT") 88 | 89 | # Get the bmesh representation of the stage 90 | bm = bmesh.from_edit_mesh(stage.data) 91 | 92 | # Get the UV layer 93 | uv_layer = bm.loops.layers.uv.verify() 94 | 95 | # convert uv rotation to radians 96 | uv_rotation = uv_rotation * 3.14159 / 180.0 97 | 98 | # create a 2x2 rotation matrix for the UVs 99 | rotation_matrix = [ 100 | [cos(uv_rotation), -sin(uv_rotation)], 101 | [sin(uv_rotation), cos(uv_rotation)], 102 | ] 103 | 104 | # Iterate over the faces and rescale the UVs 105 | for face in bm.faces: 106 | for loop in face.loops: 107 | uv = loop[uv_layer].uv 108 | 109 | # rotate the UVs by multiplying by the rotation matrix 110 | uv.x, uv.y = ( 111 | (rotation_matrix[0][0] * uv.x + rotation_matrix[0][1] * uv.y) * scale_x, 112 | (rotation_matrix[1][0] * uv.x + rotation_matrix[1][1] * uv.y) * scale_y, 113 | ) 114 | 115 | # Update the mesh and return to object mode 116 | bmesh.update_edit_mesh(stage.data) 117 | bpy.ops.object.mode_set(mode="OBJECT") 118 | return stage 119 | 120 | 121 | def apply_stage_material(stage: bpy.types.Object, combination: dict) -> None: 122 | """ 123 | Applies the stage material to the given stage object based on the combination settings. 124 | 125 | Args: 126 | stage (bpy.types.Object): The stage object to apply the material to. 127 | combination (dict): A dictionary containing the stage material settings. 128 | """ 129 | # Get the stage material settings from the combination 130 | stage_data = combination.get("stage", {}) 131 | stage_material = stage_data.get("material", {}) 132 | material_name = stage_material.get("name", "DefaultMaterial") 133 | 134 | # Create a new material for the stage 135 | material = bpy.data.materials.new(name="StageMaterial") 136 | 137 | # Assign the material to the stage object 138 | stage.data.materials.append(material) 139 | 140 | # Set the material properties based on the combination settings 141 | material.use_nodes = True 142 | nodes = material.node_tree.nodes 143 | 144 | # Clear existing nodes 145 | for node in nodes: 146 | nodes.remove(node) 147 | 148 | # Create shader nodes 149 | output = nodes.new(type="ShaderNodeOutputMaterial") 150 | principled = nodes.new(type="ShaderNodeBsdfPrincipled") 151 | 152 | # Create texture coordinate and mapping nodes 153 | tex_coord = nodes.new(type="ShaderNodeTexCoord") 154 | mapping = nodes.new(type="ShaderNodeMapping") 155 | 156 | # Connect texture coordinate to mapping 157 | links = material.node_tree.links 158 | links.new(tex_coord.outputs["UV"], mapping.inputs["Vector"]) 159 | 160 | # Load and connect diffuse texture 161 | if "Diffuse" in stage_material["maps"]: 162 | diffuse_url = stage_material["maps"]["Diffuse"] 163 | diffuse_path = download_texture(diffuse_url, material_name, "Diffuse") 164 | diffuse_tex = nodes.new(type="ShaderNodeTexImage") 165 | diffuse_tex.image = bpy.data.images.load(diffuse_path) 166 | links.new(mapping.outputs["Vector"], diffuse_tex.inputs["Vector"]) 167 | links.new(diffuse_tex.outputs["Color"], principled.inputs["Base Color"]) 168 | 169 | # Load and connect normal texture 170 | if "nor_gl" in stage_material["maps"]: 171 | normal_url = stage_material["maps"]["nor_gl"] 172 | normal_path = download_texture(normal_url, material_name, "Normal") 173 | normal_tex = nodes.new(type="ShaderNodeTexImage") 174 | normal_tex.image = bpy.data.images.load(normal_path) 175 | normal_map = nodes.new(type="ShaderNodeNormalMap") 176 | links.new(mapping.outputs["Vector"], normal_tex.inputs["Vector"]) 177 | links.new(normal_tex.outputs["Color"], normal_map.inputs["Color"]) 178 | links.new(normal_map.outputs["Normal"], principled.inputs["Normal"]) 179 | 180 | if "AO" in stage_material["maps"]: 181 | ao_url = stage_material["maps"]["AO"] 182 | ao_path = download_texture(ao_url, material_name, "AO") 183 | ao_tex = nodes.new(type="ShaderNodeTexImage") 184 | ao_tex.image = bpy.data.images.load(ao_path) 185 | mixRGB = nodes.new(type="ShaderNodeMixRGB") 186 | mixRGB.blend_type = "MULTIPLY" 187 | links.new(ao_tex.outputs["Color"], mixRGB.inputs["Color2"]) 188 | links.new(mapping.outputs["Vector"], ao_tex.inputs["Vector"]) 189 | 190 | # Connect the MixRGB node to the base color input 191 | if "Diffuse" in stage_material["maps"]: 192 | links.new(diffuse_tex.outputs["Color"], mixRGB.inputs["Color1"]) 193 | links.new(mixRGB.outputs["Color"], principled.inputs["Base Color"]) 194 | else: 195 | # If no diffuse texture, use a default base color 196 | principled.inputs["Base Color"].default_value = (1.0, 1.0, 1.0, 1.0) 197 | links.new(ao_tex.outputs["Color"], mixRGB.inputs["Color1"]) 198 | links.new(mixRGB.outputs["Color"], principled.inputs["Base Color"]) 199 | 200 | if "Rough" in stage_material["maps"]: 201 | rough_url = stage_material["maps"]["Rough"] 202 | rough_path = download_texture(rough_url, material_name, "Rough") 203 | rough_tex = nodes.new(type="ShaderNodeTexImage") 204 | rough_tex.image = bpy.data.images.load(rough_path) 205 | links.new(mapping.outputs["Vector"], rough_tex.inputs["Vector"]) 206 | links.new(rough_tex.outputs["Color"], principled.inputs["Roughness"]) 207 | 208 | if "Roughness" in stage_material["maps"]: 209 | roughness_url = stage_material["maps"]["Roughness"] 210 | roughness_path = download_texture(roughness_url, material_name, "Roughness") 211 | roughness_tex = nodes.new(type="ShaderNodeTexImage") 212 | roughness_tex.image = bpy.data.images.load(roughness_path) 213 | links.new(mapping.outputs["Vector"], roughness_tex.inputs["Vector"]) 214 | links.new(roughness_tex.outputs["Color"], principled.inputs["Roughness"]) 215 | 216 | if "arm" in stage_material["maps"]: 217 | arm_url = stage_material["maps"]["arm"] 218 | arm_path = download_texture(arm_url, material_name, "Arm") 219 | arm_tex = nodes.new(type="ShaderNodeTexImage") 220 | arm_tex.image = bpy.data.images.load(arm_path) 221 | links.new(mapping.outputs["Vector"], arm_tex.inputs["Vector"]) 222 | 223 | # Create separate RGB nodes for ambient occlusion, roughness, and metallic 224 | ao_rgb = nodes.new(type="ShaderNodeSeparateRGB") 225 | rough_rgb = nodes.new(type="ShaderNodeSeparateRGB") 226 | metal_rgb = nodes.new(type="ShaderNodeSeparateRGB") 227 | 228 | # Connect the "arm" texture to the separate RGB nodes 229 | links.new(arm_tex.outputs["Color"], ao_rgb.inputs["Image"]) 230 | links.new(arm_tex.outputs["Color"], rough_rgb.inputs["Image"]) 231 | links.new(arm_tex.outputs["Color"], metal_rgb.inputs["Image"]) 232 | 233 | # Multiply the ambient occlusion with the base color 234 | mixRGB = nodes.new(type="ShaderNodeMixRGB") 235 | mixRGB.blend_type = "MULTIPLY" 236 | links.new(ao_rgb.outputs["R"], mixRGB.inputs["Color2"]) 237 | 238 | if "Diffuse" in stage_material["maps"]: 239 | links.new(diffuse_tex.outputs["Color"], mixRGB.inputs["Color1"]) 240 | links.new(mixRGB.outputs["Color"], principled.inputs["Base Color"]) 241 | else: 242 | # If no diffuse texture, use a default base color 243 | principled.inputs["Base Color"].default_value = (0.8, 0.8, 0.8, 1.0) 244 | links.new(ao_rgb.outputs["R"], mixRGB.inputs["Color1"]) 245 | links.new(mixRGB.outputs["Color"], principled.inputs["Base Color"]) 246 | 247 | # Connect roughness and metallic to the principled BSDF node 248 | links.new(rough_rgb.outputs["G"], principled.inputs["Roughness"]) 249 | links.new(metal_rgb.outputs["B"], principled.inputs["Metallic"]) 250 | 251 | # Load and connect rough_ao texture 252 | if "rough_ao" in stage_material["maps"]: 253 | rough_ao_url = stage_material["maps"]["rough_ao"] 254 | rough_ao_path = download_texture(rough_ao_url, material_name, "RoughAO") 255 | rough_ao_tex = nodes.new(type="ShaderNodeTexImage") 256 | rough_ao_tex.image = bpy.data.images.load(rough_ao_path) 257 | links.new(mapping.outputs["Vector"], rough_ao_tex.inputs["Vector"]) 258 | links.new(rough_ao_tex.outputs["Color"], principled.inputs["Roughness"]) 259 | 260 | # Load and connect displacement texture 261 | if "Displacement" in stage_material["maps"]: 262 | disp_url = stage_material["maps"]["Displacement"] 263 | disp_path = download_texture(disp_url, material_name, "Displacement") 264 | disp_tex = nodes.new(type="ShaderNodeTexImage") 265 | disp_tex.image = bpy.data.images.load(disp_path) 266 | disp_node = nodes.new(type="ShaderNodeDisplacement") 267 | links.new(mapping.outputs["Vector"], disp_tex.inputs["Vector"]) 268 | links.new(disp_tex.outputs["Color"], disp_node.inputs["Height"]) 269 | links.new(disp_node.outputs["Displacement"], output.inputs["Displacement"]) 270 | 271 | # Connect the nodes 272 | links.new(principled.outputs["BSDF"], output.inputs["Surface"]) 273 | -------------------------------------------------------------------------------- /simian/server.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | import chromadb 4 | from chromadb.utils import embedding_functions 5 | from rich.progress import Progress, TextColumn, BarColumn, TimeRemainingColumn, MofNCompleteColumn 6 | from rich.console import Console 7 | from rich.panel import Panel 8 | from sentence_transformers import SentenceTransformer 9 | from chromadb.config import Settings 10 | 11 | 12 | def initialize_chroma_db(reset_hdri=False, reset_textures=False): 13 | """ 14 | Initialize the Chroma database with the provided data files. 15 | 16 | Args: 17 | reset_hdri (bool): Whether to reset and reprocess the HDRI backgrounds. 18 | reset_textures (bool): Whether to reset and reprocess the textures. 19 | 20 | Returns: 21 | chroma_client (chromadb.PersistentClient): The initialized Chroma client. 22 | """ 23 | db_path = "./chroma_db" 24 | 25 | # Check if the database directory exists 26 | if not os.path.exists(db_path): 27 | print("Database not found. Creating new database.") 28 | os.makedirs(db_path) 29 | 30 | # Initialize Chroma client 31 | chroma_client = chromadb.PersistentClient(path=db_path, settings=Settings(anonymized_telemetry=False)) 32 | 33 | # Create or get collections for each data type 34 | object_collection = chroma_client.get_or_create_collection(name="object_captions") 35 | hdri_collection = chroma_client.get_or_create_collection(name="hdri_backgrounds") 36 | texture_collection = chroma_client.get_or_create_collection(name="textures") 37 | 38 | # Process collections if they are empty 39 | def process_if_empty(collection, file_path, description): 40 | if collection.count() == 0: 41 | if os.path.exists(file_path): 42 | print(f"Processing {description}...") 43 | process_in_batches(file_path, collection, batch_size=1000) 44 | else: 45 | print(f"File not found: {file_path}. Please check the file path.") 46 | else: 47 | print(f"{description} already processed. Skipping.") 48 | 49 | # Adjusted paths to reflect the correct location of the datasets folder 50 | datasets_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../datasets")) 51 | 52 | def reset_and_process(collection, file_path, description): 53 | print(f"Resetting and reprocessing {description}...") 54 | # Get all IDs in the collection 55 | all_ids = collection.get(include=['embeddings'])['ids'] 56 | if all_ids: 57 | # Delete all documents if there are any 58 | collection.delete(ids=all_ids) 59 | # Process the data 60 | process_in_batches(file_path, collection, batch_size=1000) 61 | 62 | # Reset and reprocess HDRI if requested 63 | if reset_hdri: 64 | reset_and_process(hdri_collection, os.path.join(datasets_path, 'hdri_data.json'), "HDRI backgrounds") 65 | elif hdri_collection.count() == 0: 66 | process_in_batches(os.path.join(datasets_path, 'hdri_data.json'), hdri_collection, batch_size=1000) 67 | 68 | # Reset and reprocess textures if requested 69 | if reset_textures: 70 | reset_and_process(texture_collection, os.path.join(datasets_path, 'texture_data.json'), "textures") 71 | elif texture_collection.count() == 0: 72 | process_in_batches(os.path.join(datasets_path, 'texture_data.json'), texture_collection, batch_size=1000) 73 | 74 | process_if_empty(object_collection, os.path.join(datasets_path, 'cap3d_captions.json'), "object captions") 75 | process_if_empty(hdri_collection, os.path.join(datasets_path, 'hdri_data.json'), "HDRI backgrounds") 76 | process_if_empty(texture_collection, os.path.join(datasets_path, 'texture_data.json'), "textures") 77 | 78 | # Check if collections are empty and process data if needed 79 | if object_collection.count() == 0: 80 | print("Processing object captions...") 81 | process_in_batches_objects('../datasets/cap3d_captions.json', object_collection, batch_size=1000) 82 | else: 83 | print("Object captions already processed. Skipping.") 84 | 85 | if hdri_collection.count() == 0: 86 | print("Processing HDRI backgrounds...") 87 | process_in_batches('../datasets/hdri_data.json', hdri_collection, batch_size=1000) 88 | else: 89 | print("HDRI backgrounds already processed. Skipping.") 90 | 91 | if texture_collection.count() == 0: 92 | print("Processing textures...") 93 | process_in_batches('../datasets/texture_data.json', texture_collection, batch_size=1000) 94 | else: 95 | print("Textures already processed. Skipping.") 96 | 97 | print("Database initialization complete.") 98 | 99 | return chroma_client 100 | 101 | 102 | def process_in_batches_objects(file_path, collection, batch_size=1000): 103 | """ 104 | Process the data in the specified file in batches and upsert it into the collection for objects. 105 | 106 | Args: 107 | file_path (str): The path to the file containing the data. 108 | collection (chromadb.Collection): The collection to upsert the data into. 109 | batch_size (int): The size of each batch for processing. 110 | 111 | Returns: 112 | None 113 | """ 114 | sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name='all-MiniLM-L6-v2') 115 | console = Console() 116 | 117 | with open(file_path, 'r') as file: 118 | data = json.load(file) 119 | 120 | all_ids = list(data.keys()) 121 | total_items = len(all_ids) 122 | 123 | with Progress( 124 | TextColumn("[progress.description]{task.description}"), 125 | BarColumn(), 126 | MofNCompleteColumn(), 127 | TimeRemainingColumn(), 128 | TextColumn("Batch: {task.fields[current_batch]}"), 129 | console=console, 130 | expand=True 131 | ) as progress: 132 | task = progress.add_task(f"[green]Processing {file_path}", total=total_items, current_batch="0") 133 | 134 | for i in range(0, total_items, batch_size): 135 | batch_ids = all_ids[i:i+batch_size] 136 | batch_documents = [data[id] for id in batch_ids] 137 | 138 | # Compute embeddings for the batch using the embedding function 139 | embeddings = sentence_transformer_ef(batch_documents) 140 | 141 | # Upsert the batch into the collection 142 | collection.upsert( 143 | ids=batch_ids, 144 | embeddings=embeddings, 145 | documents=batch_documents 146 | ) 147 | 148 | # Update progress 149 | progress.update(task, advance=len(batch_ids), current_batch=f"{i+1}-{min(i+batch_size, total_items)}") 150 | 151 | console.print(Panel.fit(f"Data processing complete for {file_path}!", border_style="green")) 152 | 153 | 154 | def process_in_batches(file_path, collection, batch_size=1000): 155 | """ 156 | Process the data in the specified file in batches and upsert it into the collection. 157 | 158 | Args: 159 | file_path (str): The path to the file containing the data. 160 | collection (chromadb.Collection): The collection to upsert the data into. 161 | batch_size (int): The size of each batch for processing. 162 | 163 | Returns: 164 | None 165 | """ 166 | 167 | sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name='all-MiniLM-L6-v2') 168 | 169 | console = Console() 170 | 171 | with open(file_path, 'r') as file: 172 | data = json.load(file) 173 | 174 | all_ids = list(data.keys()) 175 | total_items = len(all_ids) 176 | 177 | def convert_to_chroma_compatible(value): 178 | if isinstance(value, (str, int, float, bool)): 179 | return value 180 | elif isinstance(value, list): 181 | return ', '.join(map(str, value)) 182 | elif isinstance(value, dict): 183 | return json.dumps(value) 184 | else: 185 | return str(value) 186 | 187 | with Progress( 188 | TextColumn("[progress.description]{task.description}"), 189 | BarColumn(), 190 | MofNCompleteColumn(), 191 | TimeRemainingColumn(), 192 | TextColumn("Batch: {task.fields[current_batch]}"), 193 | console=console, 194 | expand=True 195 | ) as progress: 196 | task = progress.add_task(f"[green]Processing {file_path}", total=total_items, current_batch="0") 197 | 198 | for i in range(0, total_items, batch_size): 199 | batch_ids = all_ids[i:i+batch_size] 200 | batch_documents = [] 201 | batch_metadatas = [] 202 | 203 | for id in batch_ids: 204 | item = data[id] 205 | if isinstance(item, str): 206 | # For object captions, keep as is 207 | batch_documents.append(item) 208 | batch_metadatas.append(None) 209 | else: 210 | # For HDRIs and textures 211 | name = item.get('name', id) 212 | categories = ' '.join(item.get('categories', [])) 213 | tags = ' '.join(item.get('tags', [])) 214 | 215 | # Create a searchable document string 216 | document = f"{name} {categories} {tags}".strip() 217 | batch_documents.append(document) 218 | 219 | # Convert metadata to Chroma-compatible format 220 | compatible_metadata = {k: convert_to_chroma_compatible(v) for k, v in item.items()} 221 | batch_metadatas.append(compatible_metadata) 222 | 223 | # Compute embeddings for the batch using the embedding function 224 | embeddings = sentence_transformer_ef(batch_documents) 225 | 226 | # Upsert the batch into the collection 227 | collection.upsert( 228 | ids=batch_ids, 229 | embeddings=embeddings, 230 | documents=batch_documents, 231 | metadatas=batch_metadatas 232 | ) 233 | 234 | # Update progress 235 | progress.update(task, advance=len(batch_ids), current_batch=f"{i+1}-{min(i+batch_size, total_items)}") 236 | 237 | console.print(Panel.fit(f"Data processing complete for {file_path}!", border_style="green")) 238 | 239 | 240 | def query_collection(query, collection, n_results=2): 241 | """ 242 | Query the specified collection with the given query and display the results. 243 | 244 | Args: 245 | query (str): The query string to search for. 246 | sentence_transformer_ef (embedding_functions.SentenceTransformerEmbeddingFunction): The SentenceTransformer embedding function. 247 | collection (chromadb.Collection): The collection to query. 248 | n_results (int): The number of results to display. 249 | 250 | Returns: 251 | dict: The query results. 252 | """ 253 | 254 | model = SentenceTransformer('all-MiniLM-L6-v2') # or another appropriate model 255 | sentence_transformer_ef = model.encode 256 | sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name='all-MiniLM-L6-v2') 257 | 258 | console = Console() 259 | query_embedding = sentence_transformer_ef([query]) 260 | results = collection.query( 261 | query_embeddings=query_embedding, 262 | n_results=n_results, 263 | include=["metadatas", "documents", "distances"] 264 | ) 265 | 266 | def parse_metadata_value(value): 267 | if isinstance(value, str): 268 | if value.startswith('[') and value.endswith(']'): 269 | return value.strip('[]').split(', ') 270 | elif value.startswith('{') and value.endswith('}'): 271 | return json.loads(value) 272 | return value 273 | 274 | # Parse the metadata values 275 | for i, metadata in enumerate(results['metadatas'][0]): 276 | if metadata: 277 | results['metadatas'][0][i] = {k: parse_metadata_value(v) for k, v in metadata.items()} 278 | 279 | console.print(Panel(str(results), title=f"Query Results for {collection.name}", expand=False)) 280 | return results -------------------------------------------------------------------------------- /simian/tests/__run__.py: -------------------------------------------------------------------------------- 1 | import os 2 | import subprocess 3 | import sys 4 | import argparse 5 | 6 | # Parse command-line arguments 7 | parser = argparse.ArgumentParser() 8 | parser.add_argument( 9 | "--debug", action="store_true", help="Print the command to run each subtest" 10 | ) 11 | args = parser.parse_args() 12 | 13 | files = os.listdir("simian/tests") 14 | test_files = [file for file in files if file.endswith("_test.py")] 15 | 16 | # Flag to track if any test has failed 17 | test_failed = False 18 | 19 | for test_file in test_files: 20 | print(f"Running tests in {test_file}...") 21 | # get the module name from the file by removing .py 22 | module_name = test_file[:-3] 23 | cmd = [sys.executable, "-m", f"simian.tests.{module_name}"] 24 | 25 | if args.debug: 26 | print(f"Command: {' '.join(cmd)}") 27 | 28 | # Run the test file and capture the output 29 | result = subprocess.run(cmd, capture_output=True, text=True) 30 | 31 | # Check if there are any errors in the output 32 | if ( 33 | "FAILED" in result.stderr 34 | or "ERROR" in result.stderr 35 | or "ValueError" in result.stderr 36 | or "Traceback" in result.stderr 37 | ): 38 | print(f"Tests in {test_file} failed!") 39 | print(result.stderr) 40 | test_failed = True 41 | else: 42 | print(f"Tests in {test_file} passed.") 43 | 44 | # Exit with appropriate exit code based on test results 45 | if test_failed: 46 | print("Some tests failed. Exiting with exit code 1.") 47 | sys.exit(1) 48 | else: 49 | print("All tests passed. Exiting with exit code 0.") 50 | sys.exit(0) 51 | -------------------------------------------------------------------------------- /simian/tests/background_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from unittest.mock import patch, MagicMock 4 | from ..background import ( 5 | get_hdri_path, 6 | get_background, 7 | set_background, 8 | create_photosphere, 9 | create_photosphere_material, 10 | ) 11 | 12 | import bpy 13 | from ..background import set_background 14 | 15 | 16 | def test_get_hdri_path(): 17 | """ 18 | Test the get_hdri_path function. 19 | """ 20 | combination = {"background": {"id": "123", "from": "test_dataset"}} 21 | hdri_path = "/fake/path" 22 | expected_result = f"/fake/path/test_dataset/123.hdr" 23 | result = get_hdri_path(hdri_path, combination) 24 | 25 | print("test_get_hdri_path result: ", result) 26 | assert result == expected_result 27 | print("============ Test Passed: get_hdri_path ============") 28 | 29 | 30 | def test_get_background(): 31 | """ 32 | Test the get_background function. 33 | """ 34 | combination = { 35 | "background": { 36 | "url": "http://example.com/image.hdr", 37 | "id": "123", 38 | "from": "test_dataset", 39 | } 40 | } 41 | hdri_path = "/fake/path" 42 | 43 | with patch("os.makedirs"), patch("os.path.exists", return_value=False), patch( 44 | "requests.get" 45 | ) as mock_get, patch("builtins.open", new_callable=MagicMock): 46 | mock_response = MagicMock() 47 | mock_response.content = b"fake data" 48 | mock_get.return_value = mock_response 49 | 50 | get_background(hdri_path, combination) 51 | print("get_background called") 52 | 53 | mock_get.assert_called_with("http://example.com/image.hdr") 54 | print("============ Test Passed: test_get_background ============") 55 | 56 | 57 | def test_set_background(): 58 | """ 59 | Test the set_background function. 60 | """ 61 | combination = { 62 | "background": { 63 | "id": "123", 64 | "from": "test_dataset", 65 | "url": "http://example.com/image.hdr", 66 | } 67 | } 68 | # Use a path in the user's home directory or another writable location 69 | background_base_path = os.path.join(os.path.expanduser("~"), "test_backgrounds") 70 | 71 | # Mocking dependencies 72 | with patch( 73 | "..background.get_hdri_path", 74 | return_value=os.path.join(background_base_path, "test_dataset/123.hdr"), 75 | ) as mock_get_path: 76 | with patch("os.path.exists", return_value=False): 77 | with patch("requests.get") as mock_get: 78 | mock_response = MagicMock() 79 | mock_response.content = b"Fake HDR content" 80 | mock_get.return_value = mock_response 81 | 82 | # Ensure the directory exists before writing 83 | os.makedirs( 84 | os.path.dirname( 85 | os.path.join(background_base_path, "test_dataset/123.hdr") 86 | ), 87 | exist_ok=True, 88 | ) 89 | 90 | # Function call 91 | set_background(background_base_path, combination) 92 | 93 | # Open the file to simulate writing to it 94 | with open( 95 | os.path.join(background_base_path, "test_dataset/123.hdr"), "wb" 96 | ) as file: 97 | file.write(mock_response.content) 98 | 99 | mock_get.assert_called_with("http://example.com/image.hdr") 100 | print("============ Test Passed: test_set_background ============") 101 | 102 | 103 | def create_test_photosphere(): 104 | """ 105 | Create a UV sphere in Blender with specific parameters. 106 | """ 107 | # Create a UV sphere in Blender with specific parameters 108 | bpy.ops.mesh.primitive_uv_sphere_add( 109 | segments=64, ring_count=32, radius=1.0, location=(0, 0, 3) 110 | ) 111 | sphere = bpy.context.object 112 | bpy.ops.object.shade_smooth() 113 | 114 | # Enter edit mode, select all vertices, and flip the normals 115 | bpy.ops.object.mode_set(mode="EDIT") 116 | bpy.ops.mesh.select_all(action="SELECT") 117 | bpy.ops.mesh.flip_normals() 118 | bpy.ops.object.mode_set(mode="OBJECT") 119 | 120 | # Rename the sphere for identification 121 | sphere.name = "Photosphere" 122 | print("Sphere created successfully") 123 | return sphere 124 | 125 | 126 | def test_create_photosphere(): 127 | epsilon = 0.001 # Small threshold for floating-point comparisons 128 | sphere = create_test_photosphere() 129 | 130 | # Check each component of the sphere's location to see if it matches the expected values 131 | assert abs(sphere.location.x - 0.0) < epsilon, "X coordinate is incorrect" 132 | assert abs(sphere.location.y - 0.0) < epsilon, "Y coordinate is incorrect" 133 | assert abs(sphere.location.z - 3.0) < epsilon, "Z coordinate is incorrect" 134 | print("============ Test Passed: test_create_photosphere ============") 135 | 136 | 137 | def test_create_photosphere_material(): 138 | """ 139 | Test the create_photosphere_material function. 140 | """ 141 | # Create a UV sphere in Blender with specific parameters 142 | bpy.ops.mesh.primitive_uv_sphere_add( 143 | segments=64, ring_count=32, radius=1.0, location=(0, 0, 3) 144 | ) 145 | sphere = bpy.context.object 146 | bpy.ops.object.shade_smooth() 147 | 148 | # Enter edit mode, select all vertices, and flip the normals 149 | bpy.ops.object.mode_set(mode="EDIT") 150 | bpy.ops.mesh.select_all(action="SELECT") 151 | bpy.ops.mesh.flip_normals() 152 | bpy.ops.object.mode_set(mode="OBJECT") 153 | 154 | # Rename the sphere for identification 155 | sphere.name = "Photosphere" 156 | 157 | # Create a combination dictionary with background information 158 | combination = { 159 | "background": { 160 | "id": "123", 161 | "from": "test_dataset", 162 | "url": "http://example.com/image.hdr", 163 | } 164 | } 165 | 166 | # Set the base path for background images 167 | background_base_path = os.path.join(os.path.expanduser("~"), "test_backgrounds") 168 | 169 | # Mock the get_hdri_path function to return a known path 170 | with patch( 171 | "..background.get_hdri_path", 172 | return_value=os.path.join(background_base_path, "test_dataset/123.hdr"), 173 | ) as mock_get_path: 174 | # Mock the existence of the background image file 175 | with patch("os.path.exists", return_value=True): 176 | # Mock the open function to simulate reading the background image 177 | with patch("builtins.open", MagicMock()) as mock_open: 178 | # Call the create_photosphere_material function 179 | create_photosphere_material(background_base_path, combination, sphere) 180 | 181 | # Verify that the material was created 182 | assert ( 183 | sphere.data.materials[0].name == "PhotosphereMaterial" 184 | ), "Material not created successfully" 185 | print( 186 | "============ Test Passed: test_create_photosphere_material ============" 187 | ) 188 | 189 | 190 | # Run tests if this file is executed as a script 191 | if __name__ == "__main__": 192 | test_get_hdri_path() 193 | test_get_background() 194 | # test_set_background() 195 | test_create_photosphere() 196 | # test_create_photosphere_material() 197 | print("============ ALL TESTS PASSED ============") 198 | -------------------------------------------------------------------------------- /simian/tests/batch_test.py: -------------------------------------------------------------------------------- 1 | from ..batch import render_objects 2 | 3 | 4 | def test_render_objects(): 5 | """ 6 | Test the render_objects function. 7 | """ 8 | processes = 4 9 | render_timeout = 3000 10 | width = 1920 11 | height = 1080 12 | start_index = 0 13 | end_index = 1 14 | start_frame = 1 15 | end_frame = 2 16 | 17 | # Call the function 18 | try: 19 | render_objects( 20 | processes=processes, 21 | render_timeout=render_timeout, 22 | width=width, 23 | height=height, 24 | start_index=start_index, 25 | end_index=end_index, 26 | start_frame=start_frame, 27 | end_frame=end_frame, 28 | ) 29 | print("============ Test Passed: render_objects ============") 30 | except Exception as e: 31 | print("============ Test Failed: render_objects ============") 32 | print(f"Error: {e}") 33 | raise e 34 | 35 | 36 | if __name__ == "__main__": 37 | test_render_objects() 38 | print("============ ALL TESTS PASSED ============") 39 | -------------------------------------------------------------------------------- /simian/tests/camera_test.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import bpy 4 | from ..camera import ( 5 | create_camera_rig, 6 | set_camera_settings, 7 | position_camera, 8 | ) 9 | from ..scene import initialize_scene 10 | 11 | 12 | def test_create_camera_rig(): 13 | """ 14 | Test the create_camera_rig function. 15 | """ 16 | 17 | EPSILON = 1e-6 # Small value to account for floating-point inaccuracies 18 | rig = create_camera_rig() # call the function 19 | 20 | # Check if all the expected objects are created 21 | assert isinstance(rig["camera_animation_root"], bpy.types.Object) 22 | assert isinstance(rig["camera_orientation_pivot_yaw"], bpy.types.Object) 23 | assert isinstance(rig["camera_orientation_pivot_pitch"], bpy.types.Object) 24 | assert isinstance(rig["camera_framing_pivot"], bpy.types.Object) 25 | assert isinstance(rig["camera_animation_pivot"], bpy.types.Object) 26 | assert isinstance(rig["camera_object"], bpy.types.Object) 27 | assert isinstance(rig["camera"], bpy.types.Camera) 28 | 29 | # Check if the camera is set correctly 30 | assert bpy.context.scene.camera == rig["camera_object"] 31 | 32 | # Check if the camera rotation is set correctly 33 | camera_rotation = rig["camera_object"].delta_rotation_euler 34 | 35 | # Expected rotation values 36 | expected_x_rotation = 1.5708 # 90 degrees 37 | expected_y_rotation = 0.0 38 | expected_z_rotation = 1.5708 # 90 degrees 39 | 40 | # Check rotation values with epsilon error 41 | assert math.isclose(camera_rotation[0], expected_x_rotation, abs_tol=EPSILON) 42 | assert math.isclose(camera_rotation[1], expected_y_rotation, abs_tol=EPSILON) 43 | assert math.isclose(camera_rotation[2], expected_z_rotation, abs_tol=EPSILON) 44 | print("============ Test Passed: test_create_camera_rig ============") 45 | 46 | 47 | def test_set_camera_settings(): 48 | """ 49 | Test the set_camera_settings function. 50 | """ 51 | 52 | initialize_scene() 53 | create_camera_rig() 54 | 55 | combination = { 56 | "orientation": {"yaw": 327, "pitch": 14}, 57 | "framing": {"fov": 20, "coverage_factor": 1.0}, 58 | "animation": { 59 | "name": "tilt_left", 60 | "keyframes": [ 61 | { 62 | "CameraAnimationRoot": {"rotation": [0, 0, 45]}, 63 | "Camera": {"angle_offset": 5}, 64 | }, 65 | {"CameraAnimationRoot": {"rotation": [0, 0, 0]}}, 66 | ], 67 | }, 68 | } 69 | 70 | # Retrieve the camera object from the Blender scene 71 | camera = bpy.data.objects["Camera"].data 72 | 73 | # set the camera lens_unit type to FOV 74 | camera.lens_unit = "FOV" 75 | 76 | # Call the function with full data 77 | set_camera_settings(combination) 78 | 79 | # Retrieve the orientation pivot objects 80 | camera_orientation_pivot_yaw = bpy.data.objects["CameraOrientationPivotYaw"] 81 | camera_orientation_pivot_pitch = bpy.data.objects["CameraOrientationPivotPitch"] 82 | print("camera.angle") 83 | print(camera.angle) 84 | print('combination["framing"]["fov"]') 85 | print(math.radians(combination["framing"]["fov"])) 86 | print( 87 | 'combination["framing"]["fov"] + combination["animation"]["keyframes"][0]["Camera"]["angle_offset"]' 88 | ) 89 | print( 90 | math.radians( 91 | combination["framing"]["fov"] 92 | + combination["animation"]["keyframes"][0]["Camera"]["angle_offset"] 93 | ) 94 | ) 95 | fov = math.radians( 96 | combination["framing"]["fov"] 97 | + combination["animation"]["keyframes"][0]["Camera"]["angle_offset"] 98 | ) 99 | # Assert the field of view is set correctly 100 | epsilon = 0.0001 101 | assert ( 102 | camera.angle <= fov + epsilon and camera.angle >= fov - epsilon 103 | ), "FOV is not set correctly" 104 | 105 | # Convert degrees to radians for comparison 106 | expected_yaw_radians = math.radians(combination["orientation"]["yaw"]) 107 | expected_pitch_radians = -math.radians( 108 | combination["orientation"]["pitch"] 109 | ) # Negative for Blender's coordinate system 110 | 111 | # Assert the orientation is set correctly 112 | assert math.isclose( 113 | camera_orientation_pivot_yaw.rotation_euler[2], 114 | expected_yaw_radians, 115 | abs_tol=0.001, 116 | ), "Yaw is not set correctly" 117 | assert math.isclose( 118 | camera_orientation_pivot_pitch.rotation_euler[1], 119 | expected_pitch_radians, 120 | abs_tol=0.001, 121 | ), "Pitch is not set correctly" 122 | print("============ Test Passed: test_set_camera_settings ============") 123 | 124 | 125 | def test_position_camera(): 126 | combination = { 127 | "framing": {"fov": 20, "coverage_factor": 1.0}, 128 | "animation": { 129 | "keyframes": [ 130 | { 131 | "Camera": { 132 | "position": (0, 0, 5), 133 | "rotation": (0, 0, 0), 134 | "angle_offset": 5, 135 | } 136 | }, 137 | { 138 | "Camera": { 139 | "position": (5, 0, 0), 140 | "rotation": (0, 0, 90), 141 | "angle_offset": 10, 142 | } 143 | }, 144 | ] 145 | }, 146 | } 147 | 148 | # Create a dummy object to represent the focus object 149 | bpy.ops.mesh.primitive_cube_add(size=2) 150 | focus_object = bpy.context.active_object 151 | 152 | position_camera(combination, focus_object) 153 | 154 | camera = bpy.data.objects.get("Camera") 155 | assert camera is not None, "Camera object not found" 156 | 157 | # TODO: add more asserts here 158 | 159 | print("============ Test Passed: test_position_camera ============") 160 | 161 | 162 | if __name__ == "__main__": 163 | test_create_camera_rig() 164 | test_set_camera_settings() 165 | # test_set_camera_animation() 166 | test_position_camera() 167 | print("============ ALL TESTS PASSED ============") 168 | -------------------------------------------------------------------------------- /simian/tests/new_camera_test.py: -------------------------------------------------------------------------------- 1 | import bpy 2 | import math 3 | from mathutils import Vector 4 | 5 | from ..camera import create_camera_rig 6 | 7 | 8 | def calculate_optimal_distance(camera, obj): 9 | """ 10 | Calculate the optimal distance between the camera and an object 11 | so that the entire object is within the camera's frame, considering the object's depth. 12 | 13 | Args: 14 | camera (bpy.types.Object): The camera object. 15 | obj (bpy.types.Object): The object to be framed. 16 | 17 | Returns: 18 | float: The optimal distance between the camera and the object. 19 | """ 20 | # Get the camera's field of view 21 | fov_h = camera.data.angle_x 22 | fov_v = camera.data.angle_y 23 | 24 | # Calculate the object's bounding box dimensions in world space 25 | bbox = [obj.matrix_world @ Vector(corner) for corner in obj.bound_box] 26 | bbox_dimensions = [max(coord) - min(coord) for coord in zip(*bbox)] 27 | obj_width, obj_height, obj_depth = bbox_dimensions 28 | 29 | # Find the diagonal of the bounding box which will require the farthest distance 30 | diagonal = math.sqrt(obj_width**2 + obj_height**2 + obj_depth**2) 31 | 32 | # Calculate the required distance to fit the diagonal in view 33 | fov_average = (fov_h + fov_v) / 2 34 | distance = (diagonal / 2) / math.tan(fov_average / 2) 35 | 36 | return distance 37 | 38 | 39 | def position_camera_for_object(camera, obj): 40 | """ 41 | Position the camera so that all vertices of the object are in frame or at the frame border. 42 | 43 | Args: 44 | camera (bpy.types.Object): The camera object. 45 | obj (bpy.types.Object): The object to frame. 46 | """ 47 | optimal_distance = calculate_optimal_distance(camera, obj) 48 | camera.location = obj.location + Vector((optimal_distance, 0, 0)) 49 | camera.rotation_euler = (0, 0, 0) # Reset rotation or set as needed 50 | 51 | 52 | def is_vertex_in_frame(camera, vertex): 53 | """ 54 | Check if a vertex is within the camera's frame. 55 | 56 | Args: 57 | camera (bpy.types.Object): The camera object. 58 | vertex (mathutils.Vector): The vertex to check. 59 | 60 | Returns: 61 | bool: True if the vertex is within the camera's frame, False otherwise. 62 | """ 63 | # Transform the vertex to camera space 64 | camera_matrix = camera.matrix_world.inverted() 65 | vertex_camera = camera_matrix @ vertex 66 | 67 | # Check if the vertex is within the camera's view frustum 68 | fov_h = camera.data.angle_x 69 | fov_v = camera.data.angle_y 70 | aspect_ratio = camera.data.sensor_width / camera.data.sensor_height 71 | 72 | # Calculate tangent values 73 | tan_fov_h_half = math.tan(fov_h / 2) 74 | tan_fov_v_half = math.tan(fov_v / 2) 75 | 76 | # Calculate the near plane dimensions 77 | near_plane_width = 2 * tan_fov_h_half * vertex_camera.z 78 | near_plane_height = 2 * tan_fov_v_half * vertex_camera.z 79 | 80 | # Check if the vertex is within the near plane 81 | x_relative = vertex_camera.x / near_plane_width 82 | y_relative = vertex_camera.y / near_plane_height 83 | 84 | return abs(x_relative) <= 0.5 and abs(y_relative) <= 0.5 85 | 86 | 87 | def test_optimal_distance(): 88 | # Clear existing objects 89 | bpy.ops.object.select_all(action="SELECT") 90 | bpy.ops.object.delete() 91 | 92 | # Create the camera rig 93 | dict = create_camera_rig() 94 | camera = dict["camera_object"] 95 | 96 | # Set camera properties 97 | camera.data.sensor_width = 36 98 | camera.data.sensor_height = 24 99 | camera.data.lens_unit = "FOV" 100 | camera.data.angle = math.radians(35) 101 | 102 | # Define different cube dimensions 103 | cube_dimensions = [(1, 10, 1), (10, 1, 1), (1, 1, 10)] 104 | 105 | # Iterate over each set of dimensions 106 | for i, (width, depth, height) in enumerate(cube_dimensions, start=1): 107 | # Create a test object (cube) with specific dimensions 108 | bpy.ops.mesh.primitive_cube_add(size=1, location=(0, 0, 0)) 109 | cube = bpy.context.active_object 110 | cube.scale = ( 111 | width / 2, 112 | depth / 2, 113 | height / 2, 114 | ) # Blender uses half-dimensions for scale 115 | bpy.context.view_layer.update() 116 | 117 | # Apply scale to ensure the transformations are correct 118 | bpy.ops.object.transform_apply(location=False, rotation=False, scale=True) 119 | 120 | # Position the camera for the cube 121 | position_camera_for_object(camera, cube) 122 | 123 | all_vertices_in_frame = True 124 | for vertex in cube.data.vertices: 125 | world_vertex = cube.matrix_world @ vertex.co 126 | if not is_vertex_in_frame(camera, world_vertex): 127 | all_vertices_in_frame = False 128 | print(f"Cube {i}: Vertex {vertex.co} is outside the camera's frame.") 129 | 130 | if all_vertices_in_frame: 131 | print(f"All vertices of Cube {i} are within the camera's frame.") 132 | else: 133 | print(f"Some vertices of Cube {i} are outside the camera's frame.") 134 | 135 | # Remove the cube after rendering to prepare for the next one 136 | bpy.data.objects.remove(cube, do_unlink=True) 137 | 138 | print("Test completed.") 139 | 140 | 141 | if __name__ == "__main__": 142 | test_optimal_distance() 143 | -------------------------------------------------------------------------------- /simian/tests/object_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from ..camera import create_camera_rig 4 | from ..scene import initialize_scene 5 | from ..object import ( 6 | delete_all_empties, 7 | get_hierarchy_bbox, 8 | join_objects_in_hierarchy, 9 | lock_all_objects, 10 | normalize_object_scale, 11 | optimize_meshes_in_hierarchy, 12 | get_meshes_in_hierarchy, 13 | set_pivot_to_bottom, 14 | unlock_objects, 15 | unparent_keep_transform, 16 | ) 17 | import bpy 18 | 19 | 20 | def test_hierarchy_bbox(): 21 | """ 22 | Test the get_hierarchy_bbox function. 23 | """ 24 | # Load an empty scene 25 | initialize_scene() 26 | 27 | # Create two cubes 1 meter apart 28 | bpy.ops.mesh.primitive_cube_add(size=2, location=(0, 0, 0)) 29 | cube1 = bpy.context.active_object 30 | bpy.ops.mesh.primitive_cube_add(size=2, location=(3, 0, 0)) 31 | cube2 = bpy.context.active_object 32 | 33 | # Create an empty and parent cube1 to it 34 | bpy.ops.object.empty_add(location=(0, 0, 0)) 35 | empty1 = bpy.context.active_object 36 | cube1.parent = empty1 37 | 38 | # Create a second empty, parent it to the first empty, and parent cube2 to this second empty 39 | bpy.ops.object.empty_add(location=(0, 0, 0)) 40 | empty2 = bpy.context.active_object 41 | empty2.parent = empty1 42 | cube2.parent = empty2 43 | 44 | # Call the function on the root empty 45 | min_coord, max_coord = get_hierarchy_bbox(empty1) 46 | 47 | # The expected width of the bounding box should be the distance between the centers plus the size of one cube (since both cubes have a center at their geometric center) 48 | expected_width = 5 # 3 meters apart plus 1 meter half-width of each cube 49 | calculated_width = max_coord[0] - min_coord[0] 50 | 51 | # Assert to check if the calculated width matches the expected width 52 | assert ( 53 | abs(calculated_width - expected_width) < 0.001 54 | ), "Bounding box width is incorrect" 55 | print("============ Test Passed: test_hierarchy_bbox ============") 56 | 57 | 58 | def test_remove_small_geometry(): 59 | """ 60 | Test the remove_small_geometry function. 61 | """ 62 | current_dir = os.path.dirname(__file__) 63 | 64 | # Load an empty scene 65 | initialize_scene() 66 | create_camera_rig() 67 | 68 | initial_objects = lock_all_objects() 69 | 70 | # Path to the GLB file 71 | glb_path = os.path.join(current_dir, "../", "../", "examples", "dangling_parts.glb") 72 | 73 | # Load the model 74 | bpy.ops.import_scene.gltf(filepath=glb_path) 75 | 76 | # Assume the last object added is the root object (this may need to be adjusted based on actual structure) 77 | root_obj = bpy.data.objects[0] 78 | 79 | optimize_meshes_in_hierarchy(root_obj) 80 | 81 | join_objects_in_hierarchy(root_obj) 82 | 83 | meshes = get_meshes_in_hierarchy(root_obj) 84 | obj = meshes[0] 85 | 86 | set_pivot_to_bottom(obj) 87 | unparent_keep_transform(obj) 88 | normalize_object_scale(obj) 89 | delete_all_empties() 90 | 91 | # set position to 0, 0, 0 92 | obj.location = (0, 0, 0) 93 | 94 | # go back to object mode 95 | bpy.ops.object.mode_set(mode="OBJECT") 96 | 97 | unlock_objects(initial_objects) 98 | 99 | # Assert that the resulting object has fewer vertices than the initial count 100 | assert len(obj.data.vertices) > 0 101 | assert len(meshes) == 1 102 | print("============ Test Passed: test_remove_small_geometry ============") 103 | 104 | 105 | def test_normalize_object_scale(): 106 | """ 107 | Test the normalize_object_scale function. 108 | """ 109 | # Load an empty scene 110 | initialize_scene() 111 | # Create a cube with known dimensions 112 | bpy.ops.mesh.primitive_cube_add(size=2) # Creates a cube with edge length 2 113 | cube = bpy.context.active_object 114 | 115 | # Calculate expected scale factor to normalize the cube to unit length (scale_factor / max_dimension) 116 | scale_factor = 1.0 117 | expected_dimension = 1.0 # Since we want the largest dimension to be scaled to 1.0 118 | 119 | # Normalize the cube's scale 120 | normalize_object_scale(cube, scale_factor) 121 | 122 | # Force a scene update to ensure that all transforms are applied 123 | bpy.context.view_layer.update() 124 | 125 | # Recalculate actual dimensions after scaling 126 | actual_dimensions = [cube.dimensions[i] for i in range(3)] 127 | max_dimension_after_scaling = max(actual_dimensions) 128 | 129 | # Assert that the maximum dimension is within a small epsilon of expected_dimension 130 | epsilon = 0.001 # Small threshold to account for floating point arithmetic errors 131 | assert ( 132 | abs(max_dimension_after_scaling - expected_dimension) < epsilon 133 | ), f"Expected max dimension to be close to {expected_dimension}, but got {max_dimension_after_scaling}" 134 | print("============ Test Passed: test_normalize_object_scale ============") 135 | 136 | 137 | # Run tests if this file is executed as a script 138 | if __name__ == "__main__": 139 | test_hierarchy_bbox() 140 | test_remove_small_geometry() 141 | test_normalize_object_scale() 142 | print("============ ALL TESTS PASSED ============") 143 | -------------------------------------------------------------------------------- /simian/tests/postprocessing_test.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import MagicMock 2 | 3 | from ..postprocessing import ( 4 | setup_compositor_for_black_and_white, 5 | setup_compositor_for_cel_shading, 6 | setup_compositor_for_depth, 7 | ) 8 | 9 | 10 | def mock_new_node(*args, **kwargs): 11 | """ 12 | Mock for the nodes.new function in Blender. 13 | """ 14 | node_type = kwargs.get("type", "UnknownType") 15 | node = MagicMock() 16 | node.type = node_type 17 | node.inputs = MagicMock() 18 | node.outputs = MagicMock() 19 | return node 20 | 21 | 22 | def test_setup_compositor_for_black_and_white(): 23 | """ 24 | Test the setup_compositor_for_black_and_white function. 25 | """ 26 | context = MagicMock() 27 | scene = MagicMock() 28 | nodes = MagicMock() 29 | links = MagicMock() 30 | 31 | context.scene = scene 32 | scene.node_tree.nodes = nodes 33 | scene.node_tree.links = links 34 | nodes.new.side_effect = mock_new_node # This now handles keyword arguments 35 | 36 | setup_compositor_for_black_and_white(context) 37 | 38 | # Assertions to ensure nodes are created and links are made 39 | nodes.new.assert_any_call(type="CompositorNodeRLayers") 40 | nodes.new.assert_any_call(type="CompositorNodeHueSat") 41 | nodes.new.assert_any_call(type="CompositorNodeComposite") 42 | assert links.new.called, "Links between nodes should be created" 43 | 44 | 45 | def test_setup_compositor_for_cel_shading(): 46 | """ 47 | Test the setup_compositor_for_cel_shading function. 48 | """ 49 | context = MagicMock() 50 | scene = MagicMock() 51 | nodes = MagicMock() 52 | links = MagicMock() 53 | 54 | context.scene = scene 55 | scene.node_tree.nodes = nodes 56 | scene.node_tree.links = links 57 | nodes.new.side_effect = mock_new_node 58 | 59 | setup_compositor_for_cel_shading(context) 60 | 61 | nodes.new.assert_any_call(type="CompositorNodeRLayers") 62 | nodes.new.assert_any_call(type="CompositorNodeNormal") 63 | nodes.new.assert_any_call(type="CompositorNodeValToRGB") 64 | nodes.new.assert_any_call(type="CompositorNodeMixRGB") 65 | nodes.new.assert_any_call(type="CompositorNodeAlphaOver") 66 | nodes.new.assert_any_call(type="CompositorNodeComposite") 67 | assert links.new.called 68 | print( 69 | "============ Test Passed: test_setup_compositor_for_cel_shading ============" 70 | ) 71 | 72 | 73 | def test_setup_compositor_for_depth(): 74 | """ 75 | Test the setup_compositor_for_depth function. 76 | """ 77 | context = MagicMock() 78 | scene = MagicMock() 79 | nodes = MagicMock() 80 | links = MagicMock() 81 | 82 | context.scene = scene 83 | scene.node_tree.nodes = nodes 84 | scene.node_tree.links = links 85 | nodes.new.side_effect = mock_new_node 86 | 87 | setup_compositor_for_depth(context) 88 | 89 | nodes.new.assert_any_call(type="CompositorNodeRLayers") 90 | nodes.new.assert_any_call(type="CompositorNodeNormalize") 91 | nodes.new.assert_any_call(type="CompositorNodeComposite") 92 | assert links.new.called 93 | print("============ Test Passed: test_setup_compositor_for_depth ============") 94 | 95 | 96 | if __name__ == "__main__": 97 | test_setup_compositor_for_black_and_white() 98 | test_setup_compositor_for_cel_shading() 99 | test_setup_compositor_for_depth() 100 | print("============ ALL TESTS PASSED ============") 101 | -------------------------------------------------------------------------------- /simian/tests/transform_test.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | from ..transform import ( 4 | degrees_to_radians, 5 | compute_rotation_matrix, 6 | apply_rotation, 7 | adjust_positions, 8 | determine_relationships, 9 | ) 10 | 11 | 12 | def test_degrees_to_radians(): 13 | epsilon = 1e-9 14 | assert abs(degrees_to_radians(180) - math.pi) < epsilon 15 | 16 | 17 | def test_compute_rotation_matrix(): 18 | epsilon = 1e-9 19 | theta = math.pi / 4 # 45 degrees in radians 20 | expected_matrix = [ 21 | [math.sqrt(2) / 2, -math.sqrt(2) / 2], 22 | [math.sqrt(2) / 2, math.sqrt(2) / 2], 23 | ] 24 | result_matrix = compute_rotation_matrix(theta) 25 | for i in range(2): 26 | for j in range(2): 27 | assert abs(result_matrix[i][j] - expected_matrix[i][j]) < epsilon 28 | 29 | 30 | def test_apply_rotation(): 31 | epsilon = 1e-9 32 | point = [1, 0] 33 | rotation_matrix = compute_rotation_matrix(math.pi / 2) # 90 degrees rotation 34 | expected_point = [0, 1] 35 | result_point = apply_rotation(point, rotation_matrix) 36 | for i in range(2): 37 | assert abs(result_point[i] - expected_point[i]) < epsilon 38 | 39 | 40 | def test_adjust_positions(): 41 | objects = [{"placement": 0}, {"placement": 1}] 42 | camera_yaw = 90 43 | adjusted_objects = adjust_positions(objects, camera_yaw) 44 | for obj in adjusted_objects: 45 | assert "transformed_position" in obj 46 | 47 | 48 | def test_determine_relationships(): 49 | # Define a set of objects with known transformed positions 50 | objects = [ 51 | {"name": "obj1", "transformed_position": [0, 0]}, 52 | {"name": "obj2", "transformed_position": [1, 1]}, 53 | {"name": "obj3", "transformed_position": [-1, -1]}, 54 | ] 55 | 56 | # Define the directional relationship phrases 57 | object_data = { 58 | "relationships": { 59 | "to_the_left": ["to the left of"], 60 | "to_the_right": ["to the right of"], 61 | "in_front_of": ["in front of"], 62 | "behind": ["behind"], 63 | } 64 | } 65 | 66 | # Define a known camera yaw 67 | camera_yaw = 45 68 | 69 | # Determine the relationships between the objects 70 | relationships = determine_relationships(objects, camera_yaw) 71 | 72 | # Expected relationships 73 | expected_relationships = [ 74 | "obj1 is and behind obj2.", 75 | "obj1 is and in front of obj3.", 76 | "obj2 is and in front of obj1.", 77 | "obj2 is and in front of obj3.", 78 | "obj3 is and behind obj1.", 79 | "obj3 is and behind obj2.", 80 | ] 81 | 82 | # Check if the relationships are correctly formed 83 | assert len(relationships) == len( 84 | expected_relationships 85 | ), "The number of relationships is incorrect." 86 | 87 | for relationship in expected_relationships: 88 | assert ( 89 | relationship in relationships 90 | ), f"Expected relationship '{relationship}' not found in results." 91 | 92 | 93 | if __name__ == "__main__": 94 | test_degrees_to_radians() 95 | test_compute_rotation_matrix() 96 | test_apply_rotation() 97 | test_adjust_positions() 98 | test_determine_relationships() 99 | print("============ ALL TESTS PASSED ============") 100 | -------------------------------------------------------------------------------- /simian/tests/worker_test.py: -------------------------------------------------------------------------------- 1 | import json 2 | import sys 3 | from unittest.mock import patch 4 | 5 | from ..worker import run_job 6 | 7 | 8 | @patch("..worker.subprocess.run") 9 | def test_run_job(mock_subprocess_run): 10 | combination_index = 0 11 | combination = {"objects": []} 12 | width = 1920 13 | height = 1080 14 | output_dir = "test_output" 15 | hdri_path = "test_hdri" 16 | upload_dest = "hf" 17 | start_frame = 0 18 | end_frame = 10 19 | 20 | run_job( 21 | combination_index, 22 | combination, 23 | width, 24 | height, 25 | output_dir, 26 | hdri_path, 27 | upload_dest, 28 | start_frame, 29 | end_frame 30 | ) 31 | 32 | combination_str = json.dumps(combination) 33 | combination_str = '"' + combination_str.replace('"', '\\"') + '"' 34 | 35 | command = f"{sys.executable} -m simian.render -- --width {width} --height {height} --combination_index {combination_index}" 36 | command += f" --output_dir {output_dir}" 37 | command += f" --hdri_path {hdri_path}" 38 | command += f" --start_frame {start_frame} --end_frame {end_frame}" 39 | command += f" --combination {combination_str}" 40 | 41 | mock_subprocess_run.assert_called_once_with(["bash", "-c", command], check=False) 42 | 43 | 44 | if __name__ == "__main__": 45 | # test_run_job() 46 | pass 47 | -------------------------------------------------------------------------------- /simian/vendor/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DeepAI-Research/Simverse/73a463c49ae8c6dbc6ff14ea1dcae1ba76b39ecf/simian/vendor/__init__.py -------------------------------------------------------------------------------- /simian/vendor/objaverse.py: -------------------------------------------------------------------------------- 1 | """A package for downloading and processing Objaverse.""" 2 | 3 | import glob 4 | import gzip 5 | import json 6 | import logging 7 | import multiprocessing 8 | import os 9 | import urllib.request 10 | import warnings 11 | from typing import Any, Dict, List, Optional, Tuple 12 | 13 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 14 | logger = logging.getLogger(__name__) 15 | 16 | 17 | BASE_PATH = os.path.join(os.path.expanduser("~"), ".objaverse") 18 | 19 | __version__ = "0.1.7" 20 | _VERSIONED_PATH = os.path.join(BASE_PATH, "hf-objaverse-v1") 21 | 22 | 23 | def load_annotations(uids: Optional[List[str]] = None) -> Dict[str, Any]: 24 | """Load the full metadata of all objects in the dataset. 25 | 26 | Args: 27 | uids: A list of uids with which to load metadata. If None, it loads 28 | the metadata for all uids. 29 | 30 | Returns: 31 | A dictionary mapping the uid to the metadata. 32 | """ 33 | metadata_path = os.path.join(_VERSIONED_PATH, "metadata") 34 | object_paths = _load_object_paths() 35 | dir_ids = ( 36 | set(object_paths[uid].split("/")[1] for uid in uids) 37 | if uids is not None 38 | else [f"{i // 1000:03d}-{i % 1000:03d}" for i in range(160)] 39 | ) 40 | out = {} 41 | for i_id in dir_ids: 42 | json_file = f"{i_id}.json.gz" 43 | local_path = os.path.join(metadata_path, json_file) 44 | if not os.path.exists(local_path): 45 | hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/metadata/{i_id}.json.gz" 46 | # wget the file and put it in local_path 47 | os.makedirs(os.path.dirname(local_path), exist_ok=True) 48 | urllib.request.urlretrieve(hf_url, local_path) 49 | with gzip.open(local_path, "rb") as f: 50 | data = json.load(f) 51 | if uids is not None: 52 | data = {uid: data[uid] for uid in uids if uid in data} 53 | out.update(data) 54 | if uids is not None and len(out) == len(uids): 55 | break 56 | return out 57 | 58 | 59 | def _load_object_paths() -> Dict[str, str]: 60 | """Load the object paths from the dataset. 61 | 62 | The object paths specify the location of where the object is located 63 | in the Hugging Face repo. 64 | 65 | Returns: 66 | A dictionary mapping the uid to the object path. 67 | """ 68 | object_paths_file = "object-paths.json.gz" 69 | local_path = os.path.join(_VERSIONED_PATH, object_paths_file) 70 | if not os.path.exists(local_path): 71 | hf_url = f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/{object_paths_file}" 72 | # wget the file and put it in local_path 73 | os.makedirs(os.path.dirname(local_path), exist_ok=True) 74 | urllib.request.urlretrieve(hf_url, local_path) 75 | with gzip.open(local_path, "rb") as f: 76 | object_paths = json.load(f) 77 | return object_paths 78 | 79 | 80 | def load_uids() -> List[str]: 81 | """Load the uids from the dataset. 82 | 83 | Returns: 84 | A list of uids. 85 | """ 86 | return list(_load_object_paths().keys()) 87 | 88 | 89 | def _download_object( 90 | uid: str, 91 | object_path: str, 92 | total_downloads: float, 93 | start_file_count: int, 94 | ) -> Tuple[str, str]: 95 | """Download the object for the given uid. 96 | 97 | Args: 98 | uid: The uid of the object to load. 99 | object_path: The path to the object in the Hugging Face repo. 100 | 101 | Returns: 102 | The local path of where the object was downloaded. 103 | """ 104 | local_path = os.path.join(_VERSIONED_PATH, object_path) 105 | tmp_local_path = os.path.join(_VERSIONED_PATH, object_path + ".tmp") 106 | hf_url = ( 107 | f"https://huggingface.co/datasets/allenai/objaverse/resolve/main/{object_path}" 108 | ) 109 | # wget the file and put it in local_path 110 | os.makedirs(os.path.dirname(tmp_local_path), exist_ok=True) 111 | urllib.request.urlretrieve(hf_url, tmp_local_path) 112 | 113 | os.rename(tmp_local_path, local_path) 114 | 115 | files = glob.glob(os.path.join(_VERSIONED_PATH, "glbs", "*", "*.glb")) 116 | # logger.info( 117 | # f"Downloaded {len(files) - start_file_count}/{total_downloads} objects", 118 | # ) 119 | 120 | return uid, local_path 121 | 122 | 123 | def load_objects(uids: List[str], download_processes: int = 1) -> Dict[str, str]: 124 | """Return the path to the object files for the given uids. 125 | 126 | If the object is not already downloaded, it will be downloaded. 127 | 128 | Args: 129 | uids: A list of uids. 130 | download_processes: The number of processes to use to download the objects. 131 | 132 | Returns: 133 | A dictionary mapping the object uid to the local path of where the object 134 | downloaded. 135 | """ 136 | object_paths = _load_object_paths() 137 | out = {} 138 | if download_processes == 1: 139 | uids_to_download = [] 140 | for uid in uids: 141 | if uid.endswith(".glb"): 142 | uid = uid[:-4] 143 | if uid not in object_paths: 144 | warnings.warn(f"Could not find object with uid {uid}. Skipping it.") 145 | continue 146 | object_path = object_paths[uid] 147 | local_path = os.path.join(_VERSIONED_PATH, object_path) 148 | if os.path.exists(local_path): 149 | out[uid] = local_path 150 | continue 151 | uids_to_download.append((uid, object_path)) 152 | if len(uids_to_download) == 0: 153 | return out 154 | start_file_count = len( 155 | glob.glob(os.path.join(_VERSIONED_PATH, "glbs", "*", "*.glb")) 156 | ) 157 | for uid, object_path in uids_to_download: 158 | uid, local_path = _download_object( 159 | uid, object_path, len(uids_to_download), start_file_count 160 | ) 161 | out[uid] = local_path 162 | else: 163 | args = [] 164 | for uid in uids: 165 | if uid.endswith(".glb"): 166 | uid = uid[:-4] 167 | if uid not in object_paths: 168 | warnings.warn(f"Could not find object with uid {uid}. Skipping it.") 169 | continue 170 | object_path = object_paths[uid] 171 | local_path = os.path.join(_VERSIONED_PATH, object_path) 172 | if not os.path.exists(local_path): 173 | args.append((uid, object_paths[uid])) 174 | else: 175 | out[uid] = local_path 176 | if len(args) == 0: 177 | return out 178 | # logger.info( 179 | # f"starting download of {len(args)} objects with {download_processes} processes" 180 | # ) 181 | start_file_count = len( 182 | glob.glob(os.path.join(_VERSIONED_PATH, "glbs", "*", "*.glb")) 183 | ) 184 | args_list = [(*arg, len(args), start_file_count) for arg in args] 185 | with multiprocessing.Pool(download_processes) as pool: 186 | r = pool.starmap(_download_object, args_list) 187 | for uid, local_path in r: 188 | out[uid] = local_path 189 | return out 190 | 191 | 192 | def load_lvis_annotations() -> Dict[str, List[str]]: 193 | """Load the LVIS annotations. 194 | 195 | If the annotations are not already downloaded, they will be downloaded. 196 | 197 | Returns: 198 | A dictionary mapping the LVIS category to the list of uids in that category. 199 | """ 200 | hf_url = "https://huggingface.co/datasets/allenai/objaverse/resolve/main/lvis-annotations.json.gz" 201 | local_path = os.path.join(_VERSIONED_PATH, "lvis-annotations.json.gz") 202 | os.makedirs(os.path.dirname(local_path), exist_ok=True) 203 | if not os.path.exists(local_path): 204 | urllib.request.urlretrieve(hf_url, local_path) 205 | with gzip.open(local_path, "rb") as f: 206 | lvis_annotations = json.load(f) 207 | return lvis_annotations 208 | -------------------------------------------------------------------------------- /simian/worker.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import os 4 | import sys 5 | import subprocess 6 | import boto3 7 | import shlex 8 | import time 9 | from typing import Any, Dict 10 | 11 | logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") 12 | logger = logging.getLogger(__name__) 13 | 14 | 15 | def run_job( 16 | combination_indeces: int, 17 | combinations: Dict[str, Any], 18 | width: int, 19 | height: int, 20 | output_dir: str, 21 | hdri_path: str, 22 | upload_dest: str, 23 | start_frame: int = 0, 24 | end_frame: int = 65 25 | ) -> None: 26 | """ 27 | Run a rendering job with the specified combination index and settings. 28 | 29 | Args: 30 | combination_index (int): The index of the combination to render. 31 | combination (Dict[str, Any]): The combination dictionary. 32 | width (int): The width of the rendered output. 33 | height (int): The height of the rendered output. 34 | output_dir (str): The directory to save the rendered output. 35 | hdri_path (str): The path to the HDRI file. 36 | start_frame (int, optional): The starting frame number. Defaults to 0. 37 | end_frame (int, optional): The ending frame number. Defaults to 65. 38 | 39 | Returns: 40 | None 41 | """ 42 | combination_strings = [] 43 | for combo in combinations: 44 | combination_string = json.dumps(combo) 45 | combination_string = shlex.quote(combination_string) 46 | combination_strings.append(combination_string) 47 | 48 | # upload to Hugging Face 49 | if upload_dest == "hf": 50 | 51 | # create output directory, add time to name so each new directory is unique 52 | output_dir += str(time.time()) 53 | os.makedirs(output_dir, exist_ok=True) 54 | 55 | # render images in batches (batches to handle rate limiting of uploads) 56 | batch_size = len(combination_indeces) 57 | for i in range(batch_size): 58 | 59 | args = f" --width {width} --height {height} --combination_index {combination_indeces[i]}" 60 | args += f" --output_dir {output_dir}" 61 | args += f" --hdri_path {hdri_path}" 62 | args += f" --start_frame {start_frame} --end_frame {end_frame}" 63 | args += f" --combination {combination_strings[i]}" 64 | 65 | command = f"{sys.executable} -m simian.render -- {args}" 66 | logger.info(f"Worker running simian.render") 67 | 68 | subprocess.run(["bash", "-c", command], check=True) 69 | 70 | distributask.upload_directory(output_dir) 71 | 72 | # upload to aws s3 bucket 73 | else: 74 | 75 | os.makedirs(output_dir, exist_ok=True) 76 | 77 | combination_index = combination_indeces[0] 78 | combination = combination_strings[0] 79 | 80 | args = f" --width {width} --height {height} --combination_index {combination_index}" 81 | args += f" --output_dir {output_dir}" 82 | args += f" --hdri_path {hdri_path}" 83 | args += f" --start_frame {start_frame} --end_frame {end_frame}" 84 | args += f" --combination {combination}" 85 | 86 | command = f"{sys.executable} -m simian.render -- {args}" 87 | logger.info(f"Worker running simian.render") 88 | 89 | subprocess.run(["bash", "-c", command], check=True) 90 | 91 | 92 | file_location = f"{output_dir}/{combination_index}.mp4" 93 | 94 | file_upload_name = f"{combination_index:05d}.mp4" 95 | 96 | s3_client = boto3.client('s3') 97 | s3_client.upload_file(file_location, os.getenv("S3_BUCKET_NAME"), file_upload_name) 98 | 99 | return "Task completed" 100 | 101 | 102 | # only run this is this file was started by celery or run directly 103 | # check if celery is in sys.argv, it could be sys.argv[0] but might not be 104 | 105 | if __name__ == "__main__" or any("celery" in arg for arg in sys.argv): 106 | from distributask.distributask import create_from_config 107 | 108 | distributask = create_from_config() 109 | distributask.register_function(run_job) 110 | 111 | celery = distributask.app 112 | --------------------------------------------------------------------------------