diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..e82737a31f0052da8608636b2728f0da9befd7ef 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +assets/test2.png filter=lfs diff=lfs merge=lfs -text diff --git a/README.md b/README.md index d93cfff52a23f78d2ad5217a48c9fbeee5b17ffe..0f7e183eff6eef1f672474d1e457b51b9b96adb9 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,13 @@ --- title: WiLoR -emoji: 📊 -colorFrom: purple -colorTo: indigo +emoji: 🚀 +colorFrom: red +colorTo: red sdk: gradio sdk_version: 4.44.0 app_file: app.py pinned: false -license: cc-by-nc-4.0 +license: cc-by-nc-2.0 --- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference diff --git a/app.py b/app.py new file mode 100644 index 0000000000000000000000000000000000000000..006d571fd52535d841cedeee4def9daf1f498018 --- /dev/null +++ b/app.py @@ -0,0 +1,181 @@ +import os +import sys +os.environ["PYOPENGL_PLATFORM"] = "egl" +os.environ["MESA_GL_VERSION_OVERRIDE"] = "4.1" +os.system('pip install /home/user/app/pyrender') +sys.path.append('/home/user/app/pyrender') + +import gradio as gr +import cv2 +import numpy as np +import torch +from ultralytics import YOLO +from pathlib import Path +import argparse +import json +from typing import Dict, Optional + +from wilor.models import WiLoR, load_wilor +from wilor.utils import recursive_to +from wilor.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD +from wilor.utils.renderer import Renderer, cam_crop_to_full +device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') + +LIGHT_PURPLE=(0.25098039, 0.274117647, 0.65882353) + +model, model_cfg = load_wilor(checkpoint_path = './pretrained_models/wilor_final.ckpt' , cfg_path= './pretrained_models/model_config.yaml') +# Setup the renderer +renderer = Renderer(model_cfg, faces=model.mano.faces) +renderer_side = Renderer(model_cfg, faces=model.mano.faces) +model = model.to(device) +model.eval() + +detector = YOLO('./pretrained_models/detector.pt').to(device) + +def run_wilow_model(image, conf, IoU_threshold=0.5): + img_cv2 = image[...,::-1] + img_vis = image.copy() + + detections = detector(img_cv2, conf=conf, verbose=False, iou=IoU_threshold)[0] + + bboxes = [] + is_right = [] + for det in detections: + Bbox = det.boxes.data.cpu().detach().squeeze().numpy() + Conf = det.boxes.conf.data.cpu().detach()[0].numpy().reshape(-1).astype(np.float16) + Side = det.boxes.cls.data.cpu().detach() + #Bbox[:2] -= np.int32(0.1 * Bbox[:2]) + #Bbox[2:] += np.int32(0.1 * Bbox[ 2:]) + is_right.append(det.boxes.cls.cpu().detach().squeeze().item()) + bboxes.append(Bbox[:4].tolist()) + + color = (255*0.208, 255*0.647 ,255*0.603 ) if Side==0. else (255*1, 255*0.78039, 255*0.2353) + label = f'L - {Conf[0]:.3f}' if Side==0 else f'R - {Conf[0]:.3f}' + + cv2.rectangle(img_vis, (int(Bbox[0]), int(Bbox[1])), (int(Bbox[2]), int(Bbox[3])), color , 3) + (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1) + cv2.rectangle(img_vis, (int(Bbox[0]), int(Bbox[1]) - 20), (int(Bbox[0]) + w, int(Bbox[1])), color, -1) + cv2.putText(img_vis, label, (int(Bbox[0]), int(Bbox[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,0,0), 2) + + if len(bboxes) != 0: + boxes = np.stack(bboxes) + right = np.stack(is_right) + dataset = ViTDetDataset(model_cfg, img_cv2, boxes, right, rescale_factor=2.0 ) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False, num_workers=0) + + all_verts = [] + all_cam_t = [] + all_right = [] + all_joints= [] + + for batch in dataloader: + batch = recursive_to(batch, device) + + with torch.no_grad(): + out = model(batch) + + multiplier = (2*batch['right']-1) + pred_cam = out['pred_cam'] + pred_cam[:,1] = multiplier*pred_cam[:,1] + box_center = batch["box_center"].float() + box_size = batch["box_size"].float() + img_size = batch["img_size"].float() + scaled_focal_length = model_cfg.EXTRA.FOCAL_LENGTH / model_cfg.MODEL.IMAGE_SIZE * img_size.max() + pred_cam_t_full = cam_crop_to_full(pred_cam, box_center, box_size, img_size, scaled_focal_length).detach().cpu().numpy() + + # Render the result + all_verts = [] + all_cam_t = [] + all_right = [] + all_joints = [] + + batch_size = batch['img'].shape[0] + for n in range(batch_size): + + verts = out['pred_vertices'][n].detach().cpu().numpy() + joints = out['pred_keypoints_3d'][n].detach().cpu().numpy() + + is_right = batch['right'][n].cpu().numpy() + verts[:,0] = (2*is_right-1)*verts[:,0] + joints[:,0] = (2*is_right-1)*joints[:,0] + + cam_t = pred_cam_t_full[n] + + all_verts.append(verts) + all_cam_t.append(cam_t) + all_right.append(is_right) + all_joints.append(joints) + # Render front view + + misc_args = dict( + mesh_base_color=LIGHT_PURPLE, + scene_bg_color=(1, 1, 1), + focal_length=scaled_focal_length, + ) + cam_view = renderer.render_rgba_multiple(all_verts, cam_t=all_cam_t, render_res=img_size[n], is_right=all_right, **misc_args) + + # Overlay image + + input_img = img_vis.astype(np.float32)/255.0 + input_img = np.concatenate([input_img, np.ones_like(input_img[:,:,:1])], axis=2) # Add alpha channel + input_img_overlay = input_img[:,:,:3] * (1-cam_view[:,:,3:]) + cam_view[:,:,:3] * cam_view[:,:,3:] + + image = input_img_overlay + return image, f'{len(detections)} hands detected' + + + +header = (''' +
+

WiLoR: End-to-end 3D hand localization and reconstruction in-the-wild

+

+ Rolandos Alexandros Potamias1, + Jinglei Zhang2, +
+ Jiankang Deng1, + Stefanos Zafeiriou1 +

+

+ 1Imperial College London; + 2Shanghai Jiao Tong University +

+
+
+ + + + +''') + + +with gr.Blocks(title="WiLoR: End-to-end 3D hand localization and reconstruction in-the-wild", css=".gradio-container") as demo: + + gr.Markdown(header) + + with gr.Row(): + with gr.Column(): + input_image = gr.Image(label="Input image", type="numpy") + threshold = gr.Slider(value=0.3, minimum=0.05, maximum=0.95, step=0.05, label='Detection Confidence Threshold') + #nms = gr.Slider(value=0.5, minimum=0.05, maximum=0.95, step=0.05, label='IoU NMS Threshold') + submit = gr.Button("Submit", variant="primary") + + + with gr.Column(): + reconstruction = gr.Image(label="Reconstructions", type="numpy") + hands_detected = gr.Textbox(label="Hands Detected") + + submit.click(fn=run_wilow_model, inputs=[input_image, threshold], outputs=[reconstruction, hands_detected]) + + with gr.Row(): + + example_images = gr.Examples([ + + ['/home/user/app/assets/test1.jpg'], + ['/home/user/app/assets/test2.png'], + ['/home/user/app/assets/test3.jpg'], + ['/home/user/app/assets/test4.jpg'], + ['/home/user/app/assets/test5.jpeg'] + ], + inputs=input_image) + +demo.launch() \ No newline at end of file diff --git a/assets/test1.jpg b/assets/test1.jpg new file mode 100644 index 0000000000000000000000000000000000000000..3686bcedfef98e9c671df705b324e4301430ab68 Binary files /dev/null and b/assets/test1.jpg differ diff --git a/assets/test2.png b/assets/test2.png new file mode 100644 index 0000000000000000000000000000000000000000..93f1ec0a4261e73987381685db6c4c895dfd7009 --- /dev/null +++ b/assets/test2.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:589f5d12593acbcbcb9ec07b288b04f6d7e70542e1312ceee3ea992ba0f41ff9 +size 1009481 diff --git a/assets/test3.jpg b/assets/test3.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7f291937611596d24c2fea1c3e84a57226c602a9 Binary files /dev/null and b/assets/test3.jpg differ diff --git a/assets/test4.jpg b/assets/test4.jpg new file mode 100644 index 0000000000000000000000000000000000000000..9397cd91883f48fb2a14ef22cbcb1d9d2c94686a Binary files /dev/null and b/assets/test4.jpg differ diff --git a/assets/test5.jpeg b/assets/test5.jpeg new file mode 100644 index 0000000000000000000000000000000000000000..2f75ec47037db4f8d31a3c4e1cca0ab85fdcfdf9 Binary files /dev/null and b/assets/test5.jpeg differ diff --git a/mano_data/mano/MANO_RIGHT.pkl b/mano_data/mano/MANO_RIGHT.pkl new file mode 100644 index 0000000000000000000000000000000000000000..8e7ac7faf64ad51096ec1da626ea13757ed7f665 --- /dev/null +++ b/mano_data/mano/MANO_RIGHT.pkl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:45d60aa3b27ef9107a7afd4e00808f307fd91111e1cfa35afd5c4a62de264767 +size 3821356 diff --git a/mano_data/mano_mean_params.npz b/mano_data/mano_mean_params.npz new file mode 100644 index 0000000000000000000000000000000000000000..dc294b01fb78a9cd6636c87a69b59cf82d28d15b --- /dev/null +++ b/mano_data/mano_mean_params.npz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:efc0ec58e4a5cef78f3abfb4e8f91623b8950be9eff8b8e0dbb0d036ebc63988 +size 1178 diff --git a/packages.txt b/packages.txt new file mode 100644 index 0000000000000000000000000000000000000000..90d6230bcd03a2df7688f62b96d3cf99ab98437d --- /dev/null +++ b/packages.txt @@ -0,0 +1,12 @@ +libglfw3-dev +libgles2-mesa-dev +libgl1 +freeglut3-dev +unzip +ffmpeg +libsm6 +libxext6 +libgl1-mesa-dri +libegl1-mesa +libgbm1 +build-essential \ No newline at end of file diff --git a/pretrained_models/dataset_config.yaml b/pretrained_models/dataset_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5ca6989c8c3bef738b024f762ade513a269fd672 --- /dev/null +++ b/pretrained_models/dataset_config.yaml @@ -0,0 +1,62 @@ +ARCTIC-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/arctic-train/{000000..000176}.tar + epoch_size: 177000 +BEDLAM-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/bedlam-train/{000000..000300}.tar + epoch_size: 301000 +COCOW-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/cocow-train/{000000..000036}.tar + epoch_size: 78666 +DEX-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/dex-train/{000000..000406}.tar + epoch_size: 406888 +FREIHAND-MOCAP: + DATASET_FILE: hamer_training_data/freihand_mocap.npz +FREIHAND-TEST: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/freihand-test/{000000..000003}.tar + epoch_size: 3960 +FREIHAND-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/freihand-train/{000000..000130}.tar + epoch_size: 130240 +H2O3D-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/h2o3d-train/{000000..000060}.tar + epoch_size: 121996 +HALPE-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/halpe-train/{000000..000022}.tar + epoch_size: 34289 +HO3D-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/ho3d-train/{000000..000083}.tar + epoch_size: 83325 +HOT3D-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/hot3d-train/{000000..000571}.tar + epoch_size: 572000 +INTERHAND26M-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/interhand26m-train/{000000..001056}.tar + epoch_size: 1424632 +MPIINZSL-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/mpiinzsl-train/{000000..000015}.tar + epoch_size: 15184 +MTC-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/mtc-train/{000000..000306}.tar + epoch_size: 363947 +REINTER-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/reinter-train/{000000..000418}.tar + epoch_size: 419000 +RHD-TRAIN: + TYPE: ImageDataset + URLS: hamer_training_data/dataset_tars/rhd-train/{000000..000041}.tar + epoch_size: 61705 diff --git a/pretrained_models/detector.pt b/pretrained_models/detector.pt new file mode 100644 index 0000000000000000000000000000000000000000..6c02e16cc7bc4c4fbfa3ca5f8a64e0905e6399ce --- /dev/null +++ b/pretrained_models/detector.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5ef3df44e42d2db52d4ffe91f83a22ce9925e2acc9abebf453f2c5d22e380033 +size 53582271 diff --git a/pretrained_models/model_config.yaml b/pretrained_models/model_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..823f261e6470264d077817c9c3400b0b0787e2a7 --- /dev/null +++ b/pretrained_models/model_config.yaml @@ -0,0 +1,119 @@ +task_name: train +tags: +- dev +train: true +test: false +ckpt_path: null +seed: null +DATASETS: + TRAIN: + FREIHAND-TRAIN: + WEIGHT: 0.2 + INTERHAND26M-TRAIN: + WEIGHT: 0.1 + MTC-TRAIN: + WEIGHT: 0.05 + RHD-TRAIN: + WEIGHT: 0.05 + COCOW-TRAIN: + WEIGHT: 0.05 + HALPE-TRAIN: + WEIGHT: 0.05 + MPIINZSL-TRAIN: + WEIGHT: 0.05 + HO3D-TRAIN: + WEIGHT: 0.05 + H2O3D-TRAIN: + WEIGHT: 0.05 + DEX-TRAIN: + WEIGHT: 0.05 + BEDLAM-TRAIN: + WEIGHT: 0.05 + REINTER-TRAIN: + WEIGHT: 0.1 + HOT3D-TRAIN: + WEIGHT: 0.05 + ARCTIC-TRAIN: + WEIGHT: 0.1 + VAL: + FREIHAND-TRAIN: + WEIGHT: 1.0 + MOCAP: FREIHAND-MOCAP + BETAS_REG: true + CONFIG: + SCALE_FACTOR: 0.3 + ROT_FACTOR: 30 + TRANS_FACTOR: 0.02 + COLOR_SCALE: 0.2 + ROT_AUG_RATE: 0.6 + TRANS_AUG_RATE: 0.5 + DO_FLIP: false + FLIP_AUG_RATE: 0.0 + EXTREME_CROP_AUG_RATE: 0.0 + EXTREME_CROP_AUG_LEVEL: 1 +extras: + ignore_warnings: false + enforce_tags: true + print_config: true +exp_name: WiLoR +MANO: + DATA_DIR: mano_data + MODEL_PATH: ${MANO.DATA_DIR}/mano + GENDER: neutral + NUM_HAND_JOINTS: 15 + MEAN_PARAMS: ${MANO.DATA_DIR}/mano_mean_params.npz + CREATE_BODY_POSE: false +EXTRA: + FOCAL_LENGTH: 5000 + NUM_LOG_IMAGES: 4 + NUM_LOG_SAMPLES_PER_IMAGE: 8 + PELVIS_IND: 0 +GENERAL: + TOTAL_STEPS: 1000000 + LOG_STEPS: 1000 + VAL_STEPS: 1000 + CHECKPOINT_STEPS: 1000 + CHECKPOINT_SAVE_TOP_K: 1 + NUM_WORKERS: 8 + PREFETCH_FACTOR: 2 +TRAIN: + LR: 1.0e-05 + WEIGHT_DECAY: 0.0001 + BATCH_SIZE: 32 + LOSS_REDUCTION: mean + NUM_TRAIN_SAMPLES: 2 + NUM_TEST_SAMPLES: 64 + POSE_2D_NOISE_RATIO: 0.01 + SMPL_PARAM_NOISE_RATIO: 0.005 +MODEL: + IMAGE_SIZE: 256 + IMAGE_MEAN: + - 0.485 + - 0.456 + - 0.406 + IMAGE_STD: + - 0.229 + - 0.224 + - 0.225 + BACKBONE: + TYPE: vit + PRETRAINED_WEIGHTS: hamer_training_data/vitpose_backbone.pth + MANO_HEAD: + TYPE: transformer_decoder + IN_CHANNELS: 2048 + TRANSFORMER_DECODER: + depth: 6 + heads: 8 + mlp_dim: 1024 + dim_head: 64 + dropout: 0.0 + emb_dropout: 0.0 + norm: layer + context_dim: 1280 +LOSS_WEIGHTS: + KEYPOINTS_3D: 0.05 + KEYPOINTS_2D: 0.01 + GLOBAL_ORIENT: 0.001 + HAND_POSE: 0.001 + BETAS: 0.0005 + ADVERSARIAL: 0.0005 diff --git a/pretrained_models/wilor_final.ckpt b/pretrained_models/wilor_final.ckpt new file mode 100644 index 0000000000000000000000000000000000000000..f793ea05d84fe1eb953f36f8ad840b6056a3c46f --- /dev/null +++ b/pretrained_models/wilor_final.ckpt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3e97aafc7dd08d883a4cc5a027df61fdb6fda6136dbd1319405413862ada6bb2 +size 2564989533 diff --git a/pyrender/.coveragerc b/pyrender/.coveragerc new file mode 100644 index 0000000000000000000000000000000000000000..ee31cded3509cbd991a33dd27e2525b93a1a6558 --- /dev/null +++ b/pyrender/.coveragerc @@ -0,0 +1,5 @@ +[report] +exclude_lines = + def __repr__ + def __str__ + @abc.abstractmethod diff --git a/pyrender/.flake8 b/pyrender/.flake8 new file mode 100644 index 0000000000000000000000000000000000000000..fec4bcfc3ba774b53a866d839ea15bae6ebdb4a6 --- /dev/null +++ b/pyrender/.flake8 @@ -0,0 +1,8 @@ +[flake8] +ignore = E231,W504,F405,F403 +max-line-length = 79 +select = B,C,E,F,W,T4,B9 +exclude = + docs/source/conf.py, + __pycache__, + examples/* diff --git a/pyrender/.gitignore b/pyrender/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..ae59dec631f71a23d4255aaf9c0274a699f4ba25 --- /dev/null +++ b/pyrender/.gitignore @@ -0,0 +1,106 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +docs/**/generated/** + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ diff --git a/pyrender/.pre-commit-config.yaml b/pyrender/.pre-commit-config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1817eb39bf409aff80c7d2cc79a3bc3856c70dbd --- /dev/null +++ b/pyrender/.pre-commit-config.yaml @@ -0,0 +1,6 @@ +repos: +- repo: https://gitlab.com/pycqa/flake8 + rev: 3.7.1 + hooks: + - id: flake8 + exclude: ^setup.py diff --git a/pyrender/.travis.yml b/pyrender/.travis.yml new file mode 100644 index 0000000000000000000000000000000000000000..1ad289ae1513eaf8fda74f8d5ab7840be3ef56cb --- /dev/null +++ b/pyrender/.travis.yml @@ -0,0 +1,43 @@ +language: python +sudo: required +dist: xenial + +python: +- '3.6' +- '3.7' + +before_install: + # Pre-install osmesa + - sudo apt update + - sudo wget https://github.com/mmatl/travis_debs/raw/master/xenial/mesa_18.3.3-0.deb + - sudo dpkg -i ./mesa_18.3.3-0.deb || true + - sudo apt install -f + - git clone https://github.com/mmatl/pyopengl.git + - cd pyopengl + - pip install . + - cd .. + +install: + - pip install . + # - pip install -q pytest pytest-cov coveralls + - pip install pytest pytest-cov coveralls + - pip install ./pyopengl + +script: + - PYOPENGL_PLATFORM=osmesa pytest --cov=pyrender tests + +after_success: +- coveralls || true + +deploy: + provider: pypi + skip_existing: true + user: mmatl + on: + tags: true + branch: master + password: + secure: O4WWMbTYb2eVYIO4mMOVa6/xyhX7mPvJpd96cxfNvJdyuqho8VapOhzqsI5kahMB1hFjWWr61yR4+Ru5hoDYf3XA6BQVk8eCY9+0H7qRfvoxex71lahKAqfHLMoE1xNdiVTgl+QN9hYjOnopLod24rx8I8eXfpHu/mfCpuTYGyLlNcDP5St3bXpXLPB5wg8Jo1YRRv6W/7fKoXyuWjewk9cJAS0KrEgnDnSkdwm6Pb+80B2tcbgdGvpGaByw5frndwKiMUMgVUownepDU5POQq2p29wwn9lCvRucULxjEgO+63jdbZRj5fNutLarFa2nISfYnrd72LOyDfbJubwAzzAIsy2JbFORyeHvCgloiuE9oE7a9oOQt/1QHBoIV0seiawMWn55Yp70wQ7HlJs4xSGJWCGa5+9883QRNsvj420atkb3cgO8P+PXwiwTi78Dq7Z/xHqccsU0b8poqBneQoA+pUGgNnF6V7Z8e9RsCcse2gAWSZWuOK3ua+9xCgH7I7MeL3afykr2aJ+yFCoYJMFrUjJeodMX2RbL0q+3FzIPZeGW3WdhTEAL9TSKRcJBSQTskaQlZx/OcpobxS7t3d2S68CCLG9uMTqOTYws55WZ1etalA75sRk9K2MR7ZGjZW3jdtvMViISc/t6Rrjea1GE8ZHGJC6/IeLIWA2c7nc= + distributions: sdist bdist_wheel +notifications: + email: false diff --git a/pyrender/LICENSE b/pyrender/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..4276f7d204e4d85104246df637e0e36adbef14a7 --- /dev/null +++ b/pyrender/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2019 Matthew Matl + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/pyrender/MANIFEST.in b/pyrender/MANIFEST.in new file mode 100644 index 0000000000000000000000000000000000000000..097bcca3b4fccdc39ddd63c10f710ad524898e95 --- /dev/null +++ b/pyrender/MANIFEST.in @@ -0,0 +1,5 @@ +# Include the license +include LICENSE +include README.rst +include pyrender/fonts/* +include pyrender/shaders/* diff --git a/pyrender/README.md b/pyrender/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ae88ed1c5e78f247e38291ed83cf4c81230bf976 --- /dev/null +++ b/pyrender/README.md @@ -0,0 +1,92 @@ +# Pyrender + +[![Build Status](https://travis-ci.org/mmatl/pyrender.svg?branch=master)](https://travis-ci.org/mmatl/pyrender) +[![Documentation Status](https://readthedocs.org/projects/pyrender/badge/?version=latest)](https://pyrender.readthedocs.io/en/latest/?badge=latest) +[![Coverage Status](https://coveralls.io/repos/github/mmatl/pyrender/badge.svg?branch=master)](https://coveralls.io/github/mmatl/pyrender?branch=master) +[![PyPI version](https://badge.fury.io/py/pyrender.svg)](https://badge.fury.io/py/pyrender) +[![Downloads](https://pepy.tech/badge/pyrender)](https://pepy.tech/project/pyrender) + +Pyrender is a pure Python (2.7, 3.4, 3.5, 3.6) library for physically-based +rendering and visualization. +It is designed to meet the [glTF 2.0 specification from Khronos](https://www.khronos.org/gltf/). + +Pyrender is lightweight, easy to install, and simple to use. +It comes packaged with both an intuitive scene viewer and a headache-free +offscreen renderer with support for GPU-accelerated rendering on headless +servers, which makes it perfect for machine learning applications. + +Extensive documentation, including a quickstart guide, is provided [here](https://pyrender.readthedocs.io/en/latest/). + +For a minimal working example of GPU-accelerated offscreen rendering using EGL, +check out the [EGL Google CoLab Notebook](https://colab.research.google.com/drive/1pcndwqeY8vker3bLKQNJKr3B-7-SYenE?usp=sharing). + + +

+ GIF of Viewer + Damaged Helmet +

+ +## Installation +You can install pyrender directly from pip. + +```bash +pip install pyrender +``` + +## Features + +Despite being lightweight, pyrender has lots of features, including: + +* Simple interoperation with the amazing [trimesh](https://github.com/mikedh/trimesh) project, +which enables out-of-the-box support for dozens of mesh types, including OBJ, +STL, DAE, OFF, PLY, and GLB. +* An easy-to-use scene viewer with support for animation, showing face and vertex +normals, toggling lighting conditions, and saving images and GIFs. +* An offscreen rendering module that supports OSMesa and EGL backends. +* Shadow mapping for directional and spot lights. +* Metallic-roughness materials for physically-based rendering, including several +types of texture and normal mapping. +* Transparency. +* Depth and color image generation. + +## Sample Usage + +For sample usage, check out the [quickstart +guide](https://pyrender.readthedocs.io/en/latest/examples/index.html) or one of +the Google CoLab Notebooks: + +* [EGL Google CoLab Notebook](https://colab.research.google.com/drive/1pcndwqeY8vker3bLKQNJKr3B-7-SYenE?usp=sharing) + +## Viewer Keyboard and Mouse Controls + +When using the viewer, the basic controls for moving about the scene are as follows: + +* To rotate the camera about the center of the scene, hold the left mouse button and drag the cursor. +* To rotate the camera about its viewing axis, hold `CTRL` left mouse button and drag the cursor. +* To pan the camera, do one of the following: + * Hold `SHIFT`, then hold the left mouse button and drag the cursor. + * Hold the middle mouse button and drag the cursor. +* To zoom the camera in or out, do one of the following: + * Scroll the mouse wheel. + * Hold the right mouse button and drag the cursor. + +The available keyboard commands are as follows: + +* `a`: Toggles rotational animation mode. +* `c`: Toggles backface culling. +* `f`: Toggles fullscreen mode. +* `h`: Toggles shadow rendering. +* `i`: Toggles axis display mode (no axes, world axis, mesh axes, all axes). +* `l`: Toggles lighting mode (scene lighting, Raymond lighting, or direct lighting). +* `m`: Toggles face normal visualization. +* `n`: Toggles vertex normal visualization. +* `o`: Toggles orthographic camera mode. +* `q`: Quits the viewer. +* `r`: Starts recording a GIF, and pressing again stops recording and opens a file dialog. +* `s`: Opens a file dialog to save the current view as an image. +* `w`: Toggles wireframe mode (scene default, flip wireframes, all wireframe, or all solid). +* `z`: Resets the camera to the default view. + +As a note, displaying shadows significantly slows down rendering, so if you're +experiencing low framerates, just kill shadows or reduce the number of lights in +your scene. diff --git a/pyrender/docs/Makefile b/pyrender/docs/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..b1064a04362a0c4372fae351f99ed3bd9f82ff92 --- /dev/null +++ b/pyrender/docs/Makefile @@ -0,0 +1,23 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line. +SPHINXOPTS = +SPHINXBUILD = sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +clean: + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + rm -rf ./source/generated/* + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/pyrender/docs/make.bat b/pyrender/docs/make.bat new file mode 100644 index 0000000000000000000000000000000000000000..4d9eb83d9f9309029f4b14ff09024658bb0f5563 --- /dev/null +++ b/pyrender/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +if "%1" == "" goto help + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.http://sphinx-doc.org/ + exit /b 1 +) + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% + +:end +popd diff --git a/pyrender/docs/source/api/index.rst b/pyrender/docs/source/api/index.rst new file mode 100644 index 0000000000000000000000000000000000000000..b6e473149d8f132f176e242c93406fdb84ce0b04 --- /dev/null +++ b/pyrender/docs/source/api/index.rst @@ -0,0 +1,59 @@ +Pyrender API Documentation +========================== + +Constants +--------- +.. automodapi:: pyrender.constants + :no-inheritance-diagram: + :no-main-docstr: + :no-heading: + +Cameras +------- +.. automodapi:: pyrender.camera + :no-inheritance-diagram: + :no-main-docstr: + :no-heading: + +Lighting +-------- +.. automodapi:: pyrender.light + :no-inheritance-diagram: + :no-main-docstr: + :no-heading: + +Objects +------- +.. automodapi:: pyrender + :no-inheritance-diagram: + :no-main-docstr: + :no-heading: + :skip: Camera, DirectionalLight, Light, OffscreenRenderer, Node + :skip: OrthographicCamera, PerspectiveCamera, PointLight, RenderFlags + :skip: Renderer, Scene, SpotLight, TextAlign, Viewer, GLTF + +Scenes +------ +.. automodapi:: pyrender + :no-inheritance-diagram: + :no-main-docstr: + :no-heading: + :skip: Camera, DirectionalLight, Light, OffscreenRenderer + :skip: OrthographicCamera, PerspectiveCamera, PointLight, RenderFlags + :skip: Renderer, SpotLight, TextAlign, Viewer, Sampler, Texture, Material + :skip: MetallicRoughnessMaterial, Primitive, Mesh, GLTF + +On-Screen Viewer +---------------- +.. automodapi:: pyrender.viewer + :no-inheritance-diagram: + :no-inherited-members: + :no-main-docstr: + :no-heading: + +Off-Screen Rendering +-------------------- +.. automodapi:: pyrender.offscreen + :no-inheritance-diagram: + :no-main-docstr: + :no-heading: diff --git a/pyrender/docs/source/conf.py b/pyrender/docs/source/conf.py new file mode 100644 index 0000000000000000000000000000000000000000..6bf194c375e7e789b334a838953adfeaf2eb59b6 --- /dev/null +++ b/pyrender/docs/source/conf.py @@ -0,0 +1,352 @@ +# -*- coding: utf-8 -*- +# +# core documentation build configuration file, created by +# sphinx-quickstart on Sun Oct 16 14:33:48 2016. +# +# This file is execfile()d with the current directory set to its +# containing dir. +# +# Note that not all possible configuration values are present in this +# autogenerated file. +# +# All configuration values have a default; values that are commented out +# serve to show the default. + +import sys +import os +from pyrender import __version__ +from sphinx.domains.python import PythonDomain + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +sys.path.insert(0, os.path.abspath('../../')) + +# -- General configuration ------------------------------------------------ + +# If your documentation needs a minimal Sphinx version, state it here. +#needs_sphinx = '1.0' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + 'sphinx.ext.autodoc', + 'sphinx.ext.autosummary', + 'sphinx.ext.coverage', + 'sphinx.ext.githubpages', + 'sphinx.ext.intersphinx', + 'sphinx.ext.napoleon', + 'sphinx.ext.viewcode', + 'sphinx_automodapi.automodapi', + 'sphinx_automodapi.smart_resolver' +] +numpydoc_class_members_toctree = False +automodapi_toctreedirnm = 'generated' +automodsumm_inherited_members = True + +# Add any paths that contain templates here, relative to this directory. +templates_path = ['_templates'] + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# source_suffix = ['.rst', '.md'] +source_suffix = '.rst' + +# The encoding of source files. +#source_encoding = 'utf-8-sig' + +# The master toctree document. +master_doc = 'index' + +# General information about the project. +project = u'pyrender' +copyright = u'2018, Matthew Matl' +author = u'Matthew Matl' + +# The version info for the project you're documenting, acts as replacement for +# |version| and |release|, also used in various other places throughout the +# built documents. +# +# The short X.Y version. +version = __version__ +# The full version, including alpha/beta/rc tags. +release = __version__ + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = None + +# There are two options for replacing |today|: either, you set today to some +# non-false value, then it is used: +#today = '' +# Else, today_fmt is used as the format for a strftime call. +#today_fmt = '%B %d, %Y' + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +exclude_patterns = [] + +# The reST default role (used for this markup: `text`) to use for all +# documents. +#default_role = None + +# If true, '()' will be appended to :func: etc. cross-reference text. +#add_function_parentheses = True + +# If true, the current module name will be prepended to all description +# unit titles (such as .. function::). +#add_module_names = True + +# If true, sectionauthor and moduleauthor directives will be shown in the +# output. They are ignored by default. +#show_authors = False + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = 'sphinx' + +# A list of ignored prefixes for module index sorting. +#modindex_common_prefix = [] + +# If true, keep warnings as "system message" paragraphs in the built documents. +#keep_warnings = False + +# If true, `todo` and `todoList` produce output, else they produce nothing. +todo_include_todos = False + + +# -- Options for HTML output ---------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +import sphinx_rtd_theme +html_theme = 'sphinx_rtd_theme' +html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +#html_theme_options = {} + +# Add any paths that contain custom themes here, relative to this directory. +#html_theme_path = [] + +# The name for this set of Sphinx documents. If None, it defaults to +# " v documentation". +#html_title = None + +# A shorter title for the navigation bar. Default is the same as html_title. +#html_short_title = None + +# The name of an image file (relative to this directory) to place at the top +# of the sidebar. +#html_logo = None + +# The name of an image file (relative to this directory) to use as a favicon of +# the docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 +# pixels large. +#html_favicon = None + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ['_static'] + +# Add any extra paths that contain custom files (such as robots.txt or +# .htaccess) here, relative to this directory. These files are copied +# directly to the root of the documentation. +#html_extra_path = [] + +# If not '', a 'Last updated on:' timestamp is inserted at every page bottom, +# using the given strftime format. +#html_last_updated_fmt = '%b %d, %Y' + +# If true, SmartyPants will be used to convert quotes and dashes to +# typographically correct entities. +#html_use_smartypants = True + +# Custom sidebar templates, maps document names to template names. +#html_sidebars = {} + +# Additional templates that should be rendered to pages, maps page names to +# template names. +#html_additional_pages = {} + +# If false, no module index is generated. +#html_domain_indices = True + +# If false, no index is generated. +#html_use_index = True + +# If true, the index is split into individual pages for each letter. +#html_split_index = False + +# If true, links to the reST sources are added to the pages. +#html_show_sourcelink = True + +# If true, "Created using Sphinx" is shown in the HTML footer. Default is True. +#html_show_sphinx = True + +# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True. +#html_show_copyright = True + +# If true, an OpenSearch description file will be output, and all pages will +# contain a tag referring to it. The value of this option must be the +# base URL from which the finished HTML is served. +#html_use_opensearch = '' + +# This is the file name suffix for HTML files (e.g. ".xhtml"). +#html_file_suffix = None + +# Language to be used for generating the HTML full-text search index. +# Sphinx supports the following languages: +# 'da', 'de', 'en', 'es', 'fi', 'fr', 'hu', 'it', 'ja' +# 'nl', 'no', 'pt', 'ro', 'ru', 'sv', 'tr' +#html_search_language = 'en' + +# A dictionary with options for the search language support, empty by default. +# Now only 'ja' uses this config value +#html_search_options = {'type': 'default'} + +# The name of a javascript file (relative to the configuration directory) that +# implements a search results scorer. If empty, the default will be used. +#html_search_scorer = 'scorer.js' + +# Output file base name for HTML help builder. +htmlhelp_basename = 'coredoc' + +# -- Options for LaTeX output --------------------------------------------- + +latex_elements = { +# The paper size ('letterpaper' or 'a4paper'). +#'papersize': 'letterpaper', + +# The font size ('10pt', '11pt' or '12pt'). +#'pointsize': '10pt', + +# Additional stuff for the LaTeX preamble. +#'preamble': '', + +# Latex figure (float) alignment +#'figure_align': 'htbp', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [ + (master_doc, 'pyrender.tex', u'pyrender Documentation', + u'Matthew Matl', 'manual'), +] + +# The name of an image file (relative to this directory) to place at the top of +# the title page. +#latex_logo = None + +# For "manual" documents, if this is true, then toplevel headings are parts, +# not chapters. +#latex_use_parts = False + +# If true, show page references after internal links. +#latex_show_pagerefs = False + +# If true, show URL addresses after external links. +#latex_show_urls = False + +# Documents to append as an appendix to all manuals. +#latex_appendices = [] + +# If false, no module index is generated. +#latex_domain_indices = True + + +# -- Options for manual page output --------------------------------------- + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [ + (master_doc, 'pyrender', u'pyrender Documentation', + [author], 1) +] + +# If true, show URL addresses after external links. +#man_show_urls = False + + +# -- Options for Texinfo output ------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [ + (master_doc, 'pyrender', u'pyrender Documentation', + author, 'pyrender', 'One line description of project.', + 'Miscellaneous'), +] + +# Documents to append as an appendix to all manuals. +#texinfo_appendices = [] + +# If false, no module index is generated. +#texinfo_domain_indices = True + +# How to display URL addresses: 'footnote', 'no', or 'inline'. +#texinfo_show_urls = 'footnote' + +# If true, do not generate a @detailmenu in the "Top" node's menu. +#texinfo_no_detailmenu = False + +intersphinx_mapping = { + 'python' : ('https://docs.python.org/', None), + 'pyrender' : ('https://pyrender.readthedocs.io/en/latest/', None), +} + +# Autosummary fix +autosummary_generate = True + +# Try to suppress multiple-definition warnings by always taking the shorter +# path when two or more paths have the same base module + +class MyPythonDomain(PythonDomain): + + def find_obj(self, env, modname, classname, name, type, searchmode=0): + """Ensures an object always resolves to the desired module + if defined there.""" + orig_matches = PythonDomain.find_obj( + self, env, modname, classname, name, type, searchmode + ) + + if len(orig_matches) <= 1: + return orig_matches + + # If multiple matches, try to take the shortest if all the modules are + # the same + first_match_name_sp = orig_matches[0][0].split('.') + base_name = first_match_name_sp[0] + min_len = len(first_match_name_sp) + best_match = orig_matches[0] + + for match in orig_matches[1:]: + match_name = match[0] + match_name_sp = match_name.split('.') + match_base = match_name_sp[0] + + # If we have mismatched bases, return them all to trigger warnings + if match_base != base_name: + return orig_matches + + # Otherwise, check and see if it's shorter + if len(match_name_sp) < min_len: + min_len = len(match_name_sp) + best_match = match + + return (best_match,) + + +def setup(sphinx): + """Use MyPythonDomain in place of PythonDomain""" + sphinx.override_domain(MyPythonDomain) + diff --git a/pyrender/docs/source/examples/cameras.rst b/pyrender/docs/source/examples/cameras.rst new file mode 100644 index 0000000000000000000000000000000000000000..39186b75b16584d11fd1606b92291c104e0452bd --- /dev/null +++ b/pyrender/docs/source/examples/cameras.rst @@ -0,0 +1,26 @@ +.. _camera_guide: + +Creating Cameras +================ + +Pyrender supports three camera types -- :class:`.PerspectiveCamera` and +:class:`.IntrinsicsCamera` types, +which render scenes as a human would see them, and +:class:`.OrthographicCamera` types, which preserve distances between points. + +Creating cameras is easy -- just specify their basic attributes: + +>>> pc = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.414) +>>> oc = pyrender.OrthographicCamera(xmag=1.0, ymag=1.0) + +For more information, see the Khronos group's documentation here_: + +.. _here: https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#projection-matrices + +When you add cameras to the scene, make sure that you're using OpenGL camera +coordinates to specify their pose. See the illustration below for details. +Basically, the camera z-axis points away from the scene, the x-axis points +right in image space, and the y-axis points up in image space. + +.. image:: /_static/camera_coords.png + diff --git a/pyrender/docs/source/examples/index.rst b/pyrender/docs/source/examples/index.rst new file mode 100644 index 0000000000000000000000000000000000000000..4be536cd62c1cca112228f4e114e783be77a0ab8 --- /dev/null +++ b/pyrender/docs/source/examples/index.rst @@ -0,0 +1,20 @@ +.. _guide: + +User Guide +========== + +This section contains guides on how to use Pyrender to quickly visualize +your 3D data, including a quickstart guide and more detailed descriptions +of each part of the rendering pipeline. + + +.. toctree:: + :maxdepth: 2 + + quickstart.rst + models.rst + lighting.rst + cameras.rst + scenes.rst + offscreen.rst + viewer.rst diff --git a/pyrender/docs/source/examples/lighting.rst b/pyrender/docs/source/examples/lighting.rst new file mode 100644 index 0000000000000000000000000000000000000000..f89bee7d15027a0f52711622b053b49cc6e1b410 --- /dev/null +++ b/pyrender/docs/source/examples/lighting.rst @@ -0,0 +1,21 @@ +.. _lighting_guide: + +Creating Lights +=============== + +Pyrender supports three types of punctual light: + +- :class:`.PointLight`: Point-based light sources, such as light bulbs. +- :class:`.SpotLight`: A conical light source, like a flashlight. +- :class:`.DirectionalLight`: A general light that does not attenuate with + distance. + +Creating lights is easy -- just specify their basic attributes: + +>>> pl = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=2.0) +>>> sl = pyrender.SpotLight(color=[1.0, 1.0, 1.0], intensity=2.0, +... innerConeAngle=0.05, outerConeAngle=0.5) +>>> dl = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=2.0) + +For more information about how these lighting models are implemented, +see their class documentation. diff --git a/pyrender/docs/source/examples/models.rst b/pyrender/docs/source/examples/models.rst new file mode 100644 index 0000000000000000000000000000000000000000..84e71c4ff41a8d2e0eb2dc48434caedb757ff954 --- /dev/null +++ b/pyrender/docs/source/examples/models.rst @@ -0,0 +1,143 @@ +.. _model_guide: + +Loading and Configuring Models +============================== +The first step to any rendering application is loading your models. +Pyrender implements the GLTF 2.0 specification, which means that all +models are composed of a hierarchy of objects. + +At the top level, we have a :class:`.Mesh`. The :class:`.Mesh` is +basically a wrapper of any number of :class:`.Primitive` types, +which actually represent geometry that can be drawn to the screen. + +Primitives are composed of a variety of parameters, including +vertex positions, vertex normals, color and texture information, +and triangle indices if smooth rendering is desired. +They can implement point clouds, triangular meshes, or lines +depending on how you configure their data and set their +:attr:`.Primitive.mode` parameter. + +Although you can create primitives yourself if you want to, +it's probably easier to just use the utility functions provided +in the :class:`.Mesh` class. + +Creating Triangular Meshes +-------------------------- + +Simple Construction +~~~~~~~~~~~~~~~~~~~ +Pyrender allows you to create a :class:`.Mesh` containing a +triangular mesh model directly from a :class:`~trimesh.base.Trimesh` object +using the :meth:`.Mesh.from_trimesh` static method. + +>>> import trimesh +>>> import pyrender +>>> import numpy as np +>>> tm = trimesh.load('examples/models/fuze.obj') +>>> m = pyrender.Mesh.from_trimesh(tm) +>>> m.primitives +[] + +You can also create a single :class:`.Mesh` from a list of +:class:`~trimesh.base.Trimesh` objects: + +>>> tms = [trimesh.creation.icosahedron(), trimesh.creation.cylinder()] +>>> m = pyrender.Mesh.from_trimesh(tms) +[, + ] + +Vertex Smoothing +~~~~~~~~~~~~~~~~ + +The :meth:`.Mesh.from_trimesh` method has a few additional optional parameters. +If you want to render the mesh without interpolating face normals, which can +be useful for meshes that are supposed to be angular (e.g. a cube), you +can specify ``smooth=False``. + +>>> m = pyrender.Mesh.from_trimesh(tm, smooth=False) + +Per-Face or Per-Vertex Coloration +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +If you have an untextured trimesh, you can color it in with per-face or +per-vertex colors: + +>>> tm.visual.vertex_colors = np.random.uniform(size=tm.vertices.shape) +>>> tm.visual.face_colors = np.random.uniform(size=tm.faces.shape) +>>> m = pyrender.Mesh.from_trimesh(tm) + +Instancing +~~~~~~~~~~ + +If you want to render many copies of the same mesh at different poses, +you can statically create a vast array of them in an efficient manner. +Simply specify the ``poses`` parameter to be a list of ``N`` 4x4 homogenous +transformation matrics that position the meshes relative to their common +base frame: + +>>> tfs = np.tile(np.eye(4), (3,1,1)) +>>> tfs[1,:3,3] = [0.1, 0.0, 0.0] +>>> tfs[2,:3,3] = [0.2, 0.0, 0.0] +>>> tfs +array([[[1. , 0. , 0. , 0. ], + [0. , 1. , 0. , 0. ], + [0. , 0. , 1. , 0. ], + [0. , 0. , 0. , 1. ]], + [[1. , 0. , 0. , 0.1], + [0. , 1. , 0. , 0. ], + [0. , 0. , 1. , 0. ], + [0. , 0. , 0. , 1. ]], + [[1. , 0. , 0. , 0.2], + [0. , 1. , 0. , 0. ], + [0. , 0. , 1. , 0. ], + [0. , 0. , 0. , 1. ]]]) + +>>> m = pyrender.Mesh.from_trimesh(tm, poses=tfs) + +Custom Materials +~~~~~~~~~~~~~~~~ + +You can also specify a custom material for any triangular mesh you create +in the ``material`` parameter of :meth:`.Mesh.from_trimesh`. +The main material supported by Pyrender is the +:class:`.MetallicRoughnessMaterial`. +The metallic-roughness model supports rendering highly-realistic objects across +a wide gamut of materials. + +For more information, see the documentation of the +:class:`.MetallicRoughnessMaterial` constructor or look at the Khronos_ +documentation for more information. + +.. _Khronos: https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#materials + +Creating Point Clouds +--------------------- + +Point Sprites +~~~~~~~~~~~~~ +Pyrender also allows you to create a :class:`.Mesh` containing a +point cloud directly from :class:`numpy.ndarray` instances +using the :meth:`.Mesh.from_points` static method. + +Simply provide a list of points and optional per-point colors and normals. + +>>> pts = tm.vertices.copy() +>>> colors = np.random.uniform(size=pts.shape) +>>> m = pyrender.Mesh.from_points(pts, colors=colors) + +Point clouds created in this way will be rendered as square point sprites. + +.. image:: /_static/points.png + +Point Spheres +~~~~~~~~~~~~~ +If you have a monochromatic point cloud and would like to render it with +spheres, you can render it by instancing a spherical trimesh: + +>>> sm = trimesh.creation.uv_sphere(radius=0.1) +>>> sm.visual.vertex_colors = [1.0, 0.0, 0.0] +>>> tfs = np.tile(np.eye(4), (len(pts), 1, 1)) +>>> tfs[:,:3,3] = pts +>>> m = pyrender.Mesh.from_trimesh(sm, poses=tfs) + +.. image:: /_static/points2.png diff --git a/pyrender/docs/source/examples/offscreen.rst b/pyrender/docs/source/examples/offscreen.rst new file mode 100644 index 0000000000000000000000000000000000000000..291532b6e0c0e512df35a97e3c826cc83015aeca --- /dev/null +++ b/pyrender/docs/source/examples/offscreen.rst @@ -0,0 +1,87 @@ +.. _offscreen_guide: + +Offscreen Rendering +=================== + +.. note:: + If you're using a headless server, you'll need to use either EGL (for + GPU-accelerated rendering) or OSMesa (for CPU-only software rendering). + If you're using OSMesa, be sure that you've installed it properly. See + :ref:`osmesa` for details. + +Choosing a Backend +------------------ + +Once you have a scene set up with its geometry, cameras, and lights, +you can render it using the :class:`.OffscreenRenderer`. Pyrender supports +three backends for offscreen rendering: + +- Pyglet, the same engine that runs the viewer. This requires an active + display manager, so you can't run it on a headless server. This is the + default option. +- OSMesa, a software renderer. +- EGL, which allows for GPU-accelerated rendering without a display manager. + +If you want to use OSMesa or EGL, you need to set the ``PYOPENGL_PLATFORM`` +environment variable before importing pyrender or any other OpenGL library. +You can do this at the command line: + +.. code-block:: bash + + PYOPENGL_PLATFORM=osmesa python render.py + +or at the top of your Python script: + +.. code-block:: bash + + # Top of main python script + import os + os.environ['PYOPENGL_PLATFORM'] = 'egl' + +The handle for EGL is ``egl``, and the handle for OSMesa is ``osmesa``. + +Running the Renderer +-------------------- + +Once you've set your environment variable appropriately, create your scene and +then configure the :class:`.OffscreenRenderer` object with a window width, +a window height, and a size for point-cloud points: + +>>> r = pyrender.OffscreenRenderer(viewport_width=640, +... viewport_height=480, +... point_size=1.0) + +Then, just call the :meth:`.OffscreenRenderer.render` function: + +>>> color, depth = r.render(scene) + +.. image:: /_static/scene.png + +This will return a ``(w,h,3)`` channel floating-point color image and +a ``(w,h)`` floating-point depth image rendered from the scene's main camera. + +You can customize the rendering process by using flag options from +:class:`.RenderFlags` and bitwise or-ing them together. For example, +the following code renders a color image with an alpha channel +and enables shadow mapping for all directional lights: + +>>> flags = RenderFlags.RGBA | RenderFlags.SHADOWS_DIRECTIONAL +>>> color, depth = r.render(scene, flags=flags) + +Once you're done with the offscreen renderer, you need to close it before you +can run a different renderer or open the viewer for the same scene: + +>>> r.delete() + +Google CoLab Examples +--------------------- + +For a minimal working example of offscreen rendering using OSMesa, +see the `OSMesa Google CoLab notebook`_. + +.. _OSMesa Google CoLab notebook: https://colab.research.google.com/drive/1Z71mHIc-Sqval92nK290vAsHZRUkCjUx + +For a minimal working example of offscreen rendering using EGL, +see the `EGL Google CoLab notebook`_. + +.. _EGL Google CoLab notebook: https://colab.research.google.com/drive/1rTLHk0qxh4dn8KNe-mCnN8HAWdd2_BEh diff --git a/pyrender/docs/source/examples/quickstart.rst b/pyrender/docs/source/examples/quickstart.rst new file mode 100644 index 0000000000000000000000000000000000000000..ac556419e5206c2ccd4bc985feb1a8c7347310af --- /dev/null +++ b/pyrender/docs/source/examples/quickstart.rst @@ -0,0 +1,71 @@ +.. _quickstart_guide: + +Quickstart +========== + + +Minimal Example for 3D Viewer +----------------------------- +Here is a minimal example of loading and viewing a triangular mesh model +in pyrender. + +>>> import trimesh +>>> import pyrender +>>> fuze_trimesh = trimesh.load('examples/models/fuze.obj') +>>> mesh = pyrender.Mesh.from_trimesh(fuze_trimesh) +>>> scene = pyrender.Scene() +>>> scene.add(mesh) +>>> pyrender.Viewer(scene, use_raymond_lighting=True) + +.. image:: /_static/fuze.png + + +Minimal Example for Offscreen Rendering +--------------------------------------- +.. note:: + If you're using a headless server, make sure that you followed the guide + for installing OSMesa. See :ref:`osmesa`. + +Here is a minimal example of rendering a mesh model offscreen in pyrender. +The only additional necessities are that you need to add lighting and a camera. + +>>> import numpy as np +>>> import trimesh +>>> import pyrender +>>> import matplotlib.pyplot as plt + +>>> fuze_trimesh = trimesh.load('examples/models/fuze.obj') +>>> mesh = pyrender.Mesh.from_trimesh(fuze_trimesh) +>>> scene = pyrender.Scene() +>>> scene.add(mesh) +>>> camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0) +>>> s = np.sqrt(2)/2 +>>> camera_pose = np.array([ +... [0.0, -s, s, 0.3], +... [1.0, 0.0, 0.0, 0.0], +... [0.0, s, s, 0.35], +... [0.0, 0.0, 0.0, 1.0], +... ]) +>>> scene.add(camera, pose=camera_pose) +>>> light = pyrender.SpotLight(color=np.ones(3), intensity=3.0, +... innerConeAngle=np.pi/16.0, +... outerConeAngle=np.pi/6.0) +>>> scene.add(light, pose=camera_pose) +>>> r = pyrender.OffscreenRenderer(400, 400) +>>> color, depth = r.render(scene) +>>> plt.figure() +>>> plt.subplot(1,2,1) +>>> plt.axis('off') +>>> plt.imshow(color) +>>> plt.subplot(1,2,2) +>>> plt.axis('off') +>>> plt.imshow(depth, cmap=plt.cm.gray_r) +>>> plt.show() + +.. image:: /_static/minexcolor.png + :width: 45% + :align: left +.. image:: /_static/minexdepth.png + :width: 45% + :align: right + diff --git a/pyrender/docs/source/examples/scenes.rst b/pyrender/docs/source/examples/scenes.rst new file mode 100644 index 0000000000000000000000000000000000000000..94c243f8b860b9669ac26105fd2b9906054f4568 --- /dev/null +++ b/pyrender/docs/source/examples/scenes.rst @@ -0,0 +1,78 @@ +.. _scene_guide: + +Creating Scenes +=============== + +Before you render anything, you need to put all of your lights, cameras, +and meshes into a scene. The :class:`.Scene` object keeps track of the relative +poses of these primitives by inserting them into :class:`.Node` objects and +keeping them in a directed acyclic graph. + +Adding Objects +-------------- + +To create a :class:`.Scene`, simply call the constructor. You can optionally +specify an ambient light color and a background color: + +>>> scene = pyrender.Scene(ambient_light=[0.02, 0.02, 0.02], +... bg_color=[1.0, 1.0, 1.0]) + +You can add objects to a scene by first creating a :class:`.Node` object +and adding the object and its pose to the :class:`.Node`. Poses are specified +as 4x4 homogenous transformation matrices that are stored in the node's +:attr:`.Node.matrix` attribute. Note that the :class:`.Node` +constructor requires you to specify whether you're adding a mesh, light, +or camera. + +>>> mesh = pyrender.Mesh.from_trimesh(tm) +>>> light = pyrender.PointLight(color=[1.0, 1.0, 1.0], intensity=2.0) +>>> cam = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.414) +>>> nm = pyrender.Node(mesh=mesh, matrix=np.eye(4)) +>>> nl = pyrender.Node(light=light, matrix=np.eye(4)) +>>> nc = pyrender.Node(camera=cam, matrix=np.eye(4)) +>>> scene.add_node(nm) +>>> scene.add_node(nl) +>>> scene.add_node(nc) + +You can also add objects directly to a scene with the :meth:`.Scene.add` function, +which takes care of creating a :class:`.Node` for you. + +>>> scene.add(mesh, pose=np.eye(4)) +>>> scene.add(light, pose=np.eye(4)) +>>> scene.add(cam, pose=np.eye(4)) + +Nodes can be hierarchical, in which case the node's :attr:`.Node.matrix` +specifies that node's pose relative to its parent frame. You can add nodes to +a scene hierarchically by specifying a parent node in your calls to +:meth:`.Scene.add` or :meth:`.Scene.add_node`: + +>>> scene.add_node(nl, parent_node=nc) +>>> scene.add(cam, parent_node=nm) + +If you add multiple cameras to a scene, you can specify which one to render from +by setting the :attr:`.Scene.main_camera_node` attribute. + +Updating Objects +---------------- + +You can update the poses of existing nodes with the :meth:`.Scene.set_pose` +function. Simply call it with a :class:`.Node` that is already in the scene +and the new pose of that node with respect to its parent as a 4x4 homogenous +transformation matrix: + +>>> scene.set_pose(nl, pose=np.eye(4)) + +If you want to get the local pose of a node, you can just access its +:attr:`.Node.matrix` attribute. However, if you want to the get +the pose of a node *with respect to the world frame*, you can call the +:meth:`.Scene.get_pose` method. + +>>> tf = scene.get_pose(nl) + +Removing Objects +---------------- + +Finally, you can remove a :class:`.Node` and all of its children from the +scene with the :meth:`.Scene.remove_node` function: + +>>> scene.remove_node(nl) diff --git a/pyrender/docs/source/examples/viewer.rst b/pyrender/docs/source/examples/viewer.rst new file mode 100644 index 0000000000000000000000000000000000000000..00a7973b46ec7da33b51b65581af6f25c1b1652f --- /dev/null +++ b/pyrender/docs/source/examples/viewer.rst @@ -0,0 +1,61 @@ +.. _viewer_guide: + +Live Scene Viewer +================= + +Standard Usage +-------------- +In addition to the offscreen renderer, Pyrender comes with a live scene viewer. +In its standard invocation, calling the :class:`.Viewer`'s constructor will +immediately pop a viewing window that you can navigate around in. + +>>> pyrender.Viewer(scene) + +By default, the viewer uses your scene's lighting. If you'd like to start with +some additional lighting that moves around with the camera, you can specify that +with: + +>>> pyrender.Viewer(scene, use_raymond_lighting=True) + +For a full list of the many options that the :class:`.Viewer` supports, check out its +documentation. + +.. image:: /_static/rotation.gif + +Running the Viewer in a Separate Thread +--------------------------------------- +If you'd like to animate your models, you'll want to run the viewer in a +separate thread so that you can update the scene while the viewer is running. +To do this, first pop the viewer in a separate thread by calling its constructor +with the ``run_in_thread`` option set: + +>>> v = pyrender.Viewer(scene, run_in_thread=True) + +Then, you can manipulate the :class:`.Scene` while the viewer is running to +animate things. However, be careful to acquire the viewer's +:attr:`.Viewer.render_lock` before editing the scene to prevent data corruption: + +>>> i = 0 +>>> while True: +... pose = np.eye(4) +... pose[:3,3] = [i, 0, 0] +... v.render_lock.acquire() +... scene.set_pose(mesh_node, pose) +... v.render_lock.release() +... i += 0.01 + +.. image:: /_static/scissors.gif + +You can wait on the viewer to be closed manually: + +>>> while v.is_active: +... pass + +Or you can close it from the main thread forcibly. +Make sure to still loop and block for the viewer to actually exit before using +the scene object again. + +>>> v.close_external() +>>> while v.is_active: +... pass + diff --git a/pyrender/docs/source/index.rst b/pyrender/docs/source/index.rst new file mode 100644 index 0000000000000000000000000000000000000000..baf189ede6bb3435cad5b8795e1937ef1a3c2c56 --- /dev/null +++ b/pyrender/docs/source/index.rst @@ -0,0 +1,41 @@ +.. core documentation master file, created by + sphinx-quickstart on Sun Oct 16 14:33:48 2016. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +Pyrender Documentation +======================== +Pyrender is a pure Python (2.7, 3.4, 3.5, 3.6) library for physically-based +rendering and visualization. +It is designed to meet the glTF 2.0 specification_ from Khronos + +.. _specification: https://www.khronos.org/gltf/ + +Pyrender is lightweight, easy to install, and simple to use. +It comes packaged with both an intuitive scene viewer and a headache-free +offscreen renderer with support for GPU-accelerated rendering on headless +servers, which makes it perfect for machine learning applications. +Check out the :ref:`guide` for a full tutorial, or fork me on +Github_. + +.. _Github: https://github.com/mmatl/pyrender + +.. image:: _static/rotation.gif + +.. image:: _static/damaged_helmet.png + +.. toctree:: + :maxdepth: 2 + + install/index.rst + examples/index.rst + api/index.rst + + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` + diff --git a/pyrender/docs/source/install/index.rst b/pyrender/docs/source/install/index.rst new file mode 100644 index 0000000000000000000000000000000000000000..c785f202d877f8bbaf286c21eddca1925973f75e --- /dev/null +++ b/pyrender/docs/source/install/index.rst @@ -0,0 +1,172 @@ +Installation Guide +================== + +Python Installation +------------------- + +This package is available via ``pip``. + +.. code-block:: bash + + pip install pyrender + +If you're on MacOS, you'll need +to pre-install my fork of ``pyglet``, as the version on PyPI hasn't yet included +my change that enables OpenGL contexts on MacOS. + +.. code-block:: bash + + git clone https://github.com/mmatl/pyglet.git + cd pyglet + pip install . + +.. _osmesa: + +Getting Pyrender Working with OSMesa +------------------------------------ +If you want to render scenes offscreen but don't want to have to +install a display manager or deal with the pains of trying to get +OpenGL to work over SSH, you have two options. + +The first (and preferred) option is using EGL, which enables you to perform +GPU-accelerated rendering on headless servers. +However, you'll need EGL 1.5 to get modern OpenGL contexts. +This comes packaged with NVIDIA's current drivers, but if you are having issues +getting EGL to work with your hardware, you can try using OSMesa, +a software-based offscreen renderer that is included with any Mesa +install. + +If you want to use OSMesa with pyrender, you'll have to perform two additional +installation steps: + +- :ref:`installmesa` +- :ref:`installpyopengl` + +Then, read the offscreen rendering tutorial. See :ref:`offscreen_guide`. + +.. _installmesa: + +Installing OSMesa +***************** + +As a first step, you'll need to rebuild and re-install Mesa with support +for fast offscreen rendering and OpenGL 3+ contexts. +I'd recommend installing from source, but you can also try my ``.deb`` +for Ubuntu 16.04 and up. + +Installing from a Debian Package +******************************** + +If you're running Ubuntu 16.04 or newer, you should be able to install the +required version of Mesa from my ``.deb`` file. + +.. code-block:: bash + + sudo apt update + sudo wget https://github.com/mmatl/travis_debs/raw/master/xenial/mesa_18.3.3-0.deb + sudo dpkg -i ./mesa_18.3.3-0.deb || true + sudo apt install -f + +If this doesn't work, try building from source. + +Building From Source +******************** + +First, install build dependencies via `apt` or your system's package manager. + +.. code-block:: bash + + sudo apt-get install llvm-6.0 freeglut3 freeglut3-dev + +Then, download the current release of Mesa from here_. +Unpack the source and go to the source folder: + +.. _here: https://archive.mesa3d.org/mesa-18.3.3.tar.gz + +.. code-block:: bash + + tar xfv mesa-18.3.3.tar.gz + cd mesa-18.3.3 + +Replace ``PREFIX`` with the path you want to install Mesa at. +If you're not worried about overwriting your default Mesa install, +a good place is at ``/usr/local``. + +Now, configure the installation by running the following command: + +.. code-block:: bash + + ./configure --prefix=PREFIX \ + --enable-opengl --disable-gles1 --disable-gles2 \ + --disable-va --disable-xvmc --disable-vdpau \ + --enable-shared-glapi \ + --disable-texture-float \ + --enable-gallium-llvm --enable-llvm-shared-libs \ + --with-gallium-drivers=swrast,swr \ + --disable-dri --with-dri-drivers= \ + --disable-egl --with-egl-platforms= --disable-gbm \ + --disable-glx \ + --disable-osmesa --enable-gallium-osmesa \ + ac_cv_path_LLVM_CONFIG=llvm-config-6.0 + +Finally, build and install Mesa. + +.. code-block:: bash + + make -j8 + make install + +Finally, if you didn't install Mesa in the system path, +add the following lines to your ``~/.bashrc`` file after +changing ``MESA_HOME`` to your mesa installation path (i.e. what you used as +``PREFIX`` during the configure command). + +.. code-block:: bash + + MESA_HOME=/path/to/your/mesa/installation + export LIBRARY_PATH=$LIBRARY_PATH:$MESA_HOME/lib + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$MESA_HOME/lib + export C_INCLUDE_PATH=$C_INCLUDE_PATH:$MESA_HOME/include/ + export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:$MESA_HOME/include/ + +.. _installpyopengl: + +Installing a Compatible Fork of PyOpenGL +**************************************** + +Next, install and use my fork of ``PyOpenGL``. +This fork enables getting modern OpenGL contexts with OSMesa. +My patch has been included in ``PyOpenGL``, but it has not yet been released +on PyPI. + +.. code-block:: bash + + git clone https://github.com/mmatl/pyopengl.git + pip install ./pyopengl + + +Building Documentation +---------------------- + +The online documentation for ``pyrender`` is automatically built by Read The Docs. +Building ``pyrender``'s documentation locally requires a few extra dependencies -- +specifically, `sphinx`_ and a few plugins. + +.. _sphinx: http://www.sphinx-doc.org/en/master/ + +To install the dependencies required, simply change directories into the `pyrender` source and run + +.. code-block:: bash + + $ pip install .[docs] + +Then, go to the ``docs`` directory and run ``make`` with the appropriate target. +For example, + +.. code-block:: bash + + $ cd docs/ + $ make html + +will generate a set of web pages. Any documentation files +generated in this manner can be found in ``docs/build``. diff --git a/pyrender/examples/duck.py b/pyrender/examples/duck.py new file mode 100644 index 0000000000000000000000000000000000000000..9a94bad5bfb30493f7364f2e52cbb4badbccb2c7 --- /dev/null +++ b/pyrender/examples/duck.py @@ -0,0 +1,13 @@ +from pyrender import Mesh, Scene, Viewer +from io import BytesIO +import numpy as np +import trimesh +import requests + +duck_source = "https://github.com/KhronosGroup/glTF-Sample-Models/raw/master/2.0/Duck/glTF-Binary/Duck.glb" + +duck = trimesh.load(BytesIO(requests.get(duck_source).content), file_type='glb') +duckmesh = Mesh.from_trimesh(list(duck.geometry.values())[0]) +scene = Scene(ambient_light=np.array([1.0, 1.0, 1.0, 1.0])) +scene.add(duckmesh) +Viewer(scene) diff --git a/pyrender/examples/example.py b/pyrender/examples/example.py new file mode 100644 index 0000000000000000000000000000000000000000..599a4850a5899cdeb1a76db1c5cf1c91c263cd41 --- /dev/null +++ b/pyrender/examples/example.py @@ -0,0 +1,157 @@ +"""Examples of using pyrender for viewing and offscreen rendering. +""" +import pyglet +pyglet.options['shadow_window'] = False +import os +import numpy as np +import trimesh + +from pyrender import PerspectiveCamera,\ + DirectionalLight, SpotLight, PointLight,\ + MetallicRoughnessMaterial,\ + Primitive, Mesh, Node, Scene,\ + Viewer, OffscreenRenderer, RenderFlags + +#============================================================================== +# Mesh creation +#============================================================================== + +#------------------------------------------------------------------------------ +# Creating textured meshes from trimeshes +#------------------------------------------------------------------------------ + +# Fuze trimesh +fuze_trimesh = trimesh.load('./models/fuze.obj') +fuze_mesh = Mesh.from_trimesh(fuze_trimesh) + +# Drill trimesh +drill_trimesh = trimesh.load('./models/drill.obj') +drill_mesh = Mesh.from_trimesh(drill_trimesh) +drill_pose = np.eye(4) +drill_pose[0,3] = 0.1 +drill_pose[2,3] = -np.min(drill_trimesh.vertices[:,2]) + +# Wood trimesh +wood_trimesh = trimesh.load('./models/wood.obj') +wood_mesh = Mesh.from_trimesh(wood_trimesh) + +# Water bottle trimesh +bottle_gltf = trimesh.load('./models/WaterBottle.glb') +bottle_trimesh = bottle_gltf.geometry[list(bottle_gltf.geometry.keys())[0]] +bottle_mesh = Mesh.from_trimesh(bottle_trimesh) +bottle_pose = np.array([ + [1.0, 0.0, 0.0, 0.1], + [0.0, 0.0, -1.0, -0.16], + [0.0, 1.0, 0.0, 0.13], + [0.0, 0.0, 0.0, 1.0], +]) + +#------------------------------------------------------------------------------ +# Creating meshes with per-vertex colors +#------------------------------------------------------------------------------ +boxv_trimesh = trimesh.creation.box(extents=0.1*np.ones(3)) +boxv_vertex_colors = np.random.uniform(size=(boxv_trimesh.vertices.shape)) +boxv_trimesh.visual.vertex_colors = boxv_vertex_colors +boxv_mesh = Mesh.from_trimesh(boxv_trimesh, smooth=False) + +#------------------------------------------------------------------------------ +# Creating meshes with per-face colors +#------------------------------------------------------------------------------ +boxf_trimesh = trimesh.creation.box(extents=0.1*np.ones(3)) +boxf_face_colors = np.random.uniform(size=boxf_trimesh.faces.shape) +boxf_trimesh.visual.face_colors = boxf_face_colors +boxf_mesh = Mesh.from_trimesh(boxf_trimesh, smooth=False) + +#------------------------------------------------------------------------------ +# Creating meshes from point clouds +#------------------------------------------------------------------------------ +points = trimesh.creation.icosphere(radius=0.05).vertices +point_colors = np.random.uniform(size=points.shape) +points_mesh = Mesh.from_points(points, colors=point_colors) + +#============================================================================== +# Light creation +#============================================================================== + +direc_l = DirectionalLight(color=np.ones(3), intensity=1.0) +spot_l = SpotLight(color=np.ones(3), intensity=10.0, + innerConeAngle=np.pi/16, outerConeAngle=np.pi/6) +point_l = PointLight(color=np.ones(3), intensity=10.0) + +#============================================================================== +# Camera creation +#============================================================================== + +cam = PerspectiveCamera(yfov=(np.pi / 3.0)) +cam_pose = np.array([ + [0.0, -np.sqrt(2)/2, np.sqrt(2)/2, 0.5], + [1.0, 0.0, 0.0, 0.0], + [0.0, np.sqrt(2)/2, np.sqrt(2)/2, 0.4], + [0.0, 0.0, 0.0, 1.0] +]) + +#============================================================================== +# Scene creation +#============================================================================== + +scene = Scene(ambient_light=np.array([0.02, 0.02, 0.02, 1.0])) + +#============================================================================== +# Adding objects to the scene +#============================================================================== + +#------------------------------------------------------------------------------ +# By manually creating nodes +#------------------------------------------------------------------------------ +fuze_node = Node(mesh=fuze_mesh, translation=np.array([0.1, 0.15, -np.min(fuze_trimesh.vertices[:,2])])) +scene.add_node(fuze_node) +boxv_node = Node(mesh=boxv_mesh, translation=np.array([-0.1, 0.10, 0.05])) +scene.add_node(boxv_node) +boxf_node = Node(mesh=boxf_mesh, translation=np.array([-0.1, -0.10, 0.05])) +scene.add_node(boxf_node) + +#------------------------------------------------------------------------------ +# By using the add() utility function +#------------------------------------------------------------------------------ +drill_node = scene.add(drill_mesh, pose=drill_pose) +bottle_node = scene.add(bottle_mesh, pose=bottle_pose) +wood_node = scene.add(wood_mesh) +direc_l_node = scene.add(direc_l, pose=cam_pose) +spot_l_node = scene.add(spot_l, pose=cam_pose) + +#============================================================================== +# Using the viewer with a default camera +#============================================================================== + +v = Viewer(scene, shadows=True) + +#============================================================================== +# Using the viewer with a pre-specified camera +#============================================================================== +cam_node = scene.add(cam, pose=cam_pose) +v = Viewer(scene, central_node=drill_node) + +#============================================================================== +# Rendering offscreen from that camera +#============================================================================== + +r = OffscreenRenderer(viewport_width=640*2, viewport_height=480*2) +color, depth = r.render(scene) + +import matplotlib.pyplot as plt +plt.figure() +plt.imshow(color) +plt.show() + +#============================================================================== +# Segmask rendering +#============================================================================== + +nm = {node: 20*(i + 1) for i, node in enumerate(scene.mesh_nodes)} +seg = r.render(scene, RenderFlags.SEG, nm)[0] +plt.figure() +plt.imshow(seg) +plt.show() + +r.delete() + diff --git a/pyrender/pyrender/__init__.py b/pyrender/pyrender/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ee3709846823b7c4b71b22da0e24d63d805528a8 --- /dev/null +++ b/pyrender/pyrender/__init__.py @@ -0,0 +1,24 @@ +from .camera import (Camera, PerspectiveCamera, OrthographicCamera, + IntrinsicsCamera) +from .light import Light, PointLight, DirectionalLight, SpotLight +from .sampler import Sampler +from .texture import Texture +from .material import Material, MetallicRoughnessMaterial +from .primitive import Primitive +from .mesh import Mesh +from .node import Node +from .scene import Scene +from .renderer import Renderer +from .viewer import Viewer +from .offscreen import OffscreenRenderer +from .version import __version__ +from .constants import RenderFlags, TextAlign, GLTF + +__all__ = [ + 'Camera', 'PerspectiveCamera', 'OrthographicCamera', 'IntrinsicsCamera', + 'Light', 'PointLight', 'DirectionalLight', 'SpotLight', + 'Sampler', 'Texture', 'Material', 'MetallicRoughnessMaterial', + 'Primitive', 'Mesh', 'Node', 'Scene', 'Renderer', 'Viewer', + 'OffscreenRenderer', '__version__', 'RenderFlags', 'TextAlign', + 'GLTF' +] diff --git a/pyrender/pyrender/camera.py b/pyrender/pyrender/camera.py new file mode 100644 index 0000000000000000000000000000000000000000..e019358039033c3a372c990ebad3151258c3651d --- /dev/null +++ b/pyrender/pyrender/camera.py @@ -0,0 +1,437 @@ +"""Virtual cameras compliant with the glTF 2.0 specification as described at +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-camera + +Author: Matthew Matl +""" +import abc +import numpy as np +import six +import sys + +from .constants import DEFAULT_Z_NEAR, DEFAULT_Z_FAR + + +@six.add_metaclass(abc.ABCMeta) +class Camera(object): + """Abstract base class for all cameras. + + Note + ---- + Camera poses are specified in the OpenGL format, + where the z axis points away from the view direction and the + x and y axes point to the right and up in the image plane, respectively. + + Parameters + ---------- + znear : float + The floating-point distance to the near clipping plane. + zfar : float + The floating-point distance to the far clipping plane. + ``zfar`` must be greater than ``znear``. + name : str, optional + The user-defined name of this object. + """ + + def __init__(self, + znear=DEFAULT_Z_NEAR, + zfar=DEFAULT_Z_FAR, + name=None): + self.name = name + self.znear = znear + self.zfar = zfar + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def znear(self): + """float : The distance to the near clipping plane. + """ + return self._znear + + @znear.setter + def znear(self, value): + value = float(value) + if value < 0: + raise ValueError('z-near must be >= 0.0') + self._znear = value + + @property + def zfar(self): + """float : The distance to the far clipping plane. + """ + return self._zfar + + @zfar.setter + def zfar(self, value): + value = float(value) + if value <= 0 or value <= self.znear: + raise ValueError('zfar must be >0 and >znear') + self._zfar = value + + @abc.abstractmethod + def get_projection_matrix(self, width=None, height=None): + """Return the OpenGL projection matrix for this camera. + + Parameters + ---------- + width : int + Width of the current viewport, in pixels. + height : int + Height of the current viewport, in pixels. + """ + pass + + +class PerspectiveCamera(Camera): + + """A perspective camera for perspective projection. + + Parameters + ---------- + yfov : float + The floating-point vertical field of view in radians. + znear : float + The floating-point distance to the near clipping plane. + If not specified, defaults to 0.05. + zfar : float, optional + The floating-point distance to the far clipping plane. + ``zfar`` must be greater than ``znear``. + If None, the camera uses an infinite projection matrix. + aspectRatio : float, optional + The floating-point aspect ratio of the field of view. + If not specified, the camera uses the viewport's aspect ratio. + name : str, optional + The user-defined name of this object. + """ + + def __init__(self, + yfov, + znear=DEFAULT_Z_NEAR, + zfar=None, + aspectRatio=None, + name=None): + super(PerspectiveCamera, self).__init__( + znear=znear, + zfar=zfar, + name=name, + ) + + self.yfov = yfov + self.aspectRatio = aspectRatio + + @property + def yfov(self): + """float : The vertical field of view in radians. + """ + return self._yfov + + @yfov.setter + def yfov(self, value): + value = float(value) + if value <= 0.0: + raise ValueError('Field of view must be positive') + self._yfov = value + + @property + def zfar(self): + """float : The distance to the far clipping plane. + """ + return self._zfar + + @zfar.setter + def zfar(self, value): + if value is not None: + value = float(value) + if value <= 0 or value <= self.znear: + raise ValueError('zfar must be >0 and >znear') + self._zfar = value + + @property + def aspectRatio(self): + """float : The ratio of the width to the height of the field of view. + """ + return self._aspectRatio + + @aspectRatio.setter + def aspectRatio(self, value): + if value is not None: + value = float(value) + if value <= 0.0: + raise ValueError('Aspect ratio must be positive') + self._aspectRatio = value + + def get_projection_matrix(self, width=None, height=None): + """Return the OpenGL projection matrix for this camera. + + Parameters + ---------- + width : int + Width of the current viewport, in pixels. + height : int + Height of the current viewport, in pixels. + """ + aspect_ratio = self.aspectRatio + if aspect_ratio is None: + if width is None or height is None: + raise ValueError('Aspect ratio of camera must be defined') + aspect_ratio = float(width) / float(height) + + a = aspect_ratio + t = np.tan(self.yfov / 2.0) + n = self.znear + f = self.zfar + + P = np.zeros((4,4)) + P[0][0] = 1.0 / (a * t) + P[1][1] = 1.0 / t + P[3][2] = -1.0 + + if f is None: + P[2][2] = -1.0 + P[2][3] = -2.0 * n + else: + P[2][2] = (f + n) / (n - f) + P[2][3] = (2 * f * n) / (n - f) + + return P + + +class OrthographicCamera(Camera): + """An orthographic camera for orthographic projection. + + Parameters + ---------- + xmag : float + The floating-point horizontal magnification of the view. + ymag : float + The floating-point vertical magnification of the view. + znear : float + The floating-point distance to the near clipping plane. + If not specified, defaults to 0.05. + zfar : float + The floating-point distance to the far clipping plane. + ``zfar`` must be greater than ``znear``. + If not specified, defaults to 100.0. + name : str, optional + The user-defined name of this object. + """ + + def __init__(self, + xmag, + ymag, + znear=DEFAULT_Z_NEAR, + zfar=DEFAULT_Z_FAR, + name=None): + super(OrthographicCamera, self).__init__( + znear=znear, + zfar=zfar, + name=name, + ) + + self.xmag = xmag + self.ymag = ymag + + @property + def xmag(self): + """float : The horizontal magnification of the view. + """ + return self._xmag + + @xmag.setter + def xmag(self, value): + value = float(value) + if value <= 0.0: + raise ValueError('X magnification must be positive') + self._xmag = value + + @property + def ymag(self): + """float : The vertical magnification of the view. + """ + return self._ymag + + @ymag.setter + def ymag(self, value): + value = float(value) + if value <= 0.0: + raise ValueError('Y magnification must be positive') + self._ymag = value + + @property + def znear(self): + """float : The distance to the near clipping plane. + """ + return self._znear + + @znear.setter + def znear(self, value): + value = float(value) + if value <= 0: + raise ValueError('z-near must be > 0.0') + self._znear = value + + def get_projection_matrix(self, width=None, height=None): + """Return the OpenGL projection matrix for this camera. + + Parameters + ---------- + width : int + Width of the current viewport, in pixels. + Unused in this function. + height : int + Height of the current viewport, in pixels. + Unused in this function. + """ + xmag = self.xmag + ymag = self.ymag + + # If screen width/height defined, rescale xmag + if width is not None and height is not None: + xmag = width / height * ymag + + n = self.znear + f = self.zfar + P = np.zeros((4,4)) + P[0][0] = 1.0 / xmag + P[1][1] = 1.0 / ymag + P[2][2] = 2.0 / (n - f) + P[2][3] = (f + n) / (n - f) + P[3][3] = 1.0 + return P + + +class IntrinsicsCamera(Camera): + """A perspective camera with custom intrinsics. + + Parameters + ---------- + fx : float + X-axis focal length in pixels. + fy : float + Y-axis focal length in pixels. + cx : float + X-axis optical center in pixels. + cy : float + Y-axis optical center in pixels. + znear : float + The floating-point distance to the near clipping plane. + If not specified, defaults to 0.05. + zfar : float + The floating-point distance to the far clipping plane. + ``zfar`` must be greater than ``znear``. + If not specified, defaults to 100.0. + name : str, optional + The user-defined name of this object. + """ + + def __init__(self, + fx, + fy, + cx, + cy, + znear=DEFAULT_Z_NEAR, + zfar=DEFAULT_Z_FAR, + name=None): + super(IntrinsicsCamera, self).__init__( + znear=znear, + zfar=zfar, + name=name, + ) + + self.fx = fx + self.fy = fy + self.cx = cx + self.cy = cy + + @property + def fx(self): + """float : X-axis focal length in meters. + """ + return self._fx + + @fx.setter + def fx(self, value): + self._fx = float(value) + + @property + def fy(self): + """float : Y-axis focal length in meters. + """ + return self._fy + + @fy.setter + def fy(self, value): + self._fy = float(value) + + @property + def cx(self): + """float : X-axis optical center in pixels. + """ + return self._cx + + @cx.setter + def cx(self, value): + self._cx = float(value) + + @property + def cy(self): + """float : Y-axis optical center in pixels. + """ + return self._cy + + @cy.setter + def cy(self, value): + self._cy = float(value) + + def get_projection_matrix(self, width, height): + """Return the OpenGL projection matrix for this camera. + + Parameters + ---------- + width : int + Width of the current viewport, in pixels. + height : int + Height of the current viewport, in pixels. + """ + width = float(width) + height = float(height) + + cx, cy = self.cx, self.cy + fx, fy = self.fx, self.fy + if sys.platform == 'darwin': + cx = self.cx * 2.0 + cy = self.cy * 2.0 + fx = self.fx * 2.0 + fy = self.fy * 2.0 + + P = np.zeros((4,4)) + P[0][0] = 2.0 * fx / width + P[1][1] = 2.0 * fy / height + P[0][2] = 1.0 - 2.0 * cx / width + P[1][2] = 2.0 * cy / height - 1.0 + P[3][2] = -1.0 + + n = self.znear + f = self.zfar + if f is None: + P[2][2] = -1.0 + P[2][3] = -2.0 * n + else: + P[2][2] = (f + n) / (n - f) + P[2][3] = (2 * f * n) / (n - f) + + return P + + +__all__ = ['Camera', 'PerspectiveCamera', 'OrthographicCamera', + 'IntrinsicsCamera'] diff --git a/pyrender/pyrender/constants.py b/pyrender/pyrender/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..8a5785b6fdb21910a174252c5af2f05b40ece4a5 --- /dev/null +++ b/pyrender/pyrender/constants.py @@ -0,0 +1,149 @@ +DEFAULT_Z_NEAR = 0.05 # Near clipping plane, in meters +DEFAULT_Z_FAR = 100.0 # Far clipping plane, in meters +DEFAULT_SCENE_SCALE = 2.0 # Default scene scale +MAX_N_LIGHTS = 4 # Maximum number of lights of each type allowed +TARGET_OPEN_GL_MAJOR = 4 # Target OpenGL Major Version +TARGET_OPEN_GL_MINOR = 1 # Target OpenGL Minor Version +MIN_OPEN_GL_MAJOR = 3 # Minimum OpenGL Major Version +MIN_OPEN_GL_MINOR = 3 # Minimum OpenGL Minor Version +FLOAT_SZ = 4 # Byte size of GL float32 +UINT_SZ = 4 # Byte size of GL uint32 +SHADOW_TEX_SZ = 2048 # Width and Height of Shadow Textures +TEXT_PADDING = 20 # Width of padding for rendering text (px) + + +# Flags for render type +class RenderFlags(object): + """Flags for rendering in the scene. + + Combine them with the bitwise or. For example, + + >>> flags = OFFSCREEN | SHADOWS_DIRECTIONAL | VERTEX_NORMALS + + would result in an offscreen render with directional shadows and + vertex normals enabled. + """ + NONE = 0 + """Normal PBR Render.""" + DEPTH_ONLY = 1 + """Only render the depth buffer.""" + OFFSCREEN = 2 + """Render offscreen and return the depth and (optionally) color buffers.""" + FLIP_WIREFRAME = 4 + """Invert the status of wireframe rendering for each mesh.""" + ALL_WIREFRAME = 8 + """Render all meshes as wireframes.""" + ALL_SOLID = 16 + """Render all meshes as solids.""" + SHADOWS_DIRECTIONAL = 32 + """Render shadows for directional lights.""" + SHADOWS_POINT = 64 + """Render shadows for point lights.""" + SHADOWS_SPOT = 128 + """Render shadows for spot lights.""" + SHADOWS_ALL = 32 | 64 | 128 + """Render shadows for all lights.""" + VERTEX_NORMALS = 256 + """Render vertex normals.""" + FACE_NORMALS = 512 + """Render face normals.""" + SKIP_CULL_FACES = 1024 + """Do not cull back faces.""" + RGBA = 2048 + """Render the color buffer with the alpha channel enabled.""" + FLAT = 4096 + """Render the color buffer flat, with no lighting computations.""" + SEG = 8192 + + +class TextAlign: + """Text alignment options for captions. + + Only use one at a time. + """ + CENTER = 0 + """Center the text by width and height.""" + CENTER_LEFT = 1 + """Center the text by height and left-align it.""" + CENTER_RIGHT = 2 + """Center the text by height and right-align it.""" + BOTTOM_LEFT = 3 + """Put the text in the bottom-left corner.""" + BOTTOM_RIGHT = 4 + """Put the text in the bottom-right corner.""" + BOTTOM_CENTER = 5 + """Center the text by width and fix it to the bottom.""" + TOP_LEFT = 6 + """Put the text in the top-left corner.""" + TOP_RIGHT = 7 + """Put the text in the top-right corner.""" + TOP_CENTER = 8 + """Center the text by width and fix it to the top.""" + + +class GLTF(object): + """Options for GL objects.""" + NEAREST = 9728 + """Nearest neighbor interpolation.""" + LINEAR = 9729 + """Linear interpolation.""" + NEAREST_MIPMAP_NEAREST = 9984 + """Nearest mipmapping.""" + LINEAR_MIPMAP_NEAREST = 9985 + """Linear mipmapping.""" + NEAREST_MIPMAP_LINEAR = 9986 + """Nearest mipmapping.""" + LINEAR_MIPMAP_LINEAR = 9987 + """Linear mipmapping.""" + CLAMP_TO_EDGE = 33071 + """Clamp to the edge of the texture.""" + MIRRORED_REPEAT = 33648 + """Mirror the texture.""" + REPEAT = 10497 + """Repeat the texture.""" + POINTS = 0 + """Render as points.""" + LINES = 1 + """Render as lines.""" + LINE_LOOP = 2 + """Render as a line loop.""" + LINE_STRIP = 3 + """Render as a line strip.""" + TRIANGLES = 4 + """Render as triangles.""" + TRIANGLE_STRIP = 5 + """Render as a triangle strip.""" + TRIANGLE_FAN = 6 + """Render as a triangle fan.""" + + +class BufFlags(object): + POSITION = 0 + NORMAL = 1 + TANGENT = 2 + TEXCOORD_0 = 4 + TEXCOORD_1 = 8 + COLOR_0 = 16 + JOINTS_0 = 32 + WEIGHTS_0 = 64 + + +class TexFlags(object): + NONE = 0 + NORMAL = 1 + OCCLUSION = 2 + EMISSIVE = 4 + BASE_COLOR = 8 + METALLIC_ROUGHNESS = 16 + DIFFUSE = 32 + SPECULAR_GLOSSINESS = 64 + + +class ProgramFlags: + NONE = 0 + USE_MATERIAL = 1 + VERTEX_NORMALS = 2 + FACE_NORMALS = 4 + + +__all__ = ['RenderFlags', 'TextAlign', 'GLTF'] diff --git a/pyrender/pyrender/font.py b/pyrender/pyrender/font.py new file mode 100644 index 0000000000000000000000000000000000000000..5ac530d7b949f50314a0d9cf5d744bedcace0571 --- /dev/null +++ b/pyrender/pyrender/font.py @@ -0,0 +1,272 @@ +"""Font texture loader and processor. + +Author: Matthew Matl +""" +import freetype +import numpy as np +import os + +import OpenGL +from OpenGL.GL import * + +from .constants import TextAlign, FLOAT_SZ +from .texture import Texture +from .sampler import Sampler + + +class FontCache(object): + """A cache for fonts. + """ + + def __init__(self, font_dir=None): + self._font_cache = {} + self.font_dir = font_dir + if self.font_dir is None: + base_dir, _ = os.path.split(os.path.realpath(__file__)) + self.font_dir = os.path.join(base_dir, 'fonts') + + def get_font(self, font_name, font_pt): + # If it's a file, load it directly, else, try to load from font dir. + if os.path.isfile(font_name): + font_filename = font_name + _, font_name = os.path.split(font_name) + font_name, _ = os.path.split(font_name) + else: + font_filename = os.path.join(self.font_dir, font_name) + '.ttf' + + cid = OpenGL.contextdata.getContext() + key = (cid, font_name, int(font_pt)) + + if key not in self._font_cache: + self._font_cache[key] = Font(font_filename, font_pt) + return self._font_cache[key] + + def clear(self): + for key in self._font_cache: + self._font_cache[key].delete() + self._font_cache = {} + + +class Character(object): + """A single character, with its texture and attributes. + """ + + def __init__(self, texture, size, bearing, advance): + self.texture = texture + self.size = size + self.bearing = bearing + self.advance = advance + + +class Font(object): + """A font object. + + Parameters + ---------- + font_file : str + The file to load the font from. + font_pt : int + The height of the font in pixels. + """ + + def __init__(self, font_file, font_pt=40): + self.font_file = font_file + self.font_pt = int(font_pt) + self._face = freetype.Face(font_file) + self._face.set_pixel_sizes(0, font_pt) + self._character_map = {} + + for i in range(0, 128): + + # Generate texture + face = self._face + face.load_char(chr(i)) + buf = face.glyph.bitmap.buffer + src = (np.array(buf) / 255.0).astype(np.float32) + src = src.reshape((face.glyph.bitmap.rows, + face.glyph.bitmap.width)) + tex = Texture( + sampler=Sampler( + magFilter=GL_LINEAR, + minFilter=GL_LINEAR, + wrapS=GL_CLAMP_TO_EDGE, + wrapT=GL_CLAMP_TO_EDGE + ), + source=src, + source_channels='R', + ) + character = Character( + texture=tex, + size=np.array([face.glyph.bitmap.width, + face.glyph.bitmap.rows]), + bearing=np.array([face.glyph.bitmap_left, + face.glyph.bitmap_top]), + advance=face.glyph.advance.x + ) + self._character_map[chr(i)] = character + + self._vbo = None + self._vao = None + + @property + def font_file(self): + """str : The file the font was loaded from. + """ + return self._font_file + + @font_file.setter + def font_file(self, value): + self._font_file = value + + @property + def font_pt(self): + """int : The height of the font in pixels. + """ + return self._font_pt + + @font_pt.setter + def font_pt(self, value): + self._font_pt = int(value) + + def _add_to_context(self): + + self._vao = glGenVertexArrays(1) + glBindVertexArray(self._vao) + self._vbo = glGenBuffers(1) + glBindBuffer(GL_ARRAY_BUFFER, self._vbo) + glBufferData(GL_ARRAY_BUFFER, FLOAT_SZ * 6 * 4, None, GL_DYNAMIC_DRAW) + glEnableVertexAttribArray(0) + glVertexAttribPointer( + 0, 4, GL_FLOAT, GL_FALSE, 4 * FLOAT_SZ, ctypes.c_void_p(0) + ) + glBindVertexArray(0) + + glPixelStorei(GL_UNPACK_ALIGNMENT, 1) + for c in self._character_map: + ch = self._character_map[c] + if not ch.texture._in_context(): + ch.texture._add_to_context() + + def _remove_from_context(self): + for c in self._character_map: + ch = self._character_map[c] + ch.texture.delete() + if self._vao is not None: + glDeleteVertexArrays(1, [self._vao]) + glDeleteBuffers(1, [self._vbo]) + self._vao = None + self._vbo = None + + def _in_context(self): + return self._vao is not None + + def _bind(self): + glBindVertexArray(self._vao) + + def _unbind(self): + glBindVertexArray(0) + + def delete(self): + self._unbind() + self._remove_from_context() + + def render_string(self, text, x, y, scale=1.0, + align=TextAlign.BOTTOM_LEFT): + """Render a string to the current view buffer. + + Note + ---- + Assumes correct shader program already bound w/ uniforms set. + + Parameters + ---------- + text : str + The text to render. + x : int + Horizontal pixel location of text. + y : int + Vertical pixel location of text. + scale : int + Scaling factor for text. + align : int + One of the TextAlign options which specifies where the ``x`` + and ``y`` parameters lie on the text. For example, + :attr:`.TextAlign.BOTTOM_LEFT` means that ``x`` and ``y`` indicate + the position of the bottom-left corner of the textbox. + """ + glActiveTexture(GL_TEXTURE0) + glEnable(GL_BLEND) + glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) + glDisable(GL_DEPTH_TEST) + glPolygonMode(GL_FRONT_AND_BACK, GL_FILL) + self._bind() + + # Determine width and height of text relative to x, y + width = 0.0 + height = 0.0 + for c in text: + ch = self._character_map[c] + height = max(height, ch.bearing[1] * scale) + width += (ch.advance >> 6) * scale + + # Determine offsets based on alignments + xoff = 0 + yoff = 0 + if align == TextAlign.BOTTOM_RIGHT: + xoff = -width + elif align == TextAlign.BOTTOM_CENTER: + xoff = -width / 2.0 + elif align == TextAlign.TOP_LEFT: + yoff = -height + elif align == TextAlign.TOP_RIGHT: + yoff = -height + xoff = -width + elif align == TextAlign.TOP_CENTER: + yoff = -height + xoff = -width / 2.0 + elif align == TextAlign.CENTER: + xoff = -width / 2.0 + yoff = -height / 2.0 + elif align == TextAlign.CENTER_LEFT: + yoff = -height / 2.0 + elif align == TextAlign.CENTER_RIGHT: + xoff = -width + yoff = -height / 2.0 + + x += xoff + y += yoff + + ch = None + for c in text: + ch = self._character_map[c] + xpos = x + ch.bearing[0] * scale + ypos = y - (ch.size[1] - ch.bearing[1]) * scale + w = ch.size[0] * scale + h = ch.size[1] * scale + + vertices = np.array([ + [xpos, ypos, 0.0, 0.0], + [xpos + w, ypos, 1.0, 0.0], + [xpos + w, ypos + h, 1.0, 1.0], + [xpos + w, ypos + h, 1.0, 1.0], + [xpos, ypos + h, 0.0, 1.0], + [xpos, ypos, 0.0, 0.0], + ], dtype=np.float32) + + ch.texture._bind() + + glBindBuffer(GL_ARRAY_BUFFER, self._vbo) + glBufferData( + GL_ARRAY_BUFFER, FLOAT_SZ * 6 * 4, vertices, GL_DYNAMIC_DRAW + ) + # TODO MAKE THIS MORE EFFICIENT, lgBufferSubData is broken + # glBufferSubData( + # GL_ARRAY_BUFFER, 0, 6 * 4 * FLOAT_SZ, + # np.ascontiguousarray(vertices.flatten) + # ) + glDrawArrays(GL_TRIANGLES, 0, 6) + x += (ch.advance >> 6) * scale + + self._unbind() + if ch: + ch.texture._unbind() diff --git a/pyrender/pyrender/fonts/OpenSans-Bold.ttf b/pyrender/pyrender/fonts/OpenSans-Bold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..fd79d43bea0293ac1b20e8aca1142627983d2c07 Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-Bold.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-BoldItalic.ttf b/pyrender/pyrender/fonts/OpenSans-BoldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..9bc800958a421d937fc392e00beaef4eea76dc71 Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-BoldItalic.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-ExtraBold.ttf b/pyrender/pyrender/fonts/OpenSans-ExtraBold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..21f6f84a0799946fc4ae02c52b27e61c3762c745 Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-ExtraBold.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-ExtraBoldItalic.ttf b/pyrender/pyrender/fonts/OpenSans-ExtraBoldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..31cb688340eff462dddf47efbb4dfef66cb7fbed Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-ExtraBoldItalic.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-Italic.ttf b/pyrender/pyrender/fonts/OpenSans-Italic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..c90da48ff3b8ad6167236d70c48df4d7b5de3bbb Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-Italic.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-Light.ttf b/pyrender/pyrender/fonts/OpenSans-Light.ttf new file mode 100644 index 0000000000000000000000000000000000000000..0d381897da20345fa63112f19042561f44ee3aa0 Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-Light.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-LightItalic.ttf b/pyrender/pyrender/fonts/OpenSans-LightItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..68299c4bc6b5b7adfff2c9aee4aed7c1547100ef Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-LightItalic.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-Regular.ttf b/pyrender/pyrender/fonts/OpenSans-Regular.ttf new file mode 100644 index 0000000000000000000000000000000000000000..db433349b7047f72f40072630c1bc110620bf09e Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-Regular.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-Semibold.ttf b/pyrender/pyrender/fonts/OpenSans-Semibold.ttf new file mode 100644 index 0000000000000000000000000000000000000000..1a7679e3949fb045f152f456bc4adad31e8b9f55 Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-Semibold.ttf differ diff --git a/pyrender/pyrender/fonts/OpenSans-SemiboldItalic.ttf b/pyrender/pyrender/fonts/OpenSans-SemiboldItalic.ttf new file mode 100644 index 0000000000000000000000000000000000000000..59b6d16b065f6baa6f70ddbd4322a4f44bb9636a Binary files /dev/null and b/pyrender/pyrender/fonts/OpenSans-SemiboldItalic.ttf differ diff --git a/pyrender/pyrender/light.py b/pyrender/pyrender/light.py new file mode 100644 index 0000000000000000000000000000000000000000..333d9e4e553a245c259251a89b69cb46b73b1278 --- /dev/null +++ b/pyrender/pyrender/light.py @@ -0,0 +1,385 @@ +"""Punctual light sources as defined by the glTF 2.0 KHR extension at +https://github.com/KhronosGroup/glTF/tree/master/extensions/2.0/Khronos/KHR_lights_punctual + +Author: Matthew Matl +""" +import abc +import numpy as np +import six + +from OpenGL.GL import * + +from .utils import format_color_vector +from .texture import Texture +from .constants import SHADOW_TEX_SZ +from .camera import OrthographicCamera, PerspectiveCamera + + + +@six.add_metaclass(abc.ABCMeta) +class Light(object): + """Base class for all light objects. + + Parameters + ---------- + color : (3,) float + RGB value for the light's color in linear space. + intensity : float + Brightness of light. The units that this is defined in depend on the + type of light. Point and spot lights use luminous intensity in candela + (lm/sr), while directional lights use illuminance in lux (lm/m2). + name : str, optional + Name of the light. + """ + def __init__(self, + color=None, + intensity=None, + name=None): + + if color is None: + color = np.ones(3) + if intensity is None: + intensity = 1.0 + + self.name = name + self.color = color + self.intensity = intensity + self._shadow_camera = None + self._shadow_texture = None + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def color(self): + """(3,) float : The light's color. + """ + return self._color + + @color.setter + def color(self, value): + self._color = format_color_vector(value, 3) + + @property + def intensity(self): + """float : The light's intensity in candela or lux. + """ + return self._intensity + + @intensity.setter + def intensity(self, value): + self._intensity = float(value) + + @property + def shadow_texture(self): + """:class:`.Texture` : A texture used to hold shadow maps for this light. + """ + return self._shadow_texture + + @shadow_texture.setter + def shadow_texture(self, value): + if self._shadow_texture is not None: + if self._shadow_texture._in_context(): + self._shadow_texture.delete() + self._shadow_texture = value + + @abc.abstractmethod + def _generate_shadow_texture(self, size=None): + """Generate a shadow texture for this light. + + Parameters + ---------- + size : int, optional + Size of texture map. Must be a positive power of two. + """ + pass + + @abc.abstractmethod + def _get_shadow_camera(self, scene_scale): + """Generate and return a shadow mapping camera for this light. + + Parameters + ---------- + scene_scale : float + Length of scene's bounding box diagonal. + + Returns + ------- + camera : :class:`.Camera` + The camera used to render shadowmaps for this light. + """ + pass + + +class DirectionalLight(Light): + """Directional lights are light sources that act as though they are + infinitely far away and emit light in the direction of the local -z axis. + This light type inherits the orientation of the node that it belongs to; + position and scale are ignored except for their effect on the inherited + node orientation. Because it is at an infinite distance, the light is + not attenuated. Its intensity is defined in lumens per metre squared, + or lux (lm/m2). + + Parameters + ---------- + color : (3,) float, optional + RGB value for the light's color in linear space. Defaults to white + (i.e. [1.0, 1.0, 1.0]). + intensity : float, optional + Brightness of light, in lux (lm/m^2). Defaults to 1.0 + name : str, optional + Name of the light. + """ + + def __init__(self, + color=None, + intensity=None, + name=None): + super(DirectionalLight, self).__init__( + color=color, + intensity=intensity, + name=name, + ) + + def _generate_shadow_texture(self, size=None): + """Generate a shadow texture for this light. + + Parameters + ---------- + size : int, optional + Size of texture map. Must be a positive power of two. + """ + if size is None: + size = SHADOW_TEX_SZ + self.shadow_texture = Texture(width=size, height=size, + source_channels='D', data_format=GL_FLOAT) + + def _get_shadow_camera(self, scene_scale): + """Generate and return a shadow mapping camera for this light. + + Parameters + ---------- + scene_scale : float + Length of scene's bounding box diagonal. + + Returns + ------- + camera : :class:`.Camera` + The camera used to render shadowmaps for this light. + """ + return OrthographicCamera( + znear=0.01 * scene_scale, + zfar=10 * scene_scale, + xmag=scene_scale, + ymag=scene_scale + ) + + +class PointLight(Light): + """Point lights emit light in all directions from their position in space; + rotation and scale are ignored except for their effect on the inherited + node position. The brightness of the light attenuates in a physically + correct manner as distance increases from the light's position (i.e. + brightness goes like the inverse square of the distance). Point light + intensity is defined in candela, which is lumens per square radian (lm/sr). + + Parameters + ---------- + color : (3,) float + RGB value for the light's color in linear space. + intensity : float + Brightness of light in candela (lm/sr). + range : float + Cutoff distance at which light's intensity may be considered to + have reached zero. If None, the range is assumed to be infinite. + name : str, optional + Name of the light. + """ + + def __init__(self, + color=None, + intensity=None, + range=None, + name=None): + super(PointLight, self).__init__( + color=color, + intensity=intensity, + name=name, + ) + self.range = range + + @property + def range(self): + """float : The cutoff distance for the light. + """ + return self._range + + @range.setter + def range(self, value): + if value is not None: + value = float(value) + if value <= 0: + raise ValueError('Range must be > 0') + self._range = value + self._range = value + + def _generate_shadow_texture(self, size=None): + """Generate a shadow texture for this light. + + Parameters + ---------- + size : int, optional + Size of texture map. Must be a positive power of two. + """ + raise NotImplementedError('Shadows not implemented for point lights') + + def _get_shadow_camera(self, scene_scale): + """Generate and return a shadow mapping camera for this light. + + Parameters + ---------- + scene_scale : float + Length of scene's bounding box diagonal. + + Returns + ------- + camera : :class:`.Camera` + The camera used to render shadowmaps for this light. + """ + raise NotImplementedError('Shadows not implemented for point lights') + + +class SpotLight(Light): + """Spot lights emit light in a cone in the direction of the local -z axis. + The angle and falloff of the cone is defined using two numbers, the + ``innerConeAngle`` and ``outerConeAngle``. + As with point lights, the brightness + also attenuates in a physically correct manner as distance increases from + the light's position (i.e. brightness goes like the inverse square of the + distance). Spot light intensity refers to the brightness inside the + ``innerConeAngle`` (and at the location of the light) and is defined in + candela, which is lumens per square radian (lm/sr). A spot light's position + and orientation are inherited from its node transform. Inherited scale does + not affect cone shape, and is ignored except for its effect on position + and orientation. + + Parameters + ---------- + color : (3,) float + RGB value for the light's color in linear space. + intensity : float + Brightness of light in candela (lm/sr). + range : float + Cutoff distance at which light's intensity may be considered to + have reached zero. If None, the range is assumed to be infinite. + innerConeAngle : float + Angle, in radians, from centre of spotlight where falloff begins. + Must be greater than or equal to ``0`` and less + than ``outerConeAngle``. Defaults to ``0``. + outerConeAngle : float + Angle, in radians, from centre of spotlight where falloff ends. + Must be greater than ``innerConeAngle`` and less than or equal to + ``PI / 2.0``. Defaults to ``PI / 4.0``. + name : str, optional + Name of the light. + """ + + def __init__(self, + color=None, + intensity=None, + range=None, + innerConeAngle=0.0, + outerConeAngle=(np.pi / 4.0), + name=None): + super(SpotLight, self).__init__( + name=name, + color=color, + intensity=intensity, + ) + self.outerConeAngle = outerConeAngle + self.innerConeAngle = innerConeAngle + self.range = range + + @property + def innerConeAngle(self): + """float : The inner cone angle in radians. + """ + return self._innerConeAngle + + @innerConeAngle.setter + def innerConeAngle(self, value): + if value < 0.0 or value > self.outerConeAngle: + raise ValueError('Invalid value for inner cone angle') + self._innerConeAngle = float(value) + + @property + def outerConeAngle(self): + """float : The outer cone angle in radians. + """ + return self._outerConeAngle + + @outerConeAngle.setter + def outerConeAngle(self, value): + if value < 0.0 or value > np.pi / 2.0 + 1e-9: + raise ValueError('Invalid value for outer cone angle') + self._outerConeAngle = float(value) + + @property + def range(self): + """float : The cutoff distance for the light. + """ + return self._range + + @range.setter + def range(self, value): + if value is not None: + value = float(value) + if value <= 0: + raise ValueError('Range must be > 0') + self._range = value + self._range = value + + def _generate_shadow_texture(self, size=None): + """Generate a shadow texture for this light. + + Parameters + ---------- + size : int, optional + Size of texture map. Must be a positive power of two. + """ + if size is None: + size = SHADOW_TEX_SZ + self.shadow_texture = Texture(width=size, height=size, + source_channels='D', data_format=GL_FLOAT) + + def _get_shadow_camera(self, scene_scale): + """Generate and return a shadow mapping camera for this light. + + Parameters + ---------- + scene_scale : float + Length of scene's bounding box diagonal. + + Returns + ------- + camera : :class:`.Camera` + The camera used to render shadowmaps for this light. + """ + return PerspectiveCamera( + znear=0.01 * scene_scale, + zfar=10 * scene_scale, + yfov=np.clip(2 * self.outerConeAngle + np.pi / 16.0, 0.0, np.pi), + aspectRatio=1.0 + ) + + +__all__ = ['Light', 'DirectionalLight', 'SpotLight', 'PointLight'] diff --git a/pyrender/pyrender/material.py b/pyrender/pyrender/material.py new file mode 100644 index 0000000000000000000000000000000000000000..3ce9c2d184ed213c84b015e36bea558cd1efc6b7 --- /dev/null +++ b/pyrender/pyrender/material.py @@ -0,0 +1,707 @@ +"""Material properties, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-material +and +https://github.com/KhronosGroup/glTF/tree/master/extensions/2.0/Khronos/KHR_materials_pbrSpecularGlossiness + +Author: Matthew Matl +""" +import abc +import numpy as np +import six + +from .constants import TexFlags +from .utils import format_color_vector, format_texture_source +from .texture import Texture + + +@six.add_metaclass(abc.ABCMeta) +class Material(object): + """Base for standard glTF 2.0 materials. + + Parameters + ---------- + name : str, optional + The user-defined name of this object. + normalTexture : (n,n,3) float or :class:`Texture`, optional + A tangent space normal map. The texture contains RGB components in + linear space. Each texel represents the XYZ components of a normal + vector in tangent space. Red [0 to 255] maps to X [-1 to 1]. Green + [0 to 255] maps to Y [-1 to 1]. Blue [128 to 255] maps to Z + [1/255 to 1]. The normal vectors use OpenGL conventions where +X is + right and +Y is up. +Z points toward the viewer. + occlusionTexture : (n,n,1) float or :class:`Texture`, optional + The occlusion map texture. The occlusion values are sampled from the R + channel. Higher values indicate areas that should receive full indirect + lighting and lower values indicate no indirect lighting. These values + are linear. If other channels are present (GBA), they are ignored for + occlusion calculations. + emissiveTexture : (n,n,3) float or :class:`Texture`, optional + The emissive map controls the color and intensity of the light being + emitted by the material. This texture contains RGB components in sRGB + color space. If a fourth component (A) is present, it is ignored. + emissiveFactor : (3,) float, optional + The RGB components of the emissive color of the material. These values + are linear. If an emissiveTexture is specified, this value is + multiplied with the texel values. + alphaMode : str, optional + The material's alpha rendering mode enumeration specifying the + interpretation of the alpha value of the main factor and texture. + Allowed Values: + + - `"OPAQUE"` The alpha value is ignored and the rendered output is + fully opaque. + - `"MASK"` The rendered output is either fully opaque or fully + transparent depending on the alpha value and the specified alpha + cutoff value. + - `"BLEND"` The alpha value is used to composite the source and + destination areas. The rendered output is combined with the + background using the normal painting operation (i.e. the Porter + and Duff over operator). + + alphaCutoff : float, optional + Specifies the cutoff threshold when in MASK mode. If the alpha value is + greater than or equal to this value then it is rendered as fully + opaque, otherwise, it is rendered as fully transparent. + A value greater than 1.0 will render the entire material as fully + transparent. This value is ignored for other modes. + doubleSided : bool, optional + Specifies whether the material is double sided. When this value is + false, back-face culling is enabled. When this value is true, + back-face culling is disabled and double sided lighting is enabled. + smooth : bool, optional + If True, the material is rendered smoothly by using only one normal + per vertex and face indexing. + wireframe : bool, optional + If True, the material is rendered in wireframe mode. + """ + + def __init__(self, + name=None, + normalTexture=None, + occlusionTexture=None, + emissiveTexture=None, + emissiveFactor=None, + alphaMode=None, + alphaCutoff=None, + doubleSided=False, + smooth=True, + wireframe=False): + + # Set defaults + if alphaMode is None: + alphaMode = 'OPAQUE' + + if alphaCutoff is None: + alphaCutoff = 0.5 + + if emissiveFactor is None: + emissiveFactor = np.zeros(3).astype(np.float32) + + self.name = name + self.normalTexture = normalTexture + self.occlusionTexture = occlusionTexture + self.emissiveTexture = emissiveTexture + self.emissiveFactor = emissiveFactor + self.alphaMode = alphaMode + self.alphaCutoff = alphaCutoff + self.doubleSided = doubleSided + self.smooth = smooth + self.wireframe = wireframe + + self._tex_flags = None + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def normalTexture(self): + """(n,n,3) float or :class:`Texture` : The tangent-space normal map. + """ + return self._normalTexture + + @normalTexture.setter + def normalTexture(self, value): + # TODO TMP + self._normalTexture = self._format_texture(value, 'RGB') + self._tex_flags = None + + @property + def occlusionTexture(self): + """(n,n,1) float or :class:`Texture` : The ambient occlusion map. + """ + return self._occlusionTexture + + @occlusionTexture.setter + def occlusionTexture(self, value): + self._occlusionTexture = self._format_texture(value, 'R') + self._tex_flags = None + + @property + def emissiveTexture(self): + """(n,n,3) float or :class:`Texture` : The emission map. + """ + return self._emissiveTexture + + @emissiveTexture.setter + def emissiveTexture(self, value): + self._emissiveTexture = self._format_texture(value, 'RGB') + self._tex_flags = None + + @property + def emissiveFactor(self): + """(3,) float : Base multiplier for emission colors. + """ + return self._emissiveFactor + + @emissiveFactor.setter + def emissiveFactor(self, value): + if value is None: + value = np.zeros(3) + self._emissiveFactor = format_color_vector(value, 3) + + @property + def alphaMode(self): + """str : The mode for blending. + """ + return self._alphaMode + + @alphaMode.setter + def alphaMode(self, value): + if value not in set(['OPAQUE', 'MASK', 'BLEND']): + raise ValueError('Invalid alpha mode {}'.format(value)) + self._alphaMode = value + + @property + def alphaCutoff(self): + """float : The cutoff threshold in MASK mode. + """ + return self._alphaCutoff + + @alphaCutoff.setter + def alphaCutoff(self, value): + if value < 0 or value > 1: + raise ValueError('Alpha cutoff must be in range [0,1]') + self._alphaCutoff = float(value) + + @property + def doubleSided(self): + """bool : Whether the material is double-sided. + """ + return self._doubleSided + + @doubleSided.setter + def doubleSided(self, value): + if not isinstance(value, bool): + raise TypeError('Double sided must be a boolean value') + self._doubleSided = value + + @property + def smooth(self): + """bool : Whether to render the mesh smoothly by + interpolating vertex normals. + """ + return self._smooth + + @smooth.setter + def smooth(self, value): + if not isinstance(value, bool): + raise TypeError('Double sided must be a boolean value') + self._smooth = value + + @property + def wireframe(self): + """bool : Whether to render the mesh in wireframe mode. + """ + return self._wireframe + + @wireframe.setter + def wireframe(self, value): + if not isinstance(value, bool): + raise TypeError('Wireframe must be a boolean value') + self._wireframe = value + + @property + def is_transparent(self): + """bool : If True, the object is partially transparent. + """ + return self._compute_transparency() + + @property + def tex_flags(self): + """int : Texture availability flags. + """ + if self._tex_flags is None: + self._tex_flags = self._compute_tex_flags() + return self._tex_flags + + @property + def textures(self): + """list of :class:`Texture` : The textures associated with this + material. + """ + return self._compute_textures() + + def _compute_transparency(self): + return False + + def _compute_tex_flags(self): + tex_flags = TexFlags.NONE + if self.normalTexture is not None: + tex_flags |= TexFlags.NORMAL + if self.occlusionTexture is not None: + tex_flags |= TexFlags.OCCLUSION + if self.emissiveTexture is not None: + tex_flags |= TexFlags.EMISSIVE + return tex_flags + + def _compute_textures(self): + all_textures = [ + self.normalTexture, self.occlusionTexture, self.emissiveTexture + ] + textures = set([t for t in all_textures if t is not None]) + return textures + + def _format_texture(self, texture, target_channels='RGB'): + """Format a texture as a float32 np array. + """ + if isinstance(texture, Texture) or texture is None: + return texture + else: + source = format_texture_source(texture, target_channels) + return Texture(source=source, source_channels=target_channels) + + +class MetallicRoughnessMaterial(Material): + """A material based on the metallic-roughness material model from + Physically-Based Rendering (PBR) methodology. + + Parameters + ---------- + name : str, optional + The user-defined name of this object. + normalTexture : (n,n,3) float or :class:`Texture`, optional + A tangent space normal map. The texture contains RGB components in + linear space. Each texel represents the XYZ components of a normal + vector in tangent space. Red [0 to 255] maps to X [-1 to 1]. Green + [0 to 255] maps to Y [-1 to 1]. Blue [128 to 255] maps to Z + [1/255 to 1]. The normal vectors use OpenGL conventions where +X is + right and +Y is up. +Z points toward the viewer. + occlusionTexture : (n,n,1) float or :class:`Texture`, optional + The occlusion map texture. The occlusion values are sampled from the R + channel. Higher values indicate areas that should receive full indirect + lighting and lower values indicate no indirect lighting. These values + are linear. If other channels are present (GBA), they are ignored for + occlusion calculations. + emissiveTexture : (n,n,3) float or :class:`Texture`, optional + The emissive map controls the color and intensity of the light being + emitted by the material. This texture contains RGB components in sRGB + color space. If a fourth component (A) is present, it is ignored. + emissiveFactor : (3,) float, optional + The RGB components of the emissive color of the material. These values + are linear. If an emissiveTexture is specified, this value is + multiplied with the texel values. + alphaMode : str, optional + The material's alpha rendering mode enumeration specifying the + interpretation of the alpha value of the main factor and texture. + Allowed Values: + + - `"OPAQUE"` The alpha value is ignored and the rendered output is + fully opaque. + - `"MASK"` The rendered output is either fully opaque or fully + transparent depending on the alpha value and the specified alpha + cutoff value. + - `"BLEND"` The alpha value is used to composite the source and + destination areas. The rendered output is combined with the + background using the normal painting operation (i.e. the Porter + and Duff over operator). + + alphaCutoff : float, optional + Specifies the cutoff threshold when in MASK mode. If the alpha value is + greater than or equal to this value then it is rendered as fully + opaque, otherwise, it is rendered as fully transparent. + A value greater than 1.0 will render the entire material as fully + transparent. This value is ignored for other modes. + doubleSided : bool, optional + Specifies whether the material is double sided. When this value is + false, back-face culling is enabled. When this value is true, + back-face culling is disabled and double sided lighting is enabled. + smooth : bool, optional + If True, the material is rendered smoothly by using only one normal + per vertex and face indexing. + wireframe : bool, optional + If True, the material is rendered in wireframe mode. + baseColorFactor : (4,) float, optional + The RGBA components of the base color of the material. The fourth + component (A) is the alpha coverage of the material. The alphaMode + property specifies how alpha is interpreted. These values are linear. + If a baseColorTexture is specified, this value is multiplied with the + texel values. + baseColorTexture : (n,n,4) float or :class:`Texture`, optional + The base color texture. This texture contains RGB(A) components in sRGB + color space. The first three components (RGB) specify the base color of + the material. If the fourth component (A) is present, it represents the + alpha coverage of the material. Otherwise, an alpha of 1.0 is assumed. + The alphaMode property specifies how alpha is interpreted. + The stored texels must not be premultiplied. + metallicFactor : float + The metalness of the material. A value of 1.0 means the material is a + metal. A value of 0.0 means the material is a dielectric. Values in + between are for blending between metals and dielectrics such as dirty + metallic surfaces. This value is linear. If a metallicRoughnessTexture + is specified, this value is multiplied with the metallic texel values. + roughnessFactor : float + The roughness of the material. A value of 1.0 means the material is + completely rough. A value of 0.0 means the material is completely + smooth. This value is linear. If a metallicRoughnessTexture is + specified, this value is multiplied with the roughness texel values. + metallicRoughnessTexture : (n,n,2) float or :class:`Texture`, optional + The metallic-roughness texture. The metalness values are sampled from + the B channel. The roughness values are sampled from the G channel. + These values are linear. If other channels are present (R or A), they + are ignored for metallic-roughness calculations. + """ + + def __init__(self, + name=None, + normalTexture=None, + occlusionTexture=None, + emissiveTexture=None, + emissiveFactor=None, + alphaMode=None, + alphaCutoff=None, + doubleSided=False, + smooth=True, + wireframe=False, + baseColorFactor=None, + baseColorTexture=None, + metallicFactor=1.0, + roughnessFactor=1.0, + metallicRoughnessTexture=None): + super(MetallicRoughnessMaterial, self).__init__( + name=name, + normalTexture=normalTexture, + occlusionTexture=occlusionTexture, + emissiveTexture=emissiveTexture, + emissiveFactor=emissiveFactor, + alphaMode=alphaMode, + alphaCutoff=alphaCutoff, + doubleSided=doubleSided, + smooth=smooth, + wireframe=wireframe + ) + + # Set defaults + if baseColorFactor is None: + baseColorFactor = np.ones(4).astype(np.float32) + + self.baseColorFactor = baseColorFactor + self.baseColorTexture = baseColorTexture + self.metallicFactor = metallicFactor + self.roughnessFactor = roughnessFactor + self.metallicRoughnessTexture = metallicRoughnessTexture + + @property + def baseColorFactor(self): + """(4,) float or :class:`Texture` : The RGBA base color multiplier. + """ + return self._baseColorFactor + + @baseColorFactor.setter + def baseColorFactor(self, value): + if value is None: + value = np.ones(4) + self._baseColorFactor = format_color_vector(value, 4) + + @property + def baseColorTexture(self): + """(n,n,4) float or :class:`Texture` : The diffuse texture. + """ + return self._baseColorTexture + + @baseColorTexture.setter + def baseColorTexture(self, value): + self._baseColorTexture = self._format_texture(value, 'RGBA') + self._tex_flags = None + + @property + def metallicFactor(self): + """float : The metalness of the material. + """ + return self._metallicFactor + + @metallicFactor.setter + def metallicFactor(self, value): + if value is None: + value = 1.0 + if value < 0 or value > 1: + raise ValueError('Metallic factor must be in range [0,1]') + self._metallicFactor = float(value) + + @property + def roughnessFactor(self): + """float : The roughness of the material. + """ + return self.RoughnessFactor + + @roughnessFactor.setter + def roughnessFactor(self, value): + if value is None: + value = 1.0 + if value < 0 or value > 1: + raise ValueError('Roughness factor must be in range [0,1]') + self.RoughnessFactor = float(value) + + @property + def metallicRoughnessTexture(self): + """(n,n,2) float or :class:`Texture` : The metallic-roughness texture. + """ + return self._metallicRoughnessTexture + + @metallicRoughnessTexture.setter + def metallicRoughnessTexture(self, value): + self._metallicRoughnessTexture = self._format_texture(value, 'GB') + self._tex_flags = None + + def _compute_tex_flags(self): + tex_flags = super(MetallicRoughnessMaterial, self)._compute_tex_flags() + if self.baseColorTexture is not None: + tex_flags |= TexFlags.BASE_COLOR + if self.metallicRoughnessTexture is not None: + tex_flags |= TexFlags.METALLIC_ROUGHNESS + return tex_flags + + def _compute_transparency(self): + if self.alphaMode == 'OPAQUE': + return False + cutoff = self.alphaCutoff + if self.alphaMode == 'BLEND': + cutoff = 1.0 + if self.baseColorFactor[3] < cutoff: + return True + if (self.baseColorTexture is not None and + self.baseColorTexture.is_transparent(cutoff)): + return True + return False + + def _compute_textures(self): + textures = super(MetallicRoughnessMaterial, self)._compute_textures() + all_textures = [self.baseColorTexture, self.metallicRoughnessTexture] + all_textures = {t for t in all_textures if t is not None} + textures |= all_textures + return textures + + +class SpecularGlossinessMaterial(Material): + """A material based on the specular-glossiness material model from + Physically-Based Rendering (PBR) methodology. + + Parameters + ---------- + name : str, optional + The user-defined name of this object. + normalTexture : (n,n,3) float or :class:`Texture`, optional + A tangent space normal map. The texture contains RGB components in + linear space. Each texel represents the XYZ components of a normal + vector in tangent space. Red [0 to 255] maps to X [-1 to 1]. Green + [0 to 255] maps to Y [-1 to 1]. Blue [128 to 255] maps to Z + [1/255 to 1]. The normal vectors use OpenGL conventions where +X is + right and +Y is up. +Z points toward the viewer. + occlusionTexture : (n,n,1) float or :class:`Texture`, optional + The occlusion map texture. The occlusion values are sampled from the R + channel. Higher values indicate areas that should receive full indirect + lighting and lower values indicate no indirect lighting. These values + are linear. If other channels are present (GBA), they are ignored for + occlusion calculations. + emissiveTexture : (n,n,3) float or :class:`Texture`, optional + The emissive map controls the color and intensity of the light being + emitted by the material. This texture contains RGB components in sRGB + color space. If a fourth component (A) is present, it is ignored. + emissiveFactor : (3,) float, optional + The RGB components of the emissive color of the material. These values + are linear. If an emissiveTexture is specified, this value is + multiplied with the texel values. + alphaMode : str, optional + The material's alpha rendering mode enumeration specifying the + interpretation of the alpha value of the main factor and texture. + Allowed Values: + + - `"OPAQUE"` The alpha value is ignored and the rendered output is + fully opaque. + - `"MASK"` The rendered output is either fully opaque or fully + transparent depending on the alpha value and the specified alpha + cutoff value. + - `"BLEND"` The alpha value is used to composite the source and + destination areas. The rendered output is combined with the + background using the normal painting operation (i.e. the Porter + and Duff over operator). + + alphaCutoff : float, optional + Specifies the cutoff threshold when in MASK mode. If the alpha value is + greater than or equal to this value then it is rendered as fully + opaque, otherwise, it is rendered as fully transparent. + A value greater than 1.0 will render the entire material as fully + transparent. This value is ignored for other modes. + doubleSided : bool, optional + Specifies whether the material is double sided. When this value is + false, back-face culling is enabled. When this value is true, + back-face culling is disabled and double sided lighting is enabled. + smooth : bool, optional + If True, the material is rendered smoothly by using only one normal + per vertex and face indexing. + wireframe : bool, optional + If True, the material is rendered in wireframe mode. + diffuseFactor : (4,) float + The RGBA components of the reflected diffuse color of the material. + Metals have a diffuse value of [0.0, 0.0, 0.0]. The fourth component + (A) is the opacity of the material. The values are linear. + diffuseTexture : (n,n,4) float or :class:`Texture`, optional + The diffuse texture. This texture contains RGB(A) components of the + reflected diffuse color of the material in sRGB color space. If the + fourth component (A) is present, it represents the alpha coverage of + the material. Otherwise, an alpha of 1.0 is assumed. + The alphaMode property specifies how alpha is interpreted. + The stored texels must not be premultiplied. + specularFactor : (3,) float + The specular RGB color of the material. This value is linear. + glossinessFactor : float + The glossiness or smoothness of the material. A value of 1.0 means the + material has full glossiness or is perfectly smooth. A value of 0.0 + means the material has no glossiness or is perfectly rough. This value + is linear. + specularGlossinessTexture : (n,n,4) or :class:`Texture`, optional + The specular-glossiness texture is a RGBA texture, containing the + specular color (RGB) in sRGB space and the glossiness value (A) in + linear space. + """ + + def __init__(self, + name=None, + normalTexture=None, + occlusionTexture=None, + emissiveTexture=None, + emissiveFactor=None, + alphaMode=None, + alphaCutoff=None, + doubleSided=False, + smooth=True, + wireframe=False, + diffuseFactor=None, + diffuseTexture=None, + specularFactor=None, + glossinessFactor=1.0, + specularGlossinessTexture=None): + super(SpecularGlossinessMaterial, self).__init__( + name=name, + normalTexture=normalTexture, + occlusionTexture=occlusionTexture, + emissiveTexture=emissiveTexture, + emissiveFactor=emissiveFactor, + alphaMode=alphaMode, + alphaCutoff=alphaCutoff, + doubleSided=doubleSided, + smooth=smooth, + wireframe=wireframe + ) + + # Set defaults + if diffuseFactor is None: + diffuseFactor = np.ones(4).astype(np.float32) + if specularFactor is None: + specularFactor = np.ones(3).astype(np.float32) + + self.diffuseFactor = diffuseFactor + self.diffuseTexture = diffuseTexture + self.specularFactor = specularFactor + self.glossinessFactor = glossinessFactor + self.specularGlossinessTexture = specularGlossinessTexture + + @property + def diffuseFactor(self): + """(4,) float : The diffuse base color. + """ + return self._diffuseFactor + + @diffuseFactor.setter + def diffuseFactor(self, value): + self._diffuseFactor = format_color_vector(value, 4) + + @property + def diffuseTexture(self): + """(n,n,4) float or :class:`Texture` : The diffuse map. + """ + return self._diffuseTexture + + @diffuseTexture.setter + def diffuseTexture(self, value): + self._diffuseTexture = self._format_texture(value, 'RGBA') + self._tex_flags = None + + @property + def specularFactor(self): + """(3,) float : The specular color of the material. + """ + return self._specularFactor + + @specularFactor.setter + def specularFactor(self, value): + self._specularFactor = format_color_vector(value, 3) + + @property + def glossinessFactor(self): + """float : The glossiness of the material. + """ + return self.glossinessFactor + + @glossinessFactor.setter + def glossinessFactor(self, value): + if value < 0 or value > 1: + raise ValueError('glossiness factor must be in range [0,1]') + self._glossinessFactor = float(value) + + @property + def specularGlossinessTexture(self): + """(n,n,4) or :class:`Texture` : The specular-glossiness texture. + """ + return self._specularGlossinessTexture + + @specularGlossinessTexture.setter + def specularGlossinessTexture(self, value): + self._specularGlossinessTexture = self._format_texture(value, 'GB') + self._tex_flags = None + + def _compute_tex_flags(self): + flags = super(SpecularGlossinessMaterial, self)._compute_tex_flags() + if self.diffuseTexture is not None: + flags |= TexFlags.DIFFUSE + if self.specularGlossinessTexture is not None: + flags |= TexFlags.SPECULAR_GLOSSINESS + return flags + + def _compute_transparency(self): + if self.alphaMode == 'OPAQUE': + return False + cutoff = self.alphaCutoff + if self.alphaMode == 'BLEND': + cutoff = 1.0 + if self.diffuseFactor[3] < cutoff: + return True + if (self.diffuseTexture is not None and + self.diffuseTexture.is_transparent(cutoff)): + return True + return False + + def _compute_textures(self): + textures = super(SpecularGlossinessMaterial, self)._compute_textures() + all_textures = [self.diffuseTexture, self.specularGlossinessTexture] + all_textures = {t for t in all_textures if t is not None} + textures |= all_textures + return textures diff --git a/pyrender/pyrender/mesh.py b/pyrender/pyrender/mesh.py new file mode 100644 index 0000000000000000000000000000000000000000..36833ea3dfa6c095a18fc745ff34cf106e83c95d --- /dev/null +++ b/pyrender/pyrender/mesh.py @@ -0,0 +1,328 @@ +"""Meshes, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-mesh + +Author: Matthew Matl +""" +import copy + +import numpy as np +import trimesh + +from .primitive import Primitive +from .constants import GLTF +from .material import MetallicRoughnessMaterial + + +class Mesh(object): + """A set of primitives to be rendered. + + Parameters + ---------- + name : str + The user-defined name of this object. + primitives : list of :class:`Primitive` + The primitives associated with this mesh. + weights : (k,) float + Array of weights to be applied to the Morph Targets. + is_visible : bool + If False, the mesh will not be rendered. + """ + + def __init__(self, primitives, name=None, weights=None, is_visible=True): + self.primitives = primitives + self.name = name + self.weights = weights + self.is_visible = is_visible + + self._bounds = None + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def primitives(self): + """list of :class:`Primitive` : The primitives associated + with this mesh. + """ + return self._primitives + + @primitives.setter + def primitives(self, value): + self._primitives = value + + @property + def weights(self): + """(k,) float : Weights to be applied to morph targets. + """ + return self._weights + + @weights.setter + def weights(self, value): + self._weights = value + + @property + def is_visible(self): + """bool : Whether the mesh is visible. + """ + return self._is_visible + + @is_visible.setter + def is_visible(self, value): + self._is_visible = value + + @property + def bounds(self): + """(2,3) float : The axis-aligned bounds of the mesh. + """ + if self._bounds is None: + bounds = np.array([[np.infty, np.infty, np.infty], + [-np.infty, -np.infty, -np.infty]]) + for p in self.primitives: + bounds[0] = np.minimum(bounds[0], p.bounds[0]) + bounds[1] = np.maximum(bounds[1], p.bounds[1]) + self._bounds = bounds + return self._bounds + + @property + def centroid(self): + """(3,) float : The centroid of the mesh's axis-aligned bounding box + (AABB). + """ + return np.mean(self.bounds, axis=0) + + @property + def extents(self): + """(3,) float : The lengths of the axes of the mesh's AABB. + """ + return np.diff(self.bounds, axis=0).reshape(-1) + + @property + def scale(self): + """(3,) float : The length of the diagonal of the mesh's AABB. + """ + return np.linalg.norm(self.extents) + + @property + def is_transparent(self): + """bool : If True, the mesh is partially-transparent. + """ + for p in self.primitives: + if p.is_transparent: + return True + return False + + @staticmethod + def from_points(points, colors=None, normals=None, + is_visible=True, poses=None): + """Create a Mesh from a set of points. + + Parameters + ---------- + points : (n,3) float + The point positions. + colors : (n,3) or (n,4) float, optional + RGB or RGBA colors for each point. + normals : (n,3) float, optionals + The normal vectors for each point. + is_visible : bool + If False, the points will not be rendered. + poses : (x,4,4) + Array of 4x4 transformation matrices for instancing this object. + + Returns + ------- + mesh : :class:`Mesh` + The created mesh. + """ + primitive = Primitive( + positions=points, + normals=normals, + color_0=colors, + mode=GLTF.POINTS, + poses=poses + ) + mesh = Mesh(primitives=[primitive], is_visible=is_visible) + return mesh + + @staticmethod + def from_trimesh(mesh, material=None, is_visible=True, + poses=None, wireframe=False, smooth=True): + """Create a Mesh from a :class:`~trimesh.base.Trimesh`. + + Parameters + ---------- + mesh : :class:`~trimesh.base.Trimesh` or list of them + A triangular mesh or a list of meshes. + material : :class:`Material` + The material of the object. Overrides any mesh material. + If not specified and the mesh has no material, a default material + will be used. + is_visible : bool + If False, the mesh will not be rendered. + poses : (n,4,4) float + Array of 4x4 transformation matrices for instancing this object. + wireframe : bool + If `True`, the mesh will be rendered as a wireframe object + smooth : bool + If `True`, the mesh will be rendered with interpolated vertex + normals. Otherwise, the mesh edges will stay sharp. + + Returns + ------- + mesh : :class:`Mesh` + The created mesh. + """ + + if isinstance(mesh, (list, tuple, set, np.ndarray)): + meshes = list(mesh) + elif isinstance(mesh, trimesh.Trimesh): + meshes = [mesh] + else: + raise TypeError('Expected a Trimesh or a list, got a {}' + .format(type(mesh))) + + primitives = [] + for m in meshes: + positions = None + normals = None + indices = None + + # Compute positions, normals, and indices + if smooth: + positions = m.vertices.copy() + normals = m.vertex_normals.copy() + indices = m.faces.copy() + else: + positions = m.vertices[m.faces].reshape((3 * len(m.faces), 3)) + normals = np.repeat(m.face_normals, 3, axis=0) + + # Compute colors, texture coords, and material properties + color_0, texcoord_0, primitive_material = Mesh._get_trimesh_props(m, smooth=smooth, material=material) + + # Override if material is given. + if material is not None: + #primitive_material = copy.copy(material) + primitive_material = copy.deepcopy(material) # TODO + + if primitive_material is None: + # Replace material with default if needed + primitive_material = MetallicRoughnessMaterial( + alphaMode='BLEND', + baseColorFactor=[0.3, 0.3, 0.3, 1.0], + metallicFactor=0.2, + roughnessFactor=0.8 + ) + + primitive_material.wireframe = wireframe + + # Create the primitive + primitives.append(Primitive( + positions=positions, + normals=normals, + texcoord_0=texcoord_0, + color_0=color_0, + indices=indices, + material=primitive_material, + mode=GLTF.TRIANGLES, + poses=poses + )) + + return Mesh(primitives=primitives, is_visible=is_visible) + + @staticmethod + def _get_trimesh_props(mesh, smooth=False, material=None): + """Gets the vertex colors, texture coordinates, and material properties + from a :class:`~trimesh.base.Trimesh`. + """ + colors = None + texcoords = None + + # If the trimesh visual is undefined, return none for both + if not mesh.visual.defined: + return colors, texcoords, material + + # Process vertex colors + if material is None: + if mesh.visual.kind == 'vertex': + vc = mesh.visual.vertex_colors.copy() + if smooth: + colors = vc + else: + colors = vc[mesh.faces].reshape( + (3 * len(mesh.faces), vc.shape[1]) + ) + material = MetallicRoughnessMaterial( + alphaMode='BLEND', + baseColorFactor=[1.0, 1.0, 1.0, 1.0], + metallicFactor=0.2, + roughnessFactor=0.8 + ) + # Process face colors + elif mesh.visual.kind == 'face': + if smooth: + raise ValueError('Cannot use face colors with a smooth mesh') + else: + colors = np.repeat(mesh.visual.face_colors, 3, axis=0) + + material = MetallicRoughnessMaterial( + alphaMode='BLEND', + baseColorFactor=[1.0, 1.0, 1.0, 1.0], + metallicFactor=0.2, + roughnessFactor=0.8 + ) + + # Process texture colors + if mesh.visual.kind == 'texture': + # Configure UV coordinates + if mesh.visual.uv is not None and len(mesh.visual.uv) != 0: + uv = mesh.visual.uv.copy() + if smooth: + texcoords = uv + else: + texcoords = uv[mesh.faces].reshape( + (3 * len(mesh.faces), uv.shape[1]) + ) + + if material is None: + # Configure mesh material + mat = mesh.visual.material + + if isinstance(mat, trimesh.visual.texture.PBRMaterial): + material = MetallicRoughnessMaterial( + normalTexture=mat.normalTexture, + occlusionTexture=mat.occlusionTexture, + emissiveTexture=mat.emissiveTexture, + emissiveFactor=mat.emissiveFactor, + alphaMode='BLEND', + baseColorFactor=mat.baseColorFactor, + baseColorTexture=mat.baseColorTexture, + metallicFactor=mat.metallicFactor, + roughnessFactor=mat.roughnessFactor, + metallicRoughnessTexture=mat.metallicRoughnessTexture, + doubleSided=mat.doubleSided, + alphaCutoff=mat.alphaCutoff + ) + elif isinstance(mat, trimesh.visual.texture.SimpleMaterial): + glossiness = mat.kwargs.get('Ns', 1.0) + if isinstance(glossiness, list): + glossiness = float(glossiness[0]) + roughness = (2 / (glossiness + 2)) ** (1.0 / 4.0) + material = MetallicRoughnessMaterial( + alphaMode='BLEND', + roughnessFactor=roughness, + baseColorFactor=mat.diffuse, + baseColorTexture=mat.image, + ) + elif isinstance(mat, MetallicRoughnessMaterial): + material = mat + + return colors, texcoords, material diff --git a/pyrender/pyrender/node.py b/pyrender/pyrender/node.py new file mode 100644 index 0000000000000000000000000000000000000000..1f37f7856cc732a37dc58253022a7c331489493e --- /dev/null +++ b/pyrender/pyrender/node.py @@ -0,0 +1,263 @@ +"""Nodes, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-node + +Author: Matthew Matl +""" +import numpy as np + +import trimesh.transformations as transformations + +from .camera import Camera +from .mesh import Mesh +from .light import Light + + +class Node(object): + """A node in the node hierarchy. + + Parameters + ---------- + name : str, optional + The user-defined name of this object. + camera : :class:`Camera`, optional + The camera in this node. + children : list of :class:`Node` + The children of this node. + skin : int, optional + The index of the skin referenced by this node. + matrix : (4,4) float, optional + A floating-point 4x4 transformation matrix. + mesh : :class:`Mesh`, optional + The mesh in this node. + rotation : (4,) float, optional + The node's unit quaternion in the order (x, y, z, w), where + w is the scalar. + scale : (3,) float, optional + The node's non-uniform scale, given as the scaling factors along the x, + y, and z axes. + translation : (3,) float, optional + The node's translation along the x, y, and z axes. + weights : (n,) float + The weights of the instantiated Morph Target. Number of elements must + match number of Morph Targets of used mesh. + light : :class:`Light`, optional + The light in this node. + """ + + def __init__(self, + name=None, + camera=None, + children=None, + skin=None, + matrix=None, + mesh=None, + rotation=None, + scale=None, + translation=None, + weights=None, + light=None): + # Set defaults + if children is None: + children = [] + + self._matrix = None + self._scale = None + self._rotation = None + self._translation = None + if matrix is None: + if rotation is None: + rotation = np.array([0.0, 0.0, 0.0, 1.0]) + if translation is None: + translation = np.zeros(3) + if scale is None: + scale = np.ones(3) + self.rotation = rotation + self.translation = translation + self.scale = scale + else: + self.matrix = matrix + + self.name = name + self.camera = camera + self.children = children + self.skin = skin + self.mesh = mesh + self.weights = weights + self.light = light + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def camera(self): + """:class:`Camera` : The camera in this node. + """ + return self._camera + + @camera.setter + def camera(self, value): + if value is not None and not isinstance(value, Camera): + raise TypeError('Value must be a camera') + self._camera = value + + @property + def children(self): + """list of :class:`Node` : The children of this node. + """ + return self._children + + @children.setter + def children(self, value): + self._children = value + + @property + def skin(self): + """int : The skin index for this node. + """ + return self._skin + + @skin.setter + def skin(self, value): + self._skin = value + + @property + def mesh(self): + """:class:`Mesh` : The mesh in this node. + """ + return self._mesh + + @mesh.setter + def mesh(self, value): + if value is not None and not isinstance(value, Mesh): + raise TypeError('Value must be a mesh') + self._mesh = value + + @property + def light(self): + """:class:`Light` : The light in this node. + """ + return self._light + + @light.setter + def light(self, value): + if value is not None and not isinstance(value, Light): + raise TypeError('Value must be a light') + self._light = value + + @property + def rotation(self): + """(4,) float : The xyzw quaternion for this node. + """ + return self._rotation + + @rotation.setter + def rotation(self, value): + value = np.asanyarray(value) + if value.shape != (4,): + raise ValueError('Quaternion must be a (4,) vector') + if np.abs(np.linalg.norm(value) - 1.0) > 1e-3: + raise ValueError('Quaternion must have norm == 1.0') + self._rotation = value + self._matrix = None + + @property + def translation(self): + """(3,) float : The translation for this node. + """ + return self._translation + + @translation.setter + def translation(self, value): + value = np.asanyarray(value) + if value.shape != (3,): + raise ValueError('Translation must be a (3,) vector') + self._translation = value + self._matrix = None + + @property + def scale(self): + """(3,) float : The scale for this node. + """ + return self._scale + + @scale.setter + def scale(self, value): + value = np.asanyarray(value) + if value.shape != (3,): + raise ValueError('Scale must be a (3,) vector') + self._scale = value + self._matrix = None + + @property + def matrix(self): + """(4,4) float : The homogenous transform matrix for this node. + + Note that this matrix's elements are not settable, + it's just a copy of the internal matrix. You can set the whole + matrix, but not an individual element. + """ + if self._matrix is None: + self._matrix = self._m_from_tqs( + self.translation, self.rotation, self.scale + ) + return self._matrix.copy() + + @matrix.setter + def matrix(self, value): + value = np.asanyarray(value) + if value.shape != (4,4): + raise ValueError('Matrix must be a 4x4 numpy ndarray') + if not np.allclose(value[3,:], np.array([0.0, 0.0, 0.0, 1.0])): + raise ValueError('Bottom row of matrix must be [0,0,0,1]') + self.rotation = Node._q_from_m(value) + self.scale = Node._s_from_m(value) + self.translation = Node._t_from_m(value) + self._matrix = value + + @staticmethod + def _t_from_m(m): + return m[:3,3] + + @staticmethod + def _r_from_m(m): + U = m[:3,:3] + norms = np.linalg.norm(U.T, axis=1) + return U / norms + + @staticmethod + def _q_from_m(m): + M = np.eye(4) + M[:3,:3] = Node._r_from_m(m) + q_wxyz = transformations.quaternion_from_matrix(M) + return np.roll(q_wxyz, -1) + + @staticmethod + def _s_from_m(m): + return np.linalg.norm(m[:3,:3].T, axis=1) + + @staticmethod + def _r_from_q(q): + q_wxyz = np.roll(q, 1) + return transformations.quaternion_matrix(q_wxyz)[:3,:3] + + @staticmethod + def _m_from_tqs(t, q, s): + S = np.eye(4) + S[:3,:3] = np.diag(s) + + R = np.eye(4) + R[:3,:3] = Node._r_from_q(q) + + T = np.eye(4) + T[:3,3] = t + + return T.dot(R.dot(S)) diff --git a/pyrender/pyrender/offscreen.py b/pyrender/pyrender/offscreen.py new file mode 100644 index 0000000000000000000000000000000000000000..340142983006cdc6f51b6d114e9b2b294aa4a919 --- /dev/null +++ b/pyrender/pyrender/offscreen.py @@ -0,0 +1,160 @@ +"""Wrapper for offscreen rendering. + +Author: Matthew Matl +""" +import os + +from .renderer import Renderer +from .constants import RenderFlags + + +class OffscreenRenderer(object): + """A wrapper for offscreen rendering. + + Parameters + ---------- + viewport_width : int + The width of the main viewport, in pixels. + viewport_height : int + The height of the main viewport, in pixels. + point_size : float + The size of screen-space points in pixels. + """ + + def __init__(self, viewport_width, viewport_height, point_size=1.0): + self.viewport_width = viewport_width + self.viewport_height = viewport_height + self.point_size = point_size + + self._platform = None + self._renderer = None + self._create() + + @property + def viewport_width(self): + """int : The width of the main viewport, in pixels. + """ + return self._viewport_width + + @viewport_width.setter + def viewport_width(self, value): + self._viewport_width = int(value) + + @property + def viewport_height(self): + """int : The height of the main viewport, in pixels. + """ + return self._viewport_height + + @viewport_height.setter + def viewport_height(self, value): + self._viewport_height = int(value) + + @property + def point_size(self): + """float : The pixel size of points in point clouds. + """ + return self._point_size + + @point_size.setter + def point_size(self, value): + self._point_size = float(value) + + def render(self, scene, flags=RenderFlags.NONE, seg_node_map=None): + """Render a scene with the given set of flags. + + Parameters + ---------- + scene : :class:`Scene` + A scene to render. + flags : int + A bitwise or of one or more flags from :class:`.RenderFlags`. + seg_node_map : dict + A map from :class:`.Node` objects to (3,) colors for each. + If specified along with flags set to :attr:`.RenderFlags.SEG`, + the color image will be a segmentation image. + + Returns + ------- + color_im : (h, w, 3) uint8 or (h, w, 4) uint8 + The color buffer in RGB format, or in RGBA format if + :attr:`.RenderFlags.RGBA` is set. + Not returned if flags includes :attr:`.RenderFlags.DEPTH_ONLY`. + depth_im : (h, w) float32 + The depth buffer in linear units. + """ + self._platform.make_current() + # If platform does not support dynamically-resizing framebuffers, + # destroy it and restart it + if (self._platform.viewport_height != self.viewport_height or + self._platform.viewport_width != self.viewport_width): + if not self._platform.supports_framebuffers(): + self.delete() + self._create() + + self._platform.make_current() + self._renderer.viewport_width = self.viewport_width + self._renderer.viewport_height = self.viewport_height + self._renderer.point_size = self.point_size + + if self._platform.supports_framebuffers(): + flags |= RenderFlags.OFFSCREEN + retval = self._renderer.render(scene, flags, seg_node_map) + else: + self._renderer.render(scene, flags, seg_node_map) + depth = self._renderer.read_depth_buf() + if flags & RenderFlags.DEPTH_ONLY: + retval = depth + else: + color = self._renderer.read_color_buf() + retval = color, depth + + # Make the platform not current + self._platform.make_uncurrent() + return retval + + def delete(self): + """Free all OpenGL resources. + """ + self._platform.make_current() + self._renderer.delete() + self._platform.delete_context() + del self._renderer + del self._platform + self._renderer = None + self._platform = None + import gc + gc.collect() + + def _create(self): + if 'PYOPENGL_PLATFORM' not in os.environ: + from pyrender.platforms.pyglet_platform import PygletPlatform + self._platform = PygletPlatform(self.viewport_width, + self.viewport_height) + elif os.environ['PYOPENGL_PLATFORM'] == 'egl': + from pyrender.platforms import egl + device_id = int(os.environ.get('EGL_DEVICE_ID', '0')) + egl_device = egl.get_device_by_index(device_id) + self._platform = egl.EGLPlatform(self.viewport_width, + self.viewport_height, + device=egl_device) + elif os.environ['PYOPENGL_PLATFORM'] == 'osmesa': + from pyrender.platforms.osmesa import OSMesaPlatform + self._platform = OSMesaPlatform(self.viewport_width, + self.viewport_height) + else: + raise ValueError('Unsupported PyOpenGL platform: {}'.format( + os.environ['PYOPENGL_PLATFORM'] + )) + self._platform.init_context() + self._platform.make_current() + self._renderer = Renderer(self.viewport_width, self.viewport_height) + + def __del__(self): + try: + self.delete() + except Exception: + pass + + +__all__ = ['OffscreenRenderer'] diff --git a/pyrender/pyrender/platforms/__init__.py b/pyrender/pyrender/platforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7837fd5fdeccab5e48c85e41d20b238ea7396599 --- /dev/null +++ b/pyrender/pyrender/platforms/__init__.py @@ -0,0 +1,6 @@ +"""Platforms for generating offscreen OpenGL contexts for rendering. + +Author: Matthew Matl +""" + +from .base import Platform diff --git a/pyrender/pyrender/platforms/base.py b/pyrender/pyrender/platforms/base.py new file mode 100644 index 0000000000000000000000000000000000000000..c9ecda906145e239737901809aa59db8d3e231c6 --- /dev/null +++ b/pyrender/pyrender/platforms/base.py @@ -0,0 +1,76 @@ +import abc + +import six + + +@six.add_metaclass(abc.ABCMeta) +class Platform(object): + """Base class for all OpenGL platforms. + + Parameters + ---------- + viewport_width : int + The width of the main viewport, in pixels. + viewport_height : int + The height of the main viewport, in pixels + """ + + def __init__(self, viewport_width, viewport_height): + self.viewport_width = viewport_width + self.viewport_height = viewport_height + + @property + def viewport_width(self): + """int : The width of the main viewport, in pixels. + """ + return self._viewport_width + + @viewport_width.setter + def viewport_width(self, value): + self._viewport_width = value + + @property + def viewport_height(self): + """int : The height of the main viewport, in pixels. + """ + return self._viewport_height + + @viewport_height.setter + def viewport_height(self, value): + self._viewport_height = value + + @abc.abstractmethod + def init_context(self): + """Create an OpenGL context. + """ + pass + + @abc.abstractmethod + def make_current(self): + """Make the OpenGL context current. + """ + pass + + @abc.abstractmethod + def make_uncurrent(self): + """Make the OpenGL context uncurrent. + """ + pass + + @abc.abstractmethod + def delete_context(self): + """Delete the OpenGL context. + """ + pass + + @abc.abstractmethod + def supports_framebuffers(self): + """Returns True if the method supports framebuffer rendering. + """ + pass + + def __del__(self): + try: + self.delete_context() + except Exception: + pass diff --git a/pyrender/pyrender/platforms/egl.py b/pyrender/pyrender/platforms/egl.py new file mode 100644 index 0000000000000000000000000000000000000000..ae2478d29c9a538c53ad83fa31f8e2277cd897c8 --- /dev/null +++ b/pyrender/pyrender/platforms/egl.py @@ -0,0 +1,219 @@ +import ctypes +import os + +import OpenGL.platform + +from .base import Platform + +EGL_PLATFORM_DEVICE_EXT = 0x313F +EGL_DRM_DEVICE_FILE_EXT = 0x3233 + + +def _ensure_egl_loaded(): + plugin = OpenGL.platform.PlatformPlugin.by_name('egl') + if plugin is None: + raise RuntimeError("EGL platform plugin is not available.") + + plugin_class = plugin.load() + plugin.loaded = True + # create instance of this platform implementation + plugin = plugin_class() + + plugin.install(vars(OpenGL.platform)) + + +_ensure_egl_loaded() +from OpenGL import EGL as egl + + +def _get_egl_func(func_name, res_type, *arg_types): + address = egl.eglGetProcAddress(func_name) + if address is None: + return None + + proto = ctypes.CFUNCTYPE(res_type) + proto.argtypes = arg_types + func = proto(address) + return func + + +def _get_egl_struct(struct_name): + from OpenGL._opaque import opaque_pointer_cls + return opaque_pointer_cls(struct_name) + + +# These are not defined in PyOpenGL by default. +_EGLDeviceEXT = _get_egl_struct('EGLDeviceEXT') +_eglGetPlatformDisplayEXT = _get_egl_func('eglGetPlatformDisplayEXT', egl.EGLDisplay) +_eglQueryDevicesEXT = _get_egl_func('eglQueryDevicesEXT', egl.EGLBoolean) +_eglQueryDeviceStringEXT = _get_egl_func('eglQueryDeviceStringEXT', ctypes.c_char_p) + + +def query_devices(): + if _eglQueryDevicesEXT is None: + raise RuntimeError("EGL query extension is not loaded or is not supported.") + + num_devices = egl.EGLint() + success = _eglQueryDevicesEXT(0, None, ctypes.pointer(num_devices)) + if not success or num_devices.value < 1: + return [] + + devices = (_EGLDeviceEXT * num_devices.value)() # array of size num_devices + success = _eglQueryDevicesEXT(num_devices.value, devices, ctypes.pointer(num_devices)) + if not success or num_devices.value < 1: + return [] + + return [EGLDevice(devices[i]) for i in range(num_devices.value)] + + +def get_default_device(): + # Fall back to not using query extension. + if _eglQueryDevicesEXT is None: + return EGLDevice(None) + + return query_devices()[0] + + +def get_device_by_index(device_id): + if _eglQueryDevicesEXT is None and device_id == 0: + return get_default_device() + + devices = query_devices() + if device_id >= len(devices): + raise ValueError('Invalid device ID ({})'.format(device_id, len(devices))) + return devices[device_id] + + +class EGLDevice: + + def __init__(self, display=None): + self._display = display + + def get_display(self): + if self._display is None: + return egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY) + + return _eglGetPlatformDisplayEXT(EGL_PLATFORM_DEVICE_EXT, self._display, None) + + @property + def name(self): + if self._display is None: + return 'default' + + name = _eglQueryDeviceStringEXT(self._display, EGL_DRM_DEVICE_FILE_EXT) + if name is None: + return None + + return name.decode('ascii') + + def __repr__(self): + return "".format(self.name) + + +class EGLPlatform(Platform): + """Renders using EGL. + """ + + def __init__(self, viewport_width, viewport_height, device: EGLDevice = None): + super(EGLPlatform, self).__init__(viewport_width, viewport_height) + if device is None: + device = get_default_device() + + self._egl_device = device + self._egl_display = None + self._egl_context = None + + def init_context(self): + _ensure_egl_loaded() + + from OpenGL.EGL import ( + EGL_SURFACE_TYPE, EGL_PBUFFER_BIT, EGL_BLUE_SIZE, + EGL_RED_SIZE, EGL_GREEN_SIZE, EGL_DEPTH_SIZE, + EGL_COLOR_BUFFER_TYPE, EGL_RGB_BUFFER, + EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT, EGL_CONFORMANT, + EGL_NONE, EGL_DEFAULT_DISPLAY, EGL_NO_CONTEXT, + EGL_OPENGL_API, EGL_CONTEXT_MAJOR_VERSION, + EGL_CONTEXT_MINOR_VERSION, + EGL_CONTEXT_OPENGL_PROFILE_MASK, + EGL_CONTEXT_OPENGL_CORE_PROFILE_BIT, + eglGetDisplay, eglInitialize, eglChooseConfig, + eglBindAPI, eglCreateContext, EGLConfig + ) + from OpenGL import arrays + + config_attributes = arrays.GLintArray.asArray([ + EGL_SURFACE_TYPE, EGL_PBUFFER_BIT, + EGL_BLUE_SIZE, 8, + EGL_RED_SIZE, 8, + EGL_GREEN_SIZE, 8, + EGL_DEPTH_SIZE, 24, + EGL_COLOR_BUFFER_TYPE, EGL_RGB_BUFFER, + EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT, + EGL_CONFORMANT, EGL_OPENGL_BIT, + EGL_NONE + ]) + context_attributes = arrays.GLintArray.asArray([ + EGL_CONTEXT_MAJOR_VERSION, 4, + EGL_CONTEXT_MINOR_VERSION, 1, + EGL_CONTEXT_OPENGL_PROFILE_MASK, + EGL_CONTEXT_OPENGL_CORE_PROFILE_BIT, + EGL_NONE + ]) + major, minor = ctypes.c_long(), ctypes.c_long() + num_configs = ctypes.c_long() + configs = (EGLConfig * 1)() + + # Cache DISPLAY if necessary and get an off-screen EGL display + orig_dpy = None + if 'DISPLAY' in os.environ: + orig_dpy = os.environ['DISPLAY'] + del os.environ['DISPLAY'] + + self._egl_display = self._egl_device.get_display() + if orig_dpy is not None: + os.environ['DISPLAY'] = orig_dpy + + # Initialize EGL + assert eglInitialize(self._egl_display, major, minor) + assert eglChooseConfig( + self._egl_display, config_attributes, configs, 1, num_configs + ) + + # Bind EGL to the OpenGL API + assert eglBindAPI(EGL_OPENGL_API) + + # Create an EGL context + self._egl_context = eglCreateContext( + self._egl_display, configs[0], + EGL_NO_CONTEXT, context_attributes + ) + + # Make it current + self.make_current() + + def make_current(self): + from OpenGL.EGL import eglMakeCurrent, EGL_NO_SURFACE + assert eglMakeCurrent( + self._egl_display, EGL_NO_SURFACE, EGL_NO_SURFACE, + self._egl_context + ) + + def make_uncurrent(self): + """Make the OpenGL context uncurrent. + """ + pass + + def delete_context(self): + from OpenGL.EGL import eglDestroyContext, eglTerminate + if self._egl_display is not None: + if self._egl_context is not None: + eglDestroyContext(self._egl_display, self._egl_context) + self._egl_context = None + eglTerminate(self._egl_display) + self._egl_display = None + + def supports_framebuffers(self): + return True + + +__all__ = ['EGLPlatform'] diff --git a/pyrender/pyrender/platforms/osmesa.py b/pyrender/pyrender/platforms/osmesa.py new file mode 100644 index 0000000000000000000000000000000000000000..deaa5ff44031a107883913ae9a18fc425d650f3d --- /dev/null +++ b/pyrender/pyrender/platforms/osmesa.py @@ -0,0 +1,59 @@ +from .base import Platform + + +__all__ = ['OSMesaPlatform'] + + +class OSMesaPlatform(Platform): + """Renders into a software buffer using OSMesa. Requires special versions + of OSMesa to be installed, plus PyOpenGL upgrade. + """ + + def __init__(self, viewport_width, viewport_height): + super(OSMesaPlatform, self).__init__(viewport_width, viewport_height) + self._context = None + self._buffer = None + + def init_context(self): + from OpenGL import arrays + from OpenGL.osmesa import ( + OSMesaCreateContextAttribs, OSMESA_FORMAT, + OSMESA_RGBA, OSMESA_PROFILE, OSMESA_CORE_PROFILE, + OSMESA_CONTEXT_MAJOR_VERSION, OSMESA_CONTEXT_MINOR_VERSION, + OSMESA_DEPTH_BITS + ) + + attrs = arrays.GLintArray.asArray([ + OSMESA_FORMAT, OSMESA_RGBA, + OSMESA_DEPTH_BITS, 24, + OSMESA_PROFILE, OSMESA_CORE_PROFILE, + OSMESA_CONTEXT_MAJOR_VERSION, 3, + OSMESA_CONTEXT_MINOR_VERSION, 3, + 0 + ]) + self._context = OSMesaCreateContextAttribs(attrs, None) + self._buffer = arrays.GLubyteArray.zeros( + (self.viewport_height, self.viewport_width, 4) + ) + + def make_current(self): + from OpenGL import GL as gl + from OpenGL.osmesa import OSMesaMakeCurrent + assert(OSMesaMakeCurrent( + self._context, self._buffer, gl.GL_UNSIGNED_BYTE, + self.viewport_width, self.viewport_height + )) + + def make_uncurrent(self): + """Make the OpenGL context uncurrent. + """ + pass + + def delete_context(self): + from OpenGL.osmesa import OSMesaDestroyContext + OSMesaDestroyContext(self._context) + self._context = None + self._buffer = None + + def supports_framebuffers(self): + return False diff --git a/pyrender/pyrender/platforms/pyglet_platform.py b/pyrender/pyrender/platforms/pyglet_platform.py new file mode 100644 index 0000000000000000000000000000000000000000..a70cf7b659bc85a92f6c9c8ebcc360662a068507 --- /dev/null +++ b/pyrender/pyrender/platforms/pyglet_platform.py @@ -0,0 +1,90 @@ +from pyrender.constants import (TARGET_OPEN_GL_MAJOR, TARGET_OPEN_GL_MINOR, + MIN_OPEN_GL_MAJOR, MIN_OPEN_GL_MINOR) +from .base import Platform + +import OpenGL + + +__all__ = ['PygletPlatform'] + + +class PygletPlatform(Platform): + """Renders on-screen using a 1x1 hidden Pyglet window for getting + an OpenGL context. + """ + + def __init__(self, viewport_width, viewport_height): + super(PygletPlatform, self).__init__(viewport_width, viewport_height) + self._window = None + + def init_context(self): + import pyglet + pyglet.options['shadow_window'] = False + + try: + pyglet.lib.x11.xlib.XInitThreads() + except Exception: + pass + + self._window = None + confs = [pyglet.gl.Config(sample_buffers=1, samples=4, + depth_size=24, + double_buffer=True, + major_version=TARGET_OPEN_GL_MAJOR, + minor_version=TARGET_OPEN_GL_MINOR), + pyglet.gl.Config(depth_size=24, + double_buffer=True, + major_version=TARGET_OPEN_GL_MAJOR, + minor_version=TARGET_OPEN_GL_MINOR), + pyglet.gl.Config(sample_buffers=1, samples=4, + depth_size=24, + double_buffer=True, + major_version=MIN_OPEN_GL_MAJOR, + minor_version=MIN_OPEN_GL_MINOR), + pyglet.gl.Config(depth_size=24, + double_buffer=True, + major_version=MIN_OPEN_GL_MAJOR, + minor_version=MIN_OPEN_GL_MINOR)] + for conf in confs: + try: + self._window = pyglet.window.Window(config=conf, visible=False, + resizable=False, + width=1, height=1) + break + except pyglet.window.NoSuchConfigException as e: + pass + + if not self._window: + raise ValueError( + 'Failed to initialize Pyglet window with an OpenGL >= 3+ ' + 'context. If you\'re logged in via SSH, ensure that you\'re ' + 'running your script with vglrun (i.e. VirtualGL). The ' + 'internal error message was "{}"'.format(e) + ) + + def make_current(self): + if self._window: + self._window.switch_to() + + def make_uncurrent(self): + try: + import pyglet + pyglet.gl.xlib.glx.glXMakeContextCurrent(self._window.context.x_display, 0, 0, None) + except Exception: + pass + + def delete_context(self): + if self._window is not None: + self.make_current() + cid = OpenGL.contextdata.getContext() + try: + self._window.context.destroy() + self._window.close() + except Exception: + pass + self._window = None + OpenGL.contextdata.cleanupContext(cid) + del cid + + def supports_framebuffers(self): + return True diff --git a/pyrender/pyrender/primitive.py b/pyrender/pyrender/primitive.py new file mode 100644 index 0000000000000000000000000000000000000000..7f83f46f532b126a4573e715dd03d079fef755ca --- /dev/null +++ b/pyrender/pyrender/primitive.py @@ -0,0 +1,489 @@ +"""Primitives, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-primitive + +Author: Matthew Matl +""" +import numpy as np + +from OpenGL.GL import * + +from .material import Material, MetallicRoughnessMaterial +from .constants import FLOAT_SZ, UINT_SZ, BufFlags, GLTF +from .utils import format_color_array + + +class Primitive(object): + """A primitive object which can be rendered. + + Parameters + ---------- + positions : (n, 3) float + XYZ vertex positions. + normals : (n, 3) float + Normalized XYZ vertex normals. + tangents : (n, 4) float + XYZW vertex tangents where the w component is a sign value + (either +1 or -1) indicating the handedness of the tangent basis. + texcoord_0 : (n, 2) float + The first set of UV texture coordinates. + texcoord_1 : (n, 2) float + The second set of UV texture coordinates. + color_0 : (n, 4) float + RGBA vertex colors. + joints_0 : (n, 4) float + Joint information. + weights_0 : (n, 4) float + Weight information for morphing. + indices : (m, 3) int + Face indices for triangle meshes or fans. + material : :class:`Material` + The material to apply to this primitive when rendering. + mode : int + The type of primitives to render, one of the following: + + - ``0``: POINTS + - ``1``: LINES + - ``2``: LINE_LOOP + - ``3``: LINE_STRIP + - ``4``: TRIANGLES + - ``5``: TRIANGLES_STRIP + - ``6``: TRIANGLES_FAN + targets : (k,) int + Morph target indices. + poses : (x,4,4), float + Array of 4x4 transformation matrices for instancing this object. + """ + + def __init__(self, + positions, + normals=None, + tangents=None, + texcoord_0=None, + texcoord_1=None, + color_0=None, + joints_0=None, + weights_0=None, + indices=None, + material=None, + mode=None, + targets=None, + poses=None): + + if mode is None: + mode = GLTF.TRIANGLES + + self.positions = positions + self.normals = normals + self.tangents = tangents + self.texcoord_0 = texcoord_0 + self.texcoord_1 = texcoord_1 + self.color_0 = color_0 + self.joints_0 = joints_0 + self.weights_0 = weights_0 + self.indices = indices + self.material = material + self.mode = mode + self.targets = targets + self.poses = poses + + self._bounds = None + self._vaid = None + self._buffers = [] + self._is_transparent = None + self._buf_flags = None + + @property + def positions(self): + """(n,3) float : XYZ vertex positions. + """ + return self._positions + + @positions.setter + def positions(self, value): + value = np.asanyarray(value, dtype=np.float32) + self._positions = np.ascontiguousarray(value) + self._bounds = None + + @property + def normals(self): + """(n,3) float : Normalized XYZ vertex normals. + """ + return self._normals + + @normals.setter + def normals(self, value): + if value is not None: + value = np.asanyarray(value, dtype=np.float32) + value = np.ascontiguousarray(value) + if value.shape != self.positions.shape: + raise ValueError('Incorrect normals shape') + self._normals = value + + @property + def tangents(self): + """(n,4) float : XYZW vertex tangents. + """ + return self._tangents + + @tangents.setter + def tangents(self, value): + if value is not None: + value = np.asanyarray(value, dtype=np.float32) + value = np.ascontiguousarray(value) + if value.shape != (self.positions.shape[0], 4): + raise ValueError('Incorrect tangent shape') + self._tangents = value + + @property + def texcoord_0(self): + """(n,2) float : The first set of UV texture coordinates. + """ + return self._texcoord_0 + + @texcoord_0.setter + def texcoord_0(self, value): + if value is not None: + value = np.asanyarray(value, dtype=np.float32) + value = np.ascontiguousarray(value) + if (value.ndim != 2 or value.shape[0] != self.positions.shape[0] or + value.shape[1] < 2): + raise ValueError('Incorrect texture coordinate shape') + if value.shape[1] > 2: + value = value[:,:2] + self._texcoord_0 = value + + @property + def texcoord_1(self): + """(n,2) float : The second set of UV texture coordinates. + """ + return self._texcoord_1 + + @texcoord_1.setter + def texcoord_1(self, value): + if value is not None: + value = np.asanyarray(value, dtype=np.float32) + value = np.ascontiguousarray(value) + if (value.ndim != 2 or value.shape[0] != self.positions.shape[0] or + value.shape[1] != 2): + raise ValueError('Incorrect texture coordinate shape') + self._texcoord_1 = value + + @property + def color_0(self): + """(n,4) float : RGBA vertex colors. + """ + return self._color_0 + + @color_0.setter + def color_0(self, value): + if value is not None: + value = np.ascontiguousarray( + format_color_array(value, shape=(len(self.positions), 4)) + ) + self._is_transparent = None + self._color_0 = value + + @property + def joints_0(self): + """(n,4) float : Joint information. + """ + return self._joints_0 + + @joints_0.setter + def joints_0(self, value): + self._joints_0 = value + + @property + def weights_0(self): + """(n,4) float : Weight information for morphing. + """ + return self._weights_0 + + @weights_0.setter + def weights_0(self, value): + self._weights_0 = value + + @property + def indices(self): + """(m,3) int : Face indices for triangle meshes or fans. + """ + return self._indices + + @indices.setter + def indices(self, value): + if value is not None: + value = np.asanyarray(value, dtype=np.float32) + value = np.ascontiguousarray(value) + self._indices = value + + @property + def material(self): + """:class:`Material` : The material for this primitive. + """ + return self._material + + @material.setter + def material(self, value): + # Create default material + if value is None: + value = MetallicRoughnessMaterial() + else: + if not isinstance(value, Material): + raise TypeError('Object material must be of type Material') + self._material = value + + @property + def mode(self): + """int : The type of primitive to render. + """ + return self._mode + + @mode.setter + def mode(self, value): + value = int(value) + if value < GLTF.POINTS or value > GLTF.TRIANGLE_FAN: + raise ValueError('Invalid mode') + self._mode = value + + @property + def targets(self): + """(k,) int : Morph target indices. + """ + return self._targets + + @targets.setter + def targets(self, value): + self._targets = value + + @property + def poses(self): + """(x,4,4) float : Homogenous transforms for instancing this primitive. + """ + return self._poses + + @poses.setter + def poses(self, value): + if value is not None: + value = np.asanyarray(value, dtype=np.float32) + value = np.ascontiguousarray(value) + if value.ndim == 2: + value = value[np.newaxis,:,:] + if value.shape[1] != 4 or value.shape[2] != 4: + raise ValueError('Pose matrices must be of shape (n,4,4), ' + 'got {}'.format(value.shape)) + self._poses = value + self._bounds = None + + @property + def bounds(self): + if self._bounds is None: + self._bounds = self._compute_bounds() + return self._bounds + + @property + def centroid(self): + """(3,) float : The centroid of the primitive's AABB. + """ + return np.mean(self.bounds, axis=0) + + @property + def extents(self): + """(3,) float : The lengths of the axes of the primitive's AABB. + """ + return np.diff(self.bounds, axis=0).reshape(-1) + + @property + def scale(self): + """(3,) float : The length of the diagonal of the primitive's AABB. + """ + return np.linalg.norm(self.extents) + + @property + def buf_flags(self): + """int : The flags for the render buffer. + """ + if self._buf_flags is None: + self._buf_flags = self._compute_buf_flags() + return self._buf_flags + + def delete(self): + self._unbind() + self._remove_from_context() + + @property + def is_transparent(self): + """bool : If True, the mesh is partially-transparent. + """ + return self._compute_transparency() + + def _add_to_context(self): + if self._vaid is not None: + raise ValueError('Mesh is already bound to a context') + + # Generate and bind VAO + self._vaid = glGenVertexArrays(1) + glBindVertexArray(self._vaid) + + ####################################################################### + # Fill vertex buffer + ####################################################################### + + # Generate and bind vertex buffer + vertexbuffer = glGenBuffers(1) + self._buffers.append(vertexbuffer) + glBindBuffer(GL_ARRAY_BUFFER, vertexbuffer) + + # positions + vertex_data = self.positions + attr_sizes = [3] + + # Normals + if self.normals is not None: + vertex_data = np.hstack((vertex_data, self.normals)) + attr_sizes.append(3) + + # Tangents + if self.tangents is not None: + vertex_data = np.hstack((vertex_data, self.tangents)) + attr_sizes.append(4) + + # Texture Coordinates + if self.texcoord_0 is not None: + vertex_data = np.hstack((vertex_data, self.texcoord_0)) + attr_sizes.append(2) + if self.texcoord_1 is not None: + vertex_data = np.hstack((vertex_data, self.texcoord_1)) + attr_sizes.append(2) + + # Color + if self.color_0 is not None: + vertex_data = np.hstack((vertex_data, self.color_0)) + attr_sizes.append(4) + + # TODO JOINTS AND WEIGHTS + # PASS + + # Copy data to buffer + vertex_data = np.ascontiguousarray( + vertex_data.flatten().astype(np.float32) + ) + glBufferData( + GL_ARRAY_BUFFER, FLOAT_SZ * len(vertex_data), + vertex_data, GL_STATIC_DRAW + ) + total_sz = sum(attr_sizes) + offset = 0 + for i, sz in enumerate(attr_sizes): + glVertexAttribPointer( + i, sz, GL_FLOAT, GL_FALSE, FLOAT_SZ * total_sz, + ctypes.c_void_p(FLOAT_SZ * offset) + ) + glEnableVertexAttribArray(i) + offset += sz + + ####################################################################### + # Fill model matrix buffer + ####################################################################### + + if self.poses is not None: + pose_data = np.ascontiguousarray( + np.transpose(self.poses, [0,2,1]).flatten().astype(np.float32) + ) + else: + pose_data = np.ascontiguousarray( + np.eye(4).flatten().astype(np.float32) + ) + + modelbuffer = glGenBuffers(1) + self._buffers.append(modelbuffer) + glBindBuffer(GL_ARRAY_BUFFER, modelbuffer) + glBufferData( + GL_ARRAY_BUFFER, FLOAT_SZ * len(pose_data), + pose_data, GL_STATIC_DRAW + ) + + for i in range(0, 4): + idx = i + len(attr_sizes) + glEnableVertexAttribArray(idx) + glVertexAttribPointer( + idx, 4, GL_FLOAT, GL_FALSE, FLOAT_SZ * 4 * 4, + ctypes.c_void_p(4 * FLOAT_SZ * i) + ) + glVertexAttribDivisor(idx, 1) + + ####################################################################### + # Fill element buffer + ####################################################################### + if self.indices is not None: + elementbuffer = glGenBuffers(1) + self._buffers.append(elementbuffer) + glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, elementbuffer) + glBufferData(GL_ELEMENT_ARRAY_BUFFER, UINT_SZ * self.indices.size, + self.indices.flatten().astype(np.uint32), + GL_STATIC_DRAW) + + glBindVertexArray(0) + + def _remove_from_context(self): + if self._vaid is not None: + glDeleteVertexArrays(1, [self._vaid]) + glDeleteBuffers(len(self._buffers), self._buffers) + self._vaid = None + self._buffers = [] + + def _in_context(self): + return self._vaid is not None + + def _bind(self): + if self._vaid is None: + raise ValueError('Cannot bind a Mesh that has not been added ' + 'to a context') + glBindVertexArray(self._vaid) + + def _unbind(self): + glBindVertexArray(0) + + def _compute_bounds(self): + """Compute the bounds of this object. + """ + # Compute bounds of this object + bounds = np.array([np.min(self.positions, axis=0), + np.max(self.positions, axis=0)]) + + # If instanced, compute translations for approximate bounds + if self.poses is not None: + bounds += np.array([np.min(self.poses[:,:3,3], axis=0), + np.max(self.poses[:,:3,3], axis=0)]) + return bounds + + def _compute_transparency(self): + """Compute whether or not this object is transparent. + """ + if self.material.is_transparent: + return True + if self._is_transparent is None: + self._is_transparent = False + if self.color_0 is not None: + if np.any(self._color_0[:,3] != 1.0): + self._is_transparent = True + return self._is_transparent + + def _compute_buf_flags(self): + buf_flags = BufFlags.POSITION + + if self.normals is not None: + buf_flags |= BufFlags.NORMAL + if self.tangents is not None: + buf_flags |= BufFlags.TANGENT + if self.texcoord_0 is not None: + buf_flags |= BufFlags.TEXCOORD_0 + if self.texcoord_1 is not None: + buf_flags |= BufFlags.TEXCOORD_1 + if self.color_0 is not None: + buf_flags |= BufFlags.COLOR_0 + if self.joints_0 is not None: + buf_flags |= BufFlags.JOINTS_0 + if self.weights_0 is not None: + buf_flags |= BufFlags.WEIGHTS_0 + + return buf_flags diff --git a/pyrender/pyrender/renderer.py b/pyrender/pyrender/renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..5ae14c5cdb1785226a52ae6b71b08f01de069962 --- /dev/null +++ b/pyrender/pyrender/renderer.py @@ -0,0 +1,1339 @@ +"""PBR renderer for Python. + +Author: Matthew Matl +""" +import sys + +import numpy as np +import PIL + +from .constants import (RenderFlags, TextAlign, GLTF, BufFlags, TexFlags, + ProgramFlags, DEFAULT_Z_FAR, DEFAULT_Z_NEAR, + SHADOW_TEX_SZ, MAX_N_LIGHTS) +from .shader_program import ShaderProgramCache +from .material import MetallicRoughnessMaterial, SpecularGlossinessMaterial +from .light import PointLight, SpotLight, DirectionalLight +from .font import FontCache +from .utils import format_color_vector + +from OpenGL.GL import * + + +class Renderer(object): + """Class for handling all rendering operations on a scene. + + Note + ---- + This renderer relies on the existence of an OpenGL context and + does not create one on its own. + + Parameters + ---------- + viewport_width : int + Width of the viewport in pixels. + viewport_height : int + Width of the viewport height in pixels. + point_size : float, optional + Size of points in pixels. Defaults to 1.0. + """ + + def __init__(self, viewport_width, viewport_height, point_size=1.0): + self.dpscale = 1 + # Scaling needed on retina displays + if sys.platform == 'darwin': + self.dpscale = 2 + + self.viewport_width = viewport_width + self.viewport_height = viewport_height + self.point_size = point_size + + # Optional framebuffer for offscreen renders + self._main_fb = None + self._main_cb = None + self._main_db = None + self._main_fb_ms = None + self._main_cb_ms = None + self._main_db_ms = None + self._main_fb_dims = (None, None) + self._shadow_fb = None + self._latest_znear = DEFAULT_Z_NEAR + self._latest_zfar = DEFAULT_Z_FAR + + # Shader Program Cache + self._program_cache = ShaderProgramCache() + self._font_cache = FontCache() + self._meshes = set() + self._mesh_textures = set() + self._shadow_textures = set() + self._texture_alloc_idx = 0 + + @property + def viewport_width(self): + """int : The width of the main viewport, in pixels. + """ + return self._viewport_width + + @viewport_width.setter + def viewport_width(self, value): + self._viewport_width = self.dpscale * value + + @property + def viewport_height(self): + """int : The height of the main viewport, in pixels. + """ + return self._viewport_height + + @viewport_height.setter + def viewport_height(self, value): + self._viewport_height = self.dpscale * value + + @property + def point_size(self): + """float : The size of screen-space points, in pixels. + """ + return self._point_size + + @point_size.setter + def point_size(self, value): + self._point_size = float(value) + + def render(self, scene, flags, seg_node_map=None): + """Render a scene with the given set of flags. + + Parameters + ---------- + scene : :class:`Scene` + A scene to render. + flags : int + A specification from :class:`.RenderFlags`. + seg_node_map : dict + A map from :class:`.Node` objects to (3,) colors for each. + If specified along with flags set to :attr:`.RenderFlags.SEG`, + the color image will be a segmentation image. + + Returns + ------- + color_im : (h, w, 3) uint8 or (h, w, 4) uint8 + If :attr:`RenderFlags.OFFSCREEN` is set, the color buffer. This is + normally an RGB buffer, but if :attr:`.RenderFlags.RGBA` is set, + the buffer will be a full RGBA buffer. + depth_im : (h, w) float32 + If :attr:`RenderFlags.OFFSCREEN` is set, the depth buffer + in linear units. + """ + # Update context with meshes and textures + self._update_context(scene, flags) + + # Render necessary shadow maps + if not bool(flags & RenderFlags.DEPTH_ONLY or flags & RenderFlags.SEG): + for ln in scene.light_nodes: + take_pass = False + if (isinstance(ln.light, DirectionalLight) and + bool(flags & RenderFlags.SHADOWS_DIRECTIONAL)): + take_pass = True + elif (isinstance(ln.light, SpotLight) and + bool(flags & RenderFlags.SHADOWS_SPOT)): + take_pass = True + elif (isinstance(ln.light, PointLight) and + bool(flags & RenderFlags.SHADOWS_POINT)): + take_pass = True + if take_pass: + self._shadow_mapping_pass(scene, ln, flags) + + # Make forward pass + retval = self._forward_pass(scene, flags, seg_node_map=seg_node_map) + + # If necessary, make normals pass + if flags & (RenderFlags.VERTEX_NORMALS | RenderFlags.FACE_NORMALS): + self._normals_pass(scene, flags) + + # Update camera settings for retrieving depth buffers + self._latest_znear = scene.main_camera_node.camera.znear + self._latest_zfar = scene.main_camera_node.camera.zfar + + return retval + + def render_text(self, text, x, y, font_name='OpenSans-Regular', + font_pt=40, color=None, scale=1.0, + align=TextAlign.BOTTOM_LEFT): + """Render text into the current viewport. + + Note + ---- + This cannot be done into an offscreen buffer. + + Parameters + ---------- + text : str + The text to render. + x : int + Horizontal pixel location of text. + y : int + Vertical pixel location of text. + font_name : str + Name of font, from the ``pyrender/fonts`` folder, or + a path to a ``.ttf`` file. + font_pt : int + Height of the text, in font points. + color : (4,) float + The color of the text. Default is black. + scale : int + Scaling factor for text. + align : int + One of the :class:`TextAlign` options which specifies where the + ``x`` and ``y`` parameters lie on the text. For example, + :attr:`TextAlign.BOTTOM_LEFT` means that ``x`` and ``y`` indicate + the position of the bottom-left corner of the textbox. + """ + x *= self.dpscale + y *= self.dpscale + font_pt *= self.dpscale + + if color is None: + color = np.array([0.0, 0.0, 0.0, 1.0]) + else: + color = format_color_vector(color, 4) + + # Set up viewport for render + self._configure_forward_pass_viewport(0) + + # Load font + font = self._font_cache.get_font(font_name, font_pt) + if not font._in_context(): + font._add_to_context() + + # Load program + program = self._get_text_program() + program._bind() + + # Set uniforms + p = np.eye(4) + p[0,0] = 2.0 / self.viewport_width + p[0,3] = -1.0 + p[1,1] = 2.0 / self.viewport_height + p[1,3] = -1.0 + program.set_uniform('projection', p) + program.set_uniform('text_color', color) + + # Draw text + font.render_string(text, x, y, scale, align) + + def read_color_buf(self): + """Read and return the current viewport's color buffer. + + Alpha cannot be computed for an on-screen buffer. + + Returns + ------- + color_im : (h, w, 3) uint8 + The color buffer in RGB byte format. + """ + # Extract color image from frame buffer + width, height = self.viewport_width, self.viewport_height + glBindFramebuffer(GL_READ_FRAMEBUFFER, 0) + glReadBuffer(GL_FRONT) + color_buf = glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE) + + # Re-format them into numpy arrays + color_im = np.frombuffer(color_buf, dtype=np.uint8) + color_im = color_im.reshape((height, width, 3)) + color_im = np.flip(color_im, axis=0) + + # Resize for macos if needed + if sys.platform == 'darwin': + color_im = self._resize_image(color_im, True) + + return color_im + + def read_depth_buf(self): + """Read and return the current viewport's color buffer. + + Returns + ------- + depth_im : (h, w) float32 + The depth buffer in linear units. + """ + width, height = self.viewport_width, self.viewport_height + glBindFramebuffer(GL_READ_FRAMEBUFFER, 0) + glReadBuffer(GL_FRONT) + depth_buf = glReadPixels( + 0, 0, width, height, GL_DEPTH_COMPONENT, GL_FLOAT + ) + + depth_im = np.frombuffer(depth_buf, dtype=np.float32) + depth_im = depth_im.reshape((height, width)) + depth_im = np.flip(depth_im, axis=0) + + inf_inds = (depth_im == 1.0) + depth_im = 2.0 * depth_im - 1.0 + z_near, z_far = self._latest_znear, self._latest_zfar + noninf = np.logical_not(inf_inds) + if z_far is None: + depth_im[noninf] = 2 * z_near / (1.0 - depth_im[noninf]) + else: + depth_im[noninf] = ((2.0 * z_near * z_far) / + (z_far + z_near - depth_im[noninf] * + (z_far - z_near))) + depth_im[inf_inds] = 0.0 + + # Resize for macos if needed + if sys.platform == 'darwin': + depth_im = self._resize_image(depth_im) + + return depth_im + + def delete(self): + """Free all allocated OpenGL resources. + """ + # Free shaders + self._program_cache.clear() + + # Free fonts + self._font_cache.clear() + + # Free meshes + for mesh in self._meshes: + for p in mesh.primitives: + p.delete() + + # Free textures + for mesh_texture in self._mesh_textures: + mesh_texture.delete() + + for shadow_texture in self._shadow_textures: + shadow_texture.delete() + + self._meshes = set() + self._mesh_textures = set() + self._shadow_textures = set() + self._texture_alloc_idx = 0 + + self._delete_main_framebuffer() + self._delete_shadow_framebuffer() + + def __del__(self): + try: + self.delete() + except Exception: + pass + + ########################################################################### + # Rendering passes + ########################################################################### + + def _forward_pass(self, scene, flags, seg_node_map=None): + # Set up viewport for render + self._configure_forward_pass_viewport(flags) + + # Clear it + if bool(flags & RenderFlags.SEG): + glClearColor(0.0, 0.0, 0.0, 1.0) + if seg_node_map is None: + seg_node_map = {} + else: + glClearColor(*scene.bg_color) + + glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) + + if not bool(flags & RenderFlags.SEG): + glEnable(GL_MULTISAMPLE) + else: + glDisable(GL_MULTISAMPLE) + + # Set up camera matrices + V, P = self._get_camera_matrices(scene) + + program = None + # Now, render each object in sorted order + for node in self._sorted_mesh_nodes(scene): + mesh = node.mesh + + # Skip the mesh if it's not visible + if not mesh.is_visible: + continue + + # If SEG, set color + if bool(flags & RenderFlags.SEG): + if node not in seg_node_map: + continue + color = seg_node_map[node] + if not isinstance(color, (list, tuple, np.ndarray)): + color = np.repeat(color, 3) + else: + color = np.asanyarray(color) + color = color / 255.0 + + for primitive in mesh.primitives: + + # First, get and bind the appropriate program + program = self._get_primitive_program( + primitive, flags, ProgramFlags.USE_MATERIAL + ) + program._bind() + + # Set the camera uniforms + program.set_uniform('V', V) + program.set_uniform('P', P) + program.set_uniform( + 'cam_pos', scene.get_pose(scene.main_camera_node)[:3,3] + ) + if bool(flags & RenderFlags.SEG): + program.set_uniform('color', color) + + # Next, bind the lighting + if not (flags & RenderFlags.DEPTH_ONLY or flags & RenderFlags.FLAT or + flags & RenderFlags.SEG): + self._bind_lighting(scene, program, node, flags) + + # Finally, bind and draw the primitive + self._bind_and_draw_primitive( + primitive=primitive, + pose=scene.get_pose(node), + program=program, + flags=flags + ) + self._reset_active_textures() + + # Unbind the shader and flush the output + if program is not None: + program._unbind() + glFlush() + + # If doing offscreen render, copy result from framebuffer and return + if flags & RenderFlags.OFFSCREEN: + return self._read_main_framebuffer(scene, flags) + else: + return + + def _shadow_mapping_pass(self, scene, light_node, flags): + light = light_node.light + + # Set up viewport for render + self._configure_shadow_mapping_viewport(light, flags) + + # Set up camera matrices + V, P = self._get_light_cam_matrices(scene, light_node, flags) + + # Now, render each object in sorted order + for node in self._sorted_mesh_nodes(scene): + mesh = node.mesh + + # Skip the mesh if it's not visible + if not mesh.is_visible: + continue + + for primitive in mesh.primitives: + + # First, get and bind the appropriate program + program = self._get_primitive_program( + primitive, flags, ProgramFlags.NONE + ) + program._bind() + + # Set the camera uniforms + program.set_uniform('V', V) + program.set_uniform('P', P) + program.set_uniform( + 'cam_pos', scene.get_pose(scene.main_camera_node)[:3,3] + ) + + # Finally, bind and draw the primitive + self._bind_and_draw_primitive( + primitive=primitive, + pose=scene.get_pose(node), + program=program, + flags=RenderFlags.DEPTH_ONLY + ) + self._reset_active_textures() + + # Unbind the shader and flush the output + if program is not None: + program._unbind() + glFlush() + + def _normals_pass(self, scene, flags): + # Set up viewport for render + self._configure_forward_pass_viewport(flags) + program = None + + # Set up camera matrices + V, P = self._get_camera_matrices(scene) + + # Now, render each object in sorted order + for node in self._sorted_mesh_nodes(scene): + mesh = node.mesh + + # Skip the mesh if it's not visible + if not mesh.is_visible: + continue + + for primitive in mesh.primitives: + + # Skip objects that don't have normals + if not primitive.buf_flags & BufFlags.NORMAL: + continue + + # First, get and bind the appropriate program + pf = ProgramFlags.NONE + if flags & RenderFlags.VERTEX_NORMALS: + pf = pf | ProgramFlags.VERTEX_NORMALS + if flags & RenderFlags.FACE_NORMALS: + pf = pf | ProgramFlags.FACE_NORMALS + program = self._get_primitive_program(primitive, flags, pf) + program._bind() + + # Set the camera uniforms + program.set_uniform('V', V) + program.set_uniform('P', P) + program.set_uniform('normal_magnitude', 0.05 * primitive.scale) + program.set_uniform( + 'normal_color', np.array([0.1, 0.1, 1.0, 1.0]) + ) + + # Finally, bind and draw the primitive + self._bind_and_draw_primitive( + primitive=primitive, + pose=scene.get_pose(node), + program=program, + flags=RenderFlags.DEPTH_ONLY + ) + self._reset_active_textures() + + # Unbind the shader and flush the output + if program is not None: + program._unbind() + glFlush() + + ########################################################################### + # Handlers for binding uniforms and drawing primitives + ########################################################################### + + def _bind_and_draw_primitive(self, primitive, pose, program, flags): + # Set model pose matrix + program.set_uniform('M', pose) + + # Bind mesh buffers + primitive._bind() + + # Bind mesh material + if not (flags & RenderFlags.DEPTH_ONLY or flags & RenderFlags.SEG): + material = primitive.material + + # Bind textures + tf = material.tex_flags + if tf & TexFlags.NORMAL: + self._bind_texture(material.normalTexture, + 'material.normal_texture', program) + if tf & TexFlags.OCCLUSION: + self._bind_texture(material.occlusionTexture, + 'material.occlusion_texture', program) + if tf & TexFlags.EMISSIVE: + self._bind_texture(material.emissiveTexture, + 'material.emissive_texture', program) + if tf & TexFlags.BASE_COLOR: + self._bind_texture(material.baseColorTexture, + 'material.base_color_texture', program) + if tf & TexFlags.METALLIC_ROUGHNESS: + self._bind_texture(material.metallicRoughnessTexture, + 'material.metallic_roughness_texture', + program) + if tf & TexFlags.DIFFUSE: + self._bind_texture(material.diffuseTexture, + 'material.diffuse_texture', program) + if tf & TexFlags.SPECULAR_GLOSSINESS: + self._bind_texture(material.specularGlossinessTexture, + 'material.specular_glossiness_texture', + program) + + # Bind other uniforms + b = 'material.{}' + program.set_uniform(b.format('emissive_factor'), + material.emissiveFactor) + if isinstance(material, MetallicRoughnessMaterial): + program.set_uniform(b.format('base_color_factor'), + material.baseColorFactor) + program.set_uniform(b.format('metallic_factor'), + material.metallicFactor) + program.set_uniform(b.format('roughness_factor'), + material.roughnessFactor) + elif isinstance(material, SpecularGlossinessMaterial): + program.set_uniform(b.format('diffuse_factor'), + material.diffuseFactor) + program.set_uniform(b.format('specular_factor'), + material.specularFactor) + program.set_uniform(b.format('glossiness_factor'), + material.glossinessFactor) + + # Set blending options + if material.alphaMode == 'BLEND': + glEnable(GL_BLEND) + glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA) + else: + glEnable(GL_BLEND) + glBlendFunc(GL_ONE, GL_ZERO) + + # Set wireframe mode + wf = material.wireframe + if flags & RenderFlags.FLIP_WIREFRAME: + wf = not wf + if (flags & RenderFlags.ALL_WIREFRAME) or wf: + glPolygonMode(GL_FRONT_AND_BACK, GL_LINE) + else: + glPolygonMode(GL_FRONT_AND_BACK, GL_FILL) + + # Set culling mode + if material.doubleSided or flags & RenderFlags.SKIP_CULL_FACES: + glDisable(GL_CULL_FACE) + else: + glEnable(GL_CULL_FACE) + glCullFace(GL_BACK) + else: + glEnable(GL_CULL_FACE) + glEnable(GL_BLEND) + glCullFace(GL_BACK) + glBlendFunc(GL_ONE, GL_ZERO) + glPolygonMode(GL_FRONT_AND_BACK, GL_FILL) + + # Set point size if needed + glDisable(GL_PROGRAM_POINT_SIZE) + if primitive.mode == GLTF.POINTS: + glEnable(GL_PROGRAM_POINT_SIZE) + glPointSize(self.point_size) + + # Render mesh + n_instances = 1 + if primitive.poses is not None: + n_instances = len(primitive.poses) + + if primitive.indices is not None: + glDrawElementsInstanced( + primitive.mode, primitive.indices.size, GL_UNSIGNED_INT, + ctypes.c_void_p(0), n_instances + ) + else: + glDrawArraysInstanced( + primitive.mode, 0, len(primitive.positions), n_instances + ) + + # Unbind mesh buffers + primitive._unbind() + + def _bind_lighting(self, scene, program, node, flags): + """Bind all lighting uniform values for a scene. + """ + max_n_lights = self._compute_max_n_lights(flags) + + n_d = min(len(scene.directional_light_nodes), max_n_lights[0]) + n_s = min(len(scene.spot_light_nodes), max_n_lights[1]) + n_p = min(len(scene.point_light_nodes), max_n_lights[2]) + program.set_uniform('ambient_light', scene.ambient_light) + program.set_uniform('n_directional_lights', n_d) + program.set_uniform('n_spot_lights', n_s) + program.set_uniform('n_point_lights', n_p) + plc = 0 + slc = 0 + dlc = 0 + + light_nodes = scene.light_nodes + if (len(scene.directional_light_nodes) > max_n_lights[0] or + len(scene.spot_light_nodes) > max_n_lights[1] or + len(scene.point_light_nodes) > max_n_lights[2]): + light_nodes = self._sorted_nodes_by_distance( + scene, scene.light_nodes, node + ) + + for n in light_nodes: + light = n.light + pose = scene.get_pose(n) + position = pose[:3,3] + direction = -pose[:3,2] + + if isinstance(light, PointLight): + if plc == max_n_lights[2]: + continue + b = 'point_lights[{}].'.format(plc) + plc += 1 + shadow = bool(flags & RenderFlags.SHADOWS_POINT) + program.set_uniform(b + 'position', position) + elif isinstance(light, SpotLight): + if slc == max_n_lights[1]: + continue + b = 'spot_lights[{}].'.format(slc) + slc += 1 + shadow = bool(flags & RenderFlags.SHADOWS_SPOT) + las = 1.0 / max(0.001, np.cos(light.innerConeAngle) - + np.cos(light.outerConeAngle)) + lao = -np.cos(light.outerConeAngle) * las + program.set_uniform(b + 'direction', direction) + program.set_uniform(b + 'position', position) + program.set_uniform(b + 'light_angle_scale', las) + program.set_uniform(b + 'light_angle_offset', lao) + else: + if dlc == max_n_lights[0]: + continue + b = 'directional_lights[{}].'.format(dlc) + dlc += 1 + shadow = bool(flags & RenderFlags.SHADOWS_DIRECTIONAL) + program.set_uniform(b + 'direction', direction) + + program.set_uniform(b + 'color', light.color) + program.set_uniform(b + 'intensity', light.intensity) + # if light.range is not None: + # program.set_uniform(b + 'range', light.range) + # else: + # program.set_uniform(b + 'range', 0) + + if shadow: + self._bind_texture(light.shadow_texture, + b + 'shadow_map', program) + if not isinstance(light, PointLight): + V, P = self._get_light_cam_matrices(scene, n, flags) + program.set_uniform(b + 'light_matrix', P.dot(V)) + else: + raise NotImplementedError( + 'Point light shadows not implemented' + ) + + def _sorted_mesh_nodes(self, scene): + cam_loc = scene.get_pose(scene.main_camera_node)[:3,3] + solid_nodes = [] + trans_nodes = [] + for node in scene.mesh_nodes: + mesh = node.mesh + if mesh.is_transparent: + trans_nodes.append(node) + else: + solid_nodes.append(node) + + # TODO BETTER SORTING METHOD + trans_nodes.sort( + key=lambda n: -np.linalg.norm(scene.get_pose(n)[:3,3] - cam_loc) + ) + solid_nodes.sort( + key=lambda n: -np.linalg.norm(scene.get_pose(n)[:3,3] - cam_loc) + ) + + return solid_nodes + trans_nodes + + def _sorted_nodes_by_distance(self, scene, nodes, compare_node): + nodes = list(nodes) + compare_posn = scene.get_pose(compare_node)[:3,3] + nodes.sort(key=lambda n: np.linalg.norm( + scene.get_pose(n)[:3,3] - compare_posn) + ) + return nodes + + ########################################################################### + # Context Management + ########################################################################### + + def _update_context(self, scene, flags): + + # Update meshes + scene_meshes = scene.meshes + + # Add new meshes to context + for mesh in scene_meshes - self._meshes: + for p in mesh.primitives: + p._add_to_context() + + # Remove old meshes from context + for mesh in self._meshes - scene_meshes: + for p in mesh.primitives: + p.delete() + + self._meshes = scene_meshes.copy() + + # Update mesh textures + mesh_textures = set() + for m in scene_meshes: + for p in m.primitives: + mesh_textures |= p.material.textures + + # Add new textures to context + for texture in mesh_textures - self._mesh_textures: + texture._add_to_context() + + # Remove old textures from context + for texture in self._mesh_textures - mesh_textures: + texture.delete() + + self._mesh_textures = mesh_textures.copy() + + shadow_textures = set() + for l in scene.lights: + # Create if needed + active = False + if (isinstance(l, DirectionalLight) and + flags & RenderFlags.SHADOWS_DIRECTIONAL): + active = True + elif (isinstance(l, PointLight) and + flags & RenderFlags.SHADOWS_POINT): + active = True + elif isinstance(l, SpotLight) and flags & RenderFlags.SHADOWS_SPOT: + active = True + + if active and l.shadow_texture is None: + l._generate_shadow_texture() + if l.shadow_texture is not None: + shadow_textures.add(l.shadow_texture) + + # Add new textures to context + for texture in shadow_textures - self._shadow_textures: + texture._add_to_context() + + # Remove old textures from context + for texture in self._shadow_textures - shadow_textures: + texture.delete() + + self._shadow_textures = shadow_textures.copy() + + ########################################################################### + # Texture Management + ########################################################################### + + def _bind_texture(self, texture, uniform_name, program): + """Bind a texture to an active texture unit and return + the texture unit index that was used. + """ + tex_id = self._get_next_active_texture() + glActiveTexture(GL_TEXTURE0 + tex_id) + texture._bind() + program.set_uniform(uniform_name, tex_id) + + def _get_next_active_texture(self): + val = self._texture_alloc_idx + self._texture_alloc_idx += 1 + return val + + def _reset_active_textures(self): + self._texture_alloc_idx = 0 + + ########################################################################### + # Camera Matrix Management + ########################################################################### + + def _get_camera_matrices(self, scene): + main_camera_node = scene.main_camera_node + if main_camera_node is None: + raise ValueError('Cannot render scene without a camera') + P = main_camera_node.camera.get_projection_matrix( + width=self.viewport_width, height=self.viewport_height + ) + pose = scene.get_pose(main_camera_node) + V = np.linalg.inv(pose) # V maps from world to camera + return V, P + + def _get_light_cam_matrices(self, scene, light_node, flags): + light = light_node.light + pose = scene.get_pose(light_node).copy() + s = scene.scale + camera = light._get_shadow_camera(s) + P = camera.get_projection_matrix() + if isinstance(light, DirectionalLight): + direction = -pose[:3,2] + c = scene.centroid + loc = c - direction * s + pose[:3,3] = loc + V = np.linalg.inv(pose) # V maps from world to camera + return V, P + + ########################################################################### + # Shader Program Management + ########################################################################### + + def _get_text_program(self): + program = self._program_cache.get_program( + vertex_shader='text.vert', + fragment_shader='text.frag' + ) + + if not program._in_context(): + program._add_to_context() + + return program + + def _compute_max_n_lights(self, flags): + max_n_lights = [MAX_N_LIGHTS, MAX_N_LIGHTS, MAX_N_LIGHTS] + n_tex_units = glGetIntegerv(GL_MAX_TEXTURE_IMAGE_UNITS) + + # Reserved texture units: 6 + # Normal Map + # Occlusion Map + # Emissive Map + # Base Color or Diffuse Map + # MR or SG Map + # Environment cubemap + + n_reserved_textures = 6 + n_available_textures = n_tex_units - n_reserved_textures + + # Distribute textures evenly among lights with shadows, with + # a preference for directional lights + n_shadow_types = 0 + if flags & RenderFlags.SHADOWS_DIRECTIONAL: + n_shadow_types += 1 + if flags & RenderFlags.SHADOWS_SPOT: + n_shadow_types += 1 + if flags & RenderFlags.SHADOWS_POINT: + n_shadow_types += 1 + + if n_shadow_types > 0: + tex_per_light = n_available_textures // n_shadow_types + + if flags & RenderFlags.SHADOWS_DIRECTIONAL: + max_n_lights[0] = ( + tex_per_light + + (n_available_textures - tex_per_light * n_shadow_types) + ) + if flags & RenderFlags.SHADOWS_SPOT: + max_n_lights[1] = tex_per_light + if flags & RenderFlags.SHADOWS_POINT: + max_n_lights[2] = tex_per_light + + return max_n_lights + + def _get_primitive_program(self, primitive, flags, program_flags): + vertex_shader = None + fragment_shader = None + geometry_shader = None + defines = {} + + if (bool(program_flags & ProgramFlags.USE_MATERIAL) and + not flags & RenderFlags.DEPTH_ONLY and + not flags & RenderFlags.FLAT and + not flags & RenderFlags.SEG): + vertex_shader = 'mesh.vert' + fragment_shader = 'mesh.frag' + elif bool(program_flags & (ProgramFlags.VERTEX_NORMALS | + ProgramFlags.FACE_NORMALS)): + vertex_shader = 'vertex_normals.vert' + if primitive.mode == GLTF.POINTS: + geometry_shader = 'vertex_normals_pc.geom' + else: + geometry_shader = 'vertex_normals.geom' + fragment_shader = 'vertex_normals.frag' + elif flags & RenderFlags.FLAT: + vertex_shader = 'flat.vert' + fragment_shader = 'flat.frag' + elif flags & RenderFlags.SEG: + vertex_shader = 'segmentation.vert' + fragment_shader = 'segmentation.frag' + else: + vertex_shader = 'mesh_depth.vert' + fragment_shader = 'mesh_depth.frag' + + # Set up vertex buffer DEFINES + bf = primitive.buf_flags + buf_idx = 1 + if bf & BufFlags.NORMAL: + defines['NORMAL_LOC'] = buf_idx + buf_idx += 1 + if bf & BufFlags.TANGENT: + defines['TANGENT_LOC'] = buf_idx + buf_idx += 1 + if bf & BufFlags.TEXCOORD_0: + defines['TEXCOORD_0_LOC'] = buf_idx + buf_idx += 1 + if bf & BufFlags.TEXCOORD_1: + defines['TEXCOORD_1_LOC'] = buf_idx + buf_idx += 1 + if bf & BufFlags.COLOR_0: + defines['COLOR_0_LOC'] = buf_idx + buf_idx += 1 + if bf & BufFlags.JOINTS_0: + defines['JOINTS_0_LOC'] = buf_idx + buf_idx += 1 + if bf & BufFlags.WEIGHTS_0: + defines['WEIGHTS_0_LOC'] = buf_idx + buf_idx += 1 + defines['INST_M_LOC'] = buf_idx + + # Set up shadow mapping defines + if flags & RenderFlags.SHADOWS_DIRECTIONAL: + defines['DIRECTIONAL_LIGHT_SHADOWS'] = 1 + if flags & RenderFlags.SHADOWS_SPOT: + defines['SPOT_LIGHT_SHADOWS'] = 1 + if flags & RenderFlags.SHADOWS_POINT: + defines['POINT_LIGHT_SHADOWS'] = 1 + max_n_lights = self._compute_max_n_lights(flags) + defines['MAX_DIRECTIONAL_LIGHTS'] = max_n_lights[0] + defines['MAX_SPOT_LIGHTS'] = max_n_lights[1] + defines['MAX_POINT_LIGHTS'] = max_n_lights[2] + + # Set up vertex normal defines + if program_flags & ProgramFlags.VERTEX_NORMALS: + defines['VERTEX_NORMALS'] = 1 + if program_flags & ProgramFlags.FACE_NORMALS: + defines['FACE_NORMALS'] = 1 + + # Set up material texture defines + if bool(program_flags & ProgramFlags.USE_MATERIAL): + tf = primitive.material.tex_flags + if tf & TexFlags.NORMAL: + defines['HAS_NORMAL_TEX'] = 1 + if tf & TexFlags.OCCLUSION: + defines['HAS_OCCLUSION_TEX'] = 1 + if tf & TexFlags.EMISSIVE: + defines['HAS_EMISSIVE_TEX'] = 1 + if tf & TexFlags.BASE_COLOR: + defines['HAS_BASE_COLOR_TEX'] = 1 + if tf & TexFlags.METALLIC_ROUGHNESS: + defines['HAS_METALLIC_ROUGHNESS_TEX'] = 1 + if tf & TexFlags.DIFFUSE: + defines['HAS_DIFFUSE_TEX'] = 1 + if tf & TexFlags.SPECULAR_GLOSSINESS: + defines['HAS_SPECULAR_GLOSSINESS_TEX'] = 1 + if isinstance(primitive.material, MetallicRoughnessMaterial): + defines['USE_METALLIC_MATERIAL'] = 1 + elif isinstance(primitive.material, SpecularGlossinessMaterial): + defines['USE_GLOSSY_MATERIAL'] = 1 + + program = self._program_cache.get_program( + vertex_shader=vertex_shader, + fragment_shader=fragment_shader, + geometry_shader=geometry_shader, + defines=defines + ) + + if not program._in_context(): + program._add_to_context() + + return program + + ########################################################################### + # Viewport Management + ########################################################################### + + def _configure_forward_pass_viewport(self, flags): + + # If using offscreen render, bind main framebuffer + if flags & RenderFlags.OFFSCREEN: + self._configure_main_framebuffer() + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb_ms) + else: + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, 0) + + glViewport(0, 0, self.viewport_width, self.viewport_height) + glEnable(GL_DEPTH_TEST) + glDepthMask(GL_TRUE) + glDepthFunc(GL_LESS) + glDepthRange(0.0, 1.0) + + def _configure_shadow_mapping_viewport(self, light, flags): + self._configure_shadow_framebuffer() + glBindFramebuffer(GL_FRAMEBUFFER, self._shadow_fb) + light.shadow_texture._bind() + light.shadow_texture._bind_as_depth_attachment() + glActiveTexture(GL_TEXTURE0) + light.shadow_texture._bind() + glDrawBuffer(GL_NONE) + glReadBuffer(GL_NONE) + + glClear(GL_DEPTH_BUFFER_BIT) + glViewport(0, 0, SHADOW_TEX_SZ, SHADOW_TEX_SZ) + glEnable(GL_DEPTH_TEST) + glDepthMask(GL_TRUE) + glDepthFunc(GL_LESS) + glDepthRange(0.0, 1.0) + glDisable(GL_CULL_FACE) + glDisable(GL_BLEND) + + ########################################################################### + # Framebuffer Management + ########################################################################### + + def _configure_shadow_framebuffer(self): + if self._shadow_fb is None: + self._shadow_fb = glGenFramebuffers(1) + + def _delete_shadow_framebuffer(self): + if self._shadow_fb is not None: + glDeleteFramebuffers(1, [self._shadow_fb]) + + def _configure_main_framebuffer(self): + # If mismatch with prior framebuffer, delete it + if (self._main_fb is not None and + self.viewport_width != self._main_fb_dims[0] or + self.viewport_height != self._main_fb_dims[1]): + self._delete_main_framebuffer() + + # If framebuffer doesn't exist, create it + if self._main_fb is None: + # Generate standard buffer + self._main_cb, self._main_db = glGenRenderbuffers(2) + + glBindRenderbuffer(GL_RENDERBUFFER, self._main_cb) + glRenderbufferStorage( + GL_RENDERBUFFER, GL_RGBA, + self.viewport_width, self.viewport_height + ) + + glBindRenderbuffer(GL_RENDERBUFFER, self._main_db) + glRenderbufferStorage( + GL_RENDERBUFFER, GL_DEPTH_COMPONENT24, + self.viewport_width, self.viewport_height + ) + + self._main_fb = glGenFramebuffers(1) + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb) + glFramebufferRenderbuffer( + GL_DRAW_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, + GL_RENDERBUFFER, self._main_cb + ) + glFramebufferRenderbuffer( + GL_DRAW_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, + GL_RENDERBUFFER, self._main_db + ) + + # Generate multisample buffer + self._main_cb_ms, self._main_db_ms = glGenRenderbuffers(2) + glBindRenderbuffer(GL_RENDERBUFFER, self._main_cb_ms) + # glRenderbufferStorageMultisample( + # GL_RENDERBUFFER, 4, GL_RGBA, + # self.viewport_width, self.viewport_height + # ) + # glBindRenderbuffer(GL_RENDERBUFFER, self._main_db_ms) + # glRenderbufferStorageMultisample( + # GL_RENDERBUFFER, 4, GL_DEPTH_COMPONENT24, + # self.viewport_width, self.viewport_height + # ) + # 增加这一行 + num_samples = min(glGetIntegerv(GL_MAX_SAMPLES), 4) # No more than GL_MAX_SAMPLES + + # 其实就是把 4 替换成 num_samples，其余不变 + glRenderbufferStorageMultisample(GL_RENDERBUFFER, num_samples, GL_RGBA, self.viewport_width, self.viewport_height) + + glBindRenderbuffer(GL_RENDERBUFFER, self._main_db_ms) # 这行不变 + + # 这一行也是将 4 替换成 num_samples + glRenderbufferStorageMultisample(GL_RENDERBUFFER, num_samples, GL_DEPTH_COMPONENT24, self.viewport_width, self.viewport_height) + + self._main_fb_ms = glGenFramebuffers(1) + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb_ms) + glFramebufferRenderbuffer( + GL_DRAW_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, + GL_RENDERBUFFER, self._main_cb_ms + ) + glFramebufferRenderbuffer( + GL_DRAW_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, + GL_RENDERBUFFER, self._main_db_ms + ) + + self._main_fb_dims = (self.viewport_width, self.viewport_height) + + def _delete_main_framebuffer(self): + if self._main_fb is not None: + glDeleteFramebuffers(2, [self._main_fb, self._main_fb_ms]) + if self._main_cb is not None: + glDeleteRenderbuffers(2, [self._main_cb, self._main_cb_ms]) + if self._main_db is not None: + glDeleteRenderbuffers(2, [self._main_db, self._main_db_ms]) + + self._main_fb = None + self._main_cb = None + self._main_db = None + self._main_fb_ms = None + self._main_cb_ms = None + self._main_db_ms = None + self._main_fb_dims = (None, None) + + def _read_main_framebuffer(self, scene, flags): + width, height = self._main_fb_dims[0], self._main_fb_dims[1] + + # Bind framebuffer and blit buffers + glBindFramebuffer(GL_READ_FRAMEBUFFER, self._main_fb_ms) + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, self._main_fb) + glBlitFramebuffer( + 0, 0, width, height, 0, 0, width, height, + GL_COLOR_BUFFER_BIT, GL_LINEAR + ) + glBlitFramebuffer( + 0, 0, width, height, 0, 0, width, height, + GL_DEPTH_BUFFER_BIT, GL_NEAREST + ) + glBindFramebuffer(GL_READ_FRAMEBUFFER, self._main_fb) + + # Read depth + depth_buf = glReadPixels( + 0, 0, width, height, GL_DEPTH_COMPONENT, GL_FLOAT + ) + depth_im = np.frombuffer(depth_buf, dtype=np.float32) + depth_im = depth_im.reshape((height, width)) + depth_im = np.flip(depth_im, axis=0) + inf_inds = (depth_im == 1.0) + depth_im = 2.0 * depth_im - 1.0 + z_near = scene.main_camera_node.camera.znear + z_far = scene.main_camera_node.camera.zfar + noninf = np.logical_not(inf_inds) + if z_far is None: + depth_im[noninf] = 2 * z_near / (1.0 - depth_im[noninf]) + else: + depth_im[noninf] = ((2.0 * z_near * z_far) / + (z_far + z_near - depth_im[noninf] * + (z_far - z_near))) + depth_im[inf_inds] = 0.0 + + # Resize for macos if needed + if sys.platform == 'darwin': + depth_im = self._resize_image(depth_im) + + if flags & RenderFlags.DEPTH_ONLY: + return depth_im + + # Read color + if flags & RenderFlags.RGBA: + color_buf = glReadPixels( + 0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE + ) + color_im = np.frombuffer(color_buf, dtype=np.uint8) + color_im = color_im.reshape((height, width, 4)) + else: + color_buf = glReadPixels( + 0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE + ) + color_im = np.frombuffer(color_buf, dtype=np.uint8) + color_im = color_im.reshape((height, width, 3)) + color_im = np.flip(color_im, axis=0) + + # Resize for macos if needed + if sys.platform == 'darwin': + color_im = self._resize_image(color_im, True) + + return color_im, depth_im + + def _resize_image(self, value, antialias=False): + """If needed, rescale the render for MacOS.""" + img = PIL.Image.fromarray(value) + resample = PIL.Image.NEAREST + if antialias: + resample = PIL.Image.BILINEAR + size = (self.viewport_width // self.dpscale, + self.viewport_height // self.dpscale) + img = img.resize(size, resample=resample) + return np.array(img) + + ########################################################################### + # Shadowmap Debugging + ########################################################################### + + def _forward_pass_no_reset(self, scene, flags): + # Set up camera matrices + V, P = self._get_camera_matrices(scene) + + # Now, render each object in sorted order + for node in self._sorted_mesh_nodes(scene): + mesh = node.mesh + + # Skip the mesh if it's not visible + if not mesh.is_visible: + continue + + for primitive in mesh.primitives: + + # First, get and bind the appropriate program + program = self._get_primitive_program( + primitive, flags, ProgramFlags.USE_MATERIAL + ) + program._bind() + + # Set the camera uniforms + program.set_uniform('V', V) + program.set_uniform('P', P) + program.set_uniform( + 'cam_pos', scene.get_pose(scene.main_camera_node)[:3,3] + ) + + # Next, bind the lighting + if not flags & RenderFlags.DEPTH_ONLY and not flags & RenderFlags.FLAT: + self._bind_lighting(scene, program, node, flags) + + # Finally, bind and draw the primitive + self._bind_and_draw_primitive( + primitive=primitive, + pose=scene.get_pose(node), + program=program, + flags=flags + ) + self._reset_active_textures() + + # Unbind the shader and flush the output + if program is not None: + program._unbind() + glFlush() + + def _render_light_shadowmaps(self, scene, light_nodes, flags, tile=False): + glBindFramebuffer(GL_DRAW_FRAMEBUFFER, 0) + glClearColor(*scene.bg_color) + glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) + glEnable(GL_DEPTH_TEST) + glDepthMask(GL_TRUE) + glDepthFunc(GL_LESS) + glDepthRange(0.0, 1.0) + + w = self.viewport_width + h = self.viewport_height + + num_nodes = len(light_nodes) + viewport_dims = { + (0, 2): [0, h // 2, w // 2, h], + (1, 2): [w // 2, h // 2, w, h], + (0, 3): [0, h // 2, w // 2, h], + (1, 3): [w // 2, h // 2, w, h], + (2, 3): [0, 0, w // 2, h // 2], + (0, 4): [0, h // 2, w // 2, h], + (1, 4): [w // 2, h // 2, w, h], + (2, 4): [0, 0, w // 2, h // 2], + (3, 4): [w // 2, 0, w, h // 2] + } + + if tile: + for i, ln in enumerate(light_nodes): + light = ln.light + + if light.shadow_texture is None: + raise ValueError('Light does not have a shadow texture') + + glViewport(*viewport_dims[(i, num_nodes + 1)]) + + program = self._get_debug_quad_program() + program._bind() + self._bind_texture(light.shadow_texture, 'depthMap', program) + self._render_debug_quad() + self._reset_active_textures() + glFlush() + i += 1 + glViewport(*viewport_dims[(i, num_nodes + 1)]) + self._forward_pass_no_reset(scene, flags) + else: + for i, ln in enumerate(light_nodes): + light = ln.light + + if light.shadow_texture is None: + raise ValueError('Light does not have a shadow texture') + + glViewport(0, 0, self.viewport_width, self.viewport_height) + + program = self._get_debug_quad_program() + program._bind() + self._bind_texture(light.shadow_texture, 'depthMap', program) + self._render_debug_quad() + self._reset_active_textures() + glFlush() + return + + def _get_debug_quad_program(self): + program = self._program_cache.get_program( + vertex_shader='debug_quad.vert', + fragment_shader='debug_quad.frag' + ) + if not program._in_context(): + program._add_to_context() + return program + + def _render_debug_quad(self): + x = glGenVertexArrays(1) + glBindVertexArray(x) + glDrawArrays(GL_TRIANGLES, 0, 6) + glBindVertexArray(0) + glDeleteVertexArrays(1, [x]) diff --git a/pyrender/pyrender/sampler.py b/pyrender/pyrender/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..e4784d068f808a40a56c8e748d83175f7f4e6233 --- /dev/null +++ b/pyrender/pyrender/sampler.py @@ -0,0 +1,102 @@ +"""Samplers, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-sampler + +Author: Matthew Matl +""" +from .constants import GLTF + + +class Sampler(object): + """Texture sampler properties for filtering and wrapping modes. + + Parameters + ---------- + name : str, optional + The user-defined name of this object. + magFilter : int, optional + Magnification filter. Valid values: + - :attr:`.GLTF.NEAREST` + - :attr:`.GLTF.LINEAR` + minFilter : int, optional + Minification filter. Valid values: + - :attr:`.GLTF.NEAREST` + - :attr:`.GLTF.LINEAR` + - :attr:`.GLTF.NEAREST_MIPMAP_NEAREST` + - :attr:`.GLTF.LINEAR_MIPMAP_NEAREST` + - :attr:`.GLTF.NEAREST_MIPMAP_LINEAR` + - :attr:`.GLTF.LINEAR_MIPMAP_LINEAR` + wrapS : int, optional + S (U) wrapping mode. Valid values: + - :attr:`.GLTF.CLAMP_TO_EDGE` + - :attr:`.GLTF.MIRRORED_REPEAT` + - :attr:`.GLTF.REPEAT` + wrapT : int, optional + T (V) wrapping mode. Valid values: + - :attr:`.GLTF.CLAMP_TO_EDGE` + - :attr:`.GLTF.MIRRORED_REPEAT` + - :attr:`.GLTF.REPEAT` + """ + + def __init__(self, + name=None, + magFilter=None, + minFilter=None, + wrapS=GLTF.REPEAT, + wrapT=GLTF.REPEAT): + self.name = name + self.magFilter = magFilter + self.minFilter = minFilter + self.wrapS = wrapS + self.wrapT = wrapT + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def magFilter(self): + """int : Magnification filter type. + """ + return self._magFilter + + @magFilter.setter + def magFilter(self, value): + self._magFilter = value + + @property + def minFilter(self): + """int : Minification filter type. + """ + return self._minFilter + + @minFilter.setter + def minFilter(self, value): + self._minFilter = value + + @property + def wrapS(self): + """int : S (U) wrapping mode. + """ + return self._wrapS + + @wrapS.setter + def wrapS(self, value): + self._wrapS = value + + @property + def wrapT(self): + """int : T (V) wrapping mode. + """ + return self._wrapT + + @wrapT.setter + def wrapT(self, value): + self._wrapT = value diff --git a/pyrender/pyrender/scene.py b/pyrender/pyrender/scene.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe057ec66f52f2dd9c1363aacf72a7c6cec4e6c --- /dev/null +++ b/pyrender/pyrender/scene.py @@ -0,0 +1,585 @@ +"""Scenes, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-scene + +Author: Matthew Matl +""" +import numpy as np +import networkx as nx +import trimesh + +from .mesh import Mesh +from .camera import Camera +from .light import Light, PointLight, DirectionalLight, SpotLight +from .node import Node +from .utils import format_color_vector + + +class Scene(object): + """A hierarchical scene graph. + + Parameters + ---------- + nodes : list of :class:`Node` + The set of all nodes in the scene. + bg_color : (4,) float, optional + Background color of scene. + ambient_light : (3,) float, optional + Color of ambient light. Defaults to no ambient light. + name : str, optional + The user-defined name of this object. + """ + + def __init__(self, + nodes=None, + bg_color=None, + ambient_light=None, + name=None): + + if bg_color is None: + bg_color = np.ones(4) + else: + bg_color = format_color_vector(bg_color, 4) + + if ambient_light is None: + ambient_light = np.zeros(3) + + if nodes is None: + nodes = set() + self._nodes = set() # Will be added at the end of this function + + self.bg_color = bg_color + self.ambient_light = ambient_light + self.name = name + + self._name_to_nodes = {} + self._obj_to_nodes = {} + self._obj_name_to_nodes = {} + self._mesh_nodes = set() + self._point_light_nodes = set() + self._spot_light_nodes = set() + self._directional_light_nodes = set() + self._camera_nodes = set() + self._main_camera_node = None + self._bounds = None + + # Transform tree + self._digraph = nx.DiGraph() + self._digraph.add_node('world') + self._path_cache = {} + + # Find root nodes and add them + if len(nodes) > 0: + node_parent_map = {n: None for n in nodes} + for node in nodes: + for child in node.children: + if node_parent_map[child] is not None: + raise ValueError('Nodes may not have more than ' + 'one parent') + node_parent_map[child] = node + for node in node_parent_map: + if node_parent_map[node] is None: + self.add_node(node) + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def nodes(self): + """set of :class:`Node` : Set of nodes in the scene. + """ + return self._nodes + + @property + def bg_color(self): + """(3,) float : The scene background color. + """ + return self._bg_color + + @bg_color.setter + def bg_color(self, value): + if value is None: + value = np.ones(4) + else: + value = format_color_vector(value, 4) + self._bg_color = value + + @property + def ambient_light(self): + """(3,) float : The ambient light in the scene. + """ + return self._ambient_light + + @ambient_light.setter + def ambient_light(self, value): + if value is None: + value = np.zeros(3) + else: + value = format_color_vector(value, 3) + self._ambient_light = value + + @property + def meshes(self): + """set of :class:`Mesh` : The meshes in the scene. + """ + return set([n.mesh for n in self.mesh_nodes]) + + @property + def mesh_nodes(self): + """set of :class:`Node` : The nodes containing meshes. + """ + return self._mesh_nodes + + @property + def lights(self): + """set of :class:`Light` : The lights in the scene. + """ + return self.point_lights | self.spot_lights | self.directional_lights + + @property + def light_nodes(self): + """set of :class:`Node` : The nodes containing lights. + """ + return (self.point_light_nodes | self.spot_light_nodes | + self.directional_light_nodes) + + @property + def point_lights(self): + """set of :class:`PointLight` : The point lights in the scene. + """ + return set([n.light for n in self.point_light_nodes]) + + @property + def point_light_nodes(self): + """set of :class:`Node` : The nodes containing point lights. + """ + return self._point_light_nodes + + @property + def spot_lights(self): + """set of :class:`SpotLight` : The spot lights in the scene. + """ + return set([n.light for n in self.spot_light_nodes]) + + @property + def spot_light_nodes(self): + """set of :class:`Node` : The nodes containing spot lights. + """ + return self._spot_light_nodes + + @property + def directional_lights(self): + """set of :class:`DirectionalLight` : The directional lights in + the scene. + """ + return set([n.light for n in self.directional_light_nodes]) + + @property + def directional_light_nodes(self): + """set of :class:`Node` : The nodes containing directional lights. + """ + return self._directional_light_nodes + + @property + def cameras(self): + """set of :class:`Camera` : The cameras in the scene. + """ + return set([n.camera for n in self.camera_nodes]) + + @property + def camera_nodes(self): + """set of :class:`Node` : The nodes containing cameras in the scene. + """ + return self._camera_nodes + + @property + def main_camera_node(self): + """set of :class:`Node` : The node containing the main camera in the + scene. + """ + return self._main_camera_node + + @main_camera_node.setter + def main_camera_node(self, value): + if value not in self.nodes: + raise ValueError('New main camera node must already be in scene') + self._main_camera_node = value + + @property + def bounds(self): + """(2,3) float : The axis-aligned bounds of the scene. + """ + if self._bounds is None: + # Compute corners + corners = [] + for mesh_node in self.mesh_nodes: + mesh = mesh_node.mesh + pose = self.get_pose(mesh_node) + corners_local = trimesh.bounds.corners(mesh.bounds) + corners_world = pose[:3,:3].dot(corners_local.T).T + pose[:3,3] + corners.append(corners_world) + if len(corners) == 0: + self._bounds = np.zeros((2,3)) + else: + corners = np.vstack(corners) + self._bounds = np.array([np.min(corners, axis=0), + np.max(corners, axis=0)]) + return self._bounds + + @property + def centroid(self): + """(3,) float : The centroid of the scene's axis-aligned bounding box + (AABB). + """ + return np.mean(self.bounds, axis=0) + + @property + def extents(self): + """(3,) float : The lengths of the axes of the scene's AABB. + """ + return np.diff(self.bounds, axis=0).reshape(-1) + + @property + def scale(self): + """(3,) float : The length of the diagonal of the scene's AABB. + """ + return np.linalg.norm(self.extents) + + def add(self, obj, name=None, pose=None, + parent_node=None, parent_name=None): + """Add an object (mesh, light, or camera) to the scene. + + Parameters + ---------- + obj : :class:`Mesh`, :class:`Light`, or :class:`Camera` + The object to add to the scene. + name : str + A name for the new node to be created. + pose : (4,4) float + The local pose of this node relative to its parent node. + parent_node : :class:`Node` + The parent of this Node. If None, the new node is a root node. + parent_name : str + The name of the parent node, can be specified instead of + `parent_node`. + + Returns + ------- + node : :class:`Node` + The newly-created and inserted node. + """ + if isinstance(obj, Mesh): + node = Node(name=name, matrix=pose, mesh=obj) + elif isinstance(obj, Light): + node = Node(name=name, matrix=pose, light=obj) + elif isinstance(obj, Camera): + node = Node(name=name, matrix=pose, camera=obj) + else: + raise TypeError('Unrecognized object type') + + if parent_node is None and parent_name is not None: + parent_nodes = self.get_nodes(name=parent_name) + if len(parent_nodes) == 0: + raise ValueError('No parent node with name {} found' + .format(parent_name)) + elif len(parent_nodes) > 1: + raise ValueError('More than one parent node with name {} found' + .format(parent_name)) + parent_node = list(parent_nodes)[0] + + self.add_node(node, parent_node=parent_node) + + return node + + def get_nodes(self, node=None, name=None, obj=None, obj_name=None): + """Search for existing nodes. Only nodes matching all specified + parameters is returned, or None if no such node exists. + + Parameters + ---------- + node : :class:`Node`, optional + If present, returns this node if it is in the scene. + name : str + A name for the Node. + obj : :class:`Mesh`, :class:`Light`, or :class:`Camera` + An object that is attached to the node. + obj_name : str + The name of an object that is attached to the node. + + Returns + ------- + nodes : set of :class:`.Node` + The nodes that match all query terms. + """ + if node is not None: + if node in self.nodes: + return set([node]) + else: + return set() + nodes = set(self.nodes) + if name is not None: + matches = set() + if name in self._name_to_nodes: + matches = self._name_to_nodes[name] + nodes = nodes & matches + if obj is not None: + matches = set() + if obj in self._obj_to_nodes: + matches = self._obj_to_nodes[obj] + nodes = nodes & matches + if obj_name is not None: + matches = set() + if obj_name in self._obj_name_to_nodes: + matches = self._obj_name_to_nodes[obj_name] + nodes = nodes & matches + + return nodes + + def add_node(self, node, parent_node=None): + """Add a Node to the scene. + + Parameters + ---------- + node : :class:`Node` + The node to be added. + parent_node : :class:`Node` + The parent of this Node. If None, the new node is a root node. + """ + if node in self.nodes: + raise ValueError('Node already in scene') + self.nodes.add(node) + + # Add node to sets + if node.name is not None: + if node.name not in self._name_to_nodes: + self._name_to_nodes[node.name] = set() + self._name_to_nodes[node.name].add(node) + for obj in [node.mesh, node.camera, node.light]: + if obj is not None: + if obj not in self._obj_to_nodes: + self._obj_to_nodes[obj] = set() + self._obj_to_nodes[obj].add(node) + if obj.name is not None: + if obj.name not in self._obj_name_to_nodes: + self._obj_name_to_nodes[obj.name] = set() + self._obj_name_to_nodes[obj.name].add(node) + if node.mesh is not None: + self._mesh_nodes.add(node) + if node.light is not None: + if isinstance(node.light, PointLight): + self._point_light_nodes.add(node) + if isinstance(node.light, SpotLight): + self._spot_light_nodes.add(node) + if isinstance(node.light, DirectionalLight): + self._directional_light_nodes.add(node) + if node.camera is not None: + self._camera_nodes.add(node) + if self._main_camera_node is None: + self._main_camera_node = node + + if parent_node is None: + parent_node = 'world' + elif parent_node not in self.nodes: + raise ValueError('Parent node must already be in scene') + elif node not in parent_node.children: + parent_node.children.append(node) + + # Create node in graph + self._digraph.add_node(node) + self._digraph.add_edge(node, parent_node) + + # Iterate over children + for child in node.children: + self.add_node(child, node) + + self._path_cache = {} + self._bounds = None + + def has_node(self, node): + """Check if a node is already in the scene. + + Parameters + ---------- + node : :class:`Node` + The node to be checked. + + Returns + ------- + has_node : bool + True if the node is already in the scene and false otherwise. + """ + return node in self.nodes + + def remove_node(self, node): + """Remove a node and all its children from the scene. + + Parameters + ---------- + node : :class:`Node` + The node to be removed. + """ + # Disconnect self from parent who is staying in the graph + parent = list(self._digraph.neighbors(node))[0] + self._remove_node(node) + if isinstance(parent, Node): + parent.children.remove(node) + self._path_cache = {} + self._bounds = None + + def get_pose(self, node): + """Get the world-frame pose of a node in the scene. + + Parameters + ---------- + node : :class:`Node` + The node to find the pose of. + + Returns + ------- + pose : (4,4) float + The transform matrix for this node. + """ + if node not in self.nodes: + raise ValueError('Node must already be in scene') + if node in self._path_cache: + path = self._path_cache[node] + else: + # Get path from from_frame to to_frame + path = nx.shortest_path(self._digraph, node, 'world') + self._path_cache[node] = path + + # Traverse from from_node to to_node + pose = np.eye(4) + for n in path[:-1]: + pose = np.dot(n.matrix, pose) + + return pose + + def set_pose(self, node, pose): + """Set the local-frame pose of a node in the scene. + + Parameters + ---------- + node : :class:`Node` + The node to set the pose of. + pose : (4,4) float + The pose to set the node to. + """ + if node not in self.nodes: + raise ValueError('Node must already be in scene') + node._matrix = pose + if node.mesh is not None: + self._bounds = None + + def clear(self): + """Clear out all nodes to form an empty scene. + """ + self._nodes = set() + + self._name_to_nodes = {} + self._obj_to_nodes = {} + self._obj_name_to_nodes = {} + self._mesh_nodes = set() + self._point_light_nodes = set() + self._spot_light_nodes = set() + self._directional_light_nodes = set() + self._camera_nodes = set() + self._main_camera_node = None + self._bounds = None + + # Transform tree + self._digraph = nx.DiGraph() + self._digraph.add_node('world') + self._path_cache = {} + + def _remove_node(self, node): + """Remove a node and all its children from the scene. + + Parameters + ---------- + node : :class:`Node` + The node to be removed. + """ + + # Remove self from nodes + self.nodes.remove(node) + + # Remove children + for child in node.children: + self._remove_node(child) + + # Remove self from the graph + self._digraph.remove_node(node) + + # Remove from maps + if node.name in self._name_to_nodes: + self._name_to_nodes[node.name].remove(node) + if len(self._name_to_nodes[node.name]) == 0: + self._name_to_nodes.pop(node.name) + for obj in [node.mesh, node.camera, node.light]: + if obj is None: + continue + self._obj_to_nodes[obj].remove(node) + if len(self._obj_to_nodes[obj]) == 0: + self._obj_to_nodes.pop(obj) + if obj.name is not None: + self._obj_name_to_nodes[obj.name].remove(node) + if len(self._obj_name_to_nodes[obj.name]) == 0: + self._obj_name_to_nodes.pop(obj.name) + if node.mesh is not None: + self._mesh_nodes.remove(node) + if node.light is not None: + if isinstance(node.light, PointLight): + self._point_light_nodes.remove(node) + if isinstance(node.light, SpotLight): + self._spot_light_nodes.remove(node) + if isinstance(node.light, DirectionalLight): + self._directional_light_nodes.remove(node) + if node.camera is not None: + self._camera_nodes.remove(node) + if self._main_camera_node == node: + if len(self._camera_nodes) > 0: + self._main_camera_node = next(iter(self._camera_nodes)) + else: + self._main_camera_node = None + + @staticmethod + def from_trimesh_scene(trimesh_scene, + bg_color=None, ambient_light=None): + """Create a :class:`.Scene` from a :class:`trimesh.scene.scene.Scene`. + + Parameters + ---------- + trimesh_scene : :class:`trimesh.scene.scene.Scene` + Scene with :class:~`trimesh.base.Trimesh` objects. + bg_color : (4,) float + Background color for the created scene. + ambient_light : (3,) float or None + Ambient light in the scene. + + Returns + ------- + scene_pr : :class:`Scene` + A scene containing the same geometry as the trimesh scene. + """ + # convert trimesh geometries to pyrender geometries + geometries = {name: Mesh.from_trimesh(geom) + for name, geom in trimesh_scene.geometry.items()} + + # create the pyrender scene object + scene_pr = Scene(bg_color=bg_color, ambient_light=ambient_light) + + # add every node with geometry to the pyrender scene + for node in trimesh_scene.graph.nodes_geometry: + pose, geom_name = trimesh_scene.graph[node] + scene_pr.add(geometries[geom_name], pose=pose) + + return scene_pr diff --git a/pyrender/pyrender/shader_program.py b/pyrender/pyrender/shader_program.py new file mode 100644 index 0000000000000000000000000000000000000000..c1803f280c98033abe0769771a9ad8ecfec942e3 --- /dev/null +++ b/pyrender/pyrender/shader_program.py @@ -0,0 +1,283 @@ +"""OpenGL shader program wrapper. +""" +import numpy as np +import os +import re + +import OpenGL +from OpenGL.GL import * +from OpenGL.GL import shaders as gl_shader_utils + + +class ShaderProgramCache(object): + """A cache for shader programs. + """ + + def __init__(self, shader_dir=None): + self._program_cache = {} + self.shader_dir = shader_dir + if self.shader_dir is None: + base_dir, _ = os.path.split(os.path.realpath(__file__)) + self.shader_dir = os.path.join(base_dir, 'shaders') + + def get_program(self, vertex_shader, fragment_shader, + geometry_shader=None, defines=None): + """Get a program via a list of shader files to include in the program. + + Parameters + ---------- + vertex_shader : str + The vertex shader filename. + fragment_shader : str + The fragment shader filename. + geometry_shader : str + The geometry shader filename. + defines : dict + Defines and their values for the shader. + + Returns + ------- + program : :class:`.ShaderProgram` + The program. + """ + shader_names = [] + if defines is None: + defines = {} + shader_filenames = [ + x for x in [vertex_shader, fragment_shader, geometry_shader] + if x is not None + ] + for fn in shader_filenames: + if fn is None: + continue + _, name = os.path.split(fn) + shader_names.append(name) + cid = OpenGL.contextdata.getContext() + key = tuple([cid] + sorted( + [(s,1) for s in shader_names] + [(d, defines[d]) for d in defines] + )) + + if key not in self._program_cache: + shader_filenames = [ + os.path.join(self.shader_dir, fn) for fn in shader_filenames + ] + if len(shader_filenames) == 2: + shader_filenames.append(None) + vs, fs, gs = shader_filenames + self._program_cache[key] = ShaderProgram( + vertex_shader=vs, fragment_shader=fs, + geometry_shader=gs, defines=defines + ) + return self._program_cache[key] + + def clear(self): + for key in self._program_cache: + self._program_cache[key].delete() + self._program_cache = {} + + +class ShaderProgram(object): + """A thin wrapper about OpenGL shader programs that supports easy creation, + binding, and uniform-setting. + + Parameters + ---------- + vertex_shader : str + The vertex shader filename. + fragment_shader : str + The fragment shader filename. + geometry_shader : str + The geometry shader filename. + defines : dict + Defines and their values for the shader. + """ + + def __init__(self, vertex_shader, fragment_shader, + geometry_shader=None, defines=None): + + self.vertex_shader = vertex_shader + self.fragment_shader = fragment_shader + self.geometry_shader = geometry_shader + + self.defines = defines + if self.defines is None: + self.defines = {} + + self._program_id = None + self._vao_id = None # PYOPENGL BUG + + # DEBUG + # self._unif_map = {} + + def _add_to_context(self): + if self._program_id is not None: + raise ValueError('Shader program already in context') + shader_ids = [] + + # Load vert shader + shader_ids.append(gl_shader_utils.compileShader( + self._load(self.vertex_shader), GL_VERTEX_SHADER) + ) + # Load frag shader + shader_ids.append(gl_shader_utils.compileShader( + self._load(self.fragment_shader), GL_FRAGMENT_SHADER) + ) + # Load geometry shader + if self.geometry_shader is not None: + shader_ids.append(gl_shader_utils.compileShader( + self._load(self.geometry_shader), GL_GEOMETRY_SHADER) + ) + + # Bind empty VAO PYOPENGL BUG + if self._vao_id is None: + self._vao_id = glGenVertexArrays(1) + glBindVertexArray(self._vao_id) + + # Compile program + self._program_id = gl_shader_utils.compileProgram(*shader_ids) + + # Unbind empty VAO PYOPENGL BUG + glBindVertexArray(0) + + def _in_context(self): + return self._program_id is not None + + def _remove_from_context(self): + if self._program_id is not None: + glDeleteProgram(self._program_id) + glDeleteVertexArrays(1, [self._vao_id]) + self._program_id = None + self._vao_id = None + + def _load(self, shader_filename): + path, _ = os.path.split(shader_filename) + + with open(shader_filename) as f: + text = f.read() + + def ifdef(matchobj): + if matchobj.group(1) in self.defines: + return '#if 1' + else: + return '#if 0' + + def ifndef(matchobj): + if matchobj.group(1) in self.defines: + return '#if 0' + else: + return '#if 1' + + ifdef_regex = re.compile( + '#ifdef\\s+([a-zA-Z_][a-zA-Z_0-9]*)\\s*$', re.MULTILINE + ) + ifndef_regex = re.compile( + '#ifndef\\s+([a-zA-Z_][a-zA-Z_0-9]*)\\s*$', re.MULTILINE + ) + text = re.sub(ifdef_regex, ifdef, text) + text = re.sub(ifndef_regex, ifndef, text) + + for define in self.defines: + value = str(self.defines[define]) + text = text.replace(define, value) + + return text + + def _bind(self): + """Bind this shader program to the current OpenGL context. + """ + if self._program_id is None: + raise ValueError('Cannot bind program that is not in context') + # glBindVertexArray(self._vao_id) + glUseProgram(self._program_id) + + def _unbind(self): + """Unbind this shader program from the current OpenGL context. + """ + glUseProgram(0) + + def delete(self): + """Delete this shader program from the current OpenGL context. + """ + self._remove_from_context() + + def set_uniform(self, name, value, unsigned=False): + """Set a uniform value in the current shader program. + + Parameters + ---------- + name : str + Name of the uniform to set. + value : int, float, or ndarray + Value to set the uniform to. + unsigned : bool + If True, ints will be treated as unsigned values. + """ + try: + # DEBUG + # self._unif_map[name] = 1, (1,) + loc = glGetUniformLocation(self._program_id, name) + + if loc == -1: + raise ValueError('Invalid shader variable: {}'.format(name)) + + if isinstance(value, np.ndarray): + # DEBUG + # self._unif_map[name] = value.size, value.shape + if value.ndim == 1: + if (np.issubdtype(value.dtype, np.unsignedinteger) or + unsigned): + dtype = 'u' + value = value.astype(np.uint32) + elif np.issubdtype(value.dtype, np.integer): + dtype = 'i' + value = value.astype(np.int32) + else: + dtype = 'f' + value = value.astype(np.float32) + self._FUNC_MAP[(value.shape[0], dtype)](loc, 1, value) + else: + self._FUNC_MAP[(value.shape[0], value.shape[1])]( + loc, 1, GL_TRUE, value + ) + + # Call correct uniform function + elif isinstance(value, float): + glUniform1f(loc, value) + elif isinstance(value, int): + if unsigned: + glUniform1ui(loc, value) + else: + glUniform1i(loc, value) + elif isinstance(value, bool): + if unsigned: + glUniform1ui(loc, int(value)) + else: + glUniform1i(loc, int(value)) + else: + raise ValueError('Invalid data type') + except Exception: + pass + + _FUNC_MAP = { + (1,'u'): glUniform1uiv, + (2,'u'): glUniform2uiv, + (3,'u'): glUniform3uiv, + (4,'u'): glUniform4uiv, + (1,'i'): glUniform1iv, + (2,'i'): glUniform2iv, + (3,'i'): glUniform3iv, + (4,'i'): glUniform4iv, + (1,'f'): glUniform1fv, + (2,'f'): glUniform2fv, + (3,'f'): glUniform3fv, + (4,'f'): glUniform4fv, + (2,2): glUniformMatrix2fv, + (2,3): glUniformMatrix2x3fv, + (2,4): glUniformMatrix2x4fv, + (3,2): glUniformMatrix3x2fv, + (3,3): glUniformMatrix3fv, + (3,4): glUniformMatrix3x4fv, + (4,2): glUniformMatrix4x2fv, + (4,3): glUniformMatrix4x3fv, + (4,4): glUniformMatrix4fv, + } diff --git a/pyrender/pyrender/shaders/debug_quad.frag b/pyrender/pyrender/shaders/debug_quad.frag new file mode 100644 index 0000000000000000000000000000000000000000..4647bb50dfa1e4510e2d4afb37959c7f57532eca --- /dev/null +++ b/pyrender/pyrender/shaders/debug_quad.frag @@ -0,0 +1,23 @@ +#version 330 core +out vec4 FragColor; + +in vec2 TexCoords; + +uniform sampler2D depthMap; +//uniform float near_plane; +//uniform float far_plane; +// +//// required when using a perspective projection matrix +//float LinearizeDepth(float depth) +//{ +// float z = depth * 2.0 - 1.0; // Back to NDC +// return (2.0 * near_plane * far_plane) / (far_plane + near_plane - z * (far_plane - near_plane)); +//} + +void main() +{ + float depthValue = texture(depthMap, TexCoords).r; + // FragColor = vec4(vec3(LinearizeDepth(depthValue) / far_plane), 1.0); // perspective + FragColor = vec4(vec3(depthValue), 1.0); // orthographic + //FragColor = vec4(1.0, 1.0, 0.0, 1.0); +} diff --git a/pyrender/pyrender/shaders/debug_quad.vert b/pyrender/pyrender/shaders/debug_quad.vert new file mode 100644 index 0000000000000000000000000000000000000000..d2f2fcb7626f6c22e0d52bf4d6c91251cbdb9f52 --- /dev/null +++ b/pyrender/pyrender/shaders/debug_quad.vert @@ -0,0 +1,25 @@ +#version 330 core +//layout (location = 0) in vec3 aPos; +//layout (location = 1) in vec2 aTexCoords; +// +//out vec2 TexCoords; +// +//void main() +//{ +// TexCoords = aTexCoords; +// gl_Position = vec4(aPos, 1.0); +//} +// +// +//layout(location = 0) out vec2 uv; + +out vec2 TexCoords; + +void main() +{ + float x = float(((uint(gl_VertexID) + 2u) / 3u)%2u); + float y = float(((uint(gl_VertexID) + 1u) / 3u)%2u); + + gl_Position = vec4(-1.0f + x*2.0f, -1.0f+y*2.0f, 0.0f, 1.0f); + TexCoords = vec2(x, y); +} diff --git a/pyrender/pyrender/shaders/flat.frag b/pyrender/pyrender/shaders/flat.frag new file mode 100644 index 0000000000000000000000000000000000000000..7ec01c6d095ec5dacc693accd3ad507ced61a79a --- /dev/null +++ b/pyrender/pyrender/shaders/flat.frag @@ -0,0 +1,126 @@ +#version 330 core +/////////////////////////////////////////////////////////////////////////////// +// Structs +/////////////////////////////////////////////////////////////////////////////// + +struct Material { + vec3 emissive_factor; + +#ifdef USE_METALLIC_MATERIAL + vec4 base_color_factor; + float metallic_factor; + float roughness_factor; +#endif + +#ifdef USE_GLOSSY_MATERIAL + vec4 diffuse_factor; + vec3 specular_factor; + float glossiness_factor; +#endif + +#ifdef HAS_NORMAL_TEX + sampler2D normal_texture; +#endif +#ifdef HAS_OCCLUSION_TEX + sampler2D occlusion_texture; +#endif +#ifdef HAS_EMISSIVE_TEX + sampler2D emissive_texture; +#endif +#ifdef HAS_BASE_COLOR_TEX + sampler2D base_color_texture; +#endif +#ifdef HAS_METALLIC_ROUGHNESS_TEX + sampler2D metallic_roughness_texture; +#endif +#ifdef HAS_DIFFUSE_TEX + sampler2D diffuse_texture; +#endif +#ifdef HAS_SPECULAR_GLOSSINESS_TEX + sampler2D specular_glossiness; +#endif +}; + +/////////////////////////////////////////////////////////////////////////////// +// Uniforms +/////////////////////////////////////////////////////////////////////////////// +uniform Material material; +uniform vec3 cam_pos; + +#ifdef USE_IBL +uniform samplerCube diffuse_env; +uniform samplerCube specular_env; +#endif + +/////////////////////////////////////////////////////////////////////////////// +// Inputs +/////////////////////////////////////////////////////////////////////////////// + +in vec3 frag_position; +#ifdef NORMAL_LOC +in vec3 frag_normal; +#endif +#ifdef HAS_NORMAL_TEX +#ifdef TANGENT_LOC +#ifdef NORMAL_LOC +in mat3 tbn; +#endif +#endif +#endif +#ifdef TEXCOORD_0_LOC +in vec2 uv_0; +#endif +#ifdef TEXCOORD_1_LOC +in vec2 uv_1; +#endif +#ifdef COLOR_0_LOC +in vec4 color_multiplier; +#endif + +/////////////////////////////////////////////////////////////////////////////// +// OUTPUTS +/////////////////////////////////////////////////////////////////////////////// + +out vec4 frag_color; + +/////////////////////////////////////////////////////////////////////////////// +// Constants +/////////////////////////////////////////////////////////////////////////////// +const float PI = 3.141592653589793; +const float min_roughness = 0.04; + +/////////////////////////////////////////////////////////////////////////////// +// Utility Functions +/////////////////////////////////////////////////////////////////////////////// +vec4 srgb_to_linear(vec4 srgb) +{ +#ifndef SRGB_CORRECTED + // Fast Approximation + //vec3 linOut = pow(srgbIn.xyz,vec3(2.2)); + // + vec3 b_less = step(vec3(0.04045),srgb.xyz); + vec3 lin_out = mix( srgb.xyz/vec3(12.92), pow((srgb.xyz+vec3(0.055))/vec3(1.055),vec3(2.4)), b_less ); + return vec4(lin_out, srgb.w); +#else + return srgb; +#endif +} + +/////////////////////////////////////////////////////////////////////////////// +// MAIN +/////////////////////////////////////////////////////////////////////////////// +void main() +{ + + // Compute albedo + vec4 base_color = material.base_color_factor; +#ifdef HAS_BASE_COLOR_TEX + base_color = base_color * texture(material.base_color_texture, uv_0); +#endif + +#ifdef COLOR_0_LOC + base_color *= color_multiplier; +#endif + + frag_color = clamp(base_color, 0.0, 1.0); +} diff --git a/pyrender/pyrender/shaders/flat.vert b/pyrender/pyrender/shaders/flat.vert new file mode 100644 index 0000000000000000000000000000000000000000..cfd241c3544718a261f961c3aa3c03aa13c97761 --- /dev/null +++ b/pyrender/pyrender/shaders/flat.vert @@ -0,0 +1,86 @@ +#version 330 core + +// Vertex Attributes +layout(location = 0) in vec3 position; +#ifdef NORMAL_LOC +layout(location = NORMAL_LOC) in vec3 normal; +#endif +#ifdef TANGENT_LOC +layout(location = TANGENT_LOC) in vec4 tangent; +#endif +#ifdef TEXCOORD_0_LOC +layout(location = TEXCOORD_0_LOC) in vec2 texcoord_0; +#endif +#ifdef TEXCOORD_1_LOC +layout(location = TEXCOORD_1_LOC) in vec2 texcoord_1; +#endif +#ifdef COLOR_0_LOC +layout(location = COLOR_0_LOC) in vec4 color_0; +#endif +#ifdef JOINTS_0_LOC +layout(location = JOINTS_0_LOC) in vec4 joints_0; +#endif +#ifdef WEIGHTS_0_LOC +layout(location = WEIGHTS_0_LOC) in vec4 weights_0; +#endif +layout(location = INST_M_LOC) in mat4 inst_m; + +// Uniforms +uniform mat4 M; +uniform mat4 V; +uniform mat4 P; + +// Outputs +out vec3 frag_position; +#ifdef NORMAL_LOC +out vec3 frag_normal; +#endif +#ifdef HAS_NORMAL_TEX +#ifdef TANGENT_LOC +#ifdef NORMAL_LOC +out mat3 tbn; +#endif +#endif +#endif +#ifdef TEXCOORD_0_LOC +out vec2 uv_0; +#endif +#ifdef TEXCOORD_1_LOC +out vec2 uv_1; +#endif +#ifdef COLOR_0_LOC +out vec4 color_multiplier; +#endif + + +void main() +{ + gl_Position = P * V * M * inst_m * vec4(position, 1); + frag_position = vec3(M * inst_m * vec4(position, 1.0)); + + mat4 N = transpose(inverse(M * inst_m)); + +#ifdef NORMAL_LOC + frag_normal = normalize(vec3(N * vec4(normal, 0.0))); +#endif + +#ifdef HAS_NORMAL_TEX +#ifdef TANGENT_LOC +#ifdef NORMAL_LOC + vec3 normal_w = normalize(vec3(N * vec4(normal, 0.0))); + vec3 tangent_w = normalize(vec3(N * vec4(tangent.xyz, 0.0))); + vec3 bitangent_w = cross(normal_w, tangent_w) * tangent.w; + tbn = mat3(tangent_w, bitangent_w, normal_w); +#endif +#endif +#endif +#ifdef TEXCOORD_0_LOC + uv_0 = texcoord_0; +#endif +#ifdef TEXCOORD_1_LOC + uv_1 = texcoord_1; +#endif +#ifdef COLOR_0_LOC + color_multiplier = color_0; +#endif +} diff --git a/pyrender/pyrender/shaders/mesh.frag b/pyrender/pyrender/shaders/mesh.frag new file mode 100644 index 0000000000000000000000000000000000000000..43187621b4388b18badf4e562a7ad300e59b029d --- /dev/null +++ b/pyrender/pyrender/shaders/mesh.frag @@ -0,0 +1,456 @@ +#version 330 core +/////////////////////////////////////////////////////////////////////////////// +// Structs +/////////////////////////////////////////////////////////////////////////////// + +struct SpotLight { + vec3 color; + float intensity; + float range; + vec3 position; + vec3 direction; + float light_angle_scale; + float light_angle_offset; + + #ifdef SPOT_LIGHT_SHADOWS + sampler2D shadow_map; + mat4 light_matrix; + #endif +}; + +struct DirectionalLight { + vec3 color; + float intensity; + vec3 direction; + + #ifdef DIRECTIONAL_LIGHT_SHADOWS + sampler2D shadow_map; + mat4 light_matrix; + #endif +}; + +struct PointLight { + vec3 color; + float intensity; + float range; + vec3 position; + + #ifdef POINT_LIGHT_SHADOWS + samplerCube shadow_map; + #endif +}; + +struct Material { + vec3 emissive_factor; + +#ifdef USE_METALLIC_MATERIAL + vec4 base_color_factor; + float metallic_factor; + float roughness_factor; +#endif + +#ifdef USE_GLOSSY_MATERIAL + vec4 diffuse_factor; + vec3 specular_factor; + float glossiness_factor; +#endif + +#ifdef HAS_NORMAL_TEX + sampler2D normal_texture; +#endif +#ifdef HAS_OCCLUSION_TEX + sampler2D occlusion_texture; +#endif +#ifdef HAS_EMISSIVE_TEX + sampler2D emissive_texture; +#endif +#ifdef HAS_BASE_COLOR_TEX + sampler2D base_color_texture; +#endif +#ifdef HAS_METALLIC_ROUGHNESS_TEX + sampler2D metallic_roughness_texture; +#endif +#ifdef HAS_DIFFUSE_TEX + sampler2D diffuse_texture; +#endif +#ifdef HAS_SPECULAR_GLOSSINESS_TEX + sampler2D specular_glossiness; +#endif +}; + +struct PBRInfo { + float nl; + float nv; + float nh; + float lh; + float vh; + float roughness; + float metallic; + vec3 f0; + vec3 c_diff; +}; + +/////////////////////////////////////////////////////////////////////////////// +// Uniforms +/////////////////////////////////////////////////////////////////////////////// +uniform Material material; +uniform PointLight point_lights[MAX_POINT_LIGHTS]; +uniform int n_point_lights; +uniform DirectionalLight directional_lights[MAX_DIRECTIONAL_LIGHTS]; +uniform int n_directional_lights; +uniform SpotLight spot_lights[MAX_SPOT_LIGHTS]; +uniform int n_spot_lights; +uniform vec3 cam_pos; +uniform vec3 ambient_light; + +#ifdef USE_IBL +uniform samplerCube diffuse_env; +uniform samplerCube specular_env; +#endif + +/////////////////////////////////////////////////////////////////////////////// +// Inputs +/////////////////////////////////////////////////////////////////////////////// + +in vec3 frag_position; +#ifdef NORMAL_LOC +in vec3 frag_normal; +#endif +#ifdef HAS_NORMAL_TEX +#ifdef TANGENT_LOC +#ifdef NORMAL_LOC +in mat3 tbn; +#endif +#endif +#endif +#ifdef TEXCOORD_0_LOC +in vec2 uv_0; +#endif +#ifdef TEXCOORD_1_LOC +in vec2 uv_1; +#endif +#ifdef COLOR_0_LOC +in vec4 color_multiplier; +#endif + +/////////////////////////////////////////////////////////////////////////////// +// OUTPUTS +/////////////////////////////////////////////////////////////////////////////// + +out vec4 frag_color; + +/////////////////////////////////////////////////////////////////////////////// +// Constants +/////////////////////////////////////////////////////////////////////////////// +const float PI = 3.141592653589793; +const float min_roughness = 0.04; + +/////////////////////////////////////////////////////////////////////////////// +// Utility Functions +/////////////////////////////////////////////////////////////////////////////// +vec4 srgb_to_linear(vec4 srgb) +{ +#ifndef SRGB_CORRECTED + // Fast Approximation + //vec3 linOut = pow(srgbIn.xyz,vec3(2.2)); + // + vec3 b_less = step(vec3(0.04045),srgb.xyz); + vec3 lin_out = mix( srgb.xyz/vec3(12.92), pow((srgb.xyz+vec3(0.055))/vec3(1.055),vec3(2.4)), b_less ); + return vec4(lin_out, srgb.w); +#else + return srgb; +#endif +} + +// Normal computation +vec3 get_normal() +{ +#ifdef HAS_NORMAL_TEX + +#ifndef HAS_TANGENTS + vec3 pos_dx = dFdx(frag_position); + vec3 pos_dy = dFdy(frag_position); + vec3 tex_dx = dFdx(vec3(uv_0, 0.0)); + vec3 tex_dy = dFdy(vec3(uv_0, 0.0)); + vec3 t = (tex_dy.t * pos_dx - tex_dx.t * pos_dy) / (tex_dx.s * tex_dy.t - tex_dy.s * tex_dx.t); + +#ifdef NORMAL_LOC + vec3 ng = normalize(frag_normal); +#else + vec3 = cross(pos_dx, pos_dy); +#endif + + t = normalize(t - ng * dot(ng, t)); + vec3 b = normalize(cross(ng, t)); + mat3 tbn_n = mat3(t, b, ng); + +#else + + mat3 tbn_n = tbn; + +#endif + + vec3 n = texture(material.normal_texture, uv_0).rgb; + n = normalize(tbn_n * ((2.0 * n - 1.0) * vec3(1.0, 1.0, 1.0))); + return n; // TODO NORMAL MAPPING + +#else + +#ifdef NORMAL_LOC + return frag_normal; +#else + return normalize(cam_pos - frag_position); +#endif + +#endif +} + +// Fresnel +vec3 specular_reflection(PBRInfo info) +{ + vec3 res = info.f0 + (1.0 - info.f0) * pow(clamp(1.0 - info.vh, 0.0, 1.0), 5.0); + return res; +} + +// Smith +float geometric_occlusion(PBRInfo info) +{ + float r = info.roughness + 1.0; + float k = r * r / 8.0; + float g1 = info.nv / (info.nv * (1.0 - k) + k); + float g2 = info.nl / (info.nl * (1.0 - k) + k); + //float k = info.roughness * sqrt(2.0 / PI); + //float g1 = info.lh / (info.lh * (1.0 - k) + k); + //float g2 = info.nh / (info.nh * (1.0 - k) + k); + return g1 * g2; +} + +float microfacet_distribution(PBRInfo info) +{ + float a = info.roughness * info.roughness; + float a2 = a * a; + float nh2 = info.nh * info.nh; + + float denom = (nh2 * (a2 - 1.0) + 1.0); + return a2 / (PI * denom * denom); +} + +vec3 compute_brdf(vec3 n, vec3 v, vec3 l, + float roughness, float metalness, + vec3 f0, vec3 c_diff, vec3 albedo, + vec3 radiance) +{ + vec3 h = normalize(l+v); + float nl = clamp(dot(n, l), 0.001, 1.0); + float nv = clamp(abs(dot(n, v)), 0.001, 1.0); + float nh = clamp(dot(n, h), 0.0, 1.0); + float lh = clamp(dot(l, h), 0.0, 1.0); + float vh = clamp(dot(v, h), 0.0, 1.0); + + PBRInfo info = PBRInfo(nl, nv, nh, lh, vh, roughness, metalness, f0, c_diff); + + // Compute PBR terms + vec3 F = specular_reflection(info); + float G = geometric_occlusion(info); + float D = microfacet_distribution(info); + + // Compute BRDF + vec3 diffuse_contrib = (1.0 - F) * c_diff / PI; + vec3 spec_contrib = F * G * D / (4.0 * nl * nv + 0.001); + + vec3 color = nl * radiance * (diffuse_contrib + spec_contrib); + return color; +} + +float texture2DCompare(sampler2D depths, vec2 uv, float compare) { + return compare > texture(depths, uv.xy).r ? 1.0 : 0.0; +} + +float texture2DShadowLerp(sampler2D depths, vec2 size, vec2 uv, float compare) { + vec2 texelSize = vec2(1.0)/size; + vec2 f = fract(uv*size+0.5); + vec2 centroidUV = floor(uv*size+0.5)/size; + + float lb = texture2DCompare(depths, centroidUV+texelSize*vec2(0.0, 0.0), compare); + float lt = texture2DCompare(depths, centroidUV+texelSize*vec2(0.0, 1.0), compare); + float rb = texture2DCompare(depths, centroidUV+texelSize*vec2(1.0, 0.0), compare); + float rt = texture2DCompare(depths, centroidUV+texelSize*vec2(1.0, 1.0), compare); + float a = mix(lb, lt, f.y); + float b = mix(rb, rt, f.y); + float c = mix(a, b, f.x); + return c; +} + +float PCF(sampler2D depths, vec2 size, vec2 uv, float compare){ + float result = 0.0; + for(int x=-1; x<=1; x++){ + for(int y=-1; y<=1; y++){ + vec2 off = vec2(x,y)/size; + result += texture2DShadowLerp(depths, size, uv+off, compare); + } + } + return result/9.0; +} + +float shadow_calc(mat4 light_matrix, sampler2D shadow_map, float nl) +{ + // Compute light texture UV coords + vec4 proj_coords = vec4(light_matrix * vec4(frag_position.xyz, 1.0)); + vec3 light_coords = proj_coords.xyz / proj_coords.w; + light_coords = light_coords * 0.5 + 0.5; + float current_depth = light_coords.z; + float bias = max(0.001 * (1.0 - nl), 0.0001) / proj_coords.w; + float compare = (current_depth - bias); + float shadow = PCF(shadow_map, textureSize(shadow_map, 0), light_coords.xy, compare); + if (light_coords.z > 1.0) { + shadow = 0.0; + } + return shadow; +} + +/////////////////////////////////////////////////////////////////////////////// +// MAIN +/////////////////////////////////////////////////////////////////////////////// +void main() +{ + + vec4 color = vec4(vec3(0.0), 1.0); +/////////////////////////////////////////////////////////////////////////////// +// Handle Metallic Materials +/////////////////////////////////////////////////////////////////////////////// +#ifdef USE_METALLIC_MATERIAL + + // Compute metallic/roughness factors + float roughness = material.roughness_factor; + float metallic = material.metallic_factor; +#ifdef HAS_METALLIC_ROUGHNESS_TEX + vec2 mr = texture(material.metallic_roughness_texture, uv_0).rg; + roughness = roughness * mr.r; + metallic = metallic * mr.g; +#endif + roughness = clamp(roughness, min_roughness, 1.0); + metallic = clamp(metallic, 0.0, 1.0); + // In convention, material roughness is perceputal roughness ^ 2 + float alpha_roughness = roughness * roughness; + + // Compute albedo + vec4 base_color = material.base_color_factor; +#ifdef HAS_BASE_COLOR_TEX + base_color = base_color * srgb_to_linear(texture(material.base_color_texture, uv_0)); +#endif + + // Compute specular and diffuse colors + vec3 dialectric_spec = vec3(min_roughness); + vec3 c_diff = mix(vec3(0.0), base_color.rgb * (1 - min_roughness), 1.0 - metallic); + vec3 f0 = mix(dialectric_spec, base_color.rgb, metallic); + + // Compute normal + vec3 n = normalize(get_normal()); + + // Loop over lights + for (int i = 0; i < n_directional_lights; i++) { + vec3 direction = directional_lights[i].direction; + vec3 v = normalize(cam_pos - frag_position); // Vector towards camera + vec3 l = normalize(-1.0 * direction); // Vector towards light + + // Compute attenuation and radiance + float attenuation = directional_lights[i].intensity; + vec3 radiance = attenuation * directional_lights[i].color; + + // Compute outbound color + vec3 res = compute_brdf(n, v, l, roughness, metallic, + f0, c_diff, base_color.rgb, radiance); + + // Compute shadow +#ifdef DIRECTIONAL_LIGHT_SHADOWS + float nl = clamp(dot(n,l), 0.0, 1.0); + float shadow = shadow_calc( + directional_lights[i].light_matrix, + directional_lights[i].shadow_map, + nl + ); + res = res * (1.0 - shadow); +#endif + color.xyz += res; + } + + for (int i = 0; i < n_point_lights; i++) { + vec3 position = point_lights[i].position; + vec3 v = normalize(cam_pos - frag_position); // Vector towards camera + vec3 l = normalize(position - frag_position); // Vector towards light + + // Compute attenuation and radiance + float dist = length(position - frag_position); + float attenuation = point_lights[i].intensity / (dist * dist); + vec3 radiance = attenuation * point_lights[i].color; + + // Compute outbound color + vec3 res = compute_brdf(n, v, l, roughness, metallic, + f0, c_diff, base_color.rgb, radiance); + color.xyz += res; + } + for (int i = 0; i < n_spot_lights; i++) { + vec3 position = spot_lights[i].position; + vec3 v = normalize(cam_pos - frag_position); // Vector towards camera + vec3 l = normalize(position - frag_position); // Vector towards light + + // Compute attenuation and radiance + vec3 direction = spot_lights[i].direction; + float las = spot_lights[i].light_angle_scale; + float lao = spot_lights[i].light_angle_offset; + float dist = length(position - frag_position); + float cd = clamp(dot(direction, -l), 0.0, 1.0); + float attenuation = clamp(cd * las + lao, 0.0, 1.0); + attenuation = attenuation * attenuation * spot_lights[i].intensity; + attenuation = attenuation / (dist * dist); + vec3 radiance = attenuation * spot_lights[i].color; + + // Compute outbound color + vec3 res = compute_brdf(n, v, l, roughness, metallic, + f0, c_diff, base_color.rgb, radiance); +#ifdef SPOT_LIGHT_SHADOWS + float nl = clamp(dot(n,l), 0.0, 1.0); + float shadow = shadow_calc( + spot_lights[i].light_matrix, + spot_lights[i].shadow_map, + nl + ); + res = res * (1.0 - shadow); +#endif + color.xyz += res; + } + color.xyz += base_color.xyz * ambient_light; + + // Calculate lighting from environment +#ifdef USE_IBL + // TODO +#endif + + // Apply occlusion +#ifdef HAS_OCCLUSION_TEX + float ao = texture(material.occlusion_texture, uv_0).r; + color.xyz *= ao; +#endif + + // Apply emissive map + vec3 emissive = material.emissive_factor; +#ifdef HAS_EMISSIVE_TEX + emissive *= srgb_to_linear(texture(material.emissive_texture, uv_0)).rgb; +#endif + color.xyz += emissive * material.emissive_factor; + +#ifdef COLOR_0_LOC + color *= color_multiplier; +#endif + + frag_color = clamp(vec4(pow(color.xyz, vec3(1.0/2.2)), color.a * base_color.a), 0.0, 1.0); + +#else + // TODO GLOSSY MATERIAL BRDF +#endif + +/////////////////////////////////////////////////////////////////////////////// +// Handle Glossy Materials +/////////////////////////////////////////////////////////////////////////////// + +} diff --git a/pyrender/pyrender/shaders/mesh.vert b/pyrender/pyrender/shaders/mesh.vert new file mode 100644 index 0000000000000000000000000000000000000000..cfd241c3544718a261f961c3aa3c03aa13c97761 --- /dev/null +++ b/pyrender/pyrender/shaders/mesh.vert @@ -0,0 +1,86 @@ +#version 330 core + +// Vertex Attributes +layout(location = 0) in vec3 position; +#ifdef NORMAL_LOC +layout(location = NORMAL_LOC) in vec3 normal; +#endif +#ifdef TANGENT_LOC +layout(location = TANGENT_LOC) in vec4 tangent; +#endif +#ifdef TEXCOORD_0_LOC +layout(location = TEXCOORD_0_LOC) in vec2 texcoord_0; +#endif +#ifdef TEXCOORD_1_LOC +layout(location = TEXCOORD_1_LOC) in vec2 texcoord_1; +#endif +#ifdef COLOR_0_LOC +layout(location = COLOR_0_LOC) in vec4 color_0; +#endif +#ifdef JOINTS_0_LOC +layout(location = JOINTS_0_LOC) in vec4 joints_0; +#endif +#ifdef WEIGHTS_0_LOC +layout(location = WEIGHTS_0_LOC) in vec4 weights_0; +#endif +layout(location = INST_M_LOC) in mat4 inst_m; + +// Uniforms +uniform mat4 M; +uniform mat4 V; +uniform mat4 P; + +// Outputs +out vec3 frag_position; +#ifdef NORMAL_LOC +out vec3 frag_normal; +#endif +#ifdef HAS_NORMAL_TEX +#ifdef TANGENT_LOC +#ifdef NORMAL_LOC +out mat3 tbn; +#endif +#endif +#endif +#ifdef TEXCOORD_0_LOC +out vec2 uv_0; +#endif +#ifdef TEXCOORD_1_LOC +out vec2 uv_1; +#endif +#ifdef COLOR_0_LOC +out vec4 color_multiplier; +#endif + + +void main() +{ + gl_Position = P * V * M * inst_m * vec4(position, 1); + frag_position = vec3(M * inst_m * vec4(position, 1.0)); + + mat4 N = transpose(inverse(M * inst_m)); + +#ifdef NORMAL_LOC + frag_normal = normalize(vec3(N * vec4(normal, 0.0))); +#endif + +#ifdef HAS_NORMAL_TEX +#ifdef TANGENT_LOC +#ifdef NORMAL_LOC + vec3 normal_w = normalize(vec3(N * vec4(normal, 0.0))); + vec3 tangent_w = normalize(vec3(N * vec4(tangent.xyz, 0.0))); + vec3 bitangent_w = cross(normal_w, tangent_w) * tangent.w; + tbn = mat3(tangent_w, bitangent_w, normal_w); +#endif +#endif +#endif +#ifdef TEXCOORD_0_LOC + uv_0 = texcoord_0; +#endif +#ifdef TEXCOORD_1_LOC + uv_1 = texcoord_1; +#endif +#ifdef COLOR_0_LOC + color_multiplier = color_0; +#endif +} diff --git a/pyrender/pyrender/shaders/mesh_depth.frag b/pyrender/pyrender/shaders/mesh_depth.frag new file mode 100644 index 0000000000000000000000000000000000000000..d8b1fac6091cfa457ba835ae0758e955f06d8754 --- /dev/null +++ b/pyrender/pyrender/shaders/mesh_depth.frag @@ -0,0 +1,8 @@ +#version 330 core + +out vec4 frag_color; + +void main() +{ + frag_color = vec4(1.0); +} diff --git a/pyrender/pyrender/shaders/mesh_depth.vert b/pyrender/pyrender/shaders/mesh_depth.vert new file mode 100644 index 0000000000000000000000000000000000000000..e534c058fb3e7b0efbec090513d55982db68ccaf --- /dev/null +++ b/pyrender/pyrender/shaders/mesh_depth.vert @@ -0,0 +1,13 @@ +#version 330 core +layout(location = 0) in vec3 position; +layout(location = INST_M_LOC) in mat4 inst_m; + +uniform mat4 P; +uniform mat4 V; +uniform mat4 M; + +void main() +{ + mat4 light_matrix = P * V; + gl_Position = light_matrix * M * inst_m * vec4(position, 1.0); +} diff --git a/pyrender/pyrender/shaders/segmentation.frag b/pyrender/pyrender/shaders/segmentation.frag new file mode 100644 index 0000000000000000000000000000000000000000..40deb92cbdef3ec9fd952632624cd5f4b5ce0c84 --- /dev/null +++ b/pyrender/pyrender/shaders/segmentation.frag @@ -0,0 +1,13 @@ +#version 330 core + +uniform vec3 color; +out vec4 frag_color; + +/////////////////////////////////////////////////////////////////////////////// +// MAIN +/////////////////////////////////////////////////////////////////////////////// +void main() +{ + frag_color = vec4(color, 1.0); + //frag_color = vec4(1.0, 0.5, 0.5, 1.0); +} diff --git a/pyrender/pyrender/shaders/segmentation.vert b/pyrender/pyrender/shaders/segmentation.vert new file mode 100644 index 0000000000000000000000000000000000000000..503382599dae3c9415845f35b99d6678cfc7f716 --- /dev/null +++ b/pyrender/pyrender/shaders/segmentation.vert @@ -0,0 +1,14 @@ +#version 330 core +layout(location = 0) in vec3 position; +layout(location = INST_M_LOC) in mat4 inst_m; + +uniform mat4 P; +uniform mat4 V; +uniform mat4 M; + +void main() +{ + mat4 light_matrix = P * V; + gl_Position = light_matrix * M * inst_m * vec4(position, 1.0); +} + diff --git a/pyrender/pyrender/shaders/text.frag b/pyrender/pyrender/shaders/text.frag new file mode 100644 index 0000000000000000000000000000000000000000..486c97dc94ed5e9083ae348bc1e85c5cb26c44dc --- /dev/null +++ b/pyrender/pyrender/shaders/text.frag @@ -0,0 +1,12 @@ +#version 330 core +in vec2 uv; +out vec4 color; + +uniform sampler2D text; +uniform vec4 text_color; + +void main() +{ + vec4 sampled = vec4(1.0, 1.0, 1.0, texture(text, uv).r); + color = text_color * sampled; +} diff --git a/pyrender/pyrender/shaders/text.vert b/pyrender/pyrender/shaders/text.vert new file mode 100644 index 0000000000000000000000000000000000000000..005bc439b3d63522df99e5db2088953eb8defcf4 --- /dev/null +++ b/pyrender/pyrender/shaders/text.vert @@ -0,0 +1,12 @@ +#version 330 core +layout (location = 0) in vec4 vertex; + +out vec2 uv; + +uniform mat4 projection; + +void main() +{ + gl_Position = projection * vec4(vertex.xy, 0.0, 1.0); + uv = vertex.zw; +} diff --git a/pyrender/pyrender/shaders/vertex_normals.frag b/pyrender/pyrender/shaders/vertex_normals.frag new file mode 100644 index 0000000000000000000000000000000000000000..edf5beb7f283dd67e1710bff922555539966cee4 --- /dev/null +++ b/pyrender/pyrender/shaders/vertex_normals.frag @@ -0,0 +1,10 @@ +#version 330 core + +out vec4 frag_color; + +uniform vec4 normal_color; + +void main() +{ + frag_color = normal_color; +} diff --git a/pyrender/pyrender/shaders/vertex_normals.geom b/pyrender/pyrender/shaders/vertex_normals.geom new file mode 100644 index 0000000000000000000000000000000000000000..57f0b0e645e72d41116f5767d66fc37d01ed2714 --- /dev/null +++ b/pyrender/pyrender/shaders/vertex_normals.geom @@ -0,0 +1,74 @@ +#version 330 core + +layout (triangles) in; + +#ifdef FACE_NORMALS + +#ifdef VERTEX_NORMALS + layout (line_strip, max_vertices = 8) out; +#else + layout (line_strip, max_vertices = 2) out; +#endif + +#else + + layout (line_strip, max_vertices = 6) out; + +#endif + +in VS_OUT { + vec3 position; + vec3 normal; + mat4 mvp; +} gs_in[]; + +uniform float normal_magnitude; + +void GenerateVertNormal(int index) +{ + + vec4 p0 = gs_in[index].mvp * vec4(gs_in[index].position, 1.0); + vec4 p1 = gs_in[index].mvp * vec4(normal_magnitude * normalize(gs_in[index].normal) + gs_in[index].position, 1.0); + gl_Position = p0; + EmitVertex(); + gl_Position = p1; + EmitVertex(); + EndPrimitive(); +} + +void GenerateFaceNormal() +{ + vec3 p0 = gs_in[0].position.xyz; + vec3 p1 = gs_in[1].position.xyz; + vec3 p2 = gs_in[2].position.xyz; + + vec3 v0 = p0 - p1; + vec3 v1 = p2 - p1; + + vec3 N = normalize(cross(v1, v0)); + vec3 P = (p0 + p1 + p2) / 3.0; + + vec4 np0 = gs_in[0].mvp * vec4(P, 1.0); + vec4 np1 = gs_in[0].mvp * vec4(normal_magnitude * N + P, 1.0); + + gl_Position = np0; + EmitVertex(); + gl_Position = np1; + EmitVertex(); + EndPrimitive(); +} + +void main() +{ + +#ifdef FACE_NORMALS + GenerateFaceNormal(); +#endif + +#ifdef VERTEX_NORMALS + GenerateVertNormal(0); + GenerateVertNormal(1); + GenerateVertNormal(2); +#endif + +} diff --git a/pyrender/pyrender/shaders/vertex_normals.vert b/pyrender/pyrender/shaders/vertex_normals.vert new file mode 100644 index 0000000000000000000000000000000000000000..be22eed2a0e904bcaf1ac5a4721558e574cddc62 --- /dev/null +++ b/pyrender/pyrender/shaders/vertex_normals.vert @@ -0,0 +1,27 @@ +#version 330 core + +// Inputs +layout(location = 0) in vec3 position; +layout(location = NORMAL_LOC) in vec3 normal; +layout(location = INST_M_LOC) in mat4 inst_m; + +// Output data +out VS_OUT { + vec3 position; + vec3 normal; + mat4 mvp; +} vs_out; + +// Uniform data +uniform mat4 M; +uniform mat4 V; +uniform mat4 P; + +// Render loop +void main() { + vs_out.mvp = P * V * M * inst_m; + vs_out.position = position; + vs_out.normal = normal; + + gl_Position = vec4(position, 1.0); +} diff --git a/pyrender/pyrender/shaders/vertex_normals_pc.geom b/pyrender/pyrender/shaders/vertex_normals_pc.geom new file mode 100644 index 0000000000000000000000000000000000000000..4ea4e7b8542703f64b8d28fd187e425137861fe4 --- /dev/null +++ b/pyrender/pyrender/shaders/vertex_normals_pc.geom @@ -0,0 +1,29 @@ +#version 330 core + +layout (points) in; + +layout (line_strip, max_vertices = 2) out; + +in VS_OUT { + vec3 position; + vec3 normal; + mat4 mvp; +} gs_in[]; + +uniform float normal_magnitude; + +void GenerateVertNormal(int index) +{ + vec4 p0 = gs_in[index].mvp * vec4(gs_in[index].position, 1.0); + vec4 p1 = gs_in[index].mvp * vec4(normal_magnitude * normalize(gs_in[index].normal) + gs_in[index].position, 1.0); + gl_Position = p0; + EmitVertex(); + gl_Position = p1; + EmitVertex(); + EndPrimitive(); +} + +void main() +{ + GenerateVertNormal(0); +} diff --git a/pyrender/pyrender/texture.py b/pyrender/pyrender/texture.py new file mode 100644 index 0000000000000000000000000000000000000000..477759729d7b995a4f276e81d649617d045a066e --- /dev/null +++ b/pyrender/pyrender/texture.py @@ -0,0 +1,259 @@ +"""Textures, conforming to the glTF 2.0 standards as specified in +https://github.com/KhronosGroup/glTF/tree/master/specification/2.0#reference-texture + +Author: Matthew Matl +""" +import numpy as np + +from OpenGL.GL import * + +from .utils import format_texture_source +from .sampler import Sampler + + +class Texture(object): + """A texture and its sampler. + + Parameters + ---------- + name : str, optional + The user-defined name of this object. + sampler : :class:`Sampler` + The sampler used by this texture. + source : (h,w,c) uint8 or (h,w,c) float or :class:`PIL.Image.Image` + The image used by this texture. If None, the texture is created + empty and width and height must be specified. + source_channels : str + Either `D`, `R`, `RG`, `GB`, `RGB`, or `RGBA`. Indicates the + channels to extract from `source`. Any missing channels will be filled + with `1.0`. + width : int, optional + For empty textures, the width of the texture buffer. + height : int, optional + For empty textures, the height of the texture buffer. + tex_type : int + Either GL_TEXTURE_2D or GL_TEXTURE_CUBE. + data_format : int + For now, just GL_FLOAT. + """ + + def __init__(self, + name=None, + sampler=None, + source=None, + source_channels=None, + width=None, + height=None, + tex_type=GL_TEXTURE_2D, + data_format=GL_UNSIGNED_BYTE): + self.source_channels = source_channels + self.name = name + self.sampler = sampler + self.source = source + self.width = width + self.height = height + self.tex_type = tex_type + self.data_format = data_format + + self._texid = None + self._is_transparent = False + + @property + def name(self): + """str : The user-defined name of this object. + """ + return self._name + + @name.setter + def name(self, value): + if value is not None: + value = str(value) + self._name = value + + @property + def sampler(self): + """:class:`Sampler` : The sampler used by this texture. + """ + return self._sampler + + @sampler.setter + def sampler(self, value): + if value is None: + value = Sampler() + self._sampler = value + + @property + def source(self): + """(h,w,c) uint8 or float or :class:`PIL.Image.Image` : The image + used in this texture. + """ + return self._source + + @source.setter + def source(self, value): + if value is None: + self._source = None + else: + self._source = format_texture_source(value, self.source_channels) + self._is_transparent = False + + @property + def source_channels(self): + """str : The channels that were extracted from the original source. + """ + return self._source_channels + + @source_channels.setter + def source_channels(self, value): + self._source_channels = value + + @property + def width(self): + """int : The width of the texture buffer. + """ + return self._width + + @width.setter + def width(self, value): + self._width = value + + @property + def height(self): + """int : The height of the texture buffer. + """ + return self._height + + @height.setter + def height(self, value): + self._height = value + + @property + def tex_type(self): + """int : The type of the texture. + """ + return self._tex_type + + @tex_type.setter + def tex_type(self, value): + self._tex_type = value + + @property + def data_format(self): + """int : The format of the texture data. + """ + return self._data_format + + @data_format.setter + def data_format(self, value): + self._data_format = value + + def is_transparent(self, cutoff=1.0): + """bool : If True, the texture is partially transparent. + """ + if self._is_transparent is None: + self._is_transparent = False + if self.source_channels == 'RGBA' and self.source is not None: + if np.any(self.source[:,:,3] < cutoff): + self._is_transparent = True + return self._is_transparent + + def delete(self): + """Remove this texture from the OpenGL context. + """ + self._unbind() + self._remove_from_context() + + ################## + # OpenGL code + ################## + def _add_to_context(self): + if self._texid is not None: + raise ValueError('Texture already loaded into OpenGL context') + + fmt = GL_DEPTH_COMPONENT + if self.source_channels == 'R': + fmt = GL_RED + elif self.source_channels == 'RG' or self.source_channels == 'GB': + fmt = GL_RG + elif self.source_channels == 'RGB': + fmt = GL_RGB + elif self.source_channels == 'RGBA': + fmt = GL_RGBA + + # Generate the OpenGL texture + self._texid = glGenTextures(1) + glBindTexture(self.tex_type, self._texid) + + # Flip data for OpenGL buffer + data = None + width = self.width + height = self.height + if self.source is not None: + data = np.ascontiguousarray(np.flip(self.source, axis=0).flatten()) + width = self.source.shape[1] + height = self.source.shape[0] + + # Bind texture and generate mipmaps + glTexImage2D( + self.tex_type, 0, fmt, width, height, 0, fmt, + self.data_format, data + ) + if self.source is not None: + glGenerateMipmap(self.tex_type) + + if self.sampler.magFilter is not None: + glTexParameteri( + self.tex_type, GL_TEXTURE_MAG_FILTER, self.sampler.magFilter + ) + else: + if self.source is not None: + glTexParameteri(self.tex_type, GL_TEXTURE_MAG_FILTER, GL_LINEAR) + else: + glTexParameteri(self.tex_type, GL_TEXTURE_MAG_FILTER, GL_NEAREST) + if self.sampler.minFilter is not None: + glTexParameteri( + self.tex_type, GL_TEXTURE_MIN_FILTER, self.sampler.minFilter + ) + else: + if self.source is not None: + glTexParameteri(self.tex_type, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR) + else: + glTexParameteri(self.tex_type, GL_TEXTURE_MIN_FILTER, GL_NEAREST) + + glTexParameteri(self.tex_type, GL_TEXTURE_WRAP_S, self.sampler.wrapS) + glTexParameteri(self.tex_type, GL_TEXTURE_WRAP_T, self.sampler.wrapT) + border_color = 255 * np.ones(4).astype(np.uint8) + if self.data_format == GL_FLOAT: + border_color = np.ones(4).astype(np.float32) + glTexParameterfv( + self.tex_type, GL_TEXTURE_BORDER_COLOR, + border_color + ) + + # Unbind texture + glBindTexture(self.tex_type, 0) + + def _remove_from_context(self): + if self._texid is not None: + # TODO OPENGL BUG? + # glDeleteTextures(1, [self._texid]) + glDeleteTextures([self._texid]) + self._texid = None + + def _in_context(self): + return self._texid is not None + + def _bind(self): + # TODO HANDLE INDEXING INTO OTHER UV's + glBindTexture(self.tex_type, self._texid) + + def _unbind(self): + glBindTexture(self.tex_type, 0) + + def _bind_as_depth_attachment(self): + glFramebufferTexture2D(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, + self.tex_type, self._texid, 0) + + def _bind_as_color_attachment(self): + glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, + self.tex_type, self._texid, 0) diff --git a/pyrender/pyrender/trackball.py b/pyrender/pyrender/trackball.py new file mode 100644 index 0000000000000000000000000000000000000000..3e57a0e82d3f07b80754f575c28a0e05cb73fc50 --- /dev/null +++ b/pyrender/pyrender/trackball.py @@ -0,0 +1,216 @@ +"""Trackball class for 3D manipulation of viewpoints. +""" +import numpy as np + +import trimesh.transformations as transformations + + +class Trackball(object): + """A trackball class for creating camera transforms from mouse movements. + """ + STATE_ROTATE = 0 + STATE_PAN = 1 + STATE_ROLL = 2 + STATE_ZOOM = 3 + + def __init__(self, pose, size, scale, + target=np.array([0.0, 0.0, 0.0])): + """Initialize a trackball with an initial camera-to-world pose + and the given parameters. + + Parameters + ---------- + pose : [4,4] + An initial camera-to-world pose for the trackball. + + size : (float, float) + The width and height of the camera image in pixels. + + scale : float + The diagonal of the scene's bounding box -- + used for ensuring translation motions are sufficiently + fast for differently-sized scenes. + + target : (3,) float + The center of the scene in world coordinates. + The trackball will revolve around this point. + """ + self._size = np.array(size) + self._scale = float(scale) + + self._pose = pose + self._n_pose = pose + + self._target = target + self._n_target = target + + self._state = Trackball.STATE_ROTATE + + @property + def pose(self): + """autolab_core.RigidTransform : The current camera-to-world pose. + """ + return self._n_pose + + def set_state(self, state): + """Set the state of the trackball in order to change the effect of + dragging motions. + + Parameters + ---------- + state : int + One of Trackball.STATE_ROTATE, Trackball.STATE_PAN, + Trackball.STATE_ROLL, and Trackball.STATE_ZOOM. + """ + self._state = state + + def resize(self, size): + """Resize the window. + + Parameters + ---------- + size : (float, float) + The new width and height of the camera image in pixels. + """ + self._size = np.array(size) + + def down(self, point): + """Record an initial mouse press at a given point. + + Parameters + ---------- + point : (2,) int + The x and y pixel coordinates of the mouse press. + """ + self._pdown = np.array(point, dtype=np.float32) + self._pose = self._n_pose + self._target = self._n_target + + def drag(self, point): + """Update the tracball during a drag. + + Parameters + ---------- + point : (2,) int + The current x and y pixel coordinates of the mouse during a drag. + This will compute a movement for the trackball with the relative + motion between this point and the one marked by down(). + """ + point = np.array(point, dtype=np.float32) + dx, dy = point - self._pdown + mindim = 0.3 * np.min(self._size) + + target = self._target + x_axis = self._pose[:3,0].flatten() + y_axis = self._pose[:3,1].flatten() + z_axis = self._pose[:3,2].flatten() + eye = self._pose[:3,3].flatten() + + # Interpret drag as a rotation + if self._state == Trackball.STATE_ROTATE: + x_angle = -dx / mindim + x_rot_mat = transformations.rotation_matrix( + x_angle, y_axis, target + ) + + y_angle = dy / mindim + y_rot_mat = transformations.rotation_matrix( + y_angle, x_axis, target + ) + + self._n_pose = y_rot_mat.dot(x_rot_mat.dot(self._pose)) + + # Interpret drag as a roll about the camera axis + elif self._state == Trackball.STATE_ROLL: + center = self._size / 2.0 + v_init = self._pdown - center + v_curr = point - center + v_init = v_init / np.linalg.norm(v_init) + v_curr = v_curr / np.linalg.norm(v_curr) + + theta = (-np.arctan2(v_curr[1], v_curr[0]) + + np.arctan2(v_init[1], v_init[0])) + + rot_mat = transformations.rotation_matrix(theta, z_axis, target) + + self._n_pose = rot_mat.dot(self._pose) + + # Interpret drag as a camera pan in view plane + elif self._state == Trackball.STATE_PAN: + dx = -dx / (5.0 * mindim) * self._scale + dy = -dy / (5.0 * mindim) * self._scale + + translation = dx * x_axis + dy * y_axis + self._n_target = self._target + translation + t_tf = np.eye(4) + t_tf[:3,3] = translation + self._n_pose = t_tf.dot(self._pose) + + # Interpret drag as a zoom motion + elif self._state == Trackball.STATE_ZOOM: + radius = np.linalg.norm(eye - target) + ratio = 0.0 + if dy > 0: + ratio = np.exp(abs(dy) / (0.5 * self._size[1])) - 1.0 + elif dy < 0: + ratio = 1.0 - np.exp(dy / (0.5 * (self._size[1]))) + translation = -np.sign(dy) * ratio * radius * z_axis + t_tf = np.eye(4) + t_tf[:3,3] = translation + self._n_pose = t_tf.dot(self._pose) + + def scroll(self, clicks): + """Zoom using a mouse scroll wheel motion. + + Parameters + ---------- + clicks : int + The number of clicks. Positive numbers indicate forward wheel + movement. + """ + target = self._target + ratio = 0.90 + + mult = 1.0 + if clicks > 0: + mult = ratio**clicks + elif clicks < 0: + mult = (1.0 / ratio)**abs(clicks) + + z_axis = self._n_pose[:3,2].flatten() + eye = self._n_pose[:3,3].flatten() + radius = np.linalg.norm(eye - target) + translation = (mult * radius - radius) * z_axis + t_tf = np.eye(4) + t_tf[:3,3] = translation + self._n_pose = t_tf.dot(self._n_pose) + + z_axis = self._pose[:3,2].flatten() + eye = self._pose[:3,3].flatten() + radius = np.linalg.norm(eye - target) + translation = (mult * radius - radius) * z_axis + t_tf = np.eye(4) + t_tf[:3,3] = translation + self._pose = t_tf.dot(self._pose) + + def rotate(self, azimuth, axis=None): + """Rotate the trackball about the "Up" axis by azimuth radians. + + Parameters + ---------- + azimuth : float + The number of radians to rotate. + """ + target = self._target + + y_axis = self._n_pose[:3,1].flatten() + if axis is not None: + y_axis = axis + x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target) + self._n_pose = x_rot_mat.dot(self._n_pose) + + y_axis = self._pose[:3,1].flatten() + if axis is not None: + y_axis = axis + x_rot_mat = transformations.rotation_matrix(azimuth, y_axis, target) + self._pose = x_rot_mat.dot(self._pose) diff --git a/pyrender/pyrender/utils.py b/pyrender/pyrender/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..48a11faf991606ad7fb0691582f0bc6f06101a45 --- /dev/null +++ b/pyrender/pyrender/utils.py @@ -0,0 +1,115 @@ +import numpy as np +from PIL import Image + + +def format_color_vector(value, length): + """Format a color vector. + """ + if isinstance(value, int): + value = value / 255.0 + if isinstance(value, float): + value = np.repeat(value, length) + if isinstance(value, list) or isinstance(value, tuple): + value = np.array(value) + if isinstance(value, np.ndarray): + value = value.squeeze() + if np.issubdtype(value.dtype, np.integer): + value = (value / 255.0).astype(np.float32) + if value.ndim != 1: + raise ValueError('Format vector takes only 1-D vectors') + if length > value.shape[0]: + value = np.hstack((value, np.ones(length - value.shape[0]))) + elif length < value.shape[0]: + value = value[:length] + else: + raise ValueError('Invalid vector data type') + + return value.squeeze().astype(np.float32) + + +def format_color_array(value, shape): + """Format an array of colors. + """ + # Convert uint8 to floating + value = np.asanyarray(value) + if np.issubdtype(value.dtype, np.integer): + value = (value / 255.0).astype(np.float32) + + # Match up shapes + if value.ndim == 1: + value = np.tile(value, (shape[0],1)) + if value.shape[1] < shape[1]: + nc = shape[1] - value.shape[1] + value = np.column_stack((value, np.ones((value.shape[0], nc)))) + elif value.shape[1] > shape[1]: + value = value[:,:shape[1]] + return value.astype(np.float32) + + +def format_texture_source(texture, target_channels='RGB'): + """Format a texture as a float32 np array. + """ + + # Pass through None + if texture is None: + return None + + # Convert PIL images into numpy arrays + if isinstance(texture, Image.Image): + if texture.mode == 'P' and target_channels in ('RGB', 'RGBA'): + texture = np.array(texture.convert(target_channels)) + else: + texture = np.array(texture) + + # Format numpy arrays + if isinstance(texture, np.ndarray): + if np.issubdtype(texture.dtype, np.floating): + texture = np.array(texture * 255.0, dtype=np.uint8) + elif np.issubdtype(texture.dtype, np.integer): + texture = texture.astype(np.uint8) + else: + raise TypeError('Invalid type {} for texture'.format( + type(texture) + )) + + # Format array by picking out correct texture channels or padding + if texture.ndim == 2: + texture = texture[:,:,np.newaxis] + if target_channels == 'R': + texture = texture[:,:,0] + texture = texture.squeeze() + elif target_channels == 'RG': + if texture.shape[2] == 1: + texture = np.repeat(texture, 2, axis=2) + else: + texture = texture[:,:,(0,1)] + elif target_channels == 'GB': + if texture.shape[2] == 1: + texture = np.repeat(texture, 2, axis=2) + elif texture.shape[2] > 2: + texture = texture[:,:,(1,2)] + elif target_channels == 'RGB': + if texture.shape[2] == 1: + texture = np.repeat(texture, 3, axis=2) + elif texture.shape[2] == 2: + raise ValueError('Cannot reformat 2-channel texture into RGB') + else: + texture = texture[:,:,(0,1,2)] + elif target_channels == 'RGBA': + if texture.shape[2] == 1: + texture = np.repeat(texture, 4, axis=2) + texture[:,:,3] = 255 + elif texture.shape[2] == 2: + raise ValueError('Cannot reformat 2-channel texture into RGBA') + elif texture.shape[2] == 3: + tx = np.empty((texture.shape[0], texture.shape[1], 4), dtype=np.uint8) + tx[:,:,:3] = texture + tx[:,:,3] = 255 + texture = tx + else: + raise ValueError('Invalid texture channel specification: {}' + .format(target_channels)) + else: + raise TypeError('Invalid type {} for texture'.format(type(texture))) + + return texture diff --git a/pyrender/pyrender/version.py b/pyrender/pyrender/version.py new file mode 100644 index 0000000000000000000000000000000000000000..a33fc87f61f528780e3319a5160769cc84512b1b --- /dev/null +++ b/pyrender/pyrender/version.py @@ -0,0 +1 @@ +__version__ = '0.1.45' diff --git a/pyrender/pyrender/viewer.py b/pyrender/pyrender/viewer.py new file mode 100644 index 0000000000000000000000000000000000000000..d2326c38205c6eaddb4f567e3b088329187af258 --- /dev/null +++ b/pyrender/pyrender/viewer.py @@ -0,0 +1,1160 @@ +"""A pyglet-based interactive 3D scene viewer. +""" +import copy +import os +import sys +from threading import Thread, RLock +import time + +import imageio +import numpy as np +import OpenGL +import trimesh + +try: + from Tkinter import Tk, tkFileDialog as filedialog +except Exception: + try: + from tkinter import Tk, filedialog as filedialog + except Exception: + pass + +from .constants import (TARGET_OPEN_GL_MAJOR, TARGET_OPEN_GL_MINOR, + MIN_OPEN_GL_MAJOR, MIN_OPEN_GL_MINOR, + TEXT_PADDING, DEFAULT_SCENE_SCALE, + DEFAULT_Z_FAR, DEFAULT_Z_NEAR, RenderFlags, TextAlign) +from .light import DirectionalLight +from .node import Node +from .camera import PerspectiveCamera, OrthographicCamera, IntrinsicsCamera +from .trackball import Trackball +from .renderer import Renderer +from .mesh import Mesh + +import pyglet +from pyglet import clock +pyglet.options['shadow_window'] = False + + +class Viewer(pyglet.window.Window): + """An interactive viewer for 3D scenes. + + The viewer's camera is separate from the scene's, but will take on + the parameters of the scene's main view camera and start in the same pose. + If the scene does not have a camera, a suitable default will be provided. + + Parameters + ---------- + scene : :class:`Scene` + The scene to visualize. + viewport_size : (2,) int + The width and height of the initial viewing window. + render_flags : dict + A set of flags for rendering the scene. Described in the note below. + viewer_flags : dict + A set of flags for controlling the viewer's behavior. + Described in the note below. + registered_keys : dict + A map from ASCII key characters to tuples containing: + + - A function to be called whenever the key is pressed, + whose first argument will be the viewer itself. + - (Optionally) A list of additional positional arguments + to be passed to the function. + - (Optionally) A dict of keyword arguments to be passed + to the function. + + kwargs : dict + Any keyword arguments left over will be interpreted as belonging to + either the :attr:`.Viewer.render_flags` or :attr:`.Viewer.viewer_flags` + dictionaries. Those flag sets will be updated appropriately. + + Note + ---- + The basic commands for moving about the scene are given as follows: + + - **Rotating about the scene**: Hold the left mouse button and + drag the cursor. + - **Rotating about the view axis**: Hold ``CTRL`` and the left mouse + button and drag the cursor. + - **Panning**: + + - Hold SHIFT, then hold the left mouse button and drag the cursor, or + - Hold the middle mouse button and drag the cursor. + + - **Zooming**: + + - Scroll the mouse wheel, or + - Hold the right mouse button and drag the cursor. + + Other keyboard commands are as follows: + + - ``a``: Toggles rotational animation mode. + - ``c``: Toggles backface culling. + - ``f``: Toggles fullscreen mode. + - ``h``: Toggles shadow rendering. + - ``i``: Toggles axis display mode + (no axes, world axis, mesh axes, all axes). + - ``l``: Toggles lighting mode + (scene lighting, Raymond lighting, or direct lighting). + - ``m``: Toggles face normal visualization. + - ``n``: Toggles vertex normal visualization. + - ``o``: Toggles orthographic mode. + - ``q``: Quits the viewer. + - ``r``: Starts recording a GIF, and pressing again stops recording + and opens a file dialog. + - ``s``: Opens a file dialog to save the current view as an image. + - ``w``: Toggles wireframe mode + (scene default, flip wireframes, all wireframe, or all solid). + - ``z``: Resets the camera to the initial view. + + Note + ---- + The valid keys for ``render_flags`` are as follows: + + - ``flip_wireframe``: `bool`, If `True`, all objects will have their + wireframe modes flipped from what their material indicates. + Defaults to `False`. + - ``all_wireframe``: `bool`, If `True`, all objects will be rendered + in wireframe mode. Defaults to `False`. + - ``all_solid``: `bool`, If `True`, all objects will be rendered in + solid mode. Defaults to `False`. + - ``shadows``: `bool`, If `True`, shadows will be rendered. + Defaults to `False`. + - ``vertex_normals``: `bool`, If `True`, vertex normals will be + rendered as blue lines. Defaults to `False`. + - ``face_normals``: `bool`, If `True`, face normals will be rendered as + blue lines. Defaults to `False`. + - ``cull_faces``: `bool`, If `True`, backfaces will be culled. + Defaults to `True`. + - ``point_size`` : float, The point size in pixels. Defaults to 1px. + + Note + ---- + The valid keys for ``viewer_flags`` are as follows: + + - ``rotate``: `bool`, If `True`, the scene's camera will rotate + about an axis. Defaults to `False`. + - ``rotate_rate``: `float`, The rate of rotation in radians per second. + Defaults to `PI / 3.0`. + - ``rotate_axis``: `(3,) float`, The axis in world coordinates to rotate + about. Defaults to ``[0,0,1]``. + - ``view_center``: `(3,) float`, The position to rotate the scene about. + Defaults to the scene's centroid. + - ``use_raymond_lighting``: `bool`, If `True`, an additional set of three + directional lights that move with the camera will be added to the scene. + Defaults to `False`. + - ``use_direct_lighting``: `bool`, If `True`, an additional directional + light that moves with the camera and points out of it will be added to + the scene. Defaults to `False`. + - ``lighting_intensity``: `float`, The overall intensity of the + viewer's additional lights (when they're in use). Defaults to 3.0. + - ``use_perspective_cam``: `bool`, If `True`, a perspective camera will + be used. Otherwise, an orthographic camera is used. Defaults to `True`. + - ``save_directory``: `str`, A directory to open the file dialogs in. + Defaults to `None`. + - ``window_title``: `str`, A title for the viewer's application window. + Defaults to `"Scene Viewer"`. + - ``refresh_rate``: `float`, A refresh rate for rendering, in Hertz. + Defaults to `30.0`. + - ``fullscreen``: `bool`, Whether to make viewer fullscreen. + Defaults to `False`. + - ``show_world_axis``: `bool`, Whether to show the world axis. + Defaults to `False`. + - ``show_mesh_axes``: `bool`, Whether to show the individual mesh axes. + Defaults to `False`. + - ``caption``: `list of dict`, Text caption(s) to display on the viewer. + Defaults to `None`. + + Note + ---- + Animation can be accomplished by running the viewer with ``run_in_thread`` + enabled. Then, just run a loop in your main thread, updating the scene as + needed. Before updating the scene, be sure to acquire the + :attr:`.Viewer.render_lock`, and release it when your update is done. + """ + + def __init__(self, scene, viewport_size=None, + render_flags=None, viewer_flags=None, + registered_keys=None, run_in_thread=False, + auto_start=True, + **kwargs): + + ####################################################################### + # Save attributes and flags + ####################################################################### + if viewport_size is None: + viewport_size = (640, 480) + self._scene = scene + self._viewport_size = viewport_size + self._render_lock = RLock() + self._is_active = False + self._should_close = False + self._run_in_thread = run_in_thread + self._auto_start = auto_start + + self._default_render_flags = { + 'flip_wireframe': False, + 'all_wireframe': False, + 'all_solid': False, + 'shadows': False, + 'vertex_normals': False, + 'face_normals': False, + 'cull_faces': True, + 'point_size': 1.0, + } + self._default_viewer_flags = { + 'mouse_pressed': False, + 'rotate': False, + 'rotate_rate': np.pi / 3.0, + 'rotate_axis': np.array([0.0, 0.0, 1.0]), + 'view_center': None, + 'record': False, + 'use_raymond_lighting': False, + 'use_direct_lighting': False, + 'lighting_intensity': 3.0, + 'use_perspective_cam': True, + 'save_directory': None, + 'window_title': 'Scene Viewer', + 'refresh_rate': 30.0, + 'fullscreen': False, + 'show_world_axis': False, + 'show_mesh_axes': False, + 'caption': None + } + self._render_flags = self._default_render_flags.copy() + self._viewer_flags = self._default_viewer_flags.copy() + self._viewer_flags['rotate_axis'] = ( + self._default_viewer_flags['rotate_axis'].copy() + ) + + if render_flags is not None: + self._render_flags.update(render_flags) + if viewer_flags is not None: + self._viewer_flags.update(viewer_flags) + + for key in kwargs: + if key in self.render_flags: + self._render_flags[key] = kwargs[key] + elif key in self.viewer_flags: + self._viewer_flags[key] = kwargs[key] + + # TODO MAC OS BUG FOR SHADOWS + if sys.platform == 'darwin': + self._render_flags['shadows'] = False + + self._registered_keys = {} + if registered_keys is not None: + self._registered_keys = { + ord(k.lower()): registered_keys[k] for k in registered_keys + } + + ####################################################################### + # Save internal settings + ####################################################################### + + # Set up caption stuff + self._message_text = None + self._ticks_till_fade = 2.0 / 3.0 * self.viewer_flags['refresh_rate'] + self._message_opac = 1.0 + self._ticks_till_fade + + # Set up raymond lights and direct lights + self._raymond_lights = self._create_raymond_lights() + self._direct_light = self._create_direct_light() + + # Set up axes + self._axes = {} + self._axis_mesh = Mesh.from_trimesh( + trimesh.creation.axis(origin_size=0.1, axis_radius=0.05, + axis_length=1.0), smooth=False) + if self.viewer_flags['show_world_axis']: + self._set_axes(world=self.viewer_flags['show_world_axis'], + mesh=self.viewer_flags['show_mesh_axes']) + + ####################################################################### + # Set up camera node + ####################################################################### + self._camera_node = None + self._prior_main_camera_node = None + self._default_camera_pose = None + self._default_persp_cam = None + self._default_orth_cam = None + self._trackball = None + self._saved_frames = [] + + # Extract main camera from scene and set up our mirrored copy + znear = None + zfar = None + if scene.main_camera_node is not None: + n = scene.main_camera_node + camera = copy.copy(n.camera) + if isinstance(camera, (PerspectiveCamera, IntrinsicsCamera)): + self._default_persp_cam = camera + znear = camera.znear + zfar = camera.zfar + elif isinstance(camera, OrthographicCamera): + self._default_orth_cam = camera + znear = camera.znear + zfar = camera.zfar + self._default_camera_pose = scene.get_pose(scene.main_camera_node) + self._prior_main_camera_node = n + + # Set defaults as needed + if zfar is None: + zfar = max(scene.scale * 10.0, DEFAULT_Z_FAR) + if znear is None or znear == 0: + if scene.scale == 0: + znear = DEFAULT_Z_NEAR + else: + znear = min(scene.scale / 10.0, DEFAULT_Z_NEAR) + + if self._default_persp_cam is None: + self._default_persp_cam = PerspectiveCamera( + yfov=np.pi / 3.0, znear=znear, zfar=zfar + ) + if self._default_orth_cam is None: + xmag = ymag = scene.scale + if scene.scale == 0: + xmag = ymag = 1.0 + self._default_orth_cam = OrthographicCamera( + xmag=xmag, ymag=ymag, + znear=znear, + zfar=zfar + ) + if self._default_camera_pose is None: + self._default_camera_pose = self._compute_initial_camera_pose() + + # Pick camera + if self.viewer_flags['use_perspective_cam']: + camera = self._default_persp_cam + else: + camera = self._default_orth_cam + + self._camera_node = Node( + matrix=self._default_camera_pose, camera=camera + ) + scene.add_node(self._camera_node) + scene.main_camera_node = self._camera_node + self._reset_view() + + ####################################################################### + # Initialize OpenGL context and renderer + ####################################################################### + self._renderer = Renderer( + self._viewport_size[0], self._viewport_size[1], + self.render_flags['point_size'] + ) + self._is_active = True + + if self.run_in_thread: + self._thread = Thread(target=self._init_and_start_app) + self._thread.start() + else: + if auto_start: + self._init_and_start_app() + + def start(self): + self._init_and_start_app() + + @property + def scene(self): + """:class:`.Scene` : The scene being visualized. + """ + return self._scene + + @property + def viewport_size(self): + """(2,) int : The width and height of the viewing window. + """ + return self._viewport_size + + @property + def render_lock(self): + """:class:`threading.RLock` : If acquired, prevents the viewer from + rendering until released. + + Run :meth:`.Viewer.render_lock.acquire` before making updates to + the scene in a different thread, and run + :meth:`.Viewer.render_lock.release` once you're done to let the viewer + continue. + """ + return self._render_lock + + @property + def is_active(self): + """bool : `True` if the viewer is active, or `False` if it has + been closed. + """ + return self._is_active + + @property + def run_in_thread(self): + """bool : Whether the viewer was run in a separate thread. + """ + return self._run_in_thread + + @property + def render_flags(self): + """dict : Flags for controlling the renderer's behavior. + + - ``flip_wireframe``: `bool`, If `True`, all objects will have their + wireframe modes flipped from what their material indicates. + Defaults to `False`. + - ``all_wireframe``: `bool`, If `True`, all objects will be rendered + in wireframe mode. Defaults to `False`. + - ``all_solid``: `bool`, If `True`, all objects will be rendered in + solid mode. Defaults to `False`. + - ``shadows``: `bool`, If `True`, shadows will be rendered. + Defaults to `False`. + - ``vertex_normals``: `bool`, If `True`, vertex normals will be + rendered as blue lines. Defaults to `False`. + - ``face_normals``: `bool`, If `True`, face normals will be rendered as + blue lines. Defaults to `False`. + - ``cull_faces``: `bool`, If `True`, backfaces will be culled. + Defaults to `True`. + - ``point_size`` : float, The point size in pixels. Defaults to 1px. + + """ + return self._render_flags + + @render_flags.setter + def render_flags(self, value): + self._render_flags = value + + @property + def viewer_flags(self): + """dict : Flags for controlling the viewer's behavior. + + The valid keys for ``viewer_flags`` are as follows: + + - ``rotate``: `bool`, If `True`, the scene's camera will rotate + about an axis. Defaults to `False`. + - ``rotate_rate``: `float`, The rate of rotation in radians per second. + Defaults to `PI / 3.0`. + - ``rotate_axis``: `(3,) float`, The axis in world coordinates to + rotate about. Defaults to ``[0,0,1]``. + - ``view_center``: `(3,) float`, The position to rotate the scene + about. Defaults to the scene's centroid. + - ``use_raymond_lighting``: `bool`, If `True`, an additional set of + three directional lights that move with the camera will be added to + the scene. Defaults to `False`. + - ``use_direct_lighting``: `bool`, If `True`, an additional directional + light that moves with the camera and points out of it will be + added to the scene. Defaults to `False`. + - ``lighting_intensity``: `float`, The overall intensity of the + viewer's additional lights (when they're in use). Defaults to 3.0. + - ``use_perspective_cam``: `bool`, If `True`, a perspective camera will + be used. Otherwise, an orthographic camera is used. Defaults to + `True`. + - ``save_directory``: `str`, A directory to open the file dialogs in. + Defaults to `None`. + - ``window_title``: `str`, A title for the viewer's application window. + Defaults to `"Scene Viewer"`. + - ``refresh_rate``: `float`, A refresh rate for rendering, in Hertz. + Defaults to `30.0`. + - ``fullscreen``: `bool`, Whether to make viewer fullscreen. + Defaults to `False`. + - ``show_world_axis``: `bool`, Whether to show the world axis. + Defaults to `False`. + - ``show_mesh_axes``: `bool`, Whether to show the individual mesh axes. + Defaults to `False`. + - ``caption``: `list of dict`, Text caption(s) to display on + the viewer. Defaults to `None`. + + """ + return self._viewer_flags + + @viewer_flags.setter + def viewer_flags(self, value): + self._viewer_flags = value + + @property + def registered_keys(self): + """dict : Map from ASCII key character to a handler function. + + This is a map from ASCII key characters to tuples containing: + + - A function to be called whenever the key is pressed, + whose first argument will be the viewer itself. + - (Optionally) A list of additional positional arguments + to be passed to the function. + - (Optionally) A dict of keyword arguments to be passed + to the function. + + """ + return self._registered_keys + + @registered_keys.setter + def registered_keys(self, value): + self._registered_keys = value + + def close_external(self): + """Close the viewer from another thread. + + This function will wait for the actual close, so you immediately + manipulate the scene afterwards. + """ + self._should_close = True + while self.is_active: + time.sleep(1.0 / self.viewer_flags['refresh_rate']) + + def save_gif(self, filename=None): + """Save the stored GIF frames to a file. + + To use this asynchronously, run the viewer with the ``record`` + flag and the ``run_in_thread`` flags set. + Kill the viewer after your desired time with + :meth:`.Viewer.close_external`, and then call :meth:`.Viewer.save_gif`. + + Parameters + ---------- + filename : str + The file to save the GIF to. If not specified, + a file dialog will be opened to ask the user where + to save the GIF file. + """ + if filename is None: + filename = self._get_save_filename(['gif', 'all']) + if filename is not None: + self.viewer_flags['save_directory'] = os.path.dirname(filename) + imageio.mimwrite(filename, self._saved_frames, + fps=self.viewer_flags['refresh_rate'], + palettesize=128, subrectangles=True) + self._saved_frames = [] + + def on_close(self): + """Exit the event loop when the window is closed. + """ + # Remove our camera and restore the prior one + if self._camera_node is not None: + self.scene.remove_node(self._camera_node) + if self._prior_main_camera_node is not None: + self.scene.main_camera_node = self._prior_main_camera_node + + # Delete any lighting nodes that we've attached + if self.viewer_flags['use_raymond_lighting']: + for n in self._raymond_lights: + if self.scene.has_node(n): + self.scene.remove_node(n) + if self.viewer_flags['use_direct_lighting']: + if self.scene.has_node(self._direct_light): + self.scene.remove_node(self._direct_light) + + # Delete any axis nodes that we've attached + self._remove_axes() + + # Delete renderer + if self._renderer is not None: + self._renderer.delete() + self._renderer = None + + # Force clean-up of OpenGL context data + try: + OpenGL.contextdata.cleanupContext() + self.close() + except Exception: + pass + finally: + self._is_active = False + super(Viewer, self).on_close() + pyglet.app.exit() + + def on_draw(self): + """Redraw the scene into the viewing window. + """ + if self._renderer is None: + return + + if self.run_in_thread or not self._auto_start: + self.render_lock.acquire() + + # Make OpenGL context current + self.switch_to() + + # Render the scene + self.clear() + self._render() + + if self._message_text is not None: + self._renderer.render_text( + self._message_text, + self.viewport_size[0] - TEXT_PADDING, + TEXT_PADDING, + font_pt=20, + color=np.array([0.1, 0.7, 0.2, + np.clip(self._message_opac, 0.0, 1.0)]), + align=TextAlign.BOTTOM_RIGHT + ) + + if self.viewer_flags['caption'] is not None: + for caption in self.viewer_flags['caption']: + xpos, ypos = self._location_to_x_y(caption['location']) + self._renderer.render_text( + caption['text'], + xpos, + ypos, + font_name=caption['font_name'], + font_pt=caption['font_pt'], + color=caption['color'], + scale=caption['scale'], + align=caption['location'] + ) + + if self.run_in_thread or not self._auto_start: + self.render_lock.release() + + def on_resize(self, width, height): + """Resize the camera and trackball when the window is resized. + """ + if self._renderer is None: + return + + self._viewport_size = (width, height) + self._trackball.resize(self._viewport_size) + self._renderer.viewport_width = self._viewport_size[0] + self._renderer.viewport_height = self._viewport_size[1] + self.on_draw() + + def on_mouse_press(self, x, y, buttons, modifiers): + """Record an initial mouse press. + """ + self._trackball.set_state(Trackball.STATE_ROTATE) + if (buttons == pyglet.window.mouse.LEFT): + ctrl = (modifiers & pyglet.window.key.MOD_CTRL) + shift = (modifiers & pyglet.window.key.MOD_SHIFT) + if (ctrl and shift): + self._trackball.set_state(Trackball.STATE_ZOOM) + elif ctrl: + self._trackball.set_state(Trackball.STATE_ROLL) + elif shift: + self._trackball.set_state(Trackball.STATE_PAN) + elif (buttons == pyglet.window.mouse.MIDDLE): + self._trackball.set_state(Trackball.STATE_PAN) + elif (buttons == pyglet.window.mouse.RIGHT): + self._trackball.set_state(Trackball.STATE_ZOOM) + + self._trackball.down(np.array([x, y])) + + # Stop animating while using the mouse + self.viewer_flags['mouse_pressed'] = True + + def on_mouse_drag(self, x, y, dx, dy, buttons, modifiers): + """Record a mouse drag. + """ + self._trackball.drag(np.array([x, y])) + + def on_mouse_release(self, x, y, button, modifiers): + """Record a mouse release. + """ + self.viewer_flags['mouse_pressed'] = False + + def on_mouse_scroll(self, x, y, dx, dy): + """Record a mouse scroll. + """ + if self.viewer_flags['use_perspective_cam']: + self._trackball.scroll(dy) + else: + spfc = 0.95 + spbc = 1.0 / 0.95 + sf = 1.0 + if dy > 0: + sf = spfc * dy + elif dy < 0: + sf = - spbc * dy + + c = self._camera_node.camera + xmag = max(c.xmag * sf, 1e-8) + ymag = max(c.ymag * sf, 1e-8 * c.ymag / c.xmag) + c.xmag = xmag + c.ymag = ymag + + def on_key_press(self, symbol, modifiers): + """Record a key press. + """ + # First, check for registered key callbacks + if symbol in self.registered_keys: + tup = self.registered_keys[symbol] + callback = None + args = [] + kwargs = {} + if not isinstance(tup, (list, tuple, np.ndarray)): + callback = tup + else: + callback = tup[0] + if len(tup) == 2: + args = tup[1] + if len(tup) == 3: + kwargs = tup[2] + callback(self, *args, **kwargs) + return + + # Otherwise, use default key functions + + # A causes the frame to rotate + self._message_text = None + if symbol == pyglet.window.key.A: + self.viewer_flags['rotate'] = not self.viewer_flags['rotate'] + if self.viewer_flags['rotate']: + self._message_text = 'Rotation On' + else: + self._message_text = 'Rotation Off' + + # C toggles backface culling + elif symbol == pyglet.window.key.C: + self.render_flags['cull_faces'] = ( + not self.render_flags['cull_faces'] + ) + if self.render_flags['cull_faces']: + self._message_text = 'Cull Faces On' + else: + self._message_text = 'Cull Faces Off' + + # F toggles face normals + elif symbol == pyglet.window.key.F: + self.viewer_flags['fullscreen'] = ( + not self.viewer_flags['fullscreen'] + ) + self.set_fullscreen(self.viewer_flags['fullscreen']) + self.activate() + if self.viewer_flags['fullscreen']: + self._message_text = 'Fullscreen On' + else: + self._message_text = 'Fullscreen Off' + + # S toggles shadows + elif symbol == pyglet.window.key.H and sys.platform != 'darwin': + self.render_flags['shadows'] = not self.render_flags['shadows'] + if self.render_flags['shadows']: + self._message_text = 'Shadows On' + else: + self._message_text = 'Shadows Off' + + elif symbol == pyglet.window.key.I: + if (self.viewer_flags['show_world_axis'] and not + self.viewer_flags['show_mesh_axes']): + self.viewer_flags['show_world_axis'] = False + self.viewer_flags['show_mesh_axes'] = True + self._set_axes(False, True) + self._message_text = 'Mesh Axes On' + elif (not self.viewer_flags['show_world_axis'] and + self.viewer_flags['show_mesh_axes']): + self.viewer_flags['show_world_axis'] = True + self.viewer_flags['show_mesh_axes'] = True + self._set_axes(True, True) + self._message_text = 'All Axes On' + elif (self.viewer_flags['show_world_axis'] and + self.viewer_flags['show_mesh_axes']): + self.viewer_flags['show_world_axis'] = False + self.viewer_flags['show_mesh_axes'] = False + self._set_axes(False, False) + self._message_text = 'All Axes Off' + else: + self.viewer_flags['show_world_axis'] = True + self.viewer_flags['show_mesh_axes'] = False + self._set_axes(True, False) + self._message_text = 'World Axis On' + + # L toggles the lighting mode + elif symbol == pyglet.window.key.L: + if self.viewer_flags['use_raymond_lighting']: + self.viewer_flags['use_raymond_lighting'] = False + self.viewer_flags['use_direct_lighting'] = True + self._message_text = 'Direct Lighting' + elif self.viewer_flags['use_direct_lighting']: + self.viewer_flags['use_raymond_lighting'] = False + self.viewer_flags['use_direct_lighting'] = False + self._message_text = 'Default Lighting' + else: + self.viewer_flags['use_raymond_lighting'] = True + self.viewer_flags['use_direct_lighting'] = False + self._message_text = 'Raymond Lighting' + + # M toggles face normals + elif symbol == pyglet.window.key.M: + self.render_flags['face_normals'] = ( + not self.render_flags['face_normals'] + ) + if self.render_flags['face_normals']: + self._message_text = 'Face Normals On' + else: + self._message_text = 'Face Normals Off' + + # N toggles vertex normals + elif symbol == pyglet.window.key.N: + self.render_flags['vertex_normals'] = ( + not self.render_flags['vertex_normals'] + ) + if self.render_flags['vertex_normals']: + self._message_text = 'Vert Normals On' + else: + self._message_text = 'Vert Normals Off' + + # O toggles orthographic camera mode + elif symbol == pyglet.window.key.O: + self.viewer_flags['use_perspective_cam'] = ( + not self.viewer_flags['use_perspective_cam'] + ) + if self.viewer_flags['use_perspective_cam']: + camera = self._default_persp_cam + self._message_text = 'Perspective View' + else: + camera = self._default_orth_cam + self._message_text = 'Orthographic View' + + cam_pose = self._camera_node.matrix.copy() + cam_node = Node(matrix=cam_pose, camera=camera) + self.scene.remove_node(self._camera_node) + self.scene.add_node(cam_node) + self.scene.main_camera_node = cam_node + self._camera_node = cam_node + + # Q quits the viewer + elif symbol == pyglet.window.key.Q: + self.on_close() + + # R starts recording frames + elif symbol == pyglet.window.key.R: + if self.viewer_flags['record']: + self.save_gif() + self.set_caption(self.viewer_flags['window_title']) + else: + self.set_caption( + '{} (RECORDING)'.format(self.viewer_flags['window_title']) + ) + self.viewer_flags['record'] = not self.viewer_flags['record'] + + # S saves the current frame as an image + elif symbol == pyglet.window.key.S: + self._save_image() + + # W toggles through wireframe modes + elif symbol == pyglet.window.key.W: + if self.render_flags['flip_wireframe']: + self.render_flags['flip_wireframe'] = False + self.render_flags['all_wireframe'] = True + self.render_flags['all_solid'] = False + self._message_text = 'All Wireframe' + elif self.render_flags['all_wireframe']: + self.render_flags['flip_wireframe'] = False + self.render_flags['all_wireframe'] = False + self.render_flags['all_solid'] = True + self._message_text = 'All Solid' + elif self.render_flags['all_solid']: + self.render_flags['flip_wireframe'] = False + self.render_flags['all_wireframe'] = False + self.render_flags['all_solid'] = False + self._message_text = 'Default Wireframe' + else: + self.render_flags['flip_wireframe'] = True + self.render_flags['all_wireframe'] = False + self.render_flags['all_solid'] = False + self._message_text = 'Flip Wireframe' + + # Z resets the camera viewpoint + elif symbol == pyglet.window.key.Z: + self._reset_view() + + if self._message_text is not None: + self._message_opac = 1.0 + self._ticks_till_fade + + @staticmethod + def _time_event(dt, self): + """The timer callback. + """ + # Don't run old dead events after we've already closed + if not self._is_active: + return + + if self.viewer_flags['record']: + self._record() + if (self.viewer_flags['rotate'] and not + self.viewer_flags['mouse_pressed']): + self._rotate() + + # Manage message opacity + if self._message_text is not None: + if self._message_opac > 1.0: + self._message_opac -= 1.0 + else: + self._message_opac *= 0.90 + if self._message_opac < 0.05: + self._message_opac = 1.0 + self._ticks_till_fade + self._message_text = None + + if self._should_close: + self.on_close() + else: + self.on_draw() + + def _reset_view(self): + """Reset the view to a good initial state. + + The view is initially along the positive x-axis at a + sufficient distance from the scene. + """ + scale = self.scene.scale + if scale == 0.0: + scale = DEFAULT_SCENE_SCALE + centroid = self.scene.centroid + + if self.viewer_flags['view_center'] is not None: + centroid = self.viewer_flags['view_center'] + + self._camera_node.matrix = self._default_camera_pose + self._trackball = Trackball( + self._default_camera_pose, self.viewport_size, scale, centroid + ) + + def _get_save_filename(self, file_exts): + file_types = { + 'png': ('png files', '*.png'), + 'jpg': ('jpeg files', '*.jpg'), + 'gif': ('gif files', '*.gif'), + 'all': ('all files', '*'), + } + filetypes = [file_types[x] for x in file_exts] + try: + root = Tk() + save_dir = self.viewer_flags['save_directory'] + if save_dir is None: + save_dir = os.getcwd() + filename = filedialog.asksaveasfilename( + initialdir=save_dir, title='Select file save location', + filetypes=filetypes + ) + except Exception: + return None + + root.destroy() + if filename == (): + return None + return filename + + def _save_image(self): + filename = self._get_save_filename(['png', 'jpg', 'gif', 'all']) + if filename is not None: + self.viewer_flags['save_directory'] = os.path.dirname(filename) + imageio.imwrite(filename, self._renderer.read_color_buf()) + + def _record(self): + """Save another frame for the GIF. + """ + data = self._renderer.read_color_buf() + if not np.all(data == 0.0): + self._saved_frames.append(data) + + def _rotate(self): + """Animate the scene by rotating the camera. + """ + az = (self.viewer_flags['rotate_rate'] / + self.viewer_flags['refresh_rate']) + self._trackball.rotate(az, self.viewer_flags['rotate_axis']) + + def _render(self): + """Render the scene into the framebuffer and flip. + """ + scene = self.scene + self._camera_node.matrix = self._trackball.pose.copy() + + # Set lighting + vli = self.viewer_flags['lighting_intensity'] + if self.viewer_flags['use_raymond_lighting']: + for n in self._raymond_lights: + n.light.intensity = vli / 3.0 + if not self.scene.has_node(n): + scene.add_node(n, parent_node=self._camera_node) + else: + self._direct_light.light.intensity = vli + for n in self._raymond_lights: + if self.scene.has_node(n): + self.scene.remove_node(n) + + if self.viewer_flags['use_direct_lighting']: + if not self.scene.has_node(self._direct_light): + scene.add_node( + self._direct_light, parent_node=self._camera_node + ) + elif self.scene.has_node(self._direct_light): + self.scene.remove_node(self._direct_light) + + flags = RenderFlags.NONE + if self.render_flags['flip_wireframe']: + flags |= RenderFlags.FLIP_WIREFRAME + elif self.render_flags['all_wireframe']: + flags |= RenderFlags.ALL_WIREFRAME + elif self.render_flags['all_solid']: + flags |= RenderFlags.ALL_SOLID + + if self.render_flags['shadows']: + flags |= RenderFlags.SHADOWS_DIRECTIONAL | RenderFlags.SHADOWS_SPOT + if self.render_flags['vertex_normals']: + flags |= RenderFlags.VERTEX_NORMALS + if self.render_flags['face_normals']: + flags |= RenderFlags.FACE_NORMALS + if not self.render_flags['cull_faces']: + flags |= RenderFlags.SKIP_CULL_FACES + + self._renderer.render(self.scene, flags) + + def _init_and_start_app(self): + # Try multiple configs starting with target OpenGL version + # and multisampling and removing these options if exception + # Note: multisampling not available on all hardware + from pyglet.gl import Config + confs = [Config(sample_buffers=1, samples=4, + depth_size=24, + double_buffer=True, + major_version=TARGET_OPEN_GL_MAJOR, + minor_version=TARGET_OPEN_GL_MINOR), + Config(depth_size=24, + double_buffer=True, + major_version=TARGET_OPEN_GL_MAJOR, + minor_version=TARGET_OPEN_GL_MINOR), + Config(sample_buffers=1, samples=4, + depth_size=24, + double_buffer=True, + major_version=MIN_OPEN_GL_MAJOR, + minor_version=MIN_OPEN_GL_MINOR), + Config(depth_size=24, + double_buffer=True, + major_version=MIN_OPEN_GL_MAJOR, + minor_version=MIN_OPEN_GL_MINOR)] + for conf in confs: + try: + super(Viewer, self).__init__(config=conf, resizable=True, + width=self._viewport_size[0], + height=self._viewport_size[1]) + break + except pyglet.window.NoSuchConfigException: + pass + + if not self.context: + raise ValueError('Unable to initialize an OpenGL 3+ context') + clock.schedule_interval( + Viewer._time_event, 1.0 / self.viewer_flags['refresh_rate'], self + ) + self.switch_to() + self.set_caption(self.viewer_flags['window_title']) + pyglet.app.run() + + def _compute_initial_camera_pose(self): + centroid = self.scene.centroid + if self.viewer_flags['view_center'] is not None: + centroid = self.viewer_flags['view_center'] + scale = self.scene.scale + if scale == 0.0: + scale = DEFAULT_SCENE_SCALE + + s2 = 1.0 / np.sqrt(2.0) + cp = np.eye(4) + cp[:3,:3] = np.array([ + [0.0, -s2, s2], + [1.0, 0.0, 0.0], + [0.0, s2, s2] + ]) + hfov = np.pi / 6.0 + dist = scale / (2.0 * np.tan(hfov)) + cp[:3,3] = dist * np.array([1.0, 0.0, 1.0]) + centroid + + return cp + + def _create_raymond_lights(self): + thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0]) + phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0]) + + nodes = [] + + for phi, theta in zip(phis, thetas): + xp = np.sin(theta) * np.cos(phi) + yp = np.sin(theta) * np.sin(phi) + zp = np.cos(theta) + + z = np.array([xp, yp, zp]) + z = z / np.linalg.norm(z) + x = np.array([-z[1], z[0], 0.0]) + if np.linalg.norm(x) == 0: + x = np.array([1.0, 0.0, 0.0]) + x = x / np.linalg.norm(x) + y = np.cross(z, x) + + matrix = np.eye(4) + matrix[:3,:3] = np.c_[x,y,z] + nodes.append(Node( + light=DirectionalLight(color=np.ones(3), intensity=1.0), + matrix=matrix + )) + + return nodes + + def _create_direct_light(self): + light = DirectionalLight(color=np.ones(3), intensity=1.0) + n = Node(light=light, matrix=np.eye(4)) + return n + + def _set_axes(self, world, mesh): + scale = self.scene.scale + if world: + if 'scene' not in self._axes: + n = Node(mesh=self._axis_mesh, scale=np.ones(3) * scale * 0.3) + self.scene.add_node(n) + self._axes['scene'] = n + else: + if 'scene' in self._axes: + self.scene.remove_node(self._axes['scene']) + self._axes.pop('scene') + + if mesh: + old_nodes = [] + existing_axes = set([self._axes[k] for k in self._axes]) + for node in self.scene.mesh_nodes: + if node not in existing_axes: + old_nodes.append(node) + + for node in old_nodes: + if node in self._axes: + continue + n = Node( + mesh=self._axis_mesh, + scale=np.ones(3) * node.mesh.scale * 0.5 + ) + self.scene.add_node(n, parent_node=node) + self._axes[node] = n + else: + to_remove = set() + for main_node in self._axes: + if main_node in self.scene.mesh_nodes: + self.scene.remove_node(self._axes[main_node]) + to_remove.add(main_node) + for main_node in to_remove: + self._axes.pop(main_node) + + def _remove_axes(self): + for main_node in self._axes: + axis_node = self._axes[main_node] + self.scene.remove_node(axis_node) + self._axes = {} + + def _location_to_x_y(self, location): + if location == TextAlign.CENTER: + return (self.viewport_size[0] / 2.0, self.viewport_size[1] / 2.0) + elif location == TextAlign.CENTER_LEFT: + return (TEXT_PADDING, self.viewport_size[1] / 2.0) + elif location == TextAlign.CENTER_RIGHT: + return (self.viewport_size[0] - TEXT_PADDING, + self.viewport_size[1] / 2.0) + elif location == TextAlign.BOTTOM_LEFT: + return (TEXT_PADDING, TEXT_PADDING) + elif location == TextAlign.BOTTOM_RIGHT: + return (self.viewport_size[0] - TEXT_PADDING, TEXT_PADDING) + elif location == TextAlign.BOTTOM_CENTER: + return (self.viewport_size[0] / 2.0, TEXT_PADDING) + elif location == TextAlign.TOP_LEFT: + return (TEXT_PADDING, self.viewport_size[1] - TEXT_PADDING) + elif location == TextAlign.TOP_RIGHT: + return (self.viewport_size[0] - TEXT_PADDING, + self.viewport_size[1] - TEXT_PADDING) + elif location == TextAlign.TOP_CENTER: + return (self.viewport_size[0] / 2.0, + self.viewport_size[1] - TEXT_PADDING) + + +__all__ = ['Viewer'] diff --git a/pyrender/requirements.txt b/pyrender/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..8c40b74256f0dc6697754bb8609f69a39d51beba --- /dev/null +++ b/pyrender/requirements.txt @@ -0,0 +1,14 @@ +freetype-py +imageio +networkx +numpy +Pillow +pyglet==1.4.0a1 +PyOpenGL +PyOpenGL_accelerate +six +trimesh +sphinx +sphinx_rtd_theme +sphinx-automodapi + diff --git a/pyrender/setup.py b/pyrender/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..c3b5ba0da2b0f17b759e5556597981096a80bda8 --- /dev/null +++ b/pyrender/setup.py @@ -0,0 +1,76 @@ +""" +Setup of pyrender Python codebase. + +Author: Matthew Matl +""" +import sys +from setuptools import setup + +# load __version__ +exec(open('pyrender/version.py').read()) + +def get_imageio_dep(): + if sys.version[0] == "2": + return 'imageio<=2.6.1' + return 'imageio' + +requirements = [ + 'freetype-py', # For font loading + get_imageio_dep(), # For Image I/O + 'networkx', # For the scene graph + 'numpy', # Numpy + 'Pillow', # For Trimesh texture conversions + 'pyglet>=1.4.10', # For the pyglet viewer + 'PyOpenGL~=3.1.0', # For OpenGL +# 'PyOpenGL_accelerate~=3.1.0', # For OpenGL + 'scipy', # Because of trimesh missing dep + 'six', # For Python 2/3 interop + 'trimesh', # For meshes +] + +dev_requirements = [ + 'flake8', # Code formatting checker + 'pre-commit', # Pre-commit hooks + 'pytest', # Code testing + 'pytest-cov', # Coverage testing + 'tox', # Automatic virtualenv testing +] + +docs_requirements = [ + 'sphinx', # General doc library + 'sphinx_rtd_theme', # RTD theme for sphinx + 'sphinx-automodapi' # For generating nice tables +] + + +setup( + name = 'pyrender', + version=__version__, + description='Easy-to-use Python renderer for 3D visualization', + long_description='A simple implementation of Physically-Based Rendering ' + '(PBR) in Python. Compliant with the glTF 2.0 standard.', + author='Matthew Matl', + author_email='matthewcmatl@gmail.com', + license='MIT License', + url = 'https://github.com/mmatl/pyrender', + classifiers = [ + 'Development Status :: 4 - Beta', + 'License :: OSI Approved :: MIT License', + 'Operating System :: POSIX :: Linux', + 'Operating System :: MacOS :: MacOS X', + 'Programming Language :: Python :: 2.7', + 'Programming Language :: Python :: 3.5', + 'Programming Language :: Python :: 3.6', + 'Natural Language :: English', + 'Topic :: Scientific/Engineering' + ], + keywords = 'rendering graphics opengl 3d visualization pbr gltf', + packages = ['pyrender', 'pyrender.platforms'], + setup_requires = requirements, + install_requires = requirements, + extras_require={ + 'dev': dev_requirements, + 'docs': docs_requirements, + }, + include_package_data=True +) diff --git a/pyrender/tests/__init__.py b/pyrender/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pyrender/tests/conftest.py b/pyrender/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pyrender/tests/pytest.ini b/pyrender/tests/pytest.ini new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pyrender/tests/unit/__init__.py b/pyrender/tests/unit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pyrender/tests/unit/test_cameras.py b/pyrender/tests/unit/test_cameras.py new file mode 100644 index 0000000000000000000000000000000000000000..7544ad8f8e3ee55236fd2e32dbc12065153cbe5b --- /dev/null +++ b/pyrender/tests/unit/test_cameras.py @@ -0,0 +1,164 @@ +import numpy as np +import pytest + +from pyrender import PerspectiveCamera, OrthographicCamera + + +def test_perspective_camera(): + + # Set up constants + znear = 0.05 + zfar = 100 + yfov = np.pi / 3.0 + width = 1000.0 + height = 500.0 + aspectRatio = 640.0 / 480.0 + + # Test basics + with pytest.raises(TypeError): + p = PerspectiveCamera() + + p = PerspectiveCamera(yfov=yfov) + assert p.yfov == yfov + assert p.znear == 0.05 + assert p.zfar is None + assert p.aspectRatio is None + p.name = 'asdf' + p.name = None + + with pytest.raises(ValueError): + p.yfov = 0.0 + + with pytest.raises(ValueError): + p.yfov = -1.0 + + with pytest.raises(ValueError): + p.znear = -1.0 + + p.znear = 0.0 + p.znear = 0.05 + p.zfar = 100.0 + assert p.zfar == 100.0 + + with pytest.raises(ValueError): + p.zfar = 0.03 + + with pytest.raises(ValueError): + p.zfar = 0.05 + + p.aspectRatio = 10.0 + assert p.aspectRatio == 10.0 + + with pytest.raises(ValueError): + p.aspectRatio = 0.0 + + with pytest.raises(ValueError): + p.aspectRatio = -1.0 + + # Test matrix getting/setting + + # NF + p.znear = 0.05 + p.zfar = 100 + p.aspectRatio = None + + with pytest.raises(ValueError): + p.get_projection_matrix() + + assert np.allclose( + p.get_projection_matrix(width, height), + np.array([ + [1.0 / (width / height * np.tan(yfov / 2.0)), 0.0, 0.0, 0.0], + [0.0, 1.0 / np.tan(yfov / 2.0), 0.0, 0.0], + [0.0, 0.0, (zfar + znear) / (znear - zfar), + (2 * zfar * znear) / (znear - zfar)], + [0.0, 0.0, -1.0, 0.0] + ]) + ) + + # NFA + p.aspectRatio = aspectRatio + assert np.allclose( + p.get_projection_matrix(width, height), + np.array([ + [1.0 / (aspectRatio * np.tan(yfov / 2.0)), 0.0, 0.0, 0.0], + [0.0, 1.0 / np.tan(yfov / 2.0), 0.0, 0.0], + [0.0, 0.0, (zfar + znear) / (znear - zfar), + (2 * zfar * znear) / (znear - zfar)], + [0.0, 0.0, -1.0, 0.0] + ]) + ) + assert np.allclose( + p.get_projection_matrix(), p.get_projection_matrix(width, height) + ) + + # N + p.zfar = None + p.aspectRatio = None + assert np.allclose( + p.get_projection_matrix(width, height), + np.array([ + [1.0 / (width / height * np.tan(yfov / 2.0)), 0.0, 0.0, 0.0], + [0.0, 1.0 / np.tan(yfov / 2.0), 0.0, 0.0], + [0.0, 0.0, -1.0, -2.0 * znear], + [0.0, 0.0, -1.0, 0.0] + ]) + ) + + +def test_orthographic_camera(): + xm = 1.0 + ym = 2.0 + n = 0.05 + f = 100.0 + + with pytest.raises(TypeError): + c = OrthographicCamera() + + c = OrthographicCamera(xmag=xm, ymag=ym) + + assert c.xmag == xm + assert c.ymag == ym + assert c.znear == 0.05 + assert c.zfar == 100.0 + assert c.name is None + + with pytest.raises(TypeError): + c.ymag = None + + with pytest.raises(ValueError): + c.ymag = 0.0 + + with pytest.raises(ValueError): + c.ymag = -1.0 + + with pytest.raises(TypeError): + c.xmag = None + + with pytest.raises(ValueError): + c.xmag = 0.0 + + with pytest.raises(ValueError): + c.xmag = -1.0 + + with pytest.raises(TypeError): + c.znear = None + + with pytest.raises(ValueError): + c.znear = 0.0 + + with pytest.raises(ValueError): + c.znear = -1.0 + + with pytest.raises(ValueError): + c.zfar = 0.01 + + assert np.allclose( + c.get_projection_matrix(), + np.array([ + [1.0 / xm, 0, 0, 0], + [0, 1.0 / ym, 0, 0], + [0, 0, 2.0 / (n - f), (f + n) / (n - f)], + [0, 0, 0, 1.0] + ]) + ) diff --git a/pyrender/tests/unit/test_egl.py b/pyrender/tests/unit/test_egl.py new file mode 100644 index 0000000000000000000000000000000000000000..e2f4bef39e33c2794e6837b5a1bb127d8d4dba06 --- /dev/null +++ b/pyrender/tests/unit/test_egl.py @@ -0,0 +1,16 @@ +# from pyrender.platforms import egl + + +def tmp_test_default_device(): + egl.get_default_device() + + +def tmp_test_query_device(): + devices = egl.query_devices() + assert len(devices) > 0 + + +def tmp_test_init_context(): + device = egl.query_devices()[0] + platform = egl.EGLPlatform(128, 128, device=device) + platform.init_context() diff --git a/pyrender/tests/unit/test_lights.py b/pyrender/tests/unit/test_lights.py new file mode 100644 index 0000000000000000000000000000000000000000..ffde856b21e8cce9532f0308fcd1c7eb2d1eba90 --- /dev/null +++ b/pyrender/tests/unit/test_lights.py @@ -0,0 +1,104 @@ +import numpy as np +import pytest + +from pyrender import (DirectionalLight, SpotLight, PointLight, Texture, + PerspectiveCamera, OrthographicCamera) +from pyrender.constants import SHADOW_TEX_SZ + + +def test_directional_light(): + + d = DirectionalLight() + assert d.name is None + assert np.all(d.color == 1.0) + assert d.intensity == 1.0 + + d.name = 'direc' + with pytest.raises(ValueError): + d.color = None + with pytest.raises(TypeError): + d.intensity = None + + d = DirectionalLight(color=[0.0, 0.0, 0.0]) + assert np.all(d.color == 0.0) + + d._generate_shadow_texture() + st = d.shadow_texture + assert isinstance(st, Texture) + assert st.width == st.height == SHADOW_TEX_SZ + + sc = d._get_shadow_camera(scene_scale=5.0) + assert isinstance(sc, OrthographicCamera) + assert sc.xmag == sc.ymag == 5.0 + assert sc.znear == 0.01 * 5.0 + assert sc.zfar == 10 * 5.0 + + +def test_spot_light(): + + s = SpotLight() + assert s.name is None + assert np.all(s.color == 1.0) + assert s.intensity == 1.0 + assert s.innerConeAngle == 0.0 + assert s.outerConeAngle == np.pi / 4.0 + assert s.range is None + + with pytest.raises(ValueError): + s.range = -1.0 + + with pytest.raises(ValueError): + s.range = 0.0 + + with pytest.raises(ValueError): + s.innerConeAngle = -1.0 + + with pytest.raises(ValueError): + s.innerConeAngle = np.pi / 3.0 + + with pytest.raises(ValueError): + s.outerConeAngle = -1.0 + + with pytest.raises(ValueError): + s.outerConeAngle = np.pi + + s.range = 5.0 + s.outerConeAngle = np.pi / 2 - 0.05 + s.innerConeAngle = np.pi / 3 + s.innerConeAngle = 0.0 + s.outerConeAngle = np.pi / 4.0 + + s._generate_shadow_texture() + st = s.shadow_texture + assert isinstance(st, Texture) + assert st.width == st.height == SHADOW_TEX_SZ + + sc = s._get_shadow_camera(scene_scale=5.0) + assert isinstance(sc, PerspectiveCamera) + assert sc.znear == 0.01 * 5.0 + assert sc.zfar == 10 * 5.0 + assert sc.aspectRatio == 1.0 + assert np.allclose(sc.yfov, np.pi / 16.0 * 9.0) # Plus pi / 16 + + +def test_point_light(): + + s = PointLight() + assert s.name is None + assert np.all(s.color == 1.0) + assert s.intensity == 1.0 + assert s.range is None + + with pytest.raises(ValueError): + s.range = -1.0 + + with pytest.raises(ValueError): + s.range = 0.0 + + s.range = 5.0 + + with pytest.raises(NotImplementedError): + s._generate_shadow_texture() + + with pytest.raises(NotImplementedError): + s._get_shadow_camera(scene_scale=5.0) diff --git a/pyrender/tests/unit/test_meshes.py b/pyrender/tests/unit/test_meshes.py new file mode 100644 index 0000000000000000000000000000000000000000..7070b01171c97069fa013c6eba8eee217017f08e --- /dev/null +++ b/pyrender/tests/unit/test_meshes.py @@ -0,0 +1,133 @@ +import numpy as np +import pytest +import trimesh + +from pyrender import (Mesh, Primitive) + + +def test_meshes(): + + with pytest.raises(TypeError): + x = Mesh() + with pytest.raises(TypeError): + x = Primitive() + with pytest.raises(ValueError): + x = Primitive([], mode=10) + + # Basics + x = Mesh([]) + assert x.name is None + assert x.is_visible + assert x.weights is None + + x.name = 'str' + + # From Trimesh + x = Mesh.from_trimesh(trimesh.creation.box()) + assert isinstance(x, Mesh) + assert len(x.primitives) == 1 + assert x.is_visible + assert np.allclose(x.bounds, np.array([ + [-0.5, -0.5, -0.5], + [0.5, 0.5, 0.5] + ])) + assert np.allclose(x.centroid, np.zeros(3)) + assert np.allclose(x.extents, np.ones(3)) + assert np.allclose(x.scale, np.sqrt(3)) + assert not x.is_transparent + + # Test some primitive functions + x = x.primitives[0] + with pytest.raises(ValueError): + x.normals = np.zeros(10) + with pytest.raises(ValueError): + x.tangents = np.zeros(10) + with pytest.raises(ValueError): + x.texcoord_0 = np.zeros(10) + with pytest.raises(ValueError): + x.texcoord_1 = np.zeros(10) + with pytest.raises(TypeError): + x.material = np.zeros(10) + assert x.targets is None + assert np.allclose(x.bounds, np.array([ + [-0.5, -0.5, -0.5], + [0.5, 0.5, 0.5] + ])) + assert np.allclose(x.centroid, np.zeros(3)) + assert np.allclose(x.extents, np.ones(3)) + assert np.allclose(x.scale, np.sqrt(3)) + x.material.baseColorFactor = np.array([0.0, 0.0, 0.0, 0.0]) + assert x.is_transparent + + # From two trimeshes + x = Mesh.from_trimesh([trimesh.creation.box(), + trimesh.creation.cylinder(radius=0.1, height=2.0)], + smooth=False) + assert isinstance(x, Mesh) + assert len(x.primitives) == 2 + assert x.is_visible + assert np.allclose(x.bounds, np.array([ + [-0.5, -0.5, -1.0], + [0.5, 0.5, 1.0] + ])) + assert np.allclose(x.centroid, np.zeros(3)) + assert np.allclose(x.extents, [1.0, 1.0, 2.0]) + assert np.allclose(x.scale, np.sqrt(6)) + assert not x.is_transparent + + # From bad data + with pytest.raises(TypeError): + x = Mesh.from_trimesh(None) + + # With instancing + poses = np.tile(np.eye(4), (5,1,1)) + poses[:,0,3] = np.array([0,1,2,3,4]) + x = Mesh.from_trimesh(trimesh.creation.box(), poses=poses) + assert np.allclose(x.bounds, np.array([ + [-0.5, -0.5, -0.5], + [4.5, 0.5, 0.5] + ])) + poses = np.eye(4) + x = Mesh.from_trimesh(trimesh.creation.box(), poses=poses) + poses = np.eye(3) + with pytest.raises(ValueError): + x = Mesh.from_trimesh(trimesh.creation.box(), poses=poses) + + # From textured meshes + fm = trimesh.load('tests/data/fuze.obj') + x = Mesh.from_trimesh(fm) + assert isinstance(x, Mesh) + assert len(x.primitives) == 1 + assert x.is_visible + assert not x.is_transparent + assert x.primitives[0].material.baseColorTexture is not None + + x = Mesh.from_trimesh(fm, smooth=False) + fm.visual = fm.visual.to_color() + fm.visual.face_colors = np.array([1.0, 0.0, 0.0, 1.0]) + x = Mesh.from_trimesh(fm, smooth=False) + with pytest.raises(ValueError): + x = Mesh.from_trimesh(fm, smooth=True) + + fm.visual.vertex_colors = np.array([1.0, 0.0, 0.0, 0.5]) + x = Mesh.from_trimesh(fm, smooth=False) + x = Mesh.from_trimesh(fm, smooth=True) + assert x.primitives[0].color_0 is not None + assert x.is_transparent + + bm = trimesh.load('tests/data/WaterBottle.glb').dump()[0] + x = Mesh.from_trimesh(bm) + assert x.primitives[0].material.baseColorTexture is not None + assert x.primitives[0].material.emissiveTexture is not None + assert x.primitives[0].material.metallicRoughnessTexture is not None + + # From point cloud + x = Mesh.from_points(fm.vertices) + +# def test_duck(): +# bm = trimesh.load('tests/data/Duck.glb').dump()[0] +# x = Mesh.from_trimesh(bm) +# assert x.primitives[0].material.baseColorTexture is not None +# pixel = x.primitives[0].material.baseColorTexture.source[100, 100] +# yellowish = np.array([1.0, 0.7411765, 0.0, 1.0]) +# assert np.allclose(pixel, yellowish) diff --git a/pyrender/tests/unit/test_nodes.py b/pyrender/tests/unit/test_nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..9857c8221b7f6fb8530699bdf5593f8f0b74e152 --- /dev/null +++ b/pyrender/tests/unit/test_nodes.py @@ -0,0 +1,124 @@ +import numpy as np +import pytest +from trimesh import transformations + +from pyrender import (DirectionalLight, PerspectiveCamera, Mesh, Node) + + +def test_nodes(): + + x = Node() + assert x.name is None + assert x.camera is None + assert x.children == [] + assert x.skin is None + assert np.allclose(x.matrix, np.eye(4)) + assert x.mesh is None + assert np.allclose(x.rotation, [0,0,0,1]) + assert np.allclose(x.scale, np.ones(3)) + assert np.allclose(x.translation, np.zeros(3)) + assert x.weights is None + assert x.light is None + + x.name = 'node' + + # Test node light/camera/mesh tests + c = PerspectiveCamera(yfov=2.0) + m = Mesh([]) + d = DirectionalLight() + x.camera = c + assert x.camera == c + with pytest.raises(TypeError): + x.camera = m + x.camera = d + x.camera = None + x.mesh = m + assert x.mesh == m + with pytest.raises(TypeError): + x.mesh = c + x.mesh = d + x.light = d + assert x.light == d + with pytest.raises(TypeError): + x.light = m + x.light = c + + # Test transformations getters/setters/etc... + # Set up test values + x = np.array([1.0, 0.0, 0.0]) + y = np.array([0.0, 1.0, 0.0]) + t = np.array([1.0, 2.0, 3.0]) + s = np.array([0.5, 2.0, 1.0]) + + Mx = transformations.rotation_matrix(np.pi / 2.0, x) + qx = np.roll(transformations.quaternion_about_axis(np.pi / 2.0, x), -1) + Mxt = Mx.copy() + Mxt[:3,3] = t + S = np.eye(4) + S[:3,:3] = np.diag(s) + Mxts = Mxt.dot(S) + + My = transformations.rotation_matrix(np.pi / 2.0, y) + qy = np.roll(transformations.quaternion_about_axis(np.pi / 2.0, y), -1) + Myt = My.copy() + Myt[:3,3] = t + + x = Node(matrix=Mx) + assert np.allclose(x.matrix, Mx) + assert np.allclose(x.rotation, qx) + assert np.allclose(x.translation, np.zeros(3)) + assert np.allclose(x.scale, np.ones(3)) + + x.matrix = My + assert np.allclose(x.matrix, My) + assert np.allclose(x.rotation, qy) + assert np.allclose(x.translation, np.zeros(3)) + assert np.allclose(x.scale, np.ones(3)) + x.translation = t + assert np.allclose(x.matrix, Myt) + assert np.allclose(x.rotation, qy) + x.rotation = qx + assert np.allclose(x.matrix, Mxt) + x.scale = s + assert np.allclose(x.matrix, Mxts) + + x = Node(matrix=Mxt) + assert np.allclose(x.matrix, Mxt) + assert np.allclose(x.rotation, qx) + assert np.allclose(x.translation, t) + assert np.allclose(x.scale, np.ones(3)) + + x = Node(matrix=Mxts) + assert np.allclose(x.matrix, Mxts) + assert np.allclose(x.rotation, qx) + assert np.allclose(x.translation, t) + assert np.allclose(x.scale, s) + + # Individual element getters + x.scale[0] = 0 + assert np.allclose(x.scale[0], 0) + + x.translation[0] = 0 + assert np.allclose(x.translation[0], 0) + + x.matrix = np.eye(4) + x.matrix[0,0] = 500 + assert x.matrix[0,0] == 1.0 + + # Failures + with pytest.raises(ValueError): + x.matrix = 5 * np.eye(4) + with pytest.raises(ValueError): + x.matrix = np.eye(5) + with pytest.raises(ValueError): + x.matrix = np.eye(4).dot([5,1,1,1]) + with pytest.raises(ValueError): + x.rotation = np.array([1,2]) + with pytest.raises(ValueError): + x.rotation = np.array([1,2,3]) + with pytest.raises(ValueError): + x.rotation = np.array([1,2,3,4]) + with pytest.raises(ValueError): + x.translation = np.array([1,2,3,4]) + with pytest.raises(ValueError): + x.scale = np.array([1,2,3,4]) diff --git a/pyrender/tests/unit/test_offscreen.py b/pyrender/tests/unit/test_offscreen.py new file mode 100644 index 0000000000000000000000000000000000000000..88983b0ff4e2ab6f5ef252c51f2ac669c3a0e0ca --- /dev/null +++ b/pyrender/tests/unit/test_offscreen.py @@ -0,0 +1,92 @@ +import numpy as np +import trimesh + +from pyrender import (OffscreenRenderer, PerspectiveCamera, DirectionalLight, + SpotLight, Mesh, Node, Scene) + + +def test_offscreen_renderer(tmpdir): + + # Fuze trimesh + fuze_trimesh = trimesh.load('examples/models/fuze.obj') + fuze_mesh = Mesh.from_trimesh(fuze_trimesh) + + # Drill trimesh + drill_trimesh = trimesh.load('examples/models/drill.obj') + drill_mesh = Mesh.from_trimesh(drill_trimesh) + drill_pose = np.eye(4) + drill_pose[0,3] = 0.1 + drill_pose[2,3] = -np.min(drill_trimesh.vertices[:,2]) + + # Wood trimesh + wood_trimesh = trimesh.load('examples/models/wood.obj') + wood_mesh = Mesh.from_trimesh(wood_trimesh) + + # Water bottle trimesh + bottle_gltf = trimesh.load('examples/models/WaterBottle.glb') + bottle_trimesh = bottle_gltf.geometry[list(bottle_gltf.geometry.keys())[0]] + bottle_mesh = Mesh.from_trimesh(bottle_trimesh) + bottle_pose = np.array([ + [1.0, 0.0, 0.0, 0.1], + [0.0, 0.0, -1.0, -0.16], + [0.0, 1.0, 0.0, 0.13], + [0.0, 0.0, 0.0, 1.0], + ]) + + boxv_trimesh = trimesh.creation.box(extents=0.1 * np.ones(3)) + boxv_vertex_colors = np.random.uniform(size=(boxv_trimesh.vertices.shape)) + boxv_trimesh.visual.vertex_colors = boxv_vertex_colors + boxv_mesh = Mesh.from_trimesh(boxv_trimesh, smooth=False) + boxf_trimesh = trimesh.creation.box(extents=0.1 * np.ones(3)) + boxf_face_colors = np.random.uniform(size=boxf_trimesh.faces.shape) + boxf_trimesh.visual.face_colors = boxf_face_colors + # Instanced + poses = np.tile(np.eye(4), (2,1,1)) + poses[0,:3,3] = np.array([-0.1, -0.10, 0.05]) + poses[1,:3,3] = np.array([-0.15, -0.10, 0.05]) + boxf_mesh = Mesh.from_trimesh(boxf_trimesh, poses=poses, smooth=False) + + points = trimesh.creation.icosphere(radius=0.05).vertices + point_colors = np.random.uniform(size=points.shape) + points_mesh = Mesh.from_points(points, colors=point_colors) + + direc_l = DirectionalLight(color=np.ones(3), intensity=1.0) + spot_l = SpotLight(color=np.ones(3), intensity=10.0, + innerConeAngle=np.pi / 16, outerConeAngle=np.pi / 6) + + cam = PerspectiveCamera(yfov=(np.pi / 3.0)) + cam_pose = np.array([ + [0.0, -np.sqrt(2) / 2, np.sqrt(2) / 2, 0.5], + [1.0, 0.0, 0.0, 0.0], + [0.0, np.sqrt(2) / 2, np.sqrt(2) / 2, 0.4], + [0.0, 0.0, 0.0, 1.0] + ]) + + scene = Scene(ambient_light=np.array([0.02, 0.02, 0.02])) + + fuze_node = Node(mesh=fuze_mesh, translation=np.array([ + 0.1, 0.15, -np.min(fuze_trimesh.vertices[:,2]) + ])) + scene.add_node(fuze_node) + boxv_node = Node(mesh=boxv_mesh, translation=np.array([-0.1, 0.10, 0.05])) + scene.add_node(boxv_node) + boxf_node = Node(mesh=boxf_mesh) + scene.add_node(boxf_node) + + _ = scene.add(drill_mesh, pose=drill_pose) + _ = scene.add(bottle_mesh, pose=bottle_pose) + _ = scene.add(wood_mesh) + _ = scene.add(direc_l, pose=cam_pose) + _ = scene.add(spot_l, pose=cam_pose) + _ = scene.add(points_mesh) + + _ = scene.add(cam, pose=cam_pose) + + r = OffscreenRenderer(viewport_width=640, viewport_height=480) + color, depth = r.render(scene) + + assert color.shape == (480, 640, 3) + assert depth.shape == (480, 640) + assert np.max(depth.data) > 0.05 + assert np.count_nonzero(depth.data) > (0.2 * depth.size) + r.delete() diff --git a/pyrender/tests/unit/test_scenes.py b/pyrender/tests/unit/test_scenes.py new file mode 100644 index 0000000000000000000000000000000000000000..d85dd714cb5d842ea12dee4140adfd7db55c9c01 --- /dev/null +++ b/pyrender/tests/unit/test_scenes.py @@ -0,0 +1,235 @@ +import numpy as np +import pytest +import trimesh + +from pyrender import (Mesh, PerspectiveCamera, DirectionalLight, + SpotLight, PointLight, Scene, Node, OrthographicCamera) + + +def test_scenes(): + + # Basics + s = Scene() + assert np.allclose(s.bg_color, np.ones(4)) + assert np.allclose(s.ambient_light, np.zeros(3)) + assert len(s.nodes) == 0 + assert s.name is None + s.name = 'asdf' + s.bg_color = None + s.ambient_light = None + assert np.allclose(s.bg_color, np.ones(4)) + assert np.allclose(s.ambient_light, np.zeros(3)) + + assert s.nodes == set() + assert s.cameras == set() + assert s.lights == set() + assert s.point_lights == set() + assert s.spot_lights == set() + assert s.directional_lights == set() + assert s.meshes == set() + assert s.camera_nodes == set() + assert s.light_nodes == set() + assert s.point_light_nodes == set() + assert s.spot_light_nodes == set() + assert s.directional_light_nodes == set() + assert s.mesh_nodes == set() + assert s.main_camera_node is None + assert np.all(s.bounds == 0) + assert np.all(s.centroid == 0) + assert np.all(s.extents == 0) + assert np.all(s.scale == 0) + + # From trimesh scene + tms = trimesh.load('tests/data/WaterBottle.glb') + s = Scene.from_trimesh_scene(tms) + assert len(s.meshes) == 1 + assert len(s.mesh_nodes) == 1 + + # Test bg color formatting + s = Scene(bg_color=[0, 1.0, 0]) + assert np.allclose(s.bg_color, np.array([0.0, 1.0, 0.0, 1.0])) + + # Test constructor for nodes + n1 = Node() + n2 = Node() + n3 = Node() + nodes = [n1, n2, n3] + s = Scene(nodes=nodes) + n1.children.append(n2) + s = Scene(nodes=nodes) + n3.children.append(n2) + with pytest.raises(ValueError): + s = Scene(nodes=nodes) + n3.children = [] + n2.children.append(n3) + n3.children.append(n2) + with pytest.raises(ValueError): + s = Scene(nodes=nodes) + + # Test node accessors + n1 = Node() + n2 = Node() + n3 = Node() + nodes = [n1, n2] + s = Scene(nodes=nodes) + assert s.has_node(n1) + assert s.has_node(n2) + assert not s.has_node(n3) + + # Test node poses + for n in nodes: + assert np.allclose(s.get_pose(n), np.eye(4)) + with pytest.raises(ValueError): + s.get_pose(n3) + with pytest.raises(ValueError): + s.set_pose(n3, np.eye(4)) + tf = np.eye(4) + tf[:3,3] = np.ones(3) + s.set_pose(n1, tf) + assert np.allclose(s.get_pose(n1), tf) + assert np.allclose(s.get_pose(n2), np.eye(4)) + + nodes = [n1, n2, n3] + tf2 = np.eye(4) + tf2[:3,:3] = np.diag([-1,-1,1]) + n1.children.append(n2) + n1.matrix = tf + n2.matrix = tf2 + s = Scene(nodes=nodes) + assert np.allclose(s.get_pose(n1), tf) + assert np.allclose(s.get_pose(n2), tf.dot(tf2)) + assert np.allclose(s.get_pose(n3), np.eye(4)) + + n1 = Node() + n2 = Node() + n3 = Node() + n1.children.append(n2) + s = Scene() + s.add_node(n1) + with pytest.raises(ValueError): + s.add_node(n2) + s.set_pose(n1, tf) + assert np.allclose(s.get_pose(n1), tf) + assert np.allclose(s.get_pose(n2), tf) + s.set_pose(n2, tf2) + assert np.allclose(s.get_pose(n2), tf.dot(tf2)) + + # Test node removal + n1 = Node() + n2 = Node() + n3 = Node() + n1.children.append(n2) + n2.children.append(n3) + s = Scene(nodes=[n1, n2, n3]) + s.remove_node(n2) + assert len(s.nodes) == 1 + assert n1 in s.nodes + assert len(n1.children) == 0 + assert len(n2.children) == 1 + s.add_node(n2, parent_node=n1) + assert len(n1.children) == 1 + n1.matrix = tf + n3.matrix = tf2 + assert np.allclose(s.get_pose(n3), tf.dot(tf2)) + + # Now test ADD function + s = Scene() + m = Mesh([], name='m') + cp = PerspectiveCamera(yfov=2.0) + co = OrthographicCamera(xmag=1.0, ymag=1.0) + dl = DirectionalLight() + pl = PointLight() + sl = SpotLight() + + n1 = s.add(m, name='mn') + assert n1.mesh == m + assert len(s.nodes) == 1 + assert len(s.mesh_nodes) == 1 + assert n1 in s.mesh_nodes + assert len(s.meshes) == 1 + assert m in s.meshes + assert len(s.get_nodes(node=n2)) == 0 + n2 = s.add(m, pose=tf) + assert len(s.nodes) == len(s.mesh_nodes) == 2 + assert len(s.meshes) == 1 + assert len(s.get_nodes(node=n1)) == 1 + assert len(s.get_nodes(node=n1, name='mn')) == 1 + assert len(s.get_nodes(name='mn')) == 1 + assert len(s.get_nodes(obj=m)) == 2 + assert len(s.get_nodes(obj=m, obj_name='m')) == 2 + assert len(s.get_nodes(obj=co)) == 0 + nsl = s.add(sl, name='sln') + npl = s.add(pl, parent_name='sln') + assert nsl.children[0] == npl + ndl = s.add(dl, parent_node=npl) + assert npl.children[0] == ndl + nco = s.add(co) + ncp = s.add(cp) + + assert len(s.light_nodes) == len(s.lights) == 3 + assert len(s.point_light_nodes) == len(s.point_lights) == 1 + assert npl in s.point_light_nodes + assert len(s.spot_light_nodes) == len(s.spot_lights) == 1 + assert nsl in s.spot_light_nodes + assert len(s.directional_light_nodes) == len(s.directional_lights) == 1 + assert ndl in s.directional_light_nodes + assert len(s.cameras) == len(s.camera_nodes) == 2 + assert s.main_camera_node == nco + s.main_camera_node = ncp + s.remove_node(ncp) + assert len(s.cameras) == len(s.camera_nodes) == 1 + assert s.main_camera_node == nco + s.remove_node(n2) + assert len(s.meshes) == 1 + s.remove_node(n1) + assert len(s.meshes) == 0 + s.remove_node(nsl) + assert len(s.lights) == 0 + s.remove_node(nco) + assert s.main_camera_node is None + + s.add_node(n1) + s.clear() + assert len(s.nodes) == 0 + + # Trigger final errors + with pytest.raises(ValueError): + s.main_camera_node = None + with pytest.raises(ValueError): + s.main_camera_node = ncp + with pytest.raises(ValueError): + s.add(m, parent_node=n1) + with pytest.raises(ValueError): + s.add(m, name='asdf') + s.add(m, name='asdf') + s.add(m, parent_name='asdf') + with pytest.raises(ValueError): + s.add(m, parent_name='asfd') + with pytest.raises(TypeError): + s.add(None) + + s.clear() + # Test bounds + m1 = Mesh.from_trimesh(trimesh.creation.box()) + m2 = Mesh.from_trimesh(trimesh.creation.box()) + m3 = Mesh.from_trimesh(trimesh.creation.box()) + n1 = Node(mesh=m1) + n2 = Node(mesh=m2, translation=[1.0, 0.0, 0.0]) + n3 = Node(mesh=m3, translation=[0.5, 0.0, 1.0]) + s.add_node(n1) + s.add_node(n2) + s.add_node(n3) + assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [1.5, 0.5, 1.5]]) + s.clear() + s.add_node(n1) + s.add_node(n2, parent_node=n1) + s.add_node(n3, parent_node=n2) + assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [2.0, 0.5, 1.5]]) + tf = np.eye(4) + tf[:3,3] = np.ones(3) + s.set_pose(n3, tf) + assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [2.5, 1.5, 1.5]]) + s.remove_node(n2) + assert np.allclose(s.bounds, [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]]) + s.clear() + assert np.allclose(s.bounds, 0.0) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d897caa121747d2a48a18f60be1260b6caf1153f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,25 @@ +--extra-index-url https://download.pytorch.org/whl/cu116 +torch==1.13.1+cu116 +torchvision==0.14.1+cu116 + +numpy==1.26.3 +opencv-python +pyrender +pytorch-lightning +scikit-image +smplx==0.1.28 +torch +torchvision +yacs +chumpy @ git+https://github.com/mattloper/chumpy +timm +einops +xtcocotools +pandas +hydra-core +hydra-submitit-launcher +hydra-colorlog +pyrootutils +rich +webdataset +ultralytics==8.1.34 diff --git a/wilor/configs/__init__.py b/wilor/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3bd9dcf90daa9627cb482d2c2e602288e09b1f1b --- /dev/null +++ b/wilor/configs/__init__.py @@ -0,0 +1,114 @@ +import os +from typing import Dict +from yacs.config import CfgNode as CN + +CACHE_DIR_PRETRAINED = "./pretrained_models/" + +def to_lower(x: Dict) -> Dict: + """ + Convert all dictionary keys to lowercase + Args: + x (dict): Input dictionary + Returns: + dict: Output dictionary with all keys converted to lowercase + """ + return {k.lower(): v for k, v in x.items()} + +_C = CN(new_allowed=True) + +_C.GENERAL = CN(new_allowed=True) +_C.GENERAL.RESUME = True +_C.GENERAL.TIME_TO_RUN = 3300 +_C.GENERAL.VAL_STEPS = 100 +_C.GENERAL.LOG_STEPS = 100 +_C.GENERAL.CHECKPOINT_STEPS = 20000 +_C.GENERAL.CHECKPOINT_DIR = "checkpoints" +_C.GENERAL.SUMMARY_DIR = "tensorboard" +_C.GENERAL.NUM_GPUS = 1 +_C.GENERAL.NUM_WORKERS = 4 +_C.GENERAL.MIXED_PRECISION = True +_C.GENERAL.ALLOW_CUDA = True +_C.GENERAL.PIN_MEMORY = False +_C.GENERAL.DISTRIBUTED = False +_C.GENERAL.LOCAL_RANK = 0 +_C.GENERAL.USE_SYNCBN = False +_C.GENERAL.WORLD_SIZE = 1 + +_C.TRAIN = CN(new_allowed=True) +_C.TRAIN.NUM_EPOCHS = 100 +_C.TRAIN.BATCH_SIZE = 32 +_C.TRAIN.SHUFFLE = True +_C.TRAIN.WARMUP = False +_C.TRAIN.NORMALIZE_PER_IMAGE = False +_C.TRAIN.CLIP_GRAD = False +_C.TRAIN.CLIP_GRAD_VALUE = 1.0 +_C.LOSS_WEIGHTS = CN(new_allowed=True) + +_C.DATASETS = CN(new_allowed=True) + +_C.MODEL = CN(new_allowed=True) +_C.MODEL.IMAGE_SIZE = 224 + +_C.EXTRA = CN(new_allowed=True) +_C.EXTRA.FOCAL_LENGTH = 5000 + +_C.DATASETS.CONFIG = CN(new_allowed=True) +_C.DATASETS.CONFIG.SCALE_FACTOR = 0.3 +_C.DATASETS.CONFIG.ROT_FACTOR = 30 +_C.DATASETS.CONFIG.TRANS_FACTOR = 0.02 +_C.DATASETS.CONFIG.COLOR_SCALE = 0.2 +_C.DATASETS.CONFIG.ROT_AUG_RATE = 0.6 +_C.DATASETS.CONFIG.TRANS_AUG_RATE = 0.5 +_C.DATASETS.CONFIG.DO_FLIP = False +_C.DATASETS.CONFIG.FLIP_AUG_RATE = 0.5 +_C.DATASETS.CONFIG.EXTREME_CROP_AUG_RATE = 0.10 + +def default_config() -> CN: + """ + Get a yacs CfgNode object with the default config values. + """ + # Return a clone so that the defaults will not be altered + # This is for the "local variable" use pattern + return _C.clone() + +def dataset_config(name='datasets_tar.yaml') -> CN: + """ + Get dataset config file + Returns: + CfgNode: Dataset config as a yacs CfgNode object. + """ + cfg = CN(new_allowed=True) + config_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), name) + cfg.merge_from_file(config_file) + cfg.freeze() + return cfg + +def dataset_eval_config() -> CN: + return dataset_config('datasets_eval.yaml') + +def get_config(config_file: str, merge: bool = True, update_cachedir: bool = False) -> CN: + """ + Read a config file and optionally merge it with the default config file. + Args: + config_file (str): Path to config file. + merge (bool): Whether to merge with the default config or not. + Returns: + CfgNode: Config as a yacs CfgNode object. + """ + if merge: + cfg = default_config() + else: + cfg = CN(new_allowed=True) + cfg.merge_from_file(config_file) + + if update_cachedir: + def update_path(path: str) -> str: + if os.path.isabs(path): + return path + return os.path.join(CACHE_DIR_PRETRAINED, path) + + cfg.MANO.MODEL_PATH = update_path(cfg.MANO.MODEL_PATH) + cfg.MANO.MEAN_PARAMS = update_path(cfg.MANO.MEAN_PARAMS) + + cfg.freeze() + return cfg \ No newline at end of file diff --git a/wilor/datasets/utils.py b/wilor/datasets/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2b09722ef538f25f736ab387e15df259966ce7b7 --- /dev/null +++ b/wilor/datasets/utils.py @@ -0,0 +1,994 @@ +""" +Parts of the code are taken or adapted from +https://github.com/mkocabas/EpipolarPose/blob/master/lib/utils/img_utils.py +""" +import torch +import numpy as np +from skimage.transform import rotate, resize +from skimage.filters import gaussian +import random +import cv2 +from typing import List, Dict, Tuple +from yacs.config import CfgNode + +def expand_to_aspect_ratio(input_shape, target_aspect_ratio=None): + """Increase the size of the bounding box to match the target shape.""" + if target_aspect_ratio is None: + return input_shape + + try: + w , h = input_shape + except (ValueError, TypeError): + return input_shape + + w_t, h_t = target_aspect_ratio + if h / w < h_t / w_t: + h_new = max(w * h_t / w_t, h) + w_new = w + else: + h_new = h + w_new = max(h * w_t / h_t, w) + if h_new < h or w_new < w: + breakpoint() + return np.array([w_new, h_new]) + +def do_augmentation(aug_config: CfgNode) -> Tuple: + """ + Compute random augmentation parameters. + Args: + aug_config (CfgNode): Config containing augmentation parameters. + Returns: + scale (float): Box rescaling factor. + rot (float): Random image rotation. + do_flip (bool): Whether to flip image or not. + do_extreme_crop (bool): Whether to apply extreme cropping (as proposed in EFT). + color_scale (List): Color rescaling factor + tx (float): Random translation along the x axis. + ty (float): Random translation along the y axis. + """ + + tx = np.clip(np.random.randn(), -1.0, 1.0) * aug_config.TRANS_FACTOR + ty = np.clip(np.random.randn(), -1.0, 1.0) * aug_config.TRANS_FACTOR + scale = np.clip(np.random.randn(), -1.0, 1.0) * aug_config.SCALE_FACTOR + 1.0 + rot = np.clip(np.random.randn(), -2.0, + 2.0) * aug_config.ROT_FACTOR if random.random() <= aug_config.ROT_AUG_RATE else 0 + do_flip = aug_config.DO_FLIP and random.random() <= aug_config.FLIP_AUG_RATE + do_extreme_crop = random.random() <= aug_config.EXTREME_CROP_AUG_RATE + extreme_crop_lvl = aug_config.get('EXTREME_CROP_AUG_LEVEL', 0) + # extreme_crop_lvl = 0 + c_up = 1.0 + aug_config.COLOR_SCALE + c_low = 1.0 - aug_config.COLOR_SCALE + color_scale = [random.uniform(c_low, c_up), random.uniform(c_low, c_up), random.uniform(c_low, c_up)] + return scale, rot, do_flip, do_extreme_crop, extreme_crop_lvl, color_scale, tx, ty + +def rotate_2d(pt_2d: np.array, rot_rad: float) -> np.array: + """ + Rotate a 2D point on the x-y plane. + Args: + pt_2d (np.array): Input 2D point with shape (2,). + rot_rad (float): Rotation angle + Returns: + np.array: Rotated 2D point. + """ + x = pt_2d[0] + y = pt_2d[1] + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + xx = x * cs - y * sn + yy = x * sn + y * cs + return np.array([xx, yy], dtype=np.float32) + + +def gen_trans_from_patch_cv(c_x: float, c_y: float, + src_width: float, src_height: float, + dst_width: float, dst_height: float, + scale: float, rot: float) -> np.array: + """ + Create transformation matrix for the bounding box crop. + Args: + c_x (float): Bounding box center x coordinate in the original image. + c_y (float): Bounding box center y coordinate in the original image. + src_width (float): Bounding box width. + src_height (float): Bounding box height. + dst_width (float): Output box width. + dst_height (float): Output box height. + scale (float): Rescaling factor for the bounding box (augmentation). + rot (float): Random rotation applied to the box. + Returns: + trans (np.array): Target geometric transformation. + """ + # augment size with scale + src_w = src_width * scale + src_h = src_height * scale + src_center = np.zeros(2) + src_center[0] = c_x + src_center[1] = c_y + # augment rotation + rot_rad = np.pi * rot / 180 + src_downdir = rotate_2d(np.array([0, src_h * 0.5], dtype=np.float32), rot_rad) + src_rightdir = rotate_2d(np.array([src_w * 0.5, 0], dtype=np.float32), rot_rad) + + dst_w = dst_width + dst_h = dst_height + dst_center = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32) + dst_downdir = np.array([0, dst_h * 0.5], dtype=np.float32) + dst_rightdir = np.array([dst_w * 0.5, 0], dtype=np.float32) + + src = np.zeros((3, 2), dtype=np.float32) + src[0, :] = src_center + src[1, :] = src_center + src_downdir + src[2, :] = src_center + src_rightdir + + dst = np.zeros((3, 2), dtype=np.float32) + dst[0, :] = dst_center + dst[1, :] = dst_center + dst_downdir + dst[2, :] = dst_center + dst_rightdir + + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return trans + + +def trans_point2d(pt_2d: np.array, trans: np.array): + """ + Transform a 2D point using translation matrix trans. + Args: + pt_2d (np.array): Input 2D point with shape (2,). + trans (np.array): Transformation matrix. + Returns: + np.array: Transformed 2D point. + """ + src_pt = np.array([pt_2d[0], pt_2d[1], 1.]).T + dst_pt = np.dot(trans, src_pt) + return dst_pt[0:2] + +def get_transform(center, scale, res, rot=0): + """Generate transformation matrix.""" + """Taken from PARE: https://github.com/mkocabas/PARE/blob/6e0caca86c6ab49ff80014b661350958e5b72fd8/pare/utils/image_utils.py""" + h = 200 * scale + t = np.zeros((3, 3)) + t[0, 0] = float(res[1]) / h + t[1, 1] = float(res[0]) / h + t[0, 2] = res[1] * (-float(center[0]) / h + .5) + t[1, 2] = res[0] * (-float(center[1]) / h + .5) + t[2, 2] = 1 + if not rot == 0: + rot = -rot # To match direction of rotation from cropping + rot_mat = np.zeros((3, 3)) + rot_rad = rot * np.pi / 180 + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + rot_mat[0, :2] = [cs, -sn] + rot_mat[1, :2] = [sn, cs] + rot_mat[2, 2] = 1 + # Need to rotate around center + t_mat = np.eye(3) + t_mat[0, 2] = -res[1] / 2 + t_mat[1, 2] = -res[0] / 2 + t_inv = t_mat.copy() + t_inv[:2, 2] *= -1 + t = np.dot(t_inv, np.dot(rot_mat, np.dot(t_mat, t))) + return t + + +def transform(pt, center, scale, res, invert=0, rot=0, as_int=True): + """Transform pixel location to different reference.""" + """Taken from PARE: https://github.com/mkocabas/PARE/blob/6e0caca86c6ab49ff80014b661350958e5b72fd8/pare/utils/image_utils.py""" + t = get_transform(center, scale, res, rot=rot) + if invert: + t = np.linalg.inv(t) + new_pt = np.array([pt[0] - 1, pt[1] - 1, 1.]).T + new_pt = np.dot(t, new_pt) + if as_int: + new_pt = new_pt.astype(int) + return new_pt[:2] + 1 + +def crop_img(img, ul, br, border_mode=cv2.BORDER_CONSTANT, border_value=0): + c_x = (ul[0] + br[0])/2 + c_y = (ul[1] + br[1])/2 + bb_width = patch_width = br[0] - ul[0] + bb_height = patch_height = br[1] - ul[1] + trans = gen_trans_from_patch_cv(c_x, c_y, bb_width, bb_height, patch_width, patch_height, 1.0, 0) + img_patch = cv2.warpAffine(img, trans, (int(patch_width), int(patch_height)), + flags=cv2.INTER_LINEAR, + borderMode=border_mode, + borderValue=border_value + ) + + # Force borderValue=cv2.BORDER_CONSTANT for alpha channel + if (img.shape[2] == 4) and (border_mode != cv2.BORDER_CONSTANT): + img_patch[:,:,3] = cv2.warpAffine(img[:,:,3], trans, (int(patch_width), int(patch_height)), + flags=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + ) + + return img_patch + +def generate_image_patch_skimage(img: np.array, c_x: float, c_y: float, + bb_width: float, bb_height: float, + patch_width: float, patch_height: float, + do_flip: bool, scale: float, rot: float, + border_mode=cv2.BORDER_CONSTANT, border_value=0) -> Tuple[np.array, np.array]: + """ + Crop image according to the supplied bounding box. + Args: + img (np.array): Input image of shape (H, W, 3) + c_x (float): Bounding box center x coordinate in the original image. + c_y (float): Bounding box center y coordinate in the original image. + bb_width (float): Bounding box width. + bb_height (float): Bounding box height. + patch_width (float): Output box width. + patch_height (float): Output box height. + do_flip (bool): Whether to flip image or not. + scale (float): Rescaling factor for the bounding box (augmentation). + rot (float): Random rotation applied to the box. + Returns: + img_patch (np.array): Cropped image patch of shape (patch_height, patch_height, 3) + trans (np.array): Transformation matrix. + """ + + img_height, img_width, img_channels = img.shape + if do_flip: + img = img[:, ::-1, :] + c_x = img_width - c_x - 1 + + trans = gen_trans_from_patch_cv(c_x, c_y, bb_width, bb_height, patch_width, patch_height, scale, rot) + + #img_patch = cv2.warpAffine(img, trans, (int(patch_width), int(patch_height)), flags=cv2.INTER_LINEAR) + + # skimage + center = np.zeros(2) + center[0] = c_x + center[1] = c_y + res = np.zeros(2) + res[0] = patch_width + res[1] = patch_height + # assumes bb_width = bb_height + # assumes patch_width = patch_height + assert bb_width == bb_height, f'{bb_width=} != {bb_height=}' + assert patch_width == patch_height, f'{patch_width=} != {patch_height=}' + scale1 = scale*bb_width/200. + + # Upper left point + ul = np.array(transform([1, 1], center, scale1, res, invert=1, as_int=False)) - 1 + # Bottom right point + br = np.array(transform([res[0] + 1, + res[1] + 1], center, scale1, res, invert=1, as_int=False)) - 1 + + # Padding so that when rotated proper amount of context is included + try: + pad = int(np.linalg.norm(br - ul) / 2 - float(br[1] - ul[1]) / 2) + 1 + except: + breakpoint() + if not rot == 0: + ul -= pad + br += pad + + + if False: + # Old way of cropping image + ul_int = ul.astype(int) + br_int = br.astype(int) + new_shape = [br_int[1] - ul_int[1], br_int[0] - ul_int[0]] + if len(img.shape) > 2: + new_shape += [img.shape[2]] + new_img = np.zeros(new_shape) + + # Range to fill new array + new_x = max(0, -ul_int[0]), min(br_int[0], len(img[0])) - ul_int[0] + new_y = max(0, -ul_int[1]), min(br_int[1], len(img)) - ul_int[1] + # Range to sample from original image + old_x = max(0, ul_int[0]), min(len(img[0]), br_int[0]) + old_y = max(0, ul_int[1]), min(len(img), br_int[1]) + new_img[new_y[0]:new_y[1], new_x[0]:new_x[1]] = img[old_y[0]:old_y[1], + old_x[0]:old_x[1]] + + # New way of cropping image + new_img = crop_img(img, ul, br, border_mode=border_mode, border_value=border_value).astype(np.float32) + + # print(f'{new_img.shape=}') + # print(f'{new_img1.shape=}') + # print(f'{np.allclose(new_img, new_img1)=}') + # print(f'{img.dtype=}') + + + if not rot == 0: + # Remove padding + + new_img = rotate(new_img, rot) # scipy.misc.imrotate(new_img, rot) + new_img = new_img[pad:-pad, pad:-pad] + + if new_img.shape[0] < 1 or new_img.shape[1] < 1: + print(f'{img.shape=}') + print(f'{new_img.shape=}') + print(f'{ul=}') + print(f'{br=}') + print(f'{pad=}') + print(f'{rot=}') + + breakpoint() + + # resize image + new_img = resize(new_img, res) # scipy.misc.imresize(new_img, res) + + new_img = np.clip(new_img, 0, 255).astype(np.uint8) + + return new_img, trans + + +def generate_image_patch_cv2(img: np.array, c_x: float, c_y: float, + bb_width: float, bb_height: float, + patch_width: float, patch_height: float, + do_flip: bool, scale: float, rot: float, + border_mode=cv2.BORDER_CONSTANT, border_value=0) -> Tuple[np.array, np.array]: + """ + Crop the input image and return the crop and the corresponding transformation matrix. + Args: + img (np.array): Input image of shape (H, W, 3) + c_x (float): Bounding box center x coordinate in the original image. + c_y (float): Bounding box center y coordinate in the original image. + bb_width (float): Bounding box width. + bb_height (float): Bounding box height. + patch_width (float): Output box width. + patch_height (float): Output box height. + do_flip (bool): Whether to flip image or not. + scale (float): Rescaling factor for the bounding box (augmentation). + rot (float): Random rotation applied to the box. + Returns: + img_patch (np.array): Cropped image patch of shape (patch_height, patch_height, 3) + trans (np.array): Transformation matrix. + """ + + img_height, img_width, img_channels = img.shape + if do_flip: + img = img[:, ::-1, :] + c_x = img_width - c_x - 1 + + + trans = gen_trans_from_patch_cv(c_x, c_y, bb_width, bb_height, patch_width, patch_height, scale, rot) + + img_patch = cv2.warpAffine(img, trans, (int(patch_width), int(patch_height)), + flags=cv2.INTER_LINEAR, + borderMode=border_mode, + borderValue=border_value, + ) + # Force borderValue=cv2.BORDER_CONSTANT for alpha channel + if (img.shape[2] == 4) and (border_mode != cv2.BORDER_CONSTANT): + img_patch[:,:,3] = cv2.warpAffine(img[:,:,3], trans, (int(patch_width), int(patch_height)), + flags=cv2.INTER_LINEAR, + borderMode=cv2.BORDER_CONSTANT, + ) + + return img_patch, trans + + +def convert_cvimg_to_tensor(cvimg: np.array): + """ + Convert image from HWC to CHW format. + Args: + cvimg (np.array): Image of shape (H, W, 3) as loaded by OpenCV. + Returns: + np.array: Output image of shape (3, H, W). + """ + # from h,w,c(OpenCV) to c,h,w + img = cvimg.copy() + img = np.transpose(img, (2, 0, 1)) + # from int to float + img = img.astype(np.float32) + return img + +def fliplr_params(mano_params: Dict, has_mano_params: Dict) -> Tuple[Dict, Dict]: + """ + Flip MANO parameters when flipping the image. + Args: + mano_params (Dict): MANO parameter annotations. + has_mano_params (Dict): Whether MANO annotations are valid. + Returns: + Dict, Dict: Flipped MANO parameters and valid flags. + """ + global_orient = mano_params['global_orient'].copy() + hand_pose = mano_params['hand_pose'].copy() + betas = mano_params['betas'].copy() + has_global_orient = has_mano_params['global_orient'].copy() + has_hand_pose = has_mano_params['hand_pose'].copy() + has_betas = has_mano_params['betas'].copy() + + global_orient[1::3] *= -1 + global_orient[2::3] *= -1 + hand_pose[1::3] *= -1 + hand_pose[2::3] *= -1 + + mano_params = {'global_orient': global_orient.astype(np.float32), + 'hand_pose': hand_pose.astype(np.float32), + 'betas': betas.astype(np.float32) + } + + has_mano_params = {'global_orient': has_global_orient, + 'hand_pose': has_hand_pose, + 'betas': has_betas + } + + return mano_params, has_mano_params + + +def fliplr_keypoints(joints: np.array, width: float, flip_permutation: List[int]) -> np.array: + """ + Flip 2D or 3D keypoints. + Args: + joints (np.array): Array of shape (N, 3) or (N, 4) containing 2D or 3D keypoint locations and confidence. + flip_permutation (List): Permutation to apply after flipping. + Returns: + np.array: Flipped 2D or 3D keypoints with shape (N, 3) or (N, 4) respectively. + """ + joints = joints.copy() + # Flip horizontal + joints[:, 0] = width - joints[:, 0] - 1 + joints = joints[flip_permutation, :] + + return joints + +def keypoint_3d_processing(keypoints_3d: np.array, flip_permutation: List[int], rot: float, do_flip: float) -> np.array: + """ + Process 3D keypoints (rotation/flipping). + Args: + keypoints_3d (np.array): Input array of shape (N, 4) containing the 3D keypoints and confidence. + flip_permutation (List): Permutation to apply after flipping. + rot (float): Random rotation applied to the keypoints. + do_flip (bool): Whether to flip keypoints or not. + Returns: + np.array: Transformed 3D keypoints with shape (N, 4). + """ + if do_flip: + keypoints_3d = fliplr_keypoints(keypoints_3d, 1, flip_permutation) + # in-plane rotation + rot_mat = np.eye(3) + if not rot == 0: + rot_rad = -rot * np.pi / 180 + sn,cs = np.sin(rot_rad), np.cos(rot_rad) + rot_mat[0,:2] = [cs, -sn] + rot_mat[1,:2] = [sn, cs] + keypoints_3d[:, :-1] = np.einsum('ij,kj->ki', rot_mat, keypoints_3d[:, :-1]) + # flip the x coordinates + keypoints_3d = keypoints_3d.astype('float32') + return keypoints_3d + +def rot_aa(aa: np.array, rot: float) -> np.array: + """ + Rotate axis angle parameters. + Args: + aa (np.array): Axis-angle vector of shape (3,). + rot (np.array): Rotation angle in degrees. + Returns: + np.array: Rotated axis-angle vector. + """ + # pose parameters + R = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0], + [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0], + [0, 0, 1]]) + # find the rotation of the hand in camera frame + per_rdg, _ = cv2.Rodrigues(aa) + # apply the global rotation to the global orientation + resrot, _ = cv2.Rodrigues(np.dot(R,per_rdg)) + aa = (resrot.T)[0] + return aa.astype(np.float32) + +def mano_param_processing(mano_params: Dict, has_mano_params: Dict, rot: float, do_flip: bool) -> Tuple[Dict, Dict]: + """ + Apply random augmentations to the MANO parameters. + Args: + mano_params (Dict): MANO parameter annotations. + has_mano_params (Dict): Whether mano annotations are valid. + rot (float): Random rotation applied to the keypoints. + do_flip (bool): Whether to flip keypoints or not. + Returns: + Dict, Dict: Transformed MANO parameters and valid flags. + """ + if do_flip: + mano_params, has_mano_params = fliplr_params(mano_params, has_mano_params) + mano_params['global_orient'] = rot_aa(mano_params['global_orient'], rot) + return mano_params, has_mano_params + + + +def get_example(img_path: str|np.ndarray, center_x: float, center_y: float, + width: float, height: float, + keypoints_2d: np.array, keypoints_3d: np.array, + mano_params: Dict, has_mano_params: Dict, + flip_kp_permutation: List[int], + patch_width: int, patch_height: int, + mean: np.array, std: np.array, + do_augment: bool, is_right: bool, augm_config: CfgNode, + is_bgr: bool = True, + use_skimage_antialias: bool = False, + border_mode: int = cv2.BORDER_CONSTANT, + return_trans: bool = False) -> Tuple: + """ + Get an example from the dataset and (possibly) apply random augmentations. + Args: + img_path (str): Image filename + center_x (float): Bounding box center x coordinate in the original image. + center_y (float): Bounding box center y coordinate in the original image. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array with shape (N,3) containing the 2D keypoints in the original image coordinates. + keypoints_3d (np.array): Array with shape (N,4) containing the 3D keypoints. + mano_params (Dict): MANO parameter annotations. + has_mano_params (Dict): Whether MANO annotations are valid. + flip_kp_permutation (List): Permutation to apply to the keypoints after flipping. + patch_width (float): Output box width. + patch_height (float): Output box height. + mean (np.array): Array of shape (3,) containing the mean for normalizing the input image. + std (np.array): Array of shape (3,) containing the std for normalizing the input image. + do_augment (bool): Whether to apply data augmentation or not. + aug_config (CfgNode): Config containing augmentation parameters. + Returns: + return img_patch, keypoints_2d, keypoints_3d, mano_params, has_mano_params, img_size + img_patch (np.array): Cropped image patch of shape (3, patch_height, patch_height) + keypoints_2d (np.array): Array with shape (N,3) containing the transformed 2D keypoints. + keypoints_3d (np.array): Array with shape (N,4) containing the transformed 3D keypoints. + mano_params (Dict): Transformed MANO parameters. + has_mano_params (Dict): Valid flag for transformed MANO parameters. + img_size (np.array): Image size of the original image. + """ + if isinstance(img_path, str): + # 1. load image + cvimg = cv2.imread(img_path, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION) + if not isinstance(cvimg, np.ndarray): + raise IOError("Fail to read %s" % img_path) + elif isinstance(img_path, np.ndarray): + cvimg = img_path + else: + raise TypeError('img_path must be either a string or a numpy array') + img_height, img_width, img_channels = cvimg.shape + + img_size = np.array([img_height, img_width]) + + # 2. get augmentation params + if do_augment: + scale, rot, do_flip, do_extreme_crop, extreme_crop_lvl, color_scale, tx, ty = do_augmentation(augm_config) + else: + scale, rot, do_flip, do_extreme_crop, extreme_crop_lvl, color_scale, tx, ty = 1.0, 0, False, False, 0, [1.0, 1.0, 1.0], 0., 0. + + # if it's a left hand, we flip + if not is_right: + do_flip = True + + if width < 1 or height < 1: + breakpoint() + + if do_extreme_crop: + if extreme_crop_lvl == 0: + center_x1, center_y1, width1, height1 = extreme_cropping(center_x, center_y, width, height, keypoints_2d) + elif extreme_crop_lvl == 1: + center_x1, center_y1, width1, height1 = extreme_cropping_aggressive(center_x, center_y, width, height, keypoints_2d) + + THRESH = 4 + if width1 < THRESH or height1 < THRESH: + # print(f'{do_extreme_crop=}') + # print(f'width: {width}, height: {height}') + # print(f'width1: {width1}, height1: {height1}') + # print(f'center_x: {center_x}, center_y: {center_y}') + # print(f'center_x1: {center_x1}, center_y1: {center_y1}') + # print(f'keypoints_2d: {keypoints_2d}') + # print(f'\n\n', flush=True) + # breakpoint() + pass + # print(f'skip ==> width1: {width1}, height1: {height1}, width: {width}, height: {height}') + else: + center_x, center_y, width, height = center_x1, center_y1, width1, height1 + + center_x += width * tx + center_y += height * ty + + # Process 3D keypoints + keypoints_3d = keypoint_3d_processing(keypoints_3d, flip_kp_permutation, rot, do_flip) + + # 3. generate image patch + if use_skimage_antialias: + # Blur image to avoid aliasing artifacts + downsampling_factor = (patch_width / (width*scale)) + if downsampling_factor > 1.1: + cvimg = gaussian(cvimg, sigma=(downsampling_factor-1)/2, channel_axis=2, preserve_range=True, truncate=3.0) + + img_patch_cv, trans = generate_image_patch_cv2(cvimg, + center_x, center_y, + width, height, + patch_width, patch_height, + do_flip, scale, rot, + border_mode=border_mode) + + # img_patch_cv, trans = generate_image_patch_skimage(cvimg, + # center_x, center_y, + # width, height, + # patch_width, patch_height, + # do_flip, scale, rot, + # border_mode=border_mode) + + image = img_patch_cv.copy() + if is_bgr: + image = image[:, :, ::-1] + img_patch_cv = image.copy() + img_patch = convert_cvimg_to_tensor(image) + + + mano_params, has_mano_params = mano_param_processing(mano_params, has_mano_params, rot, do_flip) + + # apply normalization + for n_c in range(min(img_channels, 3)): + img_patch[n_c, :, :] = np.clip(img_patch[n_c, :, :] * color_scale[n_c], 0, 255) + if mean is not None and std is not None: + img_patch[n_c, :, :] = (img_patch[n_c, :, :] - mean[n_c]) / std[n_c] + if do_flip: + keypoints_2d = fliplr_keypoints(keypoints_2d, img_width, flip_kp_permutation) + + + for n_jt in range(len(keypoints_2d)): + keypoints_2d[n_jt, 0:2] = trans_point2d(keypoints_2d[n_jt, 0:2], trans) + keypoints_2d[:, :-1] = keypoints_2d[:, :-1] / patch_width - 0.5 + + if not return_trans: + return img_patch, keypoints_2d, keypoints_3d, mano_params, has_mano_params, img_size + else: + return img_patch, keypoints_2d, keypoints_3d, mano_params, has_mano_params, img_size, trans + +def crop_to_hips(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array) -> Tuple: + """ + Extreme cropping: Crop the box up to the hip locations. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + lower_body_keypoints = [10, 11, 13, 14, 19, 20, 21, 22, 23, 24, 25+0, 25+1, 25+4, 25+5] + keypoints_2d[lower_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + + +def crop_to_shoulders(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box up to the shoulder locations. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + lower_body_keypoints = [3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 19, 20, 21, 22, 23, 24] + [25 + i for i in [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 14, 15, 16]] + keypoints_2d[lower_body_keypoints, :] = 0 + center, scale = get_bbox(keypoints_2d) + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.2 * scale[0] + height = 1.2 * scale[1] + return center_x, center_y, width, height + +def crop_to_head(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the head. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + lower_body_keypoints = [3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 19, 20, 21, 22, 23, 24] + [25 + i for i in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 14, 15, 16]] + keypoints_2d[lower_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.3 * scale[0] + height = 1.3 * scale[1] + return center_x, center_y, width, height + +def crop_torso_only(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the torso. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + nontorso_body_keypoints = [0, 3, 4, 6, 7, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + [25 + i for i in [0, 1, 4, 5, 6, 7, 10, 11, 13, 17, 18]] + keypoints_2d[nontorso_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + +def crop_rightarm_only(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the right arm. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + nonrightarm_body_keypoints = [0, 1, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + [25 + i for i in [0, 1, 2, 3, 4, 5, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]] + keypoints_2d[nonrightarm_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + +def crop_leftarm_only(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the left arm. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + nonleftarm_body_keypoints = [0, 1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24] + [25 + i for i in [0, 1, 2, 3, 4, 5, 6, 7, 8, 12, 13, 14, 15, 16, 17, 18]] + keypoints_2d[nonleftarm_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + +def crop_legs_only(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the legs. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + nonlegs_body_keypoints = [0, 1, 2, 3, 4, 5, 6, 7, 15, 16, 17, 18] + [25 + i for i in [6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17, 18]] + keypoints_2d[nonlegs_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + +def crop_rightleg_only(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the right leg. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + nonrightleg_body_keypoints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21] + [25 + i for i in [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]] + keypoints_2d[nonrightleg_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + +def crop_leftleg_only(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array): + """ + Extreme cropping: Crop the box and keep on only the left leg. + Args: + center_x (float): x coordinate of the bounding box center. + center_y (float): y coordinate of the bounding box center. + width (float): Bounding box width. + height (float): Bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + center_x (float): x coordinate of the new bounding box center. + center_y (float): y coordinate of the new bounding box center. + width (float): New bounding box width. + height (float): New bounding box height. + """ + keypoints_2d = keypoints_2d.copy() + nonleftleg_body_keypoints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 15, 16, 17, 18, 22, 23, 24] + [25 + i for i in [0, 1, 2, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]] + keypoints_2d[nonleftleg_body_keypoints, :] = 0 + if keypoints_2d[:, -1].sum() > 1: + center, scale = get_bbox(keypoints_2d) + center_x = center[0] + center_y = center[1] + width = 1.1 * scale[0] + height = 1.1 * scale[1] + return center_x, center_y, width, height + +def full_body(keypoints_2d: np.array) -> bool: + """ + Check if all main body joints are visible. + Args: + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + bool: True if all main body joints are visible. + """ + + body_keypoints_openpose = [2, 3, 4, 5, 6, 7, 10, 11, 13, 14] + body_keypoints = [25 + i for i in [8, 7, 6, 9, 10, 11, 1, 0, 4, 5]] + return (np.maximum(keypoints_2d[body_keypoints, -1], keypoints_2d[body_keypoints_openpose, -1]) > 0).sum() == len(body_keypoints) + +def upper_body(keypoints_2d: np.array): + """ + Check if all upper body joints are visible. + Args: + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + Returns: + bool: True if all main body joints are visible. + """ + lower_body_keypoints_openpose = [10, 11, 13, 14] + lower_body_keypoints = [25 + i for i in [1, 0, 4, 5]] + upper_body_keypoints_openpose = [0, 1, 15, 16, 17, 18] + upper_body_keypoints = [25+8, 25+9, 25+12, 25+13, 25+17, 25+18] + return ((keypoints_2d[lower_body_keypoints + lower_body_keypoints_openpose, -1] > 0).sum() == 0)\ + and ((keypoints_2d[upper_body_keypoints + upper_body_keypoints_openpose, -1] > 0).sum() >= 2) + +def get_bbox(keypoints_2d: np.array, rescale: float = 1.2) -> Tuple: + """ + Get center and scale for bounding box from openpose detections. + Args: + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + rescale (float): Scale factor to rescale bounding boxes computed from the keypoints. + Returns: + center (np.array): Array of shape (2,) containing the new bounding box center. + scale (float): New bounding box scale. + """ + valid = keypoints_2d[:,-1] > 0 + valid_keypoints = keypoints_2d[valid][:,:-1] + center = 0.5 * (valid_keypoints.max(axis=0) + valid_keypoints.min(axis=0)) + bbox_size = (valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)) + # adjust bounding box tightness + scale = bbox_size + scale *= rescale + return center, scale + +def extreme_cropping(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array) -> Tuple: + """ + Perform extreme cropping + Args: + center_x (float): x coordinate of bounding box center. + center_y (float): y coordinate of bounding box center. + width (float): bounding box width. + height (float): bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + rescale (float): Scale factor to rescale bounding boxes computed from the keypoints. + Returns: + center_x (float): x coordinate of bounding box center. + center_y (float): y coordinate of bounding box center. + width (float): bounding box width. + height (float): bounding box height. + """ + p = torch.rand(1).item() + if full_body(keypoints_2d): + if p < 0.7: + center_x, center_y, width, height = crop_to_hips(center_x, center_y, width, height, keypoints_2d) + elif p < 0.9: + center_x, center_y, width, height = crop_to_shoulders(center_x, center_y, width, height, keypoints_2d) + else: + center_x, center_y, width, height = crop_to_head(center_x, center_y, width, height, keypoints_2d) + elif upper_body(keypoints_2d): + if p < 0.9: + center_x, center_y, width, height = crop_to_shoulders(center_x, center_y, width, height, keypoints_2d) + else: + center_x, center_y, width, height = crop_to_head(center_x, center_y, width, height, keypoints_2d) + + return center_x, center_y, max(width, height), max(width, height) + +def extreme_cropping_aggressive(center_x: float, center_y: float, width: float, height: float, keypoints_2d: np.array) -> Tuple: + """ + Perform aggressive extreme cropping + Args: + center_x (float): x coordinate of bounding box center. + center_y (float): y coordinate of bounding box center. + width (float): bounding box width. + height (float): bounding box height. + keypoints_2d (np.array): Array of shape (N, 3) containing 2D keypoint locations. + rescale (float): Scale factor to rescale bounding boxes computed from the keypoints. + Returns: + center_x (float): x coordinate of bounding box center. + center_y (float): y coordinate of bounding box center. + width (float): bounding box width. + height (float): bounding box height. + """ + p = torch.rand(1).item() + if full_body(keypoints_2d): + if p < 0.2: + center_x, center_y, width, height = crop_to_hips(center_x, center_y, width, height, keypoints_2d) + elif p < 0.3: + center_x, center_y, width, height = crop_to_shoulders(center_x, center_y, width, height, keypoints_2d) + elif p < 0.4: + center_x, center_y, width, height = crop_to_head(center_x, center_y, width, height, keypoints_2d) + elif p < 0.5: + center_x, center_y, width, height = crop_torso_only(center_x, center_y, width, height, keypoints_2d) + elif p < 0.6: + center_x, center_y, width, height = crop_rightarm_only(center_x, center_y, width, height, keypoints_2d) + elif p < 0.7: + center_x, center_y, width, height = crop_leftarm_only(center_x, center_y, width, height, keypoints_2d) + elif p < 0.8: + center_x, center_y, width, height = crop_legs_only(center_x, center_y, width, height, keypoints_2d) + elif p < 0.9: + center_x, center_y, width, height = crop_rightleg_only(center_x, center_y, width, height, keypoints_2d) + else: + center_x, center_y, width, height = crop_leftleg_only(center_x, center_y, width, height, keypoints_2d) + elif upper_body(keypoints_2d): + if p < 0.2: + center_x, center_y, width, height = crop_to_shoulders(center_x, center_y, width, height, keypoints_2d) + elif p < 0.4: + center_x, center_y, width, height = crop_to_head(center_x, center_y, width, height, keypoints_2d) + elif p < 0.6: + center_x, center_y, width, height = crop_torso_only(center_x, center_y, width, height, keypoints_2d) + elif p < 0.8: + center_x, center_y, width, height = crop_rightarm_only(center_x, center_y, width, height, keypoints_2d) + else: + center_x, center_y, width, height = crop_leftarm_only(center_x, center_y, width, height, keypoints_2d) + return center_x, center_y, max(width, height), max(width, height) diff --git a/wilor/datasets/vitdet_dataset.py b/wilor/datasets/vitdet_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..1fd6e8b9bc52d927c98959174bc556bf796f7162 --- /dev/null +++ b/wilor/datasets/vitdet_dataset.py @@ -0,0 +1,95 @@ +from typing import Dict + +import cv2 +import numpy as np +from skimage.filters import gaussian +from yacs.config import CfgNode +import torch + +from .utils import (convert_cvimg_to_tensor, + expand_to_aspect_ratio, + generate_image_patch_cv2) + +DEFAULT_MEAN = 255. * np.array([0.485, 0.456, 0.406]) +DEFAULT_STD = 255. * np.array([0.229, 0.224, 0.225]) + +class ViTDetDataset(torch.utils.data.Dataset): + + def __init__(self, + cfg: CfgNode, + img_cv2: np.array, + boxes: np.array, + right: np.array, + rescale_factor=2.5, + train: bool = False, + **kwargs): + super().__init__() + self.cfg = cfg + self.img_cv2 = img_cv2 + # self.boxes = boxes + + assert train == False, "ViTDetDataset is only for inference" + self.train = train + self.img_size = cfg.MODEL.IMAGE_SIZE + self.mean = 255. * np.array(self.cfg.MODEL.IMAGE_MEAN) + self.std = 255. * np.array(self.cfg.MODEL.IMAGE_STD) + + # Preprocess annotations + boxes = boxes.astype(np.float32) + self.center = (boxes[:, 2:4] + boxes[:, 0:2]) / 2.0 + self.scale = rescale_factor * (boxes[:, 2:4] - boxes[:, 0:2]) / 200.0 + self.personid = np.arange(len(boxes), dtype=np.int32) + self.right = right.astype(np.float32) + + def __len__(self) -> int: + return len(self.personid) + + def __getitem__(self, idx: int) -> Dict[str, np.array]: + + center = self.center[idx].copy() + center_x = center[0] + center_y = center[1] + + scale = self.scale[idx] + BBOX_SHAPE = self.cfg.MODEL.get('BBOX_SHAPE', None) + bbox_size = expand_to_aspect_ratio(scale*200, target_aspect_ratio=BBOX_SHAPE).max() + + patch_width = patch_height = self.img_size + + right = self.right[idx].copy() + flip = right == 0 + + # 3. generate image patch + # if use_skimage_antialias: + cvimg = self.img_cv2.copy() + if True: + # Blur image to avoid aliasing artifacts + downsampling_factor = ((bbox_size*1.0) / patch_width) + #print(f'{downsampling_factor=}') + downsampling_factor = downsampling_factor / 2.0 + if downsampling_factor > 1.1: + cvimg = gaussian(cvimg, sigma=(downsampling_factor-1)/2, channel_axis=2, preserve_range=True) + + + img_patch_cv, trans = generate_image_patch_cv2(cvimg, + center_x, center_y, + bbox_size, bbox_size, + patch_width, patch_height, + flip, 1.0, 0, + border_mode=cv2.BORDER_CONSTANT) + img_patch_cv = img_patch_cv[:, :, ::-1] + img_patch = convert_cvimg_to_tensor(img_patch_cv) + + # apply normalization + for n_c in range(min(self.img_cv2.shape[2], 3)): + img_patch[n_c, :, :] = (img_patch[n_c, :, :] - self.mean[n_c]) / self.std[n_c] + + item = { + 'img': img_patch, + 'personid': int(self.personid[idx]), + } + item['box_center'] = self.center[idx].copy() + item['box_size'] = bbox_size + item['img_size'] = 1.0 * np.array([cvimg.shape[1], cvimg.shape[0]]) + item['right'] = self.right[idx].copy() + return item diff --git a/wilor/models/__init__.py b/wilor/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94187b99b6c63ce2113e72d5b2ffa906ab673cc6 --- /dev/null +++ b/wilor/models/__init__.py @@ -0,0 +1,36 @@ +from .mano_wrapper import MANO +from .wilor import WiLoR + +from .discriminator import Discriminator + +def load_wilor(checkpoint_path, cfg_path): + from pathlib import Path + from wilor.configs import get_config + print('Loading ', checkpoint_path) + model_cfg = get_config(cfg_path, update_cachedir=True) + + # Override some config values, to crop bbox correctly + if ('vit' in model_cfg.MODEL.BACKBONE.TYPE) and ('BBOX_SHAPE' not in model_cfg.MODEL): + + model_cfg.defrost() + assert model_cfg.MODEL.IMAGE_SIZE == 256, f"MODEL.IMAGE_SIZE ({model_cfg.MODEL.IMAGE_SIZE}) should be 256 for ViT backbone" + model_cfg.MODEL.BBOX_SHAPE = [192,256] + model_cfg.freeze() + + # Update config to be compatible with demo + if ('PRETRAINED_WEIGHTS' in model_cfg.MODEL.BACKBONE): + model_cfg.defrost() + model_cfg.MODEL.BACKBONE.pop('PRETRAINED_WEIGHTS') + model_cfg.freeze() + + # Update config to be compatible with demo + + if ('DATA_DIR' in model_cfg.MANO): + model_cfg.defrost() + model_cfg.MANO.DATA_DIR = './mano_data/' + model_cfg.MANO.MODEL_PATH = './mano_data/mano/' + model_cfg.MANO.MEAN_PARAMS = './mano_data/mano_mean_params.npz' + model_cfg.freeze() + + model = WiLoR.load_from_checkpoint(checkpoint_path, strict=False, cfg=model_cfg) + return model, model_cfg \ No newline at end of file diff --git a/wilor/models/backbones/__init__.py b/wilor/models/backbones/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..860912187d1fb6a12404f0982eede9e5d102772a --- /dev/null +++ b/wilor/models/backbones/__init__.py @@ -0,0 +1,20 @@ +from .vit import vit +import numpy as np + +def create_backbone(cfg): + if cfg.MODEL.BACKBONE.TYPE == 'vit': + return vit(cfg) + elif cfg.MODEL.BACKBONE.TYPE == 'fast_vit': + import torch + import sys + #import models + from timm.models import create_model + + + fast_vit = create_model("fastvit_ma36", drop_path_rate=0.2) + checkpoint = torch.load('./pretrained_vit/fastvit_ma36.pt') + fast_vit.load_state_dict(checkpoint['state_dict']) + return fast_vit + + else: + raise NotImplementedError('Backbone type is not implemented') diff --git a/wilor/models/backbones/vit.py b/wilor/models/backbones/vit.py new file mode 100644 index 0000000000000000000000000000000000000000..704e4922f6d9688d9434fd8598acbcda4ca6c0e7 --- /dev/null +++ b/wilor/models/backbones/vit.py @@ -0,0 +1,410 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import math +import numpy as np +import torch +from functools import partial +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from ...utils.geometry import rot6d_to_rotmat, aa_to_rotmat +from timm.models.layers import drop_path, to_2tuple, trunc_normal_ + +def vit(cfg): + return ViT( + img_size=(256, 192), + patch_size=16, + embed_dim=1280, + depth=32, + num_heads=16, + ratio=1, + use_checkpoint=False, + mlp_ratio=4, + qkv_bias=True, + drop_path_rate=0.55, + cfg = cfg + ) + +def get_abs_pos(abs_pos, h, w, ori_h, ori_w, has_cls_token=True): + """ + Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token + dimension for the original embeddings. + Args: + abs_pos (Tensor): absolute positional embeddings with (1, num_position, C). + has_cls_token (bool): If true, has 1 embedding in abs_pos for cls token. + hw (Tuple): size of input image tokens. + + Returns: + Absolute positional embeddings after processing with shape (1, H, W, C) + """ + cls_token = None + B, L, C = abs_pos.shape + if has_cls_token: + cls_token = abs_pos[:, 0:1] + abs_pos = abs_pos[:, 1:] + + if ori_h != h or ori_w != w: + new_abs_pos = F.interpolate( + abs_pos.reshape(1, ori_h, ori_w, -1).permute(0, 3, 1, 2), + size=(h, w), + mode="bicubic", + align_corners=False, + ).permute(0, 2, 3, 1).reshape(B, -1, C) + + else: + new_abs_pos = abs_pos + + if cls_token is not None: + new_abs_pos = torch.cat([cls_token, new_abs_pos], dim=1) + return new_abs_pos + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + def __init__(self, drop_prob=None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + def extra_repr(self): + return 'p={}'.format(self.drop_prob) + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + x = self.drop(x) + return x + +class Attention(nn.Module): + def __init__( + self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., + proj_drop=0., attn_head_dim=None,): + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.dim = dim + + if attn_head_dim is not None: + head_dim = attn_head_dim + all_head_dim = head_dim * self.num_heads + + self.scale = qk_scale or head_dim ** -0.5 + + self.qkv = nn.Linear(dim, all_head_dim * 3, bias=qkv_bias) + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(all_head_dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x): + B, N, C = x.shape + qkv = self.qkv(x) + qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = (q @ k.transpose(-2, -1)) + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B, N, -1) + x = self.proj(x) + x = self.proj_drop(x) + + return x + +class Block(nn.Module): + + def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, + drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, + norm_layer=nn.LayerNorm, attn_head_dim=None + ): + super().__init__() + + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, + attn_drop=attn_drop, proj_drop=drop, attn_head_dim=attn_head_dim + ) + + # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + def forward(self, x): + x = x + self.drop_path(self.attn(self.norm1(x))) + x = x + self.drop_path(self.mlp(self.norm2(x))) + return x + + +class PatchEmbed(nn.Module): + """ Image to Patch Embedding + """ + def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, ratio=1): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (ratio ** 2) + self.patch_shape = (int(img_size[0] // patch_size[0] * ratio), int(img_size[1] // patch_size[1] * ratio)) + self.origin_patch_shape = (int(img_size[0] // patch_size[0]), int(img_size[1] // patch_size[1])) + self.img_size = img_size + self.patch_size = patch_size + self.num_patches = num_patches + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=(patch_size[0] // ratio), padding=4 + 2 * (ratio//2-1)) + + def forward(self, x, **kwargs): + B, C, H, W = x.shape + x = self.proj(x) + Hp, Wp = x.shape[2], x.shape[3] + + x = x.flatten(2).transpose(1, 2) + return x, (Hp, Wp) + + +class HybridEmbed(nn.Module): + """ CNN Feature Map Embedding + Extract feature map from CNN, flatten, project to embedding dim. + """ + def __init__(self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768): + super().__init__() + assert isinstance(backbone, nn.Module) + img_size = to_2tuple(img_size) + self.img_size = img_size + self.backbone = backbone + if feature_size is None: + with torch.no_grad(): + training = backbone.training + if training: + backbone.eval() + o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))[-1] + feature_size = o.shape[-2:] + feature_dim = o.shape[1] + backbone.train(training) + else: + feature_size = to_2tuple(feature_size) + feature_dim = self.backbone.feature_info.channels()[-1] + self.num_patches = feature_size[0] * feature_size[1] + self.proj = nn.Linear(feature_dim, embed_dim) + + def forward(self, x): + x = self.backbone(x)[-1] + x = x.flatten(2).transpose(1, 2) + x = self.proj(x) + return x + + +class ViT(nn.Module): + + def __init__(self, + img_size=224, patch_size=16, in_chans=3, num_classes=80, embed_dim=768, depth=12, + num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0., + drop_path_rate=0., hybrid_backbone=None, norm_layer=None, use_checkpoint=False, + frozen_stages=-1, ratio=1, last_norm=True, + patch_padding='pad', freeze_attn=False, freeze_ffn=False,cfg=None, + ): + # Protect mutable default arguments + super(ViT, self).__init__() + norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) + self.num_classes = num_classes + self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models + self.frozen_stages = frozen_stages + self.use_checkpoint = use_checkpoint + self.patch_padding = patch_padding + self.freeze_attn = freeze_attn + self.freeze_ffn = freeze_ffn + self.depth = depth + + if hybrid_backbone is not None: + self.patch_embed = HybridEmbed( + hybrid_backbone, img_size=img_size, in_chans=in_chans, embed_dim=embed_dim) + else: + self.patch_embed = PatchEmbed( + img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, ratio=ratio) + num_patches = self.patch_embed.num_patches + + ########################################## + self.cfg = cfg + self.joint_rep_type = cfg.MODEL.MANO_HEAD.get('JOINT_REP', '6d') + self.joint_rep_dim = {'6d': 6, 'aa': 3}[self.joint_rep_type] + npose = self.joint_rep_dim * (cfg.MANO.NUM_HAND_JOINTS + 1) + self.npose = npose + mean_params = np.load(cfg.MANO.MEAN_PARAMS) + init_cam = torch.from_numpy(mean_params['cam'].astype(np.float32)).unsqueeze(0) + self.register_buffer('init_cam', init_cam) + init_hand_pose = torch.from_numpy(mean_params['pose'].astype(np.float32)).unsqueeze(0) + init_betas = torch.from_numpy(mean_params['shape'].astype('float32')).unsqueeze(0) + self.register_buffer('init_hand_pose', init_hand_pose) + self.register_buffer('init_betas', init_betas) + + self.pose_emb = nn.Linear(self.joint_rep_dim , embed_dim) + self.shape_emb = nn.Linear(10 , embed_dim) + self.cam_emb = nn.Linear(3 , embed_dim) + + self.decpose = nn.Linear(self.num_features, 6) + self.decshape = nn.Linear(self.num_features, 10) + self.deccam = nn.Linear(self.num_features, 3) + if cfg.MODEL.MANO_HEAD.get('INIT_DECODER_XAVIER', False): + # True by default in MLP. False by default in Transformer + nn.init.xavier_uniform_(self.decpose.weight, gain=0.01) + nn.init.xavier_uniform_(self.decshape.weight, gain=0.01) + nn.init.xavier_uniform_(self.deccam.weight, gain=0.01) + + + ########################################## + + # since the pretraining model has class token + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule + + self.blocks = nn.ModuleList([ + Block( + dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, + ) + for i in range(depth)]) + + self.last_norm = norm_layer(embed_dim) if last_norm else nn.Identity() + + if self.pos_embed is not None: + trunc_normal_(self.pos_embed, std=.02) + + self._freeze_stages() + + def _freeze_stages(self): + """Freeze parameters.""" + if self.frozen_stages >= 0: + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + + for i in range(1, self.frozen_stages + 1): + m = self.blocks[i] + m.eval() + for param in m.parameters(): + param.requires_grad = False + + if self.freeze_attn: + for i in range(0, self.depth): + m = self.blocks[i] + m.attn.eval() + m.norm1.eval() + for param in m.attn.parameters(): + param.requires_grad = False + for param in m.norm1.parameters(): + param.requires_grad = False + + if self.freeze_ffn: + self.pos_embed.requires_grad = False + self.patch_embed.eval() + for param in self.patch_embed.parameters(): + param.requires_grad = False + for i in range(0, self.depth): + m = self.blocks[i] + m.mlp.eval() + m.norm2.eval() + for param in m.mlp.parameters(): + param.requires_grad = False + for param in m.norm2.parameters(): + param.requires_grad = False + + def init_weights(self): + """Initialize the weights in backbone. + Args: + pretrained (str, optional): Path to pre-trained weights. + Defaults to None. + """ + def _init_weights(m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + self.apply(_init_weights) + + def get_num_layers(self): + return len(self.blocks) + + @torch.jit.ignore + def no_weight_decay(self): + return {'pos_embed', 'cls_token'} + + def forward_features(self, x): + B, C, H, W = x.shape + x, (Hp, Wp) = self.patch_embed(x) + + if self.pos_embed is not None: + # fit for multiple GPU training + # since the first element for pos embed (sin-cos manner) is zero, it will cause no difference + x = x + self.pos_embed[:, 1:] + self.pos_embed[:, :1] + # X [B, 192, 1280] + # x cat [ mean_pose, mean_shape, mean_cam] tokens + pose_tokens = self.pose_emb(self.init_hand_pose.reshape(1, self.cfg.MANO.NUM_HAND_JOINTS + 1, self.joint_rep_dim)).repeat(B, 1, 1) + shape_tokens = self.shape_emb(self.init_betas).unsqueeze(1).repeat(B, 1, 1) + cam_tokens = self.cam_emb(self.init_cam).unsqueeze(1).repeat(B, 1, 1) + + x = torch.cat([pose_tokens, shape_tokens, cam_tokens, x], 1) + for blk in self.blocks: + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x) + else: + x = blk(x) + + x = self.last_norm(x) + + + pose_feat = x[:, :(self.cfg.MANO.NUM_HAND_JOINTS + 1)] + shape_feat = x[:, (self.cfg.MANO.NUM_HAND_JOINTS + 1):1+(self.cfg.MANO.NUM_HAND_JOINTS + 1)] + cam_feat = x[:, 1+(self.cfg.MANO.NUM_HAND_JOINTS + 1):2+(self.cfg.MANO.NUM_HAND_JOINTS + 1)] + + #print(pose_feat.shape, shape_feat.shape, cam_feat.shape) + pred_hand_pose = self.decpose(pose_feat).reshape(B, -1) + self.init_hand_pose #B , 96 + pred_betas = self.decshape(shape_feat).reshape(B, -1) + self.init_betas #B , 10 + pred_cam = self.deccam(cam_feat).reshape(B, -1) + self.init_cam #B , 3 + + pred_mano_feats = {} + pred_mano_feats['hand_pose'] = pred_hand_pose + pred_mano_feats['betas'] = pred_betas + pred_mano_feats['cam'] = pred_cam + + + joint_conversion_fn = { + '6d': rot6d_to_rotmat, + 'aa': lambda x: aa_to_rotmat(x.view(-1, 3).contiguous()) + }[self.joint_rep_type] + + pred_hand_pose = joint_conversion_fn(pred_hand_pose).view(B, self.cfg.MANO.NUM_HAND_JOINTS+1, 3, 3) + pred_mano_params = {'global_orient': pred_hand_pose[:, [0]], + 'hand_pose': pred_hand_pose[:, 1:], + 'betas': pred_betas} + + img_feat = x[:, 2+(self.cfg.MANO.NUM_HAND_JOINTS + 1):].reshape(B, Hp, Wp, -1).permute(0, 3, 1, 2) + return pred_mano_params, pred_cam, pred_mano_feats, img_feat + + def forward(self, x): + x = self.forward_features(x) + return x + + def train(self, mode=True): + """Convert the model into training mode.""" + super().train(mode) + self._freeze_stages() diff --git a/wilor/models/discriminator.py b/wilor/models/discriminator.py new file mode 100644 index 0000000000000000000000000000000000000000..f1cb2d1a21fbab47e8fa10dcc603b3d2012686a7 --- /dev/null +++ b/wilor/models/discriminator.py @@ -0,0 +1,98 @@ +import torch +import torch.nn as nn + +class Discriminator(nn.Module): + + def __init__(self): + """ + Pose + Shape discriminator proposed in HMR + """ + super(Discriminator, self).__init__() + + self.num_joints = 15 + # poses_alone + self.D_conv1 = nn.Conv2d(9, 32, kernel_size=1) + nn.init.xavier_uniform_(self.D_conv1.weight) + nn.init.zeros_(self.D_conv1.bias) + self.relu = nn.ReLU(inplace=True) + self.D_conv2 = nn.Conv2d(32, 32, kernel_size=1) + nn.init.xavier_uniform_(self.D_conv2.weight) + nn.init.zeros_(self.D_conv2.bias) + pose_out = [] + for i in range(self.num_joints): + pose_out_temp = nn.Linear(32, 1) + nn.init.xavier_uniform_(pose_out_temp.weight) + nn.init.zeros_(pose_out_temp.bias) + pose_out.append(pose_out_temp) + self.pose_out = nn.ModuleList(pose_out) + + # betas + self.betas_fc1 = nn.Linear(10, 10) + nn.init.xavier_uniform_(self.betas_fc1.weight) + nn.init.zeros_(self.betas_fc1.bias) + self.betas_fc2 = nn.Linear(10, 5) + nn.init.xavier_uniform_(self.betas_fc2.weight) + nn.init.zeros_(self.betas_fc2.bias) + self.betas_out = nn.Linear(5, 1) + nn.init.xavier_uniform_(self.betas_out.weight) + nn.init.zeros_(self.betas_out.bias) + + # poses_joint + self.D_alljoints_fc1 = nn.Linear(32*self.num_joints, 1024) + nn.init.xavier_uniform_(self.D_alljoints_fc1.weight) + nn.init.zeros_(self.D_alljoints_fc1.bias) + self.D_alljoints_fc2 = nn.Linear(1024, 1024) + nn.init.xavier_uniform_(self.D_alljoints_fc2.weight) + nn.init.zeros_(self.D_alljoints_fc2.bias) + self.D_alljoints_out = nn.Linear(1024, 1) + nn.init.xavier_uniform_(self.D_alljoints_out.weight) + nn.init.zeros_(self.D_alljoints_out.bias) + + + def forward(self, poses: torch.Tensor, betas: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the discriminator. + Args: + poses (torch.Tensor): Tensor of shape (B, 23, 3, 3) containing a batch of MANO hand poses (excluding the global orientation). + betas (torch.Tensor): Tensor of shape (B, 10) containign a batch of MANO beta coefficients. + Returns: + torch.Tensor: Discriminator output with shape (B, 25) + """ + #bn = poses.shape[0] + # poses B x 207 + #poses = poses.reshape(bn, -1) + # poses B x num_joints x 1 x 9 + poses = poses.reshape(-1, self.num_joints, 1, 9) + bn = poses.shape[0] + # poses B x 9 x num_joints x 1 + poses = poses.permute(0, 3, 1, 2).contiguous() + + # poses_alone + poses = self.D_conv1(poses) + poses = self.relu(poses) + poses = self.D_conv2(poses) + poses = self.relu(poses) + + poses_out = [] + for i in range(self.num_joints): + poses_out_ = self.pose_out[i](poses[:, :, i, 0]) + poses_out.append(poses_out_) + poses_out = torch.cat(poses_out, dim=1) + + # betas + betas = self.betas_fc1(betas) + betas = self.relu(betas) + betas = self.betas_fc2(betas) + betas = self.relu(betas) + betas_out = self.betas_out(betas) + + # poses_joint + poses = poses.reshape(bn,-1) + poses_all = self.D_alljoints_fc1(poses) + poses_all = self.relu(poses_all) + poses_all = self.D_alljoints_fc2(poses_all) + poses_all = self.relu(poses_all) + poses_all_out = self.D_alljoints_out(poses_all) + + disc_out = torch.cat((poses_out, betas_out, poses_all_out), 1) + return disc_out diff --git a/wilor/models/heads/__init__.py b/wilor/models/heads/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..40279d65814d344cd6a6f452356c4ac4e6a633b3 --- /dev/null +++ b/wilor/models/heads/__init__.py @@ -0,0 +1 @@ +from .refinement_net import RefineNet \ No newline at end of file diff --git a/wilor/models/heads/refinement_net.py b/wilor/models/heads/refinement_net.py new file mode 100644 index 0000000000000000000000000000000000000000..98cd8ef17617b2f5b5efbd0731aab103be77a361 --- /dev/null +++ b/wilor/models/heads/refinement_net.py @@ -0,0 +1,204 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +from ...utils.geometry import rot6d_to_rotmat, aa_to_rotmat +from typing import Optional + +def make_linear_layers(feat_dims, relu_final=True, use_bn=False): + layers = [] + for i in range(len(feat_dims)-1): + layers.append(nn.Linear(feat_dims[i], feat_dims[i+1])) + + # Do not use ReLU for final estimation + if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and relu_final): + if use_bn: + layers.append(nn.BatchNorm1d(feat_dims[i+1])) + layers.append(nn.ReLU(inplace=True)) + + return nn.Sequential(*layers) + +def make_conv_layers(feat_dims, kernel=3, stride=1, padding=1, bnrelu_final=True): + layers = [] + for i in range(len(feat_dims)-1): + layers.append( + nn.Conv2d( + in_channels=feat_dims[i], + out_channels=feat_dims[i+1], + kernel_size=kernel, + stride=stride, + padding=padding + )) + # Do not use BN and ReLU for final estimation + if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final): + layers.append(nn.BatchNorm2d(feat_dims[i+1])) + layers.append(nn.ReLU(inplace=True)) + + return nn.Sequential(*layers) + +def make_deconv_layers(feat_dims, bnrelu_final=True): + layers = [] + for i in range(len(feat_dims)-1): + layers.append( + nn.ConvTranspose2d( + in_channels=feat_dims[i], + out_channels=feat_dims[i+1], + kernel_size=4, + stride=2, + padding=1, + output_padding=0, + bias=False)) + + # Do not use BN and ReLU for final estimation + if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final): + layers.append(nn.BatchNorm2d(feat_dims[i+1])) + layers.append(nn.ReLU(inplace=True)) + + return nn.Sequential(*layers) + +def sample_joint_features(img_feat, joint_xy): + height, width = img_feat.shape[2:] + x = joint_xy[:, :, 0] / (width - 1) * 2 - 1 + y = joint_xy[:, :, 1] / (height - 1) * 2 - 1 + grid = torch.stack((x, y), 2)[:, :, None, :] + img_feat = F.grid_sample(img_feat, grid, align_corners=True)[:, :, :, 0] # batch_size, channel_dim, joint_num + img_feat = img_feat.permute(0, 2, 1).contiguous() # batch_size, joint_num, channel_dim + return img_feat + +def perspective_projection(points: torch.Tensor, + translation: torch.Tensor, + focal_length: torch.Tensor, + camera_center: Optional[torch.Tensor] = None, + rotation: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Computes the perspective projection of a set of 3D points. + Args: + points (torch.Tensor): Tensor of shape (B, N, 3) containing the input 3D points. + translation (torch.Tensor): Tensor of shape (B, 3) containing the 3D camera translation. + focal_length (torch.Tensor): Tensor of shape (B, 2) containing the focal length in pixels. + camera_center (torch.Tensor): Tensor of shape (B, 2) containing the camera center in pixels. + rotation (torch.Tensor): Tensor of shape (B, 3, 3) containing the camera rotation. + Returns: + torch.Tensor: Tensor of shape (B, N, 2) containing the projection of the input points. + """ + batch_size = points.shape[0] + if rotation is None: + rotation = torch.eye(3, device=points.device, dtype=points.dtype).unsqueeze(0).expand(batch_size, -1, -1) + if camera_center is None: + camera_center = torch.zeros(batch_size, 2, device=points.device, dtype=points.dtype) + # Populate intrinsic camera matrix K. + K = torch.zeros([batch_size, 3, 3], device=points.device, dtype=points.dtype) + K[:,0,0] = focal_length[:,0] + K[:,1,1] = focal_length[:,1] + K[:,2,2] = 1. + K[:,:-1, -1] = camera_center + # Transform points + points = torch.einsum('bij,bkj->bki', rotation, points) + points = points + translation.unsqueeze(1) + + # Apply perspective distortion + projected_points = points / points[:,:,-1].unsqueeze(-1) + + # Apply camera intrinsics + projected_points = torch.einsum('bij,bkj->bki', K, projected_points) + + return projected_points[:, :, :-1] + +class DeConvNet(nn.Module): + def __init__(self, feat_dim=768, upscale=4): + super(DeConvNet, self).__init__() + self.first_conv = make_conv_layers([feat_dim, feat_dim//2], kernel=1, stride=1, padding=0, bnrelu_final=False) + self.deconv = nn.ModuleList([]) + for i in range(int(math.log2(upscale))+1): + if i==0: + self.deconv.append(make_deconv_layers([feat_dim//2, feat_dim//4])) + elif i==1: + self.deconv.append(make_deconv_layers([feat_dim//2, feat_dim//4, feat_dim//8])) + elif i==2: + self.deconv.append(make_deconv_layers([feat_dim//2, feat_dim//4, feat_dim//8, feat_dim//8])) + + def forward(self, img_feat): + + face_img_feats = [] + img_feat = self.first_conv(img_feat) + face_img_feats.append(img_feat) + for i, deconv in enumerate(self.deconv): + scale = 2**i + img_feat_i = deconv(img_feat) + face_img_feat = img_feat_i + face_img_feats.append(face_img_feat) + return face_img_feats[::-1] # high resolution -> low resolution + +class DeConvNet_v2(nn.Module): + def __init__(self, feat_dim=768): + super(DeConvNet_v2, self).__init__() + self.first_conv = make_conv_layers([feat_dim, feat_dim//2], kernel=1, stride=1, padding=0, bnrelu_final=False) + self.deconv = nn.Sequential(*[nn.ConvTranspose2d(in_channels=feat_dim//2, out_channels=feat_dim//4, kernel_size=4, stride=4, padding=0, output_padding=0, bias=False), + nn.BatchNorm2d(feat_dim//4), + nn.ReLU(inplace=True)]) + + def forward(self, img_feat): + + face_img_feats = [] + img_feat = self.first_conv(img_feat) + img_feat = self.deconv(img_feat) + + return [img_feat] + +class RefineNet(nn.Module): + def __init__(self, cfg, feat_dim=1280, upscale=3): + super(RefineNet, self).__init__() + #self.deconv = DeConvNet_v2(feat_dim=feat_dim) + #self.out_dim = feat_dim//4 + + self.deconv = DeConvNet(feat_dim=feat_dim, upscale=upscale) + self.out_dim = feat_dim//8 + feat_dim//4 + feat_dim//2 + self.dec_pose = nn.Linear(self.out_dim, 96) + self.dec_cam = nn.Linear(self.out_dim, 3) + self.dec_shape = nn.Linear(self.out_dim, 10) + + self.cfg = cfg + self.joint_rep_type = cfg.MODEL.MANO_HEAD.get('JOINT_REP', '6d') + self.joint_rep_dim = {'6d': 6, 'aa': 3}[self.joint_rep_type] + + def forward(self, img_feat, verts_3d, pred_cam, pred_mano_feats, focal_length): + B = img_feat.shape[0] + + img_feats = self.deconv(img_feat) + + img_feat_sizes = [img_feat.shape[2] for img_feat in img_feats] + + temp_cams = [torch.stack([pred_cam[:, 1], pred_cam[:, 2], + 2*focal_length[:, 0]/(img_feat_size * pred_cam[:, 0] +1e-9)],dim=-1) for img_feat_size in img_feat_sizes] + + verts_2d = [perspective_projection(verts_3d, + translation=temp_cams[i], + focal_length=focal_length / img_feat_sizes[i]) for i in range(len(img_feat_sizes))] + + vert_feats = [sample_joint_features(img_feats[i], verts_2d[i]).max(1).values for i in range(len(img_feat_sizes))] + + vert_feats = torch.cat(vert_feats, dim=-1) + + delta_pose = self.dec_pose(vert_feats) + delta_betas = self.dec_shape(vert_feats) + delta_cam = self.dec_cam(vert_feats) + + + pred_hand_pose = pred_mano_feats['hand_pose'] + delta_pose + pred_betas = pred_mano_feats['betas'] + delta_betas + pred_cam = pred_mano_feats['cam'] + delta_cam + + joint_conversion_fn = { + '6d': rot6d_to_rotmat, + 'aa': lambda x: aa_to_rotmat(x.view(-1, 3).contiguous()) + }[self.joint_rep_type] + + pred_hand_pose = joint_conversion_fn(pred_hand_pose).view(B, self.cfg.MANO.NUM_HAND_JOINTS+1, 3, 3) + + pred_mano_params = {'global_orient': pred_hand_pose[:, [0]], + 'hand_pose': pred_hand_pose[:, 1:], + 'betas': pred_betas} + + return pred_mano_params, pred_cam + + \ No newline at end of file diff --git a/wilor/models/losses.py b/wilor/models/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..d6e493c081a4d99b97b5641e85152c4d56072a58 --- /dev/null +++ b/wilor/models/losses.py @@ -0,0 +1,92 @@ +import torch +import torch.nn as nn + +class Keypoint2DLoss(nn.Module): + + def __init__(self, loss_type: str = 'l1'): + """ + 2D keypoint loss module. + Args: + loss_type (str): Choose between l1 and l2 losses. + """ + super(Keypoint2DLoss, self).__init__() + if loss_type == 'l1': + self.loss_fn = nn.L1Loss(reduction='none') + elif loss_type == 'l2': + self.loss_fn = nn.MSELoss(reduction='none') + else: + raise NotImplementedError('Unsupported loss function') + + def forward(self, pred_keypoints_2d: torch.Tensor, gt_keypoints_2d: torch.Tensor) -> torch.Tensor: + """ + Compute 2D reprojection loss on the keypoints. + Args: + pred_keypoints_2d (torch.Tensor): Tensor of shape [B, S, N, 2] containing projected 2D keypoints (B: batch_size, S: num_samples, N: num_keypoints) + gt_keypoints_2d (torch.Tensor): Tensor of shape [B, S, N, 3] containing the ground truth 2D keypoints and confidence. + Returns: + torch.Tensor: 2D keypoint loss. + """ + conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone() + batch_size = conf.shape[0] + loss = (conf * self.loss_fn(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).sum(dim=(1,2)) + return loss.sum() + + +class Keypoint3DLoss(nn.Module): + + def __init__(self, loss_type: str = 'l1'): + """ + 3D keypoint loss module. + Args: + loss_type (str): Choose between l1 and l2 losses. + """ + super(Keypoint3DLoss, self).__init__() + if loss_type == 'l1': + self.loss_fn = nn.L1Loss(reduction='none') + elif loss_type == 'l2': + self.loss_fn = nn.MSELoss(reduction='none') + else: + raise NotImplementedError('Unsupported loss function') + + def forward(self, pred_keypoints_3d: torch.Tensor, gt_keypoints_3d: torch.Tensor, pelvis_id: int = 0): + """ + Compute 3D keypoint loss. + Args: + pred_keypoints_3d (torch.Tensor): Tensor of shape [B, S, N, 3] containing the predicted 3D keypoints (B: batch_size, S: num_samples, N: num_keypoints) + gt_keypoints_3d (torch.Tensor): Tensor of shape [B, S, N, 4] containing the ground truth 3D keypoints and confidence. + Returns: + torch.Tensor: 3D keypoint loss. + """ + batch_size = pred_keypoints_3d.shape[0] + gt_keypoints_3d = gt_keypoints_3d.clone() + pred_keypoints_3d = pred_keypoints_3d - pred_keypoints_3d[:, pelvis_id, :].unsqueeze(dim=1) + gt_keypoints_3d[:, :, :-1] = gt_keypoints_3d[:, :, :-1] - gt_keypoints_3d[:, pelvis_id, :-1].unsqueeze(dim=1) + conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone() + gt_keypoints_3d = gt_keypoints_3d[:, :, :-1] + loss = (conf * self.loss_fn(pred_keypoints_3d, gt_keypoints_3d)).sum(dim=(1,2)) + return loss.sum() + +class ParameterLoss(nn.Module): + + def __init__(self): + """ + MANO parameter loss module. + """ + super(ParameterLoss, self).__init__() + self.loss_fn = nn.MSELoss(reduction='none') + + def forward(self, pred_param: torch.Tensor, gt_param: torch.Tensor, has_param: torch.Tensor): + """ + Compute MANO parameter loss. + Args: + pred_param (torch.Tensor): Tensor of shape [B, S, ...] containing the predicted parameters (body pose / global orientation / betas) + gt_param (torch.Tensor): Tensor of shape [B, S, ...] containing the ground truth MANO parameters. + Returns: + torch.Tensor: L2 parameter loss loss. + """ + batch_size = pred_param.shape[0] + num_dims = len(pred_param.shape) + mask_dimension = [batch_size] + [1] * (num_dims-1) + has_param = has_param.type(pred_param.type()).view(*mask_dimension) + loss_param = (has_param * self.loss_fn(pred_param, gt_param)) + return loss_param.sum() diff --git a/wilor/models/mano_wrapper.py b/wilor/models/mano_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..f6f0cc336098e9303d2514c571307c56baf3bc86 --- /dev/null +++ b/wilor/models/mano_wrapper.py @@ -0,0 +1,40 @@ +import torch +import numpy as np +import pickle +from typing import Optional +import smplx +from smplx.lbs import vertices2joints +from smplx.utils import MANOOutput, to_tensor +from smplx.vertex_ids import vertex_ids + + +class MANO(smplx.MANOLayer): + def __init__(self, *args, joint_regressor_extra: Optional[str] = None, **kwargs): + """ + Extension of the official MANO implementation to support more joints. + Args: + Same as MANOLayer. + joint_regressor_extra (str): Path to extra joint regressor. + """ + super(MANO, self).__init__(*args, **kwargs) + mano_to_openpose = [0, 13, 14, 15, 16, 1, 2, 3, 17, 4, 5, 6, 18, 10, 11, 12, 19, 7, 8, 9, 20] + + #2, 3, 5, 4, 1 + if joint_regressor_extra is not None: + self.register_buffer('joint_regressor_extra', torch.tensor(pickle.load(open(joint_regressor_extra, 'rb'), encoding='latin1'), dtype=torch.float32)) + self.register_buffer('extra_joints_idxs', to_tensor(list(vertex_ids['mano'].values()), dtype=torch.long)) + self.register_buffer('joint_map', torch.tensor(mano_to_openpose, dtype=torch.long)) + + def forward(self, *args, **kwargs) -> MANOOutput: + """ + Run forward pass. Same as MANO and also append an extra set of joints if joint_regressor_extra is specified. + """ + mano_output = super(MANO, self).forward(*args, **kwargs) + extra_joints = torch.index_select(mano_output.vertices, 1, self.extra_joints_idxs) + joints = torch.cat([mano_output.joints, extra_joints], dim=1) + joints = joints[:, self.joint_map, :] + if hasattr(self, 'joint_regressor_extra'): + extra_joints = vertices2joints(self.joint_regressor_extra, mano_output.vertices) + joints = torch.cat([joints, extra_joints], dim=1) + mano_output.joints = joints + return mano_output diff --git a/wilor/models/wilor.py b/wilor/models/wilor.py new file mode 100644 index 0000000000000000000000000000000000000000..e5306376229a56931a444e693a6b0d070cc75bfd --- /dev/null +++ b/wilor/models/wilor.py @@ -0,0 +1,376 @@ +import torch +import pytorch_lightning as pl +from typing import Any, Dict, Mapping, Tuple + +from yacs.config import CfgNode + +from ..utils import SkeletonRenderer, MeshRenderer +from ..utils.geometry import aa_to_rotmat, perspective_projection +from ..utils.pylogger import get_pylogger +from .backbones import create_backbone +from .heads import RefineNet +from .discriminator import Discriminator +from .losses import Keypoint3DLoss, Keypoint2DLoss, ParameterLoss +from . import MANO + +log = get_pylogger(__name__) + +class WiLoR(pl.LightningModule): + + def __init__(self, cfg: CfgNode, init_renderer: bool = True): + """ + Setup WiLoR model + Args: + cfg (CfgNode): Config file as a yacs CfgNode + """ + super().__init__() + + # Save hyperparameters + self.save_hyperparameters(logger=False, ignore=['init_renderer']) + + self.cfg = cfg + # Create backbone feature extractor + self.backbone = create_backbone(cfg) + if cfg.MODEL.BACKBONE.get('PRETRAINED_WEIGHTS', None): + log.info(f'Loading backbone weights from {cfg.MODEL.BACKBONE.PRETRAINED_WEIGHTS}') + self.backbone.load_state_dict(torch.load(cfg.MODEL.BACKBONE.PRETRAINED_WEIGHTS, map_location='cpu')['state_dict'], strict = False) + + # Create RefineNet head + self.refine_net = RefineNet(cfg, feat_dim=1280, upscale=3) + + # Create discriminator + if self.cfg.LOSS_WEIGHTS.ADVERSARIAL > 0: + self.discriminator = Discriminator() + + # Define loss functions + self.keypoint_3d_loss = Keypoint3DLoss(loss_type='l1') + self.keypoint_2d_loss = Keypoint2DLoss(loss_type='l1') + self.mano_parameter_loss = ParameterLoss() + + # Instantiate MANO model + mano_cfg = {k.lower(): v for k,v in dict(cfg.MANO).items()} + self.mano = MANO(**mano_cfg) + + # Buffer that shows whetheer we need to initialize ActNorm layers + self.register_buffer('initialized', torch.tensor(False)) + # Setup renderer for visualization + if init_renderer: + self.renderer = SkeletonRenderer(self.cfg) + self.mesh_renderer = MeshRenderer(self.cfg, faces=self.mano.faces) + else: + self.renderer = None + self.mesh_renderer = None + + + # Disable automatic optimization since we use adversarial training + self.automatic_optimization = False + + def on_after_backward(self): + for name, param in self.named_parameters(): + if param.grad is None: + print(param.shape) + print(name) + + + def get_parameters(self): + #all_params = list(self.mano_head.parameters()) + all_params = list(self.backbone.parameters()) + return all_params + + def configure_optimizers(self) -> Tuple[torch.optim.Optimizer, torch.optim.Optimizer]: + """ + Setup model and distriminator Optimizers + Returns: + Tuple[torch.optim.Optimizer, torch.optim.Optimizer]: Model and discriminator optimizers + """ + param_groups = [{'params': filter(lambda p: p.requires_grad, self.get_parameters()), 'lr': self.cfg.TRAIN.LR}] + + optimizer = torch.optim.AdamW(params=param_groups, + # lr=self.cfg.TRAIN.LR, + weight_decay=self.cfg.TRAIN.WEIGHT_DECAY) + optimizer_disc = torch.optim.AdamW(params=self.discriminator.parameters(), + lr=self.cfg.TRAIN.LR, + weight_decay=self.cfg.TRAIN.WEIGHT_DECAY) + + return optimizer, optimizer_disc + + def forward_step(self, batch: Dict, train: bool = False) -> Dict: + """ + Run a forward step of the network + Args: + batch (Dict): Dictionary containing batch data + train (bool): Flag indicating whether it is training or validation mode + Returns: + Dict: Dictionary containing the regression output + """ + # Use RGB image as input + x = batch['img'] + batch_size = x.shape[0] + # Compute conditioning features using the backbone + # if using ViT backbone, we need to use a different aspect ratio + temp_mano_params, pred_cam, pred_mano_feats, vit_out = self.backbone(x[:,:,:,32:-32]) # B, 1280, 16, 12 + + + # Compute camera translation + device = temp_mano_params['hand_pose'].device + dtype = temp_mano_params['hand_pose'].dtype + focal_length = self.cfg.EXTRA.FOCAL_LENGTH * torch.ones(batch_size, 2, device=device, dtype=dtype) + + + ## Temp MANO + temp_mano_params['global_orient'] = temp_mano_params['global_orient'].reshape(batch_size, -1, 3, 3) + temp_mano_params['hand_pose'] = temp_mano_params['hand_pose'].reshape(batch_size, -1, 3, 3) + temp_mano_params['betas'] = temp_mano_params['betas'].reshape(batch_size, -1) + temp_mano_output = self.mano(**{k: v.float() for k,v in temp_mano_params.items()}, pose2rot=False) + #temp_keypoints_3d = temp_mano_output.joints + temp_vertices = temp_mano_output.vertices + + pred_mano_params, pred_cam = self.refine_net(vit_out, temp_vertices, pred_cam, pred_mano_feats, focal_length) + # Store useful regression outputs to the output dict + + + output = {} + output['pred_cam'] = pred_cam + output['pred_mano_params'] = {k: v.clone() for k,v in pred_mano_params.items()} + + pred_cam_t = torch.stack([pred_cam[:, 1], + pred_cam[:, 2], + 2*focal_length[:, 0]/(self.cfg.MODEL.IMAGE_SIZE * pred_cam[:, 0] +1e-9)],dim=-1) + output['pred_cam_t'] = pred_cam_t + output['focal_length'] = focal_length + + # Compute model vertices, joints and the projected joints + pred_mano_params['global_orient'] = pred_mano_params['global_orient'].reshape(batch_size, -1, 3, 3) + pred_mano_params['hand_pose'] = pred_mano_params['hand_pose'].reshape(batch_size, -1, 3, 3) + pred_mano_params['betas'] = pred_mano_params['betas'].reshape(batch_size, -1) + mano_output = self.mano(**{k: v.float() for k,v in pred_mano_params.items()}, pose2rot=False) + pred_keypoints_3d = mano_output.joints + pred_vertices = mano_output.vertices + + output['pred_keypoints_3d'] = pred_keypoints_3d.reshape(batch_size, -1, 3) + output['pred_vertices'] = pred_vertices.reshape(batch_size, -1, 3) + pred_cam_t = pred_cam_t.reshape(-1, 3) + focal_length = focal_length.reshape(-1, 2) + + pred_keypoints_2d = perspective_projection(pred_keypoints_3d, + translation=pred_cam_t, + focal_length=focal_length / self.cfg.MODEL.IMAGE_SIZE) + output['pred_keypoints_2d'] = pred_keypoints_2d.reshape(batch_size, -1, 2) + + return output + + def compute_loss(self, batch: Dict, output: Dict, train: bool = True) -> torch.Tensor: + """ + Compute losses given the input batch and the regression output + Args: + batch (Dict): Dictionary containing batch data + output (Dict): Dictionary containing the regression output + train (bool): Flag indicating whether it is training or validation mode + Returns: + torch.Tensor : Total loss for current batch + """ + + pred_mano_params = output['pred_mano_params'] + pred_keypoints_2d = output['pred_keypoints_2d'] + pred_keypoints_3d = output['pred_keypoints_3d'] + + + batch_size = pred_mano_params['hand_pose'].shape[0] + device = pred_mano_params['hand_pose'].device + dtype = pred_mano_params['hand_pose'].dtype + + # Get annotations + gt_keypoints_2d = batch['keypoints_2d'] + gt_keypoints_3d = batch['keypoints_3d'] + gt_mano_params = batch['mano_params'] + has_mano_params = batch['has_mano_params'] + is_axis_angle = batch['mano_params_is_axis_angle'] + + # Compute 3D keypoint loss + loss_keypoints_2d = self.keypoint_2d_loss(pred_keypoints_2d, gt_keypoints_2d) + loss_keypoints_3d = self.keypoint_3d_loss(pred_keypoints_3d, gt_keypoints_3d, pelvis_id=0) + + # Compute loss on MANO parameters + loss_mano_params = {} + for k, pred in pred_mano_params.items(): + gt = gt_mano_params[k].view(batch_size, -1) + if is_axis_angle[k].all(): + gt = aa_to_rotmat(gt.reshape(-1, 3)).view(batch_size, -1, 3, 3) + has_gt = has_mano_params[k] + loss_mano_params[k] = self.mano_parameter_loss(pred.reshape(batch_size, -1), gt.reshape(batch_size, -1), has_gt) + + loss = self.cfg.LOSS_WEIGHTS['KEYPOINTS_3D'] * loss_keypoints_3d+\ + self.cfg.LOSS_WEIGHTS['KEYPOINTS_2D'] * loss_keypoints_2d+\ + sum([loss_mano_params[k] * self.cfg.LOSS_WEIGHTS[k.upper()] for k in loss_mano_params]) + + + losses = dict(loss=loss.detach(), + loss_keypoints_2d=loss_keypoints_2d.detach(), + loss_keypoints_3d=loss_keypoints_3d.detach()) + + for k, v in loss_mano_params.items(): + losses['loss_' + k] = v.detach() + + output['losses'] = losses + + return loss + + # Tensoroboard logging should run from first rank only + @pl.utilities.rank_zero.rank_zero_only + def tensorboard_logging(self, batch: Dict, output: Dict, step_count: int, train: bool = True, write_to_summary_writer: bool = True) -> None: + """ + Log results to Tensorboard + Args: + batch (Dict): Dictionary containing batch data + output (Dict): Dictionary containing the regression output + step_count (int): Global training step count + train (bool): Flag indicating whether it is training or validation mode + """ + + mode = 'train' if train else 'val' + batch_size = batch['keypoints_2d'].shape[0] + images = batch['img'] + images = images * torch.tensor([0.229, 0.224, 0.225], device=images.device).reshape(1,3,1,1) + images = images + torch.tensor([0.485, 0.456, 0.406], device=images.device).reshape(1,3,1,1) + #images = 255*images.permute(0, 2, 3, 1).cpu().numpy() + + pred_keypoints_3d = output['pred_keypoints_3d'].detach().reshape(batch_size, -1, 3) + pred_vertices = output['pred_vertices'].detach().reshape(batch_size, -1, 3) + focal_length = output['focal_length'].detach().reshape(batch_size, 2) + gt_keypoints_3d = batch['keypoints_3d'] + gt_keypoints_2d = batch['keypoints_2d'] + + losses = output['losses'] + pred_cam_t = output['pred_cam_t'].detach().reshape(batch_size, 3) + pred_keypoints_2d = output['pred_keypoints_2d'].detach().reshape(batch_size, -1, 2) + if write_to_summary_writer: + summary_writer = self.logger.experiment + for loss_name, val in losses.items(): + summary_writer.add_scalar(mode +'/' + loss_name, val.detach().item(), step_count) + num_images = min(batch_size, self.cfg.EXTRA.NUM_LOG_IMAGES) + + gt_keypoints_3d = batch['keypoints_3d'] + pred_keypoints_3d = output['pred_keypoints_3d'].detach().reshape(batch_size, -1, 3) + + # We render the skeletons instead of the full mesh because rendering a lot of meshes will make the training slow. + #predictions = self.renderer(pred_keypoints_3d[:num_images], + # gt_keypoints_3d[:num_images], + # 2 * gt_keypoints_2d[:num_images], + # images=images[:num_images], + # camera_translation=pred_cam_t[:num_images]) + predictions = self.mesh_renderer.visualize_tensorboard(pred_vertices[:num_images].cpu().numpy(), + pred_cam_t[:num_images].cpu().numpy(), + images[:num_images].cpu().numpy(), + pred_keypoints_2d[:num_images].cpu().numpy(), + gt_keypoints_2d[:num_images].cpu().numpy(), + focal_length=focal_length[:num_images].cpu().numpy()) + if write_to_summary_writer: + summary_writer.add_image('%s/predictions' % mode, predictions, step_count) + + return predictions + + def forward(self, batch: Dict) -> Dict: + """ + Run a forward step of the network in val mode + Args: + batch (Dict): Dictionary containing batch data + Returns: + Dict: Dictionary containing the regression output + """ + return self.forward_step(batch, train=False) + + def training_step_discriminator(self, batch: Dict, + hand_pose: torch.Tensor, + betas: torch.Tensor, + optimizer: torch.optim.Optimizer) -> torch.Tensor: + """ + Run a discriminator training step + Args: + batch (Dict): Dictionary containing mocap batch data + hand_pose (torch.Tensor): Regressed hand pose from current step + betas (torch.Tensor): Regressed betas from current step + optimizer (torch.optim.Optimizer): Discriminator optimizer + Returns: + torch.Tensor: Discriminator loss + """ + batch_size = hand_pose.shape[0] + gt_hand_pose = batch['hand_pose'] + gt_betas = batch['betas'] + gt_rotmat = aa_to_rotmat(gt_hand_pose.view(-1,3)).view(batch_size, -1, 3, 3) + disc_fake_out = self.discriminator(hand_pose.detach(), betas.detach()) + loss_fake = ((disc_fake_out - 0.0) ** 2).sum() / batch_size + disc_real_out = self.discriminator(gt_rotmat, gt_betas) + loss_real = ((disc_real_out - 1.0) ** 2).sum() / batch_size + loss_disc = loss_fake + loss_real + loss = self.cfg.LOSS_WEIGHTS.ADVERSARIAL * loss_disc + optimizer.zero_grad() + self.manual_backward(loss) + optimizer.step() + return loss_disc.detach() + + def training_step(self, joint_batch: Dict, batch_idx: int) -> Dict: + """ + Run a full training step + Args: + joint_batch (Dict): Dictionary containing image and mocap batch data + batch_idx (int): Unused. + batch_idx (torch.Tensor): Unused. + Returns: + Dict: Dictionary containing regression output. + """ + batch = joint_batch['img'] + mocap_batch = joint_batch['mocap'] + optimizer = self.optimizers(use_pl_optimizer=True) + if self.cfg.LOSS_WEIGHTS.ADVERSARIAL > 0: + optimizer, optimizer_disc = optimizer + + batch_size = batch['img'].shape[0] + output = self.forward_step(batch, train=True) + pred_mano_params = output['pred_mano_params'] + if self.cfg.get('UPDATE_GT_SPIN', False): + self.update_batch_gt_spin(batch, output) + loss = self.compute_loss(batch, output, train=True) + if self.cfg.LOSS_WEIGHTS.ADVERSARIAL > 0: + disc_out = self.discriminator(pred_mano_params['hand_pose'].reshape(batch_size, -1), pred_mano_params['betas'].reshape(batch_size, -1)) + loss_adv = ((disc_out - 1.0) ** 2).sum() / batch_size + loss = loss + self.cfg.LOSS_WEIGHTS.ADVERSARIAL * loss_adv + + # Error if Nan + if torch.isnan(loss): + raise ValueError('Loss is NaN') + + optimizer.zero_grad() + self.manual_backward(loss) + # Clip gradient + if self.cfg.TRAIN.get('GRAD_CLIP_VAL', 0) > 0: + gn = torch.nn.utils.clip_grad_norm_(self.get_parameters(), self.cfg.TRAIN.GRAD_CLIP_VAL, error_if_nonfinite=True) + self.log('train/grad_norm', gn, on_step=True, on_epoch=True, prog_bar=True, logger=True) + optimizer.step() + if self.cfg.LOSS_WEIGHTS.ADVERSARIAL > 0: + loss_disc = self.training_step_discriminator(mocap_batch, pred_mano_params['hand_pose'].reshape(batch_size, -1), pred_mano_params['betas'].reshape(batch_size, -1), optimizer_disc) + output['losses']['loss_gen'] = loss_adv + output['losses']['loss_disc'] = loss_disc + + if self.global_step > 0 and self.global_step % self.cfg.GENERAL.LOG_STEPS == 0: + self.tensorboard_logging(batch, output, self.global_step, train=True) + + self.log('train/loss', output['losses']['loss'], on_step=True, on_epoch=True, prog_bar=True, logger=False) + + return output + + def validation_step(self, batch: Dict, batch_idx: int, dataloader_idx=0) -> Dict: + """ + Run a validation step and log to Tensorboard + Args: + batch (Dict): Dictionary containing batch data + batch_idx (int): Unused. + Returns: + Dict: Dictionary containing regression output. + """ + # batch_size = batch['img'].shape[0] + output = self.forward_step(batch, train=False) + loss = self.compute_loss(batch, output, train=False) + output['loss'] = loss + self.tensorboard_logging(batch, output, self.global_step, train=False) + + return output diff --git a/wilor/utils/__init__.py b/wilor/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..09e47cdf8cdb303432d64902fbe58b256273f88a --- /dev/null +++ b/wilor/utils/__init__.py @@ -0,0 +1,25 @@ +import torch +from typing import Any + +from .renderer import Renderer +from .mesh_renderer import MeshRenderer +from .skeleton_renderer import SkeletonRenderer +from .pose_utils import eval_pose, Evaluator + +def recursive_to(x: Any, target: torch.device): + """ + Recursively transfer a batch of data to the target device + Args: + x (Any): Batch of data. + target (torch.device): Target device. + Returns: + Batch of data where all tensors are transfered to the target device. + """ + if isinstance(x, dict): + return {k: recursive_to(v, target) for k, v in x.items()} + elif isinstance(x, torch.Tensor): + return x.to(target) + elif isinstance(x, list): + return [recursive_to(i, target) for i in x] + else: + return x diff --git a/wilor/utils/geometry.py b/wilor/utils/geometry.py new file mode 100644 index 0000000000000000000000000000000000000000..7929ef52608618a4682788487008e73c5736101b --- /dev/null +++ b/wilor/utils/geometry.py @@ -0,0 +1,102 @@ +from typing import Optional +import torch +from torch.nn import functional as F + +def aa_to_rotmat(theta: torch.Tensor): + """ + Convert axis-angle representation to rotation matrix. + Works by first converting it to a quaternion. + Args: + theta (torch.Tensor): Tensor of shape (B, 3) containing axis-angle representations. + Returns: + torch.Tensor: Corresponding rotation matrices with shape (B, 3, 3). + """ + norm = torch.norm(theta + 1e-8, p = 2, dim = 1) + angle = torch.unsqueeze(norm, -1) + normalized = torch.div(theta, angle) + angle = angle * 0.5 + v_cos = torch.cos(angle) + v_sin = torch.sin(angle) + quat = torch.cat([v_cos, v_sin * normalized], dim = 1) + return quat_to_rotmat(quat) + +def quat_to_rotmat(quat: torch.Tensor) -> torch.Tensor: + """ + Convert quaternion representation to rotation matrix. + Args: + quat (torch.Tensor) of shape (B, 4); 4 <===> (w, x, y, z). + Returns: + torch.Tensor: Corresponding rotation matrices with shape (B, 3, 3). + """ + norm_quat = quat + norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True) + w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3] + + B = quat.size(0) + + w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2) + wx, wy, wz = w*x, w*y, w*z + xy, xz, yz = x*y, x*z, y*z + + rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz, + 2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx, + 2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3) + return rotMat + + +def rot6d_to_rotmat(x: torch.Tensor) -> torch.Tensor: + """ + Convert 6D rotation representation to 3x3 rotation matrix. + Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019 + Args: + x (torch.Tensor): (B,6) Batch of 6-D rotation representations. + Returns: + torch.Tensor: Batch of corresponding rotation matrices with shape (B,3,3). + """ + x = x.reshape(-1,2,3).permute(0, 2, 1).contiguous() + a1 = x[:, :, 0] + a2 = x[:, :, 1] + b1 = F.normalize(a1) + b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1) + b3 = torch.cross(b1, b2) + return torch.stack((b1, b2, b3), dim=-1) + +def perspective_projection(points: torch.Tensor, + translation: torch.Tensor, + focal_length: torch.Tensor, + camera_center: Optional[torch.Tensor] = None, + rotation: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Computes the perspective projection of a set of 3D points. + Args: + points (torch.Tensor): Tensor of shape (B, N, 3) containing the input 3D points. + translation (torch.Tensor): Tensor of shape (B, 3) containing the 3D camera translation. + focal_length (torch.Tensor): Tensor of shape (B, 2) containing the focal length in pixels. + camera_center (torch.Tensor): Tensor of shape (B, 2) containing the camera center in pixels. + rotation (torch.Tensor): Tensor of shape (B, 3, 3) containing the camera rotation. + Returns: + torch.Tensor: Tensor of shape (B, N, 2) containing the projection of the input points. + """ + batch_size = points.shape[0] + if rotation is None: + rotation = torch.eye(3, device=points.device, dtype=points.dtype).unsqueeze(0).expand(batch_size, -1, -1) + if camera_center is None: + camera_center = torch.zeros(batch_size, 2, device=points.device, dtype=points.dtype) + # Populate intrinsic camera matrix K. + K = torch.zeros([batch_size, 3, 3], device=points.device, dtype=points.dtype) + K[:,0,0] = focal_length[:,0] + K[:,1,1] = focal_length[:,1] + K[:,2,2] = 1. + K[:,:-1, -1] = camera_center + + # Transform points + points = torch.einsum('bij,bkj->bki', rotation, points) + points = points + translation.unsqueeze(1) + + # Apply perspective distortion + projected_points = points / points[:,:,-1].unsqueeze(-1) + + # Apply camera intrinsics + projected_points = torch.einsum('bij,bkj->bki', K, projected_points) + + return projected_points[:, :, :-1] \ No newline at end of file diff --git a/wilor/utils/mesh_renderer.py b/wilor/utils/mesh_renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..bb3e8ed2e9aed8157ec852d06d5f13e8f4ff7c54 --- /dev/null +++ b/wilor/utils/mesh_renderer.py @@ -0,0 +1,149 @@ +import os +if 'PYOPENGL_PLATFORM' not in os.environ: + os.environ['PYOPENGL_PLATFORM'] = 'egl' +import torch +from torchvision.utils import make_grid +import numpy as np +import pyrender +import trimesh +import cv2 +import torch.nn.functional as F + +from .render_openpose import render_openpose + +def create_raymond_lights(): + import pyrender + thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0]) + phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0]) + + nodes = [] + + for phi, theta in zip(phis, thetas): + xp = np.sin(theta) * np.cos(phi) + yp = np.sin(theta) * np.sin(phi) + zp = np.cos(theta) + + z = np.array([xp, yp, zp]) + z = z / np.linalg.norm(z) + x = np.array([-z[1], z[0], 0.0]) + if np.linalg.norm(x) == 0: + x = np.array([1.0, 0.0, 0.0]) + x = x / np.linalg.norm(x) + y = np.cross(z, x) + + matrix = np.eye(4) + matrix[:3,:3] = np.c_[x,y,z] + nodes.append(pyrender.Node( + light=pyrender.DirectionalLight(color=np.ones(3), intensity=1.0), + matrix=matrix + )) + + return nodes + +class MeshRenderer: + + def __init__(self, cfg, faces=None): + self.cfg = cfg + self.focal_length = cfg.EXTRA.FOCAL_LENGTH + self.img_res = cfg.MODEL.IMAGE_SIZE + self.renderer = pyrender.OffscreenRenderer(viewport_width=self.img_res, + viewport_height=self.img_res, + point_size=1.0) + + self.camera_center = [self.img_res // 2, self.img_res // 2] + self.faces = faces + + def visualize(self, vertices, camera_translation, images, focal_length=None, nrow=3, padding=2): + images_np = np.transpose(images, (0,2,3,1)) + rend_imgs = [] + for i in range(vertices.shape[0]): + fl = self.focal_length + rend_img = torch.from_numpy(np.transpose(self.__call__(vertices[i], camera_translation[i], images_np[i], focal_length=fl, side_view=False), (2,0,1))).float() + rend_img_side = torch.from_numpy(np.transpose(self.__call__(vertices[i], camera_translation[i], images_np[i], focal_length=fl, side_view=True), (2,0,1))).float() + rend_imgs.append(torch.from_numpy(images[i])) + rend_imgs.append(rend_img) + rend_imgs.append(rend_img_side) + rend_imgs = make_grid(rend_imgs, nrow=nrow, padding=padding) + return rend_imgs + + def visualize_tensorboard(self, vertices, camera_translation, images, pred_keypoints, gt_keypoints, focal_length=None, nrow=5, padding=2): + images_np = np.transpose(images, (0,2,3,1)) + rend_imgs = [] + pred_keypoints = np.concatenate((pred_keypoints, np.ones_like(pred_keypoints)[:, :, [0]]), axis=-1) + pred_keypoints = self.img_res * (pred_keypoints + 0.5) + gt_keypoints[:, :, :-1] = self.img_res * (gt_keypoints[:, :, :-1] + 0.5) + #keypoint_matches = [(1, 12), (2, 8), (3, 7), (4, 6), (5, 9), (6, 10), (7, 11), (8, 14), (9, 2), (10, 1), (11, 0), (12, 3), (13, 4), (14, 5)] + for i in range(vertices.shape[0]): + fl = self.focal_length + rend_img = torch.from_numpy(np.transpose(self.__call__(vertices[i], camera_translation[i], images_np[i], focal_length=fl, side_view=False), (2,0,1))).float() + rend_img_side = torch.from_numpy(np.transpose(self.__call__(vertices[i], camera_translation[i], images_np[i], focal_length=fl, side_view=True), (2,0,1))).float() + hand_keypoints = pred_keypoints[i, :21] + #extra_keypoints = pred_keypoints[i, -19:] + #for pair in keypoint_matches: + # hand_keypoints[pair[0], :] = extra_keypoints[pair[1], :] + pred_keypoints_img = render_openpose(255 * images_np[i].copy(), hand_keypoints) / 255 + hand_keypoints = gt_keypoints[i, :21] + #extra_keypoints = gt_keypoints[i, -19:] + #for pair in keypoint_matches: + # if extra_keypoints[pair[1], -1] > 0 and hand_keypoints[pair[0], -1] == 0: + # hand_keypoints[pair[0], :] = extra_keypoints[pair[1], :] + gt_keypoints_img = render_openpose(255*images_np[i].copy(), hand_keypoints) / 255 + rend_imgs.append(torch.from_numpy(images[i])) + rend_imgs.append(rend_img) + rend_imgs.append(rend_img_side) + rend_imgs.append(torch.from_numpy(pred_keypoints_img).permute(2,0,1)) + rend_imgs.append(torch.from_numpy(gt_keypoints_img).permute(2,0,1)) + rend_imgs = make_grid(rend_imgs, nrow=nrow, padding=padding) + return rend_imgs + + def __call__(self, vertices, camera_translation, image, focal_length=5000, text=None, resize=None, side_view=False, baseColorFactor=(1.0, 1.0, 0.9, 1.0), rot_angle=90): + renderer = pyrender.OffscreenRenderer(viewport_width=image.shape[1], + viewport_height=image.shape[0], + point_size=1.0) + material = pyrender.MetallicRoughnessMaterial( + metallicFactor=0.0, + alphaMode='OPAQUE', + baseColorFactor=baseColorFactor) + + camera_translation[0] *= -1. + + mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy()) + if side_view: + rot = trimesh.transformations.rotation_matrix( + np.radians(rot_angle), [0, 1, 0]) + mesh.apply_transform(rot) + rot = trimesh.transformations.rotation_matrix( + np.radians(180), [1, 0, 0]) + mesh.apply_transform(rot) + mesh = pyrender.Mesh.from_trimesh(mesh, material=material) + + scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0], + ambient_light=(0.3, 0.3, 0.3)) + scene.add(mesh, 'mesh') + + camera_pose = np.eye(4) + camera_pose[:3, 3] = camera_translation + camera_center = [image.shape[1] / 2., image.shape[0] / 2.] + camera = pyrender.IntrinsicsCamera(fx=focal_length, fy=focal_length, + cx=camera_center[0], cy=camera_center[1]) + scene.add(camera, pose=camera_pose) + + + light_nodes = create_raymond_lights() + for node in light_nodes: + scene.add_node(node) + + color, rend_depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA) + color = color.astype(np.float32) / 255.0 + valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis] + if not side_view: + output_img = (color[:, :, :3] * valid_mask + + (1 - valid_mask) * image) + else: + output_img = color[:, :, :3] + if resize is not None: + output_img = cv2.resize(output_img, resize) + + output_img = output_img.astype(np.float32) + renderer.delete() + return output_img diff --git a/wilor/utils/misc.py b/wilor/utils/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..ffcfe784872b305c264ce6ef67fd0a9e9ad3390f --- /dev/null +++ b/wilor/utils/misc.py @@ -0,0 +1,203 @@ +import time +import warnings +from importlib.util import find_spec +from pathlib import Path +from typing import Callable, List + +import hydra +from omegaconf import DictConfig, OmegaConf +from pytorch_lightning import Callback +from pytorch_lightning.loggers import Logger +from pytorch_lightning.utilities import rank_zero_only + +from . import pylogger, rich_utils + +log = pylogger.get_pylogger(__name__) + + +def task_wrapper(task_func: Callable) -> Callable: + """Optional decorator that wraps the task function in extra utilities. + + Makes multirun more resistant to failure. + + Utilities: + - Calling the `utils.extras()` before the task is started + - Calling the `utils.close_loggers()` after the task is finished + - Logging the exception if occurs + - Logging the task total execution time + - Logging the output dir + """ + + def wrap(cfg: DictConfig): + + # apply extra utilities + extras(cfg) + + # execute the task + try: + start_time = time.time() + ret = task_func(cfg=cfg) + except Exception as ex: + log.exception("") # save exception to `.log` file + raise ex + finally: + path = Path(cfg.paths.output_dir, "exec_time.log") + content = f"'{cfg.task_name}' execution time: {time.time() - start_time} (s)" + save_file(path, content) # save task execution time (even if exception occurs) + close_loggers() # close loggers (even if exception occurs so multirun won't fail) + + log.info(f"Output dir: {cfg.paths.output_dir}") + + return ret + + return wrap + + +def extras(cfg: DictConfig) -> None: + """Applies optional utilities before the task is started. + + Utilities: + - Ignoring python warnings + - Setting tags from command line + - Rich config printing + """ + + # return if no `extras` config + if not cfg.get("extras"): + log.warning("Extras config not found! ") + return + + # disable python warnings + if cfg.extras.get("ignore_warnings"): + log.info("Disabling python warnings! ") + warnings.filterwarnings("ignore") + + # prompt user to input tags from command line if none are provided in the config + if cfg.extras.get("enforce_tags"): + log.info("Enforcing tags! ") + rich_utils.enforce_tags(cfg, save_to_file=True) + + # pretty print config tree using Rich library + if cfg.extras.get("print_config"): + log.info("Printing config tree with Rich! ") + rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True) + + +@rank_zero_only +def save_file(path: str, content: str) -> None: + """Save file in rank zero mode (only on one process in multi-GPU setup).""" + with open(path, "w+") as file: + file.write(content) + + +def instantiate_callbacks(callbacks_cfg: DictConfig) -> List[Callback]: + """Instantiates callbacks from config.""" + callbacks: List[Callback] = [] + + if not callbacks_cfg: + log.warning("Callbacks config is empty.") + return callbacks + + if not isinstance(callbacks_cfg, DictConfig): + raise TypeError("Callbacks config must be a DictConfig!") + + for _, cb_conf in callbacks_cfg.items(): + if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf: + log.info(f"Instantiating callback <{cb_conf._target_}>") + callbacks.append(hydra.utils.instantiate(cb_conf)) + + return callbacks + + +def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]: + """Instantiates loggers from config.""" + logger: List[Logger] = [] + + if not logger_cfg: + log.warning("Logger config is empty.") + return logger + + if not isinstance(logger_cfg, DictConfig): + raise TypeError("Logger config must be a DictConfig!") + + for _, lg_conf in logger_cfg.items(): + if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf: + log.info(f"Instantiating logger <{lg_conf._target_}>") + logger.append(hydra.utils.instantiate(lg_conf)) + + return logger + + +@rank_zero_only +def log_hyperparameters(object_dict: dict) -> None: + """Controls which config parts are saved by lightning loggers. + + Additionally saves: + - Number of model parameters + """ + + hparams = {} + + cfg = object_dict["cfg"] + model = object_dict["model"] + trainer = object_dict["trainer"] + + if not trainer.logger: + log.warning("Logger not found! Skipping hyperparameter logging...") + return + + # save number of model parameters + hparams["model/params/total"] = sum(p.numel() for p in model.parameters()) + hparams["model/params/trainable"] = sum( + p.numel() for p in model.parameters() if p.requires_grad + ) + hparams["model/params/non_trainable"] = sum( + p.numel() for p in model.parameters() if not p.requires_grad + ) + + for k in cfg.keys(): + hparams[k] = cfg.get(k) + + # Resolve all interpolations + def _resolve(_cfg): + if isinstance(_cfg, DictConfig): + _cfg = OmegaConf.to_container(_cfg, resolve=True) + return _cfg + + hparams = {k: _resolve(v) for k, v in hparams.items()} + + # send hparams to all loggers + trainer.logger.log_hyperparams(hparams) + + +def get_metric_value(metric_dict: dict, metric_name: str) -> float: + """Safely retrieves value of the metric logged in LightningModule.""" + + if not metric_name: + log.info("Metric name is None! Skipping metric value retrieval...") + return None + + if metric_name not in metric_dict: + raise Exception( + f"Metric value not found! \n" + "Make sure metric name logged in LightningModule is correct!\n" + "Make sure `optimized_metric` name in `hparams_search` config is correct!" + ) + + metric_value = metric_dict[metric_name].item() + log.info(f"Retrieved metric value! <{metric_name}={metric_value}>") + + return metric_value + + +def close_loggers() -> None: + """Makes sure all loggers closed properly (prevents logging failure during multirun).""" + + log.info("Closing loggers...") + + if find_spec("wandb"): # if wandb is installed + import wandb + + if wandb.run: + log.info("Closing wandb!") + wandb.finish() diff --git a/wilor/utils/pose_utils.py b/wilor/utils/pose_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1b386e39aca64ce73ae10fea9ba2b767ce9e25b2 --- /dev/null +++ b/wilor/utils/pose_utils.py @@ -0,0 +1,352 @@ +""" +Code adapted from: https://github.com/akanazawa/hmr/blob/master/src/benchmark/eval_util.py +""" + +import torch +import numpy as np +from typing import Optional, Dict, List, Tuple + +def compute_similarity_transform(S1: torch.Tensor, S2: torch.Tensor) -> torch.Tensor: + """ + Computes a similarity transform (sR, t) in a batched way that takes + a set of 3D points S1 (B, N, 3) closest to a set of 3D points S2 (B, N, 3), + where R is a 3x3 rotation matrix, t 3x1 translation, s scale. + i.e. solves the orthogonal Procrutes problem. + Args: + S1 (torch.Tensor): First set of points of shape (B, N, 3). + S2 (torch.Tensor): Second set of points of shape (B, N, 3). + Returns: + (torch.Tensor): The first set of points after applying the similarity transformation. + """ + + batch_size = S1.shape[0] + S1 = S1.permute(0, 2, 1) + S2 = S2.permute(0, 2, 1) + # 1. Remove mean. + mu1 = S1.mean(dim=2, keepdim=True) + mu2 = S2.mean(dim=2, keepdim=True) + X1 = S1 - mu1 + X2 = S2 - mu2 + + # 2. Compute variance of X1 used for scale. + var1 = (X1**2).sum(dim=(1,2)) + + # 3. The outer product of X1 and X2. + K = torch.matmul(X1, X2.permute(0, 2, 1)) + + # 4. Solution that Maximizes trace(R'K) is R=U*V', where U, V are singular vectors of K. + U, s, V = torch.svd(K) + Vh = V.permute(0, 2, 1) + + # Construct Z that fixes the orientation of R to get det(R)=1. + Z = torch.eye(U.shape[1], device=U.device).unsqueeze(0).repeat(batch_size, 1, 1) + Z[:, -1, -1] *= torch.sign(torch.linalg.det(torch.matmul(U, Vh))) + + # Construct R. + R = torch.matmul(torch.matmul(V, Z), U.permute(0, 2, 1)) + + # 5. Recover scale. + trace = torch.matmul(R, K).diagonal(offset=0, dim1=-1, dim2=-2).sum(dim=-1) + scale = (trace / var1).unsqueeze(dim=-1).unsqueeze(dim=-1) + + # 6. Recover translation. + t = mu2 - scale*torch.matmul(R, mu1) + + # 7. Error: + S1_hat = scale*torch.matmul(R, S1) + t + + return S1_hat.permute(0, 2, 1) + +def reconstruction_error(S1, S2) -> np.array: + """ + Computes the mean Euclidean distance of 2 set of points S1, S2 after performing Procrustes alignment. + Args: + S1 (torch.Tensor): First set of points of shape (B, N, 3). + S2 (torch.Tensor): Second set of points of shape (B, N, 3). + Returns: + (np.array): Reconstruction error. + """ + S1_hat = compute_similarity_transform(S1, S2) + re = torch.sqrt( ((S1_hat - S2)** 2).sum(dim=-1)).mean(dim=-1) + return re + +def eval_pose(pred_joints, gt_joints) -> Tuple[np.array, np.array]: + """ + Compute joint errors in mm before and after Procrustes alignment. + Args: + pred_joints (torch.Tensor): Predicted 3D joints of shape (B, N, 3). + gt_joints (torch.Tensor): Ground truth 3D joints of shape (B, N, 3). + Returns: + Tuple[np.array, np.array]: Joint errors in mm before and after alignment. + """ + # Absolute error (MPJPE) + mpjpe = torch.sqrt(((pred_joints - gt_joints) ** 2).sum(dim=-1)).mean(dim=-1).cpu().numpy() + + # Reconstruction_error + r_error = reconstruction_error(pred_joints, gt_joints).cpu().numpy() + return 1000 * mpjpe, 1000 * r_error + +class Evaluator: + + def __init__(self, + dataset_length: int, + dataset: str, + keypoint_list: List, + pelvis_ind: int, + metrics: List = ['mode_mpjpe', 'mode_re', 'min_mpjpe', 'min_re'], + preds: List = ['vertices', 'keypoints_3d'], + pck_thresholds: Optional[List] = None): + """ + Class used for evaluating trained models on different 3D pose datasets. + Args: + dataset_length (int): Total dataset length. + keypoint_list [List]: List of keypoints used for evaluation. + pelvis_ind (int): Index of pelvis keypoint; used for aligning the predictions and ground truth. + metrics [List]: List of evaluation metrics to record. + """ + self.dataset_length = dataset_length + self.dataset = dataset + self.keypoint_list = keypoint_list + self.pelvis_ind = pelvis_ind + self.metrics = metrics + self.preds = preds + if self.metrics is not None: + for metric in self.metrics: + setattr(self, metric, np.zeros((dataset_length,))) + if self.preds is not None: + for pred in self.preds: + if pred == 'vertices': + self.vertices = np.zeros((dataset_length, 778, 3)) + if pred == 'keypoints_3d': + self.keypoints_3d = np.zeros((dataset_length, 21, 3)) + self.counter = 0 + if pck_thresholds is None: + self.pck_evaluator = None + else: + self.pck_evaluator = EvaluatorPCK(pck_thresholds) + + def log(self): + """ + Print current evaluation metrics + """ + if self.counter == 0: + print('Evaluation has not started') + return + print(f'{self.counter} / {self.dataset_length} samples') + if self.pck_evaluator is not None: + self.pck_evaluator.log() + if self.metrics is not None: + for metric in self.metrics: + if metric in ['mode_mpjpe', 'mode_re', 'min_mpjpe', 'min_re']: + unit = 'mm' + else: + unit = '' + print(f'{metric}: {getattr(self, metric)[:self.counter].mean()} {unit}') + print('***') + + def get_metrics_dict(self) -> Dict: + """ + Returns: + Dict: Dictionary of evaluation metrics. + """ + d1 = {metric: getattr(self, metric)[:self.counter].mean() for metric in self.metrics} + if self.pck_evaluator is not None: + d2 = self.pck_evaluator.get_metrics_dict() + d1.update(d2) + return d1 + + def get_preds_dict(self) -> Dict: + """ + Returns: + Dict: Dictionary of evaluation preds. + """ + d1 = {pred: getattr(self, pred)[:self.counter] for pred in self.preds} + return d1 + + def __call__(self, output: Dict, batch: Dict, opt_output: Optional[Dict] = None): + """ + Evaluate current batch. + Args: + output (Dict): Regression output. + batch (Dict): Dictionary containing images and their corresponding annotations. + opt_output (Dict): Optimization output. + """ + if self.pck_evaluator is not None: + self.pck_evaluator(output, batch, opt_output) + + pred_keypoints_3d = output['pred_keypoints_3d'].detach() + pred_keypoints_3d = pred_keypoints_3d[:,None,:,:] + batch_size = pred_keypoints_3d.shape[0] + num_samples = pred_keypoints_3d.shape[1] + gt_keypoints_3d = batch['keypoints_3d'][:, :, :-1].unsqueeze(1).repeat(1, num_samples, 1, 1) + pred_vertices = output['pred_vertices'].detach() + + # Align predictions and ground truth such that the pelvis location is at the origin + pred_keypoints_3d -= pred_keypoints_3d[:, :, [self.pelvis_ind]] + gt_keypoints_3d -= gt_keypoints_3d[:, :, [self.pelvis_ind]] + + # Compute joint errors + mpjpe, re = eval_pose(pred_keypoints_3d.reshape(batch_size * num_samples, -1, 3)[:, self.keypoint_list], gt_keypoints_3d.reshape(batch_size * num_samples, -1 ,3)[:, self.keypoint_list]) + mpjpe = mpjpe.reshape(batch_size, num_samples) + re = re.reshape(batch_size, num_samples) + + # Compute 2d keypoint errors + bbox_expand_factor = batch['bbox_expand_factor'][:,None,None,None].detach() + pred_keypoints_2d = output['pred_keypoints_2d'].detach() + pred_keypoints_2d = pred_keypoints_2d[:,None,:,:]*bbox_expand_factor + gt_keypoints_2d = batch['keypoints_2d'][:,None,:,:].repeat(1, num_samples, 1, 1)*bbox_expand_factor + conf = gt_keypoints_2d[:, :, :, -1].clone() + kp_err = torch.nn.functional.mse_loss( + pred_keypoints_2d, + gt_keypoints_2d[:, :, :, :-1], + reduction='none' + ).sum(dim=3) + kp_l2_loss = (conf * kp_err).mean(dim=2) + kp_l2_loss = kp_l2_loss.detach().cpu().numpy() + + # Compute joint errors after optimization, if available. + if opt_output is not None: + opt_keypoints_3d = opt_output['model_joints'] + opt_keypoints_3d -= opt_keypoints_3d[:, [self.pelvis_ind]] + opt_mpjpe, opt_re = eval_pose(opt_keypoints_3d[:, self.keypoint_list], gt_keypoints_3d[:, 0, self.keypoint_list]) + + # The 0-th sample always corresponds to the mode + if hasattr(self, 'mode_mpjpe'): + mode_mpjpe = mpjpe[:, 0] + self.mode_mpjpe[self.counter:self.counter+batch_size] = mode_mpjpe + if hasattr(self, 'mode_re'): + mode_re = re[:, 0] + self.mode_re[self.counter:self.counter+batch_size] = mode_re + if hasattr(self, 'mode_kpl2'): + mode_kpl2 = kp_l2_loss[:, 0] + self.mode_kpl2[self.counter:self.counter+batch_size] = mode_kpl2 + if hasattr(self, 'min_mpjpe'): + min_mpjpe = mpjpe.min(axis=-1) + self.min_mpjpe[self.counter:self.counter+batch_size] = min_mpjpe + if hasattr(self, 'min_re'): + min_re = re.min(axis=-1) + self.min_re[self.counter:self.counter+batch_size] = min_re + if hasattr(self, 'min_kpl2'): + min_kpl2 = kp_l2_loss.min(axis=-1) + self.min_kpl2[self.counter:self.counter+batch_size] = min_kpl2 + if hasattr(self, 'opt_mpjpe'): + self.opt_mpjpe[self.counter:self.counter+batch_size] = opt_mpjpe + if hasattr(self, 'opt_re'): + self.opt_re[self.counter:self.counter+batch_size] = opt_re + if hasattr(self, 'vertices'): + self.vertices[self.counter:self.counter+batch_size] = pred_vertices.cpu().numpy() + if hasattr(self, 'keypoints_3d'): + if self.dataset == 'HO3D-VAL': + pred_keypoints_3d = pred_keypoints_3d[:,:,[0,5,6,7,9,10,11,17,18,19,13,14,15,1,2,3,4,8,12,16,20]] + self.keypoints_3d[self.counter:self.counter+batch_size] = pred_keypoints_3d.squeeze().cpu().numpy() + + self.counter += batch_size + + if hasattr(self, 'mode_mpjpe') and hasattr(self, 'mode_re'): + return { + 'mode_mpjpe': mode_mpjpe, + 'mode_re': mode_re, + } + else: + return {} + + +class EvaluatorPCK: + + def __init__(self, thresholds: List = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5],): + """ + Class used for evaluating trained models on different 3D pose datasets. + Args: + thresholds [List]: List of PCK thresholds to evaluate. + metrics [List]: List of evaluation metrics to record. + """ + self.thresholds = thresholds + self.pred_kp_2d = [] + self.gt_kp_2d = [] + self.gt_conf_2d = [] + self.scale = [] + self.counter = 0 + + def log(self): + """ + Print current evaluation metrics + """ + if self.counter == 0: + print('Evaluation has not started') + return + print(f'{self.counter} samples') + metrics_dict = self.get_metrics_dict() + for metric in metrics_dict: + print(f'{metric}: {metrics_dict[metric]}') + print('***') + + def get_metrics_dict(self) -> Dict: + """ + Returns: + Dict: Dictionary of evaluation metrics. + """ + pcks = self.compute_pcks() + metrics = {} + for thr, (acc,avg_acc,cnt) in zip(self.thresholds, pcks): + metrics.update({f'kp{i}_pck_{thr}': float(a) for i, a in enumerate(acc) if a>=0}) + metrics.update({f'kpAvg_pck_{thr}': float(avg_acc)}) + return metrics + + def compute_pcks(self): + pred_kp_2d = np.concatenate(self.pred_kp_2d, axis=0) + gt_kp_2d = np.concatenate(self.gt_kp_2d, axis=0) + gt_conf_2d = np.concatenate(self.gt_conf_2d, axis=0) + scale = np.concatenate(self.scale, axis=0) + assert pred_kp_2d.shape == gt_kp_2d.shape + assert pred_kp_2d[..., 0].shape == gt_conf_2d.shape + assert pred_kp_2d.shape[1] == 1 # num_samples + assert scale.shape[0] == gt_conf_2d.shape[0] # num_samples + + pcks = [ + self.keypoint_pck_accuracy( + pred_kp_2d[:, 0, :, :], + gt_kp_2d[:, 0, :, :], + gt_conf_2d[:, 0, :]>0.5, + thr=thr, + scale = scale[:,None] + ) + for thr in self.thresholds + ] + return pcks + + def keypoint_pck_accuracy(self, pred, gt, conf, thr, scale): + dist = np.sqrt(np.sum((pred-gt)**2, axis=2)) + all_joints = conf>0.5 + correct_joints = np.logical_and(dist<=scale*thr, all_joints) + pck = correct_joints.sum(axis=0)/all_joints.sum(axis=0) + return pck, pck.mean(), pck.shape[0] + + def __call__(self, output: Dict, batch: Dict, opt_output: Optional[Dict] = None): + """ + Evaluate current batch. + Args: + output (Dict): Regression output. + batch (Dict): Dictionary containing images and their corresponding annotations. + opt_output (Dict): Optimization output. + """ + pred_keypoints_2d = output['pred_keypoints_2d'].detach() + num_samples = 1 + batch_size = pred_keypoints_2d.shape[0] + + right = batch['right'].detach() + pred_keypoints_2d[:,:,0] = (2*right[:,None]-1)*pred_keypoints_2d[:,:,0] + box_size = batch['box_size'].detach() + box_center = batch['box_center'].detach() + bbox_expand_factor = batch['bbox_expand_factor'].detach() + scale = box_size/bbox_expand_factor + bbox_expand_factor = bbox_expand_factor[:,None,None,None] + pred_keypoints_2d = pred_keypoints_2d*box_size[:,None,None]+box_center[:,None] + pred_keypoints_2d = pred_keypoints_2d[:,None,:,:] + gt_keypoints_2d = batch['orig_keypoints_2d'][:,None,:,:].repeat(1, num_samples, 1, 1) + + self.pred_kp_2d.append(pred_keypoints_2d[:, :, :, :2].detach().cpu().numpy()) + self.gt_conf_2d.append(gt_keypoints_2d[:, :, :, -1].detach().cpu().numpy()) + self.gt_kp_2d.append(gt_keypoints_2d[:, :, :, :2].detach().cpu().numpy()) + self.scale.append(scale.detach().cpu().numpy()) + + self.counter += batch_size \ No newline at end of file diff --git a/wilor/utils/pylogger.py b/wilor/utils/pylogger.py new file mode 100644 index 0000000000000000000000000000000000000000..92ffa71893ec20acde65e44d899334a38d8d1333 --- /dev/null +++ b/wilor/utils/pylogger.py @@ -0,0 +1,17 @@ +import logging + +from pytorch_lightning.utilities import rank_zero_only + + +def get_pylogger(name=__name__) -> logging.Logger: + """Initializes multi-GPU-friendly python command line logger.""" + + logger = logging.getLogger(name) + + # this ensures all logging levels get marked with the rank zero decorator + # otherwise logs would get multiplied for each GPU process in multi-GPU setup + logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical") + for level in logging_levels: + setattr(logger, level, rank_zero_only(getattr(logger, level))) + + return logger diff --git a/wilor/utils/render_openpose.py b/wilor/utils/render_openpose.py new file mode 100644 index 0000000000000000000000000000000000000000..8e51ee8e15f40b85e2766e1f9da42183da0d3d46 --- /dev/null +++ b/wilor/utils/render_openpose.py @@ -0,0 +1,191 @@ +""" +Render OpenPose keypoints. +Code was ported to Python from the official C++ implementation https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/utilities/keypoint.cpp +""" +import cv2 +import math +import numpy as np +from typing import List, Tuple + +def get_keypoints_rectangle(keypoints: np.array, threshold: float) -> Tuple[float, float, float]: + """ + Compute rectangle enclosing keypoints above the threshold. + Args: + keypoints (np.array): Keypoint array of shape (N, 3). + threshold (float): Confidence visualization threshold. + Returns: + Tuple[float, float, float]: Rectangle width, height and area. + """ + valid_ind = keypoints[:, -1] > threshold + if valid_ind.sum() > 0: + valid_keypoints = keypoints[valid_ind][:, :-1] + max_x = valid_keypoints[:,0].max() + max_y = valid_keypoints[:,1].max() + min_x = valid_keypoints[:,0].min() + min_y = valid_keypoints[:,1].min() + width = max_x - min_x + height = max_y - min_y + area = width * height + return width, height, area + else: + return 0,0,0 + +def render_keypoints(img: np.array, + keypoints: np.array, + pairs: List, + colors: List, + thickness_circle_ratio: float, + thickness_line_ratio_wrt_circle: float, + pose_scales: List, + threshold: float = 0.1, + alpha: float = 1.0) -> np.array: + """ + Render keypoints on input image. + Args: + img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range. + keypoints (np.array): Keypoint array of shape (N, 3). + pairs (List): List of keypoint pairs per limb. + colors: (List): List of colors per keypoint. + thickness_circle_ratio (float): Circle thickness ratio. + thickness_line_ratio_wrt_circle (float): Line thickness ratio wrt the circle. + pose_scales (List): List of pose scales. + threshold (float): Only visualize keypoints with confidence above the threshold. + Returns: + (np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image. + """ + img_orig = img.copy() + width, height = img.shape[1], img.shape[2] + area = width * height + + lineType = 8 + shift = 0 + numberColors = len(colors) + thresholdRectangle = 0.1 + + person_width, person_height, person_area = get_keypoints_rectangle(keypoints, thresholdRectangle) + if person_area > 0: + ratioAreas = min(1, max(person_width / width, person_height / height)) + thicknessRatio = np.maximum(np.round(math.sqrt(area) * thickness_circle_ratio * ratioAreas), 2) + thicknessCircle = np.maximum(1, thicknessRatio if ratioAreas > 0.05 else -np.ones_like(thicknessRatio)) + thicknessLine = np.maximum(1, np.round(thicknessRatio * thickness_line_ratio_wrt_circle)) + radius = thicknessRatio / 2 + + img = np.ascontiguousarray(img.copy()) + for i, pair in enumerate(pairs): + index1, index2 = pair + if keypoints[index1, -1] > threshold and keypoints[index2, -1] > threshold: + thicknessLineScaled = int(round(min(thicknessLine[index1], thicknessLine[index2]) * pose_scales[0])) + colorIndex = index2 + color = colors[colorIndex % numberColors] + keypoint1 = keypoints[index1, :-1].astype(np.int_) + keypoint2 = keypoints[index2, :-1].astype(np.int_) + cv2.line(img, tuple(keypoint1.tolist()), tuple(keypoint2.tolist()), tuple(color.tolist()), thicknessLineScaled, lineType, shift) + for part in range(len(keypoints)): + faceIndex = part + if keypoints[faceIndex, -1] > threshold: + radiusScaled = int(round(radius[faceIndex] * pose_scales[0])) + thicknessCircleScaled = int(round(thicknessCircle[faceIndex] * pose_scales[0])) + colorIndex = part + color = colors[colorIndex % numberColors] + center = keypoints[faceIndex, :-1].astype(np.int_) + cv2.circle(img, tuple(center.tolist()), radiusScaled, tuple(color.tolist()), thicknessCircleScaled, lineType, shift) + return img + +def render_hand_keypoints(img, right_hand_keypoints, threshold=0.1, use_confidence=False, map_fn=lambda x: np.ones_like(x), alpha=1.0): + if use_confidence and map_fn is not None: + #thicknessCircleRatioLeft = 1./50 * map_fn(left_hand_keypoints[:, -1]) + thicknessCircleRatioRight = 1./50 * map_fn(right_hand_keypoints[:, -1]) + else: + #thicknessCircleRatioLeft = 1./50 * np.ones(left_hand_keypoints.shape[0]) + thicknessCircleRatioRight = 1./50 * np.ones(right_hand_keypoints.shape[0]) + thicknessLineRatioWRTCircle = 0.75 + pairs = [0,1, 1,2, 2,3, 3,4, 0,5, 5,6, 6,7, 7,8, 0,9, 9,10, 10,11, 11,12, 0,13, 13,14, 14,15, 15,16, 0,17, 17,18, 18,19, 19,20] + pairs = np.array(pairs).reshape(-1,2) + + colors = [100., 100., 100., + 100., 0., 0., + 150., 0., 0., + 200., 0., 0., + 255., 0., 0., + 100., 100., 0., + 150., 150., 0., + 200., 200., 0., + 255., 255., 0., + 0., 100., 50., + 0., 150., 75., + 0., 200., 100., + 0., 255., 125., + 0., 50., 100., + 0., 75., 150., + 0., 100., 200., + 0., 125., 255., + 100., 0., 100., + 150., 0., 150., + 200., 0., 200., + 255., 0., 255.] + colors = np.array(colors).reshape(-1,3) + #colors = np.zeros_like(colors) + poseScales = [1] + #img = render_keypoints(img, left_hand_keypoints, pairs, colors, thicknessCircleRatioLeft, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha) + img = render_keypoints(img, right_hand_keypoints, pairs, colors, thicknessCircleRatioRight, thicknessLineRatioWRTCircle, poseScales, threshold, alpha=alpha) + #img = render_keypoints(img, right_hand_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1) + return img + +def render_body_keypoints(img: np.array, + body_keypoints: np.array) -> np.array: + """ + Render OpenPose body keypoints on input image. + Args: + img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range. + body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). + Returns: + (np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image. + """ + + thickness_circle_ratio = 1./75. * np.ones(body_keypoints.shape[0]) + thickness_line_ratio_wrt_circle = 0.75 + pairs = [] + pairs = [1,8,1,2,1,5,2,3,3,4,5,6,6,7,8,9,9,10,10,11,8,12,12,13,13,14,1,0,0,15,15,17,0,16,16,18,14,19,19,20,14,21,11,22,22,23,11,24] + pairs = np.array(pairs).reshape(-1,2) + colors = [255., 0., 85., + 255., 0., 0., + 255., 85., 0., + 255., 170., 0., + 255., 255., 0., + 170., 255., 0., + 85., 255., 0., + 0., 255., 0., + 255., 0., 0., + 0., 255., 85., + 0., 255., 170., + 0., 255., 255., + 0., 170., 255., + 0., 85., 255., + 0., 0., 255., + 255., 0., 170., + 170., 0., 255., + 255., 0., 255., + 85., 0., 255., + 0., 0., 255., + 0., 0., 255., + 0., 0., 255., + 0., 255., 255., + 0., 255., 255., + 0., 255., 255.] + colors = np.array(colors).reshape(-1,3) + pose_scales = [1] + return render_keypoints(img, body_keypoints, pairs, colors, thickness_circle_ratio, thickness_line_ratio_wrt_circle, pose_scales, 0.1) + +def render_openpose(img: np.array, + hand_keypoints: np.array) -> np.array: + """ + Render keypoints in the OpenPose format on input image. + Args: + img (np.array): Input image of shape (H, W, 3) with pixel values in the [0,255] range. + body_keypoints (np.array): Keypoint array of shape (N, 3); 3 <====> (x, y, confidence). + Returns: + (np.array): Image of shape (H, W, 3) with keypoints drawn on top of the original image. + """ + #img = render_body_keypoints(img, body_keypoints) + img = render_hand_keypoints(img, hand_keypoints) + return img diff --git a/wilor/utils/renderer.py b/wilor/utils/renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..0e161bb05921e52a684427e3eb87c4f8739a5d89 --- /dev/null +++ b/wilor/utils/renderer.py @@ -0,0 +1,423 @@ +import os +if 'PYOPENGL_PLATFORM' not in os.environ: + os.environ['PYOPENGL_PLATFORM'] = 'egl' +import torch +import numpy as np +import pyrender +import trimesh +import cv2 +from yacs.config import CfgNode +from typing import List, Optional + +def cam_crop_to_full(cam_bbox, box_center, box_size, img_size, focal_length=5000.): + # Convert cam_bbox to full image + img_w, img_h = img_size[:, 0], img_size[:, 1] + cx, cy, b = box_center[:, 0], box_center[:, 1], box_size + w_2, h_2 = img_w / 2., img_h / 2. + bs = b * cam_bbox[:, 0] + 1e-9 + tz = 2 * focal_length / bs + tx = (2 * (cx - w_2) / bs) + cam_bbox[:, 1] + ty = (2 * (cy - h_2) / bs) + cam_bbox[:, 2] + full_cam = torch.stack([tx, ty, tz], dim=-1) + return full_cam + +def get_light_poses(n_lights=5, elevation=np.pi / 3, dist=12): + # get lights in a circle around origin at elevation + thetas = elevation * np.ones(n_lights) + phis = 2 * np.pi * np.arange(n_lights) / n_lights + poses = [] + trans = make_translation(torch.tensor([0, 0, dist])) + for phi, theta in zip(phis, thetas): + rot = make_rotation(rx=-theta, ry=phi, order="xyz") + poses.append((rot @ trans).numpy()) + return poses + +def make_translation(t): + return make_4x4_pose(torch.eye(3), t) + +def make_rotation(rx=0, ry=0, rz=0, order="xyz"): + Rx = rotx(rx) + Ry = roty(ry) + Rz = rotz(rz) + if order == "xyz": + R = Rz @ Ry @ Rx + elif order == "xzy": + R = Ry @ Rz @ Rx + elif order == "yxz": + R = Rz @ Rx @ Ry + elif order == "yzx": + R = Rx @ Rz @ Ry + elif order == "zyx": + R = Rx @ Ry @ Rz + elif order == "zxy": + R = Ry @ Rx @ Rz + return make_4x4_pose(R, torch.zeros(3)) + +def make_4x4_pose(R, t): + """ + :param R (*, 3, 3) + :param t (*, 3) + return (*, 4, 4) + """ + dims = R.shape[:-2] + pose_3x4 = torch.cat([R, t.view(*dims, 3, 1)], dim=-1) + bottom = ( + torch.tensor([0, 0, 0, 1], device=R.device) + .reshape(*(1,) * len(dims), 1, 4) + .expand(*dims, 1, 4) + ) + return torch.cat([pose_3x4, bottom], dim=-2) + + +def rotx(theta): + return torch.tensor( + [ + [1, 0, 0], + [0, np.cos(theta), -np.sin(theta)], + [0, np.sin(theta), np.cos(theta)], + ], + dtype=torch.float32, + ) + + +def roty(theta): + return torch.tensor( + [ + [np.cos(theta), 0, np.sin(theta)], + [0, 1, 0], + [-np.sin(theta), 0, np.cos(theta)], + ], + dtype=torch.float32, + ) + + +def rotz(theta): + return torch.tensor( + [ + [np.cos(theta), -np.sin(theta), 0], + [np.sin(theta), np.cos(theta), 0], + [0, 0, 1], + ], + dtype=torch.float32, + ) + + +def create_raymond_lights() -> List[pyrender.Node]: + """ + Return raymond light nodes for the scene. + """ + thetas = np.pi * np.array([1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0]) + phis = np.pi * np.array([0.0, 2.0 / 3.0, 4.0 / 3.0]) + + nodes = [] + + for phi, theta in zip(phis, thetas): + xp = np.sin(theta) * np.cos(phi) + yp = np.sin(theta) * np.sin(phi) + zp = np.cos(theta) + + z = np.array([xp, yp, zp]) + z = z / np.linalg.norm(z) + x = np.array([-z[1], z[0], 0.0]) + if np.linalg.norm(x) == 0: + x = np.array([1.0, 0.0, 0.0]) + x = x / np.linalg.norm(x) + y = np.cross(z, x) + + matrix = np.eye(4) + matrix[:3,:3] = np.c_[x,y,z] + nodes.append(pyrender.Node( + light=pyrender.DirectionalLight(color=np.ones(3), intensity=1.0), + matrix=matrix + )) + + return nodes + +class Renderer: + + def __init__(self, cfg: CfgNode, faces: np.array): + """ + Wrapper around the pyrender renderer to render MANO meshes. + Args: + cfg (CfgNode): Model config file. + faces (np.array): Array of shape (F, 3) containing the mesh faces. + """ + self.cfg = cfg + self.focal_length = cfg.EXTRA.FOCAL_LENGTH + self.img_res = cfg.MODEL.IMAGE_SIZE + + # add faces that make the hand mesh watertight + faces_new = np.array([[92, 38, 234], + [234, 38, 239], + [38, 122, 239], + [239, 122, 279], + [122, 118, 279], + [279, 118, 215], + [118, 117, 215], + [215, 117, 214], + [117, 119, 214], + [214, 119, 121], + [119, 120, 121], + [121, 120, 78], + [120, 108, 78], + [78, 108, 79]]) + faces = np.concatenate([faces, faces_new], axis=0) + + self.camera_center = [self.img_res // 2, self.img_res // 2] + self.faces = faces + self.faces_left = self.faces[:,[0,2,1]] + + def __call__(self, + vertices: np.array, + camera_translation: np.array, + image: torch.Tensor, + full_frame: bool = False, + imgname: Optional[str] = None, + side_view=False, rot_angle=90, + mesh_base_color=(1.0, 1.0, 0.9), + scene_bg_color=(0,0,0), + return_rgba=False, + ) -> np.array: + """ + Render meshes on input image + Args: + vertices (np.array): Array of shape (V, 3) containing the mesh vertices. + camera_translation (np.array): Array of shape (3,) with the camera translation. + image (torch.Tensor): Tensor of shape (3, H, W) containing the image crop with normalized pixel values. + full_frame (bool): If True, then render on the full image. + imgname (Optional[str]): Contains the original image filenamee. Used only if full_frame == True. + """ + + if full_frame: + image = cv2.imread(imgname).astype(np.float32)[:, :, ::-1] / 255. + else: + image = image.clone() * torch.tensor(self.cfg.MODEL.IMAGE_STD, device=image.device).reshape(3,1,1) + image = image + torch.tensor(self.cfg.MODEL.IMAGE_MEAN, device=image.device).reshape(3,1,1) + image = image.permute(1, 2, 0).cpu().numpy() + + renderer = pyrender.OffscreenRenderer(viewport_width=image.shape[1], + viewport_height=image.shape[0], + point_size=1.0) + material = pyrender.MetallicRoughnessMaterial( + metallicFactor=0.0, + alphaMode='OPAQUE', + baseColorFactor=(*mesh_base_color, 1.0)) + + camera_translation[0] *= -1. + + mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy()) + if side_view: + rot = trimesh.transformations.rotation_matrix( + np.radians(rot_angle), [0, 1, 0]) + mesh.apply_transform(rot) + rot = trimesh.transformations.rotation_matrix( + np.radians(180), [1, 0, 0]) + mesh.apply_transform(rot) + mesh = pyrender.Mesh.from_trimesh(mesh, material=material) + + scene = pyrender.Scene(bg_color=[*scene_bg_color, 0.0], + ambient_light=(0.3, 0.3, 0.3)) + scene.add(mesh, 'mesh') + + camera_pose = np.eye(4) + camera_pose[:3, 3] = camera_translation + camera_center = [image.shape[1] / 2., image.shape[0] / 2.] + camera = pyrender.IntrinsicsCamera(fx=self.focal_length, fy=self.focal_length, + cx=camera_center[0], cy=camera_center[1], zfar=1e12) + scene.add(camera, pose=camera_pose) + + + light_nodes = create_raymond_lights() + for node in light_nodes: + scene.add_node(node) + + color, rend_depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA) + color = color.astype(np.float32) / 255.0 + renderer.delete() + + if return_rgba: + return color + + valid_mask = (color[:, :, -1])[:, :, np.newaxis] + if not side_view: + output_img = (color[:, :, :3] * valid_mask + (1 - valid_mask) * image) + else: + output_img = color[:, :, :3] + + output_img = output_img.astype(np.float32) + return output_img + + def vertices_to_trimesh(self, vertices, camera_translation, mesh_base_color=(1.0, 1.0, 0.9), + rot_axis=[1,0,0], rot_angle=0, is_right=1): + # material = pyrender.MetallicRoughnessMaterial( + # metallicFactor=0.0, + # alphaMode='OPAQUE', + # baseColorFactor=(*mesh_base_color, 1.0)) + vertex_colors = np.array([(*mesh_base_color, 1.0)] * vertices.shape[0]) + if is_right: + mesh = trimesh.Trimesh(vertices.copy() + camera_translation, self.faces.copy(), vertex_colors=vertex_colors) + else: + mesh = trimesh.Trimesh(vertices.copy() + camera_translation, self.faces_left.copy(), vertex_colors=vertex_colors) + # mesh = trimesh.Trimesh(vertices.copy(), self.faces.copy()) + + rot = trimesh.transformations.rotation_matrix( + np.radians(rot_angle), rot_axis) + mesh.apply_transform(rot) + + rot = trimesh.transformations.rotation_matrix( + np.radians(180), [1, 0, 0]) + mesh.apply_transform(rot) + return mesh + + def render_rgba( + self, + vertices: np.array, + cam_t = None, + rot=None, + rot_axis=[1,0,0], + rot_angle=0, + camera_z=3, + # camera_translation: np.array, + mesh_base_color=(1.0, 1.0, 0.9), + scene_bg_color=(0,0,0), + render_res=[256, 256], + focal_length=None, + is_right=None, + ): + + renderer = pyrender.OffscreenRenderer(viewport_width=render_res[0], + viewport_height=render_res[1], + point_size=1.0) + # material = pyrender.MetallicRoughnessMaterial( + # metallicFactor=0.0, + # alphaMode='OPAQUE', + # baseColorFactor=(*mesh_base_color, 1.0)) + + focal_length = focal_length if focal_length is not None else self.focal_length + + if cam_t is not None: + camera_translation = cam_t.copy() + camera_translation[0] *= -1. + else: + camera_translation = np.array([0, 0, camera_z * focal_length/render_res[1]]) + + mesh = self.vertices_to_trimesh(vertices, np.array([0, 0, 0]), mesh_base_color, rot_axis, rot_angle, is_right=is_right) + mesh = pyrender.Mesh.from_trimesh(mesh) + # mesh = pyrender.Mesh.from_trimesh(mesh, material=material) + + scene = pyrender.Scene(bg_color=[*scene_bg_color, 0.0], + ambient_light=(0.3, 0.3, 0.3)) + scene.add(mesh, 'mesh') + + camera_pose = np.eye(4) + camera_pose[:3, 3] = camera_translation + camera_center = [render_res[0] / 2., render_res[1] / 2.] + camera = pyrender.IntrinsicsCamera(fx=focal_length, fy=focal_length, + cx=camera_center[0], cy=camera_center[1], zfar=1e12) + + # Create camera node and add it to pyRender scene + camera_node = pyrender.Node(camera=camera, matrix=camera_pose) + scene.add_node(camera_node) + self.add_point_lighting(scene, camera_node) + self.add_lighting(scene, camera_node) + + light_nodes = create_raymond_lights() + for node in light_nodes: + scene.add_node(node) + + color, rend_depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA) + color = color.astype(np.float32) / 255.0 + renderer.delete() + + return color + + def render_rgba_multiple( + self, + vertices: List[np.array], + cam_t: List[np.array], + rot_axis=[1,0,0], + rot_angle=0, + mesh_base_color=(1.0, 1.0, 0.9), + scene_bg_color=(0,0,0), + render_res=[256, 256], + focal_length=None, + is_right=None, + ): + + renderer = pyrender.OffscreenRenderer(viewport_width=render_res[0], + viewport_height=render_res[1], + point_size=1.0) + # material = pyrender.MetallicRoughnessMaterial( + # metallicFactor=0.0, + # alphaMode='OPAQUE', + # baseColorFactor=(*mesh_base_color, 1.0)) + + if is_right is None: + is_right = [1 for _ in range(len(vertices))] + + mesh_list = [pyrender.Mesh.from_trimesh(self.vertices_to_trimesh(vvv, ttt.copy(), mesh_base_color, rot_axis, rot_angle, is_right=sss)) for vvv,ttt,sss in zip(vertices, cam_t, is_right)] + + scene = pyrender.Scene(bg_color=[*scene_bg_color, 0.0], + ambient_light=(0.3, 0.3, 0.3)) + for i,mesh in enumerate(mesh_list): + scene.add(mesh, f'mesh_{i}') + + camera_pose = np.eye(4) + # camera_pose[:3, 3] = camera_translation + camera_center = [render_res[0] / 2., render_res[1] / 2.] + focal_length = focal_length if focal_length is not None else self.focal_length + camera = pyrender.IntrinsicsCamera(fx=focal_length, fy=focal_length, + cx=camera_center[0], cy=camera_center[1], zfar=1e12) + + # Create camera node and add it to pyRender scene + camera_node = pyrender.Node(camera=camera, matrix=camera_pose) + scene.add_node(camera_node) + self.add_point_lighting(scene, camera_node) + self.add_lighting(scene, camera_node) + + light_nodes = create_raymond_lights() + for node in light_nodes: + scene.add_node(node) + + color, rend_depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA) + color = color.astype(np.float32) / 255.0 + renderer.delete() + + return color + + def add_lighting(self, scene, cam_node, color=np.ones(3), intensity=1.0): + # from phalp.visualize.py_renderer import get_light_poses + light_poses = get_light_poses() + light_poses.append(np.eye(4)) + cam_pose = scene.get_pose(cam_node) + for i, pose in enumerate(light_poses): + matrix = cam_pose @ pose + node = pyrender.Node( + name=f"light-{i:02d}", + light=pyrender.DirectionalLight(color=color, intensity=intensity), + matrix=matrix, + ) + if scene.has_node(node): + continue + scene.add_node(node) + + def add_point_lighting(self, scene, cam_node, color=np.ones(3), intensity=1.0): + # from phalp.visualize.py_renderer import get_light_poses + light_poses = get_light_poses(dist=0.5) + light_poses.append(np.eye(4)) + cam_pose = scene.get_pose(cam_node) + for i, pose in enumerate(light_poses): + matrix = cam_pose @ pose + # node = pyrender.Node( + # name=f"light-{i:02d}", + # light=pyrender.DirectionalLight(color=color, intensity=intensity), + # matrix=matrix, + # ) + node = pyrender.Node( + name=f"plight-{i:02d}", + light=pyrender.PointLight(color=color, intensity=intensity), + matrix=matrix, + ) + if scene.has_node(node): + continue + scene.add_node(node) diff --git a/wilor/utils/rich_utils.py b/wilor/utils/rich_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..19f97494ed2958ec2c3d75c772360b5367f2dc7b --- /dev/null +++ b/wilor/utils/rich_utils.py @@ -0,0 +1,105 @@ +from pathlib import Path +from typing import Sequence + +import rich +import rich.syntax +import rich.tree +from hydra.core.hydra_config import HydraConfig +from omegaconf import DictConfig, OmegaConf, open_dict +from pytorch_lightning.utilities import rank_zero_only +from rich.prompt import Prompt + +from . import pylogger + +log = pylogger.get_pylogger(__name__) + + +@rank_zero_only +def print_config_tree( + cfg: DictConfig, + print_order: Sequence[str] = ( + "datamodule", + "model", + "callbacks", + "logger", + "trainer", + "paths", + "extras", + ), + resolve: bool = False, + save_to_file: bool = False, +) -> None: + """Prints content of DictConfig using Rich library and its tree structure. + + Args: + cfg (DictConfig): Configuration composed by Hydra. + print_order (Sequence[str], optional): Determines in what order config components are printed. + resolve (bool, optional): Whether to resolve reference fields of DictConfig. + save_to_file (bool, optional): Whether to export config to the hydra output folder. + """ + + style = "dim" + tree = rich.tree.Tree("CONFIG", style=style, guide_style=style) + + queue = [] + + # add fields from `print_order` to queue + for field in print_order: + queue.append(field) if field in cfg else log.warning( + f"Field '{field}' not found in config. Skipping '{field}' config printing..." + ) + + # add all the other fields to queue (not specified in `print_order`) + for field in cfg: + if field not in queue: + queue.append(field) + + # generate config tree from queue + for field in queue: + branch = tree.add(field, style=style, guide_style=style) + + config_group = cfg[field] + if isinstance(config_group, DictConfig): + branch_content = OmegaConf.to_yaml(config_group, resolve=resolve) + else: + branch_content = str(config_group) + + branch.add(rich.syntax.Syntax(branch_content, "yaml")) + + # print config tree + rich.print(tree) + + # save config tree to file + if save_to_file: + with open(Path(cfg.paths.output_dir, "config_tree.log"), "w") as file: + rich.print(tree, file=file) + + +@rank_zero_only +def enforce_tags(cfg: DictConfig, save_to_file: bool = False) -> None: + """Prompts user to input tags from command line if no tags are provided in config.""" + + if not cfg.get("tags"): + if "id" in HydraConfig().cfg.hydra.job: + raise ValueError("Specify tags before launching a multirun!") + + log.warning("No tags provided in config. Prompting user to input tags...") + tags = Prompt.ask("Enter a list of comma separated tags", default="dev") + tags = [t.strip() for t in tags.split(",") if t != ""] + + with open_dict(cfg): + cfg.tags = tags + + log.info(f"Tags: {cfg.tags}") + + if save_to_file: + with open(Path(cfg.paths.output_dir, "tags.log"), "w") as file: + rich.print(cfg.tags, file=file) + + +if __name__ == "__main__": + from hydra import compose, initialize + + with initialize(version_base="1.2", config_path="../../configs"): + cfg = compose(config_name="train.yaml", return_hydra_config=False, overrides=[]) + print_config_tree(cfg, resolve=False, save_to_file=False) diff --git a/wilor/utils/skeleton_renderer.py b/wilor/utils/skeleton_renderer.py new file mode 100644 index 0000000000000000000000000000000000000000..46a5df75bff887eab00984eeb5be3c1f6e752960 --- /dev/null +++ b/wilor/utils/skeleton_renderer.py @@ -0,0 +1,124 @@ +import torch +import numpy as np +import trimesh +from typing import Optional +from yacs.config import CfgNode + +from .geometry import perspective_projection +from .render_openpose import render_openpose + +class SkeletonRenderer: + + def __init__(self, cfg: CfgNode): + """ + Object used to render 3D keypoints. Faster for use during training. + Args: + cfg (CfgNode): Model config file. + """ + self.cfg = cfg + + def __call__(self, + pred_keypoints_3d: torch.Tensor, + gt_keypoints_3d: torch.Tensor, + gt_keypoints_2d: torch.Tensor, + images: Optional[np.array] = None, + camera_translation: Optional[torch.Tensor] = None) -> np.array: + """ + Render batch of 3D keypoints. + Args: + pred_keypoints_3d (torch.Tensor): Tensor of shape (B, S, N, 3) containing a batch of predicted 3D keypoints, with S samples per image. + gt_keypoints_3d (torch.Tensor): Tensor of shape (B, N, 4) containing corresponding ground truth 3D keypoints; last value is the confidence. + gt_keypoints_2d (torch.Tensor): Tensor of shape (B, N, 3) containing corresponding ground truth 2D keypoints. + images (torch.Tensor): Tensor of shape (B, H, W, 3) containing images with values in the [0,255] range. + camera_translation (torch.Tensor): Tensor of shape (B, 3) containing the camera translation. + Returns: + np.array : Image with the following layout. Each row contains the a) input image, + b) image with gt 2D keypoints, + c) image with projected gt 3D keypoints, + d_1, ... , d_S) image with projected predicted 3D keypoints, + e) gt 3D keypoints rendered from a side view, + f_1, ... , f_S) predicted 3D keypoints frorm a side view + """ + batch_size = pred_keypoints_3d.shape[0] +# num_samples = pred_keypoints_3d.shape[1] + pred_keypoints_3d = pred_keypoints_3d.clone().cpu().float() + gt_keypoints_3d = gt_keypoints_3d.clone().cpu().float() + gt_keypoints_3d[:, :, :-1] = gt_keypoints_3d[:, :, :-1] - gt_keypoints_3d[:, [0], :-1] + pred_keypoints_3d[:, [0]] + gt_keypoints_2d = gt_keypoints_2d.clone().cpu().float().numpy() + gt_keypoints_2d[:, :, :-1] = self.cfg.MODEL.IMAGE_SIZE * (gt_keypoints_2d[:, :, :-1] + 1.0) / 2.0 + + #openpose_indices = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] + #gt_indices = [12, 8, 7, 6, 9, 10, 11, 14, 2, 1, 0, 3, 4, 5] + #gt_indices = [25 + i for i in gt_indices] + openpose_indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20] + gt_indices = openpose_indices + keypoints_to_render = torch.ones(batch_size, gt_keypoints_3d.shape[1], 1) + rotation = torch.eye(3).unsqueeze(0) + if camera_translation is None: + camera_translation = torch.tensor([0.0, 0.0, 2 * self.cfg.EXTRA.FOCAL_LENGTH / (0.8 * self.cfg.MODEL.IMAGE_SIZE)]).unsqueeze(0).repeat(batch_size, 1) + else: + camera_translation = camera_translation.cpu() + + if images is None: + images = np.zeros((batch_size, self.cfg.MODEL.IMAGE_SIZE, self.cfg.MODEL.IMAGE_SIZE, 3)) + focal_length = torch.tensor([self.cfg.EXTRA.FOCAL_LENGTH, self.cfg.EXTRA.FOCAL_LENGTH]).reshape(1, 2) + camera_center = torch.tensor([self.cfg.MODEL.IMAGE_SIZE, self.cfg.MODEL.IMAGE_SIZE], dtype=torch.float).reshape(1, 2) / 2. + gt_keypoints_3d_proj = perspective_projection(gt_keypoints_3d[:, :, :-1], rotation=rotation.repeat(batch_size, 1, 1), translation=camera_translation[:, :], focal_length=focal_length.repeat(batch_size, 1), camera_center=camera_center.repeat(batch_size, 1)) + pred_keypoints_3d_proj = perspective_projection(pred_keypoints_3d.reshape(batch_size, -1, 3), rotation=rotation.repeat(batch_size, 1, 1), translation=camera_translation.reshape(batch_size, -1), focal_length=focal_length.repeat(batch_size, 1), camera_center=camera_center.repeat(batch_size, 1)).reshape(batch_size, -1, 2) + gt_keypoints_3d_proj = torch.cat([gt_keypoints_3d_proj, gt_keypoints_3d[:, :, [-1]]], dim=-1).cpu().numpy() + pred_keypoints_3d_proj = torch.cat([pred_keypoints_3d_proj, keypoints_to_render.reshape(batch_size, -1, 1)], dim=-1).cpu().numpy() + rows = [] + # Rotate keypoints to visualize side view + R = torch.tensor(trimesh.transformations.rotation_matrix(np.radians(90), [0, 1, 0])[:3, :3]).float() + gt_keypoints_3d_side = gt_keypoints_3d.clone() + gt_keypoints_3d_side[:, :, :-1] = torch.einsum('bni,ij->bnj', gt_keypoints_3d_side[:, :, :-1], R) + pred_keypoints_3d_side = pred_keypoints_3d.clone() + pred_keypoints_3d_side = torch.einsum('bni,ij->bnj', pred_keypoints_3d_side, R) + gt_keypoints_3d_proj_side = perspective_projection(gt_keypoints_3d_side[:, :, :-1], rotation=rotation.repeat(batch_size, 1, 1), translation=camera_translation[:, :], focal_length=focal_length.repeat(batch_size, 1), camera_center=camera_center.repeat(batch_size, 1)) + pred_keypoints_3d_proj_side = perspective_projection(pred_keypoints_3d_side.reshape(batch_size, -1, 3), rotation=rotation.repeat(batch_size, 1, 1), translation=camera_translation.reshape(batch_size, -1), focal_length=focal_length.repeat(batch_size, 1), camera_center=camera_center.repeat(batch_size, 1)).reshape(batch_size, -1, 2) + gt_keypoints_3d_proj_side = torch.cat([gt_keypoints_3d_proj_side, gt_keypoints_3d_side[:, :, [-1]]], dim=-1).cpu().numpy() + pred_keypoints_3d_proj_side = torch.cat([pred_keypoints_3d_proj_side, keypoints_to_render.reshape(batch_size, -1, 1)], dim=-1).cpu().numpy() + for i in range(batch_size): + img = images[i] + side_img = np.zeros((self.cfg.MODEL.IMAGE_SIZE, self.cfg.MODEL.IMAGE_SIZE, 3)) + # gt 2D keypoints + body_keypoints_2d = gt_keypoints_2d[i, :21].copy() + for op, gt in zip(openpose_indices, gt_indices): + if gt_keypoints_2d[i, gt, -1] > body_keypoints_2d[op, -1]: + body_keypoints_2d[op] = gt_keypoints_2d[i, gt] + gt_keypoints_img = render_openpose(img, body_keypoints_2d) / 255. + # gt 3D keypoints + body_keypoints_3d_proj = gt_keypoints_3d_proj[i, :21].copy() + for op, gt in zip(openpose_indices, gt_indices): + if gt_keypoints_3d_proj[i, gt, -1] > body_keypoints_3d_proj[op, -1]: + body_keypoints_3d_proj[op] = gt_keypoints_3d_proj[i, gt] + gt_keypoints_3d_proj_img = render_openpose(img, body_keypoints_3d_proj) / 255. + # gt 3D keypoints from the side + body_keypoints_3d_proj = gt_keypoints_3d_proj_side[i, :21].copy() + for op, gt in zip(openpose_indices, gt_indices): + if gt_keypoints_3d_proj_side[i, gt, -1] > body_keypoints_3d_proj[op, -1]: + body_keypoints_3d_proj[op] = gt_keypoints_3d_proj_side[i, gt] + gt_keypoints_3d_proj_img_side = render_openpose(side_img, body_keypoints_3d_proj) / 255. + # pred 3D keypoints + pred_keypoints_3d_proj_imgs = [] + body_keypoints_3d_proj = pred_keypoints_3d_proj[i, :21].copy() + for op, gt in zip(openpose_indices, gt_indices): + if pred_keypoints_3d_proj[i, gt, -1] >= body_keypoints_3d_proj[op, -1]: + body_keypoints_3d_proj[op] = pred_keypoints_3d_proj[i, gt] + pred_keypoints_3d_proj_imgs.append(render_openpose(img, body_keypoints_3d_proj) / 255.) + pred_keypoints_3d_proj_img = np.concatenate(pred_keypoints_3d_proj_imgs, axis=1) + # gt 3D keypoints from the side + pred_keypoints_3d_proj_imgs_side = [] + body_keypoints_3d_proj = pred_keypoints_3d_proj_side[i, :21].copy() + for op, gt in zip(openpose_indices, gt_indices): + if pred_keypoints_3d_proj_side[i, gt, -1] >= body_keypoints_3d_proj[op, -1]: + body_keypoints_3d_proj[op] = pred_keypoints_3d_proj_side[i, gt] + pred_keypoints_3d_proj_imgs_side.append(render_openpose(side_img, body_keypoints_3d_proj) / 255.) + pred_keypoints_3d_proj_img_side = np.concatenate(pred_keypoints_3d_proj_imgs_side, axis=1) + rows.append(np.concatenate((gt_keypoints_img, gt_keypoints_3d_proj_img, pred_keypoints_3d_proj_img, gt_keypoints_3d_proj_img_side, pred_keypoints_3d_proj_img_side), axis=1)) + # Concatenate images + img = np.concatenate(rows, axis=0) + img[:, ::self.cfg.MODEL.IMAGE_SIZE, :] = 1.0 + img[::self.cfg.MODEL.IMAGE_SIZE, :, :] = 1.0 + img[:, (1+1+1)*self.cfg.MODEL.IMAGE_SIZE, :] = 0.5 + return img