diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d8448fb --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +**/__pycache__ +.DS_Store + diff --git a/README.md b/README.md new file mode 100644 index 0000000..01462e3 --- /dev/null +++ b/README.md @@ -0,0 +1,8 @@ + + + +Linux +``` +apt-get install freeglut3-dev +apt-get install libosmesa6-dev +``` \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ce0f59f --- /dev/null +++ b/__init__.py @@ -0,0 +1,29 @@ +import os + +REPO_PATH = os.path.dirname(os.path.abspath(__file__)) + +from folder_paths import models_dir + +SMPL_PATH = os.path.join(models_dir, 'smpl') +os.makedirs(SMPL_PATH, exist_ok=True) + + +from .nodes.process_humans_node import ProcessHumansNode +from .nodes.load_detectron_node import LoadDetectronNode +from .nodes.load_hmr_node import LoadHMRNode +# from .nodes.select_human_node import SelectHumanNode + + +NODE_CLASS_MAPPINGS = { + 'ProcessHumans': ProcessHumansNode, + 'LoadDetectron': LoadDetectronNode, + 'LoadHMR': LoadHMRNode, + # 'SelectHuman': SelectHumanNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + 'ProcessHumans': 'Process 4D Humans', + 'LoadDetectron': 'Load Detectron Model', + 'LoadHMR': 'Load HMR Model', + # 'SelectHuman': 'Select 4D Human' +} diff --git a/configs/cascade_mask_rcnn_vitdet_h_75ep.py b/configs/cascade_mask_rcnn_vitdet_h_75ep.py new file mode 100644 index 0000000..0c6ae0e --- /dev/null +++ b/configs/cascade_mask_rcnn_vitdet_h_75ep.py @@ -0,0 +1,129 @@ +## coco_loader_lsj.py + +import detectron2.data.transforms as T +from detectron2 import model_zoo +from detectron2.config import LazyCall as L + +# Data using LSJ +image_size = 1024 +dataloader = model_zoo.get_config("common/data/coco.py").dataloader +dataloader.train.mapper.augmentations = [ + L(T.RandomFlip)(horizontal=True), # flip first + L(T.ResizeScale)( + min_scale=0.1, max_scale=2.0, target_height=image_size, target_width=image_size + ), + L(T.FixedSizeCrop)(crop_size=(image_size, image_size), pad=False), +] +dataloader.train.mapper.image_format = "RGB" +dataloader.train.total_batch_size = 64 +# recompute boxes due to cropping +dataloader.train.mapper.recompute_boxes = True + +dataloader.test.mapper.augmentations = [ + L(T.ResizeShortestEdge)(short_edge_length=image_size, max_size=image_size), +] + +from functools import partial +from fvcore.common.param_scheduler import MultiStepParamScheduler + +from detectron2 import model_zoo +from detectron2.config import LazyCall as L +from detectron2.solver import WarmupParamScheduler +from detectron2.modeling.backbone.vit import get_vit_lr_decay_rate + +# mask_rcnn_vitdet_b_100ep.py + +model = model_zoo.get_config("common/models/mask_rcnn_vitdet.py").model + +# Initialization and trainer settings +train = model_zoo.get_config("common/train.py").train +train.amp.enabled = True +train.ddp.fp16_compression = True +train.init_checkpoint = "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth" + + +# Schedule +# 100 ep = 184375 iters * 64 images/iter / 118000 images/ep +train.max_iter = 184375 + +lr_multiplier = L(WarmupParamScheduler)( + scheduler=L(MultiStepParamScheduler)( + values=[1.0, 0.1, 0.01], + milestones=[163889, 177546], + num_updates=train.max_iter, + ), + warmup_length=250 / train.max_iter, + warmup_factor=0.001, +) + +# Optimizer +optimizer = model_zoo.get_config("common/optim.py").AdamW +optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, num_layers=12, lr_decay_rate=0.7) +optimizer.params.overrides = {"pos_embed": {"weight_decay": 0.0}} + +# cascade_mask_rcnn_vitdet_b_100ep.py + +from detectron2.config import LazyCall as L +from detectron2.layers import ShapeSpec +from detectron2.modeling.box_regression import Box2BoxTransform +from detectron2.modeling.matcher import Matcher +from detectron2.modeling.roi_heads import ( + FastRCNNOutputLayers, + FastRCNNConvFCHead, + CascadeROIHeads, +) + +# arguments that don't exist for Cascade R-CNN +[model.roi_heads.pop(k) for k in ["box_head", "box_predictor", "proposal_matcher"]] + +model.roi_heads.update( + _target_=CascadeROIHeads, + box_heads=[ + L(FastRCNNConvFCHead)( + input_shape=ShapeSpec(channels=256, height=7, width=7), + conv_dims=[256, 256, 256, 256], + fc_dims=[1024], + conv_norm="LN", + ) + for _ in range(3) + ], + box_predictors=[ + L(FastRCNNOutputLayers)( + input_shape=ShapeSpec(channels=1024), + test_score_thresh=0.05, + box2box_transform=L(Box2BoxTransform)(weights=(w1, w1, w2, w2)), + cls_agnostic_bbox_reg=True, + num_classes="${...num_classes}", + ) + for (w1, w2) in [(10, 5), (20, 10), (30, 15)] + ], + proposal_matchers=[ + L(Matcher)(thresholds=[th], labels=[0, 1], allow_low_quality_matches=False) + for th in [0.5, 0.6, 0.7] + ], +) + +# cascade_mask_rcnn_vitdet_h_75ep.py + +from functools import partial + +train.init_checkpoint = "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_huge_p14to16.pth" + +model.backbone.net.embed_dim = 1280 +model.backbone.net.depth = 32 +model.backbone.net.num_heads = 16 +model.backbone.net.drop_path_rate = 0.5 +# 7, 15, 23, 31 for global attention +model.backbone.net.window_block_indexes = ( + list(range(0, 7)) + list(range(8, 15)) + list(range(16, 23)) + list(range(24, 31)) +) + +optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, lr_decay_rate=0.9, num_layers=32) +optimizer.params.overrides = {} +optimizer.params.weight_decay_norm = None + +train.max_iter = train.max_iter * 3 // 4 # 100ep -> 75ep +lr_multiplier.scheduler.milestones = [ + milestone * 3 // 4 for milestone in lr_multiplier.scheduler.milestones +] +lr_multiplier.scheduler.num_updates = train.max_iter diff --git a/configs/model_config.yaml b/configs/model_config.yaml new file mode 100644 index 0000000..a8db0fe --- /dev/null +++ b/configs/model_config.yaml @@ -0,0 +1,148 @@ +task_name: train +tags: +- dev +train: true +test: false +ckpt_path: null +seed: null +DATASETS: + SUPPRESS_KP_CONF_THRESH: 0.3 + FILTER_NUM_KP: 4 + FILTER_NUM_KP_THRESH: 0.0 + FILTER_REPROJ_THRESH: 31000 + SUPPRESS_BETAS_THRESH: 3.0 + SUPPRESS_BAD_POSES: true + POSES_BETAS_SIMULTANEOUS: true + FILTER_NO_POSES: false + TRAIN: + H36M-TRAIN-WMASK: + WEIGHT: 0.1 + MPII-TRAIN-WMASK: + WEIGHT: 0.1 + COCO-TRAIN-2014-WMASK-PRUNED: + WEIGHT: 0.1 + COCO-TRAIN-2014-VITPOSE-REPLICATE-PRUNED12: + WEIGHT: 0.1 + MPI-INF-TRAIN-PRUNED: + WEIGHT: 0.02 + AVA-TRAIN-MIDFRAMES-1FPS-WMASK: + WEIGHT: 0.19 + AIC-TRAIN-WMASK: + WEIGHT: 0.19 + INSTA-TRAIN-WMASK: + WEIGHT: 0.2 + VAL: + COCO-VAL: + WEIGHT: 1.0 + MOCAP: CMU-MOCAP + 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: true + FLIP_AUG_RATE: 0.5 + EXTREME_CROP_AUG_RATE: 0.1 + EXTREME_CROP_AUG_LEVEL: 1 + BETAS_REG: true +trainer: + _target_: pytorch_lightning.Trainer + default_root_dir: ${paths.output_dir} + accelerator: gpu + devices: 8 + deterministic: false + num_sanity_val_steps: 0 + log_every_n_steps: ${GENERAL.LOG_STEPS} + val_check_interval: ${GENERAL.VAL_STEPS} + precision: 16 + max_steps: ${GENERAL.TOTAL_STEPS} + move_metrics_to_cpu: true + limit_val_batches: 1 + track_grad_norm: -1 + strategy: ddp + num_nodes: 1 + sync_batchnorm: true +paths: + root_dir: ${oc.env:PROJECT_ROOT} + data_dir: ${paths.root_dir}/data/ + log_dir: /fsx/shubham/code/hmr2023/logs_hydra/ + output_dir: ${hydra:runtime.output_dir} + work_dir: ${hydra:runtime.cwd} +extras: + ignore_warnings: false + enforce_tags: true + print_config: true +exp_name: hmr2 +SMPL: + MODEL_PATH: data/smpl + GENDER: neutral + NUM_BODY_JOINTS: 23 + JOINT_REGRESSOR_EXTRA: data/SMPL_to_J19.pkl + MEAN_PARAMS: data/smpl_mean_params.npz +EXTRA: + FOCAL_LENGTH: 5000 + NUM_LOG_IMAGES: 4 + NUM_LOG_SAMPLES_PER_IMAGE: 8 + PELVIS_IND: 39 +MODEL: + IMAGE_SIZE: 256 + IMAGE_MEAN: + - 0.485 + - 0.456 + - 0.406 + IMAGE_STD: + - 0.229 + - 0.224 + - 0.225 + BACKBONE: + TYPE: vit + NUM_LAYERS: 50 + OUT_CHANNELS: 2048 + ADD_NECK: false + FLOW: + DIM: 144 + NUM_LAYERS: 4 + CONTEXT_FEATURES: 2048 + LAYER_HIDDEN_FEATURES: 1024 + LAYER_DEPTH: 2 + FC_HEAD: + NUM_FEATURES: 1024 + SMPL_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 +GENERAL: + TOTAL_STEPS: 1000000 + LOG_STEPS: 1000 + VAL_STEPS: 1000 + CHECKPOINT_STEPS: 10000 + CHECKPOINT_SAVE_TOP_K: 1 + NUM_WORKERS: 6 + PREFETCH_FACTOR: 2 +TRAIN: + LR: 1.0e-05 + WEIGHT_DECAY: 0.0001 + BATCH_SIZE: 48 + LOSS_REDUCTION: mean + NUM_TRAIN_SAMPLES: 2 + NUM_TEST_SAMPLES: 64 + POSE_2D_NOISE_RATIO: 0.01 + SMPL_PARAM_NOISE_RATIO: 0.005 +LOSS_WEIGHTS: + KEYPOINTS_3D: 0.05 + KEYPOINTS_2D: 0.01 + GLOBAL_ORIENT: 0.001 + BODY_POSE: 0.001 + BETAS: 0.0005 + ADVERSARIAL: 0.0005 +local: {} diff --git a/humans4d/.gitignore b/humans4d/.gitignore new file mode 100644 index 0000000..d5b69ce --- /dev/null +++ b/humans4d/.gitignore @@ -0,0 +1,148 @@ +# Specific +/logs*/ +/results/ +/sandbox/ +*.lock +*.pt +*.npy +/example_data/downloaded* +*.tar +*.tar.gz +/discord_sandbox/ +/demo_out/ +token_channel.csv +/hmr2_training_data/ +/hmr2_evaluation_data/ +/outputs/ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-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/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# 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/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ +/checkpoints/ +/data/ diff --git a/humans4d/LICENSE.md b/humans4d/LICENSE.md new file mode 100644 index 0000000..77224b1 --- /dev/null +++ b/humans4d/LICENSE.md @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 UC Regents, Shubham Goel + +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/humans4d/README.md b/humans4d/README.md new file mode 100644 index 0000000..0f5ca7c --- /dev/null +++ b/humans4d/README.md @@ -0,0 +1,104 @@ +# 4DHumans: Reconstructing and Tracking Humans with Transformers +Code repository for the paper: +**Humans in 4D: Reconstructing and Tracking Humans with Transformers** +[Shubham Goel](https://people.eecs.berkeley.edu/~shubham-goel/), [Georgios Pavlakos](https://geopavlakos.github.io/), [Jathushan Rajasegaran](http://people.eecs.berkeley.edu/~jathushan/), [Angjoo Kanazawa](https://people.eecs.berkeley.edu/~kanazawa/)\*, [Jitendra Malik](http://people.eecs.berkeley.edu/~malik/)\* + +[![arXiv](https://img.shields.io/badge/arXiv-2305.20091-00ff00.svg)](https://arxiv.org/pdf/2305.20091.pdf) [![Website shields.io](https://img.shields.io/website-up-down-green-red/http/shields.io.svg)](https://shubham-goel.github.io/4dhumans/) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1Ex4gE5v1bPR3evfhtG7sDHxQGsWwNwby?usp=sharing) [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/brjathu/HMR2.0) + + +![teaser](assets/teaser.png) + +## Installation and Setup +First, clone the repo. Then, we recommend creating a clean [conda](https://docs.conda.io/) environment, installing all dependencies, and finally activating the environment, as follows: +```bash +git clone https://github.com/shubham-goel/4D-Humans.git +cd 4D-Humans +conda env create -f environment.yml +conda activate 4D-humans +``` + +If conda is too slow, you can use pip: +```bash +conda create --name 4D-humans python=3.10 +conda activate 4D-humans +pip install torch +pip install -e .[all] +``` + +All checkpoints and data will automatically be downloaded to `$HOME/.cache/4DHumans` the first time you run the demo code. + +Besides these files, you also need to download the *SMPL* model. You will need the [neutral model](http://smplify.is.tue.mpg.de) for training and running the demo code. Please go to the corresponding website and register to get access to the downloads section. Download the model and place `basicModel_neutral_lbs_10_207_0_v1.0.0.pkl` in `./data/`. + +## Run demo on images +The following command will run ViTDet and HMR2.0 on all images in the specified `--img_folder`, and save renderings of the reconstructions in `--out_folder`. `--batch_size` batches the images together for faster processing. The `--side_view` flags additionally renders the side view of the reconstructed mesh, `--full_frame` renders all people together in front view, `--save_mesh` saves meshes as `.obj`s. +```bash +python demo.py \ + --img_folder example_data/images \ + --out_folder demo_out \ + --batch_size=48 --side_view --save_mesh --full_frame +``` + +## Run tracking demo on videos +Our tracker builds on PHALP, please install that first: +```bash +pip install git+https://github.com/brjathu/PHALP.git +``` + +Now, run `track.py` to reconstruct and track humans in any video. Input video source may be a video file, a folder of frames, or a youtube link: +```bash +# Run on video file +python track.py video.source="example_data/videos/gymnasts.mp4" + +# Run on extracted frames +python track.py video.source="/path/to/frames_folder/" + +# Run on a youtube link (depends on pytube working properly) +python track.py video.source=\'"https://www.youtube.com/watch?v=xEH_5T9jMVU"\' +``` +The output directory (`./outputs` by default) will contain a video rendering of the tracklets and a `.pkl` file containing the tracklets with 3D pose and shape. Please see the [PHALP](https://github.com/brjathu/PHALP) repository for details. + +## Training +Download the [training data](https://www.dropbox.com/sh/mjdwu59fxuhls5h/AACQ6FCGSrggUXmRzuubRHXIa) to `./hmr2_training_data/`, then start training using the following command: +``` +bash fetch_training_data.sh +python train.py exp_name=hmr2 data=mix_all experiment=hmr_vit_transformer trainer=gpu launcher=local +``` +Checkpoints and logs will be saved to `./logs/`. We trained on 8 A100 GPUs for 7 days using PyTorch 1.13.1 and PyTorch-Lightning 1.8.1 with CUDA 11.6 on a Linux system. You may adjust batch size and number of GPUs per your convenience. + +## Evaluation +Download the [evaluation metadata](https://www.dropbox.com/scl/fi/kl79djemdgqcl6d691er7/hmr2_evaluation_data.tar.gz?rlkey=ttmbdu3x5etxwqqyzwk581zjl) to `./hmr2_evaluation_data/`. Additionally, download the Human3.6M, 3DPW, LSP-Extended, COCO, and PoseTrack dataset images and update the corresponding paths in `hmr2/configs/datasets_eval.yaml`. + +Run evaluation on multiple datasets as follows, results are stored in `results/eval_regression.csv`. +```bash +python eval.py --dataset 'H36M-VAL-P2,3DPW-TEST,LSP-EXTENDED,POSETRACK-VAL,COCO-VAL' +``` + +By default, our code uses the released checkpoint (mentioned as HMR2.0b in the paper). To use the HMR2.0a checkpoint, you may download and untar from [here](https://people.eecs.berkeley.edu/~jathushan/projects/4dhumans/hmr2a_model.tar.gz) + +## Preprocess code +To preprocess LSP Extended and Posetrack into metadata zip files for evaluation, see `hmr2/datasets/preprocess`. + +Training data preprocessing coming soon. + +## Acknowledgements +Parts of the code are taken or adapted from the following repos: +- [ProHMR](https://github.com/nkolot/ProHMR) +- [SPIN](https://github.com/nkolot/SPIN) +- [SMPLify-X](https://github.com/vchoutas/smplify-x) +- [HMR](https://github.com/akanazawa/hmr) +- [ViTPose](https://github.com/ViTAE-Transformer/ViTPose) +- [Detectron2](https://github.com/facebookresearch/detectron2) + +Additionally, we thank [StabilityAI](https://stability.ai/) for a generous compute grant that enabled this work. + +## Citing +If you find this code useful for your research, please consider citing the following paper: + +```bibtex +@inproceedings{goel2023humans, + title={Humans in 4{D}: Reconstructing and Tracking Humans with Transformers}, + author={Goel, Shubham and Pavlakos, Georgios and Rajasegaran, Jathushan and Kanazawa, Angjoo and Malik, Jitendra}, + booktitle={ICCV}, + year={2023} +} +``` diff --git a/humans4d/assets/teaser.png b/humans4d/assets/teaser.png new file mode 100644 index 0000000..ba69081 Binary files /dev/null and b/humans4d/assets/teaser.png differ diff --git a/humans4d/demo.py b/humans4d/demo.py new file mode 100644 index 0000000..f3a2251 --- /dev/null +++ b/humans4d/demo.py @@ -0,0 +1,171 @@ +from pathlib import Path +import torch +import argparse +import os +import cv2 +import numpy as np + +from hmr2.configs import CACHE_DIR_4DHUMANS +from hmr2.models import HMR2, download_models, load_hmr2, DEFAULT_CHECKPOINT +from hmr2.utils import recursive_to +from hmr2.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD +from hmr2.utils.renderer import Renderer, cam_crop_to_full + +LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353) + +def main(): + import time + start = time.time() + parser = argparse.ArgumentParser(description='HMR2 demo code') + parser.add_argument('--checkpoint', type=str, default=DEFAULT_CHECKPOINT, help='Path to pretrained model checkpoint') + parser.add_argument('--img_folder', type=str, default='example_data/images', help='Folder with input images') + parser.add_argument('--out_folder', type=str, default='demo_out', help='Output folder to save rendered results') + parser.add_argument('--side_view', dest='side_view', action='store_true', default=False, help='If set, render side view also') + parser.add_argument('--top_view', dest='top_view', action='store_true', default=False, help='If set, render top view also') + parser.add_argument('--full_frame', dest='full_frame', action='store_true', default=False, help='If set, render all people together also') + parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False, help='If set, save meshes to disk also') + parser.add_argument('--detector', type=str, default='vitdet', choices=['vitdet', 'regnety'], help='Using regnety improves runtime') + parser.add_argument('--batch_size', type=int, default=1, help='Batch size for inference/fitting') + parser.add_argument('--file_type', nargs='+', default=['*.jpg', '*.png'], help='List of file extensions to consider') + + args = parser.parse_args() + + # Download and load checkpoints + download_models(CACHE_DIR_4DHUMANS) + model, model_cfg = load_hmr2(args.checkpoint) + + # Setup HMR2.0 model + device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') + model = model.to(device) + model.eval() + + # Load detector + from hmr2.utils.utils_detectron2 import DefaultPredictor_Lazy + if args.detector == 'vitdet': + from detectron2.config import LazyConfig + import hmr2 + cfg_path = Path(hmr2.__file__).parent/'configs'/'cascade_mask_rcnn_vitdet_h_75ep.py' + detectron2_cfg = LazyConfig.load(str(cfg_path)) + detectron2_cfg.train.init_checkpoint = "https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/cascade_mask_rcnn_vitdet_h/f328730692/model_final_f05665.pkl" + for i in range(3): + detectron2_cfg.model.roi_heads.box_predictors[i].test_score_thresh = 0.25 + detector = DefaultPredictor_Lazy(detectron2_cfg) + elif args.detector == 'regnety': + from detectron2 import model_zoo + from detectron2.config import get_cfg + detectron2_cfg = model_zoo.get_config('new_baselines/mask_rcnn_regnety_4gf_dds_FPN_400ep_LSJ.py', trained=True) + detectron2_cfg.model.roi_heads.box_predictor.test_score_thresh = 0.5 + detectron2_cfg.model.roi_heads.box_predictor.test_nms_thresh = 0.4 + detector = DefaultPredictor_Lazy(detectron2_cfg) + + # Setup the renderer + renderer = Renderer(model_cfg, faces=model.smpl.faces) + + # Make output directory if it does not exist + os.makedirs(args.out_folder, exist_ok=True) + + # Get all demo images that end with .jpg or .png + img_paths = [img for end in args.file_type for img in Path(args.img_folder).glob(end)] + + # Iterate over all images in folder + for img_path in img_paths: + img_cv2 = cv2.imread(str(img_path)) + + # Detect humans in image + det_out = detector(img_cv2) + + det_instances = det_out['instances'] + valid_idx = (det_instances.pred_classes==0) & (det_instances.scores > 0.5) + boxes=det_instances.pred_boxes.tensor[valid_idx].cpu().numpy() + + # Run HMR2.0 on all detected humans + dataset = ViTDetDataset(model_cfg, img_cv2, boxes) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=False, num_workers=0) + + all_verts = [] + all_cam_t = [] + + for batch in dataloader: + batch = recursive_to(batch, device) + with torch.no_grad(): + out = model(batch) + + pred_cam = out['pred_cam'] + 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 + batch_size = batch['img'].shape[0] + for n in range(batch_size): + # Get filename from path img_path + img_fn, _ = os.path.splitext(os.path.basename(img_path)) + person_id = int(batch['personid'][n]) + white_img = (torch.ones_like(batch['img'][n]).cpu() - DEFAULT_MEAN[:,None,None]/255) / (DEFAULT_STD[:,None,None]/255) + input_patch = batch['img'][n].cpu() * (DEFAULT_STD[:,None,None]/255) + (DEFAULT_MEAN[:,None,None]/255) + input_patch = input_patch.permute(1,2,0).numpy() + + regression_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(), + out['pred_cam_t'][n].detach().cpu().numpy(), + batch['img'][n], + mesh_base_color=LIGHT_BLUE, + scene_bg_color=(1, 1, 1), + ) + + final_img = np.concatenate([input_patch, regression_img], axis=1) + + if args.side_view: + side_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(), + out['pred_cam_t'][n].detach().cpu().numpy(), + white_img, + mesh_base_color=LIGHT_BLUE, + scene_bg_color=(1, 1, 1), + side_view=True) + final_img = np.concatenate([final_img, side_img], axis=1) + + if args.top_view: + top_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(), + out['pred_cam_t'][n].detach().cpu().numpy(), + white_img, + mesh_base_color=LIGHT_BLUE, + scene_bg_color=(1, 1, 1), + top_view=True) + final_img = np.concatenate([final_img, top_img], axis=1) + + cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_{person_id}.png'), 255*final_img[:, :, ::-1]) + + # Add all verts and cams to list + verts = out['pred_vertices'][n].detach().cpu().numpy() + cam_t = pred_cam_t_full[n] + all_verts.append(verts) + all_cam_t.append(cam_t) + + # Save all meshes to disk + if args.save_mesh: + camera_translation = cam_t.copy() + tmesh = renderer.vertices_to_trimesh(verts, camera_translation, LIGHT_BLUE) + tmesh.export(os.path.join(args.out_folder, f'{img_fn}_{person_id}.obj')) + + # Render front view + if args.full_frame and len(all_verts) > 0: + misc_args = dict( + mesh_base_color=LIGHT_BLUE, + 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], **misc_args) + + # Overlay image + input_img = img_cv2.astype(np.float32)[:,:,::-1]/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:] + + cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_all.png'), 255*input_img_overlay[:, :, ::-1]) + + end = time.time() + print(end - start) + +if __name__ == '__main__': + main() diff --git a/humans4d/environment.yml b/humans4d/environment.yml new file mode 100644 index 0000000..e25626b --- /dev/null +++ b/humans4d/environment.yml @@ -0,0 +1,33 @@ +name: 4D-humans +channels: + - pytorch + - nvidia + - conda-forge +dependencies: + - python=3.10 + - numpy + - pytorch + - pytorch-cuda=11.8 + - torchvision + - pip + - pip: + - pytorch-lightning + - smplx==0.1.28 + - pyrender + - opencv-python + - yacs + - scikit-image + - einops + - timm + - dill + - pandas + - git+https://github.com/facebookresearch/detectron2 + - git+https://github.com/mattloper/chumpy + + # For training + - rich + - hydra-core + - hydra-submitit-launcher + - hydra-colorlog + - pyrootutils + - webdataset diff --git a/humans4d/eval.py b/humans4d/eval.py new file mode 100644 index 0000000..a4a54c5 --- /dev/null +++ b/humans4d/eval.py @@ -0,0 +1,131 @@ +import argparse +import os +from pathlib import Path +import traceback +from typing import Optional + +import pandas as pd +import torch +from filelock import FileLock +from hmr2.configs import dataset_eval_config +from hmr2.datasets import create_dataset +from hmr2.utils import Evaluator, recursive_to +from tqdm import tqdm + +from hmr2.configs import CACHE_DIR_4DHUMANS +from hmr2.models import HMR2, download_models, load_hmr2, DEFAULT_CHECKPOINT + +def main(): + parser = argparse.ArgumentParser(description='Evaluate trained models') + parser.add_argument('--checkpoint', type=str, default=DEFAULT_CHECKPOINT, help='Path to pretrained model checkpoint') + parser.add_argument('--results_file', type=str, default='results/eval_regression.csv', help='Path to results file.') + parser.add_argument('--dataset', type=str, default='H36M-VAL-P2,3DPW-TEST,LSP-EXTENDED,POSETRACK-VAL,COCO-VAL', help='Dataset to evaluate') # choices=['H36M-VAL-P2', '3DPW-TEST', 'MPI-INF-TEST'] + parser.add_argument('--batch_size', type=int, default=16, help='Batch size for inference') + parser.add_argument('--num_samples', type=int, default=1, help='Number of test samples to draw') + parser.add_argument('--num_workers', type=int, default=8, help='Number of workers used for data loading') + parser.add_argument('--log_freq', type=int, default=10, help='How often to log results') + parser.add_argument('--shuffle', dest='shuffle', action='store_true', default=False, help='Shuffle the dataset during evaluation') + parser.add_argument('--exp_name', type=str, default=None, help='Experiment name') + + args = parser.parse_args() + + # Download and load checkpoints + download_models(CACHE_DIR_4DHUMANS) + model, model_cfg = load_hmr2(args.checkpoint) + + # Setup HMR2.0 model + device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') + model = model.to(device) + model.eval() + + # Load config and run eval, one dataset at a time + print('Evaluating on datasets: {}'.format(args.dataset), flush=True) + for dataset in args.dataset.split(','): + dataset_cfg = dataset_eval_config()[dataset] + args.dataset = dataset + run_eval(model, model_cfg, dataset_cfg, device, args) + +def run_eval(model, model_cfg, dataset_cfg, device, args): + # Create dataset and data loader + dataset = create_dataset(model_cfg, dataset_cfg, train=False) + dataloader = torch.utils.data.DataLoader(dataset, args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers) + + # List of metrics to log + if args.dataset in ['H36M-VAL-P2','3DPW-TEST']: + metrics = ['mode_re', 'mode_mpjpe'] + pck_thresholds = None + if args.dataset in ['LSP-EXTENDED', 'POSETRACK-VAL', 'COCO-VAL']: + metrics = ['mode_kpl2'] + pck_thresholds = [0.05, 0.1] + + # Setup evaluator object + evaluator = Evaluator( + dataset_length=int(1e8), + keypoint_list=dataset_cfg.KEYPOINT_LIST, + pelvis_ind=model_cfg.EXTRA.PELVIS_IND, + metrics=metrics, + pck_thresholds=pck_thresholds, + ) + + # Go over the images in the dataset. + try: + for i, batch in enumerate(tqdm(dataloader)): + batch = recursive_to(batch, device) + with torch.no_grad(): + out = model(batch) + evaluator(out, batch) + if i % args.log_freq == args.log_freq - 1: + evaluator.log() + evaluator.log() + error = None + except (Exception, KeyboardInterrupt) as e: + traceback.print_exc() + error = repr(e) + i = 0 + + # Append results to file + metrics_dict = evaluator.get_metrics_dict() + save_eval_result(args.results_file, metrics_dict, args.checkpoint, args.dataset, error=error, iters_done=i, exp_name=args.exp_name) + + +def save_eval_result( + csv_path: str, + metric_dict: float, + checkpoint_path: str, + dataset_name: str, + # start_time: pd.Timestamp, + error: Optional[str] = None, + iters_done=None, + exp_name=None, +) -> None: + """Save evaluation results for a single scene file to a common CSV file.""" + + timestamp = pd.Timestamp.now() + exists: bool = os.path.exists(csv_path) + exp_name = exp_name or Path(checkpoint_path).parent.parent.name + + # save each metric as different row to the csv path + metric_names = list(metric_dict.keys()) + metric_values = list(metric_dict.values()) + N = len(metric_names) + df = pd.DataFrame( + dict( + timestamp=[timestamp] * N, + checkpoint_path=[checkpoint_path] * N, + exp_name=[exp_name] * N, + dataset=[dataset_name] * N, + metric_name=metric_names, + metric_value=metric_values, + error=[error] * N, + iters_done=[iters_done] * N, + ), + index=list(range(N)), + ) + + # Lock the file to prevent multiple processes from writing to it at the same time. + lock = FileLock(f"{csv_path}.lock", timeout=10) + with lock: + df.to_csv(csv_path, mode="a", header=not exists, index=False) + +if __name__ == '__main__': + main() diff --git a/humans4d/example_data/images/pexels-anete-lusina-4793258.jpg b/humans4d/example_data/images/pexels-anete-lusina-4793258.jpg new file mode 100644 index 0000000..a7c19ac Binary files /dev/null and b/humans4d/example_data/images/pexels-anete-lusina-4793258.jpg differ diff --git a/humans4d/example_data/images/skates.png b/humans4d/example_data/images/skates.png new file mode 100644 index 0000000..bbb970a Binary files /dev/null and b/humans4d/example_data/images/skates.png differ diff --git a/humans4d/example_data/videos/gymnasts.mp4 b/humans4d/example_data/videos/gymnasts.mp4 new file mode 100644 index 0000000..c08d312 Binary files /dev/null and b/humans4d/example_data/videos/gymnasts.mp4 differ diff --git a/humans4d/fetch_training_data.sh b/humans4d/fetch_training_data.sh new file mode 100644 index 0000000..d167ae5 --- /dev/null +++ b/humans4d/fetch_training_data.sh @@ -0,0 +1,21 @@ +# Google drive links to download the training data +gdown https://drive.google.com/uc?id=19e5w1x8JIJkYUT9TIPcFhQF1SXny_Qwz # hmr2_training_data_part1.tar +gdown https://drive.google.com/uc?id=10glv5LKLhqPHHpaWmvkg4wT3wSmL6E2f # hmr2_training_data_part2.tar +gdown https://drive.google.com/uc?id=19AJMbE0nKD6Lao22aAaPLngiuvrnxF5V # hmr2_training_data_part3.tar +gdown https://drive.google.com/uc?id=1ntyjT3-JbtDxoHqmCJzJ8JemMKxM78Xj # hmr2_training_data_part4a.tar +gdown https://drive.google.com/uc?id=1H1xBQqePHcayR65oRn3cnxdgau_WLrLD # hmr2_training_data_part4b.tar +gdown https://drive.google.com/uc?id=1TnYqqG-wQMfMjd0EHbMEqxDt8DXG3vbW # hmr2_training_data_part4c.tar +gdown https://drive.google.com/uc?id=1tY6N8RDtgWfHSnTLgUNfbGhPI5X9fMg9 # hmr2_training_data_part4d.tar + +# Alternatively, consider using the dropbox links: +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AACJetGWdtee9TP_ejmBMdKOa/hmr2_training_data_part1.tar +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AAA7vh5e4jubZ3cYRyOtYYQZa/hmr2_training_data_part2.tar +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AACtCMwexHgDCMxh5CYP1VEya/hmr2_training_data_part3.tar +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AAB7C9Y33MEX5G9_UCHQ2KyLa/hmr2_training_data_part4a.tar +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AABkok2N00PHFGHiKV8-RjpNa/hmr2_training_data_part4b.tar +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AADYg_LKlGbjpjvrOQWoOgrna/hmr2_training_data_part4c.tar +#wget https://www.dropbox.com/sh/4y90ay6a6nki9v7/AACARYZJMoxq9nS6cDMAL8Pya/hmr2_training_data_part4d.tar + +for f in hmr2_training_data_part*.tar; do + tar --warning=no-unknown-keyword --exclude=".*" -xvf $f +done diff --git a/humans4d/gradio_app.py b/humans4d/gradio_app.py new file mode 100644 index 0000000..2746363 --- /dev/null +++ b/humans4d/gradio_app.py @@ -0,0 +1,187 @@ +from pathlib import Path +import tempfile +import torch +import argparse +import os +import numpy as np +import gradio as gr +from PIL import Image + +from hmr2.configs import CACHE_DIR_4DHUMANS +from hmr2.models import HMR2, download_models, load_hmr2, DEFAULT_CHECKPOINT +from hmr2.utils import recursive_to +from hmr2.datasets.vitdet_dataset import ViTDetDataset, DEFAULT_MEAN, DEFAULT_STD +from hmr2.utils.renderer import Renderer, cam_crop_to_full + +LIGHT_BLUE=(0.65098039, 0.74117647, 0.85882353) + +def main(): + parser = argparse.ArgumentParser(description='HMR2 demo code') + parser.add_argument('--checkpoint', type=str, default=DEFAULT_CHECKPOINT, help='Path to pretrained model checkpoint') + # parser.add_argument('--img_folder', type=str, default='example_data/images', help='Folder with input images') + parser.add_argument('--out_folder', type=str, default='demo_out', help='Output folder to save rendered results') + # parser.add_argument('--side_view', dest='side_view', action='store_true', default=False, help='If set, render side view also') + # parser.add_argument('--full_frame', dest='full_frame', action='store_true', default=False, help='If set, render all people together also') + # parser.add_argument('--save_mesh', dest='save_mesh', action='store_true', default=False, help='If set, save meshes to disk also') + parser.add_argument('--batch_size', type=int, default=64, help='Batch size for inference/fitting') + + args = parser.parse_args() + + # Download and load checkpoints + download_models(CACHE_DIR_4DHUMANS) + model, model_cfg = load_hmr2(args.checkpoint) + + # Setup HMR2.0 model + device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') + model = model.to(device) + model.eval() + + # Load detector + from hmr2.utils.utils_detectron2 import DefaultPredictor_Lazy + from detectron2.config import LazyConfig + import hmr2 + cfg_path = Path(hmr2.__file__).parent/'configs'/'cascade_mask_rcnn_vitdet_h_75ep.py' + detectron2_cfg = LazyConfig.load(str(cfg_path)) + detectron2_cfg.train.init_checkpoint = "https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/cascade_mask_rcnn_vitdet_h/f328730692/model_final_f05665.pkl" + for i in range(3): + detectron2_cfg.model.roi_heads.box_predictors[i].test_score_thresh = 0.25 + detector = DefaultPredictor_Lazy(detectron2_cfg) + + # Setup the renderer + renderer = Renderer(model_cfg, faces=model.smpl.faces) + + # Make output directory if it does not exist + os.makedirs(args.out_folder, exist_ok=True) + + # # Iterate over all images in folder + # for img_path in Path(args.img_folder).glob('*.jpg'): + # img_cv2 = cv2.imread(str(img_path)) + def infer(in_pil_img, in_threshold=0.8, out_pil_img=None): + # Convert RGB to BGR + img_cv2 = np.array(in_pil_img)[:, :, ::-1].copy() + + # Detect humans in image + det_out = detector(img_cv2) + + det_instances = det_out['instances'] + valid_idx = (det_instances.pred_classes==0) & (det_instances.scores > 0.5) + boxes=det_instances.pred_boxes.tensor[valid_idx].cpu().numpy() + + # Run HMR2.0 on all detected humans + dataset = ViTDetDataset(model_cfg, img_cv2, boxes) + dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=False, num_workers=0) + + all_verts = [] + all_cam_t = [] + all_mesh_paths = [] + + temp_name = next(tempfile._get_candidate_names()) + + for batch in dataloader: + batch = recursive_to(batch, device) + with torch.no_grad(): + out = model(batch) + + pred_cam = out['pred_cam'] + 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 + batch_size = batch['img'].shape[0] + for n in range(batch_size): + # Get filename from path img_path + # img_fn, _ = os.path.splitext(os.path.basename(img_path)) + person_id = int(batch['personid'][n]) + white_img = (torch.ones_like(batch['img'][n]).cpu() - DEFAULT_MEAN[:,None,None]/255) / (DEFAULT_STD[:,None,None]/255) + input_patch = batch['img'][n].cpu() * (DEFAULT_STD[:,None,None]/255) + (DEFAULT_MEAN[:,None,None]/255) + input_patch = input_patch.permute(1,2,0).numpy() + + # regression_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(), + # out['pred_cam_t'][n].detach().cpu().numpy(), + # batch['img'][n], + # mesh_base_color=LIGHT_BLUE, + # scene_bg_color=(1, 1, 1), + # ) + + # if args.side_view: + # side_img = renderer(out['pred_vertices'][n].detach().cpu().numpy(), + # out['pred_cam_t'][n].detach().cpu().numpy(), + # white_img, + # mesh_base_color=LIGHT_BLUE, + # scene_bg_color=(1, 1, 1), + # side_view=True) + # final_img = np.concatenate([input_patch, regression_img, side_img], axis=1) + # else: + # final_img = np.concatenate([input_patch, regression_img], axis=1) + + # cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_{person_id}.png'), 255*final_img[:, :, ::-1]) + + # Add all verts and cams to list + verts = out['pred_vertices'][n].detach().cpu().numpy() + cam_t = pred_cam_t_full[n] + all_verts.append(verts) + all_cam_t.append(cam_t) + + # Save all meshes to disk + # if args.save_mesh: + if True: + camera_translation = cam_t.copy() + tmesh = renderer.vertices_to_trimesh(verts, camera_translation, LIGHT_BLUE) + + temp_path = os.path.join(args.out_folder, f'{temp_name}_{person_id}.obj') + tmesh.export(temp_path) + all_mesh_paths.append(temp_path) + + # Render front view + if len(all_verts) > 0: + misc_args = dict( + mesh_base_color=LIGHT_BLUE, + 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], **misc_args) + + # Overlay image + input_img = img_cv2.astype(np.float32)[:,:,::-1]/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:] + + # cv2.imwrite(os.path.join(args.out_folder, f'{img_fn}_all.png'), 255*input_img_overlay[:, :, ::-1]) + + # convert to PIL image + out_pil_img = Image.fromarray((input_img_overlay*255).astype(np.uint8)) + return out_pil_img, all_mesh_paths + else: + return None, [] + + with gr.Blocks(title="4DHumans", css=".gradio-container") as demo: + + gr.HTML("""
HMR 2.0
""") + + with gr.Row(): + with gr.Column(): + input_image = gr.Image(label="Input image", type="pil") + with gr.Column(): + output_image = gr.Image(label="Reconstructions", type="pil") + output_meshes = gr.File(label="3D meshes") + + gr.HTML("""
""") + + with gr.Row(): + threshold = gr.Slider(0, 1.0, value=0.6, label='Detection Threshold') + send_btn = gr.Button("Infer") + send_btn.click(fn=infer, inputs=[input_image, threshold], outputs=[output_image, output_meshes]) + + gr.Examples([ + ['example_data/images/pexels-anete-lusina-4793258.jpg'], + ['example_data/images/skates.png'], + ], + inputs=[input_image, 0.6]) + + demo.launch(debug=True) + +if __name__ == '__main__': + main() diff --git a/humans4d/hmr2/__init__.py b/humans4d/hmr2/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/humans4d/hmr2/configs/__init__.py b/humans4d/hmr2/configs/__init__.py new file mode 100644 index 0000000..bc2d495 --- /dev/null +++ b/humans4d/hmr2/configs/__init__.py @@ -0,0 +1,115 @@ +import os +from typing import Dict +from yacs.config import CfgNode as CN + +CACHE_DIR = os.path.join(os.environ.get("HOME"), ".cache") +CACHE_DIR_4DHUMANS = os.path.join(CACHE_DIR, "4DHumans") + +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 = True +_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_4DHUMANS, path) + + # cfg.SMPL.MODEL_PATH = update_path(cfg.SMPL.MODEL_PATH) + # cfg.SMPL.JOINT_REGRESSOR_EXTRA = update_path(cfg.SMPL.JOINT_REGRESSOR_EXTRA) + # cfg.SMPL.MEAN_PARAMS = update_path(cfg.SMPL.MEAN_PARAMS) + + return cfg diff --git a/humans4d/hmr2/configs/cascade_mask_rcnn_vitdet_h_75ep.py b/humans4d/hmr2/configs/cascade_mask_rcnn_vitdet_h_75ep.py new file mode 100644 index 0000000..0c6ae0e --- /dev/null +++ b/humans4d/hmr2/configs/cascade_mask_rcnn_vitdet_h_75ep.py @@ -0,0 +1,129 @@ +## coco_loader_lsj.py + +import detectron2.data.transforms as T +from detectron2 import model_zoo +from detectron2.config import LazyCall as L + +# Data using LSJ +image_size = 1024 +dataloader = model_zoo.get_config("common/data/coco.py").dataloader +dataloader.train.mapper.augmentations = [ + L(T.RandomFlip)(horizontal=True), # flip first + L(T.ResizeScale)( + min_scale=0.1, max_scale=2.0, target_height=image_size, target_width=image_size + ), + L(T.FixedSizeCrop)(crop_size=(image_size, image_size), pad=False), +] +dataloader.train.mapper.image_format = "RGB" +dataloader.train.total_batch_size = 64 +# recompute boxes due to cropping +dataloader.train.mapper.recompute_boxes = True + +dataloader.test.mapper.augmentations = [ + L(T.ResizeShortestEdge)(short_edge_length=image_size, max_size=image_size), +] + +from functools import partial +from fvcore.common.param_scheduler import MultiStepParamScheduler + +from detectron2 import model_zoo +from detectron2.config import LazyCall as L +from detectron2.solver import WarmupParamScheduler +from detectron2.modeling.backbone.vit import get_vit_lr_decay_rate + +# mask_rcnn_vitdet_b_100ep.py + +model = model_zoo.get_config("common/models/mask_rcnn_vitdet.py").model + +# Initialization and trainer settings +train = model_zoo.get_config("common/train.py").train +train.amp.enabled = True +train.ddp.fp16_compression = True +train.init_checkpoint = "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_base.pth" + + +# Schedule +# 100 ep = 184375 iters * 64 images/iter / 118000 images/ep +train.max_iter = 184375 + +lr_multiplier = L(WarmupParamScheduler)( + scheduler=L(MultiStepParamScheduler)( + values=[1.0, 0.1, 0.01], + milestones=[163889, 177546], + num_updates=train.max_iter, + ), + warmup_length=250 / train.max_iter, + warmup_factor=0.001, +) + +# Optimizer +optimizer = model_zoo.get_config("common/optim.py").AdamW +optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, num_layers=12, lr_decay_rate=0.7) +optimizer.params.overrides = {"pos_embed": {"weight_decay": 0.0}} + +# cascade_mask_rcnn_vitdet_b_100ep.py + +from detectron2.config import LazyCall as L +from detectron2.layers import ShapeSpec +from detectron2.modeling.box_regression import Box2BoxTransform +from detectron2.modeling.matcher import Matcher +from detectron2.modeling.roi_heads import ( + FastRCNNOutputLayers, + FastRCNNConvFCHead, + CascadeROIHeads, +) + +# arguments that don't exist for Cascade R-CNN +[model.roi_heads.pop(k) for k in ["box_head", "box_predictor", "proposal_matcher"]] + +model.roi_heads.update( + _target_=CascadeROIHeads, + box_heads=[ + L(FastRCNNConvFCHead)( + input_shape=ShapeSpec(channels=256, height=7, width=7), + conv_dims=[256, 256, 256, 256], + fc_dims=[1024], + conv_norm="LN", + ) + for _ in range(3) + ], + box_predictors=[ + L(FastRCNNOutputLayers)( + input_shape=ShapeSpec(channels=1024), + test_score_thresh=0.05, + box2box_transform=L(Box2BoxTransform)(weights=(w1, w1, w2, w2)), + cls_agnostic_bbox_reg=True, + num_classes="${...num_classes}", + ) + for (w1, w2) in [(10, 5), (20, 10), (30, 15)] + ], + proposal_matchers=[ + L(Matcher)(thresholds=[th], labels=[0, 1], allow_low_quality_matches=False) + for th in [0.5, 0.6, 0.7] + ], +) + +# cascade_mask_rcnn_vitdet_h_75ep.py + +from functools import partial + +train.init_checkpoint = "detectron2://ImageNetPretrained/MAE/mae_pretrain_vit_huge_p14to16.pth" + +model.backbone.net.embed_dim = 1280 +model.backbone.net.depth = 32 +model.backbone.net.num_heads = 16 +model.backbone.net.drop_path_rate = 0.5 +# 7, 15, 23, 31 for global attention +model.backbone.net.window_block_indexes = ( + list(range(0, 7)) + list(range(8, 15)) + list(range(16, 23)) + list(range(24, 31)) +) + +optimizer.params.lr_factor_func = partial(get_vit_lr_decay_rate, lr_decay_rate=0.9, num_layers=32) +optimizer.params.overrides = {} +optimizer.params.weight_decay_norm = None + +train.max_iter = train.max_iter * 3 // 4 # 100ep -> 75ep +lr_multiplier.scheduler.milestones = [ + milestone * 3 // 4 for milestone in lr_multiplier.scheduler.milestones +] +lr_multiplier.scheduler.num_updates = train.max_iter diff --git a/humans4d/hmr2/configs/datasets_eval.yaml b/humans4d/hmr2/configs/datasets_eval.yaml new file mode 100644 index 0000000..b0c75d6 --- /dev/null +++ b/humans4d/hmr2/configs/datasets_eval.yaml @@ -0,0 +1,31 @@ +H36M-VAL-P2: + TYPE: ImageDataset + DATASET_FILE: hmr2_evaluation_data/h36m_val_p2.npz + IMG_DIR: /shared/pavlakos/datasets/h36m/images/ + KEYPOINT_LIST: [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 43] + USE_HIPS: True + +3DPW-TEST: + TYPE: ImageDataset + DATASET_FILE: hmr2_evaluation_data/3dpw_test.npz + IMG_DIR: /shared/pavlakos/datasets/3DPW/ + KEYPOINT_LIST: [25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 43] + USE_HIPS: False + +POSETRACK-VAL: + TYPE: ImageDataset + DATASET_FILE: hmr2_evaluation_data/posetrack_2018_val.npz + IMG_DIR: /shared/pavlakos/datasets/posetrack/posetrack2018/posetrack_data/ + KEYPOINT_LIST: [0] # Dummy + +LSP-EXTENDED: + TYPE: ImageDataset + DATASET_FILE: hmr2_evaluation_data/hr-lspet_train.npz + IMG_DIR: /shared/pavlakos/datasets/hr-lspet/ + KEYPOINT_LIST: [0] # Dummy + +COCO-VAL: + TYPE: ImageDataset + DATASET_FILE: hmr2_evaluation_data/coco_val.npz + IMG_DIR: /shared/pavlakos/datasets/coco/ + KEYPOINT_LIST: [0] # Dummy diff --git a/humans4d/hmr2/configs/datasets_tar.yaml b/humans4d/hmr2/configs/datasets_tar.yaml new file mode 100644 index 0000000..2bac331 --- /dev/null +++ b/humans4d/hmr2/configs/datasets_tar.yaml @@ -0,0 +1,38 @@ +MPI-INF-TRAIN-PRUNED: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/mpi-inf-train-pruned/{000000..00006}.tar + epoch_size: 12_000 +H36M-TRAIN-WMASK: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/h36m-train/{000000..000312}.tar + epoch_size: 314_000 # Can be changed arbitrarily. Enables resampling. +MPII-TRAIN-WMASK: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/mpii-train/{000000..000009}.tar + epoch_size: 100_000 # Increase to ensure dataset doesn't end before epoch +COCO-TRAIN-2014-WMASK-PRUNED: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/coco-train-2014-pruned/{000000..000017}.tar + epoch_size: 18_000 +AVA-TRAIN-MIDFRAMES-1FPS-WMASK: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/ava-train-midframes-1fps-vitpose/{000000..000092}.tar + epoch_size: 200_000 +AIC-TRAIN-WMASK: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/aic-train-vitpose/{000000..000104}.tar + epoch_size: 200_000 +INSTA-TRAIN-WMASK: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/insta-train-vitpose-replicate/{000000..003657}.tar + epoch_size: 4_000_000 +COCO-TRAIN-2014-VITPOSE-REPLICATE-PRUNED12: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/coco-train-2014-vitpose-pruned/{000000..000044}.tar + epoch_size: 45_000 +COCO-VAL: + TYPE: ImageDataset + URLS: hmr2_training_data/dataset_tars/coco-val/{000000..000000}.tar + KEYPOINT_LIST: [0] # Dummy +CMU-MOCAP: + DATASET_FILE: hmr2_training_data/cmu_mocap.npz diff --git a/humans4d/hmr2/configs_hydra/data/mix_all.yaml b/humans4d/hmr2/configs_hydra/data/mix_all.yaml new file mode 100644 index 0000000..ed11a88 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/data/mix_all.yaml @@ -0,0 +1,27 @@ +# @package _global_ +defaults: + - /data_filtering: low1 + +DATASETS: + TRAIN: + H36M-TRAIN-WMASK: + WEIGHT: 0.1 + MPII-TRAIN-WMASK: + WEIGHT: 0.1 + COCO-TRAIN-2014-WMASK-PRUNED: + WEIGHT: 0.1 + COCO-TRAIN-2014-VITPOSE-REPLICATE-PRUNED12: + WEIGHT: 0.1 + MPI-INF-TRAIN-PRUNED: + WEIGHT: 0.02 + AVA-TRAIN-MIDFRAMES-1FPS-WMASK: + WEIGHT: 0.19 + AIC-TRAIN-WMASK: + WEIGHT: 0.19 + INSTA-TRAIN-WMASK: + WEIGHT: 0.2 + VAL: + COCO-VAL: + WEIGHT: 1.0 + + MOCAP: CMU-MOCAP diff --git a/humans4d/hmr2/configs_hydra/data_filtering/low1.yaml b/humans4d/hmr2/configs_hydra/data_filtering/low1.yaml new file mode 100644 index 0000000..7d40adf --- /dev/null +++ b/humans4d/hmr2/configs_hydra/data_filtering/low1.yaml @@ -0,0 +1,13 @@ +# @package _global_ + +DATASETS: + # Data filtering during training + SUPPRESS_KP_CONF_THRESH: 0.3 + FILTER_NUM_KP: 4 + FILTER_NUM_KP_THRESH: 0.0 + FILTER_REPROJ_THRESH: 31000 + + SUPPRESS_BETAS_THRESH: 3.0 + SUPPRESS_BAD_POSES: True + POSES_BETAS_SIMULTANEOUS: True + FILTER_NO_POSES: False # If True, filters images that don't have poses diff --git a/humans4d/hmr2/configs_hydra/experiment/default.yaml b/humans4d/hmr2/configs_hydra/experiment/default.yaml new file mode 100644 index 0000000..2a0312c --- /dev/null +++ b/humans4d/hmr2/configs_hydra/experiment/default.yaml @@ -0,0 +1,29 @@ +# @package _global_ + +SMPL: + DATA_DIR: ${oc.env:HOME}/.cache/4DHumans/data/ + MODEL_PATH: ${SMPL.DATA_DIR}/smpl + GENDER: neutral + NUM_BODY_JOINTS: 23 + JOINT_REGRESSOR_EXTRA: ${SMPL.DATA_DIR}/SMPL_to_J19.pkl + MEAN_PARAMS: ${SMPL.DATA_DIR}/smpl_mean_params.npz + +EXTRA: + FOCAL_LENGTH: 5000 + NUM_LOG_IMAGES: 4 + NUM_LOG_SAMPLES_PER_IMAGE: 8 + PELVIS_IND: 39 + +DATASETS: + 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: True + FLIP_AUG_RATE: 0.5 + EXTREME_CROP_AUG_RATE: 0.10 + EXTREME_CROP_AUG_LEVEL: 1 diff --git a/humans4d/hmr2/configs_hydra/experiment/hmr_vit_transformer.yaml b/humans4d/hmr2/configs_hydra/experiment/hmr_vit_transformer.yaml new file mode 100644 index 0000000..75df38a --- /dev/null +++ b/humans4d/hmr2/configs_hydra/experiment/hmr_vit_transformer.yaml @@ -0,0 +1,51 @@ +# @package _global_ + +defaults: + - default.yaml + +GENERAL: + TOTAL_STEPS: 1_000_000 + LOG_STEPS: 1000 + VAL_STEPS: 1000 + CHECKPOINT_STEPS: 10000 + CHECKPOINT_SAVE_TOP_K: 1 + NUM_WORKERS: 6 + PREFETCH_FACTOR: 2 + +TRAIN: + LR: 1e-5 + WEIGHT_DECAY: 1e-4 + BATCH_SIZE: 48 + 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: hmr2_training_data/vitpose_backbone.pth + SMPL_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 # from vitpose-H + +LOSS_WEIGHTS: + KEYPOINTS_3D: 0.05 + KEYPOINTS_2D: 0.01 + GLOBAL_ORIENT: 0.001 + BODY_POSE: 0.001 + BETAS: 0.0005 + ADVERSARIAL: 0.0005 diff --git a/humans4d/hmr2/configs_hydra/extras/default.yaml b/humans4d/hmr2/configs_hydra/extras/default.yaml new file mode 100644 index 0000000..b9c6b62 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/extras/default.yaml @@ -0,0 +1,8 @@ +# disable python warnings if they annoy you +ignore_warnings: False + +# ask user for tags if none are provided in the config +enforce_tags: True + +# pretty print config tree at the start of the run using Rich library +print_config: True diff --git a/humans4d/hmr2/configs_hydra/hydra/default.yaml b/humans4d/hmr2/configs_hydra/hydra/default.yaml new file mode 100644 index 0000000..c30c188 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/hydra/default.yaml @@ -0,0 +1,26 @@ +# @package _global_ +# https://hydra.cc/docs/configure_hydra/intro/ + +# enable color logging +defaults: + - override /hydra/hydra_logging: colorlog + - override /hydra/job_logging: colorlog + +# exp_name: ovrd_${hydra:job.override_dirname} +exp_name: ${now:%Y-%m-%d}_${now:%H-%M-%S} + +hydra: + run: + dir: ${paths.log_dir}/${task_name}/runs/${exp_name} + sweep: + dir: ${paths.log_dir}/${task_name}/multiruns/${exp_name} + subdir: ${hydra.job.num} + job: + config: + override_dirname: + exclude_keys: + - trainer + - trainer.devices + - trainer.num_nodes + - callbacks + - debug diff --git a/humans4d/hmr2/configs_hydra/launcher/local.yaml b/humans4d/hmr2/configs_hydra/launcher/local.yaml new file mode 100644 index 0000000..7f27718 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/launcher/local.yaml @@ -0,0 +1,13 @@ +# @package _global_ + +defaults: + - override /hydra/launcher: submitit_local + +hydra: + launcher: + timeout_min: 10_080 # 7 days + nodes: 1 + tasks_per_node: ${trainer.devices} + cpus_per_task: 6 + gpus_per_node: ${trainer.devices} + name: hmr2 diff --git a/humans4d/hmr2/configs_hydra/launcher/slurm.yaml b/humans4d/hmr2/configs_hydra/launcher/slurm.yaml new file mode 100644 index 0000000..4537f5d --- /dev/null +++ b/humans4d/hmr2/configs_hydra/launcher/slurm.yaml @@ -0,0 +1,22 @@ +# @package _global_ + +defaults: + - override /hydra/launcher: submitit_slurm + +hydra: + launcher: + timeout_min: 10_080 # 7 days + max_num_timeout: 3 + partition: g40 + qos: idle + nodes: 1 + tasks_per_node: ${trainer.devices} + gpus_per_task: null + cpus_per_task: 12 + gpus_per_node: ${trainer.devices} + cpus_per_gpu: null + comment: laion + name: hmr2 + setup: + - module load cuda openmpi libfabric-aws + - export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 diff --git a/humans4d/hmr2/configs_hydra/paths/default.yaml b/humans4d/hmr2/configs_hydra/paths/default.yaml new file mode 100644 index 0000000..b2afd22 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/paths/default.yaml @@ -0,0 +1,18 @@ +# path to root directory +# this requires PROJECT_ROOT environment variable to exist +# PROJECT_ROOT is inferred and set by pyrootutils package in `train.py` and `eval.py` +root_dir: ${oc.env:PROJECT_ROOT} + +# path to data directory +data_dir: ${paths.root_dir}/data/ + +# path to logging directory +log_dir: logs/ + +# path to output directory, created dynamically by hydra +# path generation pattern is specified in `configs/hydra/default.yaml` +# use it to store all files generated during the run, like ckpts and metrics +output_dir: ${hydra:runtime.output_dir} + +# path to working directory +work_dir: ${hydra:runtime.cwd} diff --git a/humans4d/hmr2/configs_hydra/train.yaml b/humans4d/hmr2/configs_hydra/train.yaml new file mode 100644 index 0000000..5021b4c --- /dev/null +++ b/humans4d/hmr2/configs_hydra/train.yaml @@ -0,0 +1,47 @@ +# @package _global_ + +# specify here default configuration +# order of defaults determines the order in which configs override each other +defaults: + - _self_ + - data: mix_all.yaml + - trainer: ddp.yaml + - paths: default.yaml + - extras: default.yaml + - hydra: default.yaml + + # experiment configs allow for version control of specific hyperparameters + # e.g. best hyperparameters for given model and datamodule + - experiment: null + - texture_exp: null + + # optional local config for machine/user specific settings + # it's optional since it doesn't need to exist and is excluded from version control + - optional launcher: local.yaml + # - optional launcher: slurm.yaml + + # debugging config (enable through command line, e.g. `python train.py debug=default) + - debug: null + +# task name, determines output directory path +task_name: "train" + +# tags to help you identify your experiments +# you can overwrite this in experiment configs +# overwrite from command line with `python train.py tags="[first_tag, second_tag]"` +# appending lists from command line is currently not supported :( +# https://github.com/facebookresearch/hydra/issues/1547 +tags: ["dev"] + +# set False to skip model training +train: True + +# evaluate on test set, using best model weights achieved during training +# lightning chooses best weights based on the metric specified in checkpoint callback +test: False + +# simply provide checkpoint path to resume training +ckpt_path: null + +# seed for random number generators in pytorch, numpy and python.random +seed: null diff --git a/humans4d/hmr2/configs_hydra/trainer/cpu.yaml b/humans4d/hmr2/configs_hydra/trainer/cpu.yaml new file mode 100644 index 0000000..e9ab459 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/trainer/cpu.yaml @@ -0,0 +1,6 @@ +defaults: + - default.yaml + - default_hmr.yaml + +accelerator: cpu +devices: 1 diff --git a/humans4d/hmr2/configs_hydra/trainer/ddp.yaml b/humans4d/hmr2/configs_hydra/trainer/ddp.yaml new file mode 100644 index 0000000..d8f4271 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/trainer/ddp.yaml @@ -0,0 +1,14 @@ +defaults: + - default.yaml + - default_hmr.yaml + +# use "ddp_spawn" instead of "ddp", +# it's slower but normal "ddp" currently doesn't work ideally with hydra +# https://github.com/facebookresearch/hydra/issues/2070 +# https://pytorch-lightning.readthedocs.io/en/latest/accelerators/gpu_intermediate.html#distributed-data-parallel-spawn +strategy: ddp + +accelerator: gpu +devices: 8 +num_nodes: 1 +sync_batchnorm: True diff --git a/humans4d/hmr2/configs_hydra/trainer/default.yaml b/humans4d/hmr2/configs_hydra/trainer/default.yaml new file mode 100644 index 0000000..7d444f4 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/trainer/default.yaml @@ -0,0 +1,10 @@ +_target_: pytorch_lightning.Trainer + +default_root_dir: ${paths.output_dir} + +accelerator: cpu +devices: 1 + +# set True to to ensure deterministic results +# makes training slower but gives more reproducibility than just setting seeds +deterministic: False diff --git a/humans4d/hmr2/configs_hydra/trainer/default_hmr.yaml b/humans4d/hmr2/configs_hydra/trainer/default_hmr.yaml new file mode 100644 index 0000000..963b239 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/trainer/default_hmr.yaml @@ -0,0 +1,8 @@ +num_sanity_val_steps: 0 +log_every_n_steps: ${GENERAL.LOG_STEPS} +val_check_interval: ${GENERAL.VAL_STEPS} +precision: 16 +max_steps: ${GENERAL.TOTAL_STEPS} +# move_metrics_to_cpu: True +limit_val_batches: 1 +# track_grad_norm: -1 diff --git a/humans4d/hmr2/configs_hydra/trainer/gpu.yaml b/humans4d/hmr2/configs_hydra/trainer/gpu.yaml new file mode 100644 index 0000000..f817c61 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/trainer/gpu.yaml @@ -0,0 +1,6 @@ +defaults: + - default.yaml + - default_hmr.yaml + +accelerator: gpu +devices: 1 diff --git a/humans4d/hmr2/configs_hydra/trainer/mps.yaml b/humans4d/hmr2/configs_hydra/trainer/mps.yaml new file mode 100644 index 0000000..7083026 --- /dev/null +++ b/humans4d/hmr2/configs_hydra/trainer/mps.yaml @@ -0,0 +1,6 @@ +defaults: + - default.yaml + - default_hmr.yaml + +accelerator: mps +devices: 1 diff --git a/humans4d/hmr2/datasets/__init__.py b/humans4d/hmr2/datasets/__init__.py new file mode 100644 index 0000000..4fad37d --- /dev/null +++ b/humans4d/hmr2/datasets/__init__.py @@ -0,0 +1,88 @@ +from typing import Dict, Optional + +import torch +import numpy as np +import pytorch_lightning as pl +from yacs.config import CfgNode + +import webdataset as wds +from ..configs import to_lower +from .dataset import Dataset +from .image_dataset import ImageDataset +from .mocap_dataset import MoCapDataset + +def create_dataset(cfg: CfgNode, dataset_cfg: CfgNode, train: bool = True, **kwargs) -> Dataset: + """ + Instantiate a dataset from a config file. + Args: + cfg (CfgNode): Model configuration file. + dataset_cfg (CfgNode): Dataset configuration info. + train (bool): Variable to select between train and val datasets. + """ + + dataset_type = Dataset.registry[dataset_cfg.TYPE] + return dataset_type(cfg, **to_lower(dataset_cfg), train=train, **kwargs) + +def create_webdataset(cfg: CfgNode, dataset_cfg: CfgNode, train: bool = True) -> Dataset: + """ + Like `create_dataset` but load data from tars. + """ + dataset_type = Dataset.registry[dataset_cfg.TYPE] + return dataset_type.load_tars_as_webdataset(cfg, **to_lower(dataset_cfg), train=train) + + +class MixedWebDataset(wds.WebDataset): + def __init__(self, cfg: CfgNode, dataset_cfg: CfgNode, train: bool = True) -> None: + super(wds.WebDataset, self).__init__() + dataset_list = cfg.DATASETS.TRAIN if train else cfg.DATASETS.VAL + datasets = [create_webdataset(cfg, dataset_cfg[dataset], train=train) for dataset, v in dataset_list.items()] + weights = np.array([v.WEIGHT for dataset, v in dataset_list.items()]) + weights = weights / weights.sum() # normalize + self.append(wds.RandomMix(datasets, weights)) + +class HMR2DataModule(pl.LightningDataModule): + + def __init__(self, cfg: CfgNode, dataset_cfg: CfgNode) -> None: + """ + Initialize LightningDataModule for HMR2 training + Args: + cfg (CfgNode): Config file as a yacs CfgNode containing necessary dataset info. + dataset_cfg (CfgNode): Dataset configuration file + """ + super().__init__() + self.cfg = cfg + self.dataset_cfg = dataset_cfg + self.train_dataset = None + self.val_dataset = None + self.test_dataset = None + self.mocap_dataset = None + + def setup(self, stage: Optional[str] = None) -> None: + """ + Load datasets necessary for training + Args: + cfg (CfgNode): Config file as a yacs CfgNode containing necessary dataset info. + """ + if self.train_dataset == None: + self.train_dataset = MixedWebDataset(self.cfg, self.dataset_cfg, train=True).with_epoch(100_000).shuffle(4000) + self.val_dataset = MixedWebDataset(self.cfg, self.dataset_cfg, train=False).shuffle(4000) + self.mocap_dataset = MoCapDataset(**to_lower(self.dataset_cfg[self.cfg.DATASETS.MOCAP])) + + def train_dataloader(self) -> Dict: + """ + Setup training data loader. + Returns: + Dict: Dictionary containing image and mocap data dataloaders + """ + train_dataloader = torch.utils.data.DataLoader(self.train_dataset, self.cfg.TRAIN.BATCH_SIZE, drop_last=True, num_workers=self.cfg.GENERAL.NUM_WORKERS, prefetch_factor=self.cfg.GENERAL.PREFETCH_FACTOR) + mocap_dataloader = torch.utils.data.DataLoader(self.mocap_dataset, self.cfg.TRAIN.NUM_TRAIN_SAMPLES * self.cfg.TRAIN.BATCH_SIZE, shuffle=True, drop_last=True, num_workers=1) + return {'img': train_dataloader, 'mocap': mocap_dataloader} + + def val_dataloader(self) -> torch.utils.data.DataLoader: + """ + Setup val data loader. + Returns: + torch.utils.data.DataLoader: Validation dataloader + """ + val_dataloader = torch.utils.data.DataLoader(self.val_dataset, self.cfg.TRAIN.BATCH_SIZE, drop_last=True, num_workers=self.cfg.GENERAL.NUM_WORKERS) + return val_dataloader diff --git a/humans4d/hmr2/datasets/dataset.py b/humans4d/hmr2/datasets/dataset.py new file mode 100644 index 0000000..22fc5bc --- /dev/null +++ b/humans4d/hmr2/datasets/dataset.py @@ -0,0 +1,27 @@ +""" +This file contains the defition of the base Dataset class. +""" + +class DatasetRegistration(type): + """ + Metaclass for registering different datasets + """ + def __init__(cls, name, bases, nmspc): + super().__init__(name, bases, nmspc) + if not hasattr(cls, 'registry'): + cls.registry = dict() + cls.registry[name] = cls + + # Metamethods, called on class objects: + def __iter__(cls): + return iter(cls.registry) + + def __str__(cls): + return str(cls.registry) + +class Dataset(metaclass=DatasetRegistration): + """ + Base Dataset class + """ + def __init__(self, *args, **kwargs): + pass \ No newline at end of file diff --git a/humans4d/hmr2/datasets/image_dataset.py b/humans4d/hmr2/datasets/image_dataset.py new file mode 100644 index 0000000..8af8a88 --- /dev/null +++ b/humans4d/hmr2/datasets/image_dataset.py @@ -0,0 +1,454 @@ +import copy +import os +import numpy as np +import torch +from typing import Any, Dict, List +from yacs.config import CfgNode +import braceexpand +import cv2 + +from .dataset import Dataset +from .utils import get_example, expand_to_aspect_ratio +from .smplh_prob_filter import poses_check_probable, load_amass_hist_smooth + +def expand(s): + return os.path.expanduser(os.path.expandvars(s)) +def expand_urls(urls: str|List[str]): + if isinstance(urls, str): + urls = [urls] + urls = [u for url in urls for u in braceexpand.braceexpand(expand(url))] + return urls + +AIC_TRAIN_CORRUPT_KEYS = { + '0a047f0124ae48f8eee15a9506ce1449ee1ba669', + '1a703aa174450c02fbc9cfbf578a5435ef403689', + '0394e6dc4df78042929b891dbc24f0fd7ffb6b6d', + '5c032b9626e410441544c7669123ecc4ae077058', + 'ca018a7b4c5f53494006ebeeff9b4c0917a55f07', + '4a77adb695bef75a5d34c04d589baf646fe2ba35', + 'a0689017b1065c664daef4ae2d14ea03d543217e', + '39596a45cbd21bed4a5f9c2342505532f8ec5cbb', + '3d33283b40610d87db660b62982f797d50a7366b', +} +CORRUPT_KEYS = { + *{f'aic-train/{k}' for k in AIC_TRAIN_CORRUPT_KEYS}, + *{f'aic-train-vitpose/{k}' for k in AIC_TRAIN_CORRUPT_KEYS}, +} + +body_permutation = [0, 1, 5, 6, 7, 2, 3, 4, 8, 12, 13, 14, 9, 10, 11, 16, 15, 18, 17, 22, 23, 24, 19, 20, 21] +extra_permutation = [5, 4, 3, 2, 1, 0, 11, 10, 9, 8, 7, 6, 12, 13, 14, 15, 16, 17, 18] +FLIP_KEYPOINT_PERMUTATION = body_permutation + [25 + i for i in extra_permutation] + +DEFAULT_MEAN = 255. * np.array([0.485, 0.456, 0.406]) +DEFAULT_STD = 255. * np.array([0.229, 0.224, 0.225]) +DEFAULT_IMG_SIZE = 256 + +class ImageDataset(Dataset): + + def __init__(self, + cfg: CfgNode, + dataset_file: str, + img_dir: str, + train: bool = True, + prune: Dict[str, Any] = {}, + **kwargs): + """ + Dataset class used for loading images and corresponding annotations. + Args: + cfg (CfgNode): Model config file. + dataset_file (str): Path to npz file containing dataset info. + img_dir (str): Path to image folder. + train (bool): Whether it is for training or not (enables data augmentation). + """ + super(ImageDataset, self).__init__() + self.train = train + self.cfg = cfg + + 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) + + self.img_dir = img_dir + self.data = np.load(dataset_file, allow_pickle=True) + + self.imgname = self.data['imgname'] + self.personid = np.zeros(len(self.imgname), dtype=np.int32) + self.extra_info = self.data.get('extra_info', [{} for _ in range(len(self.imgname))]) + + self.flip_keypoint_permutation = copy.copy(FLIP_KEYPOINT_PERMUTATION) + + num_pose = 3 * (self.cfg.SMPL.NUM_BODY_JOINTS + 1) + + # Bounding boxes are assumed to be in the center and scale format + self.center = self.data['center'] + self.scale = self.data['scale'].reshape(len(self.center), -1) / 200.0 + if self.scale.shape[1] == 1: + self.scale = np.tile(self.scale, (1, 2)) + assert self.scale.shape == (len(self.center), 2) + + # Get gt SMPLX parameters, if available + try: + self.body_pose = self.data['body_pose'].astype(np.float32) + self.has_body_pose = self.data['has_body_pose'].astype(np.float32) + except KeyError: + self.body_pose = np.zeros((len(self.imgname), num_pose), dtype=np.float32) + self.has_body_pose = np.zeros(len(self.imgname), dtype=np.float32) + try: + self.betas = self.data['betas'].astype(np.float32) + self.has_betas = self.data['has_betas'].astype(np.float32) + except KeyError: + self.betas = np.zeros((len(self.imgname), 10), dtype=np.float32) + self.has_betas = np.zeros(len(self.imgname), dtype=np.float32) + + # Try to get 2d keypoints, if available + try: + body_keypoints_2d = self.data['body_keypoints_2d'] + except KeyError: + body_keypoints_2d = np.zeros((len(self.center), 25, 3)) + # Try to get extra 2d keypoints, if available + try: + extra_keypoints_2d = self.data['extra_keypoints_2d'] + except KeyError: + extra_keypoints_2d = np.zeros((len(self.center), 19, 3)) + + self.keypoints_2d = np.concatenate((body_keypoints_2d, extra_keypoints_2d), axis=1).astype(np.float32) + + # Try to get 3d keypoints, if available + try: + body_keypoints_3d = self.data['body_keypoints_3d'].astype(np.float32) + except KeyError: + body_keypoints_3d = np.zeros((len(self.center), 25, 4), dtype=np.float32) + # Try to get extra 3d keypoints, if available + try: + extra_keypoints_3d = self.data['extra_keypoints_3d'].astype(np.float32) + except KeyError: + extra_keypoints_3d = np.zeros((len(self.center), 19, 4), dtype=np.float32) + + body_keypoints_3d[:, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14], -1] = 0 + + self.keypoints_3d = np.concatenate((body_keypoints_3d, extra_keypoints_3d), axis=1).astype(np.float32) + + def __len__(self) -> int: + return len(self.scale) + + def __getitem__(self, idx: int) -> Dict: + """ + Returns an example from the dataset. + """ + try: + image_file_rel = self.imgname[idx].decode('utf-8') + except AttributeError: + image_file_rel = self.imgname[idx] + image_file = os.path.join(self.img_dir, image_file_rel) + keypoints_2d = self.keypoints_2d[idx].copy() + keypoints_3d = self.keypoints_3d[idx].copy() + + 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() + bbox_expand_factor = bbox_size / ((scale*200).max()) + body_pose = self.body_pose[idx].copy().astype(np.float32) + betas = self.betas[idx].copy().astype(np.float32) + + has_body_pose = self.has_body_pose[idx].copy() + has_betas = self.has_betas[idx].copy() + + smpl_params = {'global_orient': body_pose[:3], + 'body_pose': body_pose[3:], + 'betas': betas + } + + has_smpl_params = {'global_orient': has_body_pose, + 'body_pose': has_body_pose, + 'betas': has_betas + } + + smpl_params_is_axis_angle = {'global_orient': True, + 'body_pose': True, + 'betas': False + } + + augm_config = self.cfg.DATASETS.CONFIG + # Crop image and (possibly) perform data augmentation + img_patch, keypoints_2d, keypoints_3d, smpl_params, has_smpl_params, img_size = get_example(image_file, + center_x, center_y, + bbox_size, bbox_size, + keypoints_2d, keypoints_3d, + smpl_params, has_smpl_params, + self.flip_keypoint_permutation, + self.img_size, self.img_size, + self.mean, self.std, self.train, augm_config) + + item = {} + # These are the keypoints in the original image coordinates (before cropping) + orig_keypoints_2d = self.keypoints_2d[idx].copy() + + item['img'] = img_patch + item['keypoints_2d'] = keypoints_2d.astype(np.float32) + item['keypoints_3d'] = keypoints_3d.astype(np.float32) + item['orig_keypoints_2d'] = orig_keypoints_2d + item['box_center'] = self.center[idx].copy() + item['box_size'] = bbox_size + item['bbox_expand_factor'] = bbox_expand_factor + item['img_size'] = 1.0 * img_size[::-1].copy() + item['smpl_params'] = smpl_params + item['has_smpl_params'] = has_smpl_params + item['smpl_params_is_axis_angle'] = smpl_params_is_axis_angle + item['imgname'] = image_file + item['imgname_rel'] = image_file_rel + item['personid'] = int(self.personid[idx]) + item['extra_info'] = copy.deepcopy(self.extra_info[idx]) + item['idx'] = idx + item['_scale'] = scale + return item + + + @staticmethod + def load_tars_as_webdataset(cfg: CfgNode, urls: str|List[str], train: bool, + resampled=False, + epoch_size=None, + cache_dir=None, + **kwargs) -> Dataset: + """ + Loads the dataset from a webdataset tar file. + """ + + IMG_SIZE = cfg.MODEL.IMAGE_SIZE + BBOX_SHAPE = cfg.MODEL.get('BBOX_SHAPE', None) + MEAN = 255. * np.array(cfg.MODEL.IMAGE_MEAN) + STD = 255. * np.array(cfg.MODEL.IMAGE_STD) + + def split_data(source): + for item in source: + datas = item['data.pyd'] + for data in datas: + if 'detection.npz' in item: + det_idx = data['extra_info']['detection_npz_idx'] + mask = item['detection.npz']['masks'][det_idx] + else: + mask = np.ones_like(item['jpg'][:,:,0], dtype=bool) + yield { + '__key__': item['__key__'], + 'jpg': item['jpg'], + 'data.pyd': data, + 'mask': mask, + } + + def suppress_bad_kps(item, thresh=0.0): + if thresh > 0: + kp2d = item['data.pyd']['keypoints_2d'] + kp2d_conf = np.where(kp2d[:, 2] < thresh, 0.0, kp2d[:, 2]) + item['data.pyd']['keypoints_2d'] = np.concatenate([kp2d[:,:2], kp2d_conf[:,None]], axis=1) + return item + + def filter_numkp(item, numkp=4, thresh=0.0): + kp_conf = item['data.pyd']['keypoints_2d'][:, 2] + return (kp_conf > thresh).sum() > numkp + + def filter_reproj_error(item, thresh=10**4.5): + losses = item['data.pyd'].get('extra_info', {}).get('fitting_loss', np.array({})).item() + reproj_loss = losses.get('reprojection_loss', None) + return reproj_loss is None or reproj_loss < thresh + + def filter_bbox_size(item, thresh=1): + bbox_size_min = item['data.pyd']['scale'].min().item() * 200. + return bbox_size_min > thresh + + def filter_no_poses(item): + return (item['data.pyd']['has_body_pose'] > 0) + + def supress_bad_betas(item, thresh=3): + has_betas = item['data.pyd']['has_betas'] + if thresh > 0 and has_betas: + betas_abs = np.abs(item['data.pyd']['betas']) + if (betas_abs > thresh).any(): + item['data.pyd']['has_betas'] = False + return item + + amass_poses_hist100_smooth = load_amass_hist_smooth() + def supress_bad_poses(item): + has_body_pose = item['data.pyd']['has_body_pose'] + if has_body_pose: + body_pose = item['data.pyd']['body_pose'] + pose_is_probable = poses_check_probable(torch.from_numpy(body_pose)[None, 3:], amass_poses_hist100_smooth).item() + if not pose_is_probable: + item['data.pyd']['has_body_pose'] = False + return item + + def poses_betas_simultaneous(item): + # We either have both body_pose and betas, or neither + has_betas = item['data.pyd']['has_betas'] + has_body_pose = item['data.pyd']['has_body_pose'] + item['data.pyd']['has_betas'] = item['data.pyd']['has_body_pose'] = np.array(float((has_body_pose>0) and (has_betas>0))) + return item + + def set_betas_for_reg(item): + # Always have betas set to true + has_betas = item['data.pyd']['has_betas'] + betas = item['data.pyd']['betas'] + + if not (has_betas>0): + item['data.pyd']['has_betas'] = np.array(float((True))) + item['data.pyd']['betas'] = betas * 0 + return item + + # Load the dataset + if epoch_size is not None: + resampled = True + corrupt_filter = lambda sample: (sample['__key__'] not in CORRUPT_KEYS) + import webdataset as wds + dataset = wds.WebDataset(expand_urls(urls), + nodesplitter=wds.split_by_node, + shardshuffle=True, + resampled=resampled, + cache_dir=cache_dir, + ).select(corrupt_filter) + if train: + dataset = dataset.shuffle(100) + dataset = dataset.decode('rgb8').rename(jpg='jpg;jpeg;png') + + # Process the dataset + dataset = dataset.compose(split_data) + + # Filter/clean the dataset + SUPPRESS_KP_CONF_THRESH = cfg.DATASETS.get('SUPPRESS_KP_CONF_THRESH', 0.0) + SUPPRESS_BETAS_THRESH = cfg.DATASETS.get('SUPPRESS_BETAS_THRESH', 0.0) + SUPPRESS_BAD_POSES = cfg.DATASETS.get('SUPPRESS_BAD_POSES', False) + POSES_BETAS_SIMULTANEOUS = cfg.DATASETS.get('POSES_BETAS_SIMULTANEOUS', False) + BETAS_REG = cfg.DATASETS.get('BETAS_REG', False) + FILTER_NO_POSES = cfg.DATASETS.get('FILTER_NO_POSES', False) + FILTER_NUM_KP = cfg.DATASETS.get('FILTER_NUM_KP', 4) + FILTER_NUM_KP_THRESH = cfg.DATASETS.get('FILTER_NUM_KP_THRESH', 0.0) + FILTER_REPROJ_THRESH = cfg.DATASETS.get('FILTER_REPROJ_THRESH', 0.0) + FILTER_MIN_BBOX_SIZE = cfg.DATASETS.get('FILTER_MIN_BBOX_SIZE', 0.0) + if SUPPRESS_KP_CONF_THRESH > 0: + dataset = dataset.map(lambda x: suppress_bad_kps(x, thresh=SUPPRESS_KP_CONF_THRESH)) + if SUPPRESS_BETAS_THRESH > 0: + dataset = dataset.map(lambda x: supress_bad_betas(x, thresh=SUPPRESS_BETAS_THRESH)) + if SUPPRESS_BAD_POSES: + dataset = dataset.map(lambda x: supress_bad_poses(x)) + if POSES_BETAS_SIMULTANEOUS: + dataset = dataset.map(lambda x: poses_betas_simultaneous(x)) + if FILTER_NO_POSES: + dataset = dataset.select(lambda x: filter_no_poses(x)) + if FILTER_NUM_KP > 0: + dataset = dataset.select(lambda x: filter_numkp(x, numkp=FILTER_NUM_KP, thresh=FILTER_NUM_KP_THRESH)) + if FILTER_REPROJ_THRESH > 0: + dataset = dataset.select(lambda x: filter_reproj_error(x, thresh=FILTER_REPROJ_THRESH)) + if FILTER_MIN_BBOX_SIZE > 0: + dataset = dataset.select(lambda x: filter_bbox_size(x, thresh=FILTER_MIN_BBOX_SIZE)) + if BETAS_REG: + dataset = dataset.map(lambda x: set_betas_for_reg(x)) # NOTE: Must be at the end + + use_skimage_antialias = cfg.DATASETS.get('USE_SKIMAGE_ANTIALIAS', False) + border_mode = { + 'constant': cv2.BORDER_CONSTANT, + 'replicate': cv2.BORDER_REPLICATE, + }[cfg.DATASETS.get('BORDER_MODE', 'constant')] + + # Process the dataset further + dataset = dataset.map(lambda x: ImageDataset.process_webdataset_tar_item(x, train, + augm_config=cfg.DATASETS.CONFIG, + MEAN=MEAN, STD=STD, IMG_SIZE=IMG_SIZE, + BBOX_SHAPE=BBOX_SHAPE, + use_skimage_antialias=use_skimage_antialias, + border_mode=border_mode, + )) + if epoch_size is not None: + dataset = dataset.with_epoch(epoch_size) + + return dataset + + @staticmethod + def process_webdataset_tar_item(item, train, + augm_config=None, + MEAN=DEFAULT_MEAN, + STD=DEFAULT_STD, + IMG_SIZE=DEFAULT_IMG_SIZE, + BBOX_SHAPE=None, + use_skimage_antialias=False, + border_mode=cv2.BORDER_CONSTANT, + ): + # Read data from item + key = item['__key__'] + image = item['jpg'] + data = item['data.pyd'] + mask = item['mask'] + + keypoints_2d = data['keypoints_2d'] + keypoints_3d = data['keypoints_3d'] + center = data['center'] + scale = data['scale'] + body_pose = data['body_pose'] + betas = data['betas'] + has_body_pose = data['has_body_pose'] + has_betas = data['has_betas'] + # image_file = data['image_file'] + + # Process data + orig_keypoints_2d = keypoints_2d.copy() + center_x = center[0] + center_y = center[1] + bbox_size = expand_to_aspect_ratio(scale*200, target_aspect_ratio=BBOX_SHAPE).max() + if bbox_size < 1: + breakpoint() + + + smpl_params = {'global_orient': body_pose[:3], + 'body_pose': body_pose[3:], + 'betas': betas + } + + has_smpl_params = {'global_orient': has_body_pose, + 'body_pose': has_body_pose, + 'betas': has_betas + } + + smpl_params_is_axis_angle = {'global_orient': True, + 'body_pose': True, + 'betas': False + } + + augm_config = copy.deepcopy(augm_config) + # Crop image and (possibly) perform data augmentation + img_rgba = np.concatenate([image, mask.astype(np.uint8)[:,:,None]*255], axis=2) + img_patch_rgba, keypoints_2d, keypoints_3d, smpl_params, has_smpl_params, img_size, trans = get_example(img_rgba, + center_x, center_y, + bbox_size, bbox_size, + keypoints_2d, keypoints_3d, + smpl_params, has_smpl_params, + FLIP_KEYPOINT_PERMUTATION, + IMG_SIZE, IMG_SIZE, + MEAN, STD, train, augm_config, + is_bgr=False, return_trans=True, + use_skimage_antialias=use_skimage_antialias, + border_mode=border_mode, + ) + img_patch = img_patch_rgba[:3,:,:] + mask_patch = (img_patch_rgba[3,:,:] / 255.0).clip(0,1) + if (mask_patch < 0.5).all(): + mask_patch = np.ones_like(mask_patch) + + item = {} + + item['img'] = img_patch + item['mask'] = mask_patch + # item['img_og'] = image + # item['mask_og'] = mask + item['keypoints_2d'] = keypoints_2d.astype(np.float32) + item['keypoints_3d'] = keypoints_3d.astype(np.float32) + item['orig_keypoints_2d'] = orig_keypoints_2d + item['box_center'] = center.copy() + item['box_size'] = bbox_size + item['img_size'] = 1.0 * img_size[::-1].copy() + item['smpl_params'] = smpl_params + item['has_smpl_params'] = has_smpl_params + item['smpl_params_is_axis_angle'] = smpl_params_is_axis_angle + item['_scale'] = scale + item['_trans'] = trans + item['imgname'] = key + # item['idx'] = idx + return item diff --git a/humans4d/hmr2/datasets/mocap_dataset.py b/humans4d/hmr2/datasets/mocap_dataset.py new file mode 100644 index 0000000..7ed8c94 --- /dev/null +++ b/humans4d/hmr2/datasets/mocap_dataset.py @@ -0,0 +1,25 @@ +import numpy as np +from typing import Dict + +class MoCapDataset: + + def __init__(self, dataset_file: str): + """ + Dataset class used for loading a dataset of unpaired SMPL parameter annotations + Args: + cfg (CfgNode): Model config file. + dataset_file (str): Path to npz file containing dataset info. + """ + data = np.load(dataset_file) + self.pose = data['body_pose'].astype(np.float32)[:, 3:] + self.betas = data['betas'].astype(np.float32) + self.length = len(self.pose) + + def __getitem__(self, idx: int) -> Dict: + pose = self.pose[idx].copy() + betas = self.betas[idx].copy() + item = {'body_pose': pose, 'betas': betas} + return item + + def __len__(self) -> int: + return self.length diff --git a/humans4d/hmr2/datasets/preprocess/lspet_to_npz.py b/humans4d/hmr2/datasets/preprocess/lspet_to_npz.py new file mode 100644 index 0000000..228c3df --- /dev/null +++ b/humans4d/hmr2/datasets/preprocess/lspet_to_npz.py @@ -0,0 +1,74 @@ +# Adapted from https://raw.githubusercontent.com/nkolot/SPIN/master/datasets/preprocess/hr_lspet.py +import os +import glob +import numpy as np +import scipy.io as sio +# from .read_openpose import read_openpose + +def hr_lspet_extract(dataset_path, out_path): + + # training mode + png_path = os.path.join(dataset_path, '*.png') + imgs = glob.glob(png_path) + imgs.sort() + + # structs we use + imgnames_, scales_, centers_, parts_, openposes_= [], [], [], [], [] + + # scale factor + scaleFactor = 1.2 + + # annotation files + annot_file = os.path.join(dataset_path, 'joints.mat') + joints = sio.loadmat(annot_file)['joints'] + + # main loop + for i, imgname in enumerate(imgs): + # image name + imgname = imgname.split('/')[-1] + # read keypoints + part14 = joints[:,:2,i] + # scale and center + bbox = [min(part14[:,0]), min(part14[:,1]), + max(part14[:,0]), max(part14[:,1])] + center = [(bbox[2]+bbox[0])/2, (bbox[3]+bbox[1])/2] + # scale = scaleFactor*max(bbox[2]-bbox[0], bbox[3]-bbox[1]) # Don't /200 + scale = scaleFactor*np.array([bbox[2]-bbox[0], bbox[3]-bbox[1]]) # Don't /200 + # update keypoints + part = np.zeros([24,3]) + part[:14] = np.hstack([part14, np.ones([14,1])]) + + # # read openpose detections + # json_file = os.path.join(openpose_path, 'hrlspet', + # imgname.replace('.png', '_keypoints.json')) + # openpose = read_openpose(json_file, part, 'hrlspet') + + # store the data + imgnames_.append(imgname) + centers_.append(center) + scales_.append(scale) + parts_.append(part) + # openposes_.append(openpose) + + # Populate extra_keypoints_2d: N,19,3 + # extra_keypoints_2d[:14] = parts[:14] + extra_keypoints_2d = np.zeros((len(parts_), 19, 3)) + extra_keypoints_2d[:,:14,:] = np.stack(parts_)[:,:14,:3] + + print(f'{extra_keypoints_2d.shape=}') + + # store the data struct + if not os.path.isdir(out_path): + os.makedirs(out_path) + out_file = os.path.join(out_path, 'hr-lspet_train.npz') + np.savez(out_file, imgname=imgnames_, + center=centers_, + scale=scales_, + part=parts_, + extra_keypoints_2d=extra_keypoints_2d, + # openpose=openposes_ + ) + + +if __name__ == '__main__': + hr_lspet_extract('/shared/pavlakos/datasets/hr-lspet/', 'hmr2_evaluation_data/') diff --git a/humans4d/hmr2/datasets/preprocess/posetrack_to_npz.py b/humans4d/hmr2/datasets/preprocess/posetrack_to_npz.py new file mode 100644 index 0000000..7c999b5 --- /dev/null +++ b/humans4d/hmr2/datasets/preprocess/posetrack_to_npz.py @@ -0,0 +1,92 @@ +# Adapted from https://raw.githubusercontent.com/nkolot/SPIN/master/datasets/preprocess/coco.py +import os +from os.path import join +import sys +import json +import numpy as np +from pathlib import Path +# from .read_openpose import read_openpose + +def coco_extract(dataset_path, out_path): + + # # convert joints to global order + # joints_idx = [19, 20, 21, 22, 23, 9, 8, 10, 7, 11, 6, 3, 2, 4, 1, 5, 0] + + # bbox expansion factor + scaleFactor = 1.2 + + # structs we need + imgnames_, scales_, centers_, parts_, openposes_ = [], [], [], [], [] + + # json annotation file + SPLIT='val' + json_paths = (Path(dataset_path)/'posetrack_data/annotations'/SPLIT).glob('*.json') + for json_path in json_paths: + json_data = json.load(open(json_path, 'r')) + + imgs = {} + for img in json_data['images']: + imgs[img['id']] = img + + for annot in json_data['annotations']: + # keypoints processing + keypoints = annot['keypoints'] + keypoints = np.reshape(keypoints, (17,3)) + keypoints[keypoints[:,2]>0,2] = 1 + # check if all major body joints are annotated + if sum(keypoints[5:,2]>0) < 12: + continue + # image name + image_id = annot['image_id'] + img_name = str(imgs[image_id]['file_name']) + # img_name_full = f'images/{SPLIT}/{json_path.stem}/{img_name}' + img_name_full = img_name + + # keypoints + part = np.zeros([17,3]) + # part[joints_idx] = keypoints + part = keypoints + + # scale and center + bbox = annot['bbox'] + center = [bbox[0] + bbox[2]/2, bbox[1] + bbox[3]/2] + # scale = scaleFactor*max(bbox[2], bbox[3]) # Don't do /200 + scale = scaleFactor*np.array([bbox[2], bbox[3]]) # Don't /200 + # # read openpose detections + # json_file = os.path.join(openpose_path, 'coco', + # img_name.replace('.jpg', '_keypoints.json')) + # openpose = read_openpose(json_file, part, 'coco') + + # store data + imgnames_.append(img_name_full) + centers_.append(center) + scales_.append(scale) + parts_.append(part) + # openposes_.append(openpose) + + + # NOTE: Posetrack val doesn't annotate ears (17,18) + # But Posetrack does annotate head, neck so that wil have to live in extra_kps. + posetrack_to_op_extra = [0, 37, 38, 18, 17, 5, 2, 6, 3, 7, 4, 12, 9, 13, 10, 14, 11] # Will contain 15 keypoints. + all_keypoints_2d = np.zeros((len(parts_), 44, 3)) + all_keypoints_2d[:,posetrack_to_op_extra] = np.stack(parts_)[:,:len(posetrack_to_op_extra),:3] + body_keypoints_2d = all_keypoints_2d[:,:25,:] + extra_keypoints_2d = all_keypoints_2d[:,25:,:] + + print(f'{extra_keypoints_2d.shape=}') + + # store the data struct + if not os.path.isdir(out_path): + os.makedirs(out_path) + out_file = os.path.join(out_path, f'posetrack_2018_{SPLIT}.npz') + np.savez(out_file, imgname=imgnames_, + center=centers_, + scale=scales_, + part=parts_, + body_keypoints_2d=body_keypoints_2d, + extra_keypoints_2d=extra_keypoints_2d, + # openpose=openposes_ + ) + +if __name__ == '__main__': + coco_extract('/shared/pavlakos/datasets/posetrack/posetrack2018/', 'hmr2_evaluation_data/') diff --git a/humans4d/hmr2/datasets/smplh_prob_filter.py b/humans4d/hmr2/datasets/smplh_prob_filter.py new file mode 100644 index 0000000..e3d2c5c --- /dev/null +++ b/humans4d/hmr2/datasets/smplh_prob_filter.py @@ -0,0 +1,154 @@ +import os +import numpy as np +import torch +import torch.nn.functional as F + +JOINT_NAMES = [ + 'left_hip', + 'right_hip', + 'spine1', + 'left_knee', + 'right_knee', + 'spine2', + 'left_ankle', + 'right_ankle', + 'spine3', + 'left_foot', + 'right_foot', + 'neck', + 'left_collar', + 'right_collar', + 'head', + 'left_shoulder', + 'right_shoulder', + 'left_elbow', + 'right_elbow', + 'left_wrist', + 'right_wrist' +] + +# Manually chosen probability density thresholds for each joint +# Probablities computed using SIGMA=2 gaussian blur on AMASS pose 3D histogram for range (-pi,pi) with 100x100x100 bins +JOINT_NAME_PROB_THRESHOLDS = { + 'left_hip': 5e-5, + 'right_hip': 5e-5, + 'spine1': 2e-3, + 'left_knee': 5e-6, + 'right_knee': 5e-6, + 'spine2': 0.01, + 'left_ankle': 5e-6, + 'right_ankle': 5e-6, + 'spine3': 0.025, + 'left_foot': 0, + 'right_foot': 0, + 'neck': 2e-4, + 'left_collar': 4.5e-4 , + 'right_collar': 4.5e-4, + 'head': 5e-4, + 'left_shoulder': 2e-4, + 'right_shoulder': 2e-4, + 'left_elbow': 4e-5, + 'right_elbow': 4e-5, + 'left_wrist': 1e-3, + 'right_wrist': 1e-3, +} + +JOINT_IDX_PROB_THRESHOLDS = torch.tensor([JOINT_NAME_PROB_THRESHOLDS[joint_name] for joint_name in JOINT_NAMES]) + +################################################################### +POSE_RANGE_MIN = -np.pi +POSE_RANGE_MAX = np.pi + +# Create 21x100x100x100 histogram of all 21 AMASS body pose joints using `create_pose_hist(amass_poses, nbins=100)` +AMASS_HIST100_PATH = 'hmr2_training_data/amass_poses_hist100_SMPL+H_G.npy' +if not os.path.exists(AMASS_HIST100_PATH): + AMASS_HIST100_PATH = '/shared/shubham/code/hmr2023/amass_poses_hist100_SMPL+H_G.npy' +if not os.path.exists(AMASS_HIST100_PATH): + AMASS_HIST100_PATH = '/fsx/shubham/code/stable-humans/notebooks/amass_poses_hist100_SMPL+H_G.npy' + + +def create_pose_hist(poses: np.ndarray, nbins: int = 100) -> np.ndarray: + N,K,C = poses.shape + assert C==3, poses.shape + poses_21x3 = normalize_axis_angle(torch.fromnumpy(poses).view(N*K,3)).numpy().reshape(N, K, 3) + assert (poses_21x3 > -np.pi).all() and (poses_21x3 < np.pi).all() + + Hs, Es = [], [] + for i in range(K): + H, edges = np.histogramdd(poses_21x3[:, i, :], bins=nbins, range=[(-np.pi, np.pi)]*3) + Hs.append(H) + Es.append(edges) + Hs = np.stack(Hs, axis=0) + return Hs + +def load_amass_hist_smooth(sigma=2) -> torch.Tensor: + amass_poses_hist100 = np.load(AMASS_HIST100_PATH) + amass_poses_hist100 = torch.from_numpy(amass_poses_hist100) + assert amass_poses_hist100.shape == (21,100,100,100) + + nbins = amass_poses_hist100.shape[1] + amass_poses_hist100 = amass_poses_hist100/amass_poses_hist100.sum() / (2*np.pi/nbins)**3 + + # Gaussian filter on amass_poses_hist100 + from scipy.ndimage import gaussian_filter + amass_poses_hist100_smooth = gaussian_filter(amass_poses_hist100.numpy(), sigma=sigma, mode='constant') + amass_poses_hist100_smooth = torch.from_numpy(amass_poses_hist100_smooth) + return amass_poses_hist100_smooth + +# Normalize axis angle representation s.t. angle is in [-pi, pi] +def normalize_axis_angle(poses: torch.Tensor) -> torch.Tensor: + # poses: N, 3 + # print(f'normalize_axis_angle ...') + assert poses.shape[1] == 3, poses.shape + angle = poses.norm(dim=1) + axis = F.normalize(poses, p=2, dim=1, eps=1e-8) + + angle_fixed = angle.clone() + axis_fixed = axis.clone() + + eps = 1e-6 + ii = 0 + while True: + # print(f'normalize_axis_angle iter {ii}') + ii += 1 + angle_too_big = (angle_fixed > np.pi + eps) + if not angle_too_big.any(): + break + + angle_fixed[angle_too_big] -= 2 * np.pi + angle_too_small = (angle_fixed < -eps) + axis_fixed[angle_too_small] *= -1 + angle_fixed[angle_too_small] *= -1 + + return axis_fixed * angle_fixed[:,None] + +def poses_to_joint_probs(poses: torch.Tensor, amass_poses_100_smooth: torch.Tensor) -> torch.Tensor: + # poses: Nx69 + # amass_poses_100_smooth: 21xBINSxBINSxBINS + # returns: poses_prob: Nx21 + N=poses.shape[0] + assert poses.shape == (N,69) + poses = poses[:,:63].reshape(N*21,3) + + nbins = amass_poses_100_smooth.shape[1] + assert amass_poses_100_smooth.shape == (21,nbins,nbins,nbins) + + poses_bin = (poses - POSE_RANGE_MIN) / (POSE_RANGE_MAX - POSE_RANGE_MIN) * (nbins - 1e-6) + poses_bin = poses_bin.long().clip(0, nbins-1) + joint_id = torch.arange(21, device=poses.device).view(1,21).expand(N,21).reshape(N*21) + poses_prob = amass_poses_100_smooth[joint_id, poses_bin[:,0], poses_bin[:,1], poses_bin[:,2]] + + poses_bad = ((poses < POSE_RANGE_MIN) | (poses >= POSE_RANGE_MAX)).any(dim=1) + poses_prob[poses_bad] = 0 + + return poses_prob.view(N,21) + +def poses_check_probable( + poses: torch.Tensor, + amass_poses_100_smooth: torch.Tensor, + prob_thresholds: torch.Tensor = JOINT_IDX_PROB_THRESHOLDS + ) -> torch.Tensor: + N,C=poses.shape + poses_norm = normalize_axis_angle(poses.reshape(N*(C//3),3)).reshape(N,C) + poses_prob = poses_to_joint_probs(poses_norm, amass_poses_100_smooth) + return (poses_prob > prob_thresholds).all(dim=1) diff --git a/humans4d/hmr2/datasets/utils.py b/humans4d/hmr2/datasets/utils.py new file mode 100644 index 0000000..40365fb --- /dev/null +++ b/humans4d/hmr2/datasets/utils.py @@ -0,0 +1,1025 @@ +""" +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 expand_bbox_to_aspect_ratio(bbox, target_aspect_ratio=None): + # bbox: np.array: (N,4) detectron2 bbox format + # target_aspect_ratio: (width, height) + if target_aspect_ratio is None: + return bbox + + is_singleton = (bbox.ndim == 1) + if is_singleton: + bbox = bbox[None,:] + + if bbox.shape[0] > 0: + center = np.stack(((bbox[:,0] + bbox[:,2]) / 2, (bbox[:,1] + bbox[:,3]) / 2), axis=1) + scale_wh = np.stack((bbox[:,2] - bbox[:,0], bbox[:,3] - bbox[:,1]), axis=1) + scale_wh = np.stack([expand_to_aspect_ratio(wh, target_aspect_ratio) for wh in scale_wh], axis=0) + bbox = np.stack([ + center[:,0] - scale_wh[:,0] / 2, + center[:,1] - scale_wh[:,1] / 2, + center[:,0] + scale_wh[:,0] / 2, + center[:,1] + scale_wh[:,1] / 2, + ], axis=1) + + if is_singleton: + bbox = bbox[0,:] + + return bbox + +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(smpl_params: Dict, has_smpl_params: Dict) -> Tuple[Dict, Dict]: + """ + Flip SMPL parameters when flipping the image. + Args: + smpl_params (Dict): SMPL parameter annotations. + has_smpl_params (Dict): Whether SMPL annotations are valid. + Returns: + Dict, Dict: Flipped SMPL parameters and valid flags. + """ + global_orient = smpl_params['global_orient'].copy() + body_pose = smpl_params['body_pose'].copy() + betas = smpl_params['betas'].copy() + has_global_orient = has_smpl_params['global_orient'].copy() + has_body_pose = has_smpl_params['body_pose'].copy() + has_betas = has_smpl_params['betas'].copy() + + body_pose_permutation = [6, 7, 8, 3, 4, 5, 9, 10, 11, 15, 16, 17, 12, 13, + 14 ,18, 19, 20, 24, 25, 26, 21, 22, 23, 27, 28, 29, 33, + 34, 35, 30, 31, 32, 36, 37, 38, 42, 43, 44, 39, 40, 41, + 45, 46, 47, 51, 52, 53, 48, 49, 50, 57, 58, 59, 54, 55, + 56, 63, 64, 65, 60, 61, 62, 69, 70, 71, 66, 67, 68] + body_pose_permutation = body_pose_permutation[:len(body_pose)] + body_pose_permutation = [i-3 for i in body_pose_permutation] + + body_pose = body_pose[body_pose_permutation] + + global_orient[1::3] *= -1 + global_orient[2::3] *= -1 + body_pose[1::3] *= -1 + body_pose[2::3] *= -1 + + smpl_params = {'global_orient': global_orient.astype(np.float32), + 'body_pose': body_pose.astype(np.float32), + 'betas': betas.astype(np.float32) + } + + has_smpl_params = {'global_orient': has_global_orient, + 'body_pose': has_body_pose, + 'betas': has_betas + } + + return smpl_params, has_smpl_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 body 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 smpl_param_processing(smpl_params: Dict, has_smpl_params: Dict, rot: float, do_flip: bool) -> Tuple[Dict, Dict]: + """ + Apply random augmentations to the SMPL parameters. + Args: + smpl_params (Dict): SMPL parameter annotations. + has_smpl_params (Dict): Whether SMPL annotations are valid. + rot (float): Random rotation applied to the keypoints. + do_flip (bool): Whether to flip keypoints or not. + Returns: + Dict, Dict: Transformed SMPL parameters and valid flags. + """ + if do_flip: + smpl_params, has_smpl_params = fliplr_params(smpl_params, has_smpl_params) + smpl_params['global_orient'] = rot_aa(smpl_params['global_orient'], rot) + return smpl_params, has_smpl_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, + smpl_params: Dict, has_smpl_params: Dict, + flip_kp_permutation: List[int], + patch_width: int, patch_height: int, + mean: np.array, std: np.array, + do_augment: 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. + smpl_params (Dict): SMPL parameter annotations. + has_smpl_params (Dict): Whether SMPL 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, smpl_params, has_smpl_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. + smpl_params (Dict): Transformed SMPL parameters. + has_smpl_params (Dict): Valid flag for transformed SMPL 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 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) + + + smpl_params, has_smpl_params = smpl_param_processing(smpl_params, has_smpl_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, smpl_params, has_smpl_params, img_size + else: + return img_patch, keypoints_2d, keypoints_3d, smpl_params, has_smpl_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/humans4d/hmr2/datasets/vitdet_dataset.py b/humans4d/hmr2/datasets/vitdet_dataset.py new file mode 100644 index 0000000..3b5213f --- /dev/null +++ b/humans4d/hmr2/datasets/vitdet_dataset.py @@ -0,0 +1,88 @@ +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, + 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 = (boxes[:, 2:4] - boxes[:, 0:2]) / 200.0 + self.personid = np.arange(len(boxes), dtype=np.int32) + + 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 + + # 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, + False, 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]]) + return item diff --git a/humans4d/hmr2/models/__init__.py b/humans4d/hmr2/models/__init__.py new file mode 100644 index 0000000..10a7f8e --- /dev/null +++ b/humans4d/hmr2/models/__init__.py @@ -0,0 +1,94 @@ +import os + +from .smpl_wrapper import SMPL +from .hmr2 import HMR2 +from .discriminator import Discriminator + +from ..utils.download import cache_url +from ..configs import CACHE_DIR_4DHUMANS +from .... import REPO_PATH + + +def download_models(folder=CACHE_DIR_4DHUMANS): + """Download checkpoints and files for running inference. + """ + os.makedirs(folder, exist_ok=True) + download_files = { + "hmr2_data.tar.gz" : ["https://people.eecs.berkeley.edu/~jathushan/projects/4dhumans/hmr2_data.tar.gz", folder], + } + + for file_name, url in download_files.items(): + output_path = os.path.join(url[1], file_name) + if not os.path.exists(output_path): + print("Downloading file: " + file_name) + # output = gdown.cached_download(url[0], output_path, fuzzy=True) + output = cache_url(url[0], output_path) + assert os.path.exists(output_path), f"{output} does not exist" + + # if ends with tar.gz, tar -xzf + if file_name.endswith(".tar.gz"): + print("Extracting file: " + file_name) + os.system("tar -xvf " + output_path + " -C " + url[1]) + +def check_smpl_exists(smpl_dir): + import os + candidates = [ + f'{smpl_dir}/SMPL_NEUTRAL.pkl', + f'{smpl_dir}/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl', + ] + candidates_exist = [os.path.exists(c) for c in candidates] + if not any(candidates_exist): + raise FileNotFoundError(f"SMPL model not found. Please download it from https://smplify.is.tue.mpg.de/ and place it at {candidates[1]}") + + # Code edxpects SMPL model at CACHE_DIR_4DHUMANS/data/smpl/SMPL_NEUTRAL.pkl. Copy there if needed + if (not candidates_exist[0]) and candidates_exist[1]: + convert_pkl(candidates[1], candidates[0]) + + return True + +# Convert SMPL pkl file to be compatible with Python 3 +# Script is from https://rebeccabilbro.github.io/convert-py2-pickles-to-py3/ +def convert_pkl(old_pkl, new_pkl): + """ + Convert a Python 2 pickle to Python 3 + """ + import dill + import pickle + + # Convert Python 2 "ObjectType" to Python 3 object + dill._dill._reverse_typemap["ObjectType"] = object + + # Open the pickle using latin1 encoding + with open(old_pkl, "rb") as f: + loaded = pickle.load(f, encoding="latin1") + + # Re-save as Python 3 pickle + with open(new_pkl, "wb") as outfile: + pickle.dump(loaded, outfile) + + +def load_hmr2(checkpoint_path, smpl_dir): + from pathlib import Path + from ..configs import get_config + model_cfg = os.path.join(REPO_PATH, 'configs', 'model_config.yaml') + model_cfg = get_config(model_cfg, update_cachedir=True) + + model_cfg.SMPL.MODEL_PATH = smpl_dir + model_cfg.SMPL.JOINT_REGRESSION_EXTRA = os.path.join(smpl_dir, 'SMPL_to_J19.pkl') + model_cfg.SMPL.JOINT_REGRESSOR_EXTRA = os.path.join(smpl_dir, 'SMPL_to_J19.pkl') + model_cfg.SMPL.MEAN_PARAMS = os.path.join(smpl_dir, 'smpl_mean_params.npz') + + model_cfg.freeze() + + # Override some config values, to crop bbox correctly + if (model_cfg.MODEL.BACKBONE.TYPE == 'vit') 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() + + # Ensure SMPL model exists + check_smpl_exists(smpl_dir) + + model = HMR2.load_from_checkpoint(checkpoint_path, strict=False, cfg=model_cfg) + return model, model_cfg diff --git a/humans4d/hmr2/models/backbones/__init__.py b/humans4d/hmr2/models/backbones/__init__.py new file mode 100644 index 0000000..d2b217b --- /dev/null +++ b/humans4d/hmr2/models/backbones/__init__.py @@ -0,0 +1,7 @@ +from .vit import vit + +def create_backbone(cfg): + if cfg.MODEL.BACKBONE.TYPE == 'vit': + return vit(cfg) + else: + raise NotImplementedError('Backbone type is not implemented') diff --git a/humans4d/hmr2/models/backbones/vit.py b/humans4d/hmr2/models/backbones/vit.py new file mode 100644 index 0000000..c56c718 --- /dev/null +++ b/humans4d/hmr2/models/backbones/vit.py @@ -0,0 +1,348 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import math + +import torch +from functools import partial +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint + +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, + ) + +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, + ): + # 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 + + # 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] + + for blk in self.blocks: + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x) + else: + x = blk(x) + + x = self.last_norm(x) + + xp = x.permute(0, 2, 1).reshape(B, -1, Hp, Wp).contiguous() + + return xp + + 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/humans4d/hmr2/models/backbones/vit_vitpose.py b/humans4d/hmr2/models/backbones/vit_vitpose.py new file mode 100644 index 0000000..1ef5eaf --- /dev/null +++ b/humans4d/hmr2/models/backbones/vit_vitpose.py @@ -0,0 +1,17 @@ +# import mmcv +# import mmpose +# from mmpose.models import build_posenet +# from mmcv.runner import load_checkpoint +# from pathlib import Path + +# def vit(cfg): +# vitpose_dir = Path(mmpose.__file__).parent.parent +# config = f'{vitpose_dir}/configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/coco/ViTPose_huge_coco_256x192.py' +# # checkpoint = f'{vitpose_dir}/models/vitpose-h-multi-coco.pth' + +# config = mmcv.Config.fromfile(config) +# config.model.pretrained = None +# model = build_posenet(config.model) +# # load_checkpoint(model, checkpoint, map_location='cpu') + +# return model.backbone diff --git a/humans4d/hmr2/models/components/__init__.py b/humans4d/hmr2/models/components/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/humans4d/hmr2/models/components/pose_transformer.py b/humans4d/hmr2/models/components/pose_transformer.py new file mode 100644 index 0000000..ac04971 --- /dev/null +++ b/humans4d/hmr2/models/components/pose_transformer.py @@ -0,0 +1,358 @@ +from inspect import isfunction +from typing import Callable, Optional + +import torch +from einops import rearrange +from einops.layers.torch import Rearrange +from torch import nn + +from .t_cond_mlp import ( + AdaptiveLayerNorm1D, + FrequencyEmbedder, + normalization_layer, +) +# from .vit import Attention, FeedForward + + +def exists(val): + return val is not None + + +def default(val, d): + if exists(val): + return val + return d() if isfunction(d) else d + + +class PreNorm(nn.Module): + def __init__(self, dim: int, fn: Callable, norm: str = "layer", norm_cond_dim: int = -1): + super().__init__() + self.norm = normalization_layer(norm, dim, norm_cond_dim) + self.fn = fn + + def forward(self, x: torch.Tensor, *args, **kwargs): + if isinstance(self.norm, AdaptiveLayerNorm1D): + return self.fn(self.norm(x, *args), **kwargs) + else: + return self.fn(self.norm(x), **kwargs) + + +class FeedForward(nn.Module): + def __init__(self, dim, hidden_dim, dropout=0.0): + super().__init__() + self.net = nn.Sequential( + nn.Linear(dim, hidden_dim), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_dim, dim), + nn.Dropout(dropout), + ) + + def forward(self, x): + return self.net(x) + + +class Attention(nn.Module): + def __init__(self, dim, heads=8, dim_head=64, dropout=0.0): + super().__init__() + inner_dim = dim_head * heads + project_out = not (heads == 1 and dim_head == dim) + + self.heads = heads + self.scale = dim_head**-0.5 + + self.attend = nn.Softmax(dim=-1) + self.dropout = nn.Dropout(dropout) + + self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) + + self.to_out = ( + nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout)) + if project_out + else nn.Identity() + ) + + def forward(self, x): + qkv = self.to_qkv(x).chunk(3, dim=-1) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), qkv) + + dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale + + attn = self.attend(dots) + attn = self.dropout(attn) + + out = torch.matmul(attn, v) + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class CrossAttention(nn.Module): + def __init__(self, dim, context_dim=None, heads=8, dim_head=64, dropout=0.0): + super().__init__() + inner_dim = dim_head * heads + project_out = not (heads == 1 and dim_head == dim) + + self.heads = heads + self.scale = dim_head**-0.5 + + self.attend = nn.Softmax(dim=-1) + self.dropout = nn.Dropout(dropout) + + context_dim = default(context_dim, dim) + self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=False) + self.to_q = nn.Linear(dim, inner_dim, bias=False) + + self.to_out = ( + nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout)) + if project_out + else nn.Identity() + ) + + def forward(self, x, context=None): + context = default(context, x) + k, v = self.to_kv(context).chunk(2, dim=-1) + q = self.to_q(x) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), [q, k, v]) + + dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale + + attn = self.attend(dots) + attn = self.dropout(attn) + + out = torch.matmul(attn, v) + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) + + +class Transformer(nn.Module): + def __init__( + self, + dim: int, + depth: int, + heads: int, + dim_head: int, + mlp_dim: int, + dropout: float = 0.0, + norm: str = "layer", + norm_cond_dim: int = -1, + ): + super().__init__() + self.layers = nn.ModuleList([]) + for _ in range(depth): + sa = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout) + ff = FeedForward(dim, mlp_dim, dropout=dropout) + self.layers.append( + nn.ModuleList( + [ + PreNorm(dim, sa, norm=norm, norm_cond_dim=norm_cond_dim), + PreNorm(dim, ff, norm=norm, norm_cond_dim=norm_cond_dim), + ] + ) + ) + + def forward(self, x: torch.Tensor, *args): + for attn, ff in self.layers: + x = attn(x, *args) + x + x = ff(x, *args) + x + return x + + +class TransformerCrossAttn(nn.Module): + def __init__( + self, + dim: int, + depth: int, + heads: int, + dim_head: int, + mlp_dim: int, + dropout: float = 0.0, + norm: str = "layer", + norm_cond_dim: int = -1, + context_dim: Optional[int] = None, + ): + super().__init__() + self.layers = nn.ModuleList([]) + for _ in range(depth): + sa = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout) + ca = CrossAttention( + dim, context_dim=context_dim, heads=heads, dim_head=dim_head, dropout=dropout + ) + ff = FeedForward(dim, mlp_dim, dropout=dropout) + self.layers.append( + nn.ModuleList( + [ + PreNorm(dim, sa, norm=norm, norm_cond_dim=norm_cond_dim), + PreNorm(dim, ca, norm=norm, norm_cond_dim=norm_cond_dim), + PreNorm(dim, ff, norm=norm, norm_cond_dim=norm_cond_dim), + ] + ) + ) + + def forward(self, x: torch.Tensor, *args, context=None, context_list=None): + if context_list is None: + context_list = [context] * len(self.layers) + if len(context_list) != len(self.layers): + raise ValueError(f"len(context_list) != len(self.layers) ({len(context_list)} != {len(self.layers)})") + + for i, (self_attn, cross_attn, ff) in enumerate(self.layers): + x = self_attn(x, *args) + x + x = cross_attn(x, *args, context=context_list[i]) + x + x = ff(x, *args) + x + return x + + +class DropTokenDropout(nn.Module): + def __init__(self, p: float = 0.1): + super().__init__() + if p < 0 or p > 1: + raise ValueError( + "dropout probability has to be between 0 and 1, " "but got {}".format(p) + ) + self.p = p + + def forward(self, x: torch.Tensor): + # x: (batch_size, seq_len, dim) + if self.training and self.p > 0: + zero_mask = torch.full_like(x[0, :, 0], self.p).bernoulli().bool() + # TODO: permutation idx for each batch using torch.argsort + if zero_mask.any(): + x = x[:, ~zero_mask, :] + return x + + +class ZeroTokenDropout(nn.Module): + def __init__(self, p: float = 0.1): + super().__init__() + if p < 0 or p > 1: + raise ValueError( + "dropout probability has to be between 0 and 1, " "but got {}".format(p) + ) + self.p = p + + def forward(self, x: torch.Tensor): + # x: (batch_size, seq_len, dim) + if self.training and self.p > 0: + zero_mask = torch.full_like(x[:, :, 0], self.p).bernoulli().bool() + # Zero-out the masked tokens + x[zero_mask, :] = 0 + return x + + +class TransformerEncoder(nn.Module): + def __init__( + self, + num_tokens: int, + token_dim: int, + dim: int, + depth: int, + heads: int, + mlp_dim: int, + dim_head: int = 64, + dropout: float = 0.0, + emb_dropout: float = 0.0, + emb_dropout_type: str = "drop", + emb_dropout_loc: str = "token", + norm: str = "layer", + norm_cond_dim: int = -1, + token_pe_numfreq: int = -1, + ): + super().__init__() + if token_pe_numfreq > 0: + token_dim_new = token_dim * (2 * token_pe_numfreq + 1) + self.to_token_embedding = nn.Sequential( + Rearrange("b n d -> (b n) d", n=num_tokens, d=token_dim), + FrequencyEmbedder(token_pe_numfreq, token_pe_numfreq - 1), + Rearrange("(b n) d -> b n d", n=num_tokens, d=token_dim_new), + nn.Linear(token_dim_new, dim), + ) + else: + self.to_token_embedding = nn.Linear(token_dim, dim) + self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim)) + if emb_dropout_type == "drop": + self.dropout = DropTokenDropout(emb_dropout) + elif emb_dropout_type == "zero": + self.dropout = ZeroTokenDropout(emb_dropout) + else: + raise ValueError(f"Unknown emb_dropout_type: {emb_dropout_type}") + self.emb_dropout_loc = emb_dropout_loc + + self.transformer = Transformer( + dim, depth, heads, dim_head, mlp_dim, dropout, norm=norm, norm_cond_dim=norm_cond_dim + ) + + def forward(self, inp: torch.Tensor, *args, **kwargs): + x = inp + + if self.emb_dropout_loc == "input": + x = self.dropout(x) + x = self.to_token_embedding(x) + + if self.emb_dropout_loc == "token": + x = self.dropout(x) + b, n, _ = x.shape + x += self.pos_embedding[:, :n] + + if self.emb_dropout_loc == "token_afterpos": + x = self.dropout(x) + x = self.transformer(x, *args) + return x + + +class TransformerDecoder(nn.Module): + def __init__( + self, + num_tokens: int, + token_dim: int, + dim: int, + depth: int, + heads: int, + mlp_dim: int, + dim_head: int = 64, + dropout: float = 0.0, + emb_dropout: float = 0.0, + emb_dropout_type: str = 'drop', + norm: str = "layer", + norm_cond_dim: int = -1, + context_dim: Optional[int] = None, + skip_token_embedding: bool = False, + ): + super().__init__() + if not skip_token_embedding: + self.to_token_embedding = nn.Linear(token_dim, dim) + else: + self.to_token_embedding = nn.Identity() + if token_dim != dim: + raise ValueError( + f"token_dim ({token_dim}) != dim ({dim}) when skip_token_embedding is True" + ) + + self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim)) + if emb_dropout_type == "drop": + self.dropout = DropTokenDropout(emb_dropout) + elif emb_dropout_type == "zero": + self.dropout = ZeroTokenDropout(emb_dropout) + elif emb_dropout_type == "normal": + self.dropout = nn.Dropout(emb_dropout) + + self.transformer = TransformerCrossAttn( + dim, + depth, + heads, + dim_head, + mlp_dim, + dropout, + norm=norm, + norm_cond_dim=norm_cond_dim, + context_dim=context_dim, + ) + + def forward(self, inp: torch.Tensor, *args, context=None, context_list=None): + x = self.to_token_embedding(inp) + b, n, _ = x.shape + + x = self.dropout(x) + x += self.pos_embedding[:, :n] + + x = self.transformer(x, *args, context=context, context_list=context_list) + return x + diff --git a/humans4d/hmr2/models/components/t_cond_mlp.py b/humans4d/hmr2/models/components/t_cond_mlp.py new file mode 100644 index 0000000..44d5a09 --- /dev/null +++ b/humans4d/hmr2/models/components/t_cond_mlp.py @@ -0,0 +1,199 @@ +import copy +from typing import List, Optional + +import torch + + +class AdaptiveLayerNorm1D(torch.nn.Module): + def __init__(self, data_dim: int, norm_cond_dim: int): + super().__init__() + if data_dim <= 0: + raise ValueError(f"data_dim must be positive, but got {data_dim}") + if norm_cond_dim <= 0: + raise ValueError(f"norm_cond_dim must be positive, but got {norm_cond_dim}") + self.norm = torch.nn.LayerNorm( + data_dim + ) # TODO: Check if elementwise_affine=True is correct + self.linear = torch.nn.Linear(norm_cond_dim, 2 * data_dim) + torch.nn.init.zeros_(self.linear.weight) + torch.nn.init.zeros_(self.linear.bias) + + def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + # x: (batch, ..., data_dim) + # t: (batch, norm_cond_dim) + # return: (batch, data_dim) + x = self.norm(x) + alpha, beta = self.linear(t).chunk(2, dim=-1) + + # Add singleton dimensions to alpha and beta + if x.dim() > 2: + alpha = alpha.view(alpha.shape[0], *([1] * (x.dim() - 2)), alpha.shape[1]) + beta = beta.view(beta.shape[0], *([1] * (x.dim() - 2)), beta.shape[1]) + + return x * (1 + alpha) + beta + + +class SequentialCond(torch.nn.Sequential): + def forward(self, input, *args, **kwargs): + for module in self: + if isinstance(module, (AdaptiveLayerNorm1D, SequentialCond, ResidualMLPBlock)): + # print(f'Passing on args to {module}', [a.shape for a in args]) + input = module(input, *args, **kwargs) + else: + # print(f'Skipping passing args to {module}', [a.shape for a in args]) + input = module(input) + return input + + +def normalization_layer(norm: Optional[str], dim: int, norm_cond_dim: int = -1): + if norm == "batch": + return torch.nn.BatchNorm1d(dim) + elif norm == "layer": + return torch.nn.LayerNorm(dim) + elif norm == "ada": + assert norm_cond_dim > 0, f"norm_cond_dim must be positive, got {norm_cond_dim}" + return AdaptiveLayerNorm1D(dim, norm_cond_dim) + elif norm is None: + return torch.nn.Identity() + else: + raise ValueError(f"Unknown norm: {norm}") + + +def linear_norm_activ_dropout( + input_dim: int, + output_dim: int, + activation: torch.nn.Module = torch.nn.ReLU(), + bias: bool = True, + norm: Optional[str] = "layer", # Options: ada/batch/layer + dropout: float = 0.0, + norm_cond_dim: int = -1, +) -> SequentialCond: + layers = [] + layers.append(torch.nn.Linear(input_dim, output_dim, bias=bias)) + if norm is not None: + layers.append(normalization_layer(norm, output_dim, norm_cond_dim)) + layers.append(copy.deepcopy(activation)) + if dropout > 0.0: + layers.append(torch.nn.Dropout(dropout)) + return SequentialCond(*layers) + + +def create_simple_mlp( + input_dim: int, + hidden_dims: List[int], + output_dim: int, + activation: torch.nn.Module = torch.nn.ReLU(), + bias: bool = True, + norm: Optional[str] = "layer", # Options: ada/batch/layer + dropout: float = 0.0, + norm_cond_dim: int = -1, +) -> SequentialCond: + layers = [] + prev_dim = input_dim + for hidden_dim in hidden_dims: + layers.extend( + linear_norm_activ_dropout( + prev_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim + ) + ) + prev_dim = hidden_dim + layers.append(torch.nn.Linear(prev_dim, output_dim, bias=bias)) + return SequentialCond(*layers) + + +class ResidualMLPBlock(torch.nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + num_hidden_layers: int, + output_dim: int, + activation: torch.nn.Module = torch.nn.ReLU(), + bias: bool = True, + norm: Optional[str] = "layer", # Options: ada/batch/layer + dropout: float = 0.0, + norm_cond_dim: int = -1, + ): + super().__init__() + if not (input_dim == output_dim == hidden_dim): + raise NotImplementedError( + f"input_dim {input_dim} != output_dim {output_dim} is not implemented" + ) + + layers = [] + prev_dim = input_dim + for i in range(num_hidden_layers): + layers.append( + linear_norm_activ_dropout( + prev_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim + ) + ) + prev_dim = hidden_dim + self.model = SequentialCond(*layers) + self.skip = torch.nn.Identity() + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + return x + self.model(x, *args, **kwargs) + + +class ResidualMLP(torch.nn.Module): + def __init__( + self, + input_dim: int, + hidden_dim: int, + num_hidden_layers: int, + output_dim: int, + activation: torch.nn.Module = torch.nn.ReLU(), + bias: bool = True, + norm: Optional[str] = "layer", # Options: ada/batch/layer + dropout: float = 0.0, + num_blocks: int = 1, + norm_cond_dim: int = -1, + ): + super().__init__() + self.input_dim = input_dim + self.model = SequentialCond( + linear_norm_activ_dropout( + input_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim + ), + *[ + ResidualMLPBlock( + hidden_dim, + hidden_dim, + num_hidden_layers, + hidden_dim, + activation, + bias, + norm, + dropout, + norm_cond_dim, + ) + for _ in range(num_blocks) + ], + torch.nn.Linear(hidden_dim, output_dim, bias=bias), + ) + + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: + return self.model(x, *args, **kwargs) + + +class FrequencyEmbedder(torch.nn.Module): + def __init__(self, num_frequencies, max_freq_log2): + super().__init__() + frequencies = 2 ** torch.linspace(0, max_freq_log2, steps=num_frequencies) + self.register_buffer("frequencies", frequencies) + + def forward(self, x): + # x should be of size (N,) or (N, D) + N = x.size(0) + if x.dim() == 1: # (N,) + x = x.unsqueeze(1) # (N, D) where D=1 + x_unsqueezed = x.unsqueeze(-1) # (N, D, 1) + scaled = self.frequencies.view(1, 1, -1) * x_unsqueezed # (N, D, num_frequencies) + s = torch.sin(scaled) + c = torch.cos(scaled) + embedded = torch.cat([s, c, x_unsqueezed], dim=-1).view( + N, -1 + ) # (N, D * 2 * num_frequencies + D) + return embedded + diff --git a/humans4d/hmr2/models/discriminator.py b/humans4d/hmr2/models/discriminator.py new file mode 100644 index 0000000..4fa11b9 --- /dev/null +++ b/humans4d/hmr2/models/discriminator.py @@ -0,0 +1,99 @@ +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 = 23 + # 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 SMPL body poses (excluding the global orientation). + betas (torch.Tensor): Tensor of shape (B, 10) containign a batch of SMPL beta coefficients. + Returns: + torch.Tensor: Discriminator output with shape (B, 25) + """ + #import ipdb; ipdb.set_trace() + #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/humans4d/hmr2/models/heads/__init__.py b/humans4d/hmr2/models/heads/__init__.py new file mode 100644 index 0000000..d993f8d --- /dev/null +++ b/humans4d/hmr2/models/heads/__init__.py @@ -0,0 +1 @@ +from .smpl_head import build_smpl_head diff --git a/humans4d/hmr2/models/heads/smpl_head.py b/humans4d/hmr2/models/heads/smpl_head.py new file mode 100644 index 0000000..06dd89f --- /dev/null +++ b/humans4d/hmr2/models/heads/smpl_head.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import einops + +from ...utils.geometry import rot6d_to_rotmat, aa_to_rotmat +from ..components.pose_transformer import TransformerDecoder + +def build_smpl_head(cfg): + smpl_head_type = cfg.MODEL.SMPL_HEAD.get('TYPE', 'hmr') + if smpl_head_type == 'transformer_decoder': + return SMPLTransformerDecoderHead(cfg) + else: + raise ValueError('Unknown SMPL head type: {}'.format(smpl_head_type)) + +class SMPLTransformerDecoderHead(nn.Module): + """ Cross-attention based SMPL Transformer decoder + """ + + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.joint_rep_type = cfg.MODEL.SMPL_HEAD.get('JOINT_REP', '6d') + self.joint_rep_dim = {'6d': 6, 'aa': 3}[self.joint_rep_type] + npose = self.joint_rep_dim * (cfg.SMPL.NUM_BODY_JOINTS + 1) + self.npose = npose + self.input_is_mean_shape = cfg.MODEL.SMPL_HEAD.get('TRANSFORMER_INPUT', 'zero') == 'mean_shape' + transformer_args = dict( + num_tokens=1, + token_dim=(npose + 10 + 3) if self.input_is_mean_shape else 1, + dim=1024, + ) + transformer_args = (transformer_args | dict(cfg.MODEL.SMPL_HEAD.TRANSFORMER_DECODER)) + self.transformer = TransformerDecoder( + **transformer_args + ) + dim=transformer_args['dim'] + self.decpose = nn.Linear(dim, npose) + self.decshape = nn.Linear(dim, 10) + self.deccam = nn.Linear(dim, 3) + + if cfg.MODEL.SMPL_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) + + mean_params = np.load(cfg.SMPL.MEAN_PARAMS) + init_body_pose = torch.from_numpy(mean_params['pose'].astype(np.float32)).unsqueeze(0) + init_betas = torch.from_numpy(mean_params['shape'].astype('float32')).unsqueeze(0) + init_cam = torch.from_numpy(mean_params['cam'].astype(np.float32)).unsqueeze(0) + self.register_buffer('init_body_pose', init_body_pose) + self.register_buffer('init_betas', init_betas) + self.register_buffer('init_cam', init_cam) + + def forward(self, x, **kwargs): + + batch_size = x.shape[0] + # vit pretrained backbone is channel-first. Change to token-first + x = einops.rearrange(x, 'b c h w -> b (h w) c') + + init_body_pose = self.init_body_pose.expand(batch_size, -1) + init_betas = self.init_betas.expand(batch_size, -1) + init_cam = self.init_cam.expand(batch_size, -1) + + # TODO: Convert init_body_pose to aa rep if needed + if self.joint_rep_type == 'aa': + raise NotImplementedError + + pred_body_pose = init_body_pose + pred_betas = init_betas + pred_cam = init_cam + pred_body_pose_list = [] + pred_betas_list = [] + pred_cam_list = [] + for i in range(self.cfg.MODEL.SMPL_HEAD.get('IEF_ITERS', 1)): + # Input token to transformer is zero token + if self.input_is_mean_shape: + token = torch.cat([pred_body_pose, pred_betas, pred_cam], dim=1)[:,None,:] + else: + token = torch.zeros(batch_size, 1, 1).to(x.device) + + # Pass through transformer + token_out = self.transformer(token, context=x) + token_out = token_out.squeeze(1) # (B, C) + + # Readout from token_out + pred_body_pose = self.decpose(token_out) + pred_body_pose + pred_betas = self.decshape(token_out) + pred_betas + pred_cam = self.deccam(token_out) + pred_cam + pred_body_pose_list.append(pred_body_pose) + pred_betas_list.append(pred_betas) + pred_cam_list.append(pred_cam) + + # Convert self.joint_rep_type -> rotmat + joint_conversion_fn = { + '6d': rot6d_to_rotmat, + 'aa': lambda x: aa_to_rotmat(x.view(-1, 3).contiguous()) + }[self.joint_rep_type] + + pred_smpl_params_list = {} + pred_smpl_params_list['body_pose'] = torch.cat([joint_conversion_fn(pbp).view(batch_size, -1, 3, 3)[:, 1:, :, :] for pbp in pred_body_pose_list], dim=0) + pred_smpl_params_list['betas'] = torch.cat(pred_betas_list, dim=0) + pred_smpl_params_list['cam'] = torch.cat(pred_cam_list, dim=0) + pred_body_pose = joint_conversion_fn(pred_body_pose).view(batch_size, self.cfg.SMPL.NUM_BODY_JOINTS+1, 3, 3) + + pred_smpl_params = {'global_orient': pred_body_pose[:, [0]], + 'body_pose': pred_body_pose[:, 1:], + 'betas': pred_betas} + return pred_smpl_params, pred_cam, pred_smpl_params_list diff --git a/humans4d/hmr2/models/hmr2.py b/humans4d/hmr2/models/hmr2.py new file mode 100644 index 0000000..ef08a81 --- /dev/null +++ b/humans4d/hmr2/models/hmr2.py @@ -0,0 +1,355 @@ +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 build_smpl_head +from .discriminator import Discriminator +from .losses import Keypoint3DLoss, Keypoint2DLoss, ParameterLoss +from . import SMPL + +log = get_pylogger(__name__) + +class HMR2(pl.LightningModule): + + def __init__(self, cfg: CfgNode, init_renderer: bool = True): + """ + Setup HMR2 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']) + + # Create SMPL head + self.smpl_head = build_smpl_head(cfg) + + # 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.smpl_parameter_loss = ParameterLoss() + + # Instantiate SMPL model + smpl_cfg = {k.lower(): v for k,v in dict(cfg.SMPL).items()} + self.smpl = SMPL(**smpl_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.smpl.faces) + else: + self.renderer = None + self.mesh_renderer = None + + # Disable automatic optimization since we use adversarial training + self.automatic_optimization = False + + def get_parameters(self): + all_params = list(self.smpl_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 + conditioning_feats = self.backbone(x[:,:,:,32:-32]) + + pred_smpl_params, pred_cam, _ = self.smpl_head(conditioning_feats) + + # Store useful regression outputs to the output dict + output = {} + output['pred_cam'] = pred_cam + output['pred_smpl_params'] = {k: v.clone() for k,v in pred_smpl_params.items()} + + # Compute camera translation + device = pred_smpl_params['body_pose'].device + dtype = pred_smpl_params['body_pose'].dtype + focal_length = self.cfg.EXTRA.FOCAL_LENGTH * torch.ones(batch_size, 2, device=device, dtype=dtype) + 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_smpl_params['global_orient'] = pred_smpl_params['global_orient'].reshape(batch_size, -1, 3, 3) + pred_smpl_params['body_pose'] = pred_smpl_params['body_pose'].reshape(batch_size, -1, 3, 3) + pred_smpl_params['betas'] = pred_smpl_params['betas'].reshape(batch_size, -1) + smpl_output = self.smpl(**{k: v.float() for k,v in pred_smpl_params.items()}, pose2rot=False) + pred_keypoints_3d = smpl_output.joints + pred_vertices = smpl_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_smpl_params = output['pred_smpl_params'] + pred_keypoints_2d = output['pred_keypoints_2d'] + pred_keypoints_3d = output['pred_keypoints_3d'] + + + batch_size = pred_smpl_params['body_pose'].shape[0] + device = pred_smpl_params['body_pose'].device + dtype = pred_smpl_params['body_pose'].dtype + + # Get annotations + gt_keypoints_2d = batch['keypoints_2d'] + gt_keypoints_3d = batch['keypoints_3d'] + gt_smpl_params = batch['smpl_params'] + has_smpl_params = batch['has_smpl_params'] + is_axis_angle = batch['smpl_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=25+14) + + # Compute loss on SMPL parameters + loss_smpl_params = {} + for k, pred in pred_smpl_params.items(): + gt = gt_smpl_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_smpl_params[k] + loss_smpl_params[k] = self.smpl_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_smpl_params[k] * self.cfg.LOSS_WEIGHTS[k.upper()] for k in loss_smpl_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_smpl_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, + body_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 + body_pose (torch.Tensor): Regressed body 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 = body_pose.shape[0] + gt_body_pose = batch['body_pose'] + gt_betas = batch['betas'] + gt_rotmat = aa_to_rotmat(gt_body_pose.view(-1,3)).view(batch_size, -1, 3, 3) + disc_fake_out = self.discriminator(body_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_smpl_params = output['pred_smpl_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_smpl_params['body_pose'].reshape(batch_size, -1), pred_smpl_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_smpl_params['body_pose'].reshape(batch_size, -1), pred_smpl_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/humans4d/hmr2/models/losses.py b/humans4d/hmr2/models/losses.py new file mode 100644 index 0000000..0949d06 --- /dev/null +++ b/humans4d/hmr2/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 = 39): + """ + 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): + """ + SMPL 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 SMPL 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 SMPL 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/humans4d/hmr2/models/smpl_wrapper.py b/humans4d/hmr2/models/smpl_wrapper.py new file mode 100644 index 0000000..4f98454 --- /dev/null +++ b/humans4d/hmr2/models/smpl_wrapper.py @@ -0,0 +1,41 @@ +import torch +import numpy as np +import pickle +from typing import Optional +import smplx +from smplx.lbs import vertices2joints +from smplx.utils import SMPLOutput + + +class SMPL(smplx.SMPLLayer): + def __init__(self, *args, joint_regressor_extra: Optional[str] = None, update_hips: bool = False, **kwargs): + """ + Extension of the official SMPL implementation to support more joints. + Args: + Same as SMPLLayer. + joint_regressor_extra (str): Path to extra joint regressor. + """ + super(SMPL, self).__init__(*args, **kwargs) + smpl_to_openpose = [24, 12, 17, 19, 21, 16, 18, 20, 0, 2, 5, 8, 1, 4, + 7, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34] + + 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('joint_map', torch.tensor(smpl_to_openpose, dtype=torch.long)) + self.update_hips = update_hips + + def forward(self, *args, **kwargs) -> SMPLOutput: + """ + Run forward pass. Same as SMPL and also append an extra set of joints if joint_regressor_extra is specified. + """ + smpl_output = super(SMPL, self).forward(*args, **kwargs) + joints = smpl_output.joints[:, self.joint_map, :] + if self.update_hips: + joints[:,[9,12]] = joints[:,[9,12]] + \ + 0.25*(joints[:,[9,12]]-joints[:,[12,9]]) + \ + 0.5*(joints[:,[8]] - 0.5*(joints[:,[9,12]] + joints[:,[12,9]])) + if hasattr(self, 'joint_regressor_extra'): + extra_joints = vertices2joints(self.joint_regressor_extra, smpl_output.vertices) + joints = torch.cat([joints, extra_joints], dim=1) + smpl_output.joints = joints + return smpl_output diff --git a/humans4d/hmr2/utils/__init__.py b/humans4d/hmr2/utils/__init__.py new file mode 100644 index 0000000..09e47cd --- /dev/null +++ b/humans4d/hmr2/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/humans4d/hmr2/utils/download.py b/humans4d/hmr2/utils/download.py new file mode 100644 index 0000000..84d9b34 --- /dev/null +++ b/humans4d/hmr2/utils/download.py @@ -0,0 +1,66 @@ +import os +import re +import sys +from urllib import request as urlrequest + + +def _progress_bar(count, total): + """Report download progress. Credit: + https://stackoverflow.com/questions/3173320/text-progress-bar-in-the-console/27871113 + """ + bar_len = 60 + filled_len = int(round(bar_len * count / float(total))) + percents = round(100.0 * count / float(total), 1) + bar = "=" * filled_len + "-" * (bar_len - filled_len) + sys.stdout.write( + " [{}] {}% of {:.1f}MB file \r".format(bar, percents, total / 1024 / 1024) + ) + sys.stdout.flush() + if count >= total: + sys.stdout.write("\n") + + +def download_url(url, dst_file_path, chunk_size=8192, progress_hook=_progress_bar): + """Download url and write it to dst_file_path. Credit: + https://stackoverflow.com/questions/2028517/python-urllib2-progress-hook + """ + # url = url + "?dl=1" if "dropbox" in url else url + req = urlrequest.Request(url) + response = urlrequest.urlopen(req) + total_size = response.info().get("Content-Length") + if total_size is None: + raise ValueError("Cannot determine size of download from {}".format(url)) + total_size = int(total_size.strip()) + bytes_so_far = 0 + + with open(dst_file_path, "wb") as f: + while 1: + chunk = response.read(chunk_size) + bytes_so_far += len(chunk) + if not chunk: + break + + if progress_hook: + progress_hook(bytes_so_far, total_size) + + f.write(chunk) + return bytes_so_far + + +def cache_url(url_or_file, cache_file_path, download=True): + """Download the file specified by the URL to the cache_dir and return the path to + the cached file. If the argument is not a URL, simply return it as is. + """ + is_url = re.match(r"^(?:http)s?://", url_or_file, re.IGNORECASE) is not None + if not is_url: + return url_or_file + url = url_or_file + if os.path.exists(cache_file_path): + return cache_file_path + cache_file_dir = os.path.dirname(cache_file_path) + if not os.path.exists(cache_file_dir): + os.makedirs(cache_file_dir) + if download: + print("Downloading remote file {} to {}".format(url, cache_file_path)) + download_url(url, cache_file_path) + return cache_file_path diff --git a/humans4d/hmr2/utils/geometry.py b/humans4d/hmr2/utils/geometry.py new file mode 100644 index 0000000..7929ef5 --- /dev/null +++ b/humans4d/hmr2/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/humans4d/hmr2/utils/mesh_renderer.py b/humans4d/hmr2/utils/mesh_renderer.py new file mode 100644 index 0000000..f996da9 --- /dev/null +++ b/humans4d/hmr2/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() + body_keypoints = pred_keypoints[i, :25] + extra_keypoints = pred_keypoints[i, -19:] + for pair in keypoint_matches: + body_keypoints[pair[0], :] = extra_keypoints[pair[1], :] + pred_keypoints_img = render_openpose(255 * images_np[i].copy(), body_keypoints) / 255 + body_keypoints = gt_keypoints[i, :25] + extra_keypoints = gt_keypoints[i, -19:] + for pair in keypoint_matches: + if extra_keypoints[pair[1], -1] > 0 and body_keypoints[pair[0], -1] == 0: + body_keypoints[pair[0], :] = extra_keypoints[pair[1], :] + gt_keypoints_img = render_openpose(255*images_np[i].copy(), body_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/humans4d/hmr2/utils/misc.py b/humans4d/hmr2/utils/misc.py new file mode 100644 index 0000000..ffcfe78 --- /dev/null +++ b/humans4d/hmr2/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/humans4d/hmr2/utils/pck_accuracy.py b/humans4d/hmr2/utils/pck_accuracy.py new file mode 100644 index 0000000..5fea6c3 --- /dev/null +++ b/humans4d/hmr2/utils/pck_accuracy.py @@ -0,0 +1,94 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import numpy as np + +def _calc_distances(preds, targets, mask, normalize): + """Calculate the normalized distances between preds and target. + + Note: + batch_size: N + num_keypoints: K + dimension of keypoints: D (normally, D=2 or D=3) + + Args: + preds (np.ndarray[N, K, D]): Predicted keypoint location. + targets (np.ndarray[N, K, D]): Groundtruth keypoint location. + mask (np.ndarray[N, K]): Visibility of the target. False for invisible + joints, and True for visible. Invisible joints will be ignored for + accuracy calculation. + normalize (np.ndarray[N, D]): Typical value is heatmap_size + + Returns: + np.ndarray[K, N]: The normalized distances. \ + If target keypoints are missing, the distance is -1. + """ + N, K, _ = preds.shape + # set mask=0 when normalize==0 + _mask = mask.copy() + _mask[np.where((normalize == 0).sum(1))[0], :] = False + distances = np.full((N, K), -1, dtype=np.float32) + # handle invalid values + normalize[np.where(normalize <= 0)] = 1e6 + distances[_mask] = np.linalg.norm( + ((preds - targets) / normalize[:, None, :])[_mask], axis=-1) + return distances.T + + +def _distance_acc(distances, thr=0.5): + """Return the percentage below the distance threshold, while ignoring + distances values with -1. + + Note: + batch_size: N + Args: + distances (np.ndarray[N, ]): The normalized distances. + thr (float): Threshold of the distances. + + Returns: + float: Percentage of distances below the threshold. \ + If all target keypoints are missing, return -1. + """ + distance_valid = distances != -1 + num_distance_valid = distance_valid.sum() + if num_distance_valid > 0: + return (distances[distance_valid] < thr).sum() / num_distance_valid + return -1 + + +def keypoint_pck_accuracy(pred, gt, mask, thr, normalize): + """Calculate the pose accuracy of PCK for each individual keypoint and the + averaged accuracy across all keypoints for coordinates. + + Note: + PCK metric measures accuracy of the localization of the body joints. + The distances between predicted positions and the ground-truth ones + are typically normalized by the bounding box size. + The threshold (thr) of the normalized distance is commonly set + as 0.05, 0.1 or 0.2 etc. + + - batch_size: N + - num_keypoints: K + + Args: + pred (np.ndarray[N, K, 2]): Predicted keypoint location. + gt (np.ndarray[N, K, 2]): Groundtruth keypoint location. + mask (np.ndarray[N, K]): Visibility of the target. False for invisible + joints, and True for visible. Invisible joints will be ignored for + accuracy calculation. + thr (float): Threshold of PCK calculation. + normalize (np.ndarray[N, 2]): Normalization factor for H&W. + + Returns: + tuple: A tuple containing keypoint accuracy. + + - acc (np.ndarray[K]): Accuracy of each keypoint. + - avg_acc (float): Averaged accuracy across all keypoints. + - cnt (int): Number of valid keypoints. + """ + distances = _calc_distances(pred, gt, mask, normalize) + + acc = np.array([_distance_acc(d, thr) for d in distances]) + valid_acc = acc[acc >= 0] + cnt = len(valid_acc) + avg_acc = valid_acc.mean() if cnt > 0 else 0 + return acc, avg_acc, cnt diff --git a/humans4d/hmr2/utils/pose_utils.py b/humans4d/hmr2/utils/pose_utils.py new file mode 100644 index 0000000..bcc731c --- /dev/null +++ b/humans4d/hmr2/utils/pose_utils.py @@ -0,0 +1,310 @@ +""" +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, + keypoint_list: List, + pelvis_ind: int, + metrics: List = ['mode_mpjpe', 'mode_re', 'min_mpjpe', 'min_re'], + 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.keypoint_list = keypoint_list + self.pelvis_ind = pelvis_ind + self.metrics = metrics + for metric in self.metrics: + setattr(self, metric, np.zeros((dataset_length,))) + 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() + 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 __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) + + # 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 + pred_keypoints_2d = output['pred_keypoints_2d'].detach() + pred_keypoints_2d = pred_keypoints_2d[:,None,:,:] + gt_keypoints_2d = batch['keypoints_2d'][:,None,:,:].repeat(1, num_samples, 1, 1) + 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 + + 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.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) + 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 + + from .pck_accuracy import keypoint_pck_accuracy + pcks = [ + keypoint_pck_accuracy( + pred_kp_2d[:, 0, :, :], + gt_kp_2d[:, 0, :, :], + gt_conf_2d[:, 0, :]>0.5, + thr=thr, + normalize = np.ones((len(pred_kp_2d),2)) # Already in [-0.5,0.5] range. No need to normalize + ) + for thr in self.thresholds + ] + return pcks + + 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] + + pred_keypoints_2d = pred_keypoints_2d[:,None,:,:] + gt_keypoints_2d = batch['keypoints_2d'][:,None,:,:].repeat(1, num_samples, 1, 1) + + gt_bbox_expand_factor = (batch['box_size']/(batch['_scale']*200).max(dim=-1).values) + gt_bbox_expand_factor = gt_bbox_expand_factor[:,None,None,None].repeat(1, num_samples, 1, 1) + gt_bbox_expand_factor = gt_bbox_expand_factor.detach().cpu().numpy() + + self.pred_kp_2d.append(pred_keypoints_2d[:, :, :, :2].detach().cpu().numpy() * gt_bbox_expand_factor) + self.gt_conf_2d.append(gt_keypoints_2d[:, :, :, -1].detach().cpu().numpy()) + self.gt_kp_2d.append(gt_keypoints_2d[:, :, :, :2].detach().cpu().numpy() * gt_bbox_expand_factor) + + self.counter += batch_size diff --git a/humans4d/hmr2/utils/pylogger.py b/humans4d/hmr2/utils/pylogger.py new file mode 100644 index 0000000..92ffa71 --- /dev/null +++ b/humans4d/hmr2/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/humans4d/hmr2/utils/render_openpose.py b/humans4d/hmr2/utils/render_openpose.py new file mode 100644 index 0000000..a4a26e3 --- /dev/null +++ b/humans4d/hmr2/utils/render_openpose.py @@ -0,0 +1,149 @@ +""" +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) -> 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_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, + body_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) + return img diff --git a/humans4d/hmr2/utils/renderer.py b/humans4d/hmr2/utils/renderer.py new file mode 100644 index 0000000..22be8da --- /dev/null +++ b/humans4d/hmr2/utils/renderer.py @@ -0,0 +1,400 @@ +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 SMPL 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 + + self.camera_center = [self.img_res // 2, self.img_res // 2] + self.faces = faces + + def __call__(self, + vertices: np.array, + camera_translation: np.array, + image: torch.Tensor, + full_frame: bool = False, + imgname: Optional[str] = None, + side_view=False, top_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) + elif top_view: + rot = trimesh.transformations.rotation_matrix( + np.radians(rot_angle), [1, 0, 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 and not top_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,): + # 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]) + print(vertices.shape, camera_translation.shape) + mesh = trimesh.Trimesh(vertices.copy() + camera_translation, self.faces.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], + ): + + 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 cam_t is not None: + camera_translation = cam_t.copy() + # camera_translation[0] *= -1. + else: + camera_translation = np.array([0, 0, camera_z * self.focal_length/render_res[1]]) + + mesh = self.vertices_to_trimesh(vertices, camera_translation, mesh_base_color, rot_axis, rot_angle) + 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=self.focal_length, fy=self.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, + ): + + 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)) + + mesh_list = [pyrender.Mesh.from_trimesh(self.vertices_to_trimesh(vvv, ttt.copy(), mesh_base_color, rot_axis, rot_angle)) for vvv,ttt in zip(vertices, cam_t)] + + 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/humans4d/hmr2/utils/rich_utils.py b/humans4d/hmr2/utils/rich_utils.py new file mode 100644 index 0000000..19f9749 --- /dev/null +++ b/humans4d/hmr2/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/humans4d/hmr2/utils/skeleton_renderer.py b/humans4d/hmr2/utils/skeleton_renderer.py new file mode 100644 index 0000000..d2f3f5e --- /dev/null +++ b/humans4d/hmr2/utils/skeleton_renderer.py @@ -0,0 +1,122 @@ +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[:, [25+14], :-1] + pred_keypoints_3d[:, [25+14]] + 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] + 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, :25].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, :25].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, :25].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, :25].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, :25].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 diff --git a/humans4d/hmr2/utils/texture_utils.py b/humans4d/hmr2/utils/texture_utils.py new file mode 100644 index 0000000..0d29d95 --- /dev/null +++ b/humans4d/hmr2/utils/texture_utils.py @@ -0,0 +1,85 @@ +import numpy as np +import torch +from torch.nn import functional as F +# from psbody.mesh.visibility import visibility_compute + +def uv_to_xyz_and_normals(verts, f, fmap, bmap, ftov): + vn = estimate_vertex_normals(verts, f, ftov) + pixels_to_set = torch.nonzero(fmap+1) + x_to_set = pixels_to_set[:,0] + y_to_set = pixels_to_set[:,1] + b_coords = bmap[x_to_set, y_to_set, :] + f_coords = fmap[x_to_set, y_to_set] + v_ids = f[f_coords] + points = (b_coords[:,0,None]*verts[:,v_ids[:,0]] + + b_coords[:,1,None]*verts[:,v_ids[:,1]] + + b_coords[:,2,None]*verts[:,v_ids[:,2]]) + normals = (b_coords[:,0,None]*vn[:,v_ids[:,0]] + + b_coords[:,1,None]*vn[:,v_ids[:,1]] + + b_coords[:,2,None]*vn[:,v_ids[:,2]]) + return points, normals, vn, f_coords + +def estimate_vertex_normals(v, f, ftov): + face_normals = TriNormalsScaled(v, f) + non_scaled_normals = torch.einsum('ij,bjk->bik', ftov, face_normals) + norms = torch.sum(non_scaled_normals ** 2.0, 2) ** 0.5 + norms[norms == 0] = 1.0 + return torch.div(non_scaled_normals, norms[:,:,None]) + +def TriNormalsScaled(v, f): + return torch.cross(_edges_for(v, f, 1, 0), _edges_for(v, f, 2, 0)) + +def _edges_for(v, f, cplus, cminus): + return v[:,f[:,cplus]] - v[:,f[:,cminus]] + +def psbody_get_face_visibility(v, n, f, cams, normal_threshold=0.5): + bn, nverts, _ = v.shape + nfaces, _ = f.shape + vis_f = np.zeros([bn, nfaces], dtype='float32') + for i in range(bn): + vis, n_dot_cam = visibility_compute(v=v[i], n=n[i], f=f, cams=cams) + vis_v = (vis == 1) & (n_dot_cam > normal_threshold) + vis_f[i] = np.all(vis_v[0,f],1) + return vis_f + +def compute_uvsampler(vt, ft, tex_size=6): + """ + For this mesh, pre-computes the UV coordinates for + F x T x T points. + Returns F x T x T x 2 + """ + uv = obj2nmr_uvmap(ft, vt, tex_size=tex_size) + uv = uv.reshape(-1, tex_size, tex_size, 2) + return uv + +def obj2nmr_uvmap(ft, vt, tex_size=6): + """ + Converts obj uv_map to NMR uv_map (F x T x T x 2), + where tex_size (T) is the sample rate on each face. + """ + # This is F x 3 x 2 + uv_map_for_verts = vt[ft] + + # obj's y coordinate is [1-0], but image is [0-1] + uv_map_for_verts[:, :, 1] = 1 - uv_map_for_verts[:, :, 1] + + # range [0, 1] -> [-1, 1] + uv_map_for_verts = (2 * uv_map_for_verts) - 1 + + alpha = np.arange(tex_size, dtype=float) / (tex_size - 1) + beta = np.arange(tex_size, dtype=float) / (tex_size - 1) + import itertools + # Barycentric coordinate values + coords = np.stack([p for p in itertools.product(*[alpha, beta])]) + + # Compute alpha, beta (this is the same order as NMR) + v2 = uv_map_for_verts[:, 2] + v0v2 = uv_map_for_verts[:, 0] - uv_map_for_verts[:, 2] + v1v2 = uv_map_for_verts[:, 1] - uv_map_for_verts[:, 2] + # Interpolate the vertex uv values: F x 2 x T*2 + uv_map = np.dstack([v0v2, v1v2]).dot(coords.T) + v2.reshape(-1, 2, 1) + + # F x T*2 x 2 -> F x T x T x 2 + uv_map = np.transpose(uv_map, (0, 2, 1)).reshape(-1, tex_size, tex_size, 2) + + return uv_map diff --git a/humans4d/hmr2/utils/utils_detectron2.py b/humans4d/hmr2/utils/utils_detectron2.py new file mode 100644 index 0000000..fe01e02 --- /dev/null +++ b/humans4d/hmr2/utils/utils_detectron2.py @@ -0,0 +1,93 @@ +import detectron2.data.transforms as T +import torch +from detectron2.checkpoint import DetectionCheckpointer +from detectron2.config import CfgNode, instantiate +from detectron2.data import MetadataCatalog +from omegaconf import OmegaConf + + +class DefaultPredictor_Lazy: + """Create a simple end-to-end predictor with the given config that runs on single device for a + single input image. + + Compared to using the model directly, this class does the following additions: + + 1. Load checkpoint from the weights specified in config (cfg.MODEL.WEIGHTS). + 2. Always take BGR image as the input and apply format conversion internally. + 3. Apply resizing defined by the config (`cfg.INPUT.{MIN,MAX}_SIZE_TEST`). + 4. Take one input image and produce a single output, instead of a batch. + + This is meant for simple demo purposes, so it does the above steps automatically. + This is not meant for benchmarks or running complicated inference logic. + If you'd like to do anything more complicated, please refer to its source code as + examples to build and use the model manually. + + Attributes: + metadata (Metadata): the metadata of the underlying dataset, obtained from + test dataset name in the config. + + + Examples: + :: + pred = DefaultPredictor(cfg) + inputs = cv2.imread("input.jpg") + outputs = pred(inputs) + """ + + def __init__(self, cfg): + """ + Args: + cfg: a yacs CfgNode or a omegaconf dict object. + """ + if isinstance(cfg, CfgNode): + self.cfg = cfg.clone() # cfg can be modified by model + self.model = build_model(self.cfg) # noqa: F821 + if len(cfg.DATASETS.TEST): + test_dataset = cfg.DATASETS.TEST[0] + + checkpointer = DetectionCheckpointer(self.model) + checkpointer.load(cfg.MODEL.WEIGHTS) + + self.aug = T.ResizeShortestEdge( + [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST + ) + + self.input_format = cfg.INPUT.FORMAT + else: # new LazyConfig + self.cfg = cfg + self.model = instantiate(cfg.model) + test_dataset = OmegaConf.select(cfg, "dataloader.test.dataset.names", default=None) + if isinstance(test_dataset, (list, tuple)): + test_dataset = test_dataset[0] + + checkpointer = DetectionCheckpointer(self.model) + checkpointer.load(OmegaConf.select(cfg, "train.init_checkpoint", default="")) + + mapper = instantiate(cfg.dataloader.test.mapper) + self.aug = mapper.augmentations + self.input_format = mapper.image_format + + self.model.eval().cuda() + if test_dataset: + self.metadata = MetadataCatalog.get(test_dataset) + assert self.input_format in ["RGB", "BGR"], self.input_format + + def __call__(self, original_image): + """ + Args: + original_image (np.ndarray): an image of shape (H, W, C) (in BGR order). + + Returns: + predictions (dict): + the output of the model for one image only. + See :doc:`/tutorials/models` for details about the format. + """ + with torch.no_grad(): + if self.input_format == "RGB": + original_image = original_image[:, :, ::-1] + height, width = original_image.shape[:2] + image = self.aug(T.AugInput(original_image)).apply_image(original_image) + image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1)) + inputs = {"image": image, "height": height, "width": width} + predictions = self.model([inputs])[0] + return predictions diff --git a/humans4d/setup.py b/humans4d/setup.py new file mode 100644 index 0000000..3a50c07 --- /dev/null +++ b/humans4d/setup.py @@ -0,0 +1,31 @@ +from setuptools import setup, find_packages + +print('Found packages:', find_packages()) +setup( + description='HMR2 as a package', + name='hmr2', + packages=find_packages(), + install_requires=[ + 'gdown', + 'numpy', + 'torch', + 'torchvision', + 'pytorch-lightning', + 'smplx==0.1.28', + 'pyrender', + 'opencv-python', + 'yacs', + 'scikit-image', + 'einops', + 'timm', + 'webdataset', + 'dill', + 'pandas', + 'chumpy @ git+https://github.com/mattloper/chumpy', + ], + extras_require={ + 'all': [ + 'detectron2 @ git+https://github.com/facebookresearch/detectron2', + ], + }, +) diff --git a/humans4d/track.py b/humans4d/track.py new file mode 100644 index 0000000..4e727b5 --- /dev/null +++ b/humans4d/track.py @@ -0,0 +1,185 @@ +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Tuple + +import os +import hydra +import torch +import numpy as np +from hydra.core.config_store import ConfigStore +from omegaconf import DictConfig + +from phalp.configs.base import FullConfig +from phalp.models.hmar.hmr import HMR2018Predictor +from phalp.trackers.PHALP import PHALP +from phalp.utils import get_pylogger +from phalp.configs.base import CACHE_DIR + +from hmr2.datasets.utils import expand_bbox_to_aspect_ratio + +warnings.filterwarnings('ignore') + +log = get_pylogger(__name__) + +class HMR2Predictor(HMR2018Predictor): + def __init__(self, cfg) -> None: + super().__init__(cfg) + # Setup our new model + from hmr2.models import download_models, load_hmr2 + + # Download and load checkpoints + download_models() + model, _ = load_hmr2() + + self.model = model + self.model.eval() + + def forward(self, x): + hmar_out = self.hmar_old(x) + batch = { + 'img': x[:,:3,:,:], + 'mask': (x[:,3,:,:]).clip(0,1), + } + model_out = self.model(batch) + out = hmar_out | { + 'pose_smpl': model_out['pred_smpl_params'], + 'pred_cam': model_out['pred_cam'], + } + return out + +class HMR2023TextureSampler(HMR2Predictor): + def __init__(self, cfg) -> None: + super().__init__(cfg) + + # Model's all set up. Now, load tex_bmap and tex_fmap + # Texture map atlas + bmap = np.load(os.path.join(CACHE_DIR, 'phalp/3D/bmap_256.npy')) + fmap = np.load(os.path.join(CACHE_DIR, 'phalp/3D/fmap_256.npy')) + self.register_buffer('tex_bmap', torch.tensor(bmap, dtype=torch.float)) + self.register_buffer('tex_fmap', torch.tensor(fmap, dtype=torch.long)) + + self.img_size = 256 #self.cfg.MODEL.IMAGE_SIZE + self.focal_length = 5000. #self.cfg.EXTRA.FOCAL_LENGTH + + import neural_renderer as nr + self.neural_renderer = nr.Renderer(dist_coeffs=None, orig_size=self.img_size, + image_size=self.img_size, + light_intensity_ambient=1, + light_intensity_directional=0, + anti_aliasing=False) + + def forward(self, x): + batch = { + 'img': x[:,:3,:,:], + 'mask': (x[:,3,:,:]).clip(0,1), + } + model_out = self.model(batch) + + # from hmr2.models.prohmr_texture import unproject_uvmap_to_mesh + + def unproject_uvmap_to_mesh(bmap, fmap, verts, faces): + # bmap: 256,256,3 + # fmap: 256,256 + # verts: B,V,3 + # faces: F,3 + valid_mask = (fmap >= 0) + + fmap_flat = fmap[valid_mask] # N + bmap_flat = bmap[valid_mask,:] # N,3 + + face_vids = faces[fmap_flat, :] # N,3 + face_verts = verts[:, face_vids, :] # B,N,3,3 + + bs = face_verts.shape + map_verts = torch.einsum('bnij,ni->bnj', face_verts, bmap_flat) # B,N,3 + + return map_verts, valid_mask + + pred_verts = model_out['pred_vertices'] + model_out['pred_cam_t'].unsqueeze(1) + device = pred_verts.device + face_tensor = torch.tensor(self.smpl.faces.astype(np.int64), dtype=torch.long, device=device) + map_verts, valid_mask = unproject_uvmap_to_mesh(self.tex_bmap, self.tex_fmap, pred_verts, face_tensor) # B,N,3 + + # Project map_verts to image using K,R,t + # map_verts_view = einsum('bij,bnj->bni', R, map_verts) + t # R=I t=0 + focal = self.focal_length / (self.img_size / 2) + map_verts_proj = focal * map_verts[:, :, :2] / map_verts[:, :, 2:3] # B,N,2 + map_verts_depth = map_verts[:, :, 2] # B,N + + # Render Depth. Annoying but we need to create this + K = torch.eye(3, device=device) + K[0, 0] = K[1, 1] = self.focal_length + K[1, 2] = K[0, 2] = self.img_size / 2 # Because the neural renderer only support squared images + K = K.unsqueeze(0) + R = torch.eye(3, device=device).unsqueeze(0) + t = torch.zeros(3, device=device).unsqueeze(0) + rend_depth = self.neural_renderer(pred_verts, + face_tensor[None].expand(pred_verts.shape[0], -1, -1).int(), + # textures=texture_atlas_rgb, + mode='depth', + K=K, R=R, t=t) + + rend_depth_at_proj = torch.nn.functional.grid_sample(rend_depth[:,None,:,:], map_verts_proj[:,None,:,:]) # B,1,1,N + rend_depth_at_proj = rend_depth_at_proj.squeeze(1).squeeze(1) # B,N + + img_rgba = torch.cat([batch['img'], batch['mask'][:,None,:,:]], dim=1) # B,4,H,W + img_rgba_at_proj = torch.nn.functional.grid_sample(img_rgba, map_verts_proj[:,None,:,:]) # B,4,1,N + img_rgba_at_proj = img_rgba_at_proj.squeeze(2) # B,4,N + + visibility_mask = map_verts_depth <= (rend_depth_at_proj + 1e-4) # B,N + img_rgba_at_proj[:,3,:][~visibility_mask] = 0 + + # Paste image back onto square uv_image + uv_image = torch.zeros((batch['img'].shape[0], 4, 256, 256), dtype=torch.float, device=device) + uv_image[:, :, valid_mask] = img_rgba_at_proj + + out = { + 'uv_image': uv_image, + 'uv_vector' : self.hmar_old.process_uv_image(uv_image), + 'pose_smpl': model_out['pred_smpl_params'], + 'pred_cam': model_out['pred_cam'], + } + return out + +class HMR2_4dhuman(PHALP): + def __init__(self, cfg): + super().__init__(cfg) + + def setup_hmr(self): + self.HMAR = HMR2023TextureSampler(self.cfg) + + def get_detections(self, image, frame_name, t_, additional_data=None, measurments=None): + ( + pred_bbox, pred_bbox, pred_masks, pred_scores, pred_classes, + ground_truth_track_id, ground_truth_annotations + ) = super().get_detections(image, frame_name, t_, additional_data, measurments) + + # Pad bounding boxes + pred_bbox_padded = expand_bbox_to_aspect_ratio(pred_bbox, self.cfg.expand_bbox_shape) + + return ( + pred_bbox, pred_bbox_padded, pred_masks, pred_scores, pred_classes, + ground_truth_track_id, ground_truth_annotations + ) + + +@dataclass +class Human4DConfig(FullConfig): + # override defaults if needed + expand_bbox_shape: Optional[Tuple[int]] = (192,256) + pass + +cs = ConfigStore.instance() +cs.store(name="config", node=Human4DConfig) + +@hydra.main(version_base="1.2", config_name="config") +def main(cfg: DictConfig) -> Optional[float]: + """Main function for running the PHALP tracker.""" + + phalp_tracker = HMR2_4dhuman(cfg) + + phalp_tracker.track() + +if __name__ == "__main__": + main() diff --git a/humans4d/train.py b/humans4d/train.py new file mode 100644 index 0000000..3644ed6 --- /dev/null +++ b/humans4d/train.py @@ -0,0 +1,111 @@ +from typing import Optional, Tuple +import pyrootutils + +root = pyrootutils.setup_root( + search_from=__file__, + indicator=[".git", "pyproject.toml"], + pythonpath=True, + dotenv=True, +) + +import os +from pathlib import Path + +import hydra +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf +from pytorch_lightning import Trainer +from pytorch_lightning.loggers import TensorBoardLogger +from pytorch_lightning.plugins.environments import SLURMEnvironment + +from yacs.config import CfgNode +from hmr2.configs import dataset_config +from hmr2.datasets import HMR2DataModule +from hmr2.models.hmr2 import HMR2 +from hmr2.utils.pylogger import get_pylogger +from hmr2.utils.misc import task_wrapper, log_hyperparameters + +# HACK reset the signal handling so the lightning is free to set it +# Based on https://github.com/facebookincubator/submitit/issues/1709#issuecomment-1246758283 +import signal +signal.signal(signal.SIGUSR1, signal.SIG_DFL) + +log = get_pylogger(__name__) + + +@pl.utilities.rank_zero.rank_zero_only +def save_configs(model_cfg: CfgNode, dataset_cfg: CfgNode, rootdir: str): + """Save config files to rootdir.""" + Path(rootdir).mkdir(parents=True, exist_ok=True) + OmegaConf.save(config=model_cfg, f=os.path.join(rootdir, 'model_config.yaml')) + with open(os.path.join(rootdir, 'dataset_config.yaml'), 'w') as f: + f.write(dataset_cfg.dump()) + +@task_wrapper +def train(cfg: DictConfig) -> Tuple[dict, dict]: + + # Load dataset config + dataset_cfg = dataset_config() + + # Save configs + save_configs(cfg, dataset_cfg, cfg.paths.output_dir) + + # Setup training and validation datasets + datamodule = HMR2DataModule(cfg, dataset_cfg) + + # Setup model + model = HMR2(cfg) + + # Setup Tensorboard logger + logger = TensorBoardLogger(os.path.join(cfg.paths.output_dir, 'tensorboard'), name='', version='', default_hp_metric=False) + loggers = [logger] + + # Setup checkpoint saving + checkpoint_callback = pl.callbacks.ModelCheckpoint( + dirpath=os.path.join(cfg.paths.output_dir, 'checkpoints'), + every_n_train_steps=cfg.GENERAL.CHECKPOINT_STEPS, + save_last=True, + save_top_k=cfg.GENERAL.CHECKPOINT_SAVE_TOP_K, + ) + rich_callback = pl.callbacks.RichProgressBar() + lr_monitor = pl.callbacks.LearningRateMonitor(logging_interval='step') + callbacks = [ + checkpoint_callback, + lr_monitor, + # rich_callback + ] + + log.info(f"Instantiating trainer <{cfg.trainer._target_}>") + trainer: Trainer = hydra.utils.instantiate( + cfg.trainer, + callbacks=callbacks, + logger=loggers, + plugins=(SLURMEnvironment(requeue_signal=signal.SIGUSR2) if (cfg.get('launcher',None) is not None) else None), # Submitit uses SIGUSR2 + ) + + object_dict = { + "cfg": cfg, + "datamodule": datamodule, + "model": model, + "callbacks": callbacks, + "logger": logger, + "trainer": trainer, + } + + if logger: + log.info("Logging hyperparameters!") + log_hyperparameters(object_dict) + + # Train the model + trainer.fit(model, datamodule=datamodule, ckpt_path='last') + log.info("Fitting done") + + +@hydra.main(version_base="1.2", config_path=str(root/"hmr2/configs_hydra"), config_name="train.yaml") +def main(cfg: DictConfig) -> Optional[float]: + # train the model + train(cfg) + + +if __name__ == "__main__": + main() diff --git a/nodes/load_detectron_node.py b/nodes/load_detectron_node.py new file mode 100644 index 0000000..d7f4a29 --- /dev/null +++ b/nodes/load_detectron_node.py @@ -0,0 +1,27 @@ +import os + +from folder_paths import models_dir + +from ..utils.process_humans import load_detectron + + + +DETECTRON_PATH = os.path.join(models_dir, 'detectron') +os.makedirs(DETECTRON_PATH, exist_ok=True) + + + +class LoadDetectronNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { "checkpoint": (os.listdir(DETECTRON_PATH), )}} + + RETURN_TYPES = ("DETECTRON",) + FUNCTION = "load" + + CATEGORY = "4dhumans" + + def load(self, checkpoint): + checkpoint_path = os.path.join(DETECTRON_PATH, checkpoint) + model = load_detectron(checkpoint_path) + return (model,) \ No newline at end of file diff --git a/nodes/load_hmr_node.py b/nodes/load_hmr_node.py new file mode 100644 index 0000000..72b47b1 --- /dev/null +++ b/nodes/load_hmr_node.py @@ -0,0 +1,30 @@ +import os + +import torch + +from folder_paths import models_dir + +from ..utils.process_humans import load_hmr_model + + +HMR_PATH = os.path.join(models_dir, 'hmr') +os.makedirs(HMR_PATH, exist_ok=True) + +SMPL_PATH = os.path.join(models_dir, 'smpl') +os.makedirs(SMPL_PATH, exist_ok=True) + + +class LoadHMRNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { "checkpoint": (os.listdir(HMR_PATH), )}} + + RETURN_TYPES = ("HMR",) + FUNCTION = "load" + + CATEGORY = "4dhumans" + + def load(self, checkpoint): + checkpoint_path = os.path.join(HMR_PATH, checkpoint) + model_info = load_hmr_model(checkpoint_path, SMPL_PATH) + return (model_info,) \ No newline at end of file diff --git a/nodes/process_humans_node.py b/nodes/process_humans_node.py new file mode 100644 index 0000000..b04edb1 --- /dev/null +++ b/nodes/process_humans_node.py @@ -0,0 +1,27 @@ +from folder_paths import models_dir +import comfy.model_management + +from ..utils.process_humans import process_humans + + +class ProcessHumansNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE",), + "hmr": ("HMR", ), + "detectron": ("DETECTRON", ), + }} + RETURN_TYPES = ("SMPL", "IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("smpl", "images", "masks", "semantic_maps") + FUNCTION = "process" + + CATEGORY = "4dhumans" + + def process(self, images, hmr, detectron): + intermediate_device = comfy.model_management.intermediate_device() + torch_device = comfy.model_management.get_torch_device() + + results = process_humans(images, hmr, detectron, figure_scale=None, device=torch_device) + + return (results['smpl'], results['images'], results['masks'], results['semantic_maps']) diff --git a/nodes/select_human_node.py b/nodes/select_human_node.py new file mode 100644 index 0000000..0cf82ad --- /dev/null +++ b/nodes/select_human_node.py @@ -0,0 +1,21 @@ +from ..utils.process_humans import get_human + + +class SelectHumanNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE",), + "smpl": ("SMPL", ), + "idx": ("INT", {"default": 0, "min": 0, "max": 999, "step": 1}), + }} + RETURN_TYPES = ("SMPL", "IMAGE", "MASK", "IMAGE") + RETURN_NAMES = ("smpl", "images", "masks", "semantic_maps") + FUNCTION = "process" + + CATEGORY = "4dhumans" + + def process(self, images, smpl, idx): + results = get_human(images, smpl, idx) + + return (results['smpl'], results['images'], results['masks'], results['semantic_maps']) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b2c7687 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +# git+https://github.com/shubham-goel/4D-Humans +git+https://github.com/facebookresearch/detectron2 +pyrender +yacs \ No newline at end of file diff --git a/utils/process_humans.py b/utils/process_humans.py new file mode 100644 index 0000000..38c52cd --- /dev/null +++ b/utils/process_humans.py @@ -0,0 +1,298 @@ +import os +import torch +import numpy as np +import platform +import pyrender +# from pathlib import Path +from tqdm import tqdm + + +if "PYOPENGL_PLATFORM" not in os.environ: + os.environ["PYOPENGL_PLATFORM"] = "egl" + +import comfy.utils + +from ..humans4d.hmr2.models import load_hmr2 +from ..humans4d.hmr2.utils import recursive_to +from ..humans4d.hmr2.datasets.vitdet_dataset import ViTDetDataset +from ..humans4d.hmr2.utils.renderer import cam_crop_to_full +# from ..humans4d.hmr2.configs import get_config +# from ..humans4d.hmr2.models import check_smpl_exists +from ..humans4d.hmr2.utils.utils_detectron2 import DefaultPredictor_Lazy +from detectron2.config import LazyConfig + +from .visualize_humans import SemanticRenderer +from .. import REPO_PATH + +# For Windows, remove PYOPENGL_PLATFORM to enable default rendering backend +sys_name = platform.system() +if sys_name == "Windows": + os.environ.pop("PYOPENGL_PLATFORM") + + +LIGHT_BLUE = (0.65098039, 0.74117647, 0.85882353) + + +def predict_smpl(batch, model, model_cfg, figure_scale=None): + all_verts = [] + all_cam_t = [] + with torch.no_grad(): + out = model(batch) + + pred_cam = out["pred_cam"] + pred_smpl_parameter = out["pred_smpl_params"] + if figure_scale is not None: + pred_smpl_parameter['betas'][0][1] = float(figure_scale) + smpl_output = model.smpl( + **{k: v.float() for k, v in pred_smpl_parameter.items()}, + pose2rot=False, + ) + pred_vertices = smpl_output.vertices + out["pred_vertices"] = pred_vertices.reshape( + batch["img"].shape[0], -1, 3 + ) + + 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 + batch_size = batch["img"].shape[0] + for n in range(batch_size): + # Add all verts and cams to list + verts = out["pred_vertices"][n].detach().cpu().numpy() + cam_t = pred_cam_t_full[n] + all_verts.append(verts) + all_cam_t.append(cam_t) + + misc_args = dict( + mesh_base_color=LIGHT_BLUE, + scene_bg_color=(1, 1, 1), + focal_length=scaled_focal_length, + ) + + smpl_outs = { + k: v.detach().cpu().numpy() for k, v in pred_smpl_parameter.items() + } + + results_dict_for_rendering = { + "verts": all_verts, + "cam_t": all_cam_t, + "render_res": img_size[n].cpu().numpy(), + "smpls": smpl_outs, + "scaled_focal_length": scaled_focal_length.cpu().numpy(), + } + return results_dict_for_rendering, misc_args + + +def load_image(model_cfg, image, detector): + # Detect humans in image + # image = image.permute(1,2,0)*255 + det_out = detector(image) + # image = image.numpy() + det_instances = det_out["instances"] + valid_idx = (det_instances.pred_classes == 0) & (det_instances.scores > 0.5) + boxes = det_instances.pred_boxes.tensor[valid_idx].cpu().numpy() + # Run HMR2.0 on all detected humans + dataset = ViTDetDataset(model_cfg, image, boxes) + return torch.utils.data.DataLoader( + dataset, batch_size=8, shuffle=False, num_workers=0 + ) + +def load_hmr_model(hmr2_checkpoint_path, smpl_dir_path): + model, model_cfg = load_hmr2(hmr2_checkpoint_path,smpl_dir_path) + return { + 'model': model, + 'model_cfg': model_cfg + } + + +def load_detectron(detectron_checkpoint_path): + cfg_path = os.path.join(REPO_PATH, "configs", "cascade_mask_rcnn_vitdet_h_75ep.py") + detectron2_cfg = LazyConfig.load(str(cfg_path)) + detectron2_cfg.train.init_checkpoint = detectron_checkpoint_path + for i in range(3): + detectron2_cfg.model.roi_heads.box_predictors[i].test_score_thresh = 0.25 + detector = DefaultPredictor_Lazy(detectron2_cfg) + return detector + + +def process_humans(images, hmr, detector, figure_scale, device): + images = images.permute(0, 3, 1, 2) + _, _, height, width = images.shape + model, model_cfg = hmr['model'], hmr['model_cfg'] + + + model = model.to(device) + detector.model.to(device) + # This PyRender is only used for visualizing, we use Blender after to render different conditions. + renderer = SemanticRenderer( + model_cfg, + faces=model.smpl.faces, + lbs=model.smpl.lbs_weights, + viewport_size=(720, 720), + ) + results = [] + rendered_images = [] + semantic_maps = [] + masks = [] + focal_length = None + mesh_base_color = None + scene_bg_color = None + pbar = comfy.utils.ProgressBar(len(images)) + for image in tqdm(images, desc="Processing Human Images:"): + image = image.permute(1,2,0)*255 + image = image.numpy() + renderer.renderer.delete() + renderer.renderer = pyrender.OffscreenRenderer( + viewport_width=width, + viewport_height=height, + point_size=1.0, + ) + dataloader = load_image(model_cfg, image, detector) + + + for batch in dataloader: + batch = recursive_to(batch, device) + results_dict_for_rendering, misc_args = predict_smpl(batch, model, model_cfg, figure_scale) + mesh_base_color = misc_args['mesh_base_color'] + scene_bg_color = misc_args['scene_bg_color'] + focal_length = misc_args['focal_length'].to('cpu') + + rendering_results = renderer.render_all_multiple( + results_dict_for_rendering["verts"], + cam_t=results_dict_for_rendering["cam_t"], + render_res=results_dict_for_rendering["render_res"], **misc_args + ) + # Overlay image + valid_mask = rendering_results["Image"][:, :, -1][:, :, np.newaxis] + cam_view = ( + valid_mask * rendering_results["Image"][:, :, [2, 1, 0]] + + (1 - valid_mask) * image.astype(np.float32)[:, :, ::-1] / 255 + ) + rendered_images.append(torch.from_numpy(cam_view[:, :, ::-1].copy())) + semantic_maps.append(torch.from_numpy(rendering_results.get("SemanticMap"))) + masks.append(torch.from_numpy(rendering_results.get("Mask")[:, :, 0])) + results.append(results_dict_for_rendering) # still in np + pbar.update(1) + + smpl = { + 'results': results, + 'width': width, + 'height': height, + 'model_cfg': model_cfg, + 'faces': model.smpl.faces, + 'lbs': model.smpl.lbs_weights, + 'viewport_size': (720, 720), + 'misc_args': { 'mesh_base_color': mesh_base_color, 'scene_bg_color': scene_bg_color, 'focal_length': focal_length } + } + + return { + 'smpl': smpl, + 'images': torch.stack(rendered_images), + 'masks': torch.stack(masks).to(torch.float16), + 'semantic_maps': torch.stack(semantic_maps), + } + + +def get_human(images, smpl, idx): + results = [] + rendered_images = [] + semantic_maps = [] + masks = [] + + model_cfg, faces, lbs = smpl['model_cfg'], smpl['faces'], smpl['lbs'] + viewport_size = smpl['viewport_size'] + width, height = smpl['width'], smpl['height'] + misc_args = smpl['misc_args'] + + renderer = SemanticRenderer( + model_cfg, + faces=faces, + lbs=lbs, + viewport_size=viewport_size, + ) + + renderer.renderer = pyrender.OffscreenRenderer( + viewport_width=width, + viewport_height=height, + point_size=1.0, + ) + + human_smpl = [] + keys = ['verts'] + smpl_keys = ['global_orient', 'body_pose', 'betas'] + for results_dict_for_rendering in smpl['results']: + result = { + 'scaled_focal_length': results_dict_for_rendering['scaled_focal_length'], + 'smpls': {}, + 'cam_t': results_dict_for_rendering['cam_t'], + 'render_res': results_dict_for_rendering['render_res'] + } + good = True + for key in keys: + if idx < len(results_dict_for_rendering[key]): + result[key] = [results_dict_for_rendering[key][idx]] + else: + good = False + + for key in smpl_keys: + if idx < len(results_dict_for_rendering['smpls'][key]): + result['smpls'][key] = [results_dict_for_rendering['smpls'][key][idx]] + else: + good = False + + if good: + human_smpl.append(result) + else: + human_smpl.append(None) + + for image_orig, results_dict_for_rendering in zip(images, human_smpl): + image = image_orig.permute(1,2,0)*255 + image = image.numpy() + if results_dict_for_rendering is not None: + rendering_results = renderer.render_all_multiple( + results_dict_for_rendering["verts"], + cam_t=results_dict_for_rendering["cam_t"], + render_res=results_dict_for_rendering["render_res"], **misc_args + ) + # Overlay image + valid_mask = rendering_results["Image"][:, :, -1][:, :, np.newaxis] + cam_view = ( + valid_mask * rendering_results["Image"][:, :, [2, 1, 0]] + + (1 - valid_mask) * image.astype(np.float32)[:, :, ::-1] / 255 + ) + rendered_images.append(torch.from_numpy(cam_view[:, :, ::-1].copy())) + semantic_maps.append(torch.from_numpy(rendering_results.get("SemanticMap"))) + masks.append(torch.from_numpy(rendering_results.get("Mask")[:, :, 0])) + else: + rendered_images.append(image_orig) + semantic_maps.append(torch.ones_like(image_orig)) + masks.append(torch.zeros_like(image_orig)) + results.append(results_dict_for_rendering) # still in np + + smpl = { + 'results': results, + 'width': width, + 'height': height, + 'model_cfg': model_cfg, + 'faces': faces, + 'lbs': lbs, + 'viewport_size': (720, 720), + 'misc_args': misc_args + } + + return { + 'smpl': smpl, + 'images': torch.stack(rendered_images), + 'masks': torch.stack(masks).to(torch.float16), + 'semantic_maps': torch.stack(semantic_maps), + } diff --git a/utils/visualize_humans.py b/utils/visualize_humans.py new file mode 100644 index 0000000..5fd7ef5 --- /dev/null +++ b/utils/visualize_humans.py @@ -0,0 +1,294 @@ +import os +import platform + +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 +from ..humans4d.hmr2.utils.renderer import Renderer + +# For Windows, remove PYOPENGL_PLATFORM to enable default rendering backend +sys_name = platform.system() +if sys_name == "Windows": + os.environ.pop("PYOPENGL_PLATFORM") + +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 SemanticRenderer(Renderer): + + def __init__( + self, cfg: CfgNode, faces: np.array, lbs: np.array, viewport_size=(768, 768) + ): + """ + Wrapper around the pyrender renderer to render SMPL meshes's semantic map for Champ. + 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 + + self.camera_center = [self.img_res // 2, self.img_res // 2] + self.faces = faces + self.lbs = lbs + # self.joint_names = smplx.joint_names.JOINT_NAMES + + self.vertex_labels = np.argmax(self.lbs.cpu().numpy(), axis=1) + vertex_colors = cv2.applyColorMap( + (10 * self.vertex_labels).astype(np.uint8), cv2.COLORMAP_VIRIDIS + ) + semantic_background_rgb = cv2.applyColorMap(np.uint8([0]), cv2.COLORMAP_VIRIDIS) + self.vertex_colors = np.squeeze(vertex_colors, axis=1) + self.semantic_background_rgb = ( + np.squeeze(semantic_background_rgb.astype(np.float32), axis=1) / 255 + ) + self.renderer = pyrender.OffscreenRenderer( + viewport_width=viewport_size[0], + viewport_height=viewport_size[1], + point_size=1.0, + ) + + 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, + ): + vertex_colors = np.array([(*mesh_base_color, 1.0)] * vertices.shape[0]) + mesh = trimesh.Trimesh( + vertices.copy() + camera_translation, + self.faces.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_all_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, + ): + + renderer = self.renderer + trimesh_list = [ + self.vertices_to_trimesh( + vvv, ttt.copy(), mesh_base_color, rot_axis, rot_angle + ) + for vvv, ttt in zip(vertices, cam_t) + ] + for trimesh in trimesh_list: + trimesh.visual.vertex_colors = self.vertex_colors + mesh_list = [pyrender.Mesh.from_trimesh(trimesh) for trimesh in trimesh_list] + + 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_center = [render_res[0] / 2.0, render_res[1] / 2.0] + 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.FLAT | pyrender.RenderFlags.RGBA + ) + color = color.astype(np.float32) / 255.0 + # renderer.delete() + + valid_mask = (color[:, :, -1])[:, :, np.newaxis] + semantic_map = ( + color[:, :, :3] * valid_mask + + (1 - valid_mask) * self.semantic_background_rgb + ) + semantic_map = semantic_map.astype(np.float32) + + return { + "Image": color, + "Mask": valid_mask.astype(bool), + "SemanticMap": semantic_map, + "Depth": rend_depth, + } + + 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) \ No newline at end of file