[
{
"path": ".gitignore",
"content": "# Created by .ignore support plugin (hsz.mobi)\n### Python template\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# custom\n.profile\n./*.png\n./*.jpg\n*.out\ndata/base_data\ndata/preprocessed_data\ndata/*/parses\ndata/*/data\ndata/*/*.json\ndata/*/*/*.json\ndata/*/images\ndata*/data\ndata/*/annotation\ndata/*/annotations\ndemo/*.pth.tar\noutput\n\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\n\nlib64/\nparts/\nsdist/\nvar/\nwheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# pyenv\n.python-version\n\n# celery beat schedule file\ncelerybeat-schedule\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n### macOS template\n# General\n.DS_Store\n.AppleDouble\n.LSOverride\n\n# Icon must end with two \\r\nIcon\n\n# Thumbnails\n._*\n\n# Files that might appear in the root of a volume\n.DocumentRevisions-V100\n.fseventsd\n.Spotlight-V100\n.TemporaryItems\n.Trashes\n.VolumeIcon.icns\n.com.apple.timemachine.donotpresent\n\n# Directories potentially created on remote AFP share\n.AppleDB\n.AppleDesktop\nNetwork Trash Folder\nTemporary Items\n.apdisk\n### JetBrains template\n# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and WebStorm\n# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839\n\n# User-specific stuff\n.idea/**/workspace.xml\n.idea/**/tasks.xml\n.idea/**/dictionaries\n.idea/**/shelf\n\n# Sensitive or high-churn files\n.idea/**/dataSources/\n.idea/**/dataSources.ids\n.idea/**/dataSources.local.xml\n.idea/**/sqlDataSources.xml\n.idea/**/dynamic.xml\n.idea/**/uiDesigner.xml\n.idea/**/dbnavigator.xml\n\n# Gradle\n.idea/**/gradle.xml\n.idea/**/libraries\n\n# CMake\ncmake-build-debug/\ncmake-build-release/\n\n# Mongo Explorer plugin\n.idea/**/mongoSettings.xml\n\n# File-based project format\n*.iws\n\n# IntelliJ\nout/\n\n# mpeltonen/sbt-idea plugin\n.idea_modules/\n\n# JIRA plugin\natlassian-ide-plugin.xml\n\n# Cursive Clojure plugin\n.idea/replstate.xml\n\n# Crashlytics plugin (for Android Studio and IntelliJ)\ncom_crashlytics_export_strings.xml\ncrashlytics.properties\ncrashlytics-build.properties\nfabric.properties\n\n# Editor-based Rest Client\n.idea/httpRequests\n\nexperiment/\ndata/PW3D/data/\n"
},
{
"path": "README.md",
"content": "# **3D Clothed Human Reconstruction in the Wild (ClothWild codes)**\n\n> [**3D Clothed Human Reconstruction in the Wild**](https://arxiv.org/abs/2207.10053), \n> Gyeongsik Moon*, \n> Hyeongjin Nam*,\n> Takaaki Shiratori,\n> Kyoung Mu Lee (* equal contribution) \n> *European Conference on Computer Vision (ECCV), 2022* \n\n \n![](\"assets/front_figure.png\")
\n
\n \n![](\"assets/qualitative_result.png\")
\n
\n\n## Installation\n* We recommend you to use an [Anaconda](https://www.anaconda.com/) virtual environment. Install PyTorch >=1.8.0 and Python >= 3.7.0. \n* Install Pytorch3d following [here](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md) depending on your environment.\n* Then, run `sh requirements.sh`. You should slightly change `torchgeometry` kernel code following [here](https://github.com/mks0601/I2L-MeshNet_RELEASE/issues/6#issuecomment-675152527).\n\n\n## Quick demo\n* Download the pre-trained weight from [here](https://drive.google.com/file/d/1-gafc-6V1ma7L5NS1JphzbhTdO_pQFuL/view?usp=sharing) and place it in `demo` folder.\n* Prepare `base_data` folder following below [`Directory`](./assets/directory.md/#required-data) part.\n* Prepare `input.png` and edit its `bbox` of `demo/demo.py`.\n* Prepare SMPL parameter, as `pose2pose_result.json`. You can get the SMPL parameter by running the off-the-shelf method [[code](https://github.com/mks0601/Hand4Whole_RELEASE/tree/Pose2Pose)].\n* Run `python demo.py --gpu 0`.\n\n\n## Directory\nRefer to [here](./assets/directory.md).\n\n\n## Running ClothWild\n\n### Train \nIn the `main/config.py`, you can change datasets to use.\n\n```\ncd ${ROOT}/main\npython train.py --gpu 0\n``` \n\n### Test \nPlace trained model at the `output/model_dump` and follow below.\n\nTo evaluate CD (Chamfer Distance) on 3DPW, run\n``` \ncd ${ROOT}/main\npython test.py --gpu 0 --test_epoch 7 --type cd\n``` \n\nTo evaluate BCC (Body-Cloth Correspondence) on MSCOCO, run\n``` \ncd ${ROOT}/main\npython test.py --gpu 0 --test_epoch 7 --type bcc\n``` \n\nYou can download the checkpoint trained on MSCOCO+DeepFashion2 from [here](https://drive.google.com/file/d/1-gafc-6V1ma7L5NS1JphzbhTdO_pQFuL/view?usp=sharing).\n\n\n## Result\nRefer to the [paper](https://arxiv.org/abs/2207.10053)'s main manuscript and supplementary material for diverse qualitative results!\n\n### Chamfer Distance (CD)\n\n \n![](\"assets/chamfer_distance.png\")
\n
\n\n### Body-Cloth Correspondence (BCC)\n\n \n![](\"assets/body_cloth_correspondence.png\")
\n
\n\n## Reference \n``` \n@InProceedings{Moon_2022_ECCV_ClothWild, \nauthor = {Moon, Gyeongsik and Nam, Hyeongjin and Shiratori, Takaaki and Lee, Kyoung Mu}, \ntitle = {3D Clothed Human Reconstruction in the Wild}, \nbooktitle = {European Conference on Computer Vision (ECCV)}, \nyear = {2022} \n} \n```\n"
},
{
"path": "assets/directory.md",
"content": "# Directory \n## Root\nThe `${ROOT}` is described as below. \n``` \n${ROOT} \n|-- common \n|-- data \n|-- demo\n|-- main \n|-- output \n``` \n* `common` contains kernel codes for ClothWild. \n* `data` contains required data and soft links to images and annotations directories. \n* `demo` contains demo codes.\n* `main` contains high-level codes for training or testing the network. \n* `output` contains log, trained models, visualized outputs, and test result. \n\n\n## Required data\nYou need to follow directory structure of the `data` as below. \n```\n${ROOT} \n|-- data \n| |-- base_data\n| | |-- human_models\n| | | |-- SMPL_FEMALE.pkl\n| | | |-- SMPL_MALE.pkl\n| | | |-- SMPL_NEUTRAL.pkl\n| | |-- smplicit\n| | | |-- checkpoints\n| | | | |-- hair.pth\n| | | | |-- pants.pth\n| | | | |-- shoes.pth\n| | | | |-- skirts.pth\n| | | | |-- upperclothes.pth\n| | | |-- clusters\n| | | | |-- clusters_hairs.npy\n| | | | |-- clusters_lowerbody.npy\n| | | | |-- clusters_shoes.npy\n| | | | |-- indexs_clusters_tshirt_smpl.npy\n| |-- preprocessed_data\n| | |-- densepose\n| | |-- gender\n| | |-- parse\n| | |-- smpl_param\n| |-- ...\n```\n* `base_data/human_model_files` contains `smpl` 3D model files. Download the files from [[smpl](https://smpl.is.tue.mpg.de/)].\n* `base_data/smplicit` contains 3D cloth generative model (`SMPLicit`) files. Download the files from [[smplicit](https://github.com/enriccorona/SMPLicit)].\n* `preprocessed_data` is required for training and testing stages. Download it from [[preprocessed_data](https://drive.google.com/file/d/1m5AfZt2qx90DkbXACuDK8D1tWl7LuPKe/view?usp=sharing)].\n\n\n## Dataset\n``` \n${ROOT} \n|-- data \n| |-- ...\n| |-- DeepFashion2\n| | |-- data\n| | | |-- train\n| | | |-- DeepFashion2_train.json\n| |-- MSCOCO\n| | |-- images\n| | | |-- train2017\n| | | |-- val2017\n| | |-- parses\n| | |-- annotations\n| | | |-- coco_wholebody_train_v1.0.json\n| | | |-- coco_wholebody_val_v1.0.json\n| | | |-- coco_dp_train.json\n| | | |-- coco_dp_val.json\n| |-- PW3D\n| | |-- data\n| | | |-- imageFiles\n| | | |-- sequenceFiles\n| | | |-- 3DPW_test.json\n``` \n* Download DeepFashion2 parsed data [[data](https://github.com/switchablenorms/DeepFashion2)] [[annot](https://drive.google.com/drive/folders/1P2AgxZZsq21fcGnP_RNuvaEOrye1SBkj?usp=sharing)]\n* Download MSCOCO data, parses (LIP dataset), and densepose [[data](https://github.com/jin-s13/COCO-WholeBody)] [[parses](https://drive.google.com/file/d/1i2-qbNPTtn2kmxko--riEJjBehZnvumw/view?usp=sharing)] [[densepose](https://drive.google.com/drive/folders/1P2AgxZZsq21fcGnP_RNuvaEOrye1SBkj?usp=sharing)]\n* Download 3DPW parsed data [[data](https://virtualhumans.mpi-inf.mpg.de/3DPW/)] [[annot]](https://drive.google.com/drive/folders/1P2AgxZZsq21fcGnP_RNuvaEOrye1SBkj?usp=sharing)\n* All annotation files follow [MSCOCO format](http://cocodataset.org/#format-data). If you want to add your own dataset, you have to convert it to [MSCOCO format](http://cocodataset.org/#format-data). \n\n\n### Output\n``` \n${ROOT} \n|-- output \n| |-- log \n| |-- model_dump \n| |-- result \n| |-- vis \n``` \n* Creating `output` folder as soft link form is recommended instead of folder form because it would take large storage capacity. \n* `log` folder contains training log file. \n* `model_dump` folder contains saved checkpoints for each epoch. \n* `result` folder contains final estimation files generated in the testing stage. \n* `vis` folder contains visualized results. "
},
{
"path": "common/base.py",
"content": "import os\nimport os.path as osp\nimport math\nimport time\nimport glob\nimport abc\nimport numpy as np\nfrom torch.utils.data import DataLoader\nimport torch.optim\nimport torchvision.transforms as transforms\nfrom collections import OrderedDict\nfrom timer import Timer\nfrom logger import colorlogger\nfrom torch.nn.parallel.data_parallel import DataParallel\n\nfrom config import cfg\nfrom utils.dir import make_folder\nfrom model import get_model\nfrom dataset import MultipleDatasets\nfrom utils.human_models import smpl\n\n# dynamic dataset import\nfor i in range(len(cfg.trainset_2d)):\n exec('from ' + cfg.trainset_2d[i] + ' import ' + cfg.trainset_2d[i])\nfor i in range(len(cfg.testset)):\n exec('from ' + cfg.testset[i] + ' import ' + cfg.testset[i])\n\ndef worker_init_fn(worder_id):\n np.random.seed(np.random.get_state()[1][0] + worder_id)\n\nclass Base(object):\n __metaclass__ = abc.ABCMeta\n\n def __init__(self, log_name='logs.txt'):\n self.cur_epoch = 0\n\n # timer\n self.tot_timer = Timer()\n self.gpu_timer = Timer()\n self.read_timer = Timer()\n\n # logger\n self.logger = colorlogger(cfg.log_dir, log_name=log_name)\n\n @abc.abstractmethod\n def _make_batch_generator(self):\n return\n\n @abc.abstractmethod\n def _make_model(self):\n return\n\nclass Trainer(Base):\n def __init__(self):\n super(Trainer, self).__init__(log_name = 'train_logs.txt')\n\n def get_optimizer(self, model):\n total_params = []\n for module in model.module.trainable_modules:\n total_params += list(module.parameters())\n optimizer = torch.optim.Adam(total_params, lr=cfg.lr)\n return optimizer\n\n def save_model(self, state, epoch):\n file_path = osp.join(cfg.model_dir,'snapshot_{}.pth.tar'.format(str(epoch)))\n\n # do not save smpl & smplicit layer weights\n dump_key = []\n for k in state['network'].keys():\n if 'smpl_layer' in k:\n dump_key.append(k)\n if 'smplicit_layer' in k:\n dump_key.append(k)\n for k in dump_key:\n state['network'].pop(k, None)\n\n torch.save(state, file_path)\n self.logger.info(\"Write snapshot into {}\".format(file_path))\n\n def load_model(self, model, optimizer):\n model_file_list = glob.glob(osp.join(cfg.model_dir,'*.pth.tar'))\n cur_epoch = max([int(file_name[file_name.find('snapshot_') + 9 : file_name.find('.pth.tar')]) for file_name in model_file_list])\n ckpt_path = osp.join(cfg.model_dir, 'snapshot_' + str(cur_epoch) + '.pth.tar')\n ckpt = torch.load(ckpt_path) \n start_epoch = ckpt['epoch'] + 1\n model.load_state_dict(ckpt['network'], strict=False)\n\n print(\"cur_epoch: \", cur_epoch)\n self.logger.info('Load checkpoint from {}'.format(ckpt_path))\n return start_epoch, model, optimizer\n\n def set_lr(self, epoch):\n for e in cfg.lr_dec_epoch:\n if epoch < e:\n break\n if epoch < cfg.lr_dec_epoch[-1]:\n idx = cfg.lr_dec_epoch.index(e)\n for g in self.optimizer.param_groups:\n g['lr'] = cfg.lr / (cfg.lr_dec_factor ** idx)\n else:\n for g in self.optimizer.param_groups:\n g['lr'] = cfg.lr / (cfg.lr_dec_factor ** len(cfg.lr_dec_epoch))\n\n def get_lr(self):\n for g in self.optimizer.param_groups:\n cur_lr = g['lr']\n return cur_lr\n \n def _make_batch_generator(self):\n # data load and construct batch generator\n self.logger.info(\"Creating dataset...\")\n trainset3d_loader = []\n for i in range(len(cfg.trainset_3d)):\n trainset3d_loader.append(eval(cfg.trainset_3d[i])(transforms.ToTensor(), \"train\"))\n trainset2d_loader = []\n for i in range(len(cfg.trainset_2d)):\n trainset2d_loader.append(eval(cfg.trainset_2d[i])(transforms.ToTensor(), \"train\"))\n \n valid_loader_num = 0\n if len(trainset3d_loader) > 0:\n trainset3d_loader = [MultipleDatasets(trainset3d_loader, make_same_len=False)]\n valid_loader_num += 1\n else:\n trainset3d_loader = []\n if len(trainset2d_loader) > 0:\n trainset2d_loader = [MultipleDatasets(trainset2d_loader, make_same_len=False)]\n valid_loader_num += 1\n else:\n trainset2d_loader = []\n\n if valid_loader_num > 1:\n trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader, make_same_len=True)\n else:\n trainset_loader = MultipleDatasets(trainset3d_loader + trainset2d_loader, make_same_len=False)\n\n self.itr_per_epoch = math.ceil(len(trainset_loader) / cfg.num_gpus / cfg.train_batch_size)\n self.batch_generator = DataLoader(dataset=trainset_loader, batch_size=cfg.num_gpus*cfg.train_batch_size, shuffle=True, num_workers=cfg.num_thread, pin_memory=True, drop_last=True, worker_init_fn=worker_init_fn)\n\n def _make_model(self):\n # prepare network\n self.logger.info(\"Creating graph and optimizer...\")\n model = get_model('train')\n model = DataParallel(model).cuda()\n optimizer = self.get_optimizer(model)\n if cfg.continue_train:\n start_epoch, model, optimizer = self.load_model(model, optimizer)\n else:\n start_epoch = 0\n model.train()\n\n self.start_epoch = start_epoch\n self.model = model\n self.optimizer = optimizer\n\nclass Tester(Base):\n def __init__(self, test_epoch):\n self.test_epoch = int(test_epoch)\n super(Tester, self).__init__(log_name = 'test_logs.txt')\n\n def _make_batch_generator(self):\n # data load and construct batch generator\n self.logger.info(\"Creating dataset...\")\n testset_loader = eval(cfg.testset[0])(transforms.ToTensor(), \"test\")\n batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus*cfg.test_batch_size, shuffle=False, num_workers=cfg.num_thread, pin_memory=True, worker_init_fn=worker_init_fn)\n \n self.testset = testset_loader\n self.batch_generator = batch_generator\n\n def _make_model(self):\n model_path = os.path.join(cfg.model_dir, 'snapshot_%d.pth.tar' % self.test_epoch)\n assert os.path.exists(model_path), 'Cannot find model at ' + model_path\n self.logger.info('Load checkpoint from {}'.format(model_path))\n \n # prepare network \n model = get_model('test')\n model = DataParallel(model).cuda()\n ckpt = torch.load(model_path)\n model.load_state_dict(ckpt['network'], strict=False)\n model.eval()\n\n self.model = model\n \n def _evaluate(self, outs, cur_sample_idx):\n eval_result = self.testset.evaluate(outs, cur_sample_idx)\n return eval_result\n\n def _print_eval_result(self, eval_result):\n self.testset.print_eval_result(eval_result)\n\ndef check_data_parallel(train_weight):\n new_state_dict = OrderedDict()\n for k, v in train_weight.items():\n name = k[7:] if k.startswith('module') else k # remove `module.`\n new_state_dict[name] = v\n return new_state_dict"
},
{
"path": "common/logger.py",
"content": "import logging\nimport os\n\nOK = '\\033[92m'\nWARNING = '\\033[93m'\nFAIL = '\\033[91m'\nEND = '\\033[0m'\n\nPINK = '\\033[95m'\nBLUE = '\\033[94m'\nGREEN = OK\nRED = FAIL\nWHITE = END\nYELLOW = WARNING\n\nclass colorlogger():\n def __init__(self, log_dir, log_name='train_logs.txt'):\n # set log\n self._logger = logging.getLogger(log_name)\n self._logger.setLevel(logging.INFO)\n log_file = os.path.join(log_dir, log_name)\n if not os.path.exists(log_dir):\n os.makedirs(log_dir)\n file_log = logging.FileHandler(log_file, mode='a')\n file_log.setLevel(logging.INFO)\n console_log = logging.StreamHandler()\n console_log.setLevel(logging.INFO)\n formatter = logging.Formatter(\n \"{}%(asctime)s{} %(message)s\".format(GREEN, END),\n \"%m-%d %H:%M:%S\")\n file_log.setFormatter(formatter)\n console_log.setFormatter(formatter)\n self._logger.addHandler(file_log)\n self._logger.addHandler(console_log)\n\n def debug(self, msg):\n self._logger.debug(str(msg))\n\n def info(self, msg):\n self._logger.info(str(msg))\n\n def warning(self, msg):\n self._logger.warning(WARNING + 'WRN: ' + str(msg) + END)\n\n def critical(self, msg):\n self._logger.critical(RED + 'CRI: ' + str(msg) + END)\n\n def error(self, msg):\n self._logger.error(RED + 'ERR: ' + str(msg) + END)\n\n"
},
{
"path": "common/nets/layer.py",
"content": "import torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nimport math\nfrom config import cfg\n\ndef make_linear_layers(feat_dims, relu_final=True, use_bn=False):\n layers = []\n for i in range(len(feat_dims)-1):\n layers.append(nn.Linear(feat_dims[i], feat_dims[i+1]))\n\n # Do not use ReLU for final estimation\n if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and relu_final):\n if use_bn:\n layers.append(nn.BatchNorm1d(feat_dims[i+1]))\n layers.append(nn.ReLU(inplace=True))\n\n return nn.Sequential(*layers)\n\ndef make_conv_layers(feat_dims, kernel=3, stride=1, padding=1, bnrelu_final=True):\n layers = []\n for i in range(len(feat_dims)-1):\n layers.append(\n nn.Conv2d(\n in_channels=feat_dims[i],\n out_channels=feat_dims[i+1],\n kernel_size=kernel,\n stride=stride,\n padding=padding\n ))\n # Do not use BN and ReLU for final estimation\n if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final):\n layers.append(nn.BatchNorm2d(feat_dims[i+1]))\n layers.append(nn.ReLU(inplace=True))\n\n return nn.Sequential(*layers)\n\ndef make_deconv_layers(feat_dims, bnrelu_final=True):\n layers = []\n for i in range(len(feat_dims)-1):\n layers.append(\n nn.ConvTranspose2d(\n in_channels=feat_dims[i],\n out_channels=feat_dims[i+1],\n kernel_size=4,\n stride=2,\n padding=1,\n output_padding=0,\n bias=False))\n\n # Do not use BN and ReLU for final estimation\n if i < len(feat_dims)-2 or (i == len(feat_dims)-2 and bnrelu_final):\n layers.append(nn.BatchNorm2d(feat_dims[i+1]))\n layers.append(nn.ReLU(inplace=True))\n\n return nn.Sequential(*layers)\n"
},
{
"path": "common/nets/loss.py",
"content": "import torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\n\nfrom utils.human_models import smpl\nfrom utils.vis import save_obj\nfrom config import cfg\n\nclass ClothClsLoss(nn.Module):\n def __init__(self):\n super(ClothClsLoss, self).__init__()\n self.dp_parts = {\n 'head': [22,23],\n 'upperbody': [0,1,14,15,16,17,18,19,20,21],\n 'lowerbody': [6,7,8,9,10,11,12,13],\n 'foot': [4,5]\n }\n self.part_clothes={\n 'head': ['hair'],\n 'upperbody': ['uppercloth', 'coat'],\n 'lowerbody': ['pants', 'skirts'],\n 'foot': ['shoes']\n }\n self.bce_loss = nn.BCELoss(reduction='none')\n\n def forward(self, out, patch_idx, cloth_idx):\n # valid only on visible\n valid = torch.zeros_like(out).cuda()\n index_gt = torch.zeros_like(out).cuda()\n\n for part in self.dp_parts.keys():\n valid_one_part = torch.zeros((out.shape[0],)).cuda()\n\n for part_idx in self.dp_parts[part]:\n valid_one_part += (patch_idx == part_idx).any(1)\n\n for cloth in self.part_clothes[part]:\n if cloth in cfg.cloth_types:\n valid[valid_one_part>0, cfg.cloth_types.index(cloth)] = 1\n\n for idx in range(len(cfg.cloth_types)):\n index_gt[:, idx] += (cloth_idx==idx+1).any(1)\n \n loss = self.bce_loss(out, index_gt)\n loss = loss[valid>0]\n return loss.mean()\n\n\nclass GenderClsLoss(nn.Module):\n def __init__(self):\n super(GenderClsLoss, self).__init__()\n self.bce_loss = nn.BCELoss(reduction='none')\n\n def forward(self, out, gt):\n valid = (gt != 0) # if neutral gender, set valid = 0\n gt = F.one_hot((gt.long()), num_classes=3)[:,1:].float()\n \n loss = self.bce_loss(out, gt)\n loss = loss[valid]\n return loss.mean()\n\n\nclass SdfDPLoss(nn.Module):\n def __init__(self):\n super(SdfDPLoss, self).__init__()\n\n def forward(self, sdf, cloth_meshes_unposed, smpl_cloth_idx, smpl_cloth_valid, cloth_idx, sdf_thresh, dist_thresh, v_template):\n batch_size = sdf.shape[0]\n cloth_type = cfg.cloth_types[cloth_idx[0]-1]\n\n loss_list = []\n for bid in range(batch_size):\n smpl_mask = smpl_cloth_valid[bid] > 0\n smpl_verts = v_template[bid][smpl_mask[:,None].repeat(1,3)].view(-1,3)\n cloth_verts = cloth_meshes_unposed[bid]\n\n if smpl_verts.shape[0] > 0:\n dists = torch.sqrt(torch.sum((smpl_verts[None,:,:] - cloth_verts[:,None,:])**2,2))\n else:\n loss_list.append(torch.zeros((1)).mean().float().cuda())\n continue\n\n # remove too closest query points\n dists[dists 0) * (dists < dist_thresh)\n loss_neg = torch.abs(sdf[bid,:] - sdf_thresh) * (sum([target_cloth_idx == idx for idx in cloth_idx]) == 0) * (dists < dist_thresh)\n\n cloth_exist = (sum([target_cloth_idx == idx for idx in cloth_idx]) > 0).sum() > 0\n loss = (loss_pos + loss_neg).mean() * cloth_exist\n loss_list.append(loss)\n \n loss = torch.stack(loss_list)\n return loss\n\nclass RegLoss(nn.Module):\n def __init__(self):\n super(RegLoss, self).__init__()\n self.l2_loss = nn.MSELoss(reduction='none')\n\n def forward(self, param, valid):\n zeros = torch.zeros_like(param).cuda()\n loss = self.l2_loss(param, zeros) * valid[:,None]\n return loss.mean()\n\nclass SdfParseLoss(nn.Module):\n def __init__(self):\n super(SdfParseLoss, self).__init__()\n\n def forward(self, sdf, cloth_meshes, parse_gt, sdf_thresh, cloth_meshes_unposed, parse_valid, dist_thresh, v_template):\n batch_size = sdf.shape[0]\n inf = 9999\n \n # mask invalid xy coordinatets\n x, y = cloth_meshes[:,:,0].long(), cloth_meshes[:,:,1].long()\n idx = y * cfg.input_img_shape[1] + x\n is_valid = (x >= 0) * (x < cfg.input_img_shape[1]) * (y >= 0) * (y < cfg.input_img_shape[0])\n idx[is_valid == 0] = 0\n\n # minimum sdf\n min_sdf = sdf * is_valid.float() + inf * (1 - is_valid.float())\n parse_out_min = torch.ones((batch_size, cfg.input_img_shape[0] * cfg.input_img_shape[1])).float().cuda() * inf\n \n # maximum sdf\n max_sdf = sdf * is_valid.float() - inf * (1 - is_valid.float())\n parse_out_max = torch.ones((batch_size, cfg.input_img_shape[0] * cfg.input_img_shape[1])).float().cuda() * -inf\n\n try:\n parse_out_min, _ = scatter_min(min_sdf, idx, 1, parse_out_min)\n parse_out_max, _ = scatter_max(max_sdf, idx, 1, parse_out_max)\n except:\n # some GPUs have trouble in torch_scatter, compute in CPU\n idx = idx.cpu()\n min_sdf, max_sdf = min_sdf.cpu(), max_sdf.cpu()\n parse_out_min, parse_out_max = parse_out_min.cpu(), parse_out_max.cpu()\n parse_out_min, _ = scatter_min(min_sdf, idx, 1, parse_out_min)\n parse_out_max, _ = scatter_max(max_sdf, idx, 1, parse_out_max)\n parse_out_min, parse_out_max = parse_out_min.cuda(), parse_out_max.cuda()\n\n parse_out_min = parse_out_min.view(batch_size, cfg.input_img_shape[0], cfg.input_img_shape[1])\n parse_out_min[parse_out_min == inf] = 0\n\n parse_out_max = parse_out_max.view(batch_size, cfg.input_img_shape[0], cfg.input_img_shape[1])\n parse_out_max[parse_out_max == -inf] = sdf_thresh\n\n loss_pos = torch.abs(parse_out_min) * (parse_gt == 1) * parse_valid\n loss_neg = torch.abs(parse_out_max - sdf_thresh) * (parse_gt == 0) * parse_valid\n loss = loss_pos.mean((1,2)) + loss_neg.mean((1,2))\n\n cloth_exist = (parse_gt == 1).sum((1,2)) > 0\n loss = loss * cloth_exist \n return loss"
},
{
"path": "common/nets/module.py",
"content": "import torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nfrom nets.layer import make_linear_layers, make_conv_layers, make_deconv_layers\nfrom utils.human_models import smpl\nfrom config import cfg\n\nclass ClothNet(nn.Module):\n def __init__(self):\n super(ClothNet, self).__init__()\n input_feat_dim = 2048\n if 'uppercloth' in cfg.cloth_types:\n self.z_cut_uppercloth = make_linear_layers([input_feat_dim,6], relu_final=False)\n self.z_style_uppercloth = make_linear_layers([input_feat_dim,12], relu_final=False)\n if 'coat' in cfg.cloth_types:\n self.z_cut_coat = make_linear_layers([input_feat_dim,6], relu_final=False)\n self.z_style_coat = make_linear_layers([input_feat_dim,12], relu_final=False)\n if 'pants' in cfg.cloth_types:\n self.z_cut_pants = make_linear_layers([input_feat_dim,6], relu_final=False)\n self.z_style_pants = make_linear_layers([input_feat_dim,12], relu_final=False)\n if 'skirts' in cfg.cloth_types:\n self.z_cut_skirts = make_linear_layers([input_feat_dim,6], relu_final=False)\n self.z_style_skirts = make_linear_layers([input_feat_dim,12], relu_final=False)\n if 'hair' in cfg.cloth_types:\n self.z_cut_hair = make_linear_layers([input_feat_dim,6], relu_final=False)\n self.z_style_hair = make_linear_layers([input_feat_dim,12], relu_final=False)\n if 'shoes' in cfg.cloth_types:\n self.z_style_shoes = make_linear_layers([input_feat_dim,4], relu_final=False)\n\n self.cloth_cls_layer = make_linear_layers([input_feat_dim, len(cfg.cloth_types)], relu_final=False)\n self.gender_cls_layer = make_linear_layers([input_feat_dim, 2], relu_final=False)\n \n\n def forward(self, img_feat):\n batch_size = img_feat.shape[0]\n img_feat = img_feat.mean((2,3))\n \n z_cuts, z_styles = [], []\n for cloth_type in cfg.cloth_types:\n if cloth_type == 'uppercloth':\n z_cuts.append(self.z_cut_uppercloth(img_feat))\n z_styles.append(self.z_style_uppercloth(img_feat))\n elif cloth_type == 'coat':\n z_cuts.append(self.z_cut_coat(img_feat))\n z_styles.append(self.z_style_coat(img_feat))\n elif cloth_type == 'pants':\n z_cuts.append(self.z_cut_pants(img_feat))\n z_styles.append(self.z_style_pants(img_feat))\n elif cloth_type == 'skirts':\n z_cuts.append(self.z_cut_skirts(img_feat))\n z_styles.append(self.z_style_skirts(img_feat))\n elif cloth_type == 'hair':\n z_cuts.append(self.z_cut_hair(img_feat))\n z_styles.append(self.z_style_hair(img_feat))\n elif cloth_type == 'shoes':\n z_cuts.append(torch.zeros((batch_size,0)).float().cuda())\n z_styles.append(self.z_style_shoes(img_feat))\n\n scores = self.cloth_cls_layer(img_feat)\n scores = torch.sigmoid(scores)\n\n genders = self.gender_cls_layer(img_feat)\n genders = F.softmax(genders, dim=-1)\n\n return genders, scores, z_cuts, z_styles\n"
},
{
"path": "common/nets/resnet.py",
"content": "import torch\nimport torch.nn as nn\nfrom torchvision.models.resnet import BasicBlock, Bottleneck\nfrom torchvision.models.resnet import model_urls\n\nclass ResNetBackbone(nn.Module):\n\n def __init__(self, resnet_type):\n\t\n resnet_spec = {18: (BasicBlock, [2, 2, 2, 2], [64, 64, 128, 256, 512], 'resnet18'),\n\t\t 34: (BasicBlock, [3, 4, 6, 3], [64, 64, 128, 256, 512], 'resnet34'),\n\t\t 50: (Bottleneck, [3, 4, 6, 3], [64, 256, 512, 1024, 2048], 'resnet50'),\n\t\t 101: (Bottleneck, [3, 4, 23, 3], [64, 256, 512, 1024, 2048], 'resnet101'),\n\t\t 152: (Bottleneck, [3, 8, 36, 3], [64, 256, 512, 1024, 2048], 'resnet152')}\n block, layers, channels, name = resnet_spec[resnet_type]\n \n self.name = name\n self.inplanes = 64\n super(ResNetBackbone, self).__init__()\n self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,\n bias=False)\n self.bn1 = nn.BatchNorm2d(64)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n nn.init.normal_(m.weight, mean=0, std=0.001)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.inplanes, planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes))\n\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n return x\n\n def init_weights(self):\n org_resnet = torch.utils.model_zoo.load_url(model_urls[self.name])\n # drop orginal resnet fc layer, add 'None' in case of no fc layer, that will raise error\n org_resnet.pop('fc.weight', None)\n org_resnet.pop('fc.bias', None)\n \n self.load_state_dict(org_resnet)\n print(\"Initialize resnet from model zoo\")\n\n\n"
},
{
"path": "common/timer.py",
"content": "# --------------------------------------------------------\r\n# Fast R-CNN\r\n# Copyright (c) 2015 Microsoft\r\n# Licensed under The MIT License [see LICENSE for details]\r\n# Written by Ross Girshick\r\n# --------------------------------------------------------\r\n\r\nimport time\r\n\r\nclass Timer(object):\r\n \"\"\"A simple timer.\"\"\"\r\n def __init__(self):\r\n self.total_time = 0.\r\n self.calls = 0\r\n self.start_time = 0.\r\n self.diff = 0.\r\n self.average_time = 0.\r\n self.warm_up = 0\r\n\r\n def tic(self):\r\n # using time.time instead of time.clock because time time.clock\r\n # does not normalize for multithreading\r\n self.start_time = time.time()\r\n\r\n def toc(self, average=True):\r\n self.diff = time.time() - self.start_time\r\n if self.warm_up < 10:\r\n self.warm_up += 1\r\n return self.diff\r\n else:\r\n self.total_time += self.diff\r\n self.calls += 1\r\n self.average_time = self.total_time / self.calls\r\n\r\n if average:\r\n return self.average_time\r\n else:\r\n return self.diff\r\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/SMPL.py",
"content": "import torch\nimport json\nimport sys\nimport numpy as np\nfrom .util_smpl import batch_global_rigid_transformation, batch_rodrigues, reflect_pose\nimport torch.nn as nn\nimport os\nimport trimesh\nimport pickle\n#from utils.human_models import smpl\n#from utils.vis import vis_keypoints, vis_mesh, save_obj, vis_parse, vis_dp\n\nclass SMPL(nn.Module):\n def __init__(self, model_path, joint_type = 'cocoplus', obj_saveable = False):\n super(SMPL, self).__init__()\n\n if joint_type not in ['cocoplus', 'lsp']:\n msg = 'unknow joint type: {}, it must be either \"cocoplus\" or \"lsp\"'.format(joint_type)\n sys.exit(msg)\n\n self.model_path = model_path\n self.joint_type = joint_type\n \n with open(model_path, 'rb') as reader:\n model = pickle.load(reader, encoding='latin1')\n \n if obj_saveable:\n self.faces = model['f']\n else:\n self.faces = None\n\n np_v_template = np.array(model['v_template'], dtype = np.float)\n self.register_buffer('v_template', torch.from_numpy(np_v_template).float())\n self.size = [np_v_template.shape[0], 3]\n\n np_shapedirs = np.array(model['shapedirs'], dtype = np.float)[:,:,:10]\n self.num_betas = np_shapedirs.shape[-1]\n np_shapedirs = np.reshape(np_shapedirs, [-1, self.num_betas]).T\n self.register_buffer('shapedirs', torch.from_numpy(np_shapedirs).float())\n\n np_J_regressor = np.array(model['J_regressor'].toarray().transpose(1,0), dtype = np.float)\n self.register_buffer('J_regressor', torch.from_numpy(np_J_regressor).float())\n\n np_posedirs = np.array(model['posedirs'], dtype = np.float)\n num_pose_basis = np_posedirs.shape[-1]\n np_posedirs = np.reshape(np_posedirs, [-1, num_pose_basis]).T\n self.register_buffer('posedirs', torch.from_numpy(np_posedirs).float())\n\n self.parents = np.array(model['kintree_table'])[0].astype(np.int32)\n\n np_weights = np.array(model['weights'], dtype = np.float)\n\n vertex_count = np_weights.shape[0] \n vertex_component = np_weights.shape[1]\n\n self.register_buffer('weight', torch.from_numpy(np_weights).float().reshape(-1, vertex_count, vertex_component))\n self.register_buffer('e3', torch.eye(3).float())\n \n self.cur_device = None\n\n def save_obj(self, verts, obj_mesh_name):\n if not self.faces:\n msg = 'obj not saveable!'\n sys.exit(msg)\n\n with open(obj_mesh_name, 'w') as fp:\n for v in verts:\n fp.write( 'v {:f} {:f} {:f}\\n'.format( v[0], v[1], v[2]) )\n\n for f in self.faces: # Faces are 1-based, not 0-based in obj files\n fp.write( 'f {:d} {:d} {:d}\\n'.format(f[0] + 1, f[1] + 1, f[2] + 1) )\n\n def forward(self, beta, theta, get_skin = False, theta_in_rodrigues=True):\n device, dtype = beta.device, beta.dtype\n self.cur_device = torch.device(device.type, device.index)\n num_batch = beta.shape[0]\n\n v_shaped = torch.matmul(beta, self.shapedirs).view(-1, self.size[0], self.size[1]) + self.v_template\n Jx = torch.matmul(v_shaped[:, :, 0], self.J_regressor)\n Jy = torch.matmul(v_shaped[:, :, 1], self.J_regressor)\n Jz = torch.matmul(v_shaped[:, :, 2], self.J_regressor)\n J = torch.stack([Jx, Jy, Jz], dim = 2)\n if theta_in_rodrigues:\n Rs = batch_rodrigues(theta.view(-1, 3)).view(-1, 24, 3, 3)\n else: #theta is already rotations\n Rs = theta.view(-1,24,3,3)\n\n pose_feature = (Rs[:, 1:, :, :] - torch.eye(3, dtype=dtype, device=device)).view(-1, 207)\n v_posed = torch.matmul(pose_feature, self.posedirs).view(-1, self.size[0], self.size[1]) + v_shaped\n J_transformed, A = batch_global_rigid_transformation(Rs, J, self.parents, rotate_base = False)\n\n W=self.weight.expand(num_batch,*self.weight.shape[1:])\n T = torch.matmul(W, A.view(num_batch, 24, 16)).view(num_batch, -1, 4, 4)\n \n #v_posed_homo = torch.cat([v_posed, torch.ones(num_batch, v_posed.shape[1], 1, device = self.cur_device)], dim = 2)\n v_posed_homo = torch.cat([v_posed, torch.ones(num_batch, v_posed.shape[1], 1, device=device)], dim = 2)\n v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, -1))\n\n verts = v_homo[:, :, :3, 0]\n\n joint_x = torch.matmul(verts[:, :, 0], self.J_regressor)\n joint_y = torch.matmul(verts[:, :, 1], self.J_regressor)\n joint_z = torch.matmul(verts[:, :, 2], self.J_regressor)\n\n joints = torch.stack([joint_x, joint_y, joint_z], dim = 2)\n\n if get_skin:\n return verts, joints, Rs\n else:\n return joints\n\n def deform_clothed_smpl(self, theta, J, v_smpl, v_cloth):\n num_batch = theta.shape[0]\n\n device = theta.device\n self.cur_device = torch.device(device.type, device.index)\n \n Rs = batch_rodrigues(theta.view(-1, 3)).view(-1, 24, 3, 3)\n pose_feature = (Rs[:, 1:, :, :] - torch.eye(3, device=device).float()).view(-1, 207)\n\n pose_params = torch.matmul(pose_feature, self.posedirs).view(-1, self.size[0], self.size[1])\n v_posed_smpl = pose_params + v_smpl\n \n # Calculate closest SMPL vertex for each vertex of the cloth mesh\n with torch.no_grad():\n dists = ((v_smpl.unsqueeze(1) - v_cloth.unsqueeze(2))**2).sum(-1)\n dists, correspondance = torch.min(dists, 2) # num_batch, v_cloth.shape[1]\n\n v_posed_cloth = torch.gather(pose_params, 1, correspondance[:,:,None].repeat(1,1,3)) + v_cloth\n J_transformed, A = batch_global_rigid_transformation(Rs, J, self.parents, rotate_base = False)\n\n W = self.weight.expand(num_batch,*self.weight.shape[1:])\n T = torch.matmul(W, A.view(num_batch, 24, 16)).view(num_batch, -1, 4, 4)\n\n v_posed_homo_smpl = torch.cat([v_posed_smpl, torch.ones(num_batch, v_posed_smpl.shape[1], 1, device=device)], dim = 2)\n v_posed_homo_cloth = torch.cat([v_posed_cloth, torch.ones(num_batch, v_posed_cloth.shape[1], 1, device=device)], dim = 2)\n v_homo_smpl = torch.matmul(T, torch.unsqueeze(v_posed_homo_smpl, -1))\n v_homo_cloth = torch.matmul(torch.gather(T, 1, correspondance[:,:,None,None].repeat(1,1,4,4)), torch.unsqueeze(v_posed_homo_cloth, -1))\n verts_smpl = v_homo_smpl[:, :, :3, 0]\n verts_cloth = v_homo_cloth[:, :, :3, 0]\n return verts_cloth\n \n def unpose_and_deform_cloth(self, v_cloth_posed, theta_from, theta_to, beta, Jsmpl, vsmpl, theta_in_rodrigues=True):\n ### UNPOSE:\n device = theta_from.device\n self.cur_device = torch.device(device.type, device.index)\n num_batch = beta.shape[0]\n\n v_shaped = torch.matmul(beta, self.shapedirs).view(-1, self.size[0], self.size[1]) + self.v_template\n Jx = torch.matmul(v_shaped[:, :, 0], self.J_regressor)\n Jy = torch.matmul(v_shaped[:, :, 1], self.J_regressor)\n Jz = torch.matmul(v_shaped[:, :, 2], self.J_regressor)\n J = torch.stack([Jx, Jy, Jz], dim = 2)\n if theta_in_rodrigues:\n Rs = batch_rodrigues(theta_from.view(-1, 3)).view(-1, 24, 3, 3)\n else: #theta is already rotations\n Rs = theta_from.view(-1,24,3,3)\n\n pose_feature = (Rs[:, 1:, :, :] - torch.eye(3, device=device).float()).view(-1, 207)\n\n pose_displ = torch.matmul(pose_feature, self.posedirs).view(-1, self.size[0], self.size[1])\n v_posed = pose_displ + v_shaped\n J_transformed, A = batch_global_rigid_transformation(Rs, J, self.parents, rotate_base = False)\n\n W = self.weight.expand(num_batch,*self.weight.shape[1:])\n T = torch.matmul(W, A.view(num_batch, 24, 16)).view(num_batch, -1, 4, 4)\n\n v_posed_homo = torch.cat([v_posed, torch.ones(num_batch, v_posed.shape[1], 1, device=device)], dim = 2)\n v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, -1))\n\n v_smpl = v_homo[:, :, :3, 0]\n with torch.no_grad():\n dists = ((v_smpl.unsqueeze(1) - v_cloth_posed.unsqueeze(2))**2).sum(-1)\n dists, correspondance = torch.min(dists, 2) # num_batch, v_cloth_posed.shape[1]\n\n invT = torch.inverse(torch.gather(T, 1, correspondance[:,:,None,None].repeat(1,1,4,4)).view(num_batch,-1,4,4))\n v = torch.cat([v_cloth_posed, torch.ones(num_batch, v_cloth_posed.shape[1], 1, device=device)], 2)\n v = torch.matmul(invT, v.unsqueeze(-1))[:,:, :3, 0]\n unposed_v = v - torch.gather(pose_displ, 1, correspondance[:,:,None].repeat(1,1,3))\n \n ### REPOSE:\n Rs = batch_rodrigues(theta_to.view(-1, 3)).view(-1, 24, 3, 3)\n pose_feature = (Rs[:, 1:, :, :] - torch.eye(3, device=device).float()).view(-1, 207)\n\n pose_params = torch.matmul(pose_feature, self.posedirs).view(-1, self.size[0], self.size[1])\n v_posed_cloth = torch.gather(pose_params,1,correspondance[:,:,None].repeat(1,1,3)) + unposed_v\n J_transformed, A = batch_global_rigid_transformation(Rs, Jsmpl, self.parents, rotate_base = False)\n\n W = self.weight.expand(num_batch,*self.weight.shape[1:])\n T = torch.matmul(W, A.view(num_batch, 24, 16)).view(num_batch, -1, 4, 4)\n\n v_posed_homo_cloth = torch.cat([v_posed_cloth, torch.ones(num_batch, v_posed_cloth.shape[1], 1, device=device)], dim = 2)\n v_homo_cloth = torch.matmul(torch.gather(T,1,correspondance[:,:,None,None].repeat(1,1,4,4)), torch.unsqueeze(v_posed_homo_cloth, -1))\n verts_cloth = v_homo_cloth[:, :, :3, 0]\n return verts_cloth\n\n\n def skeleton(self,beta,require_body=False):\n num_batch = beta.shape[0]\n v_shaped = torch.matmul(beta, self.shapedirs).view(-1, self.size[0], self.size[1]) + self.v_template\n Jx = torch.matmul(v_shaped[:, :, 0], self.J_regressor)\n Jy = torch.matmul(v_shaped[:, :, 1], self.J_regressor)\n Jz = torch.matmul(v_shaped[:, :, 2], self.J_regressor)\n J = torch.stack([Jx, Jy, Jz], dim = 2)\n if require_body:\n return J, v_shaped\n else:\n return J\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/SMPLicit.py",
"content": "import torch\nimport numpy as np\nimport torch.nn as nn\nimport os\nimport os.path as osp\nimport trimesh\nimport math\nimport copy\nfrom .SMPL import SMPL\nfrom .SMPLicit_options import Options\nfrom .smplicit_core_test import Model\n\nclass SMPLicit(nn.Module):\n def __init__(self, root_path, cloth_types):\n super(SMPLicit, self).__init__()\n self._opt = Options()\n\n uppercloth = Model(osp.join(root_path, self._opt.path_checkpoints, 'upperclothes.pth'),\n self._opt.upperbody_n_z_cut, \n self._opt.upperbody_n_z_style, self._opt.upperbody_num_clusters, \n osp.join(root_path, self._opt.path_cluster_files, self._opt.upperbody_clusters), \n self._opt.upperbody_b_min, self._opt.upperbody_b_max,\n self._opt.upperbody_resolution, thresh=self._opt.upperbody_thresh_occupancy)\n\n coat = Model(osp.join(root_path, self._opt.path_checkpoints, 'upperclothes.pth'),\n self._opt.upperbody_n_z_cut, \n self._opt.upperbody_n_z_style, self._opt.upperbody_num_clusters, \n osp.join(root_path, self._opt.path_cluster_files, self._opt.upperbody_clusters), \n self._opt.upperbody_b_min, self._opt.upperbody_b_max,\n self._opt.upperbody_resolution, thresh=self._opt.coat_thresh_occupancy)\n\n\n pants = Model(osp.join(root_path, self._opt.path_checkpoints, 'pants.pth'),\n self._opt.pants_n_z_cut, \n self._opt.pants_n_z_style, self._opt.pants_num_clusters, \n osp.join(root_path, self._opt.path_cluster_files, self._opt.pants_clusters), \n self._opt.pants_b_min, self._opt.pants_b_max,\n self._opt.pants_resolution, thresh=self._opt.pants_thresh_occupancy)\n\n skirts = Model(osp.join(root_path, self._opt.path_checkpoints, 'skirts.pth'),\n self._opt.skirts_n_z_cut, \n self._opt.skirts_n_z_style, self._opt.skirts_num_clusters, \n osp.join(root_path, self._opt.path_cluster_files, self._opt.skirts_clusters), \n self._opt.skirts_b_min, self._opt.skirts_b_max,\n self._opt.skirts_resolution, thresh=self._opt.skirts_thresh_occupancy)\n\n hair = Model(osp.join(root_path, self._opt.path_checkpoints, 'hair.pth'),\n self._opt.hair_n_z_cut, \n self._opt.hair_n_z_style, self._opt.hair_num_clusters, \n osp.join(root_path, self._opt.path_cluster_files, self._opt.hair_clusters), \n self._opt.hair_b_min, self._opt.hair_b_max,\n self._opt.hair_resolution, thresh=self._opt.hair_thresh_occupancy)\n\n shoes = Model(osp.join(root_path, self._opt.path_checkpoints, 'shoes.pth'),\n self._opt.shoes_n_z_cut, \n self._opt.shoes_n_z_style, self._opt.shoes_num_clusters, \n osp.join(root_path, self._opt.path_cluster_files, self._opt.shoes_clusters), \n self._opt.shoes_b_min, self._opt.shoes_b_max,\n self._opt.shoes_resolution, thresh=self._opt.shoes_thresh_occupancy)\n\n self.models = []\n for cloth_type in cloth_types:\n if cloth_type == 'uppercloth':\n self.models.append(uppercloth)\n elif cloth_type == 'coat':\n self.models.append(coat)\n elif cloth_type == 'pants':\n self.models.append(pants)\n elif cloth_type == 'skirts':\n self.models.append(skirts)\n elif cloth_type == 'hair':\n self.models.append(hair)\n elif cloth_type == 'shoes':\n self.models.append(shoes)\n else:\n assert 0, 'Not supported cloth type: ' + cloth_type\n self.cloth_types = cloth_types\n\n self.SMPL_Layers = [SMPL(osp.join(root_path, self._opt.path_SMPL, 'SMPL_NEUTRAL.pkl'), obj_saveable=True).cuda(),\\\n SMPL(osp.join(root_path, self._opt.path_SMPL, 'SMPL_MALE.pkl'), obj_saveable=True).cuda(),\\\n SMPL(osp.join(root_path, self._opt.path_SMPL, 'SMPL_FEMALE.pkl'), obj_saveable=True).cuda()]\n self.SMPL_Layer = None\n \n self.smpl_faces = self.SMPL_Layers[0].faces\n\n Astar_pose = torch.zeros(1, 72).cuda()\n Astar_pose[0, 5] = 0.04\n Astar_pose[0, 8] = -0.04\n self.register_buffer('Astar_pose', Astar_pose)\n\n # HYPERPARAMETER: Maximum number of points used when reposing.\n # This takes a lot of memory when finding the closest point in the SMPL so doing it by steps\n self.step = 1000 \n\n def get_right_shoe(self, sdf, unposed_cloth_mesh, do_marching_cube):\n \n # when not doing marching cube, mesh only contains vertices without faces\n if not do_marching_cube:\n sdf = torch.cat((sdf, sdf),1) # copy sdf\n rshoe = torch.stack((-unposed_cloth_mesh[:,:,0], unposed_cloth_mesh[:,:,1], unposed_cloth_mesh[:,:,2]),2)\n unposed_cloth_mesh = torch.cat((unposed_cloth_mesh, rshoe),1)\n return sdf, unposed_cloth_mesh\n # when doing marching cube, mesh contains both vertices and faces\n else:\n rshoe = np.stack((-unposed_cloth_mesh.vertices[:,0], unposed_cloth_mesh.vertices[:,1], unposed_cloth_mesh.vertices[:,2]),1)\n vertices = np.concatenate((unposed_cloth_mesh.vertices, rshoe))\n faces = np.concatenate((unposed_cloth_mesh.faces, unposed_cloth_mesh.faces[:,::-1] + len(rshoe)))\n unposed_cloth_mesh = trimesh.Trimesh(vertices, faces)\n return None, unposed_cloth_mesh\n\n def pose_mesh(self, unposed_cloth_mesh, pose, unposed_smpl_joint, unposed_smpl_mesh, do_marching_cube, smooth=True):\n if not do_marching_cube:\n iters = math.ceil(unposed_cloth_mesh.shape[1] / self.step)\n posed_cloth_mesh = []\n for i in range(iters):\n in_verts = unposed_cloth_mesh[:,i*self.step:(i+1)*self.step,:]\n out_verts = self.SMPL_Layer.deform_clothed_smpl(pose, unposed_smpl_joint, unposed_smpl_mesh, in_verts)\n posed_cloth_mesh.append(out_verts)\n posed_cloth_mesh = torch.cat(posed_cloth_mesh,1)\n return posed_cloth_mesh\n\n else:\n iters = math.ceil(len(unposed_cloth_mesh.vertices) / self.step)\n for i in range(iters):\n in_verts = torch.FloatTensor(unposed_cloth_mesh.vertices[None,i*self.step:(i+1)*self.step,:]).cuda()\n out_verts = self.SMPL_Layer.deform_clothed_smpl(pose, unposed_smpl_joint, unposed_smpl_mesh, in_verts)\n unposed_cloth_mesh.vertices[i*self.step:(i+1)*self.step] = out_verts.cpu().data.numpy() # replace unposed cloth mesh with posed one\n posed_cloth_mesh = unposed_cloth_mesh\n if smooth:\n posed_cloth_mesh = trimesh.smoothing.filter_laplacian(posed_cloth_mesh, lamb=0.5)\n return posed_cloth_mesh\n\n def pose_mesh_lower_body(self, unposed_cloth_mesh, pose, shape, Astar_pose, unposed_smpl_joint, unposed_smpl_mesh, do_marching_cube, smooth=True):\n if not do_marching_cube:\n iters = math.ceil(unposed_cloth_mesh.shape[1] / self.step)\n posed_cloth_mesh = []\n for i in range(iters):\n in_verts = unposed_cloth_mesh[:,i*self.step:(i+1)*self.step]\n out_verts = self.SMPL_Layer.unpose_and_deform_cloth(in_verts, Astar_pose, pose, shape, unposed_smpl_joint, unposed_smpl_mesh)\n posed_cloth_mesh.append(out_verts)\n posed_cloth_mesh = torch.cat(posed_cloth_mesh,1)\n return posed_cloth_mesh\n\n else:\n iters = math.ceil(len(unposed_cloth_mesh.vertices) / self.step)\n for i in range(iters):\n in_verts = torch.FloatTensor(unposed_cloth_mesh.vertices[None,i*self.step:(i+1)*self.step]).cuda()\n out_verts = self.SMPL_Layer.unpose_and_deform_cloth(in_verts, Astar_pose, pose, shape, unposed_smpl_joint, unposed_smpl_mesh)\n unposed_cloth_mesh.vertices[i*self.step:(i+1)*self.step] = out_verts.cpu().data.numpy() # replace unposed cloth mesh with posed one\n posed_cloth_mesh = unposed_cloth_mesh\n if smooth:\n posed_cloth_mesh = trimesh.smoothing.filter_laplacian(posed_cloth_mesh, lamb=0.5)\n return posed_cloth_mesh\n\n def forward(self, z_cuts, z_styles, pose, shape, gender=[0], do_marching_cube=False, valid=None, do_smooth=True):\n batch_size = pose.shape[0]\n \n unposed_smpl_joint, unposed_smpl_mesh = [], []\n Astar_smpl_mesh, Astar_smpl_joint = [], []\n for i in range(batch_size):\n SMPL_Layer = self.SMPL_Layers[gender[i]]\n unposed_smpl_joint_i, unposed_smpl_mesh_i = SMPL_Layer.skeleton(shape[None,i], require_body=True)\n Astar_smpl_mesh_i, Astar_smpl_joint_i, _ = SMPL_Layer.forward(beta=shape[None,i], theta=self.Astar_pose.repeat(1,1), get_skin=True)\n unposed_smpl_joint.append(unposed_smpl_joint_i); unposed_smpl_mesh.append(unposed_smpl_mesh_i)\n Astar_smpl_mesh.append(Astar_smpl_mesh_i); Astar_smpl_joint.append(Astar_smpl_joint_i)\n\n unposed_smpl_joint = torch.cat(unposed_smpl_joint); unposed_smpl_mesh = torch.cat(unposed_smpl_mesh)\n Astar_smpl_mesh = torch.cat(Astar_smpl_mesh); Astar_smpl_joint = torch.cat(Astar_smpl_joint)\n self.SMPL_Layer = self.SMPL_Layers[gender[0]]\n\n out_sdfs = []\n out_meshes = []\n out_meshes_unposed = []\n for i in range(len(self.models)):\n if ~valid[i]:\n out_sdfs.append([None])\n out_meshes.append([None])\n out_meshes_unposed.append([None])\n continue\n \n if self.cloth_types[i] in ['uppercloth', 'coat']:\n cloth_type = 'upperbody'\n else:\n cloth_type = self.cloth_types[i]\n resolution = eval(f'self._opt.{cloth_type}_resolution')\n\n if self.cloth_types[i] =='coat':\n is_coat = True\n else:\n is_coat = False\n \n if not do_marching_cube:\n resolution = 21\n\n if self.cloth_types[i] == 'pants' or self.cloth_types[i] == 'skirts':\n # forward network\n sdf, unposed_cloth_mesh = self.models[i].decode(z_cuts[i], z_styles[i], Astar_smpl_joint, Astar_smpl_mesh, resolution, do_marching_cube, do_smooth)\n\n # when not doing marching cube, all unposed_cloth_mesh have the same number of vertices\n if not do_marching_cube:\n posed_cloth_mesh = self.pose_mesh_lower_body(unposed_cloth_mesh, pose, shape, self.Astar_pose.repeat(batch_size,1), unposed_smpl_joint, unposed_smpl_mesh, do_marching_cube)\n # when doing marching cube, unposed_cloth_mesh can have different number of vertices\n else:\n posed_cloth_mesh = []\n for j in range(len(unposed_cloth_mesh)):\n if unposed_cloth_mesh[j] is None:\n posed_cloth_mesh.append(None)\n continue\n posed_cloth_mesh.append(self.pose_mesh_lower_body(unposed_cloth_mesh[j], pose[j,None], shape[j,None], self.Astar_pose, unposed_smpl_joint[j,None], unposed_smpl_mesh[j,None], do_marching_cube, do_smooth))\n else:\n # forward network \n sdf, unposed_cloth_mesh = self.models[i].decode(z_cuts[i], z_styles[i], unposed_smpl_joint, unposed_smpl_mesh, resolution, do_marching_cube, do_smooth, is_coat=is_coat)\n\n # when not doing marching cube, all unposed_cloth_mesh have the same number of vertices\n if not do_marching_cube:\n if self.cloth_types[i] == 'shoes': # duplicate left shoe\n sdf, unposed_cloth_mesh = self.get_right_shoe(sdf, unposed_cloth_mesh, do_marching_cube)\n posed_cloth_mesh = self.pose_mesh(unposed_cloth_mesh, pose, unposed_smpl_joint, unposed_smpl_mesh, do_marching_cube)\n # when doing marching cube, unposed_cloth_mesh can have different number of vertices\n else:\n posed_cloth_mesh = []\n for j in range(len(unposed_cloth_mesh)):\n if unposed_cloth_mesh[j] is None:\n posed_cloth_mesh.append(None)\n continue\n\n if self.cloth_types[i] == 'shoes': # duplicate left shoe\n _, unposed_cloth_mesh[j] = self.get_right_shoe(None, unposed_cloth_mesh[j], do_marching_cube)\n posed_cloth_mesh.append(self.pose_mesh(unposed_cloth_mesh[j], pose[j,None], unposed_smpl_joint[j,None], unposed_smpl_mesh[j,None], do_marching_cube, do_smooth))\n\n out_sdfs.append(sdf)\n out_meshes.append(posed_cloth_mesh)\n out_meshes_unposed.append(unposed_cloth_mesh)\n \n return out_sdfs, out_meshes, out_meshes_unposed\n\n\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/SMPLicit_options.py",
"content": "import torch\nimport os\nimport numpy as np\n\n# HUMAN PARSING LABELS:\n# 1 -> Hat\n# 2 -> Hair\n# 3 -> Glove\n# 4 -> Sunglasses,\n# 5 -> Upper-Clothes,\n# 6 -> Dress,\n# 7 -> Coat,\n# 8 -> Socks,\n# 9 -> Pants,\n# 10 -> Torso-Skin\n# 11 -> Scarf\n# 12 -> Skirt\n# 13 -> Face\n# 14 -> Left Arm\n# 15 -> Right Arm\n# 16 -> Left Leg\n# 17 -> Right Leg\n# 18 -> Left Shoe\n# 19 -> Right Shoe\n\nclass Options():\n def __init__(self):\n # Upper body options:\n self.upperbody_loadepoch = 11\n self.upperbody_clusters = 'indexs_clusters_tshirt_smpl.npy'\n self.upperbody_num_clusters = 500\n self.upperbody_n_z_cut = 6\n self.upperbody_n_z_style = 12\n self.upperbody_resolution = 128\n self.upperbody_thresh_occupancy = -0.03\n self.coat_thresh_occupancy = -0.08\n\n # Pants options:\n self.pants_loadepoch = 60\n self.pants_clusters = 'clusters_lowerbody.npy'\n self.pants_num_clusters = 500\n self.pants_n_z_cut = 6\n self.pants_n_z_style = 12\n self.pants_resolution = 128\n self.pants_thresh_occupancy = -0.02\n\n # Skirts options:\n self.skirts_loadepoch = 40\n self.skirts_clusters = 'clusters_lowerbody.npy'\n self.skirts_num_clusters = 500\n self.skirts_n_z_cut = 6\n self.skirts_n_z_style = 12\n self.skirts_resolution = 128\n self.skirts_thresh_occupancy = -0.05\n\n # Hair options:\n self.hair_loadepoch = 20000\n self.hair_clusters = 'clusters_hairs.npy'\n self.hair_num_clusters = 500\n self.hair_n_z_cut = 6\n self.hair_n_z_style = 12\n self.hair_resolution = 128\n self.hair_thresh_occupancy = -2.0\n\n # Shoes options\n self.shoes_loadepoch = 20000\n self.shoes_clusters = 'clusters_shoes.npy'\n self.shoes_n_z_cut = 0\n self.shoes_n_z_style = 4\n self.shoes_resolution = 64\n self.shoes_thresh_occupancy = -0.36\n self.shoes_num_clusters = 100\n\n # General options:\n self.path_checkpoints = '../../../../data/base_data/smplicit/checkpoints/'\n self.path_cluster_files = '../../../../data/base_data/smplicit/clusters/'\n self.path_SMPL = '../../../../data/base_data/human_models/smpl'\n\n self.upperbody_b_min = [-0.8, -0.4, -0.3]\n self.upperbody_b_max = [0.8, 0.6, 0.3]\n self.pants_b_min = [-0.3, -1.2, -0.3]\n self.pants_b_max = [0.3, 0.0, 0.3]\n self.skirts_b_min = [-0.3, -1.2, -0.3]\n self.skirts_b_max = [0.3, 0.0, 0.3]\n self.hair_b_min = [-0.35, -0.42, -0.33]\n self.hair_b_max = [0.35, 0.68, 0.37]\n self.shoes_b_min = [-0.1, -1.4, -0.2]\n self.shoes_b_max = [0.25, -0.6, 0.3]\n\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/__init__.py",
"content": "name = 'SMPLicit'\nfrom .SMPLicit import SMPLicit\nfrom .SMPL import SMPL\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/network.py",
"content": "import torch.nn as nn\nimport numpy as np\nimport torchvision\nimport torch\nimport torch.nn.functional as F\n\nclass Network(nn.Module):\n def __init__(self, n_z_style=1, point_pos_size=3, output_dim=1, n_z_cut=12):\n super(Network, self).__init__()\n self.point_pos_size = point_pos_size\n\n self.fc0_cloth = nn.utils.weight_norm(nn.Linear(n_z_style, 128, bias=True))\n self.fc1_cloth = nn.utils.weight_norm(nn.Linear(128, 128, bias=True))\n\n self.fc0_query = nn.utils.weight_norm(nn.Conv1d(point_pos_size, 128, kernel_size=1, bias=True))\n self.fc1_query = nn.utils.weight_norm(nn.Conv1d(128, 256, kernel_size=1, bias=True))\n\n self.fc0 = nn.utils.weight_norm(nn.Conv1d(128+256 + n_z_cut, 312, kernel_size=1, bias=True))\n self.fc1 = nn.utils.weight_norm(nn.Conv1d(312, 312, kernel_size=1, bias=True))\n self.fc2 = nn.utils.weight_norm(nn.Conv1d(312, 256, kernel_size=1, bias=True))\n self.fc3 = nn.utils.weight_norm(nn.Conv1d(256, 128, kernel_size=1, bias=True))\n self.fc4 = nn.utils.weight_norm(nn.Conv1d(128, output_dim, kernel_size=1, bias=True))\n\n self.activation = F.relu\n\n def forward(self, z_cut, z_style, query):\n batch_size = len(z_style)\n query_num = query.shape[1]\n\n x_cloth = self.activation(self.fc0_cloth(z_style))\n x_cloth = self.activation(self.fc1_cloth(x_cloth))\n x_cloth = x_cloth.unsqueeze(-1).repeat(1, 1, query_num)\n \n query = query.reshape(batch_size, query_num, self.point_pos_size).permute(0,2,1)\n x_query = self.activation(self.fc0_query(query))\n x_query = self.activation(self.fc1_query(x_query))\n \n z_cut = z_cut.unsqueeze(-1).repeat(1, 1, query_num)\n _in = torch.cat((x_cloth, x_query, z_cut), 1)\n\n x = self.fc0(_in)\n x = self.activation(x)\n x = self.fc1(x)\n x = self.activation(x)\n x = self.fc2(x)\n x = self.activation(x)\n x = self.fc3(x)\n x = self.activation(x)\n x = self.fc4(x)\n\n if x.shape[1] == 1:\n return x[:, 0]\n else:\n return x\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/smplicit_core_test.py",
"content": "import torch\nimport numpy as np\nfrom .utils.sdf import create_grid, eval_grid, eval_grid_octree\nfrom skimage import measure\nfrom .network import Network\nimport trimesh\n\nclass Model():\n def __init__(self, filename, n_z_cut, n_z_style, num_clusters, name_clusters, b_min, b_max, resolution, thresh=-0.05):\n self.filename = filename\n self.n_z_cut = n_z_cut\n self.n_z_style = n_z_style\n self.num_clusters = num_clusters\n self.clusters = np.load(name_clusters, allow_pickle=True)\n self.resolution = 128\n self.thresh = thresh\n self.load_networks()\n\n def load_networks(self):\n self._G = Network(n_z_style=self.n_z_style, point_pos_size=self.num_clusters*3, output_dim=1, n_z_cut=self.n_z_cut).cuda()\n self._G.load_state_dict(torch.load(self.filename))\n self._G.eval()\n\n def get_bbox(self, joint, mesh):\n joints_name = ('Pelvis', 'L_Hip', 'R_Hip', 'Torso', 'L_Knee', 'R_Knee', 'Spine', 'L_Ankle', 'R_Ankle', 'Chest', 'L_Toe', 'R_Toe', 'Neck', 'L_Thorax', 'R_Thorax', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Hand', 'R_Hand')\n\n if 'upper' in self.filename:\n rhand = joint[:,joints_name.index('R_Hand'),:]\n lhand = joint[:,joints_name.index('L_Hand'),:]\n xmin = rhand[:,0]; xmax = lhand[:,0];\n ycenter = joint[:,joints_name.index('Chest'),1]\n height = (ycenter - joint[:,joints_name.index('Pelvis'),1])*2*2\n ymin = ycenter - height/2; ymax = ycenter + height/2;\n zcenter = (torch.min(mesh[:,:,2],1)[0] + torch.max(mesh[:,:,2],1)[0]) / 2.\n depth = (torch.max(mesh[:,:,2],1)[0] - torch.min(mesh[:,:,2],1)[0]) * 1.5\n zmin = zcenter - depth/2.; zmax = zcenter + depth/2.\n b_min = torch.stack((xmin, ymin, zmin),1)\n b_max = torch.stack((xmax, ymax, zmax),1)\n \n elif 'pants' in self.filename:\n rankle = joint[:,joints_name.index('R_Ankle'),:]\n lankle = joint[:,joints_name.index('L_Ankle'),:]\n pelvis = joint[:,joints_name.index('Pelvis'),:]\n spine1 = joint[:,joints_name.index('Torso'),:]\n xcenter = pelvis[:,0]; width = (xcenter - rankle[:,0])*2*2.3\n xmin = xcenter - width/2; xmax = xcenter + width/2;\n ycenter = (pelvis[:,1] + rankle[:,1])/2.; height = (pelvis[:,1] - ycenter)*2*1.2\n ymin = ycenter - height/2; ymax = ycenter + height/2;\n zcenter = (torch.min(mesh[:,:,2],1)[0] + torch.max(mesh[:,:,2],1)[0]) / 2.\n depth = (torch.max(mesh[:,:,2],1)[0] - torch.min(mesh[:,:,2],1)[0]) * 1.5\n zmin = zcenter - depth/2; zmax = zcenter + depth/2\n b_min = torch.stack((xmin, ymin, zmin),1)\n b_max = torch.stack((xmax, ymax, zmax),1)\n \n elif 'skirt' in self.filename:\n rankle = joint[:,joints_name.index('R_Ankle'),:]\n lankle = joint[:,joints_name.index('L_Ankle'),:]\n pelvis = joint[:,joints_name.index('Pelvis'),:]\n spine1 = joint[:,joints_name.index('Torso'),:]\n xcenter = pelvis[:,0]; width = (xcenter - rankle[:,0])*2*3\n xmin = xcenter - width/2; xmax = xcenter + width/2;\n ycenter = (pelvis[:,1] + rankle[:,1])/2.; height = (pelvis[:,1] - ycenter)*2*1.2\n ymin = ycenter - height/2; ymax = ycenter + height/2;\n zcenter = (torch.min(mesh[:,:,2],1)[0] + torch.max(mesh[:,:,2],1)[0]) / 2.\n depth = (torch.max(mesh[:,:,2],1)[0] - torch.min(mesh[:,:,2],1)[0]) * 2\n zmin = zcenter - depth/2; zmax = zcenter + depth/2\n b_min = torch.stack((xmin, ymin, zmin),1)\n b_max = torch.stack((xmax, ymax, zmax),1)\n \n elif 'hair' in self.filename:\n lshoulder = joint[:,joints_name.index('L_Shoulder'),:]\n rshoulder = joint[:,joints_name.index('R_Shoulder'),:]\n xcenter = (lshoulder[:,0] + rshoulder[:,0])/2.\n width = (xcenter - rshoulder[:,0])*2\n xmin = xcenter - width/2; xmax = xcenter + width/2;\n head = joint[:,joints_name.index('Head'),:]\n ymax = torch.max(mesh[:,:,1],1)[0]; ymin = joint[:,joints_name.index('Spine'),1]\n ycenter = (ymin+ymax)/2.; height = (ycenter - ymin)*2*1.2\n ymin = ycenter - height/2; ymax = ycenter + height/2;\n zcenter = (torch.min(mesh[:,:,2],1)[0] + torch.max(mesh[:,:,2],1)[0]) / 2.\n depth = (torch.max(mesh[:,:,2],1)[0] - torch.min(mesh[:,:,2],1)[0]) * 1.2\n zmin = zcenter - depth*0.8; zmax = zcenter + depth/2\n b_min = torch.stack((xmin, ymin, zmin),1)\n b_max = torch.stack((xmax, ymax, zmax),1)\n\n elif 'shoes' in self.filename:\n lknee = joint[:,joints_name.index('L_Knee'),:]\n lankle = joint[:,joints_name.index('L_Ankle'),:]\n lfoot = joint[:,joints_name.index('L_Toe'),:]\n xmin = lankle[:,0] - 0.15; xmax = lankle[:,0] + 0.15;\n ycenter = lankle[:,1]\n height = (lknee[:,1] - ycenter)*1.5\n ymin = ycenter - height/2; ymax = ycenter + height/2;\n zcenter = (lankle[:,2] + lfoot[:,2])/2.\n zmin = zcenter - 0.25; zmax = zcenter + 0.25;\n b_min = torch.stack((xmin, ymin, zmin),1)\n b_max = torch.stack((xmax, ymax, zmax),1)\n \n return b_min, b_max\n\n\n def decode(self, z_cut, z_style, smpl_joint, smpl_mesh, resolution, do_marching_cube, smooth=True, is_coat=False):\n batch_size = z_cut.shape[0]\n \n # prepare query points to predict SDF\n b_min, b_max = self.get_bbox(smpl_joint, smpl_mesh)\n query_points = create_grid((resolution, resolution, resolution), b_min, b_max)\n \n # sample points in clusters from smpl mesh\n smpl_points = smpl_mesh[:,self.clusters[self.num_clusters]] # batch_size, smpl_point_num, 3\n smpl_point_num = smpl_points.shape[1]\n \n def eval_func(query_points, ref_points, z_cut, z_style, scale):\n dist = query_points[:,:,None,:] - ref_points[:,None,:,:]\n dist = dist.view(-1, query_points.shape[1], ref_points.shape[1]*3)\n pred = self._G(z_cut, z_style, dist)*scale\n return pred\n \n if not do_marching_cube:\n if not smooth:\n # remove empty 3D space\n if 'upper' in self.filename:\n query_points = query_points.view(batch_size, resolution, resolution, resolution, 3)\n is_empty = torch.zeros((resolution, resolution, resolution)).float().cuda()\n is_empty[:resolution//4,:resolution//2,:] = 1 # right\n is_empty[resolution//4*3:,:resolution//2,:] = 1 # left\n is_empty[resolution//4:resolution//4*3:,resolution//4:resolution//2,:resolution//3] = 1 # back center\n query_points = query_points[is_empty[None,:,:,:,None].repeat(batch_size,1,1,1,3)==0].view(batch_size,-1,3)\n query_point_num = query_points.shape[1]\n\n # predict SDF\n sdf = eval_grid(query_points, smpl_points, z_cut, z_style, eval_func, resolution, 1, num_samples=10000)\n cloth_points = query_points \n return sdf, cloth_points\n \n else:\n cloth_meshes = []\n sdfs = eval_grid(query_points, smpl_points, z_cut, z_style, eval_func, resolution, -100, num_samples=10000)\n sdfs = sdfs.view(batch_size,resolution,resolution,resolution)\n\n for i in range(batch_size): \n sdf = sdfs[i].cpu().numpy()\n\n if 'pant' in self.filename:\n # pant exception handling (heuristic)\n sdf[resolution*63//128:resolution*66//128,:resolution*47//64,:] = self.thresh - 0.001\n sdf[resolution*62//128:resolution*67//128,:resolution*45//64,:] = self.thresh - 0.001\n try:\n verts, faces, normals, values = measure.marching_cubes(sdf, self.thresh, method='lewiner')\n\n cloth_mesh = trimesh.Trimesh(np.float64(verts), faces[:, ::-1])\n cloth_mesh.vertices /= resolution\n cloth_mesh.vertices *= (b_max[i,None].cpu().numpy() - b_min[i,None].cpu().numpy())\n cloth_mesh.vertices += b_min[i,None].cpu().numpy()\n\n if smooth:\n smooth_mesh = trimesh.smoothing.filter_laplacian(cloth_mesh, lamb=0.5)\n if not np.isnan(smooth_mesh.vertices).any():\n cloth_mesh = smooth_mesh\n \n except ValueError:\n cloth_mesh = None\n \n cloth_meshes.append(cloth_mesh)\n\n return None, cloth_meshes\n\n\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/util_smpl.py",
"content": "# import h5py\nimport torch\nimport numpy as np\nimport json\nfrom torch.autograd import Variable\nimport torch.nn.functional as F\nimport cv2\nimport math\nimport os\n# def load_mean_theta():\n# mean = np.zeros(85, dtype = np.float)\n\n# mean_values = h5py.File(os.path.join(os.path.dirname(__file__),'model/neutral_smpl_mean_params.h5'),'r')\n# mean_pose = mean_values['pose']\n# mean_pose[:3] = 0\n# mean_shape = mean_values['shape']\n# mean_pose[0]=np.pi\n\n# #init sacle is 0.9\n# mean[0] = 0.9\n\n# mean[3:75] = mean_pose[:]\n# mean[75:] = mean_shape[:]\n\n# return mean\n\ndef batch_rodrigues(theta):\n #theta N x 3\n batch_size = theta.shape[0]\n l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)\n angle = torch.unsqueeze(l1norm, -1)\n normalized = torch.div(theta, angle)\n angle = angle * 0.5\n v_cos = torch.cos(angle)\n v_sin = torch.sin(angle)\n quat = torch.cat([v_cos, v_sin * normalized], dim = 1)\n \n return quat2mat(quat)\n\ndef quat2mat(quat):\n \"\"\"Convert quaternion coefficients to rotation matrix.\n Args:\n quat: size = [B, 4] 4 <===>(w, x, y, z)\n Returns:\n Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]\n \"\"\"\n norm_quat = quat\n norm_quat = norm_quat/norm_quat.norm(p=2, dim=1, keepdim=True)\n w, x, y, z = norm_quat[:,0], norm_quat[:,1], norm_quat[:,2], norm_quat[:,3]\n\n B = quat.size(0)\n\n w2, x2, y2, z2 = w.pow(2), x.pow(2), y.pow(2), z.pow(2)\n wx, wy, wz = w*x, w*y, w*z\n xy, xz, yz = x*y, x*z, y*z\n\n rotMat = torch.stack([w2 + x2 - y2 - z2, 2*xy - 2*wz, 2*wy + 2*xz,\n 2*wz + 2*xy, w2 - x2 + y2 - z2, 2*yz - 2*wx,\n 2*xz - 2*wy, 2*wx + 2*yz, w2 - x2 - y2 + z2], dim=1).view(B, 3, 3)\n return rotMat\n\ndef batch_global_rigid_transformation(Rs, Js, parent, rotate_base = False):\n N = Rs.shape[0]\n if rotate_base:\n np_rot_x = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]], dtype = np.float)\n np_rot_x = np.reshape(np.tile(np_rot_x, [N, 1]), [N, 3, 3])\n rot_x = Variable(torch.from_numpy(np_rot_x).float()).to(Rs.device)\n root_rotation = torch.matmul(Rs[:, 0, :, :], rot_x)\n else:\n root_rotation = Rs[:, 0, :, :]\n Js = torch.unsqueeze(Js, -1)\n\n def make_A(R, t):\n R_homo = F.pad(R, [0, 0, 0, 1, 0, 0])\n t_homo = torch.cat([t, Variable(torch.ones(N, 1, 1)).to(R.device)], dim = 1)\n return torch.cat([R_homo, t_homo], 2)\n \n A0 = make_A(root_rotation, Js[:, 0])\n results = [A0]\n\n for i in range(1, parent.shape[0]):\n j_here = Js[:, i] - Js[:, parent[i]]\n A_here = make_A(Rs[:, i], j_here)\n res_here = torch.matmul(results[parent[i]], A_here)\n results.append(res_here)\n\n results = torch.stack(results, dim = 1)\n\n new_J = results[:, :, :3, 3]\n Js_w0 = torch.cat([Js, Variable(torch.zeros(N, 24, 1, 1)).to(Rs.device)], dim = 2)\n init_bone = torch.matmul(results, Js_w0)\n init_bone = F.pad(init_bone, [3, 0, 0, 0, 0, 0, 0, 0])\n A = results - init_bone\n\n return new_J, A\n\n\ndef batch_lrotmin(theta):\n theta = theta[:,3:].contiguous()\n Rs = batch_rodrigues(theta.view(-1, 3))\n print(Rs.shape)\n e = Variable(torch.eye(3).float())\n Rs = Rs.sub(1.0, e)\n\n return Rs.view(-1, 23 * 9)\n\ndef batch_orth_proj(X, camera):\n '''\n X is N x num_points x 3\n '''\n camera = camera.view(-1, 1, 3)\n X_trans = X[:, :, :2] + camera[:, :, 1:]\n shape = X_trans.shape\n return (camera[:, :, 0] * X_trans.view(shape[0], -1)).view(shape)\n\ndef reflect_pose(poses):\n swap_inds = np.array([\n 0, 1, 2, 6, 7, 8, 3, 4, 5, 9, 10, 11, 15, 16, 17, 12, 13, 14, 18,\n 19, 20, 24, 25, 26, 21, 22, 23, 27, 28, 29, 33, 34, 35, 30, 31, 32,\n 36, 37, 38, 42, 43, 44, 39, 40, 41, 45, 46, 47, 51, 52, 53, 48, 49,\n 50, 57, 58, 59, 54, 55, 56, 63, 64, 65, 60, 61, 62, 69, 70, 71, 66,\n 67, 68\n ])\n\n sign_flip = np.array([\n 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1,\n -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1,\n -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1,\n 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1, -1, 1, -1,\n -1, 1, -1, -1\n ])\n\n return poses[swap_inds] * sign_flip\n"
},
{
"path": "common/utils/SMPLicit/SMPLicit/utils/__init__.py",
"content": ""
},
{
"path": "common/utils/SMPLicit/SMPLicit/utils/sdf.py",
"content": "import numpy as np\nimport torch\nimport math\n\ndef create_grid(resolution, b_min, b_max):\n batch_size = b_min.shape[0]\n\n # make grids\n res_x, res_y, res_z = resolution\n #zz,yy,xx = torch.meshgrid(torch.arange(res_z),torch.arange(res_y),torch.arange(res_x))\n xx,yy,zz = torch.meshgrid(torch.arange(res_x),torch.arange(res_y),torch.arange(res_z))\n coords = torch.stack((xx, yy, zz))\n coords = coords.reshape(3, -1).float()\n coords = coords[None,:,:].repeat(batch_size,1,1).float().cuda()\n \n\n # affine transform\n coords_matrix = torch.eye(4).view(1,4,4).repeat(batch_size,1,1).float().cuda()\n length = b_max - b_min\n coords_matrix[:, 0, 0] = length[:,0] / res_x\n coords_matrix[:, 1, 1] = length[:,1] / res_y\n coords_matrix[:, 2, 2] = length[:,2] / res_z\n coords_matrix[:, 0:3, 3] = b_min\n coords = torch.bmm(coords_matrix[:, :3, :3], coords) + coords_matrix[:, :3, 3:4]\n\n # return grids\n coords = coords.view(batch_size, 3, -1).transpose(2,1).contiguous() # res_x*res_y*res_z, 3\n return coords\n\ndef batch_eval(query_points, ref_points, z_cut, z_style, eval_func, scale, num_samples):\n num_pts = query_points.shape[1]\n num_batches = math.ceil(num_pts / num_samples)\n \n sdf = []\n for i in range(num_batches):\n sdf.append(eval_func(query_points[:,i * num_samples:i * num_samples + num_samples,:], ref_points, z_cut, z_style, scale))\n sdf = torch.cat(sdf,1)\n return sdf\n\ndef eval_grid(query_points, ref_points, z_cut, z_style, eval_func, resolution, scale, num_samples=512 * 512 * 512):\n sdf = batch_eval(query_points, ref_points, z_cut, z_style, eval_func, scale, num_samples=num_samples)\n return sdf\n\ndef eval_grid_octree(query_points, ref_points, z_cut, z_style, eval_func, resolution, init_resolution=64, threshold=0.01, num_samples=512 * 512 * 512):\n res_x, res_y, res_z = resolution\n sdf = np.zeros(resolution)\n dirty = np.ones(resolution, dtype=np.bool)\n grid_mask = np.zeros(resolution, dtype=np.bool)\n\n step_size = res_x // init_resolution\n while step_size > 0:\n # subdivide the grid\n grid_mask[0:res_x:step_size, 0:res_y:step_size, 0:res_z:step_size] = True\n # test samples in this iteration\n test_mask = np.logical_and(grid_mask, dirty)\n points = query_points[torch.from_numpy(test_mask).cuda().reshape(-1)==1,:]\n \n sdf[test_mask] = batch_eval(points[None,:,:], ref_points[None,:,:], z_cut[None,:], z_style[None,:], eval_func, num_samples=num_samples).detach().cpu().numpy().reshape(-1)\n dirty[test_mask] = False\n\n # do interpolation\n if step_size <= 1:\n break\n for x in range(0, res_x - step_size, step_size):\n for y in range(0, res_y - step_size, step_size):\n for z in range(0, res_z - step_size, step_size):\n # if center marked, return\n if not dirty[x + step_size // 2, y + step_size // 2, z + step_size // 2]:\n continue\n v0 = sdf[x, y, z]\n v1 = sdf[x, y, z + step_size]\n v2 = sdf[x, y + step_size, z]\n v3 = sdf[x, y + step_size, z + step_size]\n v4 = sdf[x + step_size, y, z]\n v5 = sdf[x + step_size, y, z + step_size]\n v6 = sdf[x + step_size, y + step_size, z]\n v7 = sdf[x + step_size, y + step_size, z + step_size]\n v = np.array([v0, v1, v2, v3, v4, v5, v6, v7])\n v_min = v.min()\n v_max = v.max()\n # this cell is all the same\n if (v_max - v_min) < threshold:\n sdf[x:x + step_size, y:y + step_size, z:z + step_size] = (v_max + v_min) / 2\n dirty[x:x + step_size, y:y + step_size, z:z + step_size] = False\n step_size //= 2\n\n return sdf.reshape(resolution)\n\n"
},
{
"path": "common/utils/dir.py",
"content": "import os\nimport sys\n\ndef make_folder(folder_name):\n if not os.path.exists(folder_name):\n os.makedirs(folder_name)\n\ndef add_pypath(path):\n if path not in sys.path:\n sys.path.insert(0, path)\n\n"
},
{
"path": "common/utils/human_models.py",
"content": "import numpy as np\nimport torch\nimport os.path as osp\nfrom config import cfg\nfrom utils.transforms import transform_joint_to_other_db\nimport smplx\n\nclass SMPL(object):\n def __init__(self):\n self.layer_arg = {'create_body_pose': False, 'create_betas': False, 'create_global_orient': False, 'create_transl': False}\n self.layer = {'neutral': smplx.create(cfg.human_model_path, 'smpl', gender='NEUTRAL', **self.layer_arg), 'male': smplx.create(cfg.human_model_path, 'smpl', gender='MALE', **self.layer_arg), 'female': smplx.create(cfg.human_model_path, 'smpl', gender='FEMALE', **self.layer_arg)}\n self.vertex_num = 6890\n self.face = self.layer['neutral'].faces\n self.shape_param_dim = 10\n\n # SMPL joint set\n self.joint_num = 24\n self.joints_name = ('Pelvis', 'L_Hip', 'R_Hip', 'Torso', 'L_Knee', 'R_Knee', 'Spine', 'L_Ankle', 'R_Ankle', 'Chest', 'L_Foot', 'R_Foot', 'Neck', 'L_Collar', 'R_Collar', 'Head', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Hand', 'R_Hand')\n self.flip_pairs = ( (1,2), (4,5), (7,8), (10,11), (13,14), (16,17), (18,19), (20,21), (22,23) )\n self.root_joint_idx = self.joints_name.index('Pelvis')\n self.joint_regressor = self.layer['neutral'].J_regressor.numpy().astype(np.float32)\n\n # Astar pose\n self.Astar_pose = torch.zeros(1, self.joint_num*3)\n self.Astar_pose[0, 5] = 0.04 \n self.Astar_pose[0, 8] = -0.04\n\n def get_custom_template_layer(self, v_template, gender):\n layer_arg = {'create_body_pose': False, 'create_betas': False, 'create_global_orient': False, 'create_transl': False, 'v_template': v_template}\n layer = smplx.create(cfg.human_model_path, 'smpl', gender=gender.upper(), **layer_arg)\n return layer\n\nsmpl = SMPL()\n"
},
{
"path": "common/utils/postprocessing.py",
"content": "import os\nimport os.path as osp\nimport numpy as np\nimport torch\nimport cv2\nimport json\nimport copy\nfrom pytorch3d.structures import Meshes\nfrom pytorch3d.renderer import RasterizationSettings, MeshRasterizer, TexturesVertex\nfrom pytorch3d.renderer.cameras import PerspectiveCameras\nfrom pytorch3d.renderer.lighting import AmbientLights, PointLights\nfrom pytorch3d.renderer.mesh.shader import BlendParams, HardPhongShader\nfrom pytorch3d.renderer.materials import Materials\nfrom pytorch3d.renderer.mesh.renderer import MeshRenderer\nfrom config import cfg\n\n\ndef get_face_map(pix_to_face, faces):\n face_map = torch.zeros((pix_to_face.shape[0], pix_to_face.shape[1], 3)) - 1\n for i in range(pix_to_face.shape[0]):\n for j in range(pix_to_face.shape[1]):\n if pix_to_face[i][j] != -1:\n face_map[i][j] = faces[pix_to_face[i][j]]\n return face_map\n\n\nclass Renderer:\n def __init__(self, device='cuda', focal=cfg.focal, princpt=cfg.princpt, img_shape=cfg.input_img_shape):\n self.device = device\n\n self.set_renderer(focal, princpt, img_shape)\n\n def set_renderer(self, focal, princpt, img_shape, anti_aliasing=False):\n focal, princpt = torch.FloatTensor(focal)[None,:], torch.FloatTensor(princpt)[None,:]\n self.img_shape = img_shape\n self.anti_aliasing = anti_aliasing\n\n if self.anti_aliasing:\n img_shape = (img_shape[0]*2, img_shape[1]*2)\n princpt *= 2; focal *= 2\n\n img_size = max(img_shape[0], img_shape[1])\n raster_settings = RasterizationSettings(image_size=(img_size,img_size), blur_radius=0.0, faces_per_pixel=1, bin_size=0)\n\n cameras = PerspectiveCameras(focal_length=focal, \\\n principal_point=princpt, \\\n in_ndc=False, \\\n R=torch.eye(3)[None,:,:], \\\n T=torch.zeros(3)[None,:], \\\n image_size=((img_size,img_size),),\\\n device=torch.device(self.device))\n\n lights = PointLights(device=self.device, location=[[0.0, 0.0, -10.0]])\n materials = Materials(ambient_color=((0.92, 0.92, 0.92), ), diffuse_color=((1, 1, 1), ), specular_color=((1, 1, 1), ), shininess=4, device=self.device)\n blend_params = BlendParams(sigma=1e-1, gamma=1e-4)\n shader = HardPhongShader(device=self.device, blend_params=blend_params, cameras=cameras, lights=lights, materials=materials)\n\n self.rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings).to(self.device)\n self.renderer = MeshRenderer(rasterizer=self.rasterizer, shader=shader)\n\n def rasterize_mesh(self, mesh_vert, mesh_face):\n output = self.rasterizer(Meshes(verts=[mesh_vert.to(self.device)], faces=[mesh_face.to(self.device)]))\n\n face_map = get_face_map(output.pix_to_face.squeeze(), mesh_face)\n return face_map[:, :cfg.input_img_shape[1]]\n\n def render(self, img, mesh_vert, mesh_face):\n mesh_vert, mesh_face = torch.tensor(mesh_vert), torch.tensor(mesh_face)\n\n verts_rgb = torch.ones_like(mesh_vert)[None]\n textures = TexturesVertex(verts_features=verts_rgb.to(self.device))\n\n output = self.renderer(Meshes(verts=[mesh_vert.to(self.device)], faces=[mesh_face.to(self.device)], textures=textures))\n output = (output[0]*255).cpu().numpy()\n \n if self.anti_aliasing:\n img = cv2.resize(img, (self.img_shape[1]*2, self.img_shape[0]*2))\n img_shape = (self.img_shape[0]*2, self.img_shape[1]*2)\n else:\n img_shape = self.img_shape\n\n if img_shape[0] > img_shape[1]:\n output = output[:, :img_shape[1]]\n else:\n output = output[:img_shape[0], :]\n \n valid = output[:,:,3] > 0\n img[valid] = output[:,:,:3][valid]\n\n if self.anti_aliasing:\n img = cv2.resize(img, (self.img_shape[1], self.img_shape[0]))\n\n return img\n\ndef rasterize_mesh_given_cam_param(mesh_vert, mesh_face, focal, princpt):\n device = 'cuda'\n raster_settings = RasterizationSettings(image_size=(cfg.input_img_shape[0],cfg.input_img_shape[0]), blur_radius=0.0, faces_per_pixel=1)\n \n cameras = PerspectiveCameras(focal_length=torch.FloatTensor([focal[0],focal[1]])[None,:], \\\n principal_point=torch.FloatTensor([princpt[0],princpt[1]])[None,:], \\\n in_ndc=False, \\\n R=torch.eye(3)[None,:,:], \\\n T=torch.zeros(3)[None,:], \\\n image_size=((cfg.input_img_shape[0],cfg.input_img_shape[0]),),\\\n device=torch.device(device))\n rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings).to(device)\n output = rasterizer(Meshes(verts=[mesh_vert.to(device)], faces=[mesh_face.to(device)]))\n\n face_map = get_face_map(output.pix_to_face.squeeze(), mesh_face) \n return face_map[:, :cfg.input_img_shape[1]]\n\ndef save_proj_faces(face_map, save_path):\n face_map = face_map.reshape(-1, 3)\n\n file = open(save_path, 'w')\n for idx, v in enumerate(face_map):\n file.write('%d %d %d\\n' % (v[0], v[1], v[2]))\n file.close()\n\ndef merge_mesh(verts, faces):\n vert_len = [0]\n for vert in verts:\n vert_len.append(len(vert))\n\n vert_len = np.cumsum(vert_len)\n for i, face in enumerate(faces):\n face += vert_len[i]\n \n return np.concatenate(verts), np.concatenate(faces)\n\ndef read_valid_point(verts, indexs, valid):\n valid_verts = []\n for i, val in enumerate(valid):\n if val != 0:\n idx1, idx2, idx3 = indexs[i]\n v = (verts[idx1] + verts[idx2] + verts[idx3]) / 3\n valid_verts.append(v)\n valid_verts = np.stack(valid_verts)\n return valid_verts\n \ndef pa_mpjpe(predicted, target):\n \"\"\"\n Pose error: MPJPE after rigid alignment (scale, rotation, and translation),\n often referred to as \"Protocol #2\" in many papers.\n \"\"\"\n assert predicted.shape == target.shape\n \n muX = np.mean(target, axis=1, keepdims=True)\n muY = np.mean(predicted, axis=1, keepdims=True)\n \n X0 = target - muX\n Y0 = predicted - muY\n\n normX = np.sqrt(np.sum(X0**2, axis=(1, 2), keepdims=True))\n normY = np.sqrt(np.sum(Y0**2, axis=(1, 2), keepdims=True))\n \n X0 /= normX\n Y0 /= normY\n\n H = np.matmul(X0.transpose(0, 2, 1), Y0)\n U, s, Vt = np.linalg.svd(H)\n V = Vt.transpose(0, 2, 1)\n R = np.matmul(V, U.transpose(0, 2, 1))\n\n # Avoid improper rotations (reflections), i.e. rotations with det(R) = -1\n sign_detR = np.sign(np.expand_dims(np.linalg.det(R), axis=1))\n V[:, :, -1] *= sign_detR\n s[:, -1] *= sign_detR.flatten()\n R = np.matmul(V, U.transpose(0, 2, 1)) # Rotation\n\n tr = np.expand_dims(np.sum(s, axis=1, keepdims=True), axis=2)\n\n a = tr * normX / normY # Scale\n t = muX - a*np.matmul(muY, R) # Translation\n \n return a, R, t\n\ndef pairwise_distances(a, b, p=2, inv=False, num_samples=500):\n if not inv:\n tmp = a; a = b; b= tmp\n\n a = torch.tensor(a[None, :, :]).cuda()\n b = torch.tensor(b[None, :, :]).cuda()\n num_batches = a.shape[1] // num_samples\n\n dists = []\n for i in range(num_batches):\n dist = torch.norm((a[:,i*num_samples : (i+1)*num_samples, None, :] - b[:, None, :, :]),p=2,dim=3)\n dist, _ = torch.min(dist, 2)\n dist = dist.reshape(-1)\n dists.append(dist)\n\n if a.shape[1] % num_samples > 0:\n dist = torch.norm((a[:,-1 * (a.shape[1] % num_samples):, None, :] - b[:, None, :, :]),p=2,dim=3)\n dist, _ = torch.min(dist, 2)\n dist = dist.reshape(-1)\n dists.append(dist)\n\n dist= torch.cat(dists).mean().cpu()\n return dist\n\nrenderer = Renderer()"
},
{
"path": "common/utils/preprocessing.py",
"content": "import numpy as np\nimport cv2\nimport random\nfrom config import cfg\nimport math\nfrom utils.human_models import smpl\nfrom utils.transforms import cam2pixel, transform_joint_to_other_db\nimport torch\n\ndef load_img(path, order='RGB'):\n img = cv2.imread(path, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)\n if not isinstance(img, np.ndarray):\n raise IOError(\"Fail to read %s\" % path)\n\n if order=='RGB':\n img = img[:,:,::-1].copy()\n \n img = img.astype(np.float32)\n return img\n\ndef get_bbox(joint_img, joint_valid, extend_ratio=1.2):\n\n x_img, y_img = joint_img[:,0], joint_img[:,1]\n x_img = x_img[joint_valid==1]; y_img = y_img[joint_valid==1];\n xmin = min(x_img); ymin = min(y_img); xmax = max(x_img); ymax = max(y_img);\n\n x_center = (xmin+xmax)/2.; width = xmax-xmin;\n xmin = x_center - 0.5 * width * extend_ratio\n xmax = x_center + 0.5 * width * extend_ratio\n \n y_center = (ymin+ymax)/2.; height = ymax-ymin;\n ymin = y_center - 0.5 * height * extend_ratio\n ymax = y_center + 0.5 * height * extend_ratio\n\n bbox = np.array([xmin, ymin, xmax - xmin, ymax - ymin]).astype(np.float32)\n return bbox\n\ndef process_bbox(bbox, img_width, img_height):\n # sanitize bboxes\n x, y, w, h = bbox\n x1 = np.max((0, x))\n y1 = np.max((0, y))\n x2 = np.min((img_width - 1, x1 + np.max((0, w - 1))))\n y2 = np.min((img_height - 1, y1 + np.max((0, h - 1))))\n if w*h > 0 and x2 > x1 and y2 > y1:\n bbox = np.array([x1, y1, x2-x1, y2-y1])\n else:\n return None\n\n # aspect ratio preserving bbox\n w = bbox[2]\n h = bbox[3]\n c_x = bbox[0] + w/2.\n c_y = bbox[1] + h/2.\n aspect_ratio = cfg.input_img_shape[1]/cfg.input_img_shape[0]\n if w > aspect_ratio * h:\n h = w / aspect_ratio\n elif w < aspect_ratio * h:\n w = h * aspect_ratio\n bbox[2] = w*1.25\n bbox[3] = h*1.25\n bbox[0] = c_x - bbox[2]/2.\n bbox[1] = c_y - bbox[3]/2.\n \n bbox = bbox.astype(np.float32)\n return bbox\n\ndef convert_focal_princpt(focal, princpt, img2bb_trans):\n focal = np.array([[focal[0], 0], [0, focal[1]], [0, 0]])\n princpt = np.array([[princpt[0], 0], [0, princpt[1]], [1, 1]])\n\n focal = np.dot(img2bb_trans, focal)\n princpt = np.dot(img2bb_trans, princpt)\n\n cam_param = np.array([focal[0][0], focal[1][1], princpt[0][0], princpt[1][1]])\n return cam_param\n\ndef get_aug_config():\n scale_factor = 0.25\n rot_factor = 30\n color_factor = 0.2\n \n scale = np.clip(np.random.randn(), -1.0, 1.0) * scale_factor + 1.0\n rot = np.clip(np.random.randn(), -2.0,\n 2.0) * rot_factor if random.random() <= 0.6 else 0\n c_up = 1.0 + color_factor\n c_low = 1.0 - color_factor\n color_scale = np.array([random.uniform(c_low, c_up), random.uniform(c_low, c_up), random.uniform(c_low, c_up)])\n do_flip = False\n return scale, rot, color_scale, do_flip\n\ndef augmentation(img, bbox, data_split):\n if data_split == 'train':\n scale, rot, color_scale, do_flip = get_aug_config()\n else:\n scale, rot, color_scale, do_flip = 1.0, 0.0, np.array([1,1,1]), False\n \n img, valid_mask, trans, inv_trans = generate_patch_image(img, bbox, scale, rot, do_flip, cfg.input_img_shape)\n img = np.clip(img * color_scale[None,None,:], 0, 255)\n return img, valid_mask, trans, inv_trans, rot, do_flip\n\ndef generate_patch_image(cvimg, bbox, scale, rot, do_flip, out_shape):\n img = cvimg.copy()\n img_height, img_width, img_channels = img.shape\n \n bb_c_x = float(bbox[0] + 0.5*bbox[2])\n bb_c_y = float(bbox[1] + 0.5*bbox[3])\n bb_width = float(bbox[2])\n bb_height = float(bbox[3])\n\n if do_flip:\n img = img[:, ::-1, :]\n bb_c_x = img_width - bb_c_x - 1\n\n trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot)\n img_patch = cv2.warpAffine(img, trans, (int(out_shape[1]), int(out_shape[0])), flags=cv2.INTER_LINEAR, borderValue=(-1,-1,-1))\n valid_mask = (img_patch > -1)\n if len(valid_mask.shape) == 3:\n valid_mask = valid_mask[:,:,0]\n img_patch[img_patch == -1] = 0\n img_patch = img_patch.astype(np.float32)\n inv_trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, out_shape[1], out_shape[0], scale, rot, inv=True)\n\n return img_patch, valid_mask, trans, inv_trans\n\ndef rotate_2d(pt_2d, rot_rad):\n x = pt_2d[0]\n y = pt_2d[1]\n sn, cs = np.sin(rot_rad), np.cos(rot_rad)\n xx = x * cs - y * sn\n yy = x * sn + y * cs\n return np.array([xx, yy], dtype=np.float32)\n\ndef gen_trans_from_patch_cv(c_x, c_y, src_width, src_height, dst_width, dst_height, scale, rot, inv=False):\n # augment size with scale\n src_w = src_width * scale\n src_h = src_height * scale\n src_center = np.array([c_x, c_y], dtype=np.float32)\n\n # augment rotation\n rot_rad = np.pi * rot / 180\n src_downdir = rotate_2d(np.array([0, src_h * 0.5], dtype=np.float32), rot_rad)\n src_rightdir = rotate_2d(np.array([src_w * 0.5, 0], dtype=np.float32), rot_rad)\n\n dst_w = dst_width\n dst_h = dst_height\n dst_center = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32)\n dst_downdir = np.array([0, dst_h * 0.5], dtype=np.float32)\n dst_rightdir = np.array([dst_w * 0.5, 0], dtype=np.float32)\n\n src = np.zeros((3, 2), dtype=np.float32)\n src[0, :] = src_center\n src[1, :] = src_center + src_downdir\n src[2, :] = src_center + src_rightdir\n\n dst = np.zeros((3, 2), dtype=np.float32)\n dst[0, :] = dst_center\n dst[1, :] = dst_center + dst_downdir\n dst[2, :] = dst_center + dst_rightdir\n \n if inv:\n trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))\n else:\n trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))\n\n trans = trans.astype(np.float32)\n return trans\n\ndef process_db_coord(joint_img, joint_valid, do_flip, img_shape, flip_pairs, img2bb_trans, rot, src_joints_name, target_joints_name):\n joint_img, joint_valid = joint_img.copy(), joint_valid.copy()\n\n # flip augmentation\n if do_flip:\n joint_img[:,0] = img_shape[1] - 1 - joint_img[:,0]\n for pair in flip_pairs:\n joint_img[pair[0],:], joint_img[pair[1],:] = joint_img[pair[1],:].copy(), joint_img[pair[0],:].copy()\n joint_valid[pair[0],:], joint_valid[pair[1],:] = joint_valid[pair[1],:].copy(), joint_valid[pair[0],:].copy()\n \n # affine transformation and root-relative depth\n joint_img_xy1 = np.concatenate((joint_img, np.ones_like(joint_img[:,:1])),1)\n joint_img = np.dot(img2bb_trans, joint_img_xy1.transpose(1,0)).transpose(1,0)\n joint_img[:,0] = joint_img[:,0] / cfg.input_img_shape[1] * cfg.output_joint_shape[1]\n joint_img[:,1] = joint_img[:,1] / cfg.input_img_shape[0] * cfg.output_joint_shape[0]\n \n # check truncation\n joint_trunc = joint_valid * ((joint_img[:,0] >= 0) * (joint_img[:,0] < cfg.output_joint_shape[1]) * \\\n (joint_img[:,1] >= 0) * (joint_img[:,1] < cfg.output_joint_shape[0])).reshape(-1,1).astype(np.float32)\n\n\n # transform joints to target db joints\n joint_img = transform_joint_to_other_db(joint_img, src_joints_name, target_joints_name)\n joint_valid = transform_joint_to_other_db(joint_valid, src_joints_name, target_joints_name)\n joint_trunc = transform_joint_to_other_db(joint_trunc, src_joints_name, target_joints_name)\n return joint_img, joint_valid, joint_trunc\n\ndef process_human_model_output(human_model_param, cam_param, do_flip, img_shape, img2bb_trans, rot):\n pose, shape = human_model_param['pose'], human_model_param['shape']\n\n if 'trans' in human_model_param:\n trans = human_model_param['trans']\n else:\n trans = [0,0,0]\n if 'gender' in human_model_param:\n gender = human_model_param['gender']\n else:\n gender = 'neutral'\n pose = torch.FloatTensor(pose).view(-1,3); shape = torch.FloatTensor(shape).view(1,-1); # smpl parameters (pose: 72 dimension, shape: 10 dimension)\n trans = torch.FloatTensor(trans).view(1,-1) # translation vector\n \n # apply camera extrinsic (rotation)\n # merge root pose and camera rotation \n if 'R' in cam_param:\n R = np.array(cam_param['R'], dtype=np.float32).reshape(3,3)\n root_pose = pose[smpl.root_joint_idx,:].numpy()\n root_pose, _ = cv2.Rodrigues(root_pose)\n root_pose, _ = cv2.Rodrigues(np.dot(R,root_pose))\n pose[smpl.root_joint_idx] = torch.from_numpy(root_pose).view(3)\n\n # get mesh and joint coordinates\n root_pose = pose[smpl.root_joint_idx].view(1,3)\n body_pose = torch.cat((pose[:smpl.root_joint_idx,:], pose[smpl.root_joint_idx+1:,:])).view(1,-1)\n output = smpl.layer[gender](betas=shape, body_pose=body_pose, global_orient=root_pose, transl=trans)\n mesh_coord = output.vertices[0].numpy()\n joint_coord = np.dot(smpl.joint_regressor, mesh_coord)\n\n # apply camera exrinsic (translation)\n # compenstate rotation (translation from origin to root joint was not cancled)\n if 'R' in cam_param and 't' in cam_param:\n R, t = np.array(cam_param['R'], dtype=np.float32).reshape(3,3), np.array(cam_param['t'], dtype=np.float32).reshape(1,3)\n root_coord = joint_coord[smpl.root_joint_idx,None,:]\n joint_coord = joint_coord - root_coord + np.dot(R, root_coord.transpose(1,0)).transpose(1,0) + t\n mesh_coord = mesh_coord - root_coord + np.dot(R, root_coord.transpose(1,0)).transpose(1,0) + t\n\n ## so far, data augmentations are not applied yet\n ## now, project the 3D coordinates to image space and apply data augmentations\n\n # 3D data rotation augmentation\n rot_aug_mat = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0], \n [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],\n [0, 0, 1]], dtype=np.float32)\n # flip pose parameter (axis-angle)\n if do_flip:\n for pair in smpl.flip_pairs:\n pose[pair[0], :], pose[pair[1], :] = pose[pair[1], :].clone(), pose[pair[0], :].clone()\n pose[:,1:3] *= -1 # multiply -1 to y and z axis of axis-angle\n # rotate root pose\n pose = pose.numpy()\n root_pose = pose[smpl.root_joint_idx,:]\n root_pose, _ = cv2.Rodrigues(root_pose)\n root_pose, _ = cv2.Rodrigues(np.dot(rot_aug_mat,root_pose))\n pose[smpl.root_joint_idx] = root_pose.reshape(3)\n \n # return results\n pose = pose.reshape(-1)\n # change to mean shape if beta is too far from it\n shape[(shape.abs() > 3).any(dim=1)] = 0.\n shape = shape.numpy().reshape(-1)\n return pose, shape, mesh_coord # data augmentation is not performed on mesh_coord \n\ndef bilinear_interpolate(im, x, y):\n x = np.asarray(x)\n y = np.asarray(y)\n\n x0 = np.floor(x).astype(int)\n x1 = x0 + 1\n y0 = np.floor(y).astype(int)\n y1 = y0 + 1\n\n x0 = np.clip(x0, 0, im.shape[2]-1);\n x1 = np.clip(x1, 0, im.shape[2]-1);\n y0 = np.clip(y0, 0, im.shape[1]-1);\n y1 = np.clip(y1, 0, im.shape[1]-1);\n\n Ia = im[:, y0, x0 ]\n Ib = im[:, y1, x0 ]\n Ic = im[:, y0, x1 ]\n Id = im[:, y1, x1 ]\n\n wa = (x1-x) * (y1-y)\n wb = (x1-x) * (y-y0)\n wc = (x-x0) * (y1-y)\n wd = (x-x0) * (y-y0)\n\n return wa*Ia + wb*Ib + wc*Ic + wd*Id\n\ndef iou_sil(sil_out, sil_target):\n intersect = sil_out * sil_target\n union = (sil_out + sil_target) > 0\n if np.sum(union) == 0:\n return None\n else:\n return np.sum(intersect) / np.sum(union)\n"
},
{
"path": "common/utils/transforms.py",
"content": "import torch\nimport numpy as np\nfrom config import cfg\nfrom torch.nn import functional as F\nimport torchgeometry as tgm\n\ndef cam2pixel(cam_coord, f, c):\n x = cam_coord[:,0] / cam_coord[:,2] * f[0] + c[0]\n y = cam_coord[:,1] / cam_coord[:,2] * f[1] + c[1]\n z = cam_coord[:,2]\n return np.stack((x,y,z),1)\n\ndef pixel2cam(pixel_coord, f, c):\n x = (pixel_coord[:,0] - c[0]) / f[0] * pixel_coord[:,2]\n y = (pixel_coord[:,1] - c[1]) / f[1] * pixel_coord[:,2]\n z = pixel_coord[:,2]\n return np.stack((x,y,z),1)\n\ndef world2cam(world_coord, R, t):\n cam_coord = np.dot(R, world_coord.transpose(1,0)).transpose(1,0) + t.reshape(1,3)\n return cam_coord\n\ndef transform_joint_to_other_db(src_joint, src_name, dst_name):\n src_joint_num = len(src_name)\n dst_joint_num = len(dst_name)\n\n new_joint = np.zeros(((dst_joint_num,) + src_joint.shape[1:]), dtype=np.float32)\n for src_idx in range(len(src_name)):\n name = src_name[src_idx]\n if name in dst_name:\n dst_idx = dst_name.index(name)\n new_joint[dst_idx] = src_joint[src_idx]\n\n return new_joint\n\ndef rot6d_to_axis_angle(x):\n batch_size = x.shape[0]\n\n x = x.view(-1,3,2)\n a1 = x[:, :, 0]\n a2 = x[:, :, 1]\n b1 = F.normalize(a1)\n b2 = F.normalize(a2 - torch.einsum('bi,bi->b', b1, a2).unsqueeze(-1) * b1)\n b3 = torch.cross(b1, b2)\n rot_mat = torch.stack((b1, b2, b3), dim=-1) # 3x3 rotation matrix\n \n rot_mat = torch.cat([rot_mat,torch.zeros((batch_size,3,1)).cuda().float()],2) # 3x4 rotation matrix\n axis_angle = tgm.rotation_matrix_to_angle_axis(rot_mat).reshape(-1,3) # axis-angle\n axis_angle[torch.isnan(axis_angle)] = 0.0\n return axis_angle\n\ndef unwrap_xy_to_uv(feat_xy, dp_fg, dp_I, dp_u, dp_v):\n batch_size, feat_dim, height, width = feat_xy.shape\n\n dp_fg = torch.max(dp_fg, 1)[1] # argmax\n dp_I = torch.max(dp_I, 1)[1] + 1 # argmax. add 1 to make the bkg class\n dp_I[dp_fg == 0] = 0 # bkg\n\n _dp_u, _dp_v = 0, 0\n for i in range(cfg.dp_patch_num):\n mask = (dp_I == (i+1)) # add 1 to make the bkg class\n _dp_u += dp_u[:,i,:,:] * mask\n _dp_v += dp_v[:,i,:,:] * mask\n dp_u, dp_v = _dp_u, _dp_v\n\n \n scatter_src = feat_xy.permute(1,0,2,3).reshape(feat_dim,-1)\n \n batch_idx = torch.arange(batch_size)[:,None,None].repeat(1,height,width).view(-1).to(feat_xy.device) #.cuda()\n _dp_I = dp_I.view(-1)\n _dp_u = (dp_u.view(-1) * (cfg.output_uv_shape[0]-1)).long()\n _dp_v = ((1 - dp_v.view(-1)) * (cfg.output_uv_shape[1]-1)).long() # inverse v coordinate following DensePose R-CNN\n scatter_idx = batch_idx * (cfg.dp_patch_num + 1) * cfg.output_uv_shape[0] * cfg.output_uv_shape[1] + \\\n _dp_I * cfg.output_uv_shape[0] * cfg.output_uv_shape[1] + \\\n _dp_u * cfg.output_uv_shape[1] + \\\n _dp_v\n\n is_valid = (_dp_u >= 0) * (_dp_u < cfg.output_uv_shape[0]) * (_dp_v >= 0) * (_dp_v < cfg.output_uv_shape[1])\n scatter_src = scatter_src[:,is_valid]\n scatter_idx = scatter_idx[is_valid]\n \n feat_uv = scatter_mean(scatter_src, scatter_idx, 1, dim_size = batch_size * (cfg.dp_patch_num + 1) * cfg.output_uv_shape[0] * cfg.output_uv_shape[1]).view(feat_dim, batch_size, cfg.dp_patch_num + 1, cfg.output_uv_shape[0], cfg.output_uv_shape[1]).permute(1,2,0,3,4)[:,1:,:,:,:] # remove bkg class (cfg.dp_patch_num + 1 -> cfg.dp_patch_num)\n\n return feat_uv\n\n"
},
{
"path": "common/utils/vis.py",
"content": "import os\nimport os.path as osp\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nfrom utils.human_models import smpl\nfrom utils.postprocessing import merge_mesh, renderer\nfrom config import cfg\n\ndef save_result(output, path):\n verts, faces, colors = [], [], []\n verts.append(output['smpl_mesh'].astype(np.float32))\n faces.append(smpl.face.astype(np.int32))\n colors.append(np.tile(cfg.cloth_colors['smpl_body'], (len(output['smpl_mesh']),1)))\n\n for cloth_type in cfg.cloth_types:\n if output[cloth_type + '_mesh'] is None: continue\n verts.append(output[cloth_type + '_mesh'].vertices.astype(np.float32))\n faces.append(output[cloth_type + '_mesh'].faces.astype(np.int32))\n colors.append(np.tile(cfg.cloth_colors[cloth_type], (len(output[cloth_type + '_mesh'].vertices),1)))\n \n verts, faces = merge_mesh(verts, faces)\n colors = np.concatenate(colors)\n verts[:,:2] *= -1\n save_obj_with_color(verts, faces, colors, path)\n\n return verts, faces\n \n\ndef render_result(verts, faces, img, path):\n rendered_img = renderer.render(img, verts, faces)\n cv2.imwrite(path.replace('output.obj', 'render_img.jpg'), rendered_img)\n\n\ndef vis_keypoints_with_skeleton(img, kps, kps_lines, kp_thresh=0.4, alpha=1):\n # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.\n cmap = plt.get_cmap('rainbow')\n colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]\n colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]\n\n # Perform the drawing on a copy of the image, to allow for blending.\n kp_mask = np.copy(img)\n\n # Draw the keypoints.\n for l in range(len(kps_lines)):\n i1 = kps_lines[l][0]\n i2 = kps_lines[l][1]\n p1 = kps[0, i1].astype(np.int32), kps[1, i1].astype(np.int32)\n p2 = kps[0, i2].astype(np.int32), kps[1, i2].astype(np.int32)\n if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:\n cv2.line(\n kp_mask, p1, p2,\n color=colors[l], thickness=2, lineType=cv2.LINE_AA)\n if kps[2, i1] > kp_thresh:\n cv2.circle(\n kp_mask, p1,\n radius=3, color=colors[l], thickness=-1, lineType=cv2.LINE_AA)\n if kps[2, i2] > kp_thresh:\n cv2.circle(\n kp_mask, p2,\n radius=3, color=colors[l], thickness=-1, lineType=cv2.LINE_AA)\n\n # Blend the keypoints.\n return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)\n\ndef vis_keypoints(img, kps, alpha=1):\n # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.\n cmap = plt.get_cmap('rainbow')\n colors = [cmap(i) for i in np.linspace(0, 1, len(kps) + 2)]\n colors = [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors]\n\n # Perform the drawing on a copy of the image, to allow for blending.\n kp_mask = np.copy(img)\n\n # Draw the keypoints.\n for i in range(len(kps)):\n p = kps[i][0].astype(np.int32), kps[i][1].astype(np.int32)\n cv2.circle(kp_mask, p, radius=3, color=colors[i], thickness=-1, lineType=cv2.LINE_AA)\n\n # Blend the keypoints.\n return cv2.addWeighted(img, 1.0 - alpha, kp_mask, alpha, 0)\n\ndef vis_3d_skeleton(kpt_3d, kpt_3d_vis, kps_lines, filename=None):\n\n fig = plt.figure()\n ax = fig.add_subplot(111, projection='3d')\n\n # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.\n cmap = plt.get_cmap('rainbow')\n colors = [cmap(i) for i in np.linspace(0, 1, len(kps_lines) + 2)]\n colors = [np.array((c[2], c[1], c[0])) for c in colors]\n\n for l in range(len(kps_lines)):\n i1 = kps_lines[l][0]\n i2 = kps_lines[l][1]\n x = np.array([kpt_3d[i1,0], kpt_3d[i2,0]])\n y = np.array([kpt_3d[i1,1], kpt_3d[i2,1]])\n z = np.array([kpt_3d[i1,2], kpt_3d[i2,2]])\n\n if kpt_3d_vis[i1,0] > 0 and kpt_3d_vis[i2,0] > 0:\n ax.plot(x, z, -y, c=colors[l], linewidth=2)\n if kpt_3d_vis[i1,0] > 0:\n ax.scatter(kpt_3d[i1,0], kpt_3d[i1,2], -kpt_3d[i1,1], c=colors[l], marker='o')\n if kpt_3d_vis[i2,0] > 0:\n ax.scatter(kpt_3d[i2,0], kpt_3d[i2,2], -kpt_3d[i2,1], c=colors[l], marker='o')\n\n x_r = np.array([0, cfg.input_shape[1]], dtype=np.float32)\n y_r = np.array([0, cfg.input_shape[0]], dtype=np.float32)\n z_r = np.array([0, 1], dtype=np.float32)\n \n if filename is None:\n ax.set_title('3D vis')\n else:\n ax.set_title(filename)\n\n ax.set_xlabel('X Label')\n ax.set_ylabel('Z Label')\n ax.set_zlabel('Y Label')\n ax.legend()\n\n plt.show()\n cv2.waitKey(0)\n\ndef save_obj(v, f, file_name='output.obj'):\n obj_file = open(file_name, 'w')\n for i in range(len(v)):\n obj_file.write('v ' + str(v[i][0]) + ' ' + str(v[i][1]) + ' ' + str(v[i][2]) + '\\n')\n if f is not None:\n for i in range(len(f)):\n obj_file.write('f ' + str(f[i][0]+1) + '/' + str(f[i][0]+1) + ' ' + str(f[i][1]+1) + '/' + str(f[i][1]+1) + ' ' + str(f[i][2]+1) + '/' + str(f[i][2]+1) + '\\n')\n obj_file.close()\n\ndef save_obj_with_color(v, f, c, file_name='output.obj'):\n obj_file = open(file_name, 'w')\n for i in range(len(v)):\n obj_file.write('v ' + str(v[i][0]) + ' ' + str(v[i][1]) + ' ' + str(v[i][2]) + ' ' + str(c[i][0]) + ' ' + str(c[i][1]) + ' ' + str(c[i][2]) + '\\n')\n if f is not None:\n for i in range(len(f)):\n obj_file.write('f ' + str(f[i][0]+1) + '/' + str(f[i][0]+1) + ' ' + str(f[i][1]+1) + '/' + str(f[i][1]+1) + ' ' + str(f[i][2]+1) + '/' + str(f[i][2]+1) + '\\n')\n obj_file.close()\n\ndef vis_parse(img, parse, class_num):\n # Convert from plt 0-1 RGBA colors to 0-255 BGR colors for opencv.\n cmap = plt.get_cmap('rainbow')\n colors = [cmap(i) for i in np.linspace(0, 1, class_num)]\n colors = [(0,0,0)] + [(c[2] * 255, c[1] * 255, c[0] * 255) for c in colors] # add bkg\n\n for i in range(class_num+1):\n img = img * (parse[:,:,None] != i) + np.array(colors[i]).reshape(1,1,3) * (parse[:,:,None] == i)\n return img\n\ndef vis_dp(img, dp_u, dp_v, filename):\n fig = plt.figure()\n plt.imshow(img[:,:,::-1])\n plt.contour(dp_u,10,linewidths=1)\n plt.contour(dp_v,10,linewidths=1)\n plt.axis('off')\n plt.xticks([])\n plt.yticks([])\n plt.subplots_adjust(left = 0, bottom = 0, right = 1, top = 1, hspace = 0, wspace = 0)\n plt.savefig(filename, bbox_inches='tight', pad_inches=0)\n plt.close(fig)\n return\n\n\n"
},
{
"path": "data/DeepFashion2/DeepFashion2.py",
"content": "import os\nimport os.path as osp\nimport numpy as np\nimport copy\nimport json\nimport cv2\nimport torch\nfrom pycocotools.coco import COCO\nimport pycocotools.mask as mask_util\n\nfrom utils.human_models import smpl\nfrom utils.preprocessing import load_img, process_bbox, augmentation, generate_patch_image, bilinear_interpolate\nfrom utils.vis import save_obj, vis_parse\nfrom config import cfg\n\nclass DeepFashion2(torch.utils.data.Dataset):\n def __init__(self, transform, data_split):\n self.transform = transform\n self.data_split = data_split\n\n if data_split != 'train':\n assert 0, \"Invalid train mode.\"\n\n self.img_path = osp.join('..', 'data', 'DeepFashion2', 'data')\n self.annot_path = osp.join('..', 'data', 'DeepFashion2', 'data')\n self.parse_path = osp.join('..', 'data', 'preprocessed_data', 'parse', 'DeepFashion2')\n self.preprocessed_path = osp.join('..', 'data', 'preprocessed_data')\n self.dp_path = osp.join(self.preprocessed_path, 'densepose', 'DeepFashion2')\n\n # lip parse set\n self.parse_set = {'uppercloth': (5,), 'coat': (7,), 'pants': (9,), 'skirts': (12,), 'hair': (2,), 'shoes': (18,19)}\n self.sampling_stride = 4 # subsampling for training\n\n self.datalist = self.load_data()\n print(\"Load data: \", len(self.datalist))\n \n def load_data(self):\n self.img_path = osp.join(self.img_path , 'train', 'image')\n self.dp_path = osp.join(self.dp_path, 'train')\n db = COCO(osp.join(self.annot_path, 'DeepFashion2_train.json'))\n with open(osp.join(self.preprocessed_path, 'gender', 'DeepFashion2_train_gender.json')) as f:\n genders = json.load(f)\n with open(osp.join(self.parse_path, 'train_parsing_annotation.json')) as f:\n parsing_paths = json.load(f)\n\n datalist = []\n i = 0\n for aid in db.anns.keys():\n i += 1\n if i % self.sampling_stride != 0:\n continue\n\n ann = db.anns[aid]\n img = db.loadImgs(ann['image_id'])[0]\n img_path = osp.join(self.img_path, img['file_name'])\n \n # bbox\n bbox = process_bbox(ann['bbox'], img['width'], img['height']) \n if bbox is None: continue\n \n # parse\n if parsing_paths is not None:\n if str(aid) in parsing_paths:\n parse_path = osp.join(self.parse_path, parsing_paths[str(aid)])\n else:\n continue\n else:\n parse_path = None\n\n # densepose\n if self.data_split == 'train':\n try:\n dp = np.load(osp.join(self.dp_path, str(aid) + '.npz'), allow_pickle=True)\n except:\n continue\n\n if len(dp['smpl_v_idx']) == 0: continue\n \n dp_x = np.array(dp['dp_x'], dtype=np.float32)\n dp_y = np.array(dp['dp_y'], dtype=np.float32)\n dp_xy = np.stack((dp_x, dp_y),1)\n\n dp_I = np.array(dp['dp_I'], dtype=np.int16)\n dp_u = np.array(dp['dp_U'], dtype=np.float32)\n dp_v = np.array(dp['dp_V'], dtype=np.float32)\n dp_uv = np.stack((dp_u, dp_v),1)\n \n smpl_v_idx = np.array(dp['smpl_v_idx'], dtype=np.int32)\n dp_mask = dp['dp_fg'].item()\n dp_mask = mask_util.decode(dp_mask)\n\n dp_data = {'xy': dp_xy, 'uv': dp_uv, 'I': dp_I, 'smpl_v_idx': smpl_v_idx, 'masks': dp_mask}\n else:\n dp_data = None\n\n # gender\n if str(aid) in genders:\n gender = genders[str(aid)]\n else:\n continue\n \n data_dict = {\n 'img_path': img_path, 'ann_id': aid, 'img_shape': (img['height'],img['width']), \n 'bbox': bbox, 'orig_bbox': ann['bbox'], 'gender': gender, \n 'parse_path': parse_path, 'dp': dp_data\n } \n datalist.append(data_dict)\n\n return datalist\n\n def __len__(self):\n return len(self.datalist)\n\n def __getitem__(self, idx):\n data = copy.deepcopy(self.datalist[idx])\n\n img_path, img_shape = data['img_path'], data['img_shape']\n \n # image load\n img = load_img(img_path)\n\n # affine transform\n bbox = data['bbox']\n img, valid_mask, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)\n img = self.transform(img.astype(np.float32))/255.\n\n # load parse (cloth segmentation)\n parse = cv2.imread(data['parse_path'])\n parse_list = []\n for cloth_type in ('fg',) + cfg.cloth_types:\n # get cloth indexs\n if cloth_type == 'fg':\n idxs = np.unique(parse).tolist()\n idxs.pop(idxs.index(0))\n if len(idxs) == 0:\n parse_fg = np.zeros((cfg.output_parse_shape[0], cfg.output_parse_shape[1])) > 0\n continue\n else:\n idxs = self.parse_set[cloth_type]\n \n # get masking corresponding to a cloth\n mask = [parse == i for i in idxs]\n mask = (sum(mask) > 0).astype(np.float32)\n _, _, _, lip2img_trans = generate_patch_image(mask, data['orig_bbox'], 1.0, 0.0, False, mask.shape)\n mask = cv2.warpAffine(mask, lip2img_trans, (img_shape[1], img_shape[0]), flags=cv2.INTER_LINEAR)\n mask = cv2.warpAffine(mask, img2bb_trans, (cfg.input_img_shape[1], cfg.input_img_shape[0]), flags=cv2.INTER_LINEAR)\n mask = cv2.resize(mask, (cfg.output_parse_shape[1], cfg.output_parse_shape[0]))\n \n if cloth_type == 'fg': parse_fg = mask[:,:,0] > 0\n else: parse_list.append(mask)\n\n parse = np.stack(parse_list)[:,:,:,0] # remove the last dimension (which has 3 channels)\n is_bkg = (np.prod(parse == 0, 0) == 1)\n parse = np.argmax(parse, 0) + 1 # add 1 to make the bkg class\n parse[is_bkg] = 0\n parse[valid_mask == 0] = -1\n parse_valid = valid_mask\n\n # load densepose\n dp_xy, dp_uv, dp_I, dp_vertex = data['dp']['xy'], data['dp']['uv'], data['dp']['I'], data['dp']['smpl_v_idx']\n dp_xy = np.concatenate((dp_xy, np.ones_like(dp_xy[:,:1])),1)\n dp_xy = np.dot(img2bb_trans, dp_xy.transpose(1,0)).transpose(1,0)\n dp_I = dp_I - 1 # dp_I is started wtih 1. make it zero-based index.\n\n cur_point_num = len(dp_xy)\n if cur_point_num > cfg.dp_point_num:\n idxs = np.random.choice(np.arange(cur_point_num), size=cfg.dp_point_num)\n cur_point_num = cfg.dp_point_num\n dp_xy = dp_xy[idxs]; dp_uv = dp_uv[idxs]; dp_I = dp_I[idxs]; dp_vertex = dp_vertex[idxs]\n \n # match densepose & parse\n _dp_xy = dp_xy.copy()\n _dp_xy[:,0] = _dp_xy[:,0] / cfg.input_img_shape[1] * cfg.output_parse_shape[1]\n _dp_xy[:,1] = _dp_xy[:,1] / cfg.input_img_shape[0] * cfg.output_parse_shape[0]\n\n parse_onehot = np.zeros((len(cfg.cloth_types)+1, cfg.output_parse_shape[0], cfg.output_parse_shape[1]))\n for i in range(len(cfg.cloth_types)+1):\n parse_onehot[i][parse == i] = 1.0\n\n dp_cloth_idx = np.ones((_dp_xy.shape[0]), np.int16) * -1\n dp_cloth_idx[bilinear_interpolate(parse_fg[None,:,:], _dp_xy[:,0], _dp_xy[:,1])[0] > 0.5] = 0\n for i in range(len(cfg.cloth_types)):\n dp_cloth_idx[np.argmax(bilinear_interpolate(parse_onehot, _dp_xy[:,0], _dp_xy[:,1]), 0) == (i+1)] = i+1\n\n smpl_cloth_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_cloth_idx[dp_vertex] = dp_cloth_idx\n smpl_cloth_valid = (smpl_cloth_idx != -1).astype(np.float32)\n smpl_patch_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_patch_idx[dp_vertex] = dp_I[dp_I != -1]\n\n # remove coat ambiguity\n if (smpl_cloth_idx==cfg.cloth_types.index('uppercloth')+1).sum() == 0 and (smpl_cloth_idx==cfg.cloth_types.index('coat')+1).sum() > 0:\n idxs = (smpl_cloth_idx == cfg.cloth_types.index('coat')+1)\n smpl_cloth_idx[idxs] = cfg.cloth_types.index('uppercloth')+1\n\n # gender\n if data['gender'] == 'male': gender = 1\n elif data['gender'] == 'female': gender = 2\n else: gender = 0\n\n # dummy smpl parameter\n smpl_pose = np.zeros((smpl.joint_num*3,), dtype=np.float32)\n smpl_shape = np.zeros((smpl.shape_param_dim,), dtype=np.float32)\n cam_trans = np.zeros((3,), dtype=np.float32)\n\n inputs = {'img': img}\n targets = {'gender': gender, 'parse': parse, 'smpl_cloth_idx': smpl_cloth_idx, 'smpl_patch_idx': smpl_patch_idx}\n meta_info = {'smpl_cloth_valid': smpl_cloth_valid, 'smpl_pose': smpl_pose, 'smpl_shape': smpl_shape, 'cam_trans': cam_trans}\n return inputs, targets, meta_info"
},
{
"path": "data/MSCOCO/MSCOCO.py",
"content": "import os\nimport os.path as osp\nimport numpy as np\nfrom config import cfg\nimport copy\nimport json\nimport cv2\nimport torch\nfrom pycocotools.coco import COCO\nimport pycocotools.mask as mask_util\nfrom utils.human_models import smpl\nfrom utils.preprocessing import load_img, process_bbox, augmentation, generate_patch_image, process_db_coord, process_human_model_output, bilinear_interpolate, iou_sil\nfrom utils.vis import save_obj, save_result, render_result\n\nclass MSCOCO(torch.utils.data.Dataset):\n def __init__(self, transform, data_split):\n self.transform = transform\n self.data_split = data_split\n self.img_path = osp.join('..', 'data', 'MSCOCO', 'images')\n self.annot_path = osp.join('..', 'data', 'MSCOCO', 'annotations')\n self.parse_path = osp.join('..', 'data', 'MSCOCO', 'parses')\n self.preprocessed_path = osp.join('..', 'data', 'preprocessed_data')\n self.dp_path = osp.join(self.preprocessed_path, 'densepose', 'MSCOCO')\n\n # lip parse set\n self.parse_set = {'uppercloth': (5,), 'coat': (7,), 'pants': (9,), 'skirts': (12,), 'hair': (2,), 'shoes': (18,19)}\n self.bcc_dist_threshold = 0.03\n self.eval_types = ['upper_body', 'lower_body','non_cloth']\n\n self.datalist = self.load_data()\n print(\"Load data: \", len(self.datalist))\n\n def load_data(self):\n if self.data_split == 'train':\n db = COCO(osp.join(self.annot_path, 'coco_wholebody_train_v1.0.json'))\n with open(osp.join(self.preprocessed_path, 'gender', 'MSCOCO_train_gender.json')) as f:\n genders = json.load(f)\n with open(osp.join(self.parse_path, 'LIP_trainval_parsing.json')) as f:\n parsing_paths = json.load(f)\n else:\n db = COCO(osp.join(self.annot_path, 'coco_wholebody_val_v1.0.json'))\n with open(osp.join(self.preprocessed_path, 'gender', 'MSCOCO_val_gender.json')) as f:\n genders = json.load(f)\n with open(osp.join(self.preprocessed_path, 'smpl_param', 'MSCOCO_test_Pose2Pose.json')) as f:\n smpl_params = json.load(f)\n if cfg.calculate_bcc:\n with open(osp.join(self.parse_path, 'LIP_trainval_parsing.json')) as f:\n parsing_paths = json.load(f)\n with open(osp.join(self.annot_path, 'coco_dp_val.json')) as f:\n dps = json.load(f)\n else:\n parsing_paths = None\n dps = None\n\n datalist = []\n for aid in db.anns.keys():\n ann = db.anns[aid]\n img = db.loadImgs(ann['image_id'])[0]\n\n if self.data_split == 'train':\n imgname = osp.join('train2017', img['file_name'])\n else:\n imgname = osp.join('val2017', img['file_name'])\n img_path = osp.join(self.img_path, imgname)\n \n if self.data_split == 'train':\n if ann['iscrowd'] or (ann['num_keypoints'] == 0):\n continue\n \n # bbox\n bbox = process_bbox(ann['bbox'], img['width'], img['height']) \n if bbox is None: continue\n \n # parse\n if parsing_paths is not None:\n if str(aid) in parsing_paths:\n parse_path = osp.join(self.parse_path, 'TrainVal_parsing_annotations/TrainVal_parsing_annotations/', parsing_paths[str(aid)])\n else:\n continue\n else:\n parse_path = None\n\n # filter images with few visible joints\n joint_img = np.array(ann['keypoints'], dtype=np.float32).reshape(-1,3)\n if np.sum(joint_img[:,2]>0) < 6: continue\n\n # densepose\n if self.data_split == 'train':\n try:\n dp = np.load(osp.join(self.dp_path, str(aid) + '.npz'), allow_pickle=True)\n except:\n continue\n\n if len(dp['smpl_v_idx']) == 0: continue\n\n dp_x = np.array(dp['dp_x'], dtype=np.float32)\n dp_y = np.array(dp['dp_y'], dtype=np.float32)\n dp_xy = np.stack((dp_x, dp_y),1)\n\n dp_I = np.array(dp['dp_I'], dtype=np.int16)\n dp_u = np.array(dp['dp_U'], dtype=np.float32)\n dp_v = np.array(dp['dp_V'], dtype=np.float32)\n dp_uv = np.stack((dp_u, dp_v),1)\n \n smpl_v_idx = np.array(dp['smpl_v_idx'], dtype=np.int32)\n dp_mask = dp['dp_fg'].item()\n dp_mask = mask_util.decode(dp_mask)\n\n dp_data = {'xy': dp_xy, 'uv': dp_uv, 'I': dp_I, 'smpl_v_idx': smpl_v_idx, 'masks': dp_mask}\n elif cfg.calculate_bcc:\n if str(aid) in dps:\n dp = dps[str(aid)]\n if len(dp['smpl_v_idx']) == 0: continue\n\n dp_x = np.array(dp['dp_x'], dtype=np.float32) / 256 * ann['bbox'][2] + ann['bbox'][0]\n dp_y = np.array(dp['dp_y'], dtype=np.float32) / 256 * ann['bbox'][3] + ann['bbox'][1]\n dp_xy = np.stack((dp_x, dp_y),1)\n\n dp_I = np.array(dp['dp_I'], dtype=np.int16)\n dp_u = np.array(dp['dp_U'], dtype=np.float32)\n dp_v = np.array(dp['dp_V'], dtype=np.float32)\n dp_uv = np.stack((dp_u, dp_v),1)\n \n smpl_v_idx = np.array(dp['smpl_v_idx'], dtype=np.int32)\n dp_mask = mask_util.decode(dp['dp_masks'][0])\n\n dp_data = {'xy': dp_xy, 'uv': dp_uv, 'I': dp_I, 'smpl_v_idx': smpl_v_idx, 'masks': dp_mask}\n else:\n continue\n else:\n dp_data = None\n\n # smpl params\n if self.data_split == 'test':\n if str(aid) in smpl_params:\n smpl_param = smpl_params[str(aid)]['smpl_param']\n cam_param = smpl_params[str(aid)]['cam_param']\n else:\n continue\n else:\n smpl_param, cam_param = None, None\n\n # gender\n if str(aid) in genders:\n gender = genders[str(aid)]\n else:\n continue\n\n data_dict = {\n 'img_path': img_path, 'ann_id': aid, 'img_shape': (img['height'],img['width']), \n 'bbox': bbox, 'orig_bbox': ann['bbox'], 'gender': gender, \n 'parse_path': parse_path, 'dp': dp_data,\n 'smpl_param': smpl_param, 'cam_param': cam_param\n } \n datalist.append(data_dict) \n \n \n return datalist\n\n def __len__(self):\n return len(self.datalist)\n\n def __getitem__(self, idx):\n data = copy.deepcopy(self.datalist[idx])\n\n img_path, img_shape = data['img_path'], data['img_shape']\n \n # image load\n img = load_img(img_path)\n\n # affine transform\n bbox = data['bbox']\n img, valid_mask, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)\n img = self.transform(img.astype(np.float32))/255.\n\n if self.data_split == 'train':\n # load parse (cloth segmentation)\n parse = cv2.imread(data['parse_path'])\n parse_list = []\n for cloth_type in ('fg',) + cfg.cloth_types:\n # get cloth indexs\n if cloth_type == 'fg':\n idxs = np.unique(parse).tolist()\n idxs.pop(idxs.index(0))\n if len(idxs) == 0:\n parse_fg = np.zeros((cfg.output_parse_shape[0], cfg.output_parse_shape[1])) > 0\n continue\n else:\n idxs = self.parse_set[cloth_type]\n \n # get masking corresponding to a cloth\n mask = [parse == i for i in idxs]\n mask = (sum(mask) > 0).astype(np.float32)\n _, _, _, lip2img_trans = generate_patch_image(mask, data['orig_bbox'], 1.0, 0.0, False, mask.shape)\n mask = cv2.warpAffine(mask, lip2img_trans, (img_shape[1], img_shape[0]), flags=cv2.INTER_LINEAR)\n mask = cv2.warpAffine(mask, img2bb_trans, (cfg.input_img_shape[1], cfg.input_img_shape[0]), flags=cv2.INTER_LINEAR)\n mask = cv2.resize(mask, (cfg.output_parse_shape[1], cfg.output_parse_shape[0]))\n \n if cloth_type == 'fg': parse_fg = mask[:,:,0] > 0\n else: parse_list.append(mask)\n\n parse = np.stack(parse_list)[:,:,:,0] # parse: (cloths, height, width)\n is_bkg = (np.prod(parse == 0, 0) == 1)\n parse = np.argmax(parse, 0) + 1 # add 1 for bkg class\n parse[is_bkg] = 0\n parse[valid_mask == 0 ] = -1\n parse_valid = valid_mask\n\n # load densepose\n dp_xy, dp_uv, dp_I, dp_vertex = data['dp']['xy'], data['dp']['uv'], data['dp']['I'], data['dp']['smpl_v_idx']\n dp_xy = np.concatenate((dp_xy, np.ones_like(dp_xy[:,:1])),1)\n dp_xy = np.dot(img2bb_trans, dp_xy.transpose(1,0)).transpose(1,0)\n dp_I = dp_I - 1 # dp_I is started wtih 1. make it zero-based index.\n\n cur_point_num = len(dp_xy)\n if cur_point_num > cfg.dp_point_num:\n idxs = np.random.choice(np.arange(cur_point_num), size=cfg.dp_point_num)\n cur_point_num = cfg.dp_point_num\n dp_xy = dp_xy[idxs]; dp_uv = dp_uv[idxs]; dp_I = dp_I[idxs]; dp_vertex = dp_vertex[idxs]\n \n # match densepose & parse\n _dp_xy = dp_xy.copy()\n _dp_xy[:,0] = _dp_xy[:,0] / cfg.input_img_shape[1] * cfg.output_parse_shape[1]\n _dp_xy[:,1] = _dp_xy[:,1] / cfg.input_img_shape[0] * cfg.output_parse_shape[0]\n\n parse_onehot = np.zeros((len(cfg.cloth_types)+1, cfg.output_parse_shape[0], cfg.output_parse_shape[1]))\n for i in range(len(cfg.cloth_types)+1):\n parse_onehot[i][parse == i] = 1.0\n\n dp_cloth_idx = np.ones((_dp_xy.shape[0]), np.int16) * -1\n dp_cloth_idx[bilinear_interpolate(parse_fg[None,:,:], _dp_xy[:,0], _dp_xy[:,1])[0] > 0.5] = 0\n for i in range(len(cfg.cloth_types)):\n dp_cloth_idx[np.argmax(bilinear_interpolate(parse_onehot, _dp_xy[:,0], _dp_xy[:,1]), 0) == (i+1)] = i+1\n\n smpl_cloth_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_cloth_idx[dp_vertex] = dp_cloth_idx\n smpl_cloth_valid = (smpl_cloth_idx != -1).astype(np.float32)\n smpl_patch_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_patch_idx[dp_vertex] = dp_I[dp_I != -1]\n\n # remove coat ambiguity\n if (smpl_cloth_idx==cfg.cloth_types.index('uppercloth')+1).sum() == 0 and (smpl_cloth_idx==cfg.cloth_types.index('coat')+1).sum() > 0:\n idxs = (smpl_cloth_idx == cfg.cloth_types.index('coat')+1)\n smpl_cloth_idx[idxs] = cfg.cloth_types.index('uppercloth')+1\n \n # gender\n if data['gender'] == 'male': gender = 1\n elif data['gender'] == 'female': gender = 2\n else: gender = 0\n\n # dummy smpl parameter\n smpl_pose = np.zeros((smpl.joint_num*3,), dtype=np.float32)\n smpl_shape = np.zeros((smpl.shape_param_dim,), dtype=np.float32)\n cam_trans = np.zeros((3,), dtype=np.float32)\n\n inputs = {'img': img}\n targets = {'gender': gender, 'parse': parse, 'smpl_cloth_idx': smpl_cloth_idx, 'smpl_patch_idx': smpl_patch_idx}\n meta_info = {'smpl_cloth_valid': smpl_cloth_valid, 'smpl_pose': smpl_pose, 'smpl_shape': smpl_shape, 'cam_trans': cam_trans}\n return inputs, targets, meta_info\n\n else:\n # smpl processing\n smpl_pose, smpl_shape, smpl_mesh = process_human_model_output(data['smpl_param'], data['cam_param'], do_flip, img_shape, img2bb_trans, rot)\n cam_trans = np.array(data['smpl_param']['trans'], dtype=np.float32)\n cam_param = np.array([cfg.focal[0], cfg.focal[1], cfg.princpt[0], cfg.princpt[1]])\n\n # gender\n if data['gender'] == 'male': gender = 1\n elif data['gender'] == 'female': gender = 2\n else: gender = 0\n\n if cfg.calculate_bcc:\n # load parse (cloth segmentation)\n parse = cv2.imread(data['parse_path'])\n parse_list = []\n for cloth_type in ('fg',) + cfg.cloth_types:\n # get cloth indexs\n if cloth_type == 'fg':\n idxs = np.unique(parse).tolist()\n idxs.pop(idxs.index(0))\n if len(idxs) == 0:\n parse_fg = np.zeros((cfg.output_parse_shape[0], cfg.output_parse_shape[1])) > 0\n continue\n else:\n idxs = self.parse_set[cloth_type]\n \n # get masking corresponding to a cloth\n mask = [parse == i for i in idxs]\n mask = (sum(mask) > 0).astype(np.float32)\n _, _, _, lip2img_trans = generate_patch_image(mask, data['orig_bbox'], 1.0, 0.0, False, mask.shape)\n mask = cv2.warpAffine(mask, lip2img_trans, (img_shape[1], img_shape[0]), flags=cv2.INTER_LINEAR)\n mask = cv2.warpAffine(mask, img2bb_trans, (cfg.input_img_shape[1], cfg.input_img_shape[0]), flags=cv2.INTER_LINEAR)\n mask = cv2.resize(mask, (cfg.output_parse_shape[1], cfg.output_parse_shape[0]))\n \n if cloth_type == 'fg': parse_fg = mask[:,:,0] > 0\n else: parse_list.append(mask)\n\n parse = np.stack(parse_list)[:,:,:,0] # parse: (cloths, height, width)\n is_bkg = (np.prod(parse == 0, 0) == 1)\n parse = np.argmax(parse, 0) + 1 # add 1 for bkg class\n parse[is_bkg] = 0\n parse[valid_mask == 0 ] = -1\n parse_valid = valid_mask\n\n # load densepose\n dp_xy, dp_uv, dp_I, dp_vertex = data['dp']['xy'], data['dp']['uv'], data['dp']['I'], data['dp']['smpl_v_idx']\n dp_xy = np.concatenate((dp_xy, np.ones_like(dp_xy[:,:1])),1)\n dp_xy = np.dot(img2bb_trans, dp_xy.transpose(1,0)).transpose(1,0)\n dp_I = dp_I - 1 # dp_I is started wtih 1. make it zero-based index.\n\n cur_point_num = len(dp_xy)\n if cur_point_num > cfg.dp_point_num:\n idxs = np.random.choice(np.arange(cur_point_num), size=cfg.dp_point_num)\n cur_point_num = cfg.dp_point_num\n dp_xy = dp_xy[idxs]; dp_uv = dp_uv[idxs]; dp_I = dp_I[idxs]; dp_vertex = dp_vertex[idxs]\n \n # match densepose & parse\n _dp_xy = dp_xy.copy()\n _dp_xy[:,0] = _dp_xy[:,0] / cfg.input_img_shape[1] * cfg.output_parse_shape[1]\n _dp_xy[:,1] = _dp_xy[:,1] / cfg.input_img_shape[0] * cfg.output_parse_shape[0]\n\n parse_onehot = np.zeros((len(cfg.cloth_types)+1, cfg.output_parse_shape[0], cfg.output_parse_shape[1]))\n for i in range(len(cfg.cloth_types)+1):\n parse_onehot[i][parse == i] = 1.0\n\n dp_cloth_idx = np.ones((_dp_xy.shape[0]), np.int16) * -1\n dp_cloth_idx[bilinear_interpolate(parse_fg[None,:,:], _dp_xy[:,0], _dp_xy[:,1])[0] > 0.5] = 0\n for i in range(len(cfg.cloth_types)):\n dp_cloth_idx[np.argmax(bilinear_interpolate(parse_onehot, _dp_xy[:,0], _dp_xy[:,1]), 0) == (i+1)] = i+1\n\n smpl_cloth_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_cloth_idx[dp_vertex] = dp_cloth_idx\n smpl_cloth_valid = (smpl_cloth_idx != -1).astype(np.float32)\n smpl_patch_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_patch_idx[dp_vertex] = dp_I[dp_I != -1]\n\n # For BCC calculation, T-posed mesh is used.\n smpl_pose = np.zeros((smpl.joint_num*3,), dtype=np.float32)\n smpl_shape = np.zeros((smpl.shape_param_dim,), dtype=np.float32)\n cam_trans = np.zeros((3,), dtype=np.float32)\n else:\n smpl_cloth_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n smpl_patch_idx = np.ones((smpl.vertex_num), dtype=np.int16) * -1\n\n inputs = {'img': img}\n targets = {'gender': gender, 'smpl_cloth_idx': smpl_cloth_idx, 'smpl_patch_idx': smpl_patch_idx}\n meta_info = {'smpl_pose': smpl_pose, 'smpl_shape': smpl_shape, 'cam_trans': cam_trans, 'cam_param': cam_param}\n return inputs, targets, meta_info\n\n\n def evaluate(self, outs, cur_sample_idx):\n annots = self.datalist\n sample_num = len(outs)\n\n eval_result = {'bcc_upper':[], 'bcc_lower':[], 'bcc_non_cloth':[]}\n for n in range(sample_num):\n annot = annots[cur_sample_idx + n]\n ann_id = annot['ann_id']\n out = outs[n]\n \n if cfg.calculate_bcc:\n for cloth_type in cfg.cloth_types:\n if out[cloth_type + '_mesh'] is None: \n out[cloth_type + '_mesh'] = np.array([])\n else:\n out[cloth_type + '_mesh'] = np.array(out[cloth_type + '_mesh'].vertices)\n \n total_pred_cloth_verts = np.zeros((len(out['smpl_mesh']),))\n gt_smpl_cloth_idx = out['smpl_cloth_idx_target']\n\n for cloth_type in self.eval_types:\n if cloth_type == 'upper_body':\n gt_cloth_verts = ((gt_smpl_cloth_idx == 1) |(gt_smpl_cloth_idx == 2)) # uppercloth, coat\n cloth_idx = 1\n elif cloth_type == 'lower_body':\n gt_cloth_verts = ((gt_smpl_cloth_idx == 3 )| (gt_smpl_cloth_idx == 4)) # pants, skirts\n cloth_idx = 2\n elif cloth_type == 'non_cloth':\n gt_cloth_verts = (gt_smpl_cloth_idx == 0) # non-cloth\n \n gt_idxs = np.where(gt_cloth_verts)[0] \n\n if cloth_type == 'upper_body':\n cloth_verts = torch.cat([torch.from_numpy(out['uppercloth_mesh']).cuda(), torch.from_numpy(out['coat_mesh']).cuda()])\n cloth_verts = cloth_verts[::8].float() # subsampling for memory efficiency\n\n elif cloth_type == 'lower_body':\n cloth_verts = torch.cat([torch.from_numpy(out['pants_mesh']).cuda(), torch.from_numpy(out['skirts_mesh']).cuda()])\n cloth_verts = cloth_verts[::8].float() # subsampling for memory efficiency\n\n smpl_verts = torch.tensor(out['smpl_mesh']).cuda()\n\n if cloth_type in ['upper_body', 'lower_body']:\n if len(cloth_verts) > 0:\n dists = torch.sqrt(torch.sum((smpl_verts[None,:,:] - cloth_verts[:,None,:])**2,2))\n dists = dists.min(0).values\n\n pred_verts = (dists < self.bcc_dist_threshold).cpu().numpy()\n total_pred_cloth_verts[pred_verts] = cloth_idx\n correct_verts = (pred_verts[gt_idxs] == gt_cloth_verts[gt_idxs])\n else:\n correct_verts = np.zeros_like(gt_cloth_verts[gt_idxs], dtype=bool)\n elif cloth_type == 'non_cloth':\n pred_verts = (total_pred_cloth_verts == 0)\n correct_verts = (pred_verts[gt_idxs] == gt_cloth_verts[gt_idxs])\n\n if len(gt_idxs) == 0: continue \n\n if cloth_type == 'upper_body':\n eval_result[f'bcc_upper'].append(correct_verts.sum()/len(gt_idxs))\n elif cloth_type == 'lower_body':\n eval_result[f'bcc_lower'].append(correct_verts.sum()/len(gt_idxs))\n elif cloth_type == 'non_cloth':\n eval_result[f'bcc_non_cloth'].append(correct_verts.sum()/len(gt_idxs))\n\n return eval_result\n\n def print_eval_result(self, eval_result):\n bcc_upper = np.mean(eval_result['bcc_upper'])\n bcc_lower = np.mean(eval_result['bcc_lower'])\n bcc_non_cloth = np.mean(eval_result['bcc_non_cloth'])\n bcc_average = (bcc_upper + bcc_lower + bcc_non_cloth) / 3\n\n print(\">> BCC (upper body) : %.3f\"%bcc_upper)\n print(\">> BCC (lower body) : %.3f\"%bcc_lower)\n print(\">> BCC (non-cloth) : %.3f\"%bcc_non_cloth)\n print(\">> BCC (average) : %.3f\"%bcc_average)\n\n"
},
{
"path": "data/PW3D/PW3D.py",
"content": "import os\nimport os.path as osp\nimport numpy as np\nimport torch\nimport cv2\nimport json\nimport copy\nimport pickle as pkl\nfrom pycocotools.coco import COCO\nfrom config import cfg\nfrom utils.human_models import smpl\nfrom utils.preprocessing import load_img, process_bbox, augmentation, process_human_model_output, convert_focal_princpt\nfrom utils.postprocessing import renderer, rasterize_mesh_given_cam_param, save_proj_faces, merge_mesh, read_valid_point, pa_mpjpe, pairwise_distances\nfrom utils.vis import save_obj, save_result, render_result\n\nclass PW3D(torch.utils.data.Dataset):\n def __init__(self, transform, data_split):\n self.transform = transform\n self.data_split = data_split\n self.data_path = osp.join('..', 'data', 'PW3D', 'data')\n self.sequence_path = osp.join('..', 'data', 'PW3D', 'data', 'sequenceFiles', self.data_split)\n self.preprocessed_path = osp.join('..', 'data', 'preprocessed_data')\n self.eval_stride = 25\n self.cd_inlier_threshold = 32\n\n self.pw3d_smpl_layers = {}\n self.pw3d_beta_clothes = {}\n\n self.datalist = self.load_data()\n print(f\"Load {self.data_split} data: \", len(self.datalist))\n\n def load_data(self):\n db = COCO(osp.join(self.data_path, '3DPW_' + self.data_split + '.json'))\n\n if self.data_split == 'test':\n with open(osp.join(self.preprocessed_path, 'smpl_param', '3DPW_test_Pose2Pose.json')) as f:\n smpl_params = json.load(f)\n else:\n smpl_params = None\n\n datalist = []\n for idx, aid in enumerate(db.anns.keys()):\n if idx % self.eval_stride != 0: continue\n\n ann = db.anns[aid]\n image_id = ann['image_id']\n img = db.loadImgs(image_id)[0]\n sequence_name = img['sequence']\n img_name = img['file_name']\n pid = ann['person_id']\n img_path = osp.join(self.data_path, 'imageFiles', sequence_name, img_name)\n bbox = process_bbox(np.array(ann['bbox']), img['width'], img['height'])\n if bbox is None: continue\n\n cam_param_gt = {k: np.array(v, dtype=np.float32) for k,v in img['cam_param'].items()}\n smpl_param_gt = ann['smpl_param']\n gender = smpl_param_gt['gender']\n\n if str(aid) in smpl_params:\n smpl_param = smpl_params[str(aid)]['smpl_param']\n cam_param = smpl_params[str(aid)]['cam_param']\n else:\n assert 0, \"SMPL params missed!\"\n \n # pre-save smpl layers\n if cfg.calculate_cd:\n sequence = img_path.split('/')[-2]\n index = sequence + '_' + str(pid)\n if index not in self.pw3d_smpl_layers.keys():\n data = pkl.load(open(osp.join(self.sequence_path, f'{sequence}.pkl'), 'rb'), encoding='latin1')\n v_template = data['v_template_clothed'][pid]\n betas_clothed = data['betas_clothed'][pid][:10]\n layer = smpl.get_custom_template_layer(v_template, gender)\n self.pw3d_smpl_layers[index] = layer\n self.pw3d_beta_clothes[index] = betas_clothed\n\n data_dict = {'img_path': img_path, 'ann_id': aid, 'person_id': pid, 'img_shape': (img['height'], img['width']), 'bbox': bbox, \n 'smpl_param': smpl_param, 'cam_param': cam_param, 'smpl_param_gt': smpl_param_gt, 'cam_param_gt': cam_param_gt}\n\n datalist.append(data_dict)\n\n return datalist\n\n def __len__(self):\n return len(self.datalist)\n\n def __getitem__(self, idx):\n data = copy.deepcopy(self.datalist[idx])\n img_path, img_shape = data['img_path'], data['img_shape']\n # image load\n img = load_img(img_path)\n\n # affine transform\n bbox = data['bbox']\n img, valid_mask, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)\n img = self.transform(img.astype(np.float32))/255.\n\n # smpl processing\n smpl_pose, smpl_shape, smpl_mesh = process_human_model_output(data['smpl_param'], data['cam_param'], do_flip, img_shape, img2bb_trans, rot)\n cam_trans = np.array(data['smpl_param']['trans'], dtype=np.float32)\n cam_param = np.array([cfg.focal[0], cfg.focal[1], cfg.princpt[0], cfg.princpt[1]])\n\n inputs = {'img': img}\n targets = {}\n meta_info = {'smpl_pose': smpl_pose, 'smpl_shape': smpl_shape, 'cam_trans': cam_trans, 'cam_param': cam_param}\n\n if cfg.calculate_cd:\n smpl_pose, smpl_shape, cam_trans, gender = data['smpl_param_gt']['pose'], data['smpl_param_gt']['shape'], data['smpl_param_gt']['trans'], data['smpl_param_gt']['gender']\n smpl_pose, smpl_shape, cam_trans = np.array(smpl_pose), np.array(smpl_shape), np.array(cam_trans)\n cam_param = convert_focal_princpt(data['cam_param_gt']['focal'], data['cam_param_gt']['princpt'], img2bb_trans)\n\n smpl_mesh = self.get_clothed_mesh(img_path.split('/')[-2], data['person_id'], smpl_pose, smpl_shape, cam_trans, gender)\n targets['smpl_mesh'] = smpl_mesh\n targets['cam_param'] = cam_param\n \n return inputs, targets, meta_info\n \n def get_clothed_mesh(self, sequence, pid, pose, shape, trans, gender):\n index = sequence + '_' + str(pid)\n layer = self.pw3d_smpl_layers[index]\n betas_clothed = self.pw3d_beta_clothes[index]\n\n pose = torch.FloatTensor(pose).view(1,-1); shape = torch.FloatTensor(betas_clothed).view(1,-1); trans = torch.FloatTensor(trans).view(1,-1)\n output = layer(betas=shape, body_pose=pose[:,3:], global_orient=pose[:,:3], transl=trans)\n mesh_cam = output.vertices[0].numpy()\n\n output = smpl.layer['neutral'](betas=shape, body_pose=pose[:,3:], global_orient=pose[:,:3], transl=trans)\n unclothed_mesh_cam = output.vertices[0].numpy()\n\n trans = np.mean(mesh_cam, axis=0) - np.mean(unclothed_mesh_cam, axis=0)\n mesh_cam -= trans\n return mesh_cam\n\n def evaluate(self, outs, cur_sample_idx):\n annots = self.datalist\n sample_num = len(outs)\n start_idx = 0\n\n eval_result = {'chamfer_distance': []}\n for n in range(sample_num):\n annot = annots[cur_sample_idx + n]\n ann_id = annot['ann_id']\n out = outs[n]\n\n # save ouputs for calcuting cd\n if cfg.calculate_cd:\n verts_out = []; faces_out = []\n verts = out['smpl_mesh']\n verts[:,:2] *= -1\n verts_out.append(verts)\n faces_out.append(smpl.face.astype(np.int32))\n\n for cloth_type in cfg.cloth_types:\n if out[cloth_type + '_mesh'] is None: continue\n verts = out[cloth_type + '_mesh'].vertices\n verts[:,:2] *= -1\n verts_out.append(verts)\n faces_out.append(out[cloth_type + '_mesh'].faces.astype(np.int32))\n \n # pred\n pred_verts, pred_faces = merge_mesh(verts_out, faces_out)\n pred_faces = renderer.rasterize_mesh(torch.from_numpy(pred_verts).float(), torch.from_numpy(pred_faces))\n\n # gt\n gt_verts = out['smpl_mesh_target']; gt_faces = smpl.face.astype(np.int32)\n gt_verts[:,:2] *= -1\n gt_faces = rasterize_mesh_given_cam_param(torch.from_numpy(gt_verts).float(), torch.from_numpy(gt_faces), out['cam_param_target'][:2], out['cam_param_target'][2:])\n\n # find valid pixels - exist silhouette\n pred_faces, gt_faces = pred_faces.numpy().astype(np.int32).reshape(-1, 3), gt_faces.numpy().astype(np.int32).reshape(-1, 3)\n valid = (pred_faces!=-1).sum(1) * (gt_faces!=-1).sum(1)\n valid = valid.reshape(-1)\n\n # if there are too few valid points, not evaluate\n if valid.sum() < self.cd_inlier_threshold:\n continue\n\n # set semantically matching pairs\n paired_pred_verts = read_valid_point(pred_verts, pred_faces, valid)\n paired_gt_verts = read_valid_point(gt_verts, gt_faces, valid)\n\n # rigid alignment\n a, R, t = pa_mpjpe(np.expand_dims(paired_pred_verts,0), np.expand_dims(paired_gt_verts,0))\n pred_verts = (a*np.matmul(pred_verts, R) + t)[0]\n pred_verts *= 1000; gt_verts *= 1000\n\n # pcu.pairwise_distances is too slow, approximate distance between vertices.\n dist1 = pairwise_distances(pred_verts, gt_verts)\n dist2 = pairwise_distances(pred_verts, gt_verts, inv=True)\n\n if torch.isinf(dist1) or torch.isinf(dist2): continue\n\n chamfer_dist = (dist1 + dist2) / 2\n eval_result['chamfer_distance'].append(chamfer_dist)\n \n return eval_result\n\n def print_eval_result(self, eval_result):\n print('>> CD: %.2f mm' % np.mean(eval_result['chamfer_distance']))\n\n \n"
},
{
"path": "data/dataset.py",
"content": "import random\nimport numpy as np\nfrom torch.utils.data.dataset import Dataset\nfrom config import cfg\n\nclass MultipleDatasets(Dataset):\n def __init__(self, dbs, make_same_len=True):\n self.dbs = dbs\n self.db_num = len(self.dbs)\n self.max_db_data_num = max([len(db) for db in dbs])\n self.db_len_cumsum = np.cumsum([len(db) for db in dbs])\n self.make_same_len = make_same_len\n\n def __len__(self):\n # all dbs have the same length\n if self.make_same_len:\n return self.max_db_data_num * self.db_num\n # each db has different length\n else:\n return sum([len(db) for db in self.dbs])\n\n def __getitem__(self, index):\n if self.make_same_len:\n db_idx = index // self.max_db_data_num\n data_idx = index % self.max_db_data_num \n if data_idx >= len(self.dbs[db_idx]) * (self.max_db_data_num // len(self.dbs[db_idx])): # last batch: random sampling\n data_idx = random.randint(0,len(self.dbs[db_idx])-1)\n else: # before last batch: use modular\n data_idx = data_idx % len(self.dbs[db_idx])\n else:\n for i in range(self.db_num):\n if index < self.db_len_cumsum[i]:\n db_idx = i\n break\n if db_idx == 0:\n data_idx = index\n else:\n data_idx = index - self.db_len_cumsum[db_idx-1]\n\n return self.dbs[db_idx][data_idx]\n"
},
{
"path": "demo/demo.py",
"content": "import sys\nimport os\nimport os.path as osp\n\nsys.path.insert(0, osp.join('..', 'main'))\nfrom config import cfg\n\nimport argparse\nimport json\nimport torch\nfrom tqdm import tqdm\nimport numpy as np\nimport torch.backends.cudnn as cudnn\nfrom torch.nn.parallel.data_parallel import DataParallel\nimport torchvision.transforms as transforms\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument('--gpu', type=str, default='0', dest='gpu_ids')\n parser.add_argument('--img_path', type=str, default='./input.jpg')\n parser.add_argument('--json_path', type=str, default='./pose2pose_result.json')\n parser.add_argument('--output_dir', type=str, default='./')\n args = parser.parse_args()\n\n if not args.gpu_ids:\n assert 0, \"Please set propoer gpu ids\"\n\n if '-' in args.gpu_ids:\n gpus = args.gpu_ids.split('-')\n gpus[0] = int(gpus[0])\n gpus[1] = int(gpus[1]) + 1\n args.gpu_ids = ','.join(map(lambda x: str(x), list(range(*gpus))))\n \n return args\n\nargs = parse_args()\ncfg.set_args(args.gpu_ids)\ncudnn.benchmark = True\n\nfrom model import get_model\nfrom base import check_data_parallel\nfrom utils.preprocessing import load_img, process_bbox, generate_patch_image, process_human_model_output\nfrom utils.postprocessing import renderer\nfrom utils.vis import save_result, render_result\n\nmodel_path = os.path.join('.', 'snapshot_7.pth.tar')\nassert osp.exists(model_path), 'Cannot find model at ' + model_path\nprint('Load checkpoint from {}'.format(model_path))\n\nmodel = get_model('test')\nmodel = model.cuda()\nckpt = torch.load(model_path)\nckpt = check_data_parallel(ckpt['network'])\nmodel.load_state_dict(ckpt, strict=False)\nmodel.eval()\n\ntransform = transforms.ToTensor() \noriginal_img = load_img(args.img_path)\noriginal_height, original_width = original_img.shape[:2]\n\nwith open(args.json_path, 'r') as f:\n pose2pose_result = json.load(f)\n\n# prepare bbox\nbbox = [150, 38, 244, 559]\nbbox = process_bbox(bbox, original_width, original_height)\nimg_numpy, _, img2bb_trans, bb2img_trans = generate_patch_image(original_img, bbox, 1.0, 0.0, False, cfg.input_img_shape) \nimg = transform(img_numpy.astype(np.float32))/255.\nimg = img.cuda()[None,:,:,:]\n\nsmpl_pose, smpl_shape, smpl_mesh = process_human_model_output(pose2pose_result['smpl_param'], pose2pose_result['cam_param'], False, (original_width, original_height), img2bb_trans, 0.0)\ncam_trans = np.array(pose2pose_result['smpl_param']['trans'], dtype=np.float32)\nsmpl_pose, smpl_shape, cam_trans = torch.tensor(smpl_pose)[None,:].cuda(), torch.tensor(smpl_shape)[None,:].cuda(), torch.tensor(cam_trans)[None,:].cuda()\n\n# forward\ninputs = {'img': img}\ntargets = {}\nmeta_info = {'smpl_pose': smpl_pose, 'smpl_shape': smpl_shape, 'cam_trans': cam_trans}\n\nwith torch.no_grad():\n out = model(inputs, targets, meta_info, 'test')\n\nfor k,v in out.items():\n if type(v) is torch.Tensor:\n out[k] = v[0].cpu().numpy()\n else:\n out[k] = v[0]\n\nmesh_verts, mesh_faces = save_result(out, osp.join(args.output_dir, 'output.obj'))\nrenderer.set_renderer(focal=cfg.focal, princpt=cfg.princpt, img_shape=cfg.input_img_shape, anti_aliasing=True)\nrender_result(mesh_verts, mesh_faces, img_numpy[:,:,::-1], osp.join(args.output_dir, 'render_cropped_img.jpg'))\n\nfocal = [cfg.focal[0] / cfg.input_img_shape[1] * bbox[2], cfg.focal[1] / cfg.input_img_shape[0] * bbox[3]]\nprincpt = [cfg.princpt[0] / cfg.input_img_shape[1] * bbox[2] + bbox[0], cfg.princpt[1] / cfg.input_img_shape[0] * bbox[3] + bbox[1]]\nrenderer.set_renderer(focal=focal, princpt=princpt, img_shape=(original_height, original_width), anti_aliasing=True)\nrender_result(mesh_verts, mesh_faces, original_img[:,:,::-1], osp.join(args.output_dir, 'render_original_img.jpg'))"
},
{
"path": "demo/pose2pose_result.json",
"content": "{\"smpl_param\": {\"pose\": [-2.9022064208984375, 0.14226926863193512, -0.6876412034034729, -0.37843915820121765, -0.010725350119173527, 0.24553968012332916, -0.0946338027715683, -0.09007483720779419, -0.1379311978816986, 0.3363078832626343, 0.059149932116270065, -0.03379477187991142, 0.5132664442062378, 0.0091248182579875, -0.0329560711979866, 0.5385148525238037, -0.016906920820474625, 0.08675778657197952, 0.0763339102268219, 0.055082399398088455, -0.01292260829359293, 0.006272528320550919, 0.04291265085339546, -0.11267384886741638, -0.1689642071723938, -0.09466277807950974, 0.050063978880643845, -0.1324988752603531, 0.020788447931408882, -0.0006215892499312758, 0.014061491005122662, -0.006473212502896786, 0.004046547692269087, 0.014014238491654396, 0.0007452652789652348, 0.0004654862859752029, 0.05512533709406853, 0.02992839552462101, -0.04127040132880211, -0.17167717218399048, -0.5288858413696289, -0.058588359504938126, -0.036302413791418076, 0.4928983449935913, 0.08087531477212906, -0.15405236184597015, -0.04361267015337944, -0.04372638836503029, -0.33442074060440063, -0.6012314558029175, -0.2605968117713928, 0.026607021689414978, 0.3870042562484741, 0.6654704809188843, -0.12385698407888412, -1.8147525787353516, 0.7356531620025635, 0.14879150688648224, 1.3634296655654907, -0.30997902154922485, -0.10583673417568207, -0.18372346460819244, 0.17793530225753784, -0.17634084820747375, 0.22043076157569885, 0.02467007003724575, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], \"shape\": [-0.04103930667042732, -0.0273419301956892, -0.03568168729543686, -0.05271761119365692, 0.022580627351999283, -0.06175239384174347, -0.017034463584423065, -0.038146961480379105, 0.001704443129710853, -0.040624070912599564], \"trans\": [-0.10590886324644089, 0.16711418330669403, 40.84169387817383]}, \"cam_param\": {\"focal\": [13778.80891164144, 13778.808116912842], \"princpt\": [267.12500619888306, 316.79998779296875]}}"
},
{
"path": "main/config.py",
"content": "import os\nimport os.path as osp\nimport sys\nimport numpy as np\nimport torch\n\nclass Config:\n ## dataset\n trainset_3d = []\n trainset_2d = ['MSCOCO', 'DeepFashion2']\n testset = ['MSCOCO', 'PW3D']\n\n ## model setting\n resnet_type = 50\n input_img_shape = (256, 192) \n output_parse_shape = (256, 192)\n output_dp_shape = (64, 48)\n output_uv_shape = (64, 64)\n focal = (5000, 5000) # virtual focal lengths\n princpt = (input_img_shape[1]/2, input_img_shape[0]/2) # virtual principal point position\n dp_point_num = 196\n dp_patch_num = 24\n cloth_types = ('uppercloth', 'coat', 'pants', 'skirts', 'shoes')\n\n ## training config\n lr = 1e-4\n lr_dec_factor = 10\n lr_dec_epoch = [5]\n end_epoch = 8\n train_batch_size = 8\n sdf_thresh = {'uppercloth': 0.1, 'coat': 0.1, 'pants': 0.1, 'skirts': 0.1, 'shoes': 0.01}\n dist_thresh = {'uppercloth': 0.03, 'coat': 0.1, 'pants': 0.03, 'skirts': 0.03, 'shoes': 0.03}\n min_dist_thresh = {'uppercloth': 0.0, 'coat': 0.03, 'pants': 0.0, 'skirts': 0.0, 'shoes': 0.0}\n\n cls_weight = 0.01\n gender_weight = 0.01\n dp_weight = 1.0\n reg_weight = 0.1\n cloth_reg_weight = {'uppercloth': 1.0, 'coat': 1.0, 'pants': 1.0, 'skirts': 1.0, 'shoes': 0.1}\n\n ## testing config\n test_batch_size = 1\n cls_threshold = 0.25\n calculate_cd = False\n calculate_bcc = False\n cloth_colors = {'smpl_body':(190,190,190), 'uppercloth':(140,110,160), 'coat':(170,120,60), 'pants':(110,130,100), 'skirts':(90,110,140), 'shoes':(120,60,60)}\n\n num_thread = 8\n gpu_ids = '0'\n num_gpus = 1\n continue_train = False\n \n ## directory\n cur_dir = osp.dirname(os.path.abspath(__file__))\n root_dir = osp.join(cur_dir, '..')\n data_dir = osp.join(root_dir, 'data')\n output_dir = osp.join(root_dir, 'output')\n model_dir = osp.join(output_dir, 'model_dump')\n vis_dir = osp.join(output_dir, 'vis')\n os.makedirs(vis_dir, exist_ok=True)\n\n log_dir = osp.join(output_dir, 'log')\n result_dir = osp.join(output_dir, 'result')\n human_model_path = osp.join(data_dir, 'base_data', 'human_models')\n smplicit_path = osp.join(root_dir, 'common', 'utils', 'SMPLicit', 'SMPLicit')\n\n def set_args(self, gpu_ids, continue_train=False):\n self.gpu_ids = gpu_ids\n self.num_gpus = len(self.gpu_ids.split(','))\n self.continue_train = continue_train\n os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = self.gpu_ids\n print('>>> Using GPU: {}'.format(self.gpu_ids))\n\n if self.num_gpus != 1:\n assert 0, \"Not support DataParallel.\"\n\ncfg = Config()\n\nnp.random.seed(0)\ntorch.manual_seed(0)\n\nsys.path.insert(0, osp.join(cfg.root_dir, 'common'))\nfrom utils.dir import add_pypath, make_folder\nadd_pypath(osp.join(cfg.data_dir))\nfor i in range(len(cfg.trainset_2d)):\n add_pypath(osp.join(cfg.data_dir, cfg.trainset_2d[i]))\nfor i in range(len(cfg.testset)):\n add_pypath(osp.join(cfg.data_dir, cfg.testset[i]))\nmake_folder(cfg.model_dir)\nmake_folder(cfg.vis_dir)\nmake_folder(cfg.log_dir)\nmake_folder(cfg.result_dir)"
},
{
"path": "main/model.py",
"content": "import torch\nimport torch.nn as nn\nimport numpy as np\nimport copy\nfrom torch.nn import functional as F\nfrom nets.resnet import ResNetBackbone\nfrom nets.module import ClothNet\nfrom nets.loss import ClothClsLoss, GenderClsLoss, SdfParseLoss, SdfDPLoss, RegLoss\nfrom utils.human_models import smpl\n\nfrom config import cfg\nfrom utils.SMPLicit import SMPLicit\n\n\nclass Model(nn.Module):\n def __init__(self, backbone, cloth_net, mode): \n super(Model, self).__init__()\n self.backbone = backbone\n self.cloth_net = cloth_net\n self.mode = mode\n\n self.smpl_layer = [copy.deepcopy(smpl.layer['neutral']).cuda(),\n copy.deepcopy(smpl.layer['male']).cuda(),\n copy.deepcopy(smpl.layer['female']).cuda()]\n self.smplicit_layer = SMPLicit.SMPLicit(cfg.smplicit_path, cfg.cloth_types).cuda()\n\n if mode == 'train':\n self.cloth_cls_loss = ClothClsLoss()\n self.gender_cls_loss = GenderClsLoss()\n self.sdf_dp_loss = SdfDPLoss()\n self.reg_loss = RegLoss()\n\n self.trainable_modules = [self.backbone, self.cloth_net]\n\n def forward(self, inputs, targets, meta_info, mode):\n batch_size = inputs['img'].shape[0] \n\n # feature extract & get cloth parameter\n img_feat = self.backbone(inputs['img'])\n pred_genders, pred_scores, z_cuts, z_styles = self.cloth_net(img_feat)\n \n # forward SMPL parameters to the SMPL layer\n smpl_pose = meta_info['smpl_pose']\n smpl_shape = meta_info['smpl_shape']\n cam_trans = meta_info['cam_trans']\n\n if mode == 'train':\n # forward cloth & gender parameters to the SMPLicit layer\n smpl_gender = targets['gender']\n sdfs, cloth_meshes, cloth_meshes_unposed = self.smplicit_layer(z_cuts, z_styles, smpl_pose, smpl_shape, smpl_gender, do_marching_cube=(mode=='test'), valid=torch.ones((len(z_cuts),), dtype=torch.bool), do_smooth=False)\n \n # loss functions\n loss = {}\n loss['cloth_cls'] = cfg.cls_weight * self.cloth_cls_loss(pred_scores, targets['smpl_patch_idx'], targets['smpl_cloth_idx'])\n loss['gender_cls'] = cfg.cls_weight * self.gender_cls_loss(pred_genders, smpl_gender)\n loss['sdf_dp'] = 0.0\n loss['reg'] = 0.0\n z_cut_reg, z_style_reg = 0.0, 0.0\n\n for i in range(len(cfg.cloth_types)):\n cloth_type = cfg.cloth_types[i]\n\n if cloth_type == 'uppercloth':\n target_cloth_idx = (i+1, cfg.cloth_types.index('coat')+1)\n else:\n target_cloth_idx = (i+1,)\n\n if cloth_type == 'pants' or cloth_type == 'skirts':\n body_pose = smpl.Astar_pose.float().cuda().repeat(batch_size,1)[:,3:]\n else:\n body_pose = torch.zeros((batch_size,(smpl.joint_num-1)*3)).float().cuda()\n \n # DensePose based loss\n v_template = self.smpl_layer[0](global_orient=torch.zeros((batch_size,3)).float().cuda(), body_pose=body_pose, betas=smpl_shape).vertices\n loss['sdf_dp'] += cfg.dp_weight * self.sdf_dp_loss(sdfs[i], cloth_meshes_unposed[i], targets['smpl_cloth_idx'], meta_info['smpl_cloth_valid'], target_cloth_idx, cfg.sdf_thresh[cloth_type], cfg.dist_thresh[cloth_type], v_template)\n \n # Regularization loss\n cloth_exist = (sum([targets['smpl_cloth_idx'] == idx for idx in target_cloth_idx]) > 0).sum(1) > 0\n if cloth_type != 'shoes': # shoes do not have z_cut\n z_cut_reg += cfg.cloth_reg_weight[cloth_type] * self.reg_loss(z_cuts[i], cloth_exist)\n z_style_reg += cfg.cloth_reg_weight[cloth_type] * self.reg_loss(z_styles[i], cloth_exist)\n \n loss['reg'] = cfg.reg_weight * (z_cut_reg + z_style_reg) / 2.0\n \n return loss\n else:\n pred_clothes = []\n pred_gender = []\n cloth_meshes = []\n \n for i in range(batch_size):\n z_cut = []; z_style = []\n for j in range(len(cfg.cloth_types)): \n z_cut.append(z_cuts[j][i][None,:])\n z_style.append(z_styles[j][i][None,:])\n \n valid_clothes = pred_scores[i] > cfg.cls_threshold\n gender = torch.argmax(pred_genders[i])+1 # male:1, female:2\n \n _, cloth_mesh, _ = self.smplicit_layer(z_cut, z_style, smpl_pose[None, i], smpl_shape[None, i], [gender], True, valid=valid_clothes)\n \n pred_clothes.append(valid_clothes)\n cloth_meshes.append(cloth_mesh)\n pred_gender.append(gender)\n\n cloth_meshes = [[i[0] for i in clothmesh] for clothmesh in zip(*cloth_meshes)]\n \n # add camera translations\n for i in range(len(cfg.cloth_types)):\n for j in range(batch_size):\n if cloth_meshes[i][j] is not None:\n cloth_meshes[i][j].vertices += cam_trans[j].detach().cpu().numpy()\n \n mesh_cam = self.get_coords(smpl_pose[:,:3], {'shape': smpl_shape, 'pose': smpl_pose[:,3:]}, cam_trans, pred_gender)\n \n # output\n out = {} \n out['pred_clothes'] = pred_clothes\n out['pred_gender'] = pred_gender\n out['smpl_mesh'] = mesh_cam\n\n for i,cloth_type in enumerate(cfg.cloth_types):\n out[cloth_type + '_mesh'] = cloth_meshes[i]\n \n for k,v in targets.items():\n out[f'{k}_target'] = v\n \n return out\n\n def get_coords(self, root_pose, params, cam_trans, gender):\n batch_size = root_pose.shape[0]\n\n if self.mode != 'train':\n mesh_cam = []\n for i in range(batch_size):\n output = self.smpl_layer[gender[i]](betas=params['shape'][None,i], body_pose=params['pose'][None,i], global_orient=root_pose[None,i], transl=cam_trans[None,i])\n mesh_cam.append(output.vertices)\n mesh_cam = torch.cat(mesh_cam, dim=0)\n else:\n output = self.smpl_layer[0](betas=params['shape'], body_pose=params['pose'], global_orient=root_pose, transl=cam_trans)\n mesh_cam = output.vertices\n\n return mesh_cam\n\n\ndef init_weights(m):\n if type(m) == nn.ConvTranspose2d:\n nn.init.normal_(m.weight,std=0.001)\n elif type(m) == nn.Conv2d:\n nn.init.normal_(m.weight,std=0.001)\n nn.init.constant_(m.bias, 0)\n elif type(m) == nn.BatchNorm2d:\n nn.init.constant_(m.weight,1)\n nn.init.constant_(m.bias,0)\n elif type(m) == nn.Linear:\n nn.init.normal_(m.weight,std=0.01)\n nn.init.constant_(m.bias,0)\n\n\ndef get_model(mode):\n backbone = ResNetBackbone(cfg.resnet_type)\n cloth_net = ClothNet()\n\n if mode == 'train':\n backbone.init_weights()\n cloth_net.apply(init_weights)\n\n model = Model(backbone, cloth_net, mode)\n return model\n"
},
{
"path": "main/test.py",
"content": "import os\nfrom config import cfg\nimport torch\nimport argparse\nfrom tqdm import tqdm\nimport numpy as np\nimport torch.backends.cudnn as cudnn\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument('--gpu', type=str, default='0', dest='gpu_ids')\n parser.add_argument('--test_epoch', type=str, dest='test_epoch')\n parser.add_argument('--type', type=str)\n args = parser.parse_args()\n\n if not args.gpu_ids:\n assert 0, \"Please set propoer gpu ids\"\n\n if '-' in args.gpu_ids:\n gpus = args.gpu_ids.split('-')\n gpus[0] = int(gpus[0])\n gpus[1] = int(gpus[1]) + 1\n args.gpu_ids = ','.join(map(lambda x: str(x), list(range(*gpus))))\n \n assert args.test_epoch, 'Test epoch is required.'\n return args\n\ndef main():\n args = parse_args()\n cfg.set_args(args.gpu_ids)\n cudnn.benchmark = True\n\n from base import Tester\n\n if args.type == 'cd': \n cfg.calculate_cd = True\n cfg.testset = ['PW3D']\n elif args.type == 'bcc': \n cfg.calculate_bcc = True\n cfg.testset = ['MSCOCO']\n else:\n assert 0, 'Test type is invalid.'\n\n tester = Tester(args.test_epoch)\n tester._make_batch_generator()\n tester._make_model()\n \n eval_result = {}\n cur_sample_idx = 0\n for itr, (inputs, targets, meta_info) in enumerate(tqdm(tester.batch_generator)):\n if itr < cur_sample_idx:\n continue\n \n for k,v in inputs.items():\n if type(v) is torch.Tensor: inputs[k] = v.cuda()\n for k,v in targets.items():\n if type(v) is torch.Tensor: targets[k] = v.cuda()\n for k,v in meta_info.items():\n if type(v) is torch.Tensor: meta_info[k] = v.cuda()\n \n # forward\n with torch.no_grad():\n out = tester.model(inputs, targets, meta_info, 'test')\n\n # save output\n _out = {}\n for k,v in out.items():\n if type(v) is torch.Tensor:\n _out[k] = v.cpu().numpy()\n batch_size = v.shape[0]\n else:\n _out[k] = v\n out = _out\n out = [{k: v[bid] for k,v in out.items()} for bid in range(batch_size)]\n \n # evaluate\n cur_eval_result = tester._evaluate(out, cur_sample_idx)\n for k,v in cur_eval_result.items():\n if k in eval_result: eval_result[k] += v\n else: eval_result[k] = v\n cur_sample_idx += len(out)\n \n tester._print_eval_result(eval_result)\n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "main/train.py",
"content": "import os\nimport argparse\nfrom config import cfg\nimport numpy as np\nimport torch\nimport torch.backends.cudnn as cudnn\nfrom tqdm import tqdm\n\ndef parse_args():\n parser = argparse.ArgumentParser()\n parser.add_argument('--gpu', type=str, default='0', dest='gpu_ids')\n parser.add_argument('--continue', dest='continue_train', action='store_true')\n args = parser.parse_args()\n\n if not args.gpu_ids:\n assert 0, \"Please set proper gpu ids\"\n \n if '-' in args.gpu_ids:\n gpus = args.gpu_ids.split('-')\n gpus[0] = int(gpus[0])\n gpus[1] = int(gpus[1]) + 1\n args.gpu_ids = ','.join(map(lambda x: str(x), list(range(*gpus))))\n\n return args\n\ndef main():\n # argument parse and create log\n args = parse_args()\n cfg.set_args(args.gpu_ids, args.continue_train)\n cudnn.benchmark = True\n\n from base import Trainer\n trainer = Trainer()\n trainer._make_batch_generator()\n trainer._make_model()\n\n # train\n for epoch in range(trainer.start_epoch, cfg.end_epoch):\n trainer.set_lr(epoch)\n trainer.tot_timer.tic()\n trainer.read_timer.tic()\n for itr, (inputs, targets, meta_info) in enumerate(trainer.batch_generator):\n trainer.read_timer.toc()\n trainer.gpu_timer.tic()\n \n # forward\n trainer.optimizer.zero_grad()\n loss = trainer.model(inputs, targets, meta_info, 'train')\n loss = {k:loss[k].mean() for k in loss}\n\n # backward\n sum(loss[k] for k in loss).backward()\n trainer.optimizer.step()\n trainer.gpu_timer.toc()\n screen = [\n 'Epoch %d/%d itr %d/%d:' % (epoch, cfg.end_epoch, itr, trainer.itr_per_epoch),\n 'lr: %g' % (trainer.get_lr()),\n 'speed: %.2f(%.2fs r%.2f)s/itr' % (\n trainer.tot_timer.average_time, trainer.gpu_timer.average_time, trainer.read_timer.average_time),\n '%.2fh/epoch' % (trainer.tot_timer.average_time / 3600. * trainer.itr_per_epoch),\n ]\n screen += ['%s: %.4f' % ('loss_' + k, v.detach()) for k,v in loss.items()]\n trainer.logger.info(' '.join(screen))\n\n trainer.tot_timer.toc()\n trainer.tot_timer.tic()\n trainer.read_timer.tic()\n \n trainer.save_model({\n 'epoch': epoch,\n 'network': trainer.model.state_dict(),\n 'optimizer': trainer.optimizer.state_dict(),\n }, epoch)\n \n\nif __name__ == \"__main__\":\n main()\n"
},
{
"path": "requirements.sh",
"content": "pip install chumpy tqdm torchgeometry trimesh scipy smplx scikit-image opencv-python pycocotools"
}
]