Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,26 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+detectron2/_C.cpython-39-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/overall/21744571_51588794_1000.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/overall/23962182_54027982_1000.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/person/baumu30483223c3_1719437121402_2-0._QL90_UX564_V12524t6_.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/person/mison407622250d_1719258948458_2-0._QL90_UX564_V12524t6_.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/person/mothr22044226e8_1718142523286_2-0._QL90_UX564_V12524t6_.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/upper/21514384_52353349_1000.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/upper/22790049_53294275_1000.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/condition/upper/24083449_54173465_2048.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/men/Simon_1.png filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/men/Yifeng_0.png filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/men/model_5.png filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/men/model_7.png filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/women/049713_0.jpg filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/women/1-model_3.png filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/women/2-model_4.png filter=lfs diff=lfs merge=lfs -text
+resource/demo/example/person/women/model_8.png filter=lfs diff=lfs merge=lfs -text
+resource/img/architecture.jpg filter=lfs diff=lfs merge=lfs -text
+resource/img/comfyui-1.png filter=lfs diff=lfs merge=lfs -text
+resource/img/comfyui.png filter=lfs diff=lfs merge=lfs -text
+resource/img/efficency.jpg filter=lfs diff=lfs merge=lfs -text
+resource/img/structure.jpg filter=lfs diff=lfs merge=lfs -text
+resource/img/teaser.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,14 @@
+__pycache__
+model/__pycache__
+model/DensePose/__pycache__
+model/SCHP/__pycache__
+model/SCHP/*/__pycache__
+resource/demo/output
+resource/demo/example/.DS_Store
+Models
+Datasets
+densepose_
+.vscode
+playground.py
+output
+model/cloth_masker_segformer.py

LICENSE

@@ -0,0 +1,107 @@
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an "as-is" basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses.
+
+Considerations for licensors: Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. More considerations for licensors : wiki.creativecommons.org/Considerations_for_licensors
+
+Considerations for the public: By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor's permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. More considerations for the public : wiki.creativecommons.org/Considerations_for_licensees
+
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
+
+Section 1 – Definitions.
+
+a. Adapted Material means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image.
+b. Adapter's License means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License.
+c. BY-NC-SA Compatible License means a license listed at creativecommons.org/compatiblelicenses, approved by Creative Commons as essentially the equivalent of this Public License.
+d. Copyright and Similar Rights means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights.
+e. Effective Technological Measures means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements.
+f. Exceptions and Limitations means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material.
+g. License Elements means the license attributes listed in the name of a Creative Commons Public License. The License Elements of this Public License are Attribution, NonCommercial, and ShareAlike.
+h. Licensed Material means the artistic or literary work, database, or other material to which the Licensor applied this Public License.
+i. Licensed Rights means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license.
+j. Licensor means the individual(s) or entity(ies) granting rights under this Public License.
+k. NonCommercial means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange.
+l. Share means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them.
+m. Sui Generis Database Rights means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world.
+n. You means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning.
+Section 2 – Scope.
+
+a. License grant.
+1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to:
+A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and
+B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only.
+2. Exceptions and Limitations. For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions.
+3. Term. The term of this Public License is specified in Section 6(a).
+4. Media and formats; technical modifications allowed. The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material.
+5. Downstream recipients.
+A. Offer from the Licensor – Licensed Material. Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License.
+B. Additional offer from the Licensor – Adapted Material. Every recipient of Adapted Material from You automatically receives an offer from the Licensor to exercise the Licensed Rights in the Adapted Material under the conditions of the Adapter's License You apply.
+C. No downstream restrictions. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material.
+6. No endorsement. Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i).
+b. Other rights.
+1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise.
+2. Patent and trademark rights are not licensed under this Public License.
+3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes.
+Section 3 – License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the following conditions.
+
+a. Attribution.
+1. If You Share the Licensed Material (including in modified form), You must:
+A. retain the following if it is supplied by the Licensor with the Licensed Material:
+i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated);
+ii. a copyright notice;
+iii. a notice that refers to this Public License;
+iv. a notice that refers to the disclaimer of warranties;
+v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable;
+
+B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and
+C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License.
+2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information.
+3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable.
+b. ShareAlike.In addition to the conditions in Section 3(a), if You Share Adapted Material You produce, the following conditions also apply.
+1. The Adapter's License You apply must be a Creative Commons license with the same License Elements, this version or later, or a BY-NC-SA Compatible License.
+2. You must include the text of, or the URI or hyperlink to, the Adapter's License You apply. You may satisfy this condition in any reasonable manner based on the medium, means, and context in which You Share Adapted Material.
+3. You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, Adapted Material that restrict exercise of the rights granted under the Adapter's License You apply.
+Section 4 – Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material:
+
+a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only;
+b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material, including for purposes of Section 3(b); and
+c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database.
+For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights.
+Section 5 – Disclaimer of Warranties and Limitation of Liability.
+
+a. Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.
+b. To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.
+c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability.
+Section 6 – Term and Termination.
+
+a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically.
+b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates:
+1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or
+2. upon express reinstatement by the Licensor.
+For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License.
+
+c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License.
+d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License.
+Section 7 – Other Terms and Conditions.
+
+a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed.
+b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License.
+Section 8 – Interpretation.
+
+a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License.
+b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions.
+c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor.
+d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority.
+Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the "Licensor." The text of the Creative Commons public licenses is dedicated to the public domain under the CC0 Public Domain Dedication. Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at creativecommons.org/policies, Creative Commons does not authorize the use of the trademark "Creative Commons" or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

README.md

@@ -1,12 +1,195 @@
 ---
-title: Vtontry
-emoji: 🐨
-colorFrom: purple
-colorTo: blue
-sdk: gradio
-sdk_version: 5.35.0
+title: vtontry
 app_file: app.py
-pinned: false
+sdk: gradio
+sdk_version: 4.41.0
 ---
+# [ICLR 25]🐈 CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models
+
+<div style="display: flex; justify-content: center; align-items: center;">
+  <a href="http://arxiv.org/abs/2407.15886" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/arXiv-2407.15886-red?style=flat&logo=arXiv&logoColor=red' alt='arxiv'>
+  </a>
+  <a href='https://huggingface.co/zhengchong/CatVTON' style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Hugging Face-ckpts-orange?style=flat&logo=HuggingFace&logoColor=orange' alt='huggingface'>
+  </a>
+  <a href="https://github.com/Zheng-Chong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/GitHub-Repo-blue?style=flat&logo=GitHub' alt='GitHub'>
+  </a>
+  <a href="https://huggingface.co/spaces/zhengchong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Demo-Gradio-gold?style=flat&logo=Gradio&logoColor=red' alt='Demo'>
+  </a>
+  <a href="https://huggingface.co/spaces/zhengchong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Space-ZeroGPU-orange?style=flat&logo=Gradio&logoColor=red' alt='Demo'>
+  </a>
+  <a href='https://zheng-chong.github.io/CatVTON/' style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Webpage-Project-silver?style=flat&logo=&logoColor=orange' alt='webpage'>
+  </a>
+  <a href="https://github.com/Zheng-Chong/CatVTON/LICENCE" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/License-CC BY--NC--SA--4.0-lightgreen?style=flat&logo=Lisence' alt='License'>
+  </a>
+</div>
+
+
+**CatVTON** is a simple and efficient virtual try-on diffusion model with ***1) Lightweight Network (899.06M parameters totally)***, ***2) Parameter-Efficient Training (49.57M parameters trainable)*** and ***3) Simplified Inference (< 8G VRAM for 1024X768 resolution)***.
+<div align="center">
+  <img src="resource/img/teaser.jpg" width="100%" height="100%"/>
+</div>
+
+
+
+## Updates 
+- **`2025/02/24`**: 🎉 We are excited to announce [**CatV2TON**](https://github.com/Zheng-Chong/CatV2TON), our new DiT-based model that supports both **image and video try-on**! Check it out!
+- **`2025/02/20`**: Our [**Paper on ArXiv**](http://arxiv.org/abs/2407.15886) has been updated to v2, which includes more details.
+- **`2025/01/24`**: 🥳 CatVTON has been accepted to **ICLR 2025**!
+- **`2024/12/20`**: 😄 Code for gradio app of **CatVTON-FLUX** has been released! It is not a stable version, but it is a good start!
+- **`2024/12/19`**: [**CatVTON-FLUX**](https://huggingface.co/spaces/zhengchong/CatVTON) has been released! It is a extremely lightweight LoRA (only 37.4M checkpints) for [FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev), the lora weights are available in **[huggingface repo](https://huggingface.co/zhengchong/CatVTON/tree/main/flux-lora)**, code will be released soon!
+- **`2024/11/26`**: Our **unified vision-based model for image and video try-on** will be released soon, bringing a brand-new virtual try-on experience! While our demo page will be temporarily taken offline, [**the demo on HuggingFace Space**](https://huggingface.co/spaces/zhengchong/CatVTON) will remain available for use !
+- **`2024/10/17`**:[**Mask-free version**](https://huggingface.co/zhengchong/CatVTON-MaskFree)🤗 of CatVTON is release ! 
+- **`2024/10/13`**: We have built a repo [**Awesome-Try-On-Models**](https://github.com/Zheng-Chong/Awesome-Try-On-Models) that focuses on image, video, and 3D-based try-on models published after 2023, aiming to provide insights into the latest technological trends. If you're interested, feel free to contribute or give it a 🌟 star!
+- **`2024/08/13`**: We localize DensePose & SCHP to avoid certain environment issues.
+- **`2024/08/10`**: Our 🤗 [**HuggingFace Space**](https://huggingface.co/spaces/zhengchong/CatVTON) is available now! Thanks for the grant from [**ZeroGPU**](https://huggingface.co/zero-gpu-explorers)！
+- **`2024/08/09`**: [**Evaluation code**](https://github.com/Zheng-Chong/CatVTON?tab=readme-ov-file#3-calculate-metrics) is provided to calculate metrics 📚.
+- **`2024/07/27`**: We provide code and workflow for deploying CatVTON on [**ComfyUI**](https://github.com/Zheng-Chong/CatVTON?tab=readme-ov-file#comfyui-workflow) 💥.
+- **`2024/07/24`**: Our [**Paper on ArXiv**](http://arxiv.org/abs/2407.15886) is available 🥳!
+- **`2024/07/22`**: Our [**App Code**](https://github.com/Zheng-Chong/CatVTON/blob/main/app.py) is released, deploy and enjoy CatVTON on your mechine 🎉!
+- **`2024/07/21`**: Our [**Inference Code**](https://github.com/Zheng-Chong/CatVTON/blob/main/inference.py) and [**Weights** 🤗](https://huggingface.co/zhengchong/CatVTON) are released.
+- **`2024/07/11`**: Our [**Online Demo**](https://huggingface.co/spaces/zhengchong/CatVTON) is released 😁.
+
+
+
+
+## Installation
+
+Create a conda environment & Install requirments
+```shell
+conda create -n catvton python==3.9.0
+conda activate catvton
+cd CatVTON-main  # or your path to CatVTON project dir
+pip install -r requirements.txt
+```
+
+## Deployment 
+### ComfyUI Workflow
+We have modified the main code to enable easy deployment of CatVTON on [ComfyUI](https://github.com/comfyanonymous/ComfyUI). Due to the incompatibility of the code structure, we have released this part in the [Releases](https://github.com/Zheng-Chong/CatVTON/releases/tag/ComfyUI), which includes the code placed under `custom_nodes` of ComfyUI and our workflow JSON files.
+
+To deploy CatVTON to your ComfyUI, follow these steps:
+1. Install all the requirements for both CatVTON and ComfyUI, refer to [Installation Guide for CatVTON](https://github.com/Zheng-Chong/CatVTON/blob/main/INSTALL.md) and [Installation Guide for ComfyUI](https://github.com/comfyanonymous/ComfyUI?tab=readme-ov-file#installing).
+2. Download [`ComfyUI-CatVTON.zip`](https://github.com/Zheng-Chong/CatVTON/releases/download/ComfyUI/ComfyUI-CatVTON.zip) and unzip it in the `custom_nodes` folder under your ComfyUI project (clone from [ComfyUI](https://github.com/comfyanonymous/ComfyUI)).
+3. Run the ComfyUI.
+4. Download [`catvton_workflow.json`](https://github.com/Zheng-Chong/CatVTON/releases/download/ComfyUI/catvton_workflow.json) and drag it into you ComfyUI webpage and enjoy 😆!
+
+> Problems under Windows OS, please refer to [issue#8](https://github.com/Zheng-Chong/CatVTON/issues/8).
+> 
+When you run the CatVTON workflow for the first time, the weight files will be automatically downloaded, usually taking dozens of minutes.
+
+<div align="center">
+  <img src="resource/img/comfyui-1.png" width="100%" height="100%"/>
+</div>
+
+<!-- <div align="center">
+ <img src="resource/img/comfyui.png" width="100%" height="100%"/>
+</div> -->
+
+### Gradio App
+
+To deploy the Gradio App for CatVTON on your machine, run the following command, and checkpoints will be automatically downloaded from HuggingFace.
+
+```PowerShell
+CUDA_VISIBLE_DEVICES=0 python app.py \
+--output_dir="resource/demo/output" \
+--mixed_precision="bf16" \
+--allow_tf32 
+```
+When using `bf16` precision, generating results with a resolution of `1024x768` only requires about `8G` VRAM.
+
+## Inference
+### 1. Data Preparation
+Before inference, you need to download the [VITON-HD](https://github.com/shadow2496/VITON-HD) or [DressCode](https://github.com/aimagelab/dress-code) dataset.
+Once the datasets are downloaded, the folder structures should look like these:
+```
+├── VITON-HD
+|   ├── test_pairs_unpaired.txt
+│   ├── test
+|   |   ├── image
+│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
+│   │   ├── cloth
+│   │   │   ├── [000006_00.jpg | 000008_00.jpg | ...]
+│   │   ├── agnostic-mask
+│   │   │   ├── [000006_00_mask.png | 000008_00.png | ...]
+...
+```
+
+```
+├── DressCode
+|   ├── test_pairs_paired.txt
+|   ├── test_pairs_unpaired.txt
+│   ├── [dresses | lower_body | upper_body]
+|   |   ├── test_pairs_paired.txt
+|   |   ├── test_pairs_unpaired.txt
+│   │   ├── images
+│   │   │   ├── [013563_0.jpg | 013563_1.jpg | 013564_0.jpg | 013564_1.jpg | ...]
+│   │   ├── agnostic_masks
+│   │   │   ├── [013563_0.png| 013564_0.png | ...]
+...
+```
+For the DressCode dataset, we provide script to preprocessed agnostic masks, run the following command:
+```PowerShell
+CUDA_VISIBLE_DEVICES=0 python preprocess_agnostic_mask.py \
+--data_root_path <your_path_to_DressCode> 
+```
+
+### 2. Inference on VTIONHD/DressCode
+To run the inference on the DressCode or VITON-HD dataset, run the following command, checkpoints will be automatically downloaded from HuggingFace.
+
+```PowerShell
+CUDA_VISIBLE_DEVICES=0 python inference.py \
+--dataset [dresscode | vitonhd] \
+--data_root_path <path> \
+--output_dir <path> 
+--dataloader_num_workers 8 \
+--batch_size 8 \
+--seed 555 \
+--mixed_precision [no | fp16 | bf16] \
+--allow_tf32 \
+--repaint \
+--eval_pair  
+```
+### 3. Calculate Metrics
+
+After obtaining the inference results, calculate the metrics using the following command: 
+
+```PowerShell
+CUDA_VISIBLE_DEVICES=0 python eval.py \
+--gt_folder <your_path_to_gt_image_folder> \
+--pred_folder <your_path_to_predicted_image_folder> \
+--paired \
+--batch_size=16 \
+--num_workers=16 
+```
+
+-  `--gt_folder` and `--pred_folder` should be folders that contain **only images**.
+- To evaluate the results in a paired setting, use `--paired`; for an unpaired setting, simply omit it.
+- `--batch_size` and `--num_workers` should be adjusted based on your machine.
+
+
+## Acknowledgement
+Our code is modified based on [Diffusers](https://github.com/huggingface/diffusers). We adopt [Stable Diffusion v1.5 inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting) as the base model. We use [SCHP](https://github.com/GoGoDuck912/Self-Correction-Human-Parsing/tree/master) and [DensePose](https://github.com/facebookresearch/DensePose) to automatically generate masks in our [Gradio](https://github.com/gradio-app/gradio) App and [ComfyUI](https://github.com/comfyanonymous/ComfyUI) workflow. Thanks to all the contributors!
+
+## License
+All the materials, including code, checkpoints, and demo, are made available under the [Creative Commons BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/) license. You are free to copy, redistribute, remix, transform, and build upon the project for non-commercial purposes, as long as you give appropriate credit and distribute your contributions under the same license.
+
+
+## Citation
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+```bibtex
+@misc{chong2024catvtonconcatenationneedvirtual,
+ title={CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models}, 
+ author={Zheng Chong and Xiao Dong and Haoxiang Li and Shiyue Zhang and Wenqing Zhang and Xujie Zhang and Hanqing Zhao and Xiaodan Liang},
+ year={2024},
+ eprint={2407.15886},
+ archivePrefix={arXiv},
+ primaryClass={cs.CV},
+ url={https://arxiv.org/abs/2407.15886}, 
+}
+```

app.py

@@ -0,0 +1,391 @@
+import argparse
+import os
+from datetime import datetime
+
+import gradio as gr
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+# Set memory growth for MPS to prevent out of memory errors
+os.environ['PYTORCH_MPS_HIGH_WATERMARK_RATIO'] = '0.0'
+
+from model.cloth_masker import AutoMasker, vis_mask
+from model.pipeline import CatVTONPipeline
+from utils import init_weight_dtype, resize_and_crop, resize_and_padding
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--base_model_path",
+        type=str,
+        default="booksforcharlie/stable-diffusion-inpainting",  # Change to a copy repo as runawayml delete original repo
+        help=(
+            "The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub."
+        ),
+    )
+    parser.add_argument(
+        "--resume_path",
+        type=str,
+        default="zhengchong/CatVTON",
+        help=(
+            "The Path to the checkpoint of trained tryon model."
+        ),
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="resource/demo/output",
+        help="The output directory where the model predictions will be written.",
+    )
+
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=768,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=1024,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--repaint", 
+        action="store_true", 
+        help="Whether to repaint the result image with the original background."
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        default=True,
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    
+    args = parser.parse_args()
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    return args
+
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows * cols
+
+    w, h = imgs[0].size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+args = parse_args()
+repo_path = snapshot_download(repo_id=args.resume_path)
+
+# Auto-detect device (CUDA if available, otherwise CPU)
+# Note: MPS is disabled due to memory and compatibility issues
+if torch.cuda.is_available():
+    device = 'cuda'
+else:
+    device = 'cpu'
+    print("Note: Running on CPU. This will be slower but more stable.")
+print(f"Using device: {device}")
+
+# Pipeline
+pipeline = CatVTONPipeline(
+    base_ckpt=args.base_model_path,
+    attn_ckpt=repo_path,
+    attn_ckpt_version="mix",
+    weight_dtype=init_weight_dtype(args.mixed_precision),
+    use_tf32=args.allow_tf32 and torch.cuda.is_available(),  # Only use TF32 if CUDA is available
+    device=device
+)
+# AutoMasker
+mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
+automasker = AutoMasker(
+    densepose_ckpt=os.path.join(repo_path, "DensePose"),
+    schp_ckpt=os.path.join(repo_path, "SCHP"),
+    device=device, 
+)
+
+def submit_function(
+    person_image,
+    cloth_image,
+    cloth_type,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    show_type
+):
+    person_image, mask = person_image["background"], person_image["layers"][0]
+    mask = Image.open(mask).convert("L")
+    if len(np.unique(np.array(mask))) == 1:
+        mask = None
+    else:
+        mask = np.array(mask)
+        mask[mask > 0] = 255
+        mask = Image.fromarray(mask)
+
+    tmp_folder = args.output_dir
+    date_str = datetime.now().strftime("%Y%m%d%H%M%S")
+    result_save_path = os.path.join(tmp_folder, date_str[:8], date_str[8:] + ".png")
+    if not os.path.exists(os.path.join(tmp_folder, date_str[:8])):
+        os.makedirs(os.path.join(tmp_folder, date_str[:8]))
+
+    generator = None
+    if seed != -1:
+        generator = torch.Generator(device=device).manual_seed(seed)
+
+    person_image = Image.open(person_image).convert("RGB")
+    cloth_image = Image.open(cloth_image).convert("RGB")
+    
+    # Use default resolution
+    target_width = args.width
+    target_height = args.height
+        
+    person_image = resize_and_crop(person_image, (target_width, target_height))
+    cloth_image = resize_and_padding(cloth_image, (target_width, target_height))
+    
+    # Process mask
+    if mask is not None:
+        mask = resize_and_crop(mask, (target_width, target_height))
+    else:
+        mask = automasker(
+            person_image,
+            cloth_type
+        )['mask']
+    mask = mask_processor.blur(mask, blur_factor=9)
+
+    # Inference
+    # try:
+    result_image = pipeline(
+        image=person_image,
+        condition_image=cloth_image,
+        mask=mask,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    )[0]
+    # except Exception as e:
+    #     raise gr.Error(
+    #         "An error occurred. Please try again later: {}".format(e)
+    #     )
+    
+    # Post-process
+    masked_person = vis_mask(person_image, mask)
+    save_result_image = image_grid([person_image, masked_person, cloth_image, result_image], 1, 4)
+    save_result_image.save(result_save_path)
+    if show_type == "result only":
+        return result_image
+    else:
+        width, height = person_image.size
+        if show_type == "input & result":
+            condition_width = width // 2
+            conditions = image_grid([person_image, cloth_image], 2, 1)
+        else:
+            condition_width = width // 3
+            conditions = image_grid([person_image, masked_person , cloth_image], 3, 1)
+        conditions = conditions.resize((condition_width, height), Image.NEAREST)
+        new_result_image = Image.new("RGB", (width + condition_width + 5, height))
+        new_result_image.paste(conditions, (0, 0))
+        new_result_image.paste(result_image, (condition_width + 5, 0))
+    return new_result_image
+
+
+def person_example_fn(image_path):
+    return image_path
+
+HEADER = """
+<h1 style="text-align: center;"> 🐈 CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models </h1>
+<div style="display: flex; justify-content: center; align-items: center;">
+  <a href="http://arxiv.org/abs/2407.15886" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/arXiv-2407.15886-red?style=flat&logo=arXiv&logoColor=red' alt='arxiv'>
+  </a>
+  <a href='https://huggingface.co/zhengchong/CatVTON' style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Hugging Face-ckpts-orange?style=flat&logo=HuggingFace&logoColor=orange' alt='huggingface'>
+  </a>
+  <a href="https://github.com/Zheng-Chong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/GitHub-Repo-blue?style=flat&logo=GitHub' alt='GitHub'>
+  </a>
+  <a href="http://120.76.142.206:8888" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Demo-Gradio-gold?style=flat&logo=Gradio&logoColor=red' alt='Demo'>
+  </a>
+  <a href="https://huggingface.co/spaces/zhengchong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Space-ZeroGPU-orange?style=flat&logo=Gradio&logoColor=red' alt='Demo'>
+  </a>
+  <a href='https://zheng-chong.github.io/CatVTON/' style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Webpage-Project-silver?style=flat&logo=&logoColor=orange' alt='webpage'>
+  </a>
+  <a href="https://github.com/Zheng-Chong/CatVTON/LICENCE" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/License-CC BY--NC--SA--4.0-lightgreen?style=flat&logo=Lisence' alt='License'>
+  </a>
+</div>
+<br>
+· This demo and our weights are only for <span>Non-commercial Use</span>. <br>
+· You can try CatVTON in our <a href="https://huggingface.co/spaces/zhengchong/CatVTON">HuggingFace Space</a> or our <a href="http://120.76.142.206:8888">online demo</a> (run on 3090). <br>
+· Thanks to <a href="https://huggingface.co/zero-gpu-explorers">ZeroGPU</a> for providing A100 for our <a href="https://huggingface.co/spaces/zhengchong/CatVTON">HuggingFace Space</a>. <br>
+· SafetyChecker is set to filter NSFW content, but it may block normal results too. Please adjust the <span>`seed`</span> for normal outcomes.<br> 
+"""
+
+def app_gradio():
+    with gr.Blocks(title="CatVTON") as demo:
+        gr.Markdown(HEADER)
+        with gr.Row():
+            with gr.Column(scale=1, min_width=350):
+                with gr.Row():
+                    image_path = gr.Image(
+                        type="filepath",
+                        interactive=True,
+                        visible=False,
+                    )
+                    person_image = gr.ImageEditor(
+                        interactive=True, label="Person Image", type="filepath"
+                    )
+
+                with gr.Row():
+                    with gr.Column(scale=1, min_width=230):
+                        cloth_image = gr.Image(
+                            interactive=True, label="Condition Image", type="filepath"
+                        )
+                    with gr.Column(scale=1, min_width=120):
+                        gr.Markdown(
+                            '<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `🖌️` above to draw the Mask (higher priority)<br>2. Select the `Try-On Cloth Type` to generate automatically </span>'
+                        )
+                        cloth_type = gr.Radio(
+                            label="Try-On Cloth Type",
+                            choices=["upper", "lower", "overall"],
+                            value="upper",
+                        )
+
+
+                submit = gr.Button("Submit")
+                gr.Markdown(
+                    '<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>'
+                )
+                
+                gr.Markdown(
+                    '<span style="color: #808080; font-size: small;">Advanced options can adjust details:<br>1. `Inference Step` may enhance details;<br>2. `CFG` is highly correlated with saturation;<br>3. `Random seed` may improve pseudo-shadow.</span>'
+                )
+                with gr.Accordion("Advanced Options", open=False):
+                    num_inference_steps = gr.Slider(
+                        label="Inference Step", minimum=10, maximum=100, step=5, value=50
+                    )
+                    # Guidence Scale
+                    guidance_scale = gr.Slider(
+                        label="CFG Strenth", minimum=0.0, maximum=7.5, step=0.5, value=2.5
+                    )
+                    # Random Seed
+                    seed = gr.Slider(
+                        label="Seed", minimum=-1, maximum=10000, step=1, value=42
+                    )
+                    show_type = gr.Radio(
+                        label="Show Type",
+                        choices=["result only", "input & result", "input & mask & result"],
+                        value="input & mask & result",
+                    )
+
+            with gr.Column(scale=2, min_width=500):
+                result_image = gr.Image(interactive=False, label="Result")
+                with gr.Row():
+                    # Photo Examples
+                    root_path = "resource/demo/example"
+                    with gr.Column():
+                        men_exm = gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "person", "men", _)
+                                for _ in os.listdir(os.path.join(root_path, "person", "men"))
+                            ],
+                            examples_per_page=4,
+                            inputs=image_path,
+                            label="Person Examples ①",
+                        )
+                        women_exm = gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "person", "women", _)
+                                for _ in os.listdir(os.path.join(root_path, "person", "women"))
+                            ],
+                            examples_per_page=4,
+                            inputs=image_path,
+                            label="Person Examples ②",
+                        )
+                        gr.Markdown(
+                            '<span style="color: #808080; font-size: small;">*Person examples come from the demos of <a href="https://huggingface.co/spaces/levihsu/OOTDiffusion">OOTDiffusion</a> and <a href="https://www.outfitanyone.org">OutfitAnyone</a>. </span>'
+                        )
+                    with gr.Column():
+                        condition_upper_exm = gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "condition", "upper", _)
+                                for _ in os.listdir(os.path.join(root_path, "condition", "upper"))
+                            ],
+                            examples_per_page=4,
+                            inputs=cloth_image,
+                            label="Condition Upper Examples",
+                        )
+                        condition_overall_exm = gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "condition", "overall", _)
+                                for _ in os.listdir(os.path.join(root_path, "condition", "overall"))
+                            ],
+                            examples_per_page=4,
+                            inputs=cloth_image,
+                            label="Condition Overall Examples",
+                        )
+                        condition_person_exm = gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "condition", "person", _)
+                                for _ in os.listdir(os.path.join(root_path, "condition", "person"))
+                            ],
+                            examples_per_page=4,
+                            inputs=cloth_image,
+                            label="Condition Reference Person Examples",
+                        )
+                        gr.Markdown(
+                            '<span style="color: #808080; font-size: small;">*Condition examples come from the Internet. </span>'
+                        )
+
+            image_path.change(
+                person_example_fn, inputs=image_path, outputs=person_image
+            )
+
+            submit.click(
+                submit_function,
+                [
+                    person_image,
+                    cloth_image,
+                    cloth_type,
+                    num_inference_steps,
+                    guidance_scale,
+                    seed,
+                    show_type,
+                ],
+                result_image,
+            )
+    demo.queue().launch(share=True, show_error=True)
+
+
+if __name__ == "__main__":
+    app_gradio()

app_flux.py

@@ -0,0 +1,305 @@
+import os
+import argparse
+import gradio as gr
+from datetime import datetime
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+from model.cloth_masker import AutoMasker, vis_mask
+from model.flux.pipeline_flux_tryon import FluxTryOnPipeline
+from utils import resize_and_crop, resize_and_padding
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="FLUX Try-On Demo")
+    parser.add_argument(
+        "--base_model_path",
+        type=str,
+        # default="black-forest-labs/FLUX.1-Fill-dev",
+        default="Models/FLUX.1-Fill-dev",
+        help="The path to the base model to use for evaluation."
+    )
+    parser.add_argument(
+        "--resume_path",
+        type=str,
+        default="zhengchong/CatVTON",
+        help="The Path to the checkpoint of trained tryon model."
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="resource/demo/output",
+        help="The output directory where the model predictions will be written."
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="bf16",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision."
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        default=True,
+        help="Whether or not to allow TF32 on Ampere GPUs."
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=768,
+        help="The width of the input image."
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=1024,
+        help="The height of the input image."
+    )
+    return parser.parse_args()
+
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows * cols
+    w, h = imgs[0].size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+def submit_function_flux(
+    person_image,
+    cloth_image,
+    cloth_type,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    show_type
+):
+
+    # Process image editor input
+    person_image, mask = person_image["background"], person_image["layers"][0]
+    mask = Image.open(mask).convert("L")
+    if len(np.unique(np.array(mask))) == 1:
+        mask = None
+    else:
+        mask = np.array(mask)
+        mask[mask > 0] = 255
+        mask = Image.fromarray(mask)
+
+    # Set random seed
+    generator = None
+    if seed != -1:
+        generator = torch.Generator(device='cuda').manual_seed(seed)
+
+    # Process input images
+    person_image = Image.open(person_image).convert("RGB")
+    cloth_image = Image.open(cloth_image).convert("RGB")
+    
+    # Adjust image sizes
+    person_image = resize_and_crop(person_image, (args.width, args.height))
+    cloth_image = resize_and_padding(cloth_image, (args.width, args.height))
+
+    # Process mask
+    if mask is not None:
+        mask = resize_and_crop(mask, (args.width, args.height))
+    else:
+        mask = automasker(
+            person_image,
+            cloth_type
+        )['mask']
+    mask = mask_processor.blur(mask, blur_factor=9)
+
+    # Inference
+    result_image = pipeline_flux(
+        image=person_image,
+        condition_image=cloth_image,
+        mask_image=mask,
+        height=args.height,
+        width=args.width,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    ).images[0]
+
+    # Post-processing
+    masked_person = vis_mask(person_image, mask)
+
+    # Return result based on show type
+    if show_type == "result only":
+        return result_image
+    else:
+        width, height = person_image.size
+        if show_type == "input & result":
+            condition_width = width // 2
+            conditions = image_grid([person_image, cloth_image], 2, 1)
+        else:
+            condition_width = width // 3
+            conditions = image_grid([person_image, masked_person, cloth_image], 3, 1)
+        
+        conditions = conditions.resize((condition_width, height), Image.NEAREST)
+        new_result_image = Image.new("RGB", (width + condition_width + 5, height))
+        new_result_image.paste(conditions, (0, 0))
+        new_result_image.paste(result_image, (condition_width + 5, 0))
+        return new_result_image
+
+def person_example_fn(image_path):
+    return image_path
+
+
+def app_gradio():
+    with gr.Blocks(title="CatVTON with FLUX.1-Fill-dev") as demo:
+        gr.Markdown("# CatVTON with FLUX.1-Fill-dev")
+        with gr.Row():
+            with gr.Column(scale=1, min_width=350):
+                with gr.Row():
+                    image_path_flux = gr.Image(
+                        type="filepath",
+                        interactive=True,
+                        visible=False,
+                    )
+                    person_image_flux = gr.ImageEditor(
+                        interactive=True, label="Person Image", type="filepath"
+                    )
+                
+                with gr.Row():
+                    with gr.Column(scale=1, min_width=230):
+                        cloth_image_flux = gr.Image(
+                            interactive=True, label="Condition Image", type="filepath"
+                        )
+                    with gr.Column(scale=1, min_width=120):
+                        gr.Markdown(
+                            '<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `🖌️` above to draw the Mask (higher priority)<br>2. Select the `Try-On Cloth Type` to generate automatically </span>'
+                        )
+                        cloth_type = gr.Radio(
+                            label="Try-On Cloth Type",
+                            choices=["upper", "lower", "overall"],
+                            value="upper",
+                        )
+
+                submit_flux = gr.Button("Submit")
+                gr.Markdown(
+                    '<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>'
+                )
+                
+                with gr.Accordion("Advanced Options", open=False):
+                    num_inference_steps_flux = gr.Slider(
+                        label="Inference Step", minimum=10, maximum=100, step=5, value=50
+                    )
+                    # Guidence Scale
+                    guidance_scale_flux = gr.Slider(
+                        label="CFG Strenth", minimum=0.0, maximum=50, step=0.5, value=30
+                    )
+                    # Random Seed
+                    seed_flux = gr.Slider(
+                        label="Seed", minimum=-1, maximum=10000, step=1, value=42
+                    )
+                    show_type = gr.Radio(
+                        label="Show Type",
+                        choices=["result only", "input & result", "input & mask & result"],
+                        value="input & mask & result",
+                    )
+                
+            with gr.Column(scale=2, min_width=500):
+                result_image_flux = gr.Image(interactive=False, label="Result")
+                with gr.Row():
+                    # Photo Examples
+                    root_path = "resource/demo/example"
+                    with gr.Column():
+                        gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "person", "men", _)
+                                for _ in os.listdir(os.path.join(root_path, "person", "men"))
+                            ],
+                            examples_per_page=4,
+                            inputs=image_path_flux,
+                            label="Person Examples ①",
+                        )
+                        gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "person", "women", _)
+                                for _ in os.listdir(os.path.join(root_path, "person", "women"))
+                            ],
+                            examples_per_page=4,
+                            inputs=image_path_flux,
+                            label="Person Examples ②",
+                        )
+                        gr.Markdown(
+                            '<span style="color: #808080; font-size: small;">*Person examples come from the demos of <a href="https://huggingface.co/spaces/levihsu/OOTDiffusion">OOTDiffusion</a> and <a href="https://www.outfitanyone.org">OutfitAnyone</a>. </span>'
+                        )
+                    with gr.Column():
+                        gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "condition", "upper", _)
+                                for _ in os.listdir(os.path.join(root_path, "condition", "upper"))
+                            ],
+                            examples_per_page=4,
+                            inputs=cloth_image_flux,
+                            label="Condition Upper Examples",
+                        )
+                        gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "condition", "overall", _)
+                                for _ in os.listdir(os.path.join(root_path, "condition", "overall"))
+                            ],
+                            examples_per_page=4,
+                            inputs=cloth_image_flux,
+                            label="Condition Overall Examples",
+                        )
+                        condition_person_exm = gr.Examples(
+                            examples=[
+                                os.path.join(root_path, "condition", "person", _)
+                                for _ in os.listdir(os.path.join(root_path, "condition", "person"))
+                            ],
+                            examples_per_page=4,
+                            inputs=cloth_image_flux,
+                            label="Condition Reference Person Examples",
+                        )
+                        gr.Markdown(
+                            '<span style="color: #808080; font-size: small;">*Condition examples come from the Internet. </span>'
+                        )
+
+                
+            image_path_flux.change(
+                person_example_fn, inputs=image_path_flux, outputs=person_image_flux
+            )
+
+            submit_flux.click(
+                submit_function_flux,
+                [person_image_flux, cloth_image_flux, cloth_type, num_inference_steps_flux, guidance_scale_flux, seed_flux, show_type],
+                result_image_flux,
+            )
+        
+    
+    demo.queue().launch(share=True, show_error=True)
+
+# 解析参数
+args = parse_args()
+
+# 加载模型
+repo_path = snapshot_download(repo_id=args.resume_path)
+pipeline_flux = FluxTryOnPipeline.from_pretrained(args.base_model_path)
+pipeline_flux.load_lora_weights(
+    os.path.join(repo_path, "flux-lora"), 
+    weight_name='pytorch_lora_weights.safetensors'
+)
+pipeline_flux.to("cuda", torch.bfloat16)
+
+# 初始化 AutoMasker
+mask_processor = VaeImageProcessor(
+    vae_scale_factor=8, 
+    do_normalize=False, 
+    do_binarize=True, 
+    do_convert_grayscale=True
+)
+automasker = AutoMasker(
+    densepose_ckpt=os.path.join(repo_path, "DensePose"),
+    schp_ckpt=os.path.join(repo_path, "SCHP"),
+    device='cuda'
+)
+
+if __name__ == "__main__":
+    app_gradio()

app_p2p.py

@@ -0,0 +1,567 @@
+import argparse
+import os
+from datetime import datetime
+
+import gradio as gr
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+from model.cloth_masker import AutoMasker, vis_mask
+from model.pipeline import CatVTONPipeline, CatVTONPix2PixPipeline
+from utils import init_weight_dtype, resize_and_crop, resize_and_padding
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--p2p_base_model_path",
+        type=str,
+        default="timbrooks/instruct-pix2pix", 
+        help=(
+            "The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub."
+        ),
+    )
+    parser.add_argument(
+        "--ip_base_model_path",
+        type=str,
+        default="booksforcharlie/stable-diffusion-inpainting", 
+        help=(
+            "The path to the base model to use for evaluation. This can be a local path or a model identifier from the Model Hub."
+        ),
+    )
+    parser.add_argument(
+        "--p2p_resume_path",
+        type=str,
+        default="zhengchong/CatVTON-MaskFree",
+        help=(
+            "The Path to the checkpoint of trained tryon model."
+        ),
+    )
+    parser.add_argument(
+        "--ip_resume_path",
+        type=str,
+        default="zhengchong/CatVTON",
+        help=(
+            "The Path to the checkpoint of trained tryon model."
+        ),
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="resource/demo/output",
+        help="The output directory where the model predictions will be written.",
+    )
+
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=768,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=1024,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--repaint", 
+        action="store_true", 
+        help="Whether to repaint the result image with the original background."
+    )
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        default=True,
+        help=(
+            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="bf16",
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )
+    
+    args = parser.parse_args()
+    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
+    if env_local_rank != -1 and env_local_rank != args.local_rank:
+        args.local_rank = env_local_rank
+
+    return args
+
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows * cols
+
+    w, h = imgs[0].size
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+
+args = parse_args()
+repo_path = snapshot_download(repo_id=args.ip_resume_path)
+# Pipeline
+pipeline_p2p = CatVTONPix2PixPipeline(
+    base_ckpt=args.p2p_base_model_path,
+    attn_ckpt=repo_path,
+    attn_ckpt_version="mix-48k-1024",
+    weight_dtype=init_weight_dtype(args.mixed_precision),
+    use_tf32=args.allow_tf32,
+    device='cuda'
+)
+
+# Pipeline
+repo_path = snapshot_download(repo_id=args.ip_resume_path)  
+pipeline = CatVTONPipeline(
+    base_ckpt=args.ip_base_model_path,
+    attn_ckpt=repo_path,
+    attn_ckpt_version="mix",
+    weight_dtype=init_weight_dtype(args.mixed_precision),
+    use_tf32=args.allow_tf32,
+    device='cuda'
+)
+
+# AutoMasker
+mask_processor = VaeImageProcessor(vae_scale_factor=8, do_normalize=False, do_binarize=True, do_convert_grayscale=True)
+automasker = AutoMasker(
+    densepose_ckpt=os.path.join(repo_path, "DensePose"),
+    schp_ckpt=os.path.join(repo_path, "SCHP"),
+    device='cuda', 
+)
+
+
+def submit_function_p2p(
+    person_image,
+    cloth_image,
+    num_inference_steps,
+    guidance_scale,
+    seed):
+    person_image= person_image["background"]
+
+    tmp_folder = args.output_dir
+    date_str = datetime.now().strftime("%Y%m%d%H%M%S")
+    result_save_path = os.path.join(tmp_folder, date_str[:8], date_str[8:] + ".png")
+    if not os.path.exists(os.path.join(tmp_folder, date_str[:8])):
+        os.makedirs(os.path.join(tmp_folder, date_str[:8]))
+
+    generator = None
+    if seed != -1:
+        generator = torch.Generator(device='cuda').manual_seed(seed)
+
+    person_image = Image.open(person_image).convert("RGB")
+    cloth_image = Image.open(cloth_image).convert("RGB")
+    person_image = resize_and_crop(person_image, (args.width, args.height))
+    cloth_image = resize_and_padding(cloth_image, (args.width, args.height))
+
+    # Inference
+    try:
+        result_image = pipeline_p2p(
+            image=person_image,
+            condition_image=cloth_image,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            generator=generator
+        )[0]
+    except Exception as e:
+        raise gr.Error(
+            "An error occurred. Please try again later: {}".format(e)
+        )
+    
+    # Post-process
+    save_result_image = image_grid([person_image, cloth_image, result_image], 1, 3)
+    save_result_image.save(result_save_path)
+    return result_image
+
+def submit_function(
+    person_image,
+    cloth_image,
+    cloth_type,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    show_type
+):
+    person_image, mask = person_image["background"], person_image["layers"][0]
+    mask = Image.open(mask).convert("L")
+    if len(np.unique(np.array(mask))) == 1:
+        mask = None
+    else:
+        mask = np.array(mask)
+        mask[mask > 0] = 255
+        mask = Image.fromarray(mask)
+
+    tmp_folder = args.output_dir
+    date_str = datetime.now().strftime("%Y%m%d%H%M%S")
+    result_save_path = os.path.join(tmp_folder, date_str[:8], date_str[8:] + ".png")
+    if not os.path.exists(os.path.join(tmp_folder, date_str[:8])):
+        os.makedirs(os.path.join(tmp_folder, date_str[:8]))
+
+    generator = None
+    if seed != -1:
+        generator = torch.Generator(device='cuda').manual_seed(seed)
+
+    person_image = Image.open(person_image).convert("RGB")
+    cloth_image = Image.open(cloth_image).convert("RGB")
+    person_image = resize_and_crop(person_image, (args.width, args.height))
+    cloth_image = resize_and_padding(cloth_image, (args.width, args.height))
+    
+    # Process mask
+    if mask is not None:
+        mask = resize_and_crop(mask, (args.width, args.height))
+    else:
+        mask = automasker(
+            person_image,
+            cloth_type
+        )['mask']
+    mask = mask_processor.blur(mask, blur_factor=9)
+
+    # Inference
+    # try:
+    result_image = pipeline(
+        image=person_image,
+        condition_image=cloth_image,
+        mask=mask,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator
+    )[0]
+    # except Exception as e:
+    #     raise gr.Error(
+    #         "An error occurred. Please try again later: {}".format(e)
+    #     )
+    
+    # Post-process
+    masked_person = vis_mask(person_image, mask)
+    save_result_image = image_grid([person_image, masked_person, cloth_image, result_image], 1, 4)
+    save_result_image.save(result_save_path)
+    if show_type == "result only":
+        return result_image
+    else:
+        width, height = person_image.size
+        if show_type == "input & result":
+            condition_width = width // 2
+            conditions = image_grid([person_image, cloth_image], 2, 1)
+        else:
+            condition_width = width // 3
+            conditions = image_grid([person_image, masked_person , cloth_image], 3, 1)
+        conditions = conditions.resize((condition_width, height), Image.NEAREST)
+        new_result_image = Image.new("RGB", (width + condition_width + 5, height))
+        new_result_image.paste(conditions, (0, 0))
+        new_result_image.paste(result_image, (condition_width + 5, 0))
+    return new_result_image
+
+
+
+def person_example_fn(image_path):
+    return image_path
+
+HEADER = """
+<h1 style="text-align: center;"> 🐈 CatVTON: Concatenation Is All You Need for Virtual Try-On with Diffusion Models </h1>
+<div style="display: flex; justify-content: center; align-items: center;">
+  <a href="http://arxiv.org/abs/2407.15886" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/arXiv-2407.15886-red?style=flat&logo=arXiv&logoColor=red' alt='arxiv'>
+  </a>
+  <a href='https://huggingface.co/zhengchong/CatVTON' style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Hugging Face-ckpts-orange?style=flat&logo=HuggingFace&logoColor=orange' alt='huggingface'>
+  </a>
+  <a href="https://github.com/Zheng-Chong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/GitHub-Repo-blue?style=flat&logo=GitHub' alt='GitHub'>
+  </a>
+  <a href="http://120.76.142.206:8888" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Demo-Gradio-gold?style=flat&logo=Gradio&logoColor=red' alt='Demo'>
+  </a>
+  <a href="https://huggingface.co/spaces/zhengchong/CatVTON" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Space-ZeroGPU-orange?style=flat&logo=Gradio&logoColor=red' alt='Demo'>
+  </a>
+  <a href='https://zheng-chong.github.io/CatVTON/' style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/Webpage-Project-silver?style=flat&logo=&logoColor=orange' alt='webpage'>
+  </a>
+  <a href="https://github.com/Zheng-Chong/CatVTON/LICENCE" style="margin: 0 2px;">
+    <img src='https://img.shields.io/badge/License-CC BY--NC--SA--4.0-lightgreen?style=flat&logo=Lisence' alt='License'>
+  </a>
+</div>
+<br>
+· This demo and our weights are only for <span>Non-commercial Use</span>. <br>
+· You can try CatVTON in our <a href="https://huggingface.co/spaces/zhengchong/CatVTON">HuggingFace Space</a> or our <a href="http://120.76.142.206:8888">online demo</a> (run on 3090). <br>
+· Thanks to <a href="https://huggingface.co/zero-gpu-explorers">ZeroGPU</a> for providing A100 for our <a href="https://huggingface.co/spaces/zhengchong/CatVTON">HuggingFace Space</a>. <br>
+· SafetyChecker is set to filter NSFW content, but it may block normal results too. Please adjust the <span>`seed`</span> for normal outcomes.<br> 
+"""
+
+def app_gradio():
+    with gr.Blocks(title="CatVTON") as demo:
+        gr.Markdown(HEADER)
+        with gr.Tab("Mask-based Virtual Try-On"):
+            with gr.Row():
+                with gr.Column(scale=1, min_width=350):
+                    with gr.Row():
+                        image_path = gr.Image(
+                            type="filepath",
+                            interactive=True,
+                            visible=False,
+                        )
+                        person_image = gr.ImageEditor(
+                            interactive=True, label="Person Image", type="filepath"
+                        )
+
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=230):
+                            cloth_image = gr.Image(
+                                interactive=True, label="Condition Image", type="filepath"
+                            )
+                        with gr.Column(scale=1, min_width=120):
+                            gr.Markdown(
+                                '<span style="color: #808080; font-size: small;">Two ways to provide Mask:<br>1. Upload the person image and use the `🖌️` above to draw the Mask (higher priority)<br>2. Select the `Try-On Cloth Type` to generate automatically </span>'
+                            )
+                            cloth_type = gr.Radio(
+                                label="Try-On Cloth Type",
+                                choices=["upper", "lower", "overall"],
+                                value="upper",
+                            )
+
+
+                    submit = gr.Button("Submit")
+                    gr.Markdown(
+                        '<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>'
+                    )
+                    
+                    gr.Markdown(
+                        '<span style="color: #808080; font-size: small;">Advanced options can adjust details:<br>1. `Inference Step` may enhance details;<br>2. `CFG` is highly correlated with saturation;<br>3. `Random seed` may improve pseudo-shadow.</span>'
+                    )
+                    with gr.Accordion("Advanced Options", open=False):
+                        num_inference_steps = gr.Slider(
+                            label="Inference Step", minimum=10, maximum=100, step=5, value=50
+                        )
+                        # Guidence Scale
+                        guidance_scale = gr.Slider(
+                            label="CFG Strenth", minimum=0.0, maximum=7.5, step=0.5, value=2.5
+                        )
+                        # Random Seed
+                        seed = gr.Slider(
+                            label="Seed", minimum=-1, maximum=10000, step=1, value=42
+                        )
+                        show_type = gr.Radio(
+                            label="Show Type",
+                            choices=["result only", "input & result", "input & mask & result"],
+                            value="input & mask & result",
+                        )
+
+                with gr.Column(scale=2, min_width=500):
+                    result_image = gr.Image(interactive=False, label="Result")
+                    with gr.Row():
+                        # Photo Examples
+                        root_path = "resource/demo/example"
+                        with gr.Column():
+                            men_exm = gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "person", "men", _)
+                                    for _ in os.listdir(os.path.join(root_path, "person", "men"))
+                                ],
+                                examples_per_page=4,
+                                inputs=image_path,
+                                label="Person Examples ①",
+                            )
+                            women_exm = gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "person", "women", _)
+                                    for _ in os.listdir(os.path.join(root_path, "person", "women"))
+                                ],
+                                examples_per_page=4,
+                                inputs=image_path,
+                                label="Person Examples ②",
+                            )
+                            gr.Markdown(
+                                '<span style="color: #808080; font-size: small;">*Person examples come from the demos of <a href="https://huggingface.co/spaces/levihsu/OOTDiffusion">OOTDiffusion</a> and <a href="https://www.outfitanyone.org">OutfitAnyone</a>. </span>'
+                            )
+                        with gr.Column():
+                            condition_upper_exm = gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "condition", "upper", _)
+                                    for _ in os.listdir(os.path.join(root_path, "condition", "upper"))
+                                ],
+                                examples_per_page=4,
+                                inputs=cloth_image,
+                                label="Condition Upper Examples",
+                            )
+                            condition_overall_exm = gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "condition", "overall", _)
+                                    for _ in os.listdir(os.path.join(root_path, "condition", "overall"))
+                                ],
+                                examples_per_page=4,
+                                inputs=cloth_image,
+                                label="Condition Overall Examples",
+                            )
+                            condition_person_exm = gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "condition", "person", _)
+                                    for _ in os.listdir(os.path.join(root_path, "condition", "person"))
+                                ],
+                                examples_per_page=4,
+                                inputs=cloth_image,
+                                label="Condition Reference Person Examples",
+                            )
+                            gr.Markdown(
+                                '<span style="color: #808080; font-size: small;">*Condition examples come from the Internet. </span>'
+                            )
+
+                image_path.change(
+                    person_example_fn, inputs=image_path, outputs=person_image
+                )
+
+                submit.click(
+                    submit_function,
+                    [
+                        person_image,
+                        cloth_image,
+                        cloth_type,
+                        num_inference_steps,
+                        guidance_scale,
+                        seed,
+                        show_type,
+                    ],
+                    result_image,
+                )
+    
+        with gr.Tab("Mask-Free Virtual Try-On"):
+            with gr.Row():
+                with gr.Column(scale=1, min_width=350):
+                    with gr.Row():
+                        image_path_p2p = gr.Image(
+                            type="filepath",
+                            interactive=True,
+                            visible=False,
+                        )
+                        person_image_p2p = gr.ImageEditor(
+                            interactive=True, label="Person Image", type="filepath"
+                        )
+
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=230):
+                            cloth_image_p2p = gr.Image(
+                                interactive=True, label="Condition Image", type="filepath"
+                            )
+
+                    submit_p2p = gr.Button("Submit")
+                    gr.Markdown(
+                        '<center><span style="color: #FF0000">!!! Click only Once, Wait for Delay !!!</span></center>'
+                    )
+                    
+                    gr.Markdown(
+                        '<span style="color: #808080; font-size: small;">Advanced options can adjust details:<br>1. `Inference Step` may enhance details;<br>2. `CFG` is highly correlated with saturation;<br>3. `Random seed` may improve pseudo-shadow.</span>'
+                    )
+                    with gr.Accordion("Advanced Options", open=False):
+                        num_inference_steps_p2p = gr.Slider(
+                            label="Inference Step", minimum=10, maximum=100, step=5, value=50
+                        )
+                        # Guidence Scale
+                        guidance_scale_p2p = gr.Slider(
+                            label="CFG Strenth", minimum=0.0, maximum=7.5, step=0.5, value=2.5
+                        )
+                        # Random Seed
+                        seed_p2p = gr.Slider(
+                            label="Seed", minimum=-1, maximum=10000, step=1, value=42
+                        )
+                        # show_type = gr.Radio(
+                        #     label="Show Type",
+                        #     choices=["result only", "input & result", "input & mask & result"],
+                        #     value="input & mask & result",
+                        # )
+
+                with gr.Column(scale=2, min_width=500):
+                    result_image_p2p = gr.Image(interactive=False, label="Result")
+                    with gr.Row():
+                        # Photo Examples
+                        root_path = "resource/demo/example"
+                        with gr.Column():
+                            gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "person", "men", _)
+                                    for _ in os.listdir(os.path.join(root_path, "person", "men"))
+                                ],
+                                examples_per_page=4,
+                                inputs=image_path_p2p,
+                                label="Person Examples ①",
+                            )
+                            gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "person", "women", _)
+                                    for _ in os.listdir(os.path.join(root_path, "person", "women"))
+                                ],
+                                examples_per_page=4,
+                                inputs=image_path_p2p,
+                                label="Person Examples ②",
+                            )
+                            gr.Markdown(
+                                '<span style="color: #808080; font-size: small;">*Person examples come from the demos of <a href="https://huggingface.co/spaces/levihsu/OOTDiffusion">OOTDiffusion</a> and <a href="https://www.outfitanyone.org">OutfitAnyone</a>. </span>'
+                            )
+                        with gr.Column():
+                            gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "condition", "upper", _)
+                                    for _ in os.listdir(os.path.join(root_path, "condition", "upper"))
+                                ],
+                                examples_per_page=4,
+                                inputs=cloth_image_p2p,
+                                label="Condition Upper Examples",
+                            )
+                            gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "condition", "overall", _)
+                                    for _ in os.listdir(os.path.join(root_path, "condition", "overall"))
+                                ],
+                                examples_per_page=4,
+                                inputs=cloth_image_p2p,
+                                label="Condition Overall Examples",
+                            )
+                            condition_person_exm = gr.Examples(
+                                examples=[
+                                    os.path.join(root_path, "condition", "person", _)
+                                    for _ in os.listdir(os.path.join(root_path, "condition", "person"))
+                                ],
+                                examples_per_page=4,
+                                inputs=cloth_image_p2p,
+                                label="Condition Reference Person Examples",
+                            )
+                            gr.Markdown(
+                                '<span style="color: #808080; font-size: small;">*Condition examples come from the Internet. </span>'
+                            )
+
+                image_path_p2p.change(
+                    person_example_fn, inputs=image_path_p2p, outputs=person_image_p2p
+                )
+
+                submit_p2p.click(
+                    submit_function_p2p,
+                    [
+                        person_image_p2p,
+                        cloth_image_p2p,
+                        num_inference_steps_p2p,
+                        guidance_scale_p2p,
+                        seed_p2p],
+                    result_image_p2p,
+                )
+        
+    demo.queue().launch(share=True, show_error=True)
+
+
+if __name__ == "__main__":
+    app_gradio()

densepose/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+from .data.datasets import builtin  # just to register data
+from .converters import builtin as builtin_converters  # register converters
+from .config import (
+    add_densepose_config,
+    add_densepose_head_config,
+    add_hrnet_config,
+    add_dataset_category_config,
+    add_bootstrap_config,
+    load_bootstrap_config,
+)
+from .structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
+from .evaluation import DensePoseCOCOEvaluator
+from .modeling.roi_heads import DensePoseROIHeads
+from .modeling.test_time_augmentation import (
+    DensePoseGeneralizedRCNNWithTTA,
+    DensePoseDatasetMapperTTA,
+)
+from .utils.transform import load_from_cfg
+from .modeling.hrfpn import build_hrfpn_backbone

densepose/config.py

@@ -0,0 +1,277 @@
+# -*- coding = utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+# pyre-ignore-all-errors
+
+from detectron2.config import CfgNode as CN
+
+
+def add_dataset_category_config(cfg: CN) -> None:
+    """
+    Add config for additional category-related dataset options
+     - category whitelisting
+     - category mapping
+    """
+    _C = cfg
+    _C.DATASETS.CATEGORY_MAPS = CN(new_allowed=True)
+    _C.DATASETS.WHITELISTED_CATEGORIES = CN(new_allowed=True)
+    # class to mesh mapping
+    _C.DATASETS.CLASS_TO_MESH_NAME_MAPPING = CN(new_allowed=True)
+
+
+def add_evaluation_config(cfg: CN) -> None:
+    _C = cfg
+    _C.DENSEPOSE_EVALUATION = CN()
+    # evaluator type, possible values:
+    #  - "iou": evaluator for models that produce iou data
+    #  - "cse": evaluator for models that produce cse data
+    _C.DENSEPOSE_EVALUATION.TYPE = "iou"
+    # storage for DensePose results, possible values:
+    #  - "none": no explicit storage, all the results are stored in the
+    #            dictionary with predictions, memory intensive;
+    #            historically the default storage type
+    #  - "ram": RAM storage, uses per-process RAM storage, which is
+    #           reduced to a single process storage on later stages,
+    #           less memory intensive
+    #  - "file": file storage, uses per-process file-based storage,
+    #            the least memory intensive, but may create bottlenecks
+    #            on file system accesses
+    _C.DENSEPOSE_EVALUATION.STORAGE = "none"
+    # minimum threshold for IOU values: the lower its values is,
+    # the more matches are produced (and the higher the AP score)
+    _C.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD = 0.5
+    # Non-distributed inference is slower (at inference time) but can avoid RAM OOM
+    _C.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE = True
+    # evaluate mesh alignment based on vertex embeddings, only makes sense in CSE context
+    _C.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT = False
+    # meshes to compute mesh alignment for
+    _C.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES = []
+
+
+def add_bootstrap_config(cfg: CN) -> None:
+    """ """
+    _C = cfg
+    _C.BOOTSTRAP_DATASETS = []
+    _C.BOOTSTRAP_MODEL = CN()
+    _C.BOOTSTRAP_MODEL.WEIGHTS = ""
+    _C.BOOTSTRAP_MODEL.DEVICE = "cuda"
+
+
+def get_bootstrap_dataset_config() -> CN:
+    _C = CN()
+    _C.DATASET = ""
+    # ratio used to mix data loaders
+    _C.RATIO = 0.1
+    # image loader
+    _C.IMAGE_LOADER = CN(new_allowed=True)
+    _C.IMAGE_LOADER.TYPE = ""
+    _C.IMAGE_LOADER.BATCH_SIZE = 4
+    _C.IMAGE_LOADER.NUM_WORKERS = 4
+    _C.IMAGE_LOADER.CATEGORIES = []
+    _C.IMAGE_LOADER.MAX_COUNT_PER_CATEGORY = 1_000_000
+    _C.IMAGE_LOADER.CATEGORY_TO_CLASS_MAPPING = CN(new_allowed=True)
+    # inference
+    _C.INFERENCE = CN()
+    # batch size for model inputs
+    _C.INFERENCE.INPUT_BATCH_SIZE = 4
+    # batch size to group model outputs
+    _C.INFERENCE.OUTPUT_BATCH_SIZE = 2
+    # sampled data
+    _C.DATA_SAMPLER = CN(new_allowed=True)
+    _C.DATA_SAMPLER.TYPE = ""
+    _C.DATA_SAMPLER.USE_GROUND_TRUTH_CATEGORIES = False
+    # filter
+    _C.FILTER = CN(new_allowed=True)
+    _C.FILTER.TYPE = ""
+    return _C
+
+
+def load_bootstrap_config(cfg: CN) -> None:
+    """
+    Bootstrap datasets are given as a list of `dict` that are not automatically
+    converted into CfgNode. This method processes all bootstrap dataset entries
+    and ensures that they are in CfgNode format and comply with the specification
+    """
+    if not cfg.BOOTSTRAP_DATASETS:
+        return
+
+    bootstrap_datasets_cfgnodes = []
+    for dataset_cfg in cfg.BOOTSTRAP_DATASETS:
+        _C = get_bootstrap_dataset_config().clone()
+        _C.merge_from_other_cfg(CN(dataset_cfg))
+        bootstrap_datasets_cfgnodes.append(_C)
+    cfg.BOOTSTRAP_DATASETS = bootstrap_datasets_cfgnodes
+
+
+def add_densepose_head_cse_config(cfg: CN) -> None:
+    """
+    Add configuration options for Continuous Surface Embeddings (CSE)
+    """
+    _C = cfg
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE = CN()
+    # Dimensionality D of the embedding space
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE = 16
+    # Embedder specifications for various mesh IDs
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS = CN(new_allowed=True)
+    # normalization coefficient for embedding distances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_DIST_GAUSS_SIGMA = 0.01
+    # normalization coefficient for geodesic distances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.GEODESIC_DIST_GAUSS_SIGMA = 0.01
+    # embedding loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_WEIGHT = 0.6
+    # embedding loss name, currently the following options are supported:
+    # - EmbeddingLoss: cross-entropy on vertex labels
+    # - SoftEmbeddingLoss: cross-entropy on vertex label combined with
+    #    Gaussian penalty on distance between vertices
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_NAME = "EmbeddingLoss"
+    # optimizer hyperparameters
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR = 1.0
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR = 1.0
+    # Shape to shape cycle consistency loss parameters:
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
+    # shape to shape cycle consistency loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.025
+    # norm type used for loss computation
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
+    # normalization term for embedding similarity matrices
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.TEMPERATURE = 0.05
+    # maximum number of vertices to include into shape to shape cycle loss
+    # if negative or zero, all vertices are considered
+    # if positive, random subset of vertices of given size is considered
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.MAX_NUM_VERTICES = 4936
+    # Pixel to shape cycle consistency loss parameters:
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
+    # pixel to shape cycle consistency loss weight
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.0001
+    # norm type used for loss computation
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
+    # map images to all meshes and back (if false, use only gt meshes from the batch)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.USE_ALL_MESHES_NOT_GT_ONLY = False
+    # Randomly select at most this number of pixels from every instance
+    # if negative or zero, all vertices are considered
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NUM_PIXELS_TO_SAMPLE = 100
+    # normalization factor for pixel to pixel distances (higher value = smoother distribution)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.PIXEL_SIGMA = 5.0
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_PIXEL_TO_VERTEX = 0.05
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_VERTEX_TO_PIXEL = 0.05
+
+
+def add_densepose_head_config(cfg: CN) -> None:
+    """
+    Add config for densepose head.
+    """
+    _C = cfg
+
+    _C.MODEL.DENSEPOSE_ON = True
+
+    _C.MODEL.ROI_DENSEPOSE_HEAD = CN()
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NAME = ""
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS = 8
+    # Number of parts used for point labels
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES = 24
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL = 4
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM = 512
+    _C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL = 3
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UP_SCALE = 2
+    _C.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE = 112
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE = "ROIAlignV2"
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION = 28
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO = 2
+    _C.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS = 2  # 15 or 2
+    # Overlap threshold for an RoI to be considered foreground (if >= FG_IOU_THRESHOLD)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.FG_IOU_THRESHOLD = 0.7
+    # Loss weights for annotation masks.(14 Parts)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.INDEX_WEIGHTS = 5.0
+    # Loss weights for surface parts. (24 Parts)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.PART_WEIGHTS = 1.0
+    # Loss weights for UV regression.
+    _C.MODEL.ROI_DENSEPOSE_HEAD.POINT_REGRESSION_WEIGHTS = 0.01
+    # Coarse segmentation is trained using instance segmentation task data
+    _C.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS = False
+    # For Decoder
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_ON = True
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NUM_CLASSES = 256
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_CONV_DIMS = 256
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NORM = ""
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_COMMON_STRIDE = 4
+    # For DeepLab head
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB = CN()
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NORM = "GN"
+    _C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NONLOCAL_ON = 0
+    # Predictor class name, must be registered in DENSEPOSE_PREDICTOR_REGISTRY
+    # Some registered predictors:
+    #   "DensePoseChartPredictor": predicts segmentation and UV coordinates for predefined charts
+    #   "DensePoseChartWithConfidencePredictor": predicts segmentation, UV coordinates
+    #       and associated confidences for predefined charts (default)
+    #   "DensePoseEmbeddingWithConfidencePredictor": predicts segmentation, embeddings
+    #       and associated confidences for CSE
+    _C.MODEL.ROI_DENSEPOSE_HEAD.PREDICTOR_NAME = "DensePoseChartWithConfidencePredictor"
+    # Loss class name, must be registered in DENSEPOSE_LOSS_REGISTRY
+    # Some registered losses:
+    #   "DensePoseChartLoss": loss for chart-based models that estimate
+    #      segmentation and UV coordinates
+    #   "DensePoseChartWithConfidenceLoss": loss for chart-based models that estimate
+    #      segmentation, UV coordinates and the corresponding confidences (default)
+    _C.MODEL.ROI_DENSEPOSE_HEAD.LOSS_NAME = "DensePoseChartWithConfidenceLoss"
+    # Confidences
+    # Enable learning UV confidences (variances) along with the actual values
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE = CN({"ENABLED": False})
+    # UV confidence lower bound
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.EPSILON = 0.01
+    # Enable learning segmentation confidences (variances) along with the actual values
+    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE = CN({"ENABLED": False})
+    # Segmentation confidence lower bound
+    _C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE.EPSILON = 0.01
+    # Statistical model type for confidence learning, possible values:
+    # - "iid_iso": statistically independent identically distributed residuals
+    #    with isotropic covariance
+    # - "indep_aniso": statistically independent residuals with anisotropic
+    #    covariances
+    _C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.TYPE = "iid_iso"
+    # List of angles for rotation in data augmentation during training
+    _C.INPUT.ROTATION_ANGLES = [0]
+    _C.TEST.AUG.ROTATION_ANGLES = ()  # Rotation TTA
+
+    add_densepose_head_cse_config(cfg)
+
+
+def add_hrnet_config(cfg: CN) -> None:
+    """
+    Add config for HRNet backbone.
+    """
+    _C = cfg
+
+    # For HigherHRNet w32
+    _C.MODEL.HRNET = CN()
+    _C.MODEL.HRNET.STEM_INPLANES = 64
+    _C.MODEL.HRNET.STAGE2 = CN()
+    _C.MODEL.HRNET.STAGE2.NUM_MODULES = 1
+    _C.MODEL.HRNET.STAGE2.NUM_BRANCHES = 2
+    _C.MODEL.HRNET.STAGE2.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE2.NUM_BLOCKS = [4, 4]
+    _C.MODEL.HRNET.STAGE2.NUM_CHANNELS = [32, 64]
+    _C.MODEL.HRNET.STAGE2.FUSE_METHOD = "SUM"
+    _C.MODEL.HRNET.STAGE3 = CN()
+    _C.MODEL.HRNET.STAGE3.NUM_MODULES = 4
+    _C.MODEL.HRNET.STAGE3.NUM_BRANCHES = 3
+    _C.MODEL.HRNET.STAGE3.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE3.NUM_BLOCKS = [4, 4, 4]
+    _C.MODEL.HRNET.STAGE3.NUM_CHANNELS = [32, 64, 128]
+    _C.MODEL.HRNET.STAGE3.FUSE_METHOD = "SUM"
+    _C.MODEL.HRNET.STAGE4 = CN()
+    _C.MODEL.HRNET.STAGE4.NUM_MODULES = 3
+    _C.MODEL.HRNET.STAGE4.NUM_BRANCHES = 4
+    _C.MODEL.HRNET.STAGE4.BLOCK = "BASIC"
+    _C.MODEL.HRNET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
+    _C.MODEL.HRNET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
+    _C.MODEL.HRNET.STAGE4.FUSE_METHOD = "SUM"
+
+    _C.MODEL.HRNET.HRFPN = CN()
+    _C.MODEL.HRNET.HRFPN.OUT_CHANNELS = 256
+
+
+def add_densepose_config(cfg: CN) -> None:
+    add_densepose_head_config(cfg)
+    add_hrnet_config(cfg)
+    add_bootstrap_config(cfg)
+    add_dataset_category_config(cfg)
+    add_evaluation_config(cfg)

densepose/converters/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .hflip import HFlipConverter
+from .to_mask import ToMaskConverter
+from .to_chart_result import ToChartResultConverter, ToChartResultConverterWithConfidences
+from .segm_to_mask import (
+    predictor_output_with_fine_and_coarse_segm_to_mask,
+    predictor_output_with_coarse_segm_to_mask,
+    resample_fine_and_coarse_segm_to_bbox,
+)
+from .chart_output_to_chart_result import (
+    densepose_chart_predictor_output_to_result,
+    densepose_chart_predictor_output_to_result_with_confidences,
+)
+from .chart_output_hflip import densepose_chart_predictor_output_hflip

densepose/converters/base.py

@@ -0,0 +1,95 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Tuple, Type
+import torch
+
+
+class BaseConverter:
+    """
+    Converter base class to be reused by various converters.
+    Converter allows one to convert data from various source types to a particular
+    destination type. Each source type needs to register its converter. The
+    registration for each source type is valid for all descendants of that type.
+    """
+
+    @classmethod
+    def register(cls, from_type: Type, converter: Any = None):
+        """
+        Registers a converter for the specified type.
+        Can be used as a decorator (if converter is None), or called as a method.
+
+        Args:
+            from_type (type): type to register the converter for;
+                all instances of this type will use the same converter
+            converter (callable): converter to be registered for the given
+                type; if None, this method is assumed to be a decorator for the converter
+        """
+
+        if converter is not None:
+            cls._do_register(from_type, converter)
+
+        def wrapper(converter: Any) -> Any:
+            cls._do_register(from_type, converter)
+            return converter
+
+        return wrapper
+
+    @classmethod
+    def _do_register(cls, from_type: Type, converter: Any):
+        cls.registry[from_type] = converter  # pyre-ignore[16]
+
+    @classmethod
+    def _lookup_converter(cls, from_type: Type) -> Any:
+        """
+        Perform recursive lookup for the given type
+        to find registered converter. If a converter was found for some base
+        class, it gets registered for this class to save on further lookups.
+
+        Args:
+            from_type: type for which to find a converter
+        Return:
+            callable or None - registered converter or None
+                if no suitable entry was found in the registry
+        """
+        if from_type in cls.registry:  # pyre-ignore[16]
+            return cls.registry[from_type]
+        for base in from_type.__bases__:
+            converter = cls._lookup_converter(base)
+            if converter is not None:
+                cls._do_register(from_type, converter)
+                return converter
+        return None
+
+    @classmethod
+    def convert(cls, instance: Any, *args, **kwargs):
+        """
+        Convert an instance to the destination type using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            instance: source instance to convert to the destination type
+        Return:
+            An instance of the destination type obtained from the source instance
+            Raises KeyError, if no suitable converter found
+        """
+        instance_type = type(instance)
+        converter = cls._lookup_converter(instance_type)
+        if converter is None:
+            if cls.dst_type is None:  # pyre-ignore[16]
+                output_type_str = "itself"
+            else:
+                output_type_str = cls.dst_type
+            raise KeyError(f"Could not find converter from {instance_type} to {output_type_str}")
+        return converter(instance, *args, **kwargs)
+
+
+IntTupleBox = Tuple[int, int, int, int]
+
+
+def make_int_box(box: torch.Tensor) -> IntTupleBox:
+    int_box = [0, 0, 0, 0]
+    int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
+    return int_box[0], int_box[1], int_box[2], int_box[3]

densepose/converters/builtin.py

@@ -0,0 +1,33 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from ..structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
+from . import (
+    HFlipConverter,
+    ToChartResultConverter,
+    ToChartResultConverterWithConfidences,
+    ToMaskConverter,
+    densepose_chart_predictor_output_hflip,
+    densepose_chart_predictor_output_to_result,
+    densepose_chart_predictor_output_to_result_with_confidences,
+    predictor_output_with_coarse_segm_to_mask,
+    predictor_output_with_fine_and_coarse_segm_to_mask,
+)
+
+ToMaskConverter.register(
+    DensePoseChartPredictorOutput, predictor_output_with_fine_and_coarse_segm_to_mask
+)
+ToMaskConverter.register(
+    DensePoseEmbeddingPredictorOutput, predictor_output_with_coarse_segm_to_mask
+)
+
+ToChartResultConverter.register(
+    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result
+)
+
+ToChartResultConverterWithConfidences.register(
+    DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result_with_confidences
+)
+
+HFlipConverter.register(DensePoseChartPredictorOutput, densepose_chart_predictor_output_hflip)

densepose/converters/chart_output_hflip.py

@@ -0,0 +1,73 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+from dataclasses import fields
+import torch
+
+from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
+
+
+def densepose_chart_predictor_output_hflip(
+    densepose_predictor_output: DensePoseChartPredictorOutput,
+    transform_data: DensePoseTransformData,
+) -> DensePoseChartPredictorOutput:
+    """
+    Change  to take into account a Horizontal flip.
+    """
+    if len(densepose_predictor_output) > 0:
+
+        PredictorOutput = type(densepose_predictor_output)
+        output_dict = {}
+
+        for field in fields(densepose_predictor_output):
+            field_value = getattr(densepose_predictor_output, field.name)
+            # flip tensors
+            if isinstance(field_value, torch.Tensor):
+                setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
+
+        densepose_predictor_output = _flip_iuv_semantics_tensor(
+            densepose_predictor_output, transform_data
+        )
+        densepose_predictor_output = _flip_segm_semantics_tensor(
+            densepose_predictor_output, transform_data
+        )
+
+        for field in fields(densepose_predictor_output):
+            output_dict[field.name] = getattr(densepose_predictor_output, field.name)
+
+        return PredictorOutput(**output_dict)
+    else:
+        return densepose_predictor_output
+
+
+def _flip_iuv_semantics_tensor(
+    densepose_predictor_output: DensePoseChartPredictorOutput,
+    dp_transform_data: DensePoseTransformData,
+) -> DensePoseChartPredictorOutput:
+    point_label_symmetries = dp_transform_data.point_label_symmetries
+    uv_symmetries = dp_transform_data.uv_symmetries
+
+    N, C, H, W = densepose_predictor_output.u.shape
+    u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
+    v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
+    Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
+        None, :, None, None
+    ].expand(N, C - 1, H, W)
+    densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
+    densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
+
+    for el in ["fine_segm", "u", "v"]:
+        densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
+            :, point_label_symmetries, :, :
+        ]
+    return densepose_predictor_output
+
+
+def _flip_segm_semantics_tensor(
+    densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
+):
+    if densepose_predictor_output.coarse_segm.shape[1] > 2:
+        densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
+            :, dp_transform_data.mask_label_symmetries, :, :
+        ]
+    return densepose_predictor_output

densepose/converters/chart_output_to_chart_result.py

@@ -0,0 +1,190 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Dict
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures.boxes import Boxes, BoxMode
+
+from ..structures import (
+    DensePoseChartPredictorOutput,
+    DensePoseChartResult,
+    DensePoseChartResultWithConfidences,
+)
+from . import resample_fine_and_coarse_segm_to_bbox
+from .base import IntTupleBox, make_int_box
+
+
+def resample_uv_tensors_to_bbox(
+    u: torch.Tensor,
+    v: torch.Tensor,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        u (tensor [1, C, H, W] of float): U coordinates
+        v (tensor [1, C, H, W] of float): V coordinates
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
+    v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
+    uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
+    for part_id in range(1, u_bbox.size(1)):
+        uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
+        uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
+    return uv
+
+
+def resample_uv_to_bbox(
+    predictor_output: DensePoseChartPredictorOutput,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be resampled
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    return resample_uv_tensors_to_bbox(
+        predictor_output.u,
+        predictor_output.v,
+        labels,
+        box_xywh_abs,
+    )
+
+
+def densepose_chart_predictor_output_to_result(
+    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
+) -> DensePoseChartResult:
+    """
+    Convert densepose chart predictor outputs to results
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be converted to results, must contain only 1 output
+        boxes (Boxes): bounding box that corresponds to the predictor output,
+            must contain only 1 bounding box
+    Return:
+       DensePose chart-based result (DensePoseChartResult)
+    """
+    assert len(predictor_output) == 1 and len(boxes) == 1, (
+        f"Predictor output to result conversion can operate only single outputs"
+        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
+    )
+
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
+    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
+    return DensePoseChartResult(labels=labels, uv=uv)
+
+
+def resample_confidences_to_bbox(
+    predictor_output: DensePoseChartPredictorOutput,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> Dict[str, torch.Tensor]:
+    """
+    Resamples confidences for the given bounding box
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output to be resampled
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled confidences - a dict of [H, W] tensors of float
+    """
+
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+
+    confidence_names = [
+        "sigma_1",
+        "sigma_2",
+        "kappa_u",
+        "kappa_v",
+        "fine_segm_confidence",
+        "coarse_segm_confidence",
+    ]
+    confidence_results = {key: None for key in confidence_names}
+    confidence_names = [
+        key for key in confidence_names if getattr(predictor_output, key) is not None
+    ]
+    confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device)
+
+    # assign data from channels that correspond to the labels
+    for key in confidence_names:
+        resampled_confidence = F.interpolate(
+            getattr(predictor_output, key),
+            (h, w),
+            mode="bilinear",
+            align_corners=False,
+        )
+        result = confidence_base.clone()
+        for part_id in range(1, predictor_output.u.size(1)):
+            if resampled_confidence.size(1) != predictor_output.u.size(1):
+                # confidence is not part-based, don't try to fill it part by part
+                continue
+            result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id]
+
+        if resampled_confidence.size(1) != predictor_output.u.size(1):
+            # confidence is not part-based, fill the data with the first channel
+            # (targeted for segmentation confidences that have only 1 channel)
+            result = resampled_confidence[0, 0]
+
+        confidence_results[key] = result
+
+    return confidence_results  # pyre-ignore[7]
+
+
+def densepose_chart_predictor_output_to_result_with_confidences(
+    predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
+) -> DensePoseChartResultWithConfidences:
+    """
+    Convert densepose chart predictor outputs to results
+
+    Args:
+        predictor_output (DensePoseChartPredictorOutput): DensePose predictor
+            output with confidences to be converted to results, must contain only 1 output
+        boxes (Boxes): bounding box that corresponds to the predictor output,
+            must contain only 1 bounding box
+    Return:
+       DensePose chart-based result with confidences (DensePoseChartResultWithConfidences)
+    """
+    assert len(predictor_output) == 1 and len(boxes) == 1, (
+        f"Predictor output to result conversion can operate only single outputs"
+        f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
+    )
+
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
+    uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
+    confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh)
+    return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences)

densepose/converters/hflip.py

@@ -0,0 +1,36 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any
+
+from .base import BaseConverter
+
+
+class HFlipConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = None
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(cls, predictor_outputs: Any, transform_data: Any, *args, **kwargs):
+        """
+        Performs an horizontal flip on DensePose predictor outputs.
+        Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            predictor_outputs: DensePose predictor output to be converted to BitMasks
+            transform_data: Anything useful for the flip
+        Return:
+            An instance of the same type as predictor_outputs
+        """
+        return super(HFlipConverter, cls).convert(
+            predictor_outputs, transform_data, *args, **kwargs
+        )

densepose/converters/segm_to_mask.py

@@ -0,0 +1,152 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures import BitMasks, Boxes, BoxMode
+
+from .base import IntTupleBox, make_int_box
+from .to_mask import ImageSizeType
+
+
+def resample_coarse_segm_tensor_to_bbox(coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox):
+    """
+    Resample coarse segmentation tensor to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    labels = F.interpolate(coarse_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+    return labels
+
+
+def resample_fine_and_coarse_segm_tensors_to_bbox(
+    fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
+):
+    """
+    Resample fine and coarse segmentation tensors to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        fine_segm: float tensor of shape [1, C, Hout, Wout]
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    # coarse segmentation
+    coarse_segm_bbox = F.interpolate(
+        coarse_segm,
+        (h, w),
+        mode="bilinear",
+        align_corners=False,
+    ).argmax(dim=1)
+    # combined coarse and fine segmentation
+    labels = (
+        F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+        * (coarse_segm_bbox > 0).long()
+    )
+    return labels
+
+
+def resample_fine_and_coarse_segm_to_bbox(predictor_output: Any, box_xywh_abs: IntTupleBox):
+    """
+    Resample fine and coarse segmentation outputs from a predictor to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        predictor_output: DensePose predictor output that contains segmentation
+            results to be resampled
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    return resample_fine_and_coarse_segm_tensors_to_bbox(
+        predictor_output.fine_segm,
+        predictor_output.coarse_segm,
+        box_xywh_abs,
+    )
+
+
+def predictor_output_with_coarse_segm_to_mask(
+    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
+) -> BitMasks:
+    """
+    Convert predictor output with coarse and fine segmentation to a mask.
+    Assumes that predictor output has the following attributes:
+     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
+         unnormalized scores for N instances; D is the number of coarse
+         segmentation labels, H and W is the resolution of the estimate
+
+    Args:
+        predictor_output: DensePose predictor output to be converted to mask
+        boxes (Boxes): bounding boxes that correspond to the DensePose
+            predictor outputs
+        image_size_hw (tuple [int, int]): image height Himg and width Wimg
+    Return:
+        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
+        a mask of the size of the image for each instance
+    """
+    H, W = image_size_hw
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    N = len(boxes_xywh_abs)
+    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
+    for i in range(len(boxes_xywh_abs)):
+        box_xywh = make_int_box(boxes_xywh_abs[i])
+        box_mask = resample_coarse_segm_tensor_to_bbox(predictor_output[i].coarse_segm, box_xywh)
+        x, y, w, h = box_xywh
+        masks[i, y : y + h, x : x + w] = box_mask
+
+    return BitMasks(masks)
+
+
+def predictor_output_with_fine_and_coarse_segm_to_mask(
+    predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
+) -> BitMasks:
+    """
+    Convert predictor output with coarse and fine segmentation to a mask.
+    Assumes that predictor output has the following attributes:
+     - coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
+         unnormalized scores for N instances; D is the number of coarse
+         segmentation labels, H and W is the resolution of the estimate
+     - fine_segm (tensor of size [N, C, H, W]): fine segmentation
+         unnormalized scores for N instances; C is the number of fine
+         segmentation labels, H and W is the resolution of the estimate
+
+    Args:
+        predictor_output: DensePose predictor output to be converted to mask
+        boxes (Boxes): bounding boxes that correspond to the DensePose
+            predictor outputs
+        image_size_hw (tuple [int, int]): image height Himg and width Wimg
+    Return:
+        BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
+        a mask of the size of the image for each instance
+    """
+    H, W = image_size_hw
+    boxes_xyxy_abs = boxes.tensor.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    N = len(boxes_xywh_abs)
+    masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
+    for i in range(len(boxes_xywh_abs)):
+        box_xywh = make_int_box(boxes_xywh_abs[i])
+        labels_i = resample_fine_and_coarse_segm_to_bbox(predictor_output[i], box_xywh)
+        x, y, w, h = box_xywh
+        masks[i, y : y + h, x : x + w] = labels_i > 0
+    return BitMasks(masks)

densepose/converters/to_chart_result.py

@@ -0,0 +1,72 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any
+
+from detectron2.structures import Boxes
+
+from ..structures import DensePoseChartResult, DensePoseChartResultWithConfidences
+from .base import BaseConverter
+
+
+class ToChartResultConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = DensePoseChartResult
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs) -> DensePoseChartResult:
+        """
+        Convert DensePose predictor outputs to DensePoseResult using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output to be
+                converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+        Return:
+            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
+        """
+        return super(ToChartResultConverter, cls).convert(predictor_outputs, boxes, *args, **kwargs)
+
+
+class ToChartResultConverterWithConfidences(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to DensePose results.
+    Each DensePose predictor output type has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = DensePoseChartResultWithConfidences
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(
+        cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs
+    ) -> DensePoseChartResultWithConfidences:
+        """
+        Convert DensePose predictor outputs to DensePoseResult with confidences
+        using some registered converter. Does recursive lookup for base classes,
+        so there's no need for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output with confidences
+                to be converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+        Return:
+            An instance of DensePoseResult. If no suitable converter was found, raises KeyError
+        """
+        return super(ToChartResultConverterWithConfidences, cls).convert(
+            predictor_outputs, boxes, *args, **kwargs
+        )

densepose/converters/to_mask.py

@@ -0,0 +1,51 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Tuple
+
+from detectron2.structures import BitMasks, Boxes
+
+from .base import BaseConverter
+
+ImageSizeType = Tuple[int, int]
+
+
+class ToMaskConverter(BaseConverter):
+    """
+    Converts various DensePose predictor outputs to masks
+    in bit mask format (see `BitMasks`). Each DensePose predictor output type
+    has to register its convertion strategy.
+    """
+
+    registry = {}
+    dst_type = BitMasks
+
+    @classmethod
+    # pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
+    #  inconsistently.
+    def convert(
+        cls,
+        densepose_predictor_outputs: Any,
+        boxes: Boxes,
+        image_size_hw: ImageSizeType,
+        *args,
+        **kwargs
+    ) -> BitMasks:
+        """
+        Convert DensePose predictor outputs to BitMasks using some registered
+        converter. Does recursive lookup for base classes, so there's no need
+        for explicit registration for derived classes.
+
+        Args:
+            densepose_predictor_outputs: DensePose predictor output to be
+                converted to BitMasks
+            boxes (Boxes): bounding boxes that correspond to the DensePose
+                predictor outputs
+            image_size_hw (tuple [int, int]): image height and width
+        Return:
+            An instance of `BitMasks`. If no suitable converter was found, raises KeyError
+        """
+        return super(ToMaskConverter, cls).convert(
+            densepose_predictor_outputs, boxes, image_size_hw, *args, **kwargs
+        )

densepose/data/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .meshes import builtin
+from .build import (
+    build_detection_test_loader,
+    build_detection_train_loader,
+    build_combined_loader,
+    build_frame_selector,
+    build_inference_based_loaders,
+    has_inference_based_loaders,
+    BootstrapDatasetFactoryCatalog,
+)
+from .combined_loader import CombinedDataLoader
+from .dataset_mapper import DatasetMapper
+from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
+from .image_list_dataset import ImageListDataset
+from .utils import is_relative_local_path, maybe_prepend_base_path
+
+# ensure the builtin datasets are registered
+from . import datasets
+
+# ensure the bootstrap datasets builders are registered
+from . import build
+
+__all__ = [k for k in globals().keys() if not k.startswith("_")]

densepose/data/build.py

@@ -0,0 +1,738 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import itertools
+import logging
+import numpy as np
+from collections import UserDict, defaultdict
+from dataclasses import dataclass
+from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Sequence, Tuple
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.config import CfgNode
+from detectron2.data.build import build_detection_test_loader as d2_build_detection_test_loader
+from detectron2.data.build import build_detection_train_loader as d2_build_detection_train_loader
+from detectron2.data.build import (
+    load_proposals_into_dataset,
+    print_instances_class_histogram,
+    trivial_batch_collator,
+    worker_init_reset_seed,
+)
+from detectron2.data.catalog import DatasetCatalog, Metadata, MetadataCatalog
+from detectron2.data.samplers import TrainingSampler
+from detectron2.utils.comm import get_world_size
+
+from densepose.config import get_bootstrap_dataset_config
+from densepose.modeling import build_densepose_embedder
+
+from .combined_loader import CombinedDataLoader, Loader
+from .dataset_mapper import DatasetMapper
+from .datasets.coco import DENSEPOSE_CSE_KEYS_WITHOUT_MASK, DENSEPOSE_IUV_KEYS_WITHOUT_MASK
+from .datasets.dataset_type import DatasetType
+from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
+from .samplers import (
+    DensePoseConfidenceBasedSampler,
+    DensePoseCSEConfidenceBasedSampler,
+    DensePoseCSEUniformSampler,
+    DensePoseUniformSampler,
+    MaskFromDensePoseSampler,
+    PredictionToGroundTruthSampler,
+)
+from .transform import ImageResizeTransform
+from .utils import get_category_to_class_mapping, get_class_to_mesh_name_mapping
+from .video import (
+    FirstKFramesSelector,
+    FrameSelectionStrategy,
+    LastKFramesSelector,
+    RandomKFramesSelector,
+    VideoKeyframeDataset,
+    video_list_from_file,
+)
+
+__all__ = ["build_detection_train_loader", "build_detection_test_loader"]
+
+
+Instance = Dict[str, Any]
+InstancePredicate = Callable[[Instance], bool]
+
+
+def _compute_num_images_per_worker(cfg: CfgNode) -> int:
+    num_workers = get_world_size()
+    images_per_batch = cfg.SOLVER.IMS_PER_BATCH
+    assert (
+        images_per_batch % num_workers == 0
+    ), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format(
+        images_per_batch, num_workers
+    )
+    assert (
+        images_per_batch >= num_workers
+    ), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format(
+        images_per_batch, num_workers
+    )
+    images_per_worker = images_per_batch // num_workers
+    return images_per_worker
+
+
+def _map_category_id_to_contiguous_id(dataset_name: str, dataset_dicts: Iterable[Instance]) -> None:
+    meta = MetadataCatalog.get(dataset_name)
+    for dataset_dict in dataset_dicts:
+        for ann in dataset_dict["annotations"]:
+            ann["category_id"] = meta.thing_dataset_id_to_contiguous_id[ann["category_id"]]
+
+
+@dataclass
+class _DatasetCategory:
+    """
+    Class representing category data in a dataset:
+     - id: category ID, as specified in the dataset annotations file
+     - name: category name, as specified in the dataset annotations file
+     - mapped_id: category ID after applying category maps (DATASETS.CATEGORY_MAPS config option)
+     - mapped_name: category name after applying category maps
+     - dataset_name: dataset in which the category is defined
+
+    For example, when training models in a class-agnostic manner, one could take LVIS 1.0
+    dataset and map the animal categories to the same category as human data from COCO:
+     id = 225
+     name = "cat"
+     mapped_id = 1
+     mapped_name = "person"
+     dataset_name = "lvis_v1_animals_dp_train"
+    """
+
+    id: int
+    name: str
+    mapped_id: int
+    mapped_name: str
+    dataset_name: str
+
+
+_MergedCategoriesT = Dict[int, List[_DatasetCategory]]
+
+
+def _add_category_id_to_contiguous_id_maps_to_metadata(
+    merged_categories: _MergedCategoriesT,
+) -> None:
+    merged_categories_per_dataset = {}
+    for contiguous_cat_id, cat_id in enumerate(sorted(merged_categories.keys())):
+        for cat in merged_categories[cat_id]:
+            if cat.dataset_name not in merged_categories_per_dataset:
+                merged_categories_per_dataset[cat.dataset_name] = defaultdict(list)
+            merged_categories_per_dataset[cat.dataset_name][cat_id].append(
+                (
+                    contiguous_cat_id,
+                    cat,
+                )
+            )
+
+    logger = logging.getLogger(__name__)
+    for dataset_name, merged_categories in merged_categories_per_dataset.items():
+        meta = MetadataCatalog.get(dataset_name)
+        if not hasattr(meta, "thing_classes"):
+            meta.thing_classes = []
+            meta.thing_dataset_id_to_contiguous_id = {}
+            meta.thing_dataset_id_to_merged_id = {}
+        else:
+            meta.thing_classes.clear()
+            meta.thing_dataset_id_to_contiguous_id.clear()
+            meta.thing_dataset_id_to_merged_id.clear()
+        logger.info(f"Dataset {dataset_name}: category ID to contiguous ID mapping:")
+        for _cat_id, categories in sorted(merged_categories.items()):
+            added_to_thing_classes = False
+            for contiguous_cat_id, cat in categories:
+                if not added_to_thing_classes:
+                    meta.thing_classes.append(cat.mapped_name)
+                    added_to_thing_classes = True
+                meta.thing_dataset_id_to_contiguous_id[cat.id] = contiguous_cat_id
+                meta.thing_dataset_id_to_merged_id[cat.id] = cat.mapped_id
+                logger.info(f"{cat.id} ({cat.name}) -> {contiguous_cat_id}")
+
+
+def _maybe_create_general_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    def has_annotations(instance: Instance) -> bool:
+        return "annotations" in instance
+
+    def has_only_crowd_anotations(instance: Instance) -> bool:
+        for ann in instance["annotations"]:
+            if ann.get("is_crowd", 0) == 0:
+                return False
+        return True
+
+    def general_keep_instance_predicate(instance: Instance) -> bool:
+        return has_annotations(instance) and not has_only_crowd_anotations(instance)
+
+    if not cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS:
+        return None
+    return general_keep_instance_predicate
+
+
+def _maybe_create_keypoints_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+
+    min_num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
+
+    def has_sufficient_num_keypoints(instance: Instance) -> bool:
+        num_kpts = sum(
+            (np.array(ann["keypoints"][2::3]) > 0).sum()
+            for ann in instance["annotations"]
+            if "keypoints" in ann
+        )
+        return num_kpts >= min_num_keypoints
+
+    if cfg.MODEL.KEYPOINT_ON and (min_num_keypoints > 0):
+        return has_sufficient_num_keypoints
+    return None
+
+
+def _maybe_create_mask_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    if not cfg.MODEL.MASK_ON:
+        return None
+
+    def has_mask_annotations(instance: Instance) -> bool:
+        return any("segmentation" in ann for ann in instance["annotations"])
+
+    return has_mask_annotations
+
+
+def _maybe_create_densepose_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    if not cfg.MODEL.DENSEPOSE_ON:
+        return None
+
+    use_masks = cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
+
+    def has_densepose_annotations(instance: Instance) -> bool:
+        for ann in instance["annotations"]:
+            if all(key in ann for key in DENSEPOSE_IUV_KEYS_WITHOUT_MASK) or all(
+                key in ann for key in DENSEPOSE_CSE_KEYS_WITHOUT_MASK
+            ):
+                return True
+            if use_masks and "segmentation" in ann:
+                return True
+        return False
+
+    return has_densepose_annotations
+
+
+def _maybe_create_specific_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
+    specific_predicate_creators = [
+        _maybe_create_keypoints_keep_instance_predicate,
+        _maybe_create_mask_keep_instance_predicate,
+        _maybe_create_densepose_keep_instance_predicate,
+    ]
+    predicates = [creator(cfg) for creator in specific_predicate_creators]
+    predicates = [p for p in predicates if p is not None]
+    if not predicates:
+        return None
+
+    def combined_predicate(instance: Instance) -> bool:
+        return any(p(instance) for p in predicates)
+
+    return combined_predicate
+
+
+def _get_train_keep_instance_predicate(cfg: CfgNode):
+    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
+    combined_specific_keep_predicate = _maybe_create_specific_keep_instance_predicate(cfg)
+
+    def combined_general_specific_keep_predicate(instance: Instance) -> bool:
+        return general_keep_predicate(instance) and combined_specific_keep_predicate(instance)
+
+    if (general_keep_predicate is None) and (combined_specific_keep_predicate is None):
+        return None
+    if general_keep_predicate is None:
+        return combined_specific_keep_predicate
+    if combined_specific_keep_predicate is None:
+        return general_keep_predicate
+    return combined_general_specific_keep_predicate
+
+
+def _get_test_keep_instance_predicate(cfg: CfgNode):
+    general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
+    return general_keep_predicate
+
+
+def _maybe_filter_and_map_categories(
+    dataset_name: str, dataset_dicts: List[Instance]
+) -> List[Instance]:
+    meta = MetadataCatalog.get(dataset_name)
+    category_id_map = meta.thing_dataset_id_to_contiguous_id
+    filtered_dataset_dicts = []
+    for dataset_dict in dataset_dicts:
+        anns = []
+        for ann in dataset_dict["annotations"]:
+            cat_id = ann["category_id"]
+            if cat_id not in category_id_map:
+                continue
+            ann["category_id"] = category_id_map[cat_id]
+            anns.append(ann)
+        dataset_dict["annotations"] = anns
+        filtered_dataset_dicts.append(dataset_dict)
+    return filtered_dataset_dicts
+
+
+def _add_category_whitelists_to_metadata(cfg: CfgNode) -> None:
+    for dataset_name, whitelisted_cat_ids in cfg.DATASETS.WHITELISTED_CATEGORIES.items():
+        meta = MetadataCatalog.get(dataset_name)
+        meta.whitelisted_categories = whitelisted_cat_ids
+        logger = logging.getLogger(__name__)
+        logger.info(
+            "Whitelisted categories for dataset {}: {}".format(
+                dataset_name, meta.whitelisted_categories
+            )
+        )
+
+
+def _add_category_maps_to_metadata(cfg: CfgNode) -> None:
+    for dataset_name, category_map in cfg.DATASETS.CATEGORY_MAPS.items():
+        category_map = {
+            int(cat_id_src): int(cat_id_dst) for cat_id_src, cat_id_dst in category_map.items()
+        }
+        meta = MetadataCatalog.get(dataset_name)
+        meta.category_map = category_map
+        logger = logging.getLogger(__name__)
+        logger.info("Category maps for dataset {}: {}".format(dataset_name, meta.category_map))
+
+
+def _add_category_info_to_bootstrapping_metadata(dataset_name: str, dataset_cfg: CfgNode) -> None:
+    meta = MetadataCatalog.get(dataset_name)
+    meta.category_to_class_mapping = get_category_to_class_mapping(dataset_cfg)
+    meta.categories = dataset_cfg.CATEGORIES
+    meta.max_count_per_category = dataset_cfg.MAX_COUNT_PER_CATEGORY
+    logger = logging.getLogger(__name__)
+    logger.info(
+        "Category to class mapping for dataset {}: {}".format(
+            dataset_name, meta.category_to_class_mapping
+        )
+    )
+
+
+def _maybe_add_class_to_mesh_name_map_to_metadata(dataset_names: List[str], cfg: CfgNode) -> None:
+    for dataset_name in dataset_names:
+        meta = MetadataCatalog.get(dataset_name)
+        if not hasattr(meta, "class_to_mesh_name"):
+            meta.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
+
+
+def _merge_categories(dataset_names: Collection[str]) -> _MergedCategoriesT:
+    merged_categories = defaultdict(list)
+    category_names = {}
+    for dataset_name in dataset_names:
+        meta = MetadataCatalog.get(dataset_name)
+        whitelisted_categories = meta.get("whitelisted_categories")
+        category_map = meta.get("category_map", {})
+        cat_ids = (
+            whitelisted_categories if whitelisted_categories is not None else meta.categories.keys()
+        )
+        for cat_id in cat_ids:
+            cat_name = meta.categories[cat_id]
+            cat_id_mapped = category_map.get(cat_id, cat_id)
+            if cat_id_mapped == cat_id or cat_id_mapped in cat_ids:
+                category_names[cat_id] = cat_name
+            else:
+                category_names[cat_id] = str(cat_id_mapped)
+            # assign temporary mapped category name, this name can be changed
+            # during the second pass, since mapped ID can correspond to a category
+            # from a different dataset
+            cat_name_mapped = meta.categories[cat_id_mapped]
+            merged_categories[cat_id_mapped].append(
+                _DatasetCategory(
+                    id=cat_id,
+                    name=cat_name,
+                    mapped_id=cat_id_mapped,
+                    mapped_name=cat_name_mapped,
+                    dataset_name=dataset_name,
+                )
+            )
+    # second pass to assign proper mapped category names
+    for cat_id, categories in merged_categories.items():
+        for cat in categories:
+            if cat_id in category_names and cat.mapped_name != category_names[cat_id]:
+                cat.mapped_name = category_names[cat_id]
+
+    return merged_categories
+
+
+def _warn_if_merged_different_categories(merged_categories: _MergedCategoriesT) -> None:
+    logger = logging.getLogger(__name__)
+    for cat_id in merged_categories:
+        merged_categories_i = merged_categories[cat_id]
+        first_cat_name = merged_categories_i[0].name
+        if len(merged_categories_i) > 1 and not all(
+            cat.name == first_cat_name for cat in merged_categories_i[1:]
+        ):
+            cat_summary_str = ", ".join(
+                [f"{cat.id} ({cat.name}) from {cat.dataset_name}" for cat in merged_categories_i]
+            )
+            logger.warning(
+                f"Merged category {cat_id} corresponds to the following categories: "
+                f"{cat_summary_str}"
+            )
+
+
+def combine_detection_dataset_dicts(
+    dataset_names: Collection[str],
+    keep_instance_predicate: Optional[InstancePredicate] = None,
+    proposal_files: Optional[Collection[str]] = None,
+) -> List[Instance]:
+    """
+    Load and prepare dataset dicts for training / testing
+
+    Args:
+        dataset_names (Collection[str]): a list of dataset names
+        keep_instance_predicate (Callable: Dict[str, Any] -> bool): predicate
+            applied to instance dicts which defines whether to keep the instance
+        proposal_files (Collection[str]): if given, a list of object proposal files
+            that match each dataset in `dataset_names`.
+    """
+    assert len(dataset_names)
+    if proposal_files is None:
+        proposal_files = [None] * len(dataset_names)
+    assert len(dataset_names) == len(proposal_files)
+    # load datasets and metadata
+    dataset_name_to_dicts = {}
+    for dataset_name in dataset_names:
+        dataset_name_to_dicts[dataset_name] = DatasetCatalog.get(dataset_name)
+        assert len(dataset_name_to_dicts), f"Dataset '{dataset_name}' is empty!"
+    # merge categories, requires category metadata to be loaded
+    # cat_id -> [(orig_cat_id, cat_name, dataset_name)]
+    merged_categories = _merge_categories(dataset_names)
+    _warn_if_merged_different_categories(merged_categories)
+    merged_category_names = [
+        merged_categories[cat_id][0].mapped_name for cat_id in sorted(merged_categories)
+    ]
+    # map to contiguous category IDs
+    _add_category_id_to_contiguous_id_maps_to_metadata(merged_categories)
+    # load annotations and dataset metadata
+    for dataset_name, proposal_file in zip(dataset_names, proposal_files):
+        dataset_dicts = dataset_name_to_dicts[dataset_name]
+        assert len(dataset_dicts), f"Dataset '{dataset_name}' is empty!"
+        if proposal_file is not None:
+            dataset_dicts = load_proposals_into_dataset(dataset_dicts, proposal_file)
+        dataset_dicts = _maybe_filter_and_map_categories(dataset_name, dataset_dicts)
+        print_instances_class_histogram(dataset_dicts, merged_category_names)
+        dataset_name_to_dicts[dataset_name] = dataset_dicts
+
+    if keep_instance_predicate is not None:
+        all_datasets_dicts_plain = [
+            d
+            for d in itertools.chain.from_iterable(dataset_name_to_dicts.values())
+            if keep_instance_predicate(d)
+        ]
+    else:
+        all_datasets_dicts_plain = list(
+            itertools.chain.from_iterable(dataset_name_to_dicts.values())
+        )
+    return all_datasets_dicts_plain
+
+
+def build_detection_train_loader(cfg: CfgNode, mapper=None):
+    """
+    A data loader is created in a way similar to that of Detectron2.
+    The main differences are:
+     - it allows to combine datasets with different but compatible object category sets
+
+    The data loader is created by the following steps:
+    1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
+    2. Start workers to work on the dicts. Each worker will:
+        * Map each metadata dict into another format to be consumed by the model.
+        * Batch them by simply putting dicts into a list.
+    The batched ``list[mapped_dict]`` is what this dataloader will return.
+
+    Args:
+        cfg (CfgNode): the config
+        mapper (callable): a callable which takes a sample (dict) from dataset and
+            returns the format to be consumed by the model.
+            By default it will be `DatasetMapper(cfg, True)`.
+
+    Returns:
+        an infinite iterator of training data
+    """
+
+    _add_category_whitelists_to_metadata(cfg)
+    _add_category_maps_to_metadata(cfg)
+    _maybe_add_class_to_mesh_name_map_to_metadata(cfg.DATASETS.TRAIN, cfg)
+    dataset_dicts = combine_detection_dataset_dicts(
+        cfg.DATASETS.TRAIN,
+        keep_instance_predicate=_get_train_keep_instance_predicate(cfg),
+        proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
+    )
+    if mapper is None:
+        mapper = DatasetMapper(cfg, True)
+    return d2_build_detection_train_loader(cfg, dataset=dataset_dicts, mapper=mapper)
+
+
+def build_detection_test_loader(cfg, dataset_name, mapper=None):
+    """
+    Similar to `build_detection_train_loader`.
+    But this function uses the given `dataset_name` argument (instead of the names in cfg),
+    and uses batch size 1.
+
+    Args:
+        cfg: a detectron2 CfgNode
+        dataset_name (str): a name of the dataset that's available in the DatasetCatalog
+        mapper (callable): a callable which takes a sample (dict) from dataset
+            and returns the format to be consumed by the model.
+            By default it will be `DatasetMapper(cfg, False)`.
+
+    Returns:
+        DataLoader: a torch DataLoader, that loads the given detection
+            dataset, with test-time transformation and batching.
+    """
+    _add_category_whitelists_to_metadata(cfg)
+    _add_category_maps_to_metadata(cfg)
+    _maybe_add_class_to_mesh_name_map_to_metadata([dataset_name], cfg)
+    dataset_dicts = combine_detection_dataset_dicts(
+        [dataset_name],
+        keep_instance_predicate=_get_test_keep_instance_predicate(cfg),
+        proposal_files=(
+            [cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]]
+            if cfg.MODEL.LOAD_PROPOSALS
+            else None
+        ),
+    )
+    sampler = None
+    if not cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE:
+        sampler = torch.utils.data.SequentialSampler(dataset_dicts)
+    if mapper is None:
+        mapper = DatasetMapper(cfg, False)
+    return d2_build_detection_test_loader(
+        dataset_dicts, mapper=mapper, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler
+    )
+
+
+def build_frame_selector(cfg: CfgNode):
+    strategy = FrameSelectionStrategy(cfg.STRATEGY)
+    if strategy == FrameSelectionStrategy.RANDOM_K:
+        frame_selector = RandomKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.FIRST_K:
+        frame_selector = FirstKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.LAST_K:
+        frame_selector = LastKFramesSelector(cfg.NUM_IMAGES)
+    elif strategy == FrameSelectionStrategy.ALL:
+        frame_selector = None
+    # pyre-fixme[61]: `frame_selector` may not be initialized here.
+    return frame_selector
+
+
+def build_transform(cfg: CfgNode, data_type: str):
+    if cfg.TYPE == "resize":
+        if data_type == "image":
+            return ImageResizeTransform(cfg.MIN_SIZE, cfg.MAX_SIZE)
+    raise ValueError(f"Unknown transform {cfg.TYPE} for data type {data_type}")
+
+
+def build_combined_loader(cfg: CfgNode, loaders: Collection[Loader], ratios: Sequence[float]):
+    images_per_worker = _compute_num_images_per_worker(cfg)
+    return CombinedDataLoader(loaders, images_per_worker, ratios)
+
+
+def build_bootstrap_dataset(dataset_name: str, cfg: CfgNode) -> Sequence[torch.Tensor]:
+    """
+    Build dataset that provides data to bootstrap on
+
+    Args:
+        dataset_name (str): Name of the dataset, needs to have associated metadata
+            to load the data
+        cfg (CfgNode): bootstrapping config
+    Returns:
+        Sequence[Tensor] - dataset that provides image batches, Tensors of size
+            [N, C, H, W] of type float32
+    """
+    logger = logging.getLogger(__name__)
+    _add_category_info_to_bootstrapping_metadata(dataset_name, cfg)
+    meta = MetadataCatalog.get(dataset_name)
+    factory = BootstrapDatasetFactoryCatalog.get(meta.dataset_type)
+    dataset = None
+    if factory is not None:
+        dataset = factory(meta, cfg)
+    if dataset is None:
+        logger.warning(f"Failed to create dataset {dataset_name} of type {meta.dataset_type}")
+    return dataset
+
+
+def build_data_sampler(cfg: CfgNode, sampler_cfg: CfgNode, embedder: Optional[torch.nn.Module]):
+    if sampler_cfg.TYPE == "densepose_uniform":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseUniformSampler(count_per_class=sampler_cfg.COUNT_PER_CLASS),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_UV_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="sigma_2",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_fine_segm_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="fine_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_coarse_segm_confidence":
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseConfidenceBasedSampler(
+                confidence_channel="coarse_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_cse_uniform":
+        assert embedder is not None
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseCSEUniformSampler(
+                cfg=cfg,
+                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
+                embedder=embedder,
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+    elif sampler_cfg.TYPE == "densepose_cse_coarse_segm_confidence":
+        assert embedder is not None
+        data_sampler = PredictionToGroundTruthSampler()
+        # transform densepose pred -> gt
+        data_sampler.register_sampler(
+            "pred_densepose",
+            "gt_densepose",
+            DensePoseCSEConfidenceBasedSampler(
+                cfg=cfg,
+                use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
+                embedder=embedder,
+                confidence_channel="coarse_segm_confidence",
+                count_per_class=sampler_cfg.COUNT_PER_CLASS,
+                search_proportion=0.5,
+            ),
+        )
+        data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
+        return data_sampler
+
+    raise ValueError(f"Unknown data sampler type {sampler_cfg.TYPE}")
+
+
+def build_data_filter(cfg: CfgNode):
+    if cfg.TYPE == "detection_score":
+        min_score = cfg.MIN_VALUE
+        return ScoreBasedFilter(min_score=min_score)
+    raise ValueError(f"Unknown data filter type {cfg.TYPE}")
+
+
+def build_inference_based_loader(
+    cfg: CfgNode,
+    dataset_cfg: CfgNode,
+    model: torch.nn.Module,
+    embedder: Optional[torch.nn.Module] = None,
+) -> InferenceBasedLoader:
+    """
+    Constructs data loader based on inference results of a model.
+    """
+    dataset = build_bootstrap_dataset(dataset_cfg.DATASET, dataset_cfg.IMAGE_LOADER)
+    meta = MetadataCatalog.get(dataset_cfg.DATASET)
+    training_sampler = TrainingSampler(len(dataset))
+    data_loader = torch.utils.data.DataLoader(
+        dataset,  # pyre-ignore[6]
+        batch_size=dataset_cfg.IMAGE_LOADER.BATCH_SIZE,
+        sampler=training_sampler,
+        num_workers=dataset_cfg.IMAGE_LOADER.NUM_WORKERS,
+        collate_fn=trivial_batch_collator,
+        worker_init_fn=worker_init_reset_seed,
+    )
+    return InferenceBasedLoader(
+        model,
+        data_loader=data_loader,
+        data_sampler=build_data_sampler(cfg, dataset_cfg.DATA_SAMPLER, embedder),
+        data_filter=build_data_filter(dataset_cfg.FILTER),
+        shuffle=True,
+        batch_size=dataset_cfg.INFERENCE.OUTPUT_BATCH_SIZE,
+        inference_batch_size=dataset_cfg.INFERENCE.INPUT_BATCH_SIZE,
+        category_to_class_mapping=meta.category_to_class_mapping,
+    )
+
+
+def has_inference_based_loaders(cfg: CfgNode) -> bool:
+    """
+    Returns True, if at least one inferense-based loader must
+    be instantiated for training
+    """
+    return len(cfg.BOOTSTRAP_DATASETS) > 0
+
+
+def build_inference_based_loaders(
+    cfg: CfgNode, model: torch.nn.Module
+) -> Tuple[List[InferenceBasedLoader], List[float]]:
+    loaders = []
+    ratios = []
+    embedder = build_densepose_embedder(cfg).to(device=model.device)  # pyre-ignore[16]
+    for dataset_spec in cfg.BOOTSTRAP_DATASETS:
+        dataset_cfg = get_bootstrap_dataset_config().clone()
+        dataset_cfg.merge_from_other_cfg(CfgNode(dataset_spec))
+        loader = build_inference_based_loader(cfg, dataset_cfg, model, embedder)
+        loaders.append(loader)
+        ratios.append(dataset_cfg.RATIO)
+    return loaders, ratios
+
+
+def build_video_list_dataset(meta: Metadata, cfg: CfgNode):
+    video_list_fpath = meta.video_list_fpath
+    video_base_path = meta.video_base_path
+    category = meta.category
+    if cfg.TYPE == "video_keyframe":
+        frame_selector = build_frame_selector(cfg.SELECT)
+        transform = build_transform(cfg.TRANSFORM, data_type="image")
+        video_list = video_list_from_file(video_list_fpath, video_base_path)
+        keyframe_helper_fpath = getattr(cfg, "KEYFRAME_HELPER", None)
+        return VideoKeyframeDataset(
+            video_list, category, frame_selector, transform, keyframe_helper_fpath
+        )
+
+
+class _BootstrapDatasetFactoryCatalog(UserDict):
+    """
+    A global dictionary that stores information about bootstrapped datasets creation functions
+    from metadata and config, for diverse DatasetType
+    """
+
+    def register(self, dataset_type: DatasetType, factory: Callable[[Metadata, CfgNode], Dataset]):
+        """
+        Args:
+            dataset_type (DatasetType): a DatasetType e.g. DatasetType.VIDEO_LIST
+            factory (Callable[Metadata, CfgNode]): a callable which takes Metadata and cfg
+            arguments and returns a dataset object.
+        """
+        assert dataset_type not in self, "Dataset '{}' is already registered!".format(dataset_type)
+        self[dataset_type] = factory
+
+
+BootstrapDatasetFactoryCatalog = _BootstrapDatasetFactoryCatalog()
+BootstrapDatasetFactoryCatalog.register(DatasetType.VIDEO_LIST, build_video_list_dataset)

densepose/data/combined_loader.py

@@ -0,0 +1,46 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from collections import deque
+from typing import Any, Collection, Deque, Iterable, Iterator, List, Sequence
+
+Loader = Iterable[Any]
+
+
+def _pooled_next(iterator: Iterator[Any], pool: Deque[Any]):
+    if not pool:
+        pool.extend(next(iterator))
+    return pool.popleft()
+
+
+class CombinedDataLoader:
+    """
+    Combines data loaders using the provided sampling ratios
+    """
+
+    BATCH_COUNT = 100
+
+    def __init__(self, loaders: Collection[Loader], batch_size: int, ratios: Sequence[float]):
+        self.loaders = loaders
+        self.batch_size = batch_size
+        self.ratios = ratios
+
+    def __iter__(self) -> Iterator[List[Any]]:
+        iters = [iter(loader) for loader in self.loaders]
+        indices = []
+        pool = [deque()] * len(iters)
+        # infinite iterator, as in D2
+        while True:
+            if not indices:
+                # just a buffer of indices, its size doesn't matter
+                # as long as it's a multiple of batch_size
+                k = self.batch_size * self.BATCH_COUNT
+                indices = random.choices(range(len(self.loaders)), self.ratios, k=k)
+            try:
+                batch = [_pooled_next(iters[i], pool[i]) for i in indices[: self.batch_size]]
+            except StopIteration:
+                break
+            indices = indices[self.batch_size :]
+            yield batch

densepose/data/dataset_mapper.py

@@ -0,0 +1,170 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import copy
+import logging
+from typing import Any, Dict, List, Tuple
+import torch
+
+from detectron2.data import MetadataCatalog
+from detectron2.data import detection_utils as utils
+from detectron2.data import transforms as T
+from detectron2.layers import ROIAlign
+from detectron2.structures import BoxMode
+from detectron2.utils.file_io import PathManager
+
+from densepose.structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
+
+
+def build_augmentation(cfg, is_train):
+    logger = logging.getLogger(__name__)
+    result = utils.build_augmentation(cfg, is_train)
+    if is_train:
+        random_rotation = T.RandomRotation(
+            cfg.INPUT.ROTATION_ANGLES, expand=False, sample_style="choice"
+        )
+        result.append(random_rotation)
+        logger.info("DensePose-specific augmentation used in training: " + str(random_rotation))
+    return result
+
+
+class DatasetMapper:
+    """
+    A customized version of `detectron2.data.DatasetMapper`
+    """
+
+    def __init__(self, cfg, is_train=True):
+        self.augmentation = build_augmentation(cfg, is_train)
+
+        # fmt: off
+        self.img_format     = cfg.INPUT.FORMAT
+        self.mask_on        = (
+            cfg.MODEL.MASK_ON or (
+                cfg.MODEL.DENSEPOSE_ON
+                and cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS)
+        )
+        self.keypoint_on    = cfg.MODEL.KEYPOINT_ON
+        self.densepose_on   = cfg.MODEL.DENSEPOSE_ON
+        assert not cfg.MODEL.LOAD_PROPOSALS, "not supported yet"
+        # fmt: on
+        if self.keypoint_on and is_train:
+            # Flip only makes sense in training
+            self.keypoint_hflip_indices = utils.create_keypoint_hflip_indices(cfg.DATASETS.TRAIN)
+        else:
+            self.keypoint_hflip_indices = None
+
+        if self.densepose_on:
+            densepose_transform_srcs = [
+                MetadataCatalog.get(ds).densepose_transform_src
+                for ds in cfg.DATASETS.TRAIN + cfg.DATASETS.TEST
+            ]
+            assert len(densepose_transform_srcs) > 0
+            # TODO: check that DensePose transformation data is the same for
+            # all the datasets. Otherwise one would have to pass DB ID with
+            # each entry to select proper transformation data. For now, since
+            # all DensePose annotated data uses the same data semantics, we
+            # omit this check.
+            densepose_transform_data_fpath = PathManager.get_local_path(densepose_transform_srcs[0])
+            self.densepose_transform_data = DensePoseTransformData.load(
+                densepose_transform_data_fpath
+            )
+
+        self.is_train = is_train
+
+    def __call__(self, dataset_dict):
+        """
+        Args:
+            dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
+
+        Returns:
+            dict: a format that builtin models in detectron2 accept
+        """
+        dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
+        image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
+        utils.check_image_size(dataset_dict, image)
+
+        image, transforms = T.apply_transform_gens(self.augmentation, image)
+        image_shape = image.shape[:2]  # h, w
+        dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
+
+        if not self.is_train:
+            dataset_dict.pop("annotations", None)
+            return dataset_dict
+
+        for anno in dataset_dict["annotations"]:
+            if not self.mask_on:
+                anno.pop("segmentation", None)
+            if not self.keypoint_on:
+                anno.pop("keypoints", None)
+
+        # USER: Implement additional transformations if you have other types of data
+        # USER: Don't call transpose_densepose if you don't need
+        annos = [
+            self._transform_densepose(
+                utils.transform_instance_annotations(
+                    obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices
+                ),
+                transforms,
+            )
+            for obj in dataset_dict.pop("annotations")
+            if obj.get("iscrowd", 0) == 0
+        ]
+
+        if self.mask_on:
+            self._add_densepose_masks_as_segmentation(annos, image_shape)
+
+        instances = utils.annotations_to_instances(annos, image_shape, mask_format="bitmask")
+        densepose_annotations = [obj.get("densepose") for obj in annos]
+        if densepose_annotations and not all(v is None for v in densepose_annotations):
+            instances.gt_densepose = DensePoseList(
+                densepose_annotations, instances.gt_boxes, image_shape
+            )
+
+        dataset_dict["instances"] = instances[instances.gt_boxes.nonempty()]
+        return dataset_dict
+
+    def _transform_densepose(self, annotation, transforms):
+        if not self.densepose_on:
+            return annotation
+
+        # Handle densepose annotations
+        is_valid, reason_not_valid = DensePoseDataRelative.validate_annotation(annotation)
+        if is_valid:
+            densepose_data = DensePoseDataRelative(annotation, cleanup=True)
+            densepose_data.apply_transform(transforms, self.densepose_transform_data)
+            annotation["densepose"] = densepose_data
+        else:
+            # logger = logging.getLogger(__name__)
+            # logger.debug("Could not load DensePose annotation: {}".format(reason_not_valid))
+            DensePoseDataRelative.cleanup_annotation(annotation)
+            # NOTE: annotations for certain instances may be unavailable.
+            # 'None' is accepted by the DensePostList data structure.
+            annotation["densepose"] = None
+        return annotation
+
+    def _add_densepose_masks_as_segmentation(
+        self, annotations: List[Dict[str, Any]], image_shape_hw: Tuple[int, int]
+    ):
+        for obj in annotations:
+            if ("densepose" not in obj) or ("segmentation" in obj):
+                continue
+            # DP segmentation: torch.Tensor [S, S] of float32, S=256
+            segm_dp = torch.zeros_like(obj["densepose"].segm)
+            segm_dp[obj["densepose"].segm > 0] = 1
+            segm_h, segm_w = segm_dp.shape
+            bbox_segm_dp = torch.tensor((0, 0, segm_h - 1, segm_w - 1), dtype=torch.float32)
+            # image bbox
+            x0, y0, x1, y1 = (
+                v.item() for v in BoxMode.convert(obj["bbox"], obj["bbox_mode"], BoxMode.XYXY_ABS)
+            )
+            segm_aligned = (
+                ROIAlign((y1 - y0, x1 - x0), 1.0, 0, aligned=True)
+                .forward(segm_dp.view(1, 1, *segm_dp.shape), bbox_segm_dp)
+                .squeeze()
+            )
+            image_mask = torch.zeros(*image_shape_hw, dtype=torch.float32)
+            image_mask[y0:y1, x0:x1] = segm_aligned
+            # segmentation for BitMask: np.array [H, W] of bool
+            obj["segmentation"] = image_mask >= 0.5

densepose/data/image_list_dataset.py

@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import logging
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Union
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.data.detection_utils import read_image
+
+ImageTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+class ImageListDataset(Dataset):
+    """
+    Dataset that provides images from a list.
+    """
+
+    _EMPTY_IMAGE = torch.empty((0, 3, 1, 1))
+
+    def __init__(
+        self,
+        image_list: List[str],
+        category_list: Union[str, List[str], None] = None,
+        transform: Optional[ImageTransform] = None,
+    ):
+        """
+        Args:
+            image_list (List[str]): list of paths to image files
+            category_list (Union[str, List[str], None]): list of animal categories for
+                each image. If it is a string, or None, this applies to all images
+        """
+        if type(category_list) is list:
+            self.category_list = category_list
+        else:
+            self.category_list = [category_list] * len(image_list)
+        assert len(image_list) == len(
+            self.category_list
+        ), "length of image and category lists must be equal"
+        self.image_list = image_list
+        self.transform = transform
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Gets selected images from the list
+
+        Args:
+            idx (int): video index in the video list file
+        Returns:
+            A dictionary containing two keys:
+                images (torch.Tensor): tensor of size [N, 3, H, W] (N = 1, or 0 for _EMPTY_IMAGE)
+                categories (List[str]): categories of the frames
+        """
+        categories = [self.category_list[idx]]
+        fpath = self.image_list[idx]
+        transform = self.transform
+
+        try:
+            image = torch.from_numpy(np.ascontiguousarray(read_image(fpath, format="BGR")))
+            image = image.permute(2, 0, 1).unsqueeze(0).float()  # HWC -> NCHW
+            if transform is not None:
+                image = transform(image)
+            return {"images": image, "categories": categories}
+        except (OSError, RuntimeError) as e:
+            logger = logging.getLogger(__name__)
+            logger.warning(f"Error opening image file container {fpath}: {e}")
+
+        return {"images": self._EMPTY_IMAGE, "categories": []}
+
+    def __len__(self):
+        return len(self.image_list)

densepose/data/inference_based_loader.py

@@ -0,0 +1,174 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple
+import torch
+from torch import nn
+
+SampledData = Any
+ModelOutput = Any
+
+
+def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]:
+    """
+    Group elements of an iterable by chunks of size `n`, e.g.
+    grouper(range(9), 4) ->
+        (0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None)
+    """
+    it = iter(iterable)
+    while True:
+        values = []
+        for _ in range(n):
+            try:
+                value = next(it)
+            except StopIteration:
+                if values:
+                    values.extend([fillvalue] * (n - len(values)))
+                    yield tuple(values)
+                return
+            values.append(value)
+        yield tuple(values)
+
+
+class ScoreBasedFilter:
+    """
+    Filters entries in model output based on their scores
+    Discards all entries with score less than the specified minimum
+    """
+
+    def __init__(self, min_score: float = 0.8):
+        self.min_score = min_score
+
+    def __call__(self, model_output: ModelOutput) -> ModelOutput:
+        for model_output_i in model_output:
+            instances = model_output_i["instances"]
+            if not instances.has("scores"):
+                continue
+            instances_filtered = instances[instances.scores >= self.min_score]
+            model_output_i["instances"] = instances_filtered
+        return model_output
+
+
+class InferenceBasedLoader:
+    """
+    Data loader based on results inferred by a model. Consists of:
+     - a data loader that provides batches of images
+     - a model that is used to infer the results
+     - a data sampler that converts inferred results to annotations
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        data_loader: Iterable[List[Dict[str, Any]]],
+        data_sampler: Optional[Callable[[ModelOutput], List[SampledData]]] = None,
+        data_filter: Optional[Callable[[ModelOutput], ModelOutput]] = None,
+        shuffle: bool = True,
+        batch_size: int = 4,
+        inference_batch_size: int = 4,
+        drop_last: bool = False,
+        category_to_class_mapping: Optional[dict] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          model (torch.nn.Module): model used to produce data
+          data_loader (Iterable[List[Dict[str, Any]]]): iterable that provides
+            dictionaries with "images" and "categories" fields to perform inference on
+          data_sampler (Callable: ModelOutput -> SampledData): functor
+              that produces annotation data from inference results;
+              (optional, default: None)
+          data_filter (Callable: ModelOutput -> ModelOutput): filter
+              that selects model outputs for further processing
+              (optional, default: None)
+          shuffle (bool): if True, the input images get shuffled
+          batch_size (int): batch size for the produced annotation data
+          inference_batch_size (int): batch size for input images
+          drop_last (bool): if True, drop the last batch if it is undersized
+          category_to_class_mapping (dict): category to class mapping
+        """
+        self.model = model
+        self.model.eval()
+        self.data_loader = data_loader
+        self.data_sampler = data_sampler
+        self.data_filter = data_filter
+        self.shuffle = shuffle
+        self.batch_size = batch_size
+        self.inference_batch_size = inference_batch_size
+        self.drop_last = drop_last
+        if category_to_class_mapping is not None:
+            self.category_to_class_mapping = category_to_class_mapping
+        else:
+            self.category_to_class_mapping = {}
+
+    def __iter__(self) -> Iterator[List[SampledData]]:
+        for batch in self.data_loader:
+            # batch : List[Dict[str: Tensor[N, C, H, W], str: Optional[str]]]
+            # images_batch : Tensor[N, C, H, W]
+            # image : Tensor[C, H, W]
+            images_and_categories = [
+                {"image": image, "category": category}
+                for element in batch
+                for image, category in zip(element["images"], element["categories"])
+            ]
+            if not images_and_categories:
+                continue
+            if self.shuffle:
+                random.shuffle(images_and_categories)
+            yield from self._produce_data(images_and_categories)  # pyre-ignore[6]
+
+    def _produce_data(
+        self, images_and_categories: List[Tuple[torch.Tensor, Optional[str]]]
+    ) -> Iterator[List[SampledData]]:
+        """
+        Produce batches of data from images
+
+        Args:
+          images_and_categories (List[Tuple[torch.Tensor, Optional[str]]]):
+            list of images and corresponding categories to process
+
+        Returns:
+          Iterator over batches of data sampled from model outputs
+        """
+        data_batches: List[SampledData] = []
+        category_to_class_mapping = self.category_to_class_mapping
+        batched_images_and_categories = _grouper(images_and_categories, self.inference_batch_size)
+        for batch in batched_images_and_categories:
+            batch = [
+                {
+                    "image": image_and_category["image"].to(self.model.device),
+                    "category": image_and_category["category"],
+                }
+                for image_and_category in batch
+                if image_and_category is not None
+            ]
+            if not batch:
+                continue
+            with torch.no_grad():
+                model_output = self.model(batch)
+            for model_output_i, batch_i in zip(model_output, batch):
+                assert len(batch_i["image"].shape) == 3
+                model_output_i["image"] = batch_i["image"]
+                instance_class = category_to_class_mapping.get(batch_i["category"], 0)
+                model_output_i["instances"].dataset_classes = torch.tensor(
+                    [instance_class] * len(model_output_i["instances"])
+                )
+            model_output_filtered = (
+                model_output if self.data_filter is None else self.data_filter(model_output)
+            )
+            data = (
+                model_output_filtered
+                if self.data_sampler is None
+                else self.data_sampler(model_output_filtered)
+            )
+            for data_i in data:
+                if len(data_i["instances"]):
+                    data_batches.append(data_i)
+            if len(data_batches) >= self.batch_size:
+                yield data_batches[: self.batch_size]
+                data_batches = data_batches[self.batch_size :]
+        if not self.drop_last and data_batches:
+            yield data_batches

densepose/data/meshes/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+from . import builtin
+
+__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]

densepose/data/meshes/builtin.py

@@ -0,0 +1,103 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+from .catalog import MeshInfo, register_meshes
+
+DENSEPOSE_MESHES_DIR = "https://dl.fbaipublicfiles.com/densepose/meshes/"
+
+MESHES = [
+    MeshInfo(
+        name="smpl_27554",
+        data="smpl_27554.pkl",
+        geodists="geodists/geodists_smpl_27554.pkl",
+        symmetry="symmetry/symmetry_smpl_27554.pkl",
+        texcoords="texcoords/texcoords_smpl_27554.pkl",
+    ),
+    MeshInfo(
+        name="chimp_5029",
+        data="chimp_5029.pkl",
+        geodists="geodists/geodists_chimp_5029.pkl",
+        symmetry="symmetry/symmetry_chimp_5029.pkl",
+        texcoords="texcoords/texcoords_chimp_5029.pkl",
+    ),
+    MeshInfo(
+        name="cat_5001",
+        data="cat_5001.pkl",
+        geodists="geodists/geodists_cat_5001.pkl",
+        symmetry="symmetry/symmetry_cat_5001.pkl",
+        texcoords="texcoords/texcoords_cat_5001.pkl",
+    ),
+    MeshInfo(
+        name="cat_7466",
+        data="cat_7466.pkl",
+        geodists="geodists/geodists_cat_7466.pkl",
+        symmetry="symmetry/symmetry_cat_7466.pkl",
+        texcoords="texcoords/texcoords_cat_7466.pkl",
+    ),
+    MeshInfo(
+        name="sheep_5004",
+        data="sheep_5004.pkl",
+        geodists="geodists/geodists_sheep_5004.pkl",
+        symmetry="symmetry/symmetry_sheep_5004.pkl",
+        texcoords="texcoords/texcoords_sheep_5004.pkl",
+    ),
+    MeshInfo(
+        name="zebra_5002",
+        data="zebra_5002.pkl",
+        geodists="geodists/geodists_zebra_5002.pkl",
+        symmetry="symmetry/symmetry_zebra_5002.pkl",
+        texcoords="texcoords/texcoords_zebra_5002.pkl",
+    ),
+    MeshInfo(
+        name="horse_5004",
+        data="horse_5004.pkl",
+        geodists="geodists/geodists_horse_5004.pkl",
+        symmetry="symmetry/symmetry_horse_5004.pkl",
+        texcoords="texcoords/texcoords_zebra_5002.pkl",
+    ),
+    MeshInfo(
+        name="giraffe_5002",
+        data="giraffe_5002.pkl",
+        geodists="geodists/geodists_giraffe_5002.pkl",
+        symmetry="symmetry/symmetry_giraffe_5002.pkl",
+        texcoords="texcoords/texcoords_giraffe_5002.pkl",
+    ),
+    MeshInfo(
+        name="elephant_5002",
+        data="elephant_5002.pkl",
+        geodists="geodists/geodists_elephant_5002.pkl",
+        symmetry="symmetry/symmetry_elephant_5002.pkl",
+        texcoords="texcoords/texcoords_elephant_5002.pkl",
+    ),
+    MeshInfo(
+        name="dog_5002",
+        data="dog_5002.pkl",
+        geodists="geodists/geodists_dog_5002.pkl",
+        symmetry="symmetry/symmetry_dog_5002.pkl",
+        texcoords="texcoords/texcoords_dog_5002.pkl",
+    ),
+    MeshInfo(
+        name="dog_7466",
+        data="dog_7466.pkl",
+        geodists="geodists/geodists_dog_7466.pkl",
+        symmetry="symmetry/symmetry_dog_7466.pkl",
+        texcoords="texcoords/texcoords_dog_7466.pkl",
+    ),
+    MeshInfo(
+        name="cow_5002",
+        data="cow_5002.pkl",
+        geodists="geodists/geodists_cow_5002.pkl",
+        symmetry="symmetry/symmetry_cow_5002.pkl",
+        texcoords="texcoords/texcoords_cow_5002.pkl",
+    ),
+    MeshInfo(
+        name="bear_4936",
+        data="bear_4936.pkl",
+        geodists="geodists/geodists_bear_4936.pkl",
+        symmetry="symmetry/symmetry_bear_4936.pkl",
+        texcoords="texcoords/texcoords_bear_4936.pkl",
+    ),
+]
+
+register_meshes(MESHES, DENSEPOSE_MESHES_DIR)

densepose/data/meshes/catalog.py

@@ -0,0 +1,73 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+import logging
+from collections import UserDict
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+from ..utils import maybe_prepend_base_path
+
+
+@dataclass
+class MeshInfo:
+    name: str
+    data: str
+    geodists: Optional[str] = None
+    symmetry: Optional[str] = None
+    texcoords: Optional[str] = None
+
+
+class _MeshCatalog(UserDict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mesh_ids = {}
+        self.mesh_names = {}
+        self.max_mesh_id = -1
+
+    def __setitem__(self, key, value):
+        if key in self:
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                f"Overwriting mesh catalog entry '{key}': old value {self[key]}"
+                f", new value {value}"
+            )
+            mesh_id = self.mesh_ids[key]
+        else:
+            self.max_mesh_id += 1
+            mesh_id = self.max_mesh_id
+        super().__setitem__(key, value)
+        self.mesh_ids[key] = mesh_id
+        self.mesh_names[mesh_id] = key
+
+    def get_mesh_id(self, shape_name: str) -> int:
+        return self.mesh_ids[shape_name]
+
+    def get_mesh_name(self, mesh_id: int) -> str:
+        return self.mesh_names[mesh_id]
+
+
+MeshCatalog = _MeshCatalog()
+
+
+def register_mesh(mesh_info: MeshInfo, base_path: Optional[str]) -> None:
+    geodists, symmetry, texcoords = mesh_info.geodists, mesh_info.symmetry, mesh_info.texcoords
+    if geodists:
+        geodists = maybe_prepend_base_path(base_path, geodists)
+    if symmetry:
+        symmetry = maybe_prepend_base_path(base_path, symmetry)
+    if texcoords:
+        texcoords = maybe_prepend_base_path(base_path, texcoords)
+    MeshCatalog[mesh_info.name] = MeshInfo(
+        name=mesh_info.name,
+        data=maybe_prepend_base_path(base_path, mesh_info.data),
+        geodists=geodists,
+        symmetry=symmetry,
+        texcoords=texcoords,
+    )
+
+
+def register_meshes(mesh_infos: Iterable[MeshInfo], base_path: Optional[str]) -> None:
+    for mesh_info in mesh_infos:
+        register_mesh(mesh_info, base_path)

densepose/data/samplers/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .densepose_uniform import DensePoseUniformSampler
+from .densepose_confidence_based import DensePoseConfidenceBasedSampler
+from .densepose_cse_uniform import DensePoseCSEUniformSampler
+from .densepose_cse_confidence_based import DensePoseCSEConfidenceBasedSampler
+from .mask_from_densepose import MaskFromDensePoseSampler
+from .prediction_to_gt import PredictionToGroundTruthSampler

densepose/data/samplers/densepose_base.py

@@ -0,0 +1,205 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Dict, List, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.structures import BoxMode, Instances
+
+from densepose.converters import ToChartResultConverter
+from densepose.converters.base import IntTupleBox, make_int_box
+from densepose.structures import DensePoseDataRelative, DensePoseList
+
+
+class DensePoseBaseSampler:
+    """
+    Base DensePose sampler to produce DensePose data from DensePose predictions.
+    Samples for each class are drawn according to some distribution over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(self, count_per_class: int = 8):
+        """
+        Constructor
+
+        Args:
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        self.count_per_class = count_per_class
+
+    def __call__(self, instances: Instances) -> DensePoseList:
+        """
+        Convert DensePose predictions (an instance of `DensePoseChartPredictorOutput`)
+        into DensePose annotations data (an instance of `DensePoseList`)
+        """
+        boxes_xyxy_abs = instances.pred_boxes.tensor.clone().cpu()
+        boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+        dp_datas = []
+        for i in range(len(boxes_xywh_abs)):
+            annotation_i = self._sample(instances[i], make_int_box(boxes_xywh_abs[i]))
+            annotation_i[DensePoseDataRelative.S_KEY] = self._resample_mask(  # pyre-ignore[6]
+                instances[i].pred_densepose
+            )
+            dp_datas.append(DensePoseDataRelative(annotation_i))
+        # create densepose annotations on CPU
+        dp_list = DensePoseList(dp_datas, boxes_xyxy_abs, instances.image_size)
+        return dp_list
+
+    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
+        """
+        Sample DensPoseDataRelative from estimation results
+        """
+        labels, dp_result = self._produce_labels_and_results(instance)
+        annotation = {
+            DensePoseDataRelative.X_KEY: [],
+            DensePoseDataRelative.Y_KEY: [],
+            DensePoseDataRelative.U_KEY: [],
+            DensePoseDataRelative.V_KEY: [],
+            DensePoseDataRelative.I_KEY: [],
+        }
+        n, h, w = dp_result.shape
+        for part_id in range(1, DensePoseDataRelative.N_PART_LABELS + 1):
+            # indices - tuple of 3 1D tensors of size k
+            # 0: index along the first dimension N
+            # 1: index along H dimension
+            # 2: index along W dimension
+            indices = torch.nonzero(labels.expand(n, h, w) == part_id, as_tuple=True)
+            # values - an array of size [n, k]
+            # n: number of channels (U, V, confidences)
+            # k: number of points labeled with part_id
+            values = dp_result[indices].view(n, -1)
+            k = values.shape[1]
+            count = min(self.count_per_class, k)
+            if count <= 0:
+                continue
+            index_sample = self._produce_index_sample(values, count)
+            sampled_values = values[:, index_sample]
+            sampled_y = indices[1][index_sample] + 0.5
+            sampled_x = indices[2][index_sample] + 0.5
+            # prepare / normalize data
+            x = (sampled_x / w * 256.0).cpu().tolist()
+            y = (sampled_y / h * 256.0).cpu().tolist()
+            u = sampled_values[0].clamp(0, 1).cpu().tolist()
+            v = sampled_values[1].clamp(0, 1).cpu().tolist()
+            fine_segm_labels = [part_id] * count
+            # extend annotations
+            annotation[DensePoseDataRelative.X_KEY].extend(x)
+            annotation[DensePoseDataRelative.Y_KEY].extend(y)
+            annotation[DensePoseDataRelative.U_KEY].extend(u)
+            annotation[DensePoseDataRelative.V_KEY].extend(v)
+            annotation[DensePoseDataRelative.I_KEY].extend(fine_segm_labels)
+        return annotation
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Abstract method to produce a sample of indices to select data
+        To be implemented in descendants
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        raise NotImplementedError
+
+    def _produce_labels_and_results(self, instance: Instances) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of `DensePoseChartPredictorOutput`
+
+        Return:
+            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
+            dp_result (torch.Tensor): shape [2, H, W], stacked DensePose results u and v
+        """
+        converter = ToChartResultConverter
+        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
+        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
+        return labels, dp_result
+
+    def _resample_mask(self, output: Any) -> torch.Tensor:
+        """
+        Convert DensePose predictor output to segmentation annotation - tensors of size
+        (256, 256) and type `int64`.
+
+        Args:
+            output: DensePose predictor output with the following attributes:
+             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
+               segmentation scores
+             - fine_segm: tensor of size [N, C, H, W] with unnormalized fine
+               segmentation scores
+        Return:
+            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
+            where S = DensePoseDataRelative.MASK_SIZE
+        """
+        sz = DensePoseDataRelative.MASK_SIZE
+        S = (
+            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
+            .argmax(dim=1)
+            .long()
+        )
+        I = (
+            (
+                F.interpolate(
+                    output.fine_segm,
+                    (sz, sz),
+                    mode="bilinear",
+                    align_corners=False,
+                ).argmax(dim=1)
+                * (S > 0).long()
+            )
+            .squeeze()
+            .cpu()
+        )
+        # Map fine segmentation results to coarse segmentation ground truth
+        # TODO: extract this into separate classes
+        # coarse segmentation: 1 = Torso, 2 = Right Hand, 3 = Left Hand,
+        # 4 = Left Foot, 5 = Right Foot, 6 = Upper Leg Right, 7 = Upper Leg Left,
+        # 8 = Lower Leg Right, 9 = Lower Leg Left, 10 = Upper Arm Left,
+        # 11 = Upper Arm Right, 12 = Lower Arm Left, 13 = Lower Arm Right,
+        # 14 = Head
+        # fine segmentation: 1, 2 = Torso, 3 = Right Hand, 4 = Left Hand,
+        # 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right,
+        # 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right,
+        # 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left,
+        # 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left,
+        # 20, 22 = Lower Arm Right, 23, 24 = Head
+        FINE_TO_COARSE_SEGMENTATION = {
+            1: 1,
+            2: 1,
+            3: 2,
+            4: 3,
+            5: 4,
+            6: 5,
+            7: 6,
+            8: 7,
+            9: 6,
+            10: 7,
+            11: 8,
+            12: 9,
+            13: 8,
+            14: 9,
+            15: 10,
+            16: 11,
+            17: 10,
+            18: 11,
+            19: 12,
+            20: 13,
+            21: 12,
+            22: 13,
+            23: 14,
+            24: 14,
+        }
+        mask = torch.zeros((sz, sz), dtype=torch.int64, device=torch.device("cpu"))
+        for i in range(DensePoseDataRelative.N_PART_LABELS):
+            mask[I == i + 1] = FINE_TO_COARSE_SEGMENTATION[i + 1]
+        return mask

densepose/data/samplers/densepose_confidence_based.py

@@ -0,0 +1,110 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from typing import Optional, Tuple
+import torch
+
+from densepose.converters import ToChartResultConverterWithConfidences
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseConfidenceBasedSampler(DensePoseBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn using confidence value estimates.
+    """
+
+    def __init__(
+        self,
+        confidence_channel: str,
+        count_per_class: int = 8,
+        search_count_multiplier: Optional[float] = None,
+        search_proportion: Optional[float] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          confidence_channel (str): confidence channel to use for sampling;
+            possible values:
+              "sigma_2": confidences for UV values
+              "fine_segm_confidence": confidences for fine segmentation
+              "coarse_segm_confidence": confidences for coarse segmentation
+            (default: "sigma_2")
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category (default: 8)
+          search_count_multiplier (float or None): if not None, the total number
+              of the most confident estimates of a given class to consider is
+              defined as `min(search_count_multiplier * count_per_class, N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_proportion` (default: None)
+          search_proportion (float or None): if not None, the total number of the
+              of the most confident estimates of a given class to consider is
+              defined as `min(max(search_proportion * N, count_per_class), N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_count_multiplier` (default: None)
+        """
+        super().__init__(count_per_class)
+        self.confidence_channel = confidence_channel
+        self.search_count_multiplier = search_count_multiplier
+        self.search_proportion = search_proportion
+        assert (search_count_multiplier is None) or (search_proportion is None), (
+            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
+            f"and search_proportion (={search_proportion})"
+        )
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a sample of indices to select data based on confidences
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        if k == count:
+            index_sample = list(range(k))
+        else:
+            # take the best count * search_count_multiplier pixels,
+            # sample from them uniformly
+            # (here best = smallest variance)
+            _, sorted_confidence_indices = torch.sort(values[2])
+            if self.search_count_multiplier is not None:
+                search_count = min(int(count * self.search_count_multiplier), k)
+            elif self.search_proportion is not None:
+                search_count = min(max(int(k * self.search_proportion), count), k)
+            else:
+                search_count = min(count, k)
+            sample_from_top = random.sample(range(search_count), count)
+            index_sample = sorted_confidence_indices[:search_count][sample_from_top]
+        return index_sample
+
+    def _produce_labels_and_results(self, instance) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance, with confidences
+
+        Args:
+            instance (Instances): an instance of `DensePoseChartPredictorOutputWithConfidences`
+
+        Return:
+            labels (torch.Tensor): shape [H, W], DensePose segmentation labels
+            dp_result (torch.Tensor): shape [3, H, W], DensePose results u and v
+                stacked with the confidence channel
+        """
+        converter = ToChartResultConverterWithConfidences
+        chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
+        labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
+        dp_result = torch.cat(
+            (dp_result, getattr(chart_result, self.confidence_channel)[None].cpu())
+        )
+
+        return labels, dp_result

densepose/data/samplers/densepose_cse_base.py

@@ -0,0 +1,141 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from typing import Any, Dict, List, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.config import CfgNode
+from detectron2.structures import Instances
+
+from densepose.converters.base import IntTupleBox
+from densepose.data.utils import get_class_to_mesh_name_mapping
+from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
+from densepose.structures import DensePoseDataRelative
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseCSEBaseSampler(DensePoseBaseSampler):
+    """
+    Base DensePose sampler to produce DensePose data from DensePose predictions.
+    Samples for each class are drawn according to some distribution over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(
+        self,
+        cfg: CfgNode,
+        use_gt_categories: bool,
+        embedder: torch.nn.Module,
+        count_per_class: int = 8,
+    ):
+        """
+        Constructor
+
+        Args:
+          cfg (CfgNode): the config of the model
+          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        super().__init__(count_per_class)
+        self.embedder = embedder
+        self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
+        self.use_gt_categories = use_gt_categories
+
+    def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
+        """
+        Sample DensPoseDataRelative from estimation results
+        """
+        if self.use_gt_categories:
+            instance_class = instance.dataset_classes.tolist()[0]
+        else:
+            instance_class = instance.pred_classes.tolist()[0]
+        mesh_name = self.class_to_mesh_name[instance_class]
+
+        annotation = {
+            DensePoseDataRelative.X_KEY: [],
+            DensePoseDataRelative.Y_KEY: [],
+            DensePoseDataRelative.VERTEX_IDS_KEY: [],
+            DensePoseDataRelative.MESH_NAME_KEY: mesh_name,
+        }
+
+        mask, embeddings, other_values = self._produce_mask_and_results(instance, bbox_xywh)
+        indices = torch.nonzero(mask, as_tuple=True)
+        selected_embeddings = embeddings.permute(1, 2, 0)[indices].cpu()
+        values = other_values[:, indices[0], indices[1]]
+        k = values.shape[1]
+
+        count = min(self.count_per_class, k)
+        if count <= 0:
+            return annotation
+
+        index_sample = self._produce_index_sample(values, count)
+        closest_vertices = squared_euclidean_distance_matrix(
+            selected_embeddings[index_sample], self.embedder(mesh_name)
+        )
+        closest_vertices = torch.argmin(closest_vertices, dim=1)
+
+        sampled_y = indices[0][index_sample] + 0.5
+        sampled_x = indices[1][index_sample] + 0.5
+        # prepare / normalize data
+        _, _, w, h = bbox_xywh
+        x = (sampled_x / w * 256.0).cpu().tolist()
+        y = (sampled_y / h * 256.0).cpu().tolist()
+        # extend annotations
+        annotation[DensePoseDataRelative.X_KEY].extend(x)
+        annotation[DensePoseDataRelative.Y_KEY].extend(y)
+        annotation[DensePoseDataRelative.VERTEX_IDS_KEY].extend(closest_vertices.cpu().tolist())
+        return annotation
+
+    def _produce_mask_and_results(
+        self, instance: Instances, bbox_xywh: IntTupleBox
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of `DensePoseEmbeddingPredictorOutput`
+            bbox_xywh (IntTupleBox): the corresponding bounding box
+
+        Return:
+            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
+            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W],
+                DensePose CSE Embeddings
+            other_values (Tuple[torch.Tensor]): a tensor of shape [0, H, W],
+                for potential other values
+        """
+        densepose_output = instance.pred_densepose
+        S = densepose_output.coarse_segm
+        E = densepose_output.embedding
+        _, _, w, h = bbox_xywh
+        embeddings = F.interpolate(E, size=(h, w), mode="bilinear")[0]
+        coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0]
+        mask = coarse_segm_resized.argmax(0) > 0
+        other_values = torch.empty((0, h, w), device=E.device)
+        return mask, embeddings, other_values
+
+    def _resample_mask(self, output: Any) -> torch.Tensor:
+        """
+        Convert DensePose predictor output to segmentation annotation - tensors of size
+        (256, 256) and type `int64`.
+
+        Args:
+            output: DensePose predictor output with the following attributes:
+             - coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
+               segmentation scores
+        Return:
+            Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
+            where S = DensePoseDataRelative.MASK_SIZE
+        """
+        sz = DensePoseDataRelative.MASK_SIZE
+        mask = (
+            F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
+            .argmax(dim=1)
+            .long()
+            .squeeze()
+            .cpu()
+        )
+        return mask

densepose/data/samplers/densepose_cse_confidence_based.py

@@ -0,0 +1,121 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from typing import Optional, Tuple
+import torch
+from torch.nn import functional as F
+
+from detectron2.config import CfgNode
+from detectron2.structures import Instances
+
+from densepose.converters.base import IntTupleBox
+
+from .densepose_cse_base import DensePoseCSEBaseSampler
+
+
+class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn using confidence value estimates.
+    """
+
+    def __init__(
+        self,
+        cfg: CfgNode,
+        use_gt_categories: bool,
+        embedder: torch.nn.Module,
+        confidence_channel: str,
+        count_per_class: int = 8,
+        search_count_multiplier: Optional[float] = None,
+        search_proportion: Optional[float] = None,
+    ):
+        """
+        Constructor
+
+        Args:
+          cfg (CfgNode): the config of the model
+          embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
+          confidence_channel (str): confidence channel to use for sampling;
+            possible values:
+              "coarse_segm_confidence": confidences for coarse segmentation
+            (default: "coarse_segm_confidence")
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category (default: 8)
+          search_count_multiplier (float or None): if not None, the total number
+              of the most confident estimates of a given class to consider is
+              defined as `min(search_count_multiplier * count_per_class, N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_proportion` (default: None)
+          search_proportion (float or None): if not None, the total number of the
+              of the most confident estimates of a given class to consider is
+              defined as `min(max(search_proportion * N, count_per_class), N)`,
+              where `N` is the total number of estimates of the class; cannot be
+              specified together with `search_count_multiplier` (default: None)
+        """
+        super().__init__(cfg, use_gt_categories, embedder, count_per_class)
+        self.confidence_channel = confidence_channel
+        self.search_count_multiplier = search_count_multiplier
+        self.search_proportion = search_proportion
+        assert (search_count_multiplier is None) or (search_proportion is None), (
+            f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
+            f"and search_proportion (={search_proportion})"
+        )
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a sample of indices to select data based on confidences
+
+        Args:
+            values (torch.Tensor): a tensor of length k that contains confidences
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        if k == count:
+            index_sample = list(range(k))
+        else:
+            # take the best count * search_count_multiplier pixels,
+            # sample from them uniformly
+            # (here best = smallest variance)
+            _, sorted_confidence_indices = torch.sort(values[0])
+            if self.search_count_multiplier is not None:
+                search_count = min(int(count * self.search_count_multiplier), k)
+            elif self.search_proportion is not None:
+                search_count = min(max(int(k * self.search_proportion), count), k)
+            else:
+                search_count = min(count, k)
+            sample_from_top = random.sample(range(search_count), count)
+            index_sample = sorted_confidence_indices[-search_count:][sample_from_top]
+        return index_sample
+
+    def _produce_mask_and_results(
+        self, instance: Instances, bbox_xywh: IntTupleBox
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Method to get labels and DensePose results from an instance
+
+        Args:
+            instance (Instances): an instance of
+                `DensePoseEmbeddingPredictorOutputWithConfidences`
+            bbox_xywh (IntTupleBox): the corresponding bounding box
+
+        Return:
+            mask (torch.Tensor): shape [H, W], DensePose segmentation mask
+            embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W]
+                DensePose CSE Embeddings
+            other_values: a tensor of shape [1, H, W], DensePose CSE confidence
+        """
+        _, _, w, h = bbox_xywh
+        densepose_output = instance.pred_densepose
+        mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh)
+        other_values = F.interpolate(
+            getattr(densepose_output, self.confidence_channel),
+            size=(h, w),
+            mode="bilinear",
+        )[0].cpu()
+        return mask, embeddings, other_values

densepose/data/samplers/densepose_cse_uniform.py

@@ -0,0 +1,14 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .densepose_cse_base import DensePoseCSEBaseSampler
+from .densepose_uniform import DensePoseUniformSampler
+
+
+class DensePoseCSEUniformSampler(DensePoseCSEBaseSampler, DensePoseUniformSampler):
+    """
+    Uniform Sampler for CSE
+    """
+
+    pass

densepose/data/samplers/densepose_uniform.py

@@ -0,0 +1,43 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+import torch
+
+from .densepose_base import DensePoseBaseSampler
+
+
+class DensePoseUniformSampler(DensePoseBaseSampler):
+    """
+    Samples DensePose data from DensePose predictions.
+    Samples for each class are drawn uniformly over all pixels estimated
+    to belong to that class.
+    """
+
+    def __init__(self, count_per_class: int = 8):
+        """
+        Constructor
+
+        Args:
+          count_per_class (int): the sampler produces at most `count_per_class`
+              samples for each category
+        """
+        super().__init__(count_per_class)
+
+    def _produce_index_sample(self, values: torch.Tensor, count: int):
+        """
+        Produce a uniform sample of indices to select data
+
+        Args:
+            values (torch.Tensor): an array of size [n, k] that contains
+                estimated values (U, V, confidences);
+                n: number of channels (U, V, confidences)
+                k: number of points labeled with part_id
+            count (int): number of samples to produce, should be positive and <= k
+
+        Return:
+            list(int): indices of values (along axis 1) selected as a sample
+        """
+        k = values.shape[1]
+        return random.sample(range(k), count)

densepose/data/samplers/mask_from_densepose.py

@@ -0,0 +1,30 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from detectron2.structures import BitMasks, Instances
+
+from densepose.converters import ToMaskConverter
+
+
+class MaskFromDensePoseSampler:
+    """
+    Produce mask GT from DensePose predictions
+    This sampler simply converts DensePose predictions to BitMasks
+    that a contain a bool tensor of the size of the input image
+    """
+
+    def __call__(self, instances: Instances) -> BitMasks:
+        """
+        Converts predicted data from `instances` into the GT mask data
+
+        Args:
+            instances (Instances): predicted results, expected to have `pred_densepose` field
+
+        Returns:
+            Boolean Tensor of the size of the input image that has non-zero
+            values at pixels that are estimated to belong to the detected object
+        """
+        return ToMaskConverter.convert(
+            instances.pred_densepose, instances.pred_boxes, instances.image_size
+        )

densepose/data/samplers/prediction_to_gt.py

@@ -0,0 +1,100 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional
+
+from detectron2.structures import Instances
+
+ModelOutput = Dict[str, Any]
+SampledData = Dict[str, Any]
+
+
+@dataclass
+class _Sampler:
+    """
+    Sampler registry entry that contains:
+     - src (str): source field to sample from (deleted after sampling)
+     - dst (Optional[str]): destination field to sample to, if not None
+     - func (Optional[Callable: Any -> Any]): function that performs sampling,
+         if None, reference copy is performed
+    """
+
+    src: str
+    dst: Optional[str]
+    func: Optional[Callable[[Any], Any]]
+
+
+class PredictionToGroundTruthSampler:
+    """
+    Sampler implementation that converts predictions to GT using registered
+    samplers for different fields of `Instances`.
+    """
+
+    def __init__(self, dataset_name: str = ""):
+        self.dataset_name = dataset_name
+        self._samplers = {}
+        self.register_sampler("pred_boxes", "gt_boxes", None)
+        self.register_sampler("pred_classes", "gt_classes", None)
+        # delete scores
+        self.register_sampler("scores")
+
+    def __call__(self, model_output: List[ModelOutput]) -> List[SampledData]:
+        """
+        Transform model output into ground truth data through sampling
+
+        Args:
+          model_output (Dict[str, Any]): model output
+        Returns:
+          Dict[str, Any]: sampled data
+        """
+        for model_output_i in model_output:
+            instances: Instances = model_output_i["instances"]
+            # transform data in each field
+            for _, sampler in self._samplers.items():
+                if not instances.has(sampler.src) or sampler.dst is None:
+                    continue
+                if sampler.func is None:
+                    instances.set(sampler.dst, instances.get(sampler.src))
+                else:
+                    instances.set(sampler.dst, sampler.func(instances))
+            # delete model output data that was transformed
+            for _, sampler in self._samplers.items():
+                if sampler.src != sampler.dst and instances.has(sampler.src):
+                    instances.remove(sampler.src)
+            model_output_i["dataset"] = self.dataset_name
+        return model_output
+
+    def register_sampler(
+        self,
+        prediction_attr: str,
+        gt_attr: Optional[str] = None,
+        func: Optional[Callable[[Any], Any]] = None,
+    ):
+        """
+        Register sampler for a field
+
+        Args:
+          prediction_attr (str): field to replace with a sampled value
+          gt_attr (Optional[str]): field to store the sampled value to, if not None
+          func (Optional[Callable: Any -> Any]): sampler function
+        """
+        self._samplers[(prediction_attr, gt_attr)] = _Sampler(
+            src=prediction_attr, dst=gt_attr, func=func
+        )
+
+    def remove_sampler(
+        self,
+        prediction_attr: str,
+        gt_attr: Optional[str] = None,
+    ):
+        """
+        Remove sampler for a field
+
+        Args:
+          prediction_attr (str): field to replace with a sampled value
+          gt_attr (Optional[str]): field to store the sampled value to, if not None
+        """
+        assert (prediction_attr, gt_attr) in self._samplers
+        del self._samplers[(prediction_attr, gt_attr)]

densepose/data/transform/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .image import ImageResizeTransform

densepose/data/transform/image.py

@@ -0,0 +1,41 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import torch
+
+
+class ImageResizeTransform:
+    """
+    Transform that resizes images loaded from a dataset
+    (BGR data in NCHW channel order, typically uint8) to a format ready to be
+    consumed by DensePose training (BGR float32 data in NCHW channel order)
+    """
+
+    def __init__(self, min_size: int = 800, max_size: int = 1333):
+        self.min_size = min_size
+        self.max_size = max_size
+
+    def __call__(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            images (torch.Tensor): tensor of size [N, 3, H, W] that contains
+                BGR data (typically in uint8)
+        Returns:
+            images (torch.Tensor): tensor of size [N, 3, H1, W1] where
+                H1 and W1 are chosen to respect the specified min and max sizes
+                and preserve the original aspect ratio, the data channels
+                follow BGR order and the data type is `torch.float32`
+        """
+        # resize with min size
+        images = images.float()
+        min_size = min(images.shape[-2:])
+        max_size = max(images.shape[-2:])
+        scale = min(self.min_size / min_size, self.max_size / max_size)
+        images = torch.nn.functional.interpolate(
+            images,
+            scale_factor=scale,
+            mode="bilinear",
+            align_corners=False,
+        )
+        return images

densepose/data/utils.py

@@ -0,0 +1,40 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import os
+from typing import Dict, Optional
+
+from detectron2.config import CfgNode
+
+
+def is_relative_local_path(path: str) -> bool:
+    path_str = os.fsdecode(path)
+    return ("://" not in path_str) and not os.path.isabs(path)
+
+
+def maybe_prepend_base_path(base_path: Optional[str], path: str):
+    """
+    Prepends the provided path with a base path prefix if:
+    1) base path is not None;
+    2) path is a local path
+    """
+    if base_path is None:
+        return path
+    if is_relative_local_path(path):
+        return os.path.join(base_path, path)
+    return path
+
+
+def get_class_to_mesh_name_mapping(cfg: CfgNode) -> Dict[int, str]:
+    return {
+        int(class_id): mesh_name
+        for class_id, mesh_name in cfg.DATASETS.CLASS_TO_MESH_NAME_MAPPING.items()
+    }
+
+
+def get_category_to_class_mapping(dataset_cfg: CfgNode) -> Dict[str, int]:
+    return {
+        category: int(class_id)
+        for category, class_id in dataset_cfg.CATEGORY_TO_CLASS_MAPPING.items()
+    }

densepose/data/video/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .frame_selector import (
+    FrameSelectionStrategy,
+    RandomKFramesSelector,
+    FirstKFramesSelector,
+    LastKFramesSelector,
+    FrameTsList,
+    FrameSelector,
+)
+
+from .video_keyframe_dataset import (
+    VideoKeyframeDataset,
+    video_list_from_file,
+    list_keyframes,
+    read_keyframes,
+)

densepose/data/video/frame_selector.py

@@ -0,0 +1,89 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import random
+from collections.abc import Callable
+from enum import Enum
+from typing import Callable as TCallable
+from typing import List
+
+FrameTsList = List[int]
+FrameSelector = TCallable[[FrameTsList], FrameTsList]
+
+
+class FrameSelectionStrategy(Enum):
+    """
+    Frame selection strategy used with videos:
+     - "random_k": select k random frames
+     - "first_k": select k first frames
+     - "last_k": select k last frames
+     - "all": select all frames
+    """
+
+    # fmt: off
+    RANDOM_K = "random_k"
+    FIRST_K  = "first_k"
+    LAST_K   = "last_k"
+    ALL      = "all"
+    # fmt: on
+
+
+class RandomKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` random frames
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` random frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return random.sample(frame_tss, min(self.k, len(frame_tss)))
+
+
+class FirstKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` first frames
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` first frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return frame_tss[: self.k]
+
+
+class LastKFramesSelector(Callable):  # pyre-ignore[39]
+    """
+    Selector that retains at most `k` last frames from video data
+    """
+
+    def __init__(self, k: int):
+        self.k = k
+
+    def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
+        """
+        Select `k` last frames
+
+        Args:
+          frames_tss (List[int]): timestamps of input frames
+        Returns:
+          List[int]: timestamps of selected frames
+        """
+        return frame_tss[-self.k :]

densepose/data/video/video_keyframe_dataset.py

@@ -0,0 +1,304 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import csv
+import logging
+import numpy as np
+from typing import Any, Callable, Dict, List, Optional, Union
+import av
+import torch
+from torch.utils.data.dataset import Dataset
+
+from detectron2.utils.file_io import PathManager
+
+from ..utils import maybe_prepend_base_path
+from .frame_selector import FrameSelector, FrameTsList
+
+FrameList = List[av.frame.Frame]  # pyre-ignore[16]
+FrameTransform = Callable[[torch.Tensor], torch.Tensor]
+
+
+def list_keyframes(video_fpath: str, video_stream_idx: int = 0) -> FrameTsList:
+    """
+    Traverses all keyframes of a video file. Returns a list of keyframe
+    timestamps. Timestamps are counts in timebase units.
+
+    Args:
+       video_fpath (str): Video file path
+       video_stream_idx (int): Video stream index (default: 0)
+    Returns:
+       List[int]: list of keyframe timestaps (timestamp is a count in timebase
+           units)
+    """
+    try:
+        with PathManager.open(video_fpath, "rb") as io:
+            # pyre-fixme[16]: Module `av` has no attribute `open`.
+            container = av.open(io, mode="r")
+            stream = container.streams.video[video_stream_idx]
+            keyframes = []
+            pts = -1
+            # Note: even though we request forward seeks for keyframes, sometimes
+            # a keyframe in backwards direction is returned. We introduce tolerance
+            # as a max count of ignored backward seeks
+            tolerance_backward_seeks = 2
+            while True:
+                try:
+                    container.seek(pts + 1, backward=False, any_frame=False, stream=stream)
+                except av.AVError as e:
+                    # the exception occurs when the video length is exceeded,
+                    # we then return whatever data we've already collected
+                    logger = logging.getLogger(__name__)
+                    logger.debug(
+                        f"List keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts + 1}, AV error: {e}"
+                    )
+                    return keyframes
+                except OSError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"List keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts + 1}, OS error: {e}"
+                    )
+                    return []
+                packet = next(container.demux(video=video_stream_idx))
+                if packet.pts is not None and packet.pts <= pts:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Video file {video_fpath}, stream {video_stream_idx}: "
+                        f"bad seek for packet {pts + 1} (got packet {packet.pts}), "
+                        f"tolerance {tolerance_backward_seeks}."
+                    )
+                    tolerance_backward_seeks -= 1
+                    if tolerance_backward_seeks == 0:
+                        return []
+                    pts += 1
+                    continue
+                tolerance_backward_seeks = 2
+                pts = packet.pts
+                if pts is None:
+                    return keyframes
+                if packet.is_keyframe:
+                    keyframes.append(pts)
+            return keyframes
+    except OSError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"List keyframes: Error opening video file container {video_fpath}, " f"OS error: {e}"
+        )
+    except RuntimeError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"List keyframes: Error opening video file container {video_fpath}, "
+            f"Runtime error: {e}"
+        )
+    return []
+
+
+def read_keyframes(
+    video_fpath: str, keyframes: FrameTsList, video_stream_idx: int = 0
+) -> FrameList:  # pyre-ignore[11]
+    """
+    Reads keyframe data from a video file.
+
+    Args:
+        video_fpath (str): Video file path
+        keyframes (List[int]): List of keyframe timestamps (as counts in
+            timebase units to be used in container seek operations)
+        video_stream_idx (int): Video stream index (default: 0)
+    Returns:
+        List[Frame]: list of frames that correspond to the specified timestamps
+    """
+    try:
+        with PathManager.open(video_fpath, "rb") as io:
+            # pyre-fixme[16]: Module `av` has no attribute `open`.
+            container = av.open(io)
+            stream = container.streams.video[video_stream_idx]
+            frames = []
+            for pts in keyframes:
+                try:
+                    container.seek(pts, any_frame=False, stream=stream)
+                    frame = next(container.decode(video=0))
+                    frames.append(frame)
+                except av.AVError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}, AV error: {e}"
+                    )
+                    container.close()
+                    return frames
+                except OSError as e:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error seeking video file {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}, OS error: {e}"
+                    )
+                    container.close()
+                    return frames
+                except StopIteration:
+                    logger = logging.getLogger(__name__)
+                    logger.warning(
+                        f"Read keyframes: Error decoding frame from {video_fpath}, "
+                        f"video stream {video_stream_idx}, pts {pts}"
+                    )
+                    container.close()
+                    return frames
+
+            container.close()
+            return frames
+    except OSError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Read keyframes: Error opening video file container {video_fpath}, OS error: {e}"
+        )
+    except RuntimeError as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            f"Read keyframes: Error opening video file container {video_fpath}, Runtime error: {e}"
+        )
+    return []
+
+
+def video_list_from_file(video_list_fpath: str, base_path: Optional[str] = None):
+    """
+    Create a list of paths to video files from a text file.
+
+    Args:
+        video_list_fpath (str): path to a plain text file with the list of videos
+        base_path (str): base path for entries from the video list (default: None)
+    """
+    video_list = []
+    with PathManager.open(video_list_fpath, "r") as io:
+        for line in io:
+            video_list.append(maybe_prepend_base_path(base_path, str(line.strip())))
+    return video_list
+
+
+def read_keyframe_helper_data(fpath: str):
+    """
+    Read keyframe data from a file in CSV format: the header should contain
+    "video_id" and "keyframes" fields. Value specifications are:
+      video_id: int
+      keyframes: list(int)
+    Example of contents:
+      video_id,keyframes
+      2,"[1,11,21,31,41,51,61,71,81]"
+
+    Args:
+        fpath (str): File containing keyframe data
+
+    Return:
+        video_id_to_keyframes (dict: int -> list(int)): for a given video ID it
+          contains a list of keyframes for that video
+    """
+    video_id_to_keyframes = {}
+    try:
+        with PathManager.open(fpath, "r") as io:
+            csv_reader = csv.reader(io)
+            header = next(csv_reader)
+            video_id_idx = header.index("video_id")
+            keyframes_idx = header.index("keyframes")
+            for row in csv_reader:
+                video_id = int(row[video_id_idx])
+                assert (
+                    video_id not in video_id_to_keyframes
+                ), f"Duplicate keyframes entry for video {fpath}"
+                video_id_to_keyframes[video_id] = (
+                    [int(v) for v in row[keyframes_idx][1:-1].split(",")]
+                    if len(row[keyframes_idx]) > 2
+                    else []
+                )
+    except Exception as e:
+        logger = logging.getLogger(__name__)
+        logger.warning(f"Error reading keyframe helper data from {fpath}: {e}")
+    return video_id_to_keyframes
+
+
+class VideoKeyframeDataset(Dataset):
+    """
+    Dataset that provides keyframes for a set of videos.
+    """
+
+    _EMPTY_FRAMES = torch.empty((0, 3, 1, 1))
+
+    def __init__(
+        self,
+        video_list: List[str],
+        category_list: Union[str, List[str], None] = None,
+        frame_selector: Optional[FrameSelector] = None,
+        transform: Optional[FrameTransform] = None,
+        keyframe_helper_fpath: Optional[str] = None,
+    ):
+        """
+        Dataset constructor
+
+        Args:
+            video_list (List[str]): list of paths to video files
+            category_list (Union[str, List[str], None]): list of animal categories for each
+                video file. If it is a string, or None, this applies to all videos
+            frame_selector (Callable: KeyFrameList -> KeyFrameList):
+                selects keyframes to process, keyframes are given by
+                packet timestamps in timebase counts. If None, all keyframes
+                are selected (default: None)
+            transform (Callable: torch.Tensor -> torch.Tensor):
+                transforms a batch of RGB images (tensors of size [B, 3, H, W]),
+                returns a tensor of the same size. If None, no transform is
+                applied (default: None)
+
+        """
+        if type(category_list) is list:
+            self.category_list = category_list
+        else:
+            self.category_list = [category_list] * len(video_list)
+        assert len(video_list) == len(
+            self.category_list
+        ), "length of video and category lists must be equal"
+        self.video_list = video_list
+        self.frame_selector = frame_selector
+        self.transform = transform
+        self.keyframe_helper_data = (
+            read_keyframe_helper_data(keyframe_helper_fpath)
+            if keyframe_helper_fpath is not None
+            else None
+        )
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """
+        Gets selected keyframes from a given video
+
+        Args:
+            idx (int): video index in the video list file
+        Returns:
+            A dictionary containing two keys:
+                images (torch.Tensor): tensor of size [N, H, W, 3] or of size
+                    defined by the transform that contains keyframes data
+                categories (List[str]): categories of the frames
+        """
+        categories = [self.category_list[idx]]
+        fpath = self.video_list[idx]
+        keyframes = (
+            list_keyframes(fpath)
+            if self.keyframe_helper_data is None or idx not in self.keyframe_helper_data
+            else self.keyframe_helper_data[idx]
+        )
+        transform = self.transform
+        frame_selector = self.frame_selector
+        if not keyframes:
+            return {"images": self._EMPTY_FRAMES, "categories": []}
+        if frame_selector is not None:
+            keyframes = frame_selector(keyframes)
+        frames = read_keyframes(fpath, keyframes)
+        if not frames:
+            return {"images": self._EMPTY_FRAMES, "categories": []}
+        frames = np.stack([frame.to_rgb().to_ndarray() for frame in frames])
+        frames = torch.as_tensor(frames, device=torch.device("cpu"))
+        frames = frames[..., [2, 1, 0]]  # RGB -> BGR
+        frames = frames.permute(0, 3, 1, 2).float()  # NHWC -> NCHW
+        if transform is not None:
+            frames = transform(frames)
+        return {"images": frames, "categories": categories}
+
+    def __len__(self):
+        return len(self.video_list)

densepose/engine/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .trainer import Trainer

densepose/engine/trainer.py

@@ -0,0 +1,260 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+import logging
+import os
+from collections import OrderedDict
+from typing import List, Optional, Union
+import torch
+from torch import nn
+
+from detectron2.checkpoint import DetectionCheckpointer
+from detectron2.config import CfgNode
+from detectron2.engine import DefaultTrainer
+from detectron2.evaluation import (
+    DatasetEvaluator,
+    DatasetEvaluators,
+    inference_on_dataset,
+    print_csv_format,
+)
+from detectron2.solver.build import get_default_optimizer_params, maybe_add_gradient_clipping
+from detectron2.utils import comm
+from detectron2.utils.events import EventWriter, get_event_storage
+
+from densepose import DensePoseDatasetMapperTTA, DensePoseGeneralizedRCNNWithTTA, load_from_cfg
+from densepose.data import (
+    DatasetMapper,
+    build_combined_loader,
+    build_detection_test_loader,
+    build_detection_train_loader,
+    build_inference_based_loaders,
+    has_inference_based_loaders,
+)
+from densepose.evaluation.d2_evaluator_adapter import Detectron2COCOEvaluatorAdapter
+from densepose.evaluation.evaluator import DensePoseCOCOEvaluator, build_densepose_evaluator_storage
+from densepose.modeling.cse import Embedder
+
+
+class SampleCountingLoader:
+    def __init__(self, loader):
+        self.loader = loader
+
+    def __iter__(self):
+        it = iter(self.loader)
+        storage = get_event_storage()
+        while True:
+            try:
+                batch = next(it)
+                num_inst_per_dataset = {}
+                for data in batch:
+                    dataset_name = data["dataset"]
+                    if dataset_name not in num_inst_per_dataset:
+                        num_inst_per_dataset[dataset_name] = 0
+                    num_inst = len(data["instances"])
+                    num_inst_per_dataset[dataset_name] += num_inst
+                for dataset_name in num_inst_per_dataset:
+                    storage.put_scalar(f"batch/{dataset_name}", num_inst_per_dataset[dataset_name])
+                yield batch
+            except StopIteration:
+                break
+
+
+class SampleCountMetricPrinter(EventWriter):
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+
+    def write(self):
+        storage = get_event_storage()
+        batch_stats_strs = []
+        for key, buf in storage.histories().items():
+            if key.startswith("batch/"):
+                batch_stats_strs.append(f"{key} {buf.avg(20)}")
+        self.logger.info(", ".join(batch_stats_strs))
+
+
+class Trainer(DefaultTrainer):
+    @classmethod
+    def extract_embedder_from_model(cls, model: nn.Module) -> Optional[Embedder]:
+        if isinstance(model, nn.parallel.DistributedDataParallel):
+            model = model.module
+        if hasattr(model, "roi_heads") and hasattr(model.roi_heads, "embedder"):
+            return model.roi_heads.embedder
+        return None
+
+    # TODO: the only reason to copy the base class code here is to pass the embedder from
+    # the model to the evaluator; that should be refactored to avoid unnecessary copy-pasting
+    @classmethod
+    def test(
+        cls,
+        cfg: CfgNode,
+        model: nn.Module,
+        evaluators: Optional[Union[DatasetEvaluator, List[DatasetEvaluator]]] = None,
+    ):
+        """
+        Args:
+            cfg (CfgNode):
+            model (nn.Module):
+            evaluators (DatasetEvaluator, list[DatasetEvaluator] or None): if None, will call
+                :meth:`build_evaluator`. Otherwise, must have the same length as
+                ``cfg.DATASETS.TEST``.
+
+        Returns:
+            dict: a dict of result metrics
+        """
+        logger = logging.getLogger(__name__)
+        if isinstance(evaluators, DatasetEvaluator):
+            evaluators = [evaluators]
+        if evaluators is not None:
+            assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
+                len(cfg.DATASETS.TEST), len(evaluators)
+            )
+
+        results = OrderedDict()
+        for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
+            data_loader = cls.build_test_loader(cfg, dataset_name)
+            # When evaluators are passed in as arguments,
+            # implicitly assume that evaluators can be created before data_loader.
+            if evaluators is not None:
+                evaluator = evaluators[idx]
+            else:
+                try:
+                    embedder = cls.extract_embedder_from_model(model)
+                    evaluator = cls.build_evaluator(cfg, dataset_name, embedder=embedder)
+                except NotImplementedError:
+                    logger.warn(
+                        "No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
+                        "or implement its `build_evaluator` method."
+                    )
+                    results[dataset_name] = {}
+                    continue
+            if cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE or comm.is_main_process():
+                results_i = inference_on_dataset(model, data_loader, evaluator)
+            else:
+                results_i = {}
+            results[dataset_name] = results_i
+            if comm.is_main_process():
+                assert isinstance(
+                    results_i, dict
+                ), "Evaluator must return a dict on the main process. Got {} instead.".format(
+                    results_i
+                )
+                logger.info("Evaluation results for {} in csv format:".format(dataset_name))
+                print_csv_format(results_i)
+
+        if len(results) == 1:
+            results = list(results.values())[0]
+        return results
+
+    @classmethod
+    def build_evaluator(
+        cls,
+        cfg: CfgNode,
+        dataset_name: str,
+        output_folder: Optional[str] = None,
+        embedder: Optional[Embedder] = None,
+    ) -> DatasetEvaluators:
+        if output_folder is None:
+            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
+        evaluators = []
+        distributed = cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE
+        # Note: we currently use COCO evaluator for both COCO and LVIS datasets
+        # to have compatible metrics. LVIS bbox evaluator could also be used
+        # with an adapter to properly handle filtered / mapped categories
+        # evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
+        # if evaluator_type == "coco":
+        #     evaluators.append(COCOEvaluator(dataset_name, output_dir=output_folder))
+        # elif evaluator_type == "lvis":
+        #     evaluators.append(LVISEvaluator(dataset_name, output_dir=output_folder))
+        evaluators.append(
+            Detectron2COCOEvaluatorAdapter(
+                dataset_name, output_dir=output_folder, distributed=distributed
+            )
+        )
+        if cfg.MODEL.DENSEPOSE_ON:
+            storage = build_densepose_evaluator_storage(cfg, output_folder)
+            evaluators.append(
+                DensePoseCOCOEvaluator(
+                    dataset_name,
+                    distributed,
+                    output_folder,
+                    evaluator_type=cfg.DENSEPOSE_EVALUATION.TYPE,
+                    min_iou_threshold=cfg.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD,
+                    storage=storage,
+                    embedder=embedder,
+                    should_evaluate_mesh_alignment=cfg.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT,
+                    mesh_alignment_mesh_names=cfg.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES,
+                )
+            )
+        return DatasetEvaluators(evaluators)
+
+    @classmethod
+    def build_optimizer(cls, cfg: CfgNode, model: nn.Module):
+        params = get_default_optimizer_params(
+            model,
+            base_lr=cfg.SOLVER.BASE_LR,
+            weight_decay_norm=cfg.SOLVER.WEIGHT_DECAY_NORM,
+            bias_lr_factor=cfg.SOLVER.BIAS_LR_FACTOR,
+            weight_decay_bias=cfg.SOLVER.WEIGHT_DECAY_BIAS,
+            overrides={
+                "features": {
+                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR,
+                },
+                "embeddings": {
+                    "lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR,
+                },
+            },
+        )
+        optimizer = torch.optim.SGD(
+            params,
+            cfg.SOLVER.BASE_LR,
+            momentum=cfg.SOLVER.MOMENTUM,
+            nesterov=cfg.SOLVER.NESTEROV,
+            weight_decay=cfg.SOLVER.WEIGHT_DECAY,
+        )
+        # pyre-fixme[6]: For 2nd param expected `Type[Optimizer]` but got `SGD`.
+        return maybe_add_gradient_clipping(cfg, optimizer)
+
+    @classmethod
+    def build_test_loader(cls, cfg: CfgNode, dataset_name):
+        return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
+
+    @classmethod
+    def build_train_loader(cls, cfg: CfgNode):
+        data_loader = build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
+        if not has_inference_based_loaders(cfg):
+            return data_loader
+        model = cls.build_model(cfg)
+        model.to(cfg.BOOTSTRAP_MODEL.DEVICE)
+        DetectionCheckpointer(model).resume_or_load(cfg.BOOTSTRAP_MODEL.WEIGHTS, resume=False)
+        inference_based_loaders, ratios = build_inference_based_loaders(cfg, model)
+        loaders = [data_loader] + inference_based_loaders
+        ratios = [1.0] + ratios
+        combined_data_loader = build_combined_loader(cfg, loaders, ratios)
+        sample_counting_loader = SampleCountingLoader(combined_data_loader)
+        return sample_counting_loader
+
+    def build_writers(self):
+        writers = super().build_writers()
+        writers.append(SampleCountMetricPrinter())
+        return writers
+
+    @classmethod
+    def test_with_TTA(cls, cfg: CfgNode, model):
+        logger = logging.getLogger("detectron2.trainer")
+        # In the end of training, run an evaluation with TTA
+        # Only support some R-CNN models.
+        logger.info("Running inference with test-time augmentation ...")
+        transform_data = load_from_cfg(cfg)
+        model = DensePoseGeneralizedRCNNWithTTA(
+            cfg, model, transform_data, DensePoseDatasetMapperTTA(cfg)
+        )
+        evaluators = [
+            cls.build_evaluator(
+                cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
+            )
+            for name in cfg.DATASETS.TEST
+        ]
+        res = cls.test(cfg, model, evaluators)  # pyre-ignore[6]
+        res = OrderedDict({k + "_TTA": v for k, v in res.items()})
+        return res

densepose/evaluation/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from .evaluator import DensePoseCOCOEvaluator

densepose/evaluation/d2_evaluator_adapter.py

@@ -0,0 +1,52 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+from detectron2.data.catalog import Metadata
+from detectron2.evaluation import COCOEvaluator
+
+from densepose.data.datasets.coco import (
+    get_contiguous_id_to_category_id_map,
+    maybe_filter_categories_cocoapi,
+)
+
+
+def _maybe_add_iscrowd_annotations(cocoapi) -> None:
+    for ann in cocoapi.dataset["annotations"]:
+        if "iscrowd" not in ann:
+            ann["iscrowd"] = 0
+
+
+class Detectron2COCOEvaluatorAdapter(COCOEvaluator):
+    def __init__(
+        self,
+        dataset_name,
+        output_dir=None,
+        distributed=True,
+    ):
+        super().__init__(dataset_name, output_dir=output_dir, distributed=distributed)
+        maybe_filter_categories_cocoapi(dataset_name, self._coco_api)
+        _maybe_add_iscrowd_annotations(self._coco_api)
+        # substitute category metadata to account for categories
+        # that are mapped to the same contiguous id
+        if hasattr(self._metadata, "thing_dataset_id_to_contiguous_id"):
+            self._maybe_substitute_metadata()
+
+    def _maybe_substitute_metadata(self):
+        cont_id_2_cat_id = get_contiguous_id_to_category_id_map(self._metadata)
+        cat_id_2_cont_id = self._metadata.thing_dataset_id_to_contiguous_id
+        if len(cont_id_2_cat_id) == len(cat_id_2_cont_id):
+            return
+
+        cat_id_2_cont_id_injective = {}
+        for cat_id, cont_id in cat_id_2_cont_id.items():
+            if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
+                cat_id_2_cont_id_injective[cat_id] = cont_id
+
+        metadata_new = Metadata(name=self._metadata.name)
+        for key, value in self._metadata.__dict__.items():
+            if key == "thing_dataset_id_to_contiguous_id":
+                setattr(metadata_new, key, cat_id_2_cont_id_injective)
+            else:
+                setattr(metadata_new, key, value)
+        self._metadata = metadata_new

densepose/evaluation/densepose_coco_evaluation.py

@@ -0,0 +1,1305 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# This is a modified version of cocoeval.py where we also have the densepose evaluation.
+
+# pyre-unsafe
+
+__author__ = "tsungyi"
+
+import copy
+import datetime
+import logging
+import numpy as np
+import pickle
+import time
+from collections import defaultdict
+from enum import Enum
+from typing import Any, Dict, Tuple
+import scipy.spatial.distance as ssd
+import torch
+import torch.nn.functional as F
+from pycocotools import mask as maskUtils
+from scipy.io import loadmat
+from scipy.ndimage import zoom as spzoom
+
+from detectron2.utils.file_io import PathManager
+
+from densepose.converters.chart_output_to_chart_result import resample_uv_tensors_to_bbox
+from densepose.converters.segm_to_mask import (
+    resample_coarse_segm_tensor_to_bbox,
+    resample_fine_and_coarse_segm_tensors_to_bbox,
+)
+from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
+from densepose.structures import DensePoseDataRelative
+from densepose.structures.mesh import create_mesh
+
+logger = logging.getLogger(__name__)
+
+
+class DensePoseEvalMode(str, Enum):
+    # use both masks and geodesic distances (GPS * IOU) to compute scores
+    GPSM = "gpsm"
+    # use only geodesic distances (GPS)  to compute scores
+    GPS = "gps"
+    # use only masks (IOU) to compute scores
+    IOU = "iou"
+
+
+class DensePoseDataMode(str, Enum):
+    # use estimated IUV data (default mode)
+    IUV_DT = "iuvdt"
+    # use ground truth IUV data
+    IUV_GT = "iuvgt"
+    # use ground truth labels I and set UV to 0
+    I_GT_UV_0 = "igtuv0"
+    # use ground truth labels I and estimated UV coordinates
+    I_GT_UV_DT = "igtuvdt"
+    # use estimated labels I and set UV to 0
+    I_DT_UV_0 = "idtuv0"
+
+
+class DensePoseCocoEval:
+    # Interface for evaluating detection on the Microsoft COCO dataset.
+    #
+    # The usage for CocoEval is as follows:
+    #  cocoGt=..., cocoDt=...       # load dataset and results
+    #  E = CocoEval(cocoGt,cocoDt); # initialize CocoEval object
+    #  E.params.recThrs = ...;      # set parameters as desired
+    #  E.evaluate();                # run per image evaluation
+    #  E.accumulate();              # accumulate per image results
+    #  E.summarize();               # display summary metrics of results
+    # For example usage see evalDemo.m and http://mscoco.org/.
+    #
+    # The evaluation parameters are as follows (defaults in brackets):
+    #  imgIds     - [all] N img ids to use for evaluation
+    #  catIds     - [all] K cat ids to use for evaluation
+    #  iouThrs    - [.5:.05:.95] T=10 IoU thresholds for evaluation
+    #  recThrs    - [0:.01:1] R=101 recall thresholds for evaluation
+    #  areaRng    - [...] A=4 object area ranges for evaluation
+    #  maxDets    - [1 10 100] M=3 thresholds on max detections per image
+    #  iouType    - ['segm'] set iouType to 'segm', 'bbox', 'keypoints' or 'densepose'
+    #  iouType replaced the now DEPRECATED useSegm parameter.
+    #  useCats    - [1] if true use category labels for evaluation
+    # Note: if useCats=0 category labels are ignored as in proposal scoring.
+    # Note: multiple areaRngs [Ax2] and maxDets [Mx1] can be specified.
+    #
+    # evaluate(): evaluates detections on every image and every category and
+    # concats the results into the "evalImgs" with fields:
+    #  dtIds      - [1xD] id for each of the D detections (dt)
+    #  gtIds      - [1xG] id for each of the G ground truths (gt)
+    #  dtMatches  - [TxD] matching gt id at each IoU or 0
+    #  gtMatches  - [TxG] matching dt id at each IoU or 0
+    #  dtScores   - [1xD] confidence of each dt
+    #  gtIgnore   - [1xG] ignore flag for each gt
+    #  dtIgnore   - [TxD] ignore flag for each dt at each IoU
+    #
+    # accumulate(): accumulates the per-image, per-category evaluation
+    # results in "evalImgs" into the dictionary "eval" with fields:
+    #  params     - parameters used for evaluation
+    #  date       - date evaluation was performed
+    #  counts     - [T,R,K,A,M] parameter dimensions (see above)
+    #  precision  - [TxRxKxAxM] precision for every evaluation setting
+    #  recall     - [TxKxAxM] max recall for every evaluation setting
+    # Note: precision and recall==-1 for settings with no gt objects.
+    #
+    # See also coco, mask, pycocoDemo, pycocoEvalDemo
+    #
+    # Microsoft COCO Toolbox.      version 2.0
+    # Data, paper, and tutorials available at:  http://mscoco.org/
+    # Code written by Piotr Dollar and Tsung-Yi Lin, 2015.
+    # Licensed under the Simplified BSD License [see coco/license.txt]
+    def __init__(
+        self,
+        cocoGt=None,
+        cocoDt=None,
+        iouType: str = "densepose",
+        multi_storage=None,
+        embedder=None,
+        dpEvalMode: DensePoseEvalMode = DensePoseEvalMode.GPS,
+        dpDataMode: DensePoseDataMode = DensePoseDataMode.IUV_DT,
+    ):
+        """
+        Initialize CocoEval using coco APIs for gt and dt
+        :param cocoGt: coco object with ground truth annotations
+        :param cocoDt: coco object with detection results
+        :return: None
+        """
+        self.cocoGt = cocoGt  # ground truth COCO API
+        self.cocoDt = cocoDt  # detections COCO API
+        self.multi_storage = multi_storage
+        self.embedder = embedder
+        self._dpEvalMode = dpEvalMode
+        self._dpDataMode = dpDataMode
+        self.evalImgs = defaultdict(list)  # per-image per-category eval results [KxAxI]
+        self.eval = {}  # accumulated evaluation results
+        self._gts = defaultdict(list)  # gt for evaluation
+        self._dts = defaultdict(list)  # dt for evaluation
+        self.params = Params(iouType=iouType)  # parameters
+        self._paramsEval = {}  # parameters for evaluation
+        self.stats = []  # result summarization
+        self.ious = {}  # ious between all gts and dts
+        if cocoGt is not None:
+            self.params.imgIds = sorted(cocoGt.getImgIds())
+            self.params.catIds = sorted(cocoGt.getCatIds())
+        self.ignoreThrBB = 0.7
+        self.ignoreThrUV = 0.9
+
+    def _loadGEval(self):
+        smpl_subdiv_fpath = PathManager.get_local_path(
+            "https://dl.fbaipublicfiles.com/densepose/data/SMPL_subdiv.mat"
+        )
+        pdist_transform_fpath = PathManager.get_local_path(
+            "https://dl.fbaipublicfiles.com/densepose/data/SMPL_SUBDIV_TRANSFORM.mat"
+        )
+        pdist_matrix_fpath = PathManager.get_local_path(
+            "https://dl.fbaipublicfiles.com/densepose/data/Pdist_matrix.pkl", timeout_sec=120
+        )
+        SMPL_subdiv = loadmat(smpl_subdiv_fpath)
+        self.PDIST_transform = loadmat(pdist_transform_fpath)
+        self.PDIST_transform = self.PDIST_transform["index"].squeeze()
+        UV = np.array([SMPL_subdiv["U_subdiv"], SMPL_subdiv["V_subdiv"]]).squeeze()
+        ClosestVertInds = np.arange(UV.shape[1]) + 1
+        self.Part_UVs = []
+        self.Part_ClosestVertInds = []
+        for i in np.arange(24):
+            self.Part_UVs.append(UV[:, SMPL_subdiv["Part_ID_subdiv"].squeeze() == (i + 1)])
+            self.Part_ClosestVertInds.append(
+                ClosestVertInds[SMPL_subdiv["Part_ID_subdiv"].squeeze() == (i + 1)]
+            )
+
+        with open(pdist_matrix_fpath, "rb") as hFile:
+            arrays = pickle.load(hFile, encoding="latin1")
+        self.Pdist_matrix = arrays["Pdist_matrix"]
+        self.Part_ids = np.array(SMPL_subdiv["Part_ID_subdiv"].squeeze())
+        # Mean geodesic distances for parts.
+        self.Mean_Distances = np.array([0, 0.351, 0.107, 0.126, 0.237, 0.173, 0.142, 0.128, 0.150])
+        # Coarse Part labels.
+        self.CoarseParts = np.array(
+            [0, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7, 7, 8, 8]
+        )
+
+    def _prepare(self):
+        """
+        Prepare ._gts and ._dts for evaluation based on params
+        :return: None
+        """
+
+        def _toMask(anns, coco):
+            # modify ann['segmentation'] by reference
+            for ann in anns:
+                # safeguard for invalid segmentation annotation;
+                # annotations containing empty lists exist in the posetrack
+                # dataset. This is not a correct segmentation annotation
+                # in terms of COCO format; we need to deal with it somehow
+                segm = ann["segmentation"]
+                if type(segm) is list and len(segm) == 0:
+                    ann["segmentation"] = None
+                    continue
+                rle = coco.annToRLE(ann)
+                ann["segmentation"] = rle
+
+        def _getIgnoreRegion(iid, coco):
+            img = coco.imgs[iid]
+
+            if "ignore_regions_x" not in img.keys():
+                return None
+
+            if len(img["ignore_regions_x"]) == 0:
+                return None
+
+            rgns_merged = [
+                [v for xy in zip(region_x, region_y) for v in xy]
+                for region_x, region_y in zip(img["ignore_regions_x"], img["ignore_regions_y"])
+            ]
+            rles = maskUtils.frPyObjects(rgns_merged, img["height"], img["width"])
+            rle = maskUtils.merge(rles)
+            return maskUtils.decode(rle)
+
+        def _checkIgnore(dt, iregion):
+            if iregion is None:
+                return True
+
+            bb = np.array(dt["bbox"]).astype(int)
+            x1, y1, x2, y2 = bb[0], bb[1], bb[0] + bb[2], bb[1] + bb[3]
+            x2 = min([x2, iregion.shape[1]])
+            y2 = min([y2, iregion.shape[0]])
+
+            if bb[2] * bb[3] == 0:
+                return False
+
+            crop_iregion = iregion[y1:y2, x1:x2]
+
+            if crop_iregion.sum() == 0:
+                return True
+
+            if "densepose" not in dt.keys():  # filtering boxes
+                return crop_iregion.sum() / bb[2] / bb[3] < self.ignoreThrBB
+
+            # filtering UVs
+            ignoremask = np.require(crop_iregion, requirements=["F"])
+            mask = self._extract_mask(dt)
+            uvmask = np.require(np.asarray(mask > 0), dtype=np.uint8, requirements=["F"])
+            uvmask_ = maskUtils.encode(uvmask)
+            ignoremask_ = maskUtils.encode(ignoremask)
+            uviou = maskUtils.iou([uvmask_], [ignoremask_], [1])[0]
+            return uviou < self.ignoreThrUV
+
+        p = self.params
+
+        if p.useCats:
+            gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds))
+            dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds))
+        else:
+            gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds))
+            dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds))
+
+        imns = self.cocoGt.loadImgs(p.imgIds)
+        self.size_mapping = {}
+        for im in imns:
+            self.size_mapping[im["id"]] = [im["height"], im["width"]]
+
+        # if iouType == 'uv', add point gt annotations
+        if p.iouType == "densepose":
+            self._loadGEval()
+
+        # convert ground truth to mask if iouType == 'segm'
+        if p.iouType == "segm":
+            _toMask(gts, self.cocoGt)
+            _toMask(dts, self.cocoDt)
+
+        # set ignore flag
+        for gt in gts:
+            gt["ignore"] = gt["ignore"] if "ignore" in gt else 0
+            gt["ignore"] = "iscrowd" in gt and gt["iscrowd"]
+            if p.iouType == "keypoints":
+                gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"]
+            if p.iouType == "densepose":
+                gt["ignore"] = ("dp_x" in gt) == 0
+            if p.iouType == "segm":
+                gt["ignore"] = gt["segmentation"] is None
+
+        self._gts = defaultdict(list)  # gt for evaluation
+        self._dts = defaultdict(list)  # dt for evaluation
+        self._igrgns = defaultdict(list)
+
+        for gt in gts:
+            iid = gt["image_id"]
+            if iid not in self._igrgns.keys():
+                self._igrgns[iid] = _getIgnoreRegion(iid, self.cocoGt)
+            if _checkIgnore(gt, self._igrgns[iid]):
+                self._gts[iid, gt["category_id"]].append(gt)
+        for dt in dts:
+            iid = dt["image_id"]
+            if (iid not in self._igrgns) or _checkIgnore(dt, self._igrgns[iid]):
+                self._dts[iid, dt["category_id"]].append(dt)
+
+        self.evalImgs = defaultdict(list)  # per-image per-category evaluation results
+        self.eval = {}  # accumulated evaluation results
+
+    def evaluate(self):
+        """
+        Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
+        :return: None
+        """
+        tic = time.time()
+        logger.info("Running per image DensePose evaluation... {}".format(self.params.iouType))
+        p = self.params
+        # add backward compatibility if useSegm is specified in params
+        if p.useSegm is not None:
+            p.iouType = "segm" if p.useSegm == 1 else "bbox"
+            logger.info("useSegm (deprecated) is not None. Running DensePose evaluation")
+        p.imgIds = list(np.unique(p.imgIds))
+        if p.useCats:
+            p.catIds = list(np.unique(p.catIds))
+        p.maxDets = sorted(p.maxDets)
+        self.params = p
+
+        self._prepare()
+        # loop through images, area range, max detection number
+        catIds = p.catIds if p.useCats else [-1]
+
+        if p.iouType in ["segm", "bbox"]:
+            computeIoU = self.computeIoU
+        elif p.iouType == "keypoints":
+            computeIoU = self.computeOks
+        elif p.iouType == "densepose":
+            computeIoU = self.computeOgps
+            if self._dpEvalMode in {DensePoseEvalMode.GPSM, DensePoseEvalMode.IOU}:
+                self.real_ious = {
+                    (imgId, catId): self.computeDPIoU(imgId, catId)
+                    for imgId in p.imgIds
+                    for catId in catIds
+                }
+
+        self.ious = {
+            (imgId, catId): computeIoU(imgId, catId) for imgId in p.imgIds for catId in catIds
+        }
+
+        evaluateImg = self.evaluateImg
+        maxDet = p.maxDets[-1]
+        self.evalImgs = [
+            evaluateImg(imgId, catId, areaRng, maxDet)
+            for catId in catIds
+            for areaRng in p.areaRng
+            for imgId in p.imgIds
+        ]
+        self._paramsEval = copy.deepcopy(self.params)
+        toc = time.time()
+        logger.info("DensePose evaluation DONE (t={:0.2f}s).".format(toc - tic))
+
+    def getDensePoseMask(self, polys):
+        maskGen = np.zeros([256, 256])
+        stop = min(len(polys) + 1, 15)
+        for i in range(1, stop):
+            if polys[i - 1]:
+                currentMask = maskUtils.decode(polys[i - 1])
+                maskGen[currentMask > 0] = i
+        return maskGen
+
+    def _generate_rlemask_on_image(self, mask, imgId, data):
+        bbox_xywh = np.array(data["bbox"])
+        x, y, w, h = bbox_xywh
+        im_h, im_w = self.size_mapping[imgId]
+        im_mask = np.zeros((im_h, im_w), dtype=np.uint8)
+        if mask is not None:
+            x0 = max(int(x), 0)
+            x1 = min(int(x + w), im_w, int(x) + mask.shape[1])
+            y0 = max(int(y), 0)
+            y1 = min(int(y + h), im_h, int(y) + mask.shape[0])
+            y = int(y)
+            x = int(x)
+            im_mask[y0:y1, x0:x1] = mask[y0 - y : y1 - y, x0 - x : x1 - x]
+        im_mask = np.require(np.asarray(im_mask > 0), dtype=np.uint8, requirements=["F"])
+        rle_mask = maskUtils.encode(np.array(im_mask[:, :, np.newaxis], order="F"))[0]
+        return rle_mask
+
+    def computeDPIoU(self, imgId, catId):
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        if len(gt) == 0 and len(dt) == 0:
+            return []
+        inds = np.argsort([-d["score"] for d in dt], kind="mergesort")
+        dt = [dt[i] for i in inds]
+        if len(dt) > p.maxDets[-1]:
+            dt = dt[0 : p.maxDets[-1]]
+
+        gtmasks = []
+        for g in gt:
+            if DensePoseDataRelative.S_KEY in g:
+                # convert DensePose mask to a binary mask
+                mask = np.minimum(self.getDensePoseMask(g[DensePoseDataRelative.S_KEY]), 1.0)
+                _, _, w, h = g["bbox"]
+                scale_x = float(max(w, 1)) / mask.shape[1]
+                scale_y = float(max(h, 1)) / mask.shape[0]
+                mask = spzoom(mask, (scale_y, scale_x), order=1, prefilter=False)
+                mask = np.array(mask > 0.5, dtype=np.uint8)
+                rle_mask = self._generate_rlemask_on_image(mask, imgId, g)
+            elif "segmentation" in g:
+                segmentation = g["segmentation"]
+                if isinstance(segmentation, list) and segmentation:
+                    # polygons
+                    im_h, im_w = self.size_mapping[imgId]
+                    rles = maskUtils.frPyObjects(segmentation, im_h, im_w)
+                    rle_mask = maskUtils.merge(rles)
+                elif isinstance(segmentation, dict):
+                    if isinstance(segmentation["counts"], list):
+                        # uncompressed RLE
+                        im_h, im_w = self.size_mapping[imgId]
+                        rle_mask = maskUtils.frPyObjects(segmentation, im_h, im_w)
+                    else:
+                        # compressed RLE
+                        rle_mask = segmentation
+                else:
+                    rle_mask = self._generate_rlemask_on_image(None, imgId, g)
+            else:
+                rle_mask = self._generate_rlemask_on_image(None, imgId, g)
+            gtmasks.append(rle_mask)
+
+        dtmasks = []
+        for d in dt:
+            mask = self._extract_mask(d)
+            mask = np.require(np.asarray(mask > 0), dtype=np.uint8, requirements=["F"])
+            rle_mask = self._generate_rlemask_on_image(mask, imgId, d)
+            dtmasks.append(rle_mask)
+
+        # compute iou between each dt and gt region
+        iscrowd = [int(o.get("iscrowd", 0)) for o in gt]
+        iousDP = maskUtils.iou(dtmasks, gtmasks, iscrowd)
+        return iousDP
+
+    def computeIoU(self, imgId, catId):
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        if len(gt) == 0 and len(dt) == 0:
+            return []
+        inds = np.argsort([-d["score"] for d in dt], kind="mergesort")
+        dt = [dt[i] for i in inds]
+        if len(dt) > p.maxDets[-1]:
+            dt = dt[0 : p.maxDets[-1]]
+
+        if p.iouType == "segm":
+            g = [g["segmentation"] for g in gt if g["segmentation"] is not None]
+            d = [d["segmentation"] for d in dt if d["segmentation"] is not None]
+        elif p.iouType == "bbox":
+            g = [g["bbox"] for g in gt]
+            d = [d["bbox"] for d in dt]
+        else:
+            raise Exception("unknown iouType for iou computation")
+
+        # compute iou between each dt and gt region
+        iscrowd = [int(o.get("iscrowd", 0)) for o in gt]
+        ious = maskUtils.iou(d, g, iscrowd)
+        return ious
+
+    def computeOks(self, imgId, catId):
+        p = self.params
+        # dimension here should be Nxm
+        gts = self._gts[imgId, catId]
+        dts = self._dts[imgId, catId]
+        inds = np.argsort([-d["score"] for d in dts], kind="mergesort")
+        dts = [dts[i] for i in inds]
+        if len(dts) > p.maxDets[-1]:
+            dts = dts[0 : p.maxDets[-1]]
+        # if len(gts) == 0 and len(dts) == 0:
+        if len(gts) == 0 or len(dts) == 0:
+            return []
+        ious = np.zeros((len(dts), len(gts)))
+        sigmas = (
+            np.array(
+                [
+                    0.26,
+                    0.25,
+                    0.25,
+                    0.35,
+                    0.35,
+                    0.79,
+                    0.79,
+                    0.72,
+                    0.72,
+                    0.62,
+                    0.62,
+                    1.07,
+                    1.07,
+                    0.87,
+                    0.87,
+                    0.89,
+                    0.89,
+                ]
+            )
+            / 10.0
+        )
+        vars = (sigmas * 2) ** 2
+        k = len(sigmas)
+        # compute oks between each detection and ground truth object
+        for j, gt in enumerate(gts):
+            # create bounds for ignore regions(double the gt bbox)
+            g = np.array(gt["keypoints"])
+            xg = g[0::3]
+            yg = g[1::3]
+            vg = g[2::3]
+            k1 = np.count_nonzero(vg > 0)
+            bb = gt["bbox"]
+            x0 = bb[0] - bb[2]
+            x1 = bb[0] + bb[2] * 2
+            y0 = bb[1] - bb[3]
+            y1 = bb[1] + bb[3] * 2
+            for i, dt in enumerate(dts):
+                d = np.array(dt["keypoints"])
+                xd = d[0::3]
+                yd = d[1::3]
+                if k1 > 0:
+                    # measure the per-keypoint distance if keypoints visible
+                    dx = xd - xg
+                    dy = yd - yg
+                else:
+                    # measure minimum distance to keypoints in (x0,y0) & (x1,y1)
+                    z = np.zeros(k)
+                    dx = np.max((z, x0 - xd), axis=0) + np.max((z, xd - x1), axis=0)
+                    dy = np.max((z, y0 - yd), axis=0) + np.max((z, yd - y1), axis=0)
+                e = (dx**2 + dy**2) / vars / (gt["area"] + np.spacing(1)) / 2
+                if k1 > 0:
+                    e = e[vg > 0]
+                ious[i, j] = np.sum(np.exp(-e)) / e.shape[0]
+        return ious
+
+    def _extract_mask(self, dt: Dict[str, Any]) -> np.ndarray:
+        if "densepose" in dt:
+            densepose_results_quantized = dt["densepose"]
+            return densepose_results_quantized.labels_uv_uint8[0].numpy()
+        elif "cse_mask" in dt:
+            return dt["cse_mask"]
+        elif "coarse_segm" in dt:
+            dy = max(int(dt["bbox"][3]), 1)
+            dx = max(int(dt["bbox"][2]), 1)
+            return (
+                F.interpolate(
+                    dt["coarse_segm"].unsqueeze(0),
+                    (dy, dx),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                .squeeze(0)
+                .argmax(0)
+                .numpy()
+                .astype(np.uint8)
+            )
+        elif "record_id" in dt:
+            assert (
+                self.multi_storage is not None
+            ), f"Storage record id encountered in a detection {dt}, but no storage provided!"
+            record = self.multi_storage.get(dt["rank"], dt["record_id"])
+            coarse_segm = record["coarse_segm"]
+            dy = max(int(dt["bbox"][3]), 1)
+            dx = max(int(dt["bbox"][2]), 1)
+            return (
+                F.interpolate(
+                    coarse_segm.unsqueeze(0),
+                    (dy, dx),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                .squeeze(0)
+                .argmax(0)
+                .numpy()
+                .astype(np.uint8)
+            )
+        else:
+            raise Exception(f"No mask data in the detection: {dt}")
+        raise ValueError('The prediction dict needs to contain either "densepose" or "cse_mask"')
+
+    def _extract_iuv(
+        self, densepose_data: np.ndarray, py: np.ndarray, px: np.ndarray, gt: Dict[str, Any]
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Extract arrays of I, U and V values at given points as numpy arrays
+        given the data mode stored in self._dpDataMode
+        """
+        if self._dpDataMode == DensePoseDataMode.IUV_DT:
+            # estimated labels and UV (default)
+            ipoints = densepose_data[0, py, px]
+            upoints = densepose_data[1, py, px] / 255.0  # convert from uint8 by /255.
+            vpoints = densepose_data[2, py, px] / 255.0
+        elif self._dpDataMode == DensePoseDataMode.IUV_GT:
+            # ground truth
+            ipoints = np.array(gt["dp_I"])
+            upoints = np.array(gt["dp_U"])
+            vpoints = np.array(gt["dp_V"])
+        elif self._dpDataMode == DensePoseDataMode.I_GT_UV_0:
+            # ground truth labels, UV = 0
+            ipoints = np.array(gt["dp_I"])
+            upoints = upoints * 0.0
+            vpoints = vpoints * 0.0
+        elif self._dpDataMode == DensePoseDataMode.I_GT_UV_DT:
+            # ground truth labels, estimated UV
+            ipoints = np.array(gt["dp_I"])
+            upoints = densepose_data[1, py, px] / 255.0  # convert from uint8 by /255.
+            vpoints = densepose_data[2, py, px] / 255.0
+        elif self._dpDataMode == DensePoseDataMode.I_DT_UV_0:
+            # estimated labels, UV = 0
+            ipoints = densepose_data[0, py, px]
+            upoints = upoints * 0.0
+            vpoints = vpoints * 0.0
+        else:
+            raise ValueError(f"Unknown data mode: {self._dpDataMode}")
+        return ipoints, upoints, vpoints
+
+    def computeOgps_single_pair(self, dt, gt, py, px, pt_mask):
+        if "densepose" in dt:
+            ipoints, upoints, vpoints = self.extract_iuv_from_quantized(dt, gt, py, px, pt_mask)
+            return self.computeOgps_single_pair_iuv(dt, gt, ipoints, upoints, vpoints)
+        elif "u" in dt:
+            ipoints, upoints, vpoints = self.extract_iuv_from_raw(dt, gt, py, px, pt_mask)
+            return self.computeOgps_single_pair_iuv(dt, gt, ipoints, upoints, vpoints)
+        elif "record_id" in dt:
+            assert (
+                self.multi_storage is not None
+            ), f"Storage record id encountered in detection {dt}, but no storage provided!"
+            record = self.multi_storage.get(dt["rank"], dt["record_id"])
+            record["bbox"] = dt["bbox"]
+            if "u" in record:
+                ipoints, upoints, vpoints = self.extract_iuv_from_raw(record, gt, py, px, pt_mask)
+                return self.computeOgps_single_pair_iuv(dt, gt, ipoints, upoints, vpoints)
+            elif "embedding" in record:
+                return self.computeOgps_single_pair_cse(
+                    dt,
+                    gt,
+                    py,
+                    px,
+                    pt_mask,
+                    record["coarse_segm"],
+                    record["embedding"],
+                    record["bbox"],
+                )
+            else:
+                raise Exception(f"Unknown record format: {record}")
+        elif "embedding" in dt:
+            return self.computeOgps_single_pair_cse(
+                dt, gt, py, px, pt_mask, dt["coarse_segm"], dt["embedding"], dt["bbox"]
+            )
+        raise Exception(f"Unknown detection format: {dt}")
+
+    def extract_iuv_from_quantized(self, dt, gt, py, px, pt_mask):
+        densepose_results_quantized = dt["densepose"]
+        ipoints, upoints, vpoints = self._extract_iuv(
+            densepose_results_quantized.labels_uv_uint8.numpy(), py, px, gt
+        )
+        ipoints[pt_mask == -1] = 0
+        return ipoints, upoints, vpoints
+
+    def extract_iuv_from_raw(self, dt, gt, py, px, pt_mask):
+        labels_dt = resample_fine_and_coarse_segm_tensors_to_bbox(
+            dt["fine_segm"].unsqueeze(0),
+            dt["coarse_segm"].unsqueeze(0),
+            dt["bbox"],
+        )
+        uv = resample_uv_tensors_to_bbox(
+            dt["u"].unsqueeze(0), dt["v"].unsqueeze(0), labels_dt.squeeze(0), dt["bbox"]
+        )
+        labels_uv_uint8 = torch.cat((labels_dt.byte(), (uv * 255).clamp(0, 255).byte()))
+        ipoints, upoints, vpoints = self._extract_iuv(labels_uv_uint8.numpy(), py, px, gt)
+        ipoints[pt_mask == -1] = 0
+        return ipoints, upoints, vpoints
+
+    def computeOgps_single_pair_iuv(self, dt, gt, ipoints, upoints, vpoints):
+        cVertsGT, ClosestVertsGTTransformed = self.findAllClosestVertsGT(gt)
+        cVerts = self.findAllClosestVertsUV(upoints, vpoints, ipoints)
+        # Get pairwise geodesic distances between gt and estimated mesh points.
+        dist = self.getDistancesUV(ClosestVertsGTTransformed, cVerts)
+        # Compute the Ogps measure.
+        # Find the mean geodesic normalization distance for
+        # each GT point, based on which part it is on.
+        Current_Mean_Distances = self.Mean_Distances[
+            self.CoarseParts[self.Part_ids[cVertsGT[cVertsGT > 0].astype(int) - 1]]
+        ]
+        return dist, Current_Mean_Distances
+
+    def computeOgps_single_pair_cse(
+        self, dt, gt, py, px, pt_mask, coarse_segm, embedding, bbox_xywh_abs
+    ):
+        # 0-based mesh vertex indices
+        cVertsGT = torch.as_tensor(gt["dp_vertex"], dtype=torch.int64)
+        # label for each pixel of the bbox, [H, W] tensor of long
+        labels_dt = resample_coarse_segm_tensor_to_bbox(
+            coarse_segm.unsqueeze(0), bbox_xywh_abs
+        ).squeeze(0)
+        x, y, w, h = bbox_xywh_abs
+        # embedding for each pixel of the bbox, [D, H, W] tensor of float32
+        embedding = F.interpolate(
+            embedding.unsqueeze(0), (int(h), int(w)), mode="bilinear", align_corners=False
+        ).squeeze(0)
+        # valid locations py, px
+        py_pt = torch.from_numpy(py[pt_mask > -1])
+        px_pt = torch.from_numpy(px[pt_mask > -1])
+        cVerts = torch.ones_like(cVertsGT) * -1
+        cVerts[pt_mask > -1] = self.findClosestVertsCse(
+            embedding, py_pt, px_pt, labels_dt, gt["ref_model"]
+        )
+        # Get pairwise geodesic distances between gt and estimated mesh points.
+        dist = self.getDistancesCse(cVertsGT, cVerts, gt["ref_model"])
+        # normalize distances
+        if (gt["ref_model"] == "smpl_27554") and ("dp_I" in gt):
+            Current_Mean_Distances = self.Mean_Distances[
+                self.CoarseParts[np.array(gt["dp_I"], dtype=int)]
+            ]
+        else:
+            Current_Mean_Distances = 0.255
+        return dist, Current_Mean_Distances
+
+    def computeOgps(self, imgId, catId):
+        p = self.params
+        # dimension here should be Nxm
+        g = self._gts[imgId, catId]
+        d = self._dts[imgId, catId]
+        inds = np.argsort([-d_["score"] for d_ in d], kind="mergesort")
+        d = [d[i] for i in inds]
+        if len(d) > p.maxDets[-1]:
+            d = d[0 : p.maxDets[-1]]
+        # if len(gts) == 0 and len(dts) == 0:
+        if len(g) == 0 or len(d) == 0:
+            return []
+        ious = np.zeros((len(d), len(g)))
+        # compute opgs between each detection and ground truth object
+        # sigma = self.sigma #0.255 # dist = 0.3m corresponds to ogps = 0.5
+        # 1 # dist = 0.3m corresponds to ogps = 0.96
+        # 1.45 # dist = 1.7m (person height) corresponds to ogps = 0.5)
+        for j, gt in enumerate(g):
+            if not gt["ignore"]:
+                g_ = gt["bbox"]
+                for i, dt in enumerate(d):
+                    #
+                    dy = int(dt["bbox"][3])
+                    dx = int(dt["bbox"][2])
+                    dp_x = np.array(gt["dp_x"]) * g_[2] / 255.0
+                    dp_y = np.array(gt["dp_y"]) * g_[3] / 255.0
+                    py = (dp_y + g_[1] - dt["bbox"][1]).astype(int)
+                    px = (dp_x + g_[0] - dt["bbox"][0]).astype(int)
+                    #
+                    pts = np.zeros(len(px))
+                    pts[px >= dx] = -1
+                    pts[py >= dy] = -1
+                    pts[px < 0] = -1
+                    pts[py < 0] = -1
+                    if len(pts) < 1:
+                        ogps = 0.0
+                    elif np.max(pts) == -1:
+                        ogps = 0.0
+                    else:
+                        px[pts == -1] = 0
+                        py[pts == -1] = 0
+                        dists_between_matches, dist_norm_coeffs = self.computeOgps_single_pair(
+                            dt, gt, py, px, pts
+                        )
+                        # Compute gps
+                        ogps_values = np.exp(
+                            -(dists_between_matches**2) / (2 * (dist_norm_coeffs**2))
+                        )
+                        #
+                        ogps = np.mean(ogps_values) if len(ogps_values) > 0 else 0.0
+                    ious[i, j] = ogps
+
+        gbb = [gt["bbox"] for gt in g]
+        dbb = [dt["bbox"] for dt in d]
+
+        # compute iou between each dt and gt region
+        iscrowd = [int(o.get("iscrowd", 0)) for o in g]
+        ious_bb = maskUtils.iou(dbb, gbb, iscrowd)
+        return ious, ious_bb
+
+    def evaluateImg(self, imgId, catId, aRng, maxDet):
+        """
+        perform evaluation for single category and image
+        :return: dict (single image results)
+        """
+
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+        if len(gt) == 0 and len(dt) == 0:
+            return None
+
+        for g in gt:
+            # g['_ignore'] = g['ignore']
+            if g["ignore"] or (g["area"] < aRng[0] or g["area"] > aRng[1]):
+                g["_ignore"] = True
+            else:
+                g["_ignore"] = False
+
+        # sort dt highest score first, sort gt ignore last
+        gtind = np.argsort([g["_ignore"] for g in gt], kind="mergesort")
+        gt = [gt[i] for i in gtind]
+        dtind = np.argsort([-d["score"] for d in dt], kind="mergesort")
+        dt = [dt[i] for i in dtind[0:maxDet]]
+        iscrowd = [int(o.get("iscrowd", 0)) for o in gt]
+        # load computed ious
+        if p.iouType == "densepose":
+            # print('Checking the length', len(self.ious[imgId, catId]))
+            # if len(self.ious[imgId, catId]) == 0:
+            #    print(self.ious[imgId, catId])
+            ious = (
+                self.ious[imgId, catId][0][:, gtind]
+                if len(self.ious[imgId, catId]) > 0
+                else self.ious[imgId, catId]
+            )
+            ioubs = (
+                self.ious[imgId, catId][1][:, gtind]
+                if len(self.ious[imgId, catId]) > 0
+                else self.ious[imgId, catId]
+            )
+            if self._dpEvalMode in {DensePoseEvalMode.GPSM, DensePoseEvalMode.IOU}:
+                iousM = (
+                    self.real_ious[imgId, catId][:, gtind]
+                    if len(self.real_ious[imgId, catId]) > 0
+                    else self.real_ious[imgId, catId]
+                )
+        else:
+            ious = (
+                self.ious[imgId, catId][:, gtind]
+                if len(self.ious[imgId, catId]) > 0
+                else self.ious[imgId, catId]
+            )
+
+        T = len(p.iouThrs)
+        G = len(gt)
+        D = len(dt)
+        gtm = np.zeros((T, G))
+        dtm = np.zeros((T, D))
+        gtIg = np.array([g["_ignore"] for g in gt])
+        dtIg = np.zeros((T, D))
+        if np.all(gtIg) and p.iouType == "densepose":
+            dtIg = np.logical_or(dtIg, True)
+
+        if len(ious) > 0:  # and not p.iouType == 'densepose':
+            for tind, t in enumerate(p.iouThrs):
+                for dind, d in enumerate(dt):
+                    # information about best match so far (m=-1 -> unmatched)
+                    iou = min([t, 1 - 1e-10])
+                    m = -1
+                    for gind, _g in enumerate(gt):
+                        # if this gt already matched, and not a crowd, continue
+                        if gtm[tind, gind] > 0 and not iscrowd[gind]:
+                            continue
+                        # if dt matched to reg gt, and on ignore gt, stop
+                        if m > -1 and gtIg[m] == 0 and gtIg[gind] == 1:
+                            break
+                        if p.iouType == "densepose":
+                            if self._dpEvalMode == DensePoseEvalMode.GPSM:
+                                new_iou = np.sqrt(iousM[dind, gind] * ious[dind, gind])
+                            elif self._dpEvalMode == DensePoseEvalMode.IOU:
+                                new_iou = iousM[dind, gind]
+                            elif self._dpEvalMode == DensePoseEvalMode.GPS:
+                                new_iou = ious[dind, gind]
+                        else:
+                            new_iou = ious[dind, gind]
+                        if new_iou < iou:
+                            continue
+                        if new_iou == 0.0:
+                            continue
+                        # if match successful and best so far, store appropriately
+                        iou = new_iou
+                        m = gind
+                    # if match made store id of match for both dt and gt
+                    if m == -1:
+                        continue
+                    dtIg[tind, dind] = gtIg[m]
+                    dtm[tind, dind] = gt[m]["id"]
+                    gtm[tind, m] = d["id"]
+
+        if p.iouType == "densepose":
+            if not len(ioubs) == 0:
+                for dind, d in enumerate(dt):
+                    # information about best match so far (m=-1 -> unmatched)
+                    if dtm[tind, dind] == 0:
+                        ioub = 0.8
+                        m = -1
+                        for gind, _g in enumerate(gt):
+                            # if this gt already matched, and not a crowd, continue
+                            if gtm[tind, gind] > 0 and not iscrowd[gind]:
+                                continue
+                            # continue to next gt unless better match made
+                            if ioubs[dind, gind] < ioub:
+                                continue
+                            # if match successful and best so far, store appropriately
+                            ioub = ioubs[dind, gind]
+                            m = gind
+                            # if match made store id of match for both dt and gt
+                        if m > -1:
+                            dtIg[:, dind] = gtIg[m]
+                            if gtIg[m]:
+                                dtm[tind, dind] = gt[m]["id"]
+                                gtm[tind, m] = d["id"]
+        # set unmatched detections outside of area range to ignore
+        a = np.array([d["area"] < aRng[0] or d["area"] > aRng[1] for d in dt]).reshape((1, len(dt)))
+        dtIg = np.logical_or(dtIg, np.logical_and(dtm == 0, np.repeat(a, T, 0)))
+        # store results for given image and category
+        # print('Done with the function', len(self.ious[imgId, catId]))
+        return {
+            "image_id": imgId,
+            "category_id": catId,
+            "aRng": aRng,
+            "maxDet": maxDet,
+            "dtIds": [d["id"] for d in dt],
+            "gtIds": [g["id"] for g in gt],
+            "dtMatches": dtm,
+            "gtMatches": gtm,
+            "dtScores": [d["score"] for d in dt],
+            "gtIgnore": gtIg,
+            "dtIgnore": dtIg,
+        }
+
+    def accumulate(self, p=None):
+        """
+        Accumulate per image evaluation results and store the result in self.eval
+        :param p: input params for evaluation
+        :return: None
+        """
+        logger.info("Accumulating evaluation results...")
+        tic = time.time()
+        if not self.evalImgs:
+            logger.info("Please run evaluate() first")
+        # allows input customized parameters
+        if p is None:
+            p = self.params
+        p.catIds = p.catIds if p.useCats == 1 else [-1]
+        T = len(p.iouThrs)
+        R = len(p.recThrs)
+        K = len(p.catIds) if p.useCats else 1
+        A = len(p.areaRng)
+        M = len(p.maxDets)
+        precision = -(np.ones((T, R, K, A, M)))  # -1 for the precision of absent categories
+        recall = -(np.ones((T, K, A, M)))
+
+        # create dictionary for future indexing
+        logger.info("Categories: {}".format(p.catIds))
+        _pe = self._paramsEval
+        catIds = _pe.catIds if _pe.useCats else [-1]
+        setK = set(catIds)
+        setA = set(map(tuple, _pe.areaRng))
+        setM = set(_pe.maxDets)
+        setI = set(_pe.imgIds)
+        # get inds to evaluate
+        k_list = [n for n, k in enumerate(p.catIds) if k in setK]
+        m_list = [m for n, m in enumerate(p.maxDets) if m in setM]
+        a_list = [n for n, a in enumerate(map(lambda x: tuple(x), p.areaRng)) if a in setA]
+        i_list = [n for n, i in enumerate(p.imgIds) if i in setI]
+        I0 = len(_pe.imgIds)
+        A0 = len(_pe.areaRng)
+        # retrieve E at each category, area range, and max number of detections
+        for k, k0 in enumerate(k_list):
+            Nk = k0 * A0 * I0
+            for a, a0 in enumerate(a_list):
+                Na = a0 * I0
+                for m, maxDet in enumerate(m_list):
+                    E = [self.evalImgs[Nk + Na + i] for i in i_list]
+                    E = [e for e in E if e is not None]
+                    if len(E) == 0:
+                        continue
+                    dtScores = np.concatenate([e["dtScores"][0:maxDet] for e in E])
+
+                    # different sorting method generates slightly different results.
+                    # mergesort is used to be consistent as Matlab implementation.
+                    inds = np.argsort(-dtScores, kind="mergesort")
+
+                    dtm = np.concatenate([e["dtMatches"][:, 0:maxDet] for e in E], axis=1)[:, inds]
+                    dtIg = np.concatenate([e["dtIgnore"][:, 0:maxDet] for e in E], axis=1)[:, inds]
+                    gtIg = np.concatenate([e["gtIgnore"] for e in E])
+                    npig = np.count_nonzero(gtIg == 0)
+                    if npig == 0:
+                        continue
+                    tps = np.logical_and(dtm, np.logical_not(dtIg))
+                    fps = np.logical_and(np.logical_not(dtm), np.logical_not(dtIg))
+                    tp_sum = np.cumsum(tps, axis=1).astype(dtype=float)
+                    fp_sum = np.cumsum(fps, axis=1).astype(dtype=float)
+                    for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
+                        tp = np.array(tp)
+                        fp = np.array(fp)
+                        nd = len(tp)
+                        rc = tp / npig
+                        pr = tp / (fp + tp + np.spacing(1))
+                        q = np.zeros((R,))
+
+                        if nd:
+                            recall[t, k, a, m] = rc[-1]
+                        else:
+                            recall[t, k, a, m] = 0
+
+                        # numpy is slow without cython optimization for accessing elements
+                        # use python array gets significant speed improvement
+                        pr = pr.tolist()
+                        q = q.tolist()
+
+                        for i in range(nd - 1, 0, -1):
+                            if pr[i] > pr[i - 1]:
+                                pr[i - 1] = pr[i]
+
+                        inds = np.searchsorted(rc, p.recThrs, side="left")
+                        try:
+                            for ri, pi in enumerate(inds):
+                                q[ri] = pr[pi]
+                        except Exception:
+                            pass
+                        precision[t, :, k, a, m] = np.array(q)
+        logger.info(
+            "Final: max precision {}, min precision {}".format(np.max(precision), np.min(precision))
+        )
+        self.eval = {
+            "params": p,
+            "counts": [T, R, K, A, M],
+            "date": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "precision": precision,
+            "recall": recall,
+        }
+        toc = time.time()
+        logger.info("DONE (t={:0.2f}s).".format(toc - tic))
+
+    def summarize(self):
+        """
+        Compute and display summary metrics for evaluation results.
+        Note this function can *only* be applied on the default parameter setting
+        """
+
+        def _summarize(ap=1, iouThr=None, areaRng="all", maxDets=100):
+            p = self.params
+            iStr = " {:<18} {} @[ {}={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}"
+            titleStr = "Average Precision" if ap == 1 else "Average Recall"
+            typeStr = "(AP)" if ap == 1 else "(AR)"
+            measure = "IoU"
+            if self.params.iouType == "keypoints":
+                measure = "OKS"
+            elif self.params.iouType == "densepose":
+                measure = "OGPS"
+            iouStr = (
+                "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1])
+                if iouThr is None
+                else "{:0.2f}".format(iouThr)
+            )
+
+            aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
+            mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
+            if ap == 1:
+                # dimension of precision: [TxRxKxAxM]
+                s = self.eval["precision"]
+                # IoU
+                if iouThr is not None:
+                    t = np.where(np.abs(iouThr - p.iouThrs) < 0.001)[0]
+                    s = s[t]
+                s = s[:, :, :, aind, mind]
+            else:
+                # dimension of recall: [TxKxAxM]
+                s = self.eval["recall"]
+                if iouThr is not None:
+                    t = np.where(np.abs(iouThr - p.iouThrs) < 0.001)[0]
+                    s = s[t]
+                s = s[:, :, aind, mind]
+            if len(s[s > -1]) == 0:
+                mean_s = -1
+            else:
+                mean_s = np.mean(s[s > -1])
+            logger.info(iStr.format(titleStr, typeStr, measure, iouStr, areaRng, maxDets, mean_s))
+            return mean_s
+
+        def _summarizeDets():
+            stats = np.zeros((12,))
+            stats[0] = _summarize(1)
+            stats[1] = _summarize(1, iouThr=0.5, maxDets=self.params.maxDets[2])
+            stats[2] = _summarize(1, iouThr=0.75, maxDets=self.params.maxDets[2])
+            stats[3] = _summarize(1, areaRng="small", maxDets=self.params.maxDets[2])
+            stats[4] = _summarize(1, areaRng="medium", maxDets=self.params.maxDets[2])
+            stats[5] = _summarize(1, areaRng="large", maxDets=self.params.maxDets[2])
+            stats[6] = _summarize(0, maxDets=self.params.maxDets[0])
+            stats[7] = _summarize(0, maxDets=self.params.maxDets[1])
+            stats[8] = _summarize(0, maxDets=self.params.maxDets[2])
+            stats[9] = _summarize(0, areaRng="small", maxDets=self.params.maxDets[2])
+            stats[10] = _summarize(0, areaRng="medium", maxDets=self.params.maxDets[2])
+            stats[11] = _summarize(0, areaRng="large", maxDets=self.params.maxDets[2])
+            return stats
+
+        def _summarizeKps():
+            stats = np.zeros((10,))
+            stats[0] = _summarize(1, maxDets=20)
+            stats[1] = _summarize(1, maxDets=20, iouThr=0.5)
+            stats[2] = _summarize(1, maxDets=20, iouThr=0.75)
+            stats[3] = _summarize(1, maxDets=20, areaRng="medium")
+            stats[4] = _summarize(1, maxDets=20, areaRng="large")
+            stats[5] = _summarize(0, maxDets=20)
+            stats[6] = _summarize(0, maxDets=20, iouThr=0.5)
+            stats[7] = _summarize(0, maxDets=20, iouThr=0.75)
+            stats[8] = _summarize(0, maxDets=20, areaRng="medium")
+            stats[9] = _summarize(0, maxDets=20, areaRng="large")
+            return stats
+
+        def _summarizeUvs():
+            stats = [_summarize(1, maxDets=self.params.maxDets[0])]
+            min_threshold = self.params.iouThrs.min()
+            if min_threshold <= 0.201:
+                stats += [_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.2)]
+            if min_threshold <= 0.301:
+                stats += [_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.3)]
+            if min_threshold <= 0.401:
+                stats += [_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.4)]
+            stats += [
+                _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.5),
+                _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.75),
+                _summarize(1, maxDets=self.params.maxDets[0], areaRng="medium"),
+                _summarize(1, maxDets=self.params.maxDets[0], areaRng="large"),
+                _summarize(0, maxDets=self.params.maxDets[0]),
+                _summarize(0, maxDets=self.params.maxDets[0], iouThr=0.5),
+                _summarize(0, maxDets=self.params.maxDets[0], iouThr=0.75),
+                _summarize(0, maxDets=self.params.maxDets[0], areaRng="medium"),
+                _summarize(0, maxDets=self.params.maxDets[0], areaRng="large"),
+            ]
+            return np.array(stats)
+
+        def _summarizeUvsOld():
+            stats = np.zeros((18,))
+            stats[0] = _summarize(1, maxDets=self.params.maxDets[0])
+            stats[1] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.5)
+            stats[2] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.55)
+            stats[3] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.60)
+            stats[4] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.65)
+            stats[5] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.70)
+            stats[6] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.75)
+            stats[7] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.80)
+            stats[8] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.85)
+            stats[9] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.90)
+            stats[10] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.95)
+            stats[11] = _summarize(1, maxDets=self.params.maxDets[0], areaRng="medium")
+            stats[12] = _summarize(1, maxDets=self.params.maxDets[0], areaRng="large")
+            stats[13] = _summarize(0, maxDets=self.params.maxDets[0])
+            stats[14] = _summarize(0, maxDets=self.params.maxDets[0], iouThr=0.5)
+            stats[15] = _summarize(0, maxDets=self.params.maxDets[0], iouThr=0.75)
+            stats[16] = _summarize(0, maxDets=self.params.maxDets[0], areaRng="medium")
+            stats[17] = _summarize(0, maxDets=self.params.maxDets[0], areaRng="large")
+            return stats
+
+        if not self.eval:
+            raise Exception("Please run accumulate() first")
+        iouType = self.params.iouType
+        if iouType in ["segm", "bbox"]:
+            summarize = _summarizeDets
+        elif iouType in ["keypoints"]:
+            summarize = _summarizeKps
+        elif iouType in ["densepose"]:
+            summarize = _summarizeUvs
+        self.stats = summarize()
+
+    def __str__(self):
+        self.summarize()
+
+    # ================ functions for dense pose ==============================
+    def findAllClosestVertsUV(self, U_points, V_points, Index_points):
+        ClosestVerts = np.ones(Index_points.shape) * -1
+        for i in np.arange(24):
+            #
+            if (i + 1) in Index_points:
+                UVs = np.array(
+                    [U_points[Index_points == (i + 1)], V_points[Index_points == (i + 1)]]
+                )
+                Current_Part_UVs = self.Part_UVs[i]
+                Current_Part_ClosestVertInds = self.Part_ClosestVertInds[i]
+                D = ssd.cdist(Current_Part_UVs.transpose(), UVs.transpose()).squeeze()
+                ClosestVerts[Index_points == (i + 1)] = Current_Part_ClosestVertInds[
+                    np.argmin(D, axis=0)
+                ]
+        ClosestVertsTransformed = self.PDIST_transform[ClosestVerts.astype(int) - 1]
+        ClosestVertsTransformed[ClosestVerts < 0] = 0
+        return ClosestVertsTransformed
+
+    def findClosestVertsCse(self, embedding, py, px, mask, mesh_name):
+        mesh_vertex_embeddings = self.embedder(mesh_name)
+        pixel_embeddings = embedding[:, py, px].t().to(device="cuda")
+        mask_vals = mask[py, px]
+        edm = squared_euclidean_distance_matrix(pixel_embeddings, mesh_vertex_embeddings)
+        vertex_indices = edm.argmin(dim=1).cpu()
+        vertex_indices[mask_vals <= 0] = -1
+        return vertex_indices
+
+    def findAllClosestVertsGT(self, gt):
+        #
+        I_gt = np.array(gt["dp_I"])
+        U_gt = np.array(gt["dp_U"])
+        V_gt = np.array(gt["dp_V"])
+        #
+        # print(I_gt)
+        #
+        ClosestVertsGT = np.ones(I_gt.shape) * -1
+        for i in np.arange(24):
+            if (i + 1) in I_gt:
+                UVs = np.array([U_gt[I_gt == (i + 1)], V_gt[I_gt == (i + 1)]])
+                Current_Part_UVs = self.Part_UVs[i]
+                Current_Part_ClosestVertInds = self.Part_ClosestVertInds[i]
+                D = ssd.cdist(Current_Part_UVs.transpose(), UVs.transpose()).squeeze()
+                ClosestVertsGT[I_gt == (i + 1)] = Current_Part_ClosestVertInds[np.argmin(D, axis=0)]
+        #
+        ClosestVertsGTTransformed = self.PDIST_transform[ClosestVertsGT.astype(int) - 1]
+        ClosestVertsGTTransformed[ClosestVertsGT < 0] = 0
+        return ClosestVertsGT, ClosestVertsGTTransformed
+
+    def getDistancesCse(self, cVertsGT, cVerts, mesh_name):
+        geodists_vertices = torch.ones_like(cVertsGT) * float("inf")
+        selected = (cVertsGT >= 0) * (cVerts >= 0)
+        mesh = create_mesh(mesh_name, "cpu")
+        geodists_vertices[selected] = mesh.geodists[cVertsGT[selected], cVerts[selected]]
+        return geodists_vertices.numpy()
+
+    def getDistancesUV(self, cVertsGT, cVerts):
+        #
+        n = 27554
+        dists = []
+        for d in range(len(cVertsGT)):
+            if cVertsGT[d] > 0:
+                if cVerts[d] > 0:
+                    i = cVertsGT[d] - 1
+                    j = cVerts[d] - 1
+                    if j == i:
+                        dists.append(0)
+                    elif j > i:
+                        ccc = i
+                        i = j
+                        j = ccc
+                        i = n - i - 1
+                        j = n - j - 1
+                        k = (n * (n - 1) / 2) - (n - i) * ((n - i) - 1) / 2 + j - i - 1
+                        k = (n * n - n) / 2 - k - 1
+                        dists.append(self.Pdist_matrix[int(k)][0])
+                    else:
+                        i = n - i - 1
+                        j = n - j - 1
+                        k = (n * (n - 1) / 2) - (n - i) * ((n - i) - 1) / 2 + j - i - 1
+                        k = (n * n - n) / 2 - k - 1
+                        dists.append(self.Pdist_matrix[int(k)][0])
+                else:
+                    dists.append(np.inf)
+        return np.atleast_1d(np.array(dists).squeeze())
+
+
+class Params:
+    """
+    Params for coco evaluation api
+    """
+
+    def setDetParams(self):
+        self.imgIds = []
+        self.catIds = []
+        # np.arange causes trouble.  the data point on arange is slightly larger than the true value
+        self.iouThrs = np.linspace(0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True)
+        self.recThrs = np.linspace(0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True)
+        self.maxDets = [1, 10, 100]
+        self.areaRng = [
+            [0**2, 1e5**2],
+            [0**2, 32**2],
+            [32**2, 96**2],
+            [96**2, 1e5**2],
+        ]
+        self.areaRngLbl = ["all", "small", "medium", "large"]
+        self.useCats = 1
+
+    def setKpParams(self):
+        self.imgIds = []
+        self.catIds = []
+        # np.arange causes trouble.  the data point on arange is slightly larger than the true value
+        self.iouThrs = np.linspace(0.5, 0.95, np.round((0.95 - 0.5) / 0.05) + 1, endpoint=True)
+        self.recThrs = np.linspace(0.0, 1.00, np.round((1.00 - 0.0) / 0.01) + 1, endpoint=True)
+        self.maxDets = [20]
+        self.areaRng = [[0**2, 1e5**2], [32**2, 96**2], [96**2, 1e5**2]]
+        self.areaRngLbl = ["all", "medium", "large"]
+        self.useCats = 1
+
+    def setUvParams(self):
+        self.imgIds = []
+        self.catIds = []
+        self.iouThrs = np.linspace(0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True)
+        self.recThrs = np.linspace(0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True)
+        self.maxDets = [20]
+        self.areaRng = [[0**2, 1e5**2], [32**2, 96**2], [96**2, 1e5**2]]
+        self.areaRngLbl = ["all", "medium", "large"]
+        self.useCats = 1
+
+    def __init__(self, iouType="segm"):
+        if iouType == "segm" or iouType == "bbox":
+            self.setDetParams()
+        elif iouType == "keypoints":
+            self.setKpParams()
+        elif iouType == "densepose":
+            self.setUvParams()
+        else:
+            raise Exception("iouType not supported")
+        self.iouType = iouType
+        # useSegm is deprecated
+        self.useSegm = None

densepose/evaluation/evaluator.py

@@ -0,0 +1,423 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import contextlib
+import copy
+import io
+import itertools
+import logging
+import numpy as np
+import os
+from collections import OrderedDict
+from typing import Dict, Iterable, List, Optional
+import pycocotools.mask as mask_utils
+import torch
+from pycocotools.coco import COCO
+from tabulate import tabulate
+
+from detectron2.config import CfgNode
+from detectron2.data import MetadataCatalog
+from detectron2.evaluation import DatasetEvaluator
+from detectron2.structures import BoxMode
+from detectron2.utils.comm import gather, get_rank, is_main_process, synchronize
+from detectron2.utils.file_io import PathManager
+from detectron2.utils.logger import create_small_table
+
+from densepose.converters import ToChartResultConverter, ToMaskConverter
+from densepose.data.datasets.coco import maybe_filter_and_map_categories_cocoapi
+from densepose.structures import (
+    DensePoseChartPredictorOutput,
+    DensePoseEmbeddingPredictorOutput,
+    quantize_densepose_chart_result,
+)
+
+from .densepose_coco_evaluation import DensePoseCocoEval, DensePoseEvalMode
+from .mesh_alignment_evaluator import MeshAlignmentEvaluator
+from .tensor_storage import (
+    SingleProcessFileTensorStorage,
+    SingleProcessRamTensorStorage,
+    SingleProcessTensorStorage,
+    SizeData,
+    storage_gather,
+)
+
+
+class DensePoseCOCOEvaluator(DatasetEvaluator):
+    def __init__(
+        self,
+        dataset_name,
+        distributed,
+        output_dir=None,
+        evaluator_type: str = "iuv",
+        min_iou_threshold: float = 0.5,
+        storage: Optional[SingleProcessTensorStorage] = None,
+        embedder=None,
+        should_evaluate_mesh_alignment: bool = False,
+        mesh_alignment_mesh_names: Optional[List[str]] = None,
+    ):
+        self._embedder = embedder
+        self._distributed = distributed
+        self._output_dir = output_dir
+        self._evaluator_type = evaluator_type
+        self._storage = storage
+        self._should_evaluate_mesh_alignment = should_evaluate_mesh_alignment
+
+        assert not (
+            should_evaluate_mesh_alignment and embedder is None
+        ), "Mesh alignment evaluation is activated, but no vertex embedder provided!"
+        if should_evaluate_mesh_alignment:
+            self._mesh_alignment_evaluator = MeshAlignmentEvaluator(
+                embedder,
+                mesh_alignment_mesh_names,
+            )
+
+        self._cpu_device = torch.device("cpu")
+        self._logger = logging.getLogger(__name__)
+
+        self._metadata = MetadataCatalog.get(dataset_name)
+        self._min_threshold = min_iou_threshold
+        json_file = PathManager.get_local_path(self._metadata.json_file)
+        with contextlib.redirect_stdout(io.StringIO()):
+            self._coco_api = COCO(json_file)
+        maybe_filter_and_map_categories_cocoapi(dataset_name, self._coco_api)
+
+    def reset(self):
+        self._predictions = []
+
+    def process(self, inputs, outputs):
+        """
+        Args:
+            inputs: the inputs to a COCO model (e.g., GeneralizedRCNN).
+                It is a list of dict. Each dict corresponds to an image and
+                contains keys like "height", "width", "file_name", "image_id".
+            outputs: the outputs of a COCO model. It is a list of dicts with key
+                "instances" that contains :class:`Instances`.
+                The :class:`Instances` object needs to have `densepose` field.
+        """
+        for input, output in zip(inputs, outputs):
+            instances = output["instances"].to(self._cpu_device)
+            if not instances.has("pred_densepose"):
+                continue
+            prediction_list = prediction_to_dict(
+                instances,
+                input["image_id"],
+                self._embedder,
+                self._metadata.class_to_mesh_name,
+                self._storage is not None,
+            )
+            if self._storage is not None:
+                for prediction_dict in prediction_list:
+                    dict_to_store = {}
+                    for field_name in self._storage.data_schema:
+                        dict_to_store[field_name] = prediction_dict[field_name]
+                    record_id = self._storage.put(dict_to_store)
+                    prediction_dict["record_id"] = record_id
+                    prediction_dict["rank"] = get_rank()
+                    for field_name in self._storage.data_schema:
+                        del prediction_dict[field_name]
+            self._predictions.extend(prediction_list)
+
+    def evaluate(self, img_ids=None):
+        if self._distributed:
+            synchronize()
+            predictions = gather(self._predictions)
+            predictions = list(itertools.chain(*predictions))
+        else:
+            predictions = self._predictions
+
+        multi_storage = storage_gather(self._storage) if self._storage is not None else None
+
+        if not is_main_process():
+            return
+        return copy.deepcopy(self._eval_predictions(predictions, multi_storage, img_ids))
+
+    def _eval_predictions(self, predictions, multi_storage=None, img_ids=None):
+        """
+        Evaluate predictions on densepose.
+        Return results with the metrics of the tasks.
+        """
+        self._logger.info("Preparing results for COCO format ...")
+
+        if self._output_dir:
+            PathManager.mkdirs(self._output_dir)
+            file_path = os.path.join(self._output_dir, "coco_densepose_predictions.pth")
+            with PathManager.open(file_path, "wb") as f:
+                torch.save(predictions, f)
+
+        self._logger.info("Evaluating predictions ...")
+        res = OrderedDict()
+        results_gps, results_gpsm, results_segm = _evaluate_predictions_on_coco(
+            self._coco_api,
+            predictions,
+            multi_storage,
+            self._embedder,
+            class_names=self._metadata.get("thing_classes"),
+            min_threshold=self._min_threshold,
+            img_ids=img_ids,
+        )
+        res["densepose_gps"] = results_gps
+        res["densepose_gpsm"] = results_gpsm
+        res["densepose_segm"] = results_segm
+        if self._should_evaluate_mesh_alignment:
+            res["densepose_mesh_alignment"] = self._evaluate_mesh_alignment()
+        return res
+
+    def _evaluate_mesh_alignment(self):
+        self._logger.info("Mesh alignment evaluation ...")
+        mean_ge, mean_gps, per_mesh_metrics = self._mesh_alignment_evaluator.evaluate()
+        results = {
+            "GE": mean_ge * 100,
+            "GPS": mean_gps * 100,
+        }
+        mesh_names = set()
+        for metric_name in per_mesh_metrics:
+            for mesh_name, value in per_mesh_metrics[metric_name].items():
+                results[f"{metric_name}-{mesh_name}"] = value * 100
+                mesh_names.add(mesh_name)
+        self._print_mesh_alignment_results(results, mesh_names)
+        return results
+
+    def _print_mesh_alignment_results(self, results: Dict[str, float], mesh_names: Iterable[str]):
+        self._logger.info("Evaluation results for densepose, mesh alignment:")
+        self._logger.info(f'| {"Mesh":13s} | {"GErr":7s} | {"GPS":7s} |')
+        self._logger.info("| :-----------: | :-----: | :-----: |")
+        for mesh_name in mesh_names:
+            ge_key = f"GE-{mesh_name}"
+            ge_str = f"{results[ge_key]:.4f}" if ge_key in results else " "
+            gps_key = f"GPS-{mesh_name}"
+            gps_str = f"{results[gps_key]:.4f}" if gps_key in results else " "
+            self._logger.info(f"| {mesh_name:13s} | {ge_str:7s} | {gps_str:7s} |")
+        self._logger.info("| :-------------------------------: |")
+        ge_key = "GE"
+        ge_str = f"{results[ge_key]:.4f}" if ge_key in results else " "
+        gps_key = "GPS"
+        gps_str = f"{results[gps_key]:.4f}" if gps_key in results else " "
+        self._logger.info(f'| {"MEAN":13s} | {ge_str:7s} | {gps_str:7s} |')
+
+
+def prediction_to_dict(instances, img_id, embedder, class_to_mesh_name, use_storage):
+    """
+    Args:
+        instances (Instances): the output of the model
+        img_id (str): the image id in COCO
+
+    Returns:
+        list[dict]: the results in densepose evaluation format
+    """
+    scores = instances.scores.tolist()
+    classes = instances.pred_classes.tolist()
+    raw_boxes_xywh = BoxMode.convert(
+        instances.pred_boxes.tensor.clone(), BoxMode.XYXY_ABS, BoxMode.XYWH_ABS
+    )
+
+    if isinstance(instances.pred_densepose, DensePoseEmbeddingPredictorOutput):
+        results_densepose = densepose_cse_predictions_to_dict(
+            instances, embedder, class_to_mesh_name, use_storage
+        )
+    elif isinstance(instances.pred_densepose, DensePoseChartPredictorOutput):
+        if not use_storage:
+            results_densepose = densepose_chart_predictions_to_dict(instances)
+        else:
+            results_densepose = densepose_chart_predictions_to_storage_dict(instances)
+
+    results = []
+    for k in range(len(instances)):
+        result = {
+            "image_id": img_id,
+            "category_id": classes[k],
+            "bbox": raw_boxes_xywh[k].tolist(),
+            "score": scores[k],
+        }
+        results.append({**result, **results_densepose[k]})
+    return results
+
+
+def densepose_chart_predictions_to_dict(instances):
+    segmentations = ToMaskConverter.convert(
+        instances.pred_densepose, instances.pred_boxes, instances.image_size
+    )
+
+    results = []
+    for k in range(len(instances)):
+        densepose_results_quantized = quantize_densepose_chart_result(
+            ToChartResultConverter.convert(instances.pred_densepose[k], instances.pred_boxes[k])
+        )
+        densepose_results_quantized.labels_uv_uint8 = (
+            densepose_results_quantized.labels_uv_uint8.cpu()
+        )
+        segmentation = segmentations.tensor[k]
+        segmentation_encoded = mask_utils.encode(
+            np.require(segmentation.numpy(), dtype=np.uint8, requirements=["F"])
+        )
+        segmentation_encoded["counts"] = segmentation_encoded["counts"].decode("utf-8")
+        result = {
+            "densepose": densepose_results_quantized,
+            "segmentation": segmentation_encoded,
+        }
+        results.append(result)
+    return results
+
+
+def densepose_chart_predictions_to_storage_dict(instances):
+    results = []
+    for k in range(len(instances)):
+        densepose_predictor_output = instances.pred_densepose[k]
+        result = {
+            "coarse_segm": densepose_predictor_output.coarse_segm.squeeze(0).cpu(),
+            "fine_segm": densepose_predictor_output.fine_segm.squeeze(0).cpu(),
+            "u": densepose_predictor_output.u.squeeze(0).cpu(),
+            "v": densepose_predictor_output.v.squeeze(0).cpu(),
+        }
+        results.append(result)
+    return results
+
+
+def densepose_cse_predictions_to_dict(instances, embedder, class_to_mesh_name, use_storage):
+    results = []
+    for k in range(len(instances)):
+        cse = instances.pred_densepose[k]
+        results.append(
+            {
+                "coarse_segm": cse.coarse_segm[0].cpu(),
+                "embedding": cse.embedding[0].cpu(),
+            }
+        )
+    return results
+
+
+def _evaluate_predictions_on_coco(
+    coco_gt,
+    coco_results,
+    multi_storage=None,
+    embedder=None,
+    class_names=None,
+    min_threshold: float = 0.5,
+    img_ids=None,
+):
+    logger = logging.getLogger(__name__)
+
+    densepose_metrics = _get_densepose_metrics(min_threshold)
+    if len(coco_results) == 0:  # cocoapi does not handle empty results very well
+        logger.warn("No predictions from the model! Set scores to -1")
+        results_gps = {metric: -1 for metric in densepose_metrics}
+        results_gpsm = {metric: -1 for metric in densepose_metrics}
+        results_segm = {metric: -1 for metric in densepose_metrics}
+        return results_gps, results_gpsm, results_segm
+
+    coco_dt = coco_gt.loadRes(coco_results)
+
+    results = []
+    for eval_mode_name in ["GPS", "GPSM", "IOU"]:
+        eval_mode = getattr(DensePoseEvalMode, eval_mode_name)
+        coco_eval = DensePoseCocoEval(
+            coco_gt, coco_dt, "densepose", multi_storage, embedder, dpEvalMode=eval_mode
+        )
+        result = _derive_results_from_coco_eval(
+            coco_eval, eval_mode_name, densepose_metrics, class_names, min_threshold, img_ids
+        )
+        results.append(result)
+    return results
+
+
+def _get_densepose_metrics(min_threshold: float = 0.5):
+    metrics = ["AP"]
+    if min_threshold <= 0.201:
+        metrics += ["AP20"]
+    if min_threshold <= 0.301:
+        metrics += ["AP30"]
+    if min_threshold <= 0.401:
+        metrics += ["AP40"]
+    metrics.extend(["AP50", "AP75", "APm", "APl", "AR", "AR50", "AR75", "ARm", "ARl"])
+    return metrics
+
+
+def _derive_results_from_coco_eval(
+    coco_eval, eval_mode_name, metrics, class_names, min_threshold: float, img_ids
+):
+    if img_ids is not None:
+        coco_eval.params.imgIds = img_ids
+    coco_eval.params.iouThrs = np.linspace(
+        min_threshold, 0.95, int(np.round((0.95 - min_threshold) / 0.05)) + 1, endpoint=True
+    )
+    coco_eval.evaluate()
+    coco_eval.accumulate()
+    coco_eval.summarize()
+    results = {metric: float(coco_eval.stats[idx] * 100) for idx, metric in enumerate(metrics)}
+    logger = logging.getLogger(__name__)
+    logger.info(
+        f"Evaluation results for densepose, {eval_mode_name} metric: \n"
+        + create_small_table(results)
+    )
+    if class_names is None or len(class_names) <= 1:
+        return results
+
+    # Compute per-category AP, the same way as it is done in D2
+    # (see detectron2/evaluation/coco_evaluation.py):
+    precisions = coco_eval.eval["precision"]
+    # precision has dims (iou, recall, cls, area range, max dets)
+    assert len(class_names) == precisions.shape[2]
+
+    results_per_category = []
+    for idx, name in enumerate(class_names):
+        # area range index 0: all area ranges
+        # max dets index -1: typically 100 per image
+        precision = precisions[:, :, idx, 0, -1]
+        precision = precision[precision > -1]
+        ap = np.mean(precision) if precision.size else float("nan")
+        results_per_category.append((f"{name}", float(ap * 100)))
+
+    # tabulate it
+    n_cols = min(6, len(results_per_category) * 2)
+    results_flatten = list(itertools.chain(*results_per_category))
+    results_2d = itertools.zip_longest(*[results_flatten[i::n_cols] for i in range(n_cols)])
+    table = tabulate(
+        results_2d,
+        tablefmt="pipe",
+        floatfmt=".3f",
+        headers=["category", "AP"] * (n_cols // 2),
+        numalign="left",
+    )
+    logger.info(f"Per-category {eval_mode_name} AP: \n" + table)
+
+    results.update({"AP-" + name: ap for name, ap in results_per_category})
+    return results
+
+
+def build_densepose_evaluator_storage(cfg: CfgNode, output_folder: str):
+    storage_spec = cfg.DENSEPOSE_EVALUATION.STORAGE
+    if storage_spec == "none":
+        return None
+    evaluator_type = cfg.DENSEPOSE_EVALUATION.TYPE
+    # common output tensor sizes
+    hout = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE
+    wout = cfg.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE
+    n_csc = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS
+    # specific output tensors
+    if evaluator_type == "iuv":
+        n_fsc = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES + 1
+        schema = {
+            "coarse_segm": SizeData(dtype="float32", shape=(n_csc, hout, wout)),
+            "fine_segm": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
+            "u": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
+            "v": SizeData(dtype="float32", shape=(n_fsc, hout, wout)),
+        }
+    elif evaluator_type == "cse":
+        embed_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE
+        schema = {
+            "coarse_segm": SizeData(dtype="float32", shape=(n_csc, hout, wout)),
+            "embedding": SizeData(dtype="float32", shape=(embed_size, hout, wout)),
+        }
+    else:
+        raise ValueError(f"Unknown evaluator type: {evaluator_type}")
+    # storage types
+    if storage_spec == "ram":
+        storage = SingleProcessRamTensorStorage(schema, io.BytesIO())
+    elif storage_spec == "file":
+        fpath = os.path.join(output_folder, f"DensePoseEvaluatorStorage.{get_rank()}.bin")
+        PathManager.mkdirs(output_folder)
+        storage = SingleProcessFileTensorStorage(schema, fpath, "wb")
+    else:
+        raise ValueError(f"Unknown storage specification: {storage_spec}")
+    return storage

densepose/evaluation/mesh_alignment_evaluator.py

@@ -0,0 +1,68 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# pyre-unsafe
+
+import json
+import logging
+from typing import List, Optional
+import torch
+from torch import nn
+
+from detectron2.utils.file_io import PathManager
+
+from densepose.structures.mesh import create_mesh
+
+
+class MeshAlignmentEvaluator:
+    """
+    Class for evaluation of 3D mesh alignment based on the learned vertex embeddings
+    """
+
+    def __init__(self, embedder: nn.Module, mesh_names: Optional[List[str]]):
+        self.embedder = embedder
+        # use the provided mesh names if not None and not an empty list
+        self.mesh_names = mesh_names if mesh_names else embedder.mesh_names
+        self.logger = logging.getLogger(__name__)
+        with PathManager.open(
+            "https://dl.fbaipublicfiles.com/densepose/data/cse/mesh_keyvertices_v0.json", "r"
+        ) as f:
+            self.mesh_keyvertices = json.load(f)
+
+    def evaluate(self):
+        ge_per_mesh = {}
+        gps_per_mesh = {}
+        for mesh_name_1 in self.mesh_names:
+            avg_errors = []
+            avg_gps = []
+            embeddings_1 = self.embedder(mesh_name_1)
+            keyvertices_1 = self.mesh_keyvertices[mesh_name_1]
+            keyvertex_names_1 = list(keyvertices_1.keys())
+            keyvertex_indices_1 = [keyvertices_1[name] for name in keyvertex_names_1]
+            for mesh_name_2 in self.mesh_names:
+                if mesh_name_1 == mesh_name_2:
+                    continue
+                embeddings_2 = self.embedder(mesh_name_2)
+                keyvertices_2 = self.mesh_keyvertices[mesh_name_2]
+                sim_matrix_12 = embeddings_1[keyvertex_indices_1].mm(embeddings_2.T)
+                vertices_2_matching_keyvertices_1 = sim_matrix_12.argmax(axis=1)
+                mesh_2 = create_mesh(mesh_name_2, embeddings_2.device)
+                geodists = mesh_2.geodists[
+                    vertices_2_matching_keyvertices_1,
+                    [keyvertices_2[name] for name in keyvertex_names_1],
+                ]
+                Current_Mean_Distances = 0.255
+                gps = (-(geodists**2) / (2 * (Current_Mean_Distances**2))).exp()
+                avg_errors.append(geodists.mean().item())
+                avg_gps.append(gps.mean().item())
+
+            ge_mean = torch.as_tensor(avg_errors).mean().item()
+            gps_mean = torch.as_tensor(avg_gps).mean().item()
+            ge_per_mesh[mesh_name_1] = ge_mean
+            gps_per_mesh[mesh_name_1] = gps_mean
+        ge_mean_global = torch.as_tensor(list(ge_per_mesh.values())).mean().item()
+        gps_mean_global = torch.as_tensor(list(gps_per_mesh.values())).mean().item()
+        per_mesh_metrics = {
+            "GE": ge_per_mesh,
+            "GPS": gps_per_mesh,
+        }
+        return ge_mean_global, gps_mean_global, per_mesh_metrics

densepose/evaluation/tensor_storage.py

@@ -0,0 +1,241 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# pyre-unsafe
+
+import io
+import numpy as np
+import os
+from dataclasses import dataclass
+from functools import reduce
+from operator import mul
+from typing import BinaryIO, Dict, Optional, Tuple
+import torch
+
+from detectron2.utils.comm import gather, get_rank
+from detectron2.utils.file_io import PathManager
+
+
+@dataclass
+class SizeData:
+    dtype: str
+    shape: Tuple[int]
+
+
+def _calculate_record_field_size_b(data_schema: Dict[str, SizeData], field_name: str) -> int:
+    schema = data_schema[field_name]
+    element_size_b = np.dtype(schema.dtype).itemsize
+    record_field_size_b = reduce(mul, schema.shape) * element_size_b
+    return record_field_size_b
+
+
+def _calculate_record_size_b(data_schema: Dict[str, SizeData]) -> int:
+    record_size_b = 0
+    for field_name in data_schema:
+        record_field_size_b = _calculate_record_field_size_b(data_schema, field_name)
+        record_size_b += record_field_size_b
+    return record_size_b
+
+
+def _calculate_record_field_sizes_b(data_schema: Dict[str, SizeData]) -> Dict[str, int]:
+    field_sizes_b = {}
+    for field_name in data_schema:
+        field_sizes_b[field_name] = _calculate_record_field_size_b(data_schema, field_name)
+    return field_sizes_b
+
+
+class SingleProcessTensorStorage:
+    """
+    Compact tensor storage to keep tensor data of predefined size and type.
+    """
+
+    def __init__(self, data_schema: Dict[str, SizeData], storage_impl: BinaryIO):
+        """
+        Construct tensor storage based on information on data shape and size.
+        Internally uses numpy to interpret the type specification.
+        The storage must support operations `seek(offset, whence=os.SEEK_SET)` and
+        `read(size)` to be able to perform the `get` operation.
+        The storage must support operation `write(bytes)` to be able to perform
+        the `put` operation.
+
+        Args:
+            data_schema (dict: str -> SizeData): dictionary which maps tensor name
+                to its size data (shape and data type), e.g.
+                ```
+                {
+                  "coarse_segm": SizeData(dtype="float32", shape=(112, 112)),
+                  "embedding": SizeData(dtype="float32", shape=(16, 112, 112)),
+                }
+                ```
+            storage_impl (BinaryIO): io instance that handles file-like seek, read
+                and write operations, e.g. a file handle or a memory buffer like io.BytesIO
+        """
+        self.data_schema = data_schema
+        self.record_size_b = _calculate_record_size_b(data_schema)
+        self.record_field_sizes_b = _calculate_record_field_sizes_b(data_schema)
+        self.storage_impl = storage_impl
+        self.next_record_id = 0
+
+    def get(self, record_id: int) -> Dict[str, torch.Tensor]:
+        """
+        Load tensors from the storage by record ID
+
+        Args:
+            record_id (int): Record ID, for which to load the data
+
+        Return:
+            dict: str -> tensor: tensor name mapped to tensor data, recorded under the provided ID
+        """
+        self.storage_impl.seek(record_id * self.record_size_b, os.SEEK_SET)
+        data_bytes = self.storage_impl.read(self.record_size_b)
+        assert len(data_bytes) == self.record_size_b, (
+            f"Expected data size {self.record_size_b} B could not be read: "
+            f"got {len(data_bytes)} B"
+        )
+        record = {}
+        cur_idx = 0
+        # it's important to read and write in the same order
+        for field_name in sorted(self.data_schema):
+            schema = self.data_schema[field_name]
+            field_size_b = self.record_field_sizes_b[field_name]
+            chunk = data_bytes[cur_idx : cur_idx + field_size_b]
+            data_np = np.frombuffer(
+                chunk, dtype=schema.dtype, count=reduce(mul, schema.shape)
+            ).reshape(schema.shape)
+            record[field_name] = torch.from_numpy(data_np)
+            cur_idx += field_size_b
+        return record
+
+    def put(self, data: Dict[str, torch.Tensor]) -> int:
+        """
+        Store tensors in the storage
+
+        Args:
+            data (dict: str -> tensor): data to store, a dictionary which maps
+                tensor names into tensors; tensor shapes must match those specified
+                in data schema.
+        Return:
+            int: record ID, under which the data is stored
+        """
+        # it's important to read and write in the same order
+        for field_name in sorted(self.data_schema):
+            assert (
+                field_name in data
+            ), f"Field '{field_name}' not present in data: data keys are {data.keys()}"
+            value = data[field_name]
+            assert value.shape == self.data_schema[field_name].shape, (
+                f"Mismatched tensor shapes for field '{field_name}': "
+                f"expected {self.data_schema[field_name].shape}, got {value.shape}"
+            )
+            data_bytes = value.cpu().numpy().tobytes()
+            assert len(data_bytes) == self.record_field_sizes_b[field_name], (
+                f"Expected field {field_name} to be of size "
+                f"{self.record_field_sizes_b[field_name]} B, got {len(data_bytes)} B"
+            )
+            self.storage_impl.write(data_bytes)
+        record_id = self.next_record_id
+        self.next_record_id += 1
+        return record_id
+
+
+class SingleProcessFileTensorStorage(SingleProcessTensorStorage):
+    """
+    Implementation of a single process tensor storage which stores data in a file
+    """
+
+    def __init__(self, data_schema: Dict[str, SizeData], fpath: str, mode: str):
+        self.fpath = fpath
+        assert "b" in mode, f"Tensor storage should be opened in binary mode, got '{mode}'"
+        if "w" in mode:
+            # pyre-fixme[6]: For 2nd argument expected `Union[typing_extensions.Liter...
+            file_h = PathManager.open(fpath, mode)
+        elif "r" in mode:
+            local_fpath = PathManager.get_local_path(fpath)
+            file_h = open(local_fpath, mode)
+        else:
+            raise ValueError(f"Unsupported file mode {mode}, supported modes: rb, wb")
+        super().__init__(data_schema, file_h)  # pyre-ignore[6]
+
+
+class SingleProcessRamTensorStorage(SingleProcessTensorStorage):
+    """
+    Implementation of a single process tensor storage which stores data in RAM
+    """
+
+    def __init__(self, data_schema: Dict[str, SizeData], buf: io.BytesIO):
+        super().__init__(data_schema, buf)
+
+
+class MultiProcessTensorStorage:
+    """
+    Representation of a set of tensor storages created by individual processes,
+    allows to access those storages from a single owner process. The storages
+    should either be shared or broadcasted to the owner process.
+    The processes are identified by their rank, data is uniquely defined by
+    the rank of the process and the record ID.
+    """
+
+    def __init__(self, rank_to_storage: Dict[int, SingleProcessTensorStorage]):
+        self.rank_to_storage = rank_to_storage
+
+    def get(self, rank: int, record_id: int) -> Dict[str, torch.Tensor]:
+        storage = self.rank_to_storage[rank]
+        return storage.get(record_id)
+
+    def put(self, rank: int, data: Dict[str, torch.Tensor]) -> int:
+        storage = self.rank_to_storage[rank]
+        return storage.put(data)
+
+
+class MultiProcessFileTensorStorage(MultiProcessTensorStorage):
+    def __init__(self, data_schema: Dict[str, SizeData], rank_to_fpath: Dict[int, str], mode: str):
+        rank_to_storage = {
+            rank: SingleProcessFileTensorStorage(data_schema, fpath, mode)
+            for rank, fpath in rank_to_fpath.items()
+        }
+        super().__init__(rank_to_storage)  # pyre-ignore[6]
+
+
+class MultiProcessRamTensorStorage(MultiProcessTensorStorage):
+    def __init__(self, data_schema: Dict[str, SizeData], rank_to_buffer: Dict[int, io.BytesIO]):
+        rank_to_storage = {
+            rank: SingleProcessRamTensorStorage(data_schema, buf)
+            for rank, buf in rank_to_buffer.items()
+        }
+        super().__init__(rank_to_storage)  # pyre-ignore[6]
+
+
+def _ram_storage_gather(
+    storage: SingleProcessRamTensorStorage, dst_rank: int = 0
+) -> Optional[MultiProcessRamTensorStorage]:
+    storage.storage_impl.seek(0, os.SEEK_SET)
+    # TODO: overhead, pickling a bytes object, can just pass bytes in a tensor directly
+    # see detectron2/utils.comm.py
+    data_list = gather(storage.storage_impl.read(), dst=dst_rank)
+    if get_rank() != dst_rank:
+        return None
+    rank_to_buffer = {i: io.BytesIO(data_list[i]) for i in range(len(data_list))}
+    multiprocess_storage = MultiProcessRamTensorStorage(storage.data_schema, rank_to_buffer)
+    return multiprocess_storage
+
+
+def _file_storage_gather(
+    storage: SingleProcessFileTensorStorage,
+    dst_rank: int = 0,
+    mode: str = "rb",
+) -> Optional[MultiProcessFileTensorStorage]:
+    storage.storage_impl.close()
+    fpath_list = gather(storage.fpath, dst=dst_rank)
+    if get_rank() != dst_rank:
+        return None
+    rank_to_fpath = {i: fpath_list[i] for i in range(len(fpath_list))}
+    return MultiProcessFileTensorStorage(storage.data_schema, rank_to_fpath, mode)
+
+
+def storage_gather(
+    storage: SingleProcessTensorStorage, dst_rank: int = 0
+) -> Optional[MultiProcessTensorStorage]:
+    if isinstance(storage, SingleProcessRamTensorStorage):
+        return _ram_storage_gather(storage, dst_rank)
+    elif isinstance(storage, SingleProcessFileTensorStorage):
+        return _file_storage_gather(storage, dst_rank)
+    raise Exception(f"Unsupported storage for gather operation: {storage}")