[ { "path": ".flake8", "content": "[flake8]\nignore =\n # At least two spaces before inline comment\n E261,\n # Line lengths are recommended to be no greater than 79 characters\n E501,\n # Missing whitespace around arithmetic operator \n E226,\n # Blank line contains whitespace\n W293,\n # Do not use bare 'except'\n E722,\n # Line break after binary operator\n W504,\n # isort found an import in the wrong position\n I001\nmax-line-length = 79\nexclude = __init__.py, build, torchreid/metrics/rank_cylib/" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/f\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# celery beat schedule file\ncelerybeat-schedule\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# Cython eval code\n*.c\n*.html\n\n# OS X\n.DS_Store\n.Spotlight-V100\n.Trashes\n._*\n\n# ReID\nreid-data/\nlog/\nsaved-models/\nmodel-zoo/\npretrained_models/\ndebug*\n/.idea/\n/configs_user/\n" }, { "path": ".isort.cfg", "content": "[isort]\nline_length=79\nmulti_line_output=3\nlength_sort=true\nknown_standard_library=numpy,setuptools\nknown_myself=torchreid\nknown_third_party=matplotlib,cv2,torch,torchvision,PIL,yacs\nno_lines_before=STDLIB,THIRDPARTY\nsections=FUTURE,STDLIB,THIRDPARTY,myself,FIRSTPARTY,LOCALFOLDER\ndefault_section=FIRSTPARTY" }, { "path": ".style.yapf", "content": "[style]\nBASED_ON_STYLE = pep8\nBLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true\nSPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true\nDEDENT_CLOSING_BRACKETS = true\nSPACES_BEFORE_COMMENT = 1\nARITHMETIC_PRECEDENCE_INDICATION = true" }, { "path": "LICENSE", "content": "**HIPPOCRATIC LICENSE**\n\n**Version 3.0, October 2021**\n\n\n\n**TERMS AND CONDITIONS**\n\nTERMS AND CONDITIONS FOR USE, COPY, MODIFICATION, PREPARATION OF DERIVATIVE WORK, REPRODUCTION, AND DISTRIBUTION:\n\n**[1.](#1) DEFINITIONS:**\n\n_This section defines certain terms used throughout this license agreement._\n\n[1.1.](#1.1) “License” means the terms and conditions, as stated herein, for use, copy, modification, preparation of derivative work, reproduction, and distribution of Software (as defined below).\n\n[1.2.](#1.2) “Licensor” means the copyright and/or patent owner or entity authorized by the copyright and/or patent owner that is granting the License.\n\n[1.3.](#1.3) “Licensee” means the individual or entity exercising permissions granted by this License, including the use, copy, modification, preparation of derivative work, reproduction, and distribution of Software (as defined below).\n\n[1.4.](#1.4) “Software” means any copyrighted work, including but not limited to software code, authored by Licensor and made available under this License.\n\n[1.5.](#1.5) “Supply Chain” means the sequence of processes involved in the production and/or distribution of a commodity, good, or service offered by the Licensee.\n\n[1.6.](#1.6) “Supply Chain Impacted Party” or “Supply Chain Impacted Parties” means any person(s) directly impacted by any of Licensee’s Supply Chain, including the practices of all persons or entities within the Supply Chain prior to a good or service reaching the Licensee.\n\n[1.7.](#1.7) “Duty of Care” is defined by its use in tort law, delict law, and/or similar bodies of law closely related to tort and/or delict law, including without limitation, a requirement to act with the watchfulness, attention, caution, and prudence that a reasonable person in the same or similar circumstances would use towards any Supply Chain Impacted Party.\n\n[1.8.](#1.8) “Worker” is defined to include any and all permanent, temporary, and agency workers, as well as piece-rate, salaried, hourly paid, legal young (minors), part-time, night, and migrant workers.\n\n**[2.](#2) INTELLECTUAL PROPERTY GRANTS:**\n\n_This section identifies intellectual property rights granted to a Licensee_.\n\n[2.1.](#2.1) _Grant of Copyright License_: Subject to the terms and conditions of this License, Licensor hereby grants to Licensee a worldwide, non-exclusive, no-charge, royalty-free copyright license to use, copy, modify, prepare derivative work, reproduce, or distribute the Software, Licensor authored modified software, or other work derived from the Software.\n\n[2.2.](#2.2) _Grant of Patent License_: Subject to the terms and conditions of this License, Licensor hereby grants Licensee a worldwide, non-exclusive, no-charge, royalty-free patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer Software.\n\n**[3.](#3) ETHICAL STANDARDS:**\n\n_This section lists conditions the Licensee must comply with in order to have rights under this License._\n\nThe rights granted to the Licensee by this License are expressly made subject to the Licensee’s ongoing compliance with the following conditions:\n\n* [3.1.](#3.1) The Licensee SHALL NOT, whether directly or indirectly, through agents or assigns:\n * [3.1.1.](#3.1.1) Infringe upon any person’s right to life or security of person, engage in extrajudicial killings, or commit murder, without lawful cause (See Article 3, _United Nations Universal Declaration of Human Rights_; Article 6, _International Covenant on Civil and Political Rights_)\n * [3.1.2.](#3.1.2) Hold any person in slavery, servitude, or forced labor (See Article 4, _United Nations Universal Declaration of Human Rights_; Article 8, _International Covenant on Civil and Political Rights_);\n * [3.1.3.](#3.1.3) Contribute to the institution of slavery, slave trading, forced labor, or unlawful child labor (See Article 4, _United Nations Universal Declaration of Human Rights_; Article 8, _International Covenant on Civil and Political Rights_);\n * [3.1.4.](#3.1.4) Torture or subject any person to cruel, inhumane, or degrading treatment or punishment (See Article 5, _United Nations Universal Declaration of Human Rights_; Article 7, _International Covenant on Civil and Political Rights_);\n * [3.1.5.](#3.1.5) Discriminate on the basis of sex, gender, sexual orientation, race, ethnicity, nationality, religion, caste, age, medical disability or impairment, and/or any other like circumstances (See Article 7, _United Nations Universal Declaration of Human Rights_; Article 2, _International Covenant on Economic, Social and Cultural Rights_; Article 26, _International Covenant on Civil and Political Rights_);\n * [3.1.6.](#3.1.6) Prevent any person from exercising his/her/their right to seek an effective remedy by a competent court or national tribunal (including domestic judicial systems, international courts, arbitration bodies, and other adjudicating bodies) for actions violating the fundamental rights granted to him/her/them by applicable constitutions, applicable laws, or by this License (See Article 8, _United Nations Universal Declaration of Human Rights_; Articles 9 and 14, _International Covenant on Civil and Political Rights_);\n * [3.1.7.](#3.1.7) Subject any person to arbitrary arrest, detention, or exile (See Article 9, _United Nations Universal Declaration of Human Rights_; Article 9, _International Covenant on Civil and Political Rights_);\n * [3.1.8.](#3.1.8) Subject any person to arbitrary interference with a person’s privacy, family, home, or correspondence without the express written consent of the person (See Article 12, _United Nations Universal Declaration of Human Rights_; Article 17, _International Covenant on Civil and Political Rights_);\n * [3.1.9.](#3.1.9) Arbitrarily deprive any person of his/her/their property (See Article 17, _United Nations Universal Declaration of Human Rights_);\n * [3.1.10.](#3.1.10) Forcibly remove indigenous peoples from their lands or territories or take any action with the aim or effect of dispossessing indigenous peoples from their lands, territories, or resources, including without limitation the intellectual property or traditional knowledge of indigenous peoples, without the free, prior, and informed consent of indigenous peoples concerned (See Articles 8 and 10, _United Nations Declaration on the Rights of Indigenous Peoples_);\n * [3.1.11.](#3.1.11) _Mass Surveillance_: Be a government agency or multinational corporation, or a representative, agent, affiliate, successor, attorney, or assign of a government or multinational corporation, which participates in mass surveillance programs;\n * [3.1.12.](#3.1.12) _Military Activities_: Be an entity or a representative, agent, affiliate, successor, attorney, or assign of an entity which conducts military activities;\n * [3.1.13.](#3.1.13) _Law Enforcement_: Be an individual or entity, or a or a representative, agent, affiliate, successor, attorney, or assign of an individual or entity, that provides good or services to, or otherwise enters into any commercial contracts with, any local, state, or federal law enforcement agency;\n * [3.1.14.](#3.1.14) _Media_: Be an individual or entity, or a or a representative, agent, affiliate, successor, attorney, or assign of an individual or entity, that broadcasts messages promoting killing, torture, or other forms of extreme violence;\n * [3.1.15.](#3.1.15) Interfere with Workers' free exercise of the right to organize and associate (See Article 20, United Nations Universal Declaration of Human Rights; C087 - Freedom of Association and Protection of the Right to Organise Convention, 1948 (No. 87), International Labour Organization; Article 8, International Covenant on Economic, Social and Cultural Rights); and\n * [3.1.16.](#3.1.16) Harm the environment in a manner inconsistent with local, state, national, or international law.\n* [3.2.](#3.2) The Licensee SHALL:\n * [3.2.1.](#3.2.1) _Social Auditing_: Only use social auditing mechanisms that adhere to Worker-Driven Social Responsibility Network’s Statement of Principles () over traditional social auditing mechanisms, to the extent the Licensee uses any social auditing mechanisms at all;\n * [3.2.2.](#3.2.2) Provide equal pay for equal work where the performance of such work requires equal skill, effort, and responsibility, and which are performed under similar working conditions, except where such payment is made pursuant to:\n * [3.2.2.1.](#3.2.2.1) A seniority system;\n * [3.2.2.2.](#3.2.2.2) A merit system;\n * [3.2.2.3.](#3.2.2.3) A system which measures earnings by quantity or quality of production; or\n * [3.2.2.4.](#3.2.2.4) A differential based on any other factor other than sex, gender, sexual orientation, race, ethnicity, nationality, religion, caste, age, medical disability or impairment, and/or any other like circumstances (See 29 U.S.C.A. § 206(d)(1); Article 23, _United Nations Universal Declaration of Human Rights_; Article 7, _International Covenant on Economic, Social and Cultural Rights_; Article 26, _International Covenant on Civil and Political Rights_); and\n * [3.2.3.](#3.2.3) Allow for reasonable limitation of working hours and periodic holidays with pay (See Article 24, _United Nations Universal Declaration of Human Rights_; Article 7, _International Covenant on Economic, Social and Cultural Rights_).\n\n**[4.](#4) SUPPLY CHAIN IMPACTED PARTIES:**\n\n_This section identifies additional individuals or entities that a Licensee could harm as a result of violating the Ethical Standards section, the condition that the Licensee must voluntarily accept a Duty of Care for those individuals or entities, and the right to a private right of action that those individuals or entities possess as a result of violations of the Ethical Standards section._\n\n[4.1.](#4.1) In addition to the above Ethical Standards, Licensee voluntarily accepts a Duty of Care for Supply Chain Impacted Parties of this License, including individuals and communities impacted by violations of the Ethical Standards. The Duty of Care is breached when a provision within the Ethical Standards section is violated by a Licensee, one of its successors or assigns, or by an individual or entity that exists within the Supply Chain prior to a good or service reaching the Licensee.\n\n[4.2.](#4.2) Breaches of the Duty of Care, as stated within this section, shall create a private right of action, allowing any Supply Chain Impacted Party harmed by the Licensee to take legal action against the Licensee in accordance with applicable negligence laws, whether they be in tort law, delict law, and/or similar bodies of law closely related to tort and/or delict law, regardless if Licensee is directly responsible for the harms suffered by a Supply Chain Impacted Party. Nothing in this section shall be interpreted to include acts committed by individuals outside of the scope of his/her/their employment.\n\n[5.](#5) **NOTICE:** _This section explains when a Licensee must notify others of the License._\n\n[5.1.](#5.1) _Distribution of Notice_: Licensee must ensure that everyone who receives a copy of or uses any part of Software from Licensee, with or without changes, also receives the License and the copyright notice included with Software (and if included by the Licensor, patent, trademark, and attribution notice). Licensee must ensure that License is prominently displayed so that any individual or entity seeking to download, copy, use, or otherwise receive any part of Software from Licensee is notified of this License and its terms and conditions. Licensee must cause any modified versions of the Software to carry prominent notices stating that Licensee changed the Software.\n\n[5.2.](#5.2) _Modified Software_: Licensee is free to create modifications of the Software and distribute only the modified portion created by Licensee, however, any derivative work stemming from the Software or its code must be distributed pursuant to this License, including this Notice provision.\n\n[5.3.](#5.3) _Recipients as Licensees_: Any individual or entity that uses, copies, modifies, reproduces, distributes, or prepares derivative work based upon the Software, all or part of the Software’s code, or a derivative work developed by using the Software, including a portion of its code, is a Licensee as defined above and is subject to the terms and conditions of this License.\n\n**[6.](#6) REPRESENTATIONS AND WARRANTIES:**\n\n[6.1.](#6.1) _Disclaimer of Warranty_: TO THE FULL EXTENT ALLOWED BY LAW, THIS SOFTWARE COMES “AS IS,” WITHOUT ANY WARRANTY, EXPRESS OR IMPLIED, AND LICENSOR SHALL NOT BE LIABLE TO ANY PERSON OR ENTITY FOR ANY DAMAGES OR OTHER LIABILITY ARISING FROM, OUT OF, OR IN CONNECTION WITH THE SOFTWARE OR THIS LICENSE, UNDER ANY LEGAL CLAIM.\n\n[6.2.](#6.2) _Limitation of Liability_: LICENSEE SHALL HOLD LICENSOR HARMLESS AGAINST ANY AND ALL CLAIMS, DEBTS, DUES, LIABILITIES, LIENS, CAUSES OF ACTION, DEMANDS, OBLIGATIONS, DISPUTES, DAMAGES, LOSSES, EXPENSES, ATTORNEYS' FEES, COSTS, LIABILITIES, AND ALL OTHER CLAIMS OF EVERY KIND AND NATURE WHATSOEVER, WHETHER KNOWN OR UNKNOWN, ANTICIPATED OR UNANTICIPATED, FORESEEN OR UNFORESEEN, ACCRUED OR UNACCRUED, DISCLOSED OR UNDISCLOSED, ARISING OUT OF OR RELATING TO LICENSEE’S USE OF THE SOFTWARE. NOTHING IN THIS SECTION SHOULD BE INTERPRETED TO REQUIRE LICENSEE TO INDEMNIFY LICENSOR, NOR REQUIRE LICENSOR TO INDEMNIFY LICENSEE.\n\n**[7.](#7) TERMINATION**\n\n[7.1.](#7.1) _Violations of Ethical Standards or Breaching Duty of Care_: If Licensee violates the Ethical Standards section or Licensee, or any other person or entity within the Supply Chain prior to a good or service reaching the Licensee, breaches its Duty of Care to Supply Chain Impacted Parties, Licensee must remedy the violation or harm caused by Licensee within 30 days of being notified of the violation or harm. If Licensee fails to remedy the violation or harm within 30 days, all rights in the Software granted to Licensee by License will be null and void as between Licensor and Licensee.\n\n[7.2.](#7.2) _Failure of Notice_: If any person or entity notifies Licensee in writing that Licensee has not complied with the Notice section of this License, Licensee can keep this License by taking all practical steps to comply within 30 days after the notice of noncompliance. If Licensee does not do so, Licensee’s License (and all rights licensed hereunder) will end immediately.\n\n[7.3.](#7.3) _Judicial Findings_: In the event Licensee is found by a civil, criminal, administrative, or other court of competent jurisdiction, or some other adjudicating body with legal authority, to have committed actions which are in violation of the Ethical Standards or Supply Chain Impacted Party sections of this License, all rights granted to Licensee by this License will terminate immediately.\n\n[7.4.](#7.4) _Patent Litigation_: If Licensee institutes patent litigation against any entity (including a cross-claim or counterclaim in a suit) alleging that the Software, all or part of the Software’s code, or a derivative work developed using the Software, including a portion of its code, constitutes direct or contributory patent infringement, then any patent license, along with all other rights, granted to Licensee under this License will terminate as of the date such litigation is filed.\n\n[7.5.](#7.5) _Additional Remedies_: Termination of the License by failing to remedy harms in no way prevents Licensor or Supply Chain Impacted Party from seeking appropriate remedies at law or in equity.\n\n**[8.](#8) MISCELLANEOUS:**\n\n[8.1.](#8.1) _Conditions_: Sections 3, 4.1, 5.1, 5.2, 7.1, 7.2, 7.3, and 7.4 are conditions of the rights granted to Licensee in the License.\n\n[8.2.](#8.2) _Equitable Relief_: Licensor and any Supply Chain Impacted Party shall be entitled to equitable relief, including injunctive relief or specific performance of the terms hereof, in addition to any other remedy to which they are entitled at law or in equity.\n\n[8.3.](#8.3) _Severability_: If any term or provision of this License is determined to be invalid, illegal, or unenforceable by a court of competent jurisdiction, any such determination of invalidity, illegality, or unenforceability shall not affect any other term or provision of this License or invalidate or render unenforceable such term or provision in any other jurisdiction. If the determination of invalidity, illegality, or unenforceability by a court of competent jurisdiction pertains to the terms or provisions contained in the Ethical Standards section of this License, all rights in the Software granted to Licensee shall be deemed null and void as between Licensor and Licensee.\n\n[8.4.](#8.4) _Section Titles_: Section titles are solely written for organizational purposes and should not be used to interpret the language within each section.\n\n[8.5.](#8.5) _Citations_: Citations are solely written to provide context for the source of the provisions in the Ethical Standards.\n\n[8.6.](#8.6) _Section Summaries_: Some sections have a brief _italicized description_ which is provided for the sole purpose of briefly describing the section and should not be used to interpret the terms of the License.\n\n[8.7.](#8.7) _Entire License_: This is the entire License between the Licensor and Licensee with respect to the claims released herein and that the consideration stated herein is the only consideration or compensation to be paid or exchanged between them for this License. This License cannot be modified or amended except in a writing signed by Licensor and Licensee.\n\n[8.8.](#8.8) _Successors and Assigns_: This License shall be binding upon and inure to the benefit of the Licensor’s and Licensee’s respective heirs, successors, and assigns.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "\n\n# BPBReID: Body Part-based (Occluded) Re-Identification\n\n**A strong baseline for body part-based person re-identification** \n\n🔥 *Our new work on [Keypoint Promptable ReID](https://github.com/VlSomers/keypoint_promptable_reidentification) was accepted at ECCV24* 🔥\n\n[[Paper](https://arxiv.org/abs/2211.03679)] [[Video](https://www.youtube.com/watch?v=4NQump-vg_A&ab_channel=VladimirSomers)] [[Poster](docs/figures/bpbreid/wacv23_poster_bpbreid.pdf)]\n\n[![arXiv](https://img.shields.io/badge/arXiv-2211.03679-.svg)](https://arxiv.org/abs/2211.03679) [![Hippocratic License HL3-LAW-MEDIA-MIL-SOC-SV](https://img.shields.io/static/v1?label=Hippocratic%20License&message=HL3-LAW-MEDIA-MIL-SOC-SV&labelColor=5e2751&color=bc8c3d)](https://firstdonoharm.dev/version/3/0/law-media-mil-soc-sv.html)\n\n>**[Body Part-Based Representation Learning for Occluded Person Re-Identification, WACV23](https://arxiv.org/abs/2211.03679)**\n>\n>Vladimir Somers, Christophe De Vleeschouwer, Alexandre Alahi\n>\n>[*arxiv 2211.03679*](https://arxiv.org/abs/2211.03679)\n>\n\n### State-of-the-art performance on 5 datasets:\n\nOccluded-Duke: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/body-part-based-representation-learning-for/person-re-identification-on-occluded-dukemtmc)](https://paperswithcode.com/sota/person-re-identification-on-occluded-dukemtmc?p=body-part-based-representation-learning-for)\n\nOccluded ReID: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/body-part-based-representation-learning-for/person-re-identification-on-occluded-reid-1)](https://paperswithcode.com/sota/person-re-identification-on-occluded-reid-1?p=body-part-based-representation-learning-for)\n\nP-DukeMTMC: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/body-part-based-representation-learning-for/person-re-identification-on-p-dukemtmc-reid)](https://paperswithcode.com/sota/person-re-identification-on-p-dukemtmc-reid?p=body-part-based-representation-learning-for)\n\nDukeMTMC-ReID: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/body-part-based-representation-learning-for/person-re-identification-on-dukemtmc-reid)](https://paperswithcode.com/sota/person-re-identification-on-dukemtmc-reid?p=body-part-based-representation-learning-for)\n\nMarket1501: [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/body-part-based-representation-learning-for/person-re-identification-on-market-1501)](https://paperswithcode.com/sota/person-re-identification-on-market-1501?p=body-part-based-representation-learning-for)\n\n \n

\n \n\n## News\n- [2025.10.25] Added a discussion on the Theoretical Limitations of Global Embeddings in ReID Under Partial Observations\n- [2024.08.23] 🚀🔥 Our new work on [Keypoint Promptable ReID](https://arxiv.org/abs/2407.18112) was accepted to ECCV24, full codebase available [here](https://github.com/VlSomers/keypoint_promptable_reidentification).\n- [2023.09.20] New paper and big update coming soon 🚀 ...\n- [2023.07.26] The Python script from @samihormi to generate human parsing labels based on PifPaf and MaskRCNN has been released, have a look at the \"Generate human parsing labels\" section below. This script is different from the one used by the authors (especially when facing multiple pedestrians in a single image): resulting human parsing labels will not be exactly the same.\n- [2023.06.28] Please find a non-official script to generate human parsing labels from PifPaf and MaskRCNN in this [Pull Request](https://github.com/VlSomers/bpbreid/pull/18). The PR will be merged soon.\n- [2022.12.02] We release the first version of our codebase. Please update frequently as we will add more documentation during the next few weeks.\n\n## What's next\nWe plan on extending BPBReID in the near future, put a star and stay updated for future changes:\n- part-based video/tracklet reid\n- part-based reid for multi-object tracking\n- ...\n\n\n\n## Table of content\n- [BPBReID: Body Part-based Re-Identification](#bpbreid--body-part-based-re-identification)\n * [News](#news)\n * [What's next](#what-s-next)\n * [Table of content](#table-of-content)\n * [Introduction](#introduction)\n * [What to find in this repository](#what-to-find-in-this-repository)\n * [Discussion on the Theoretical Limitations of Global Embeddings in ReID Under Partial Observations](#discussion-on-the-theoretical-limitations-of-global-embeddings-in-reid-under-partial-observationsRetry)\n * [Instructions](#instructions)\n + [Installation](#installation)\n + [Download human parsing labels](#download-human-parsing-labels)\n + [Generate human parsing labels](#generate-human-parsing-labels)\n + [Download the pre-trained models](#download-the-pre-trained-models)\n + [Inference](#inference)\n + [Training](#training)\n + [Visualization tools](#visualization-tools)\n * [Other works](#other-works)\n * [Questions and suggestions](#questions-and-suggestions)\n * [Citation](#citation)\n * [Acknowledgement](#acknowledgement)\n\n## Introduction\nWelcome to the official repository for our WACV23 paper \"_Body Part-Based Representation Learning for Occluded Person Re-Identification_\".\nIn this work, we propose BPBReID, a part-based method for person re-identification using body part feature representations to compute to similarity between two samples.\nAs illustrated in the figure below, **part-based** ReID methods output multiple features per input sample, i.e. one for each part, whereas standard global methods only output a single feature.\nCompared to global methods, part-based ones come with some advantages:\n\n1. They achieve explicit appearance feature alignement for better ReID accuracy.\n2. They are robust to occlusions, since only mutually visible parts are used when comparing two samples.\n\nOur model BPBreID uses pseudo human parsing labels at training time to learn an attention mechanism. This attention mechanism has K branches to pool the global spatial feature map into K body part-based embeddings. Based on the attention maps activations, visibility scores are computed for each part.\nAt test time, no human parsing labels is required.\nThe final similarity score between two samples is computed using the average distance of all mutually visible part-based embeddings.\nPlease refer to [our paper](https://arxiv.org/abs/2211.03679) for more information.\n\n\n\n\n## What to find in this repository\nIn this repository, we propose a framework and a strong baseline to support further research on part-based ReID methods.\nOur code is based on the popular [Torchreid](https://github.com/KaiyangZhou/deep-person-reid) framework for person re-identification.\nIn this codebase, we provide several adaptations to the original framework to support part-based ReID methods: \n- The [ImagePartBasedEngine](torchreid/engine/image/part_based_engine.py) to train/test part-based models, compute query-gallery distance matrix using multiple features per test sample with support for visibility scores.\n- The fully configurable [GiLt loss](/torchreid/losses/GiLt_loss.py) to selectively apply id/triplet loss on holistics (global) and part-based features.\n- The [BodyPartAttentionLoss](torchreid/losses/body_part_attention_loss.py) to train the attention mechanism.\n- The [BPBreID](torchreid/models/bpbreid.py) part-based model to compute part-based features with support for body-part learnable attention, fixed attention heatmaps from an external model, PCB-like horizontal stripes, etc.\n- The [Albumentation](https://albumentations.ai/) data augmentation library used for data augmentation, with support for external heatmaps/masks transforms.\n- Support for [Weights & Biases](https://wandb.ai/site) and other logging tools in the [Logger](torchreid/utils/logging/logger.py) class.\n- An [EngineState](torchreid/utils/engine_state.py) class to keep track of training epoch, etc.\n- A new [ranking visualization](torchreid/utils/visualization/visualize_query_gallery_rankings.py) tool to display part heatmaps, local distance for each part and other metrics.\n- For more information about all available configuration and parameters, please have a look at the [default config file](torchreid/scripts/default_config.py).\n\nYou can also have a look at the original [Torchreid README](Torchreid_original_README.rst) for additional information, such as documentation, how-to instructions, etc.\nBe aware that some of the original Torchreid functionnality and models might be broken (for example, we don't support video re-id yet).\n\n\n## Discussion on the Theoretical Limitations of Global Embeddings in ReID Under Partial Observations (Occlusions) \nCurrent ReID models aim to learn a single global embedding where images of the same person cluster together. \nHowever, this paradigm faces an inherent paradox when dealing with partial observations. \nConsider three images of the same person: a full-body image (A) and two occluded variations - an upper-body-only image (B) and a lower-body-only image (C).\nThis becomes particularly problematic when comparing B with C - they share no common visible features despite representing the same person. \nThe global embedding approach implicitly assumes transitivity (if A=B and B=C, then A=C), but this property breaks down under partial observations, revealing a fundamental flaw in the single embedding space paradigm.\nThis observation suggests that the traditional approach of mapping all images to a unified embedding space may be fundamentally flawed. ReID models do not learn truly identity-centric representations, but rather appearance-centric representations that serve as a proxy to model identity. These appearance-centric representations are learned to be invariant to pose, lighting, and viewpoint changes, but when faced with partial observations, expecting such models to produce consistent embeddings becomes theoretically questionable.\nThis theoretical limitation is reflected in practice, where part-based representation learning, which compares only mutually visible regions, offer a more principled solution to this ambiguity.\nRather than being a technical hack to improve image retrieval performance, they may represent a fundamental requirement for representation learning under partial observation constraints.\nNevertheless, this apparent limitation of deep metric learning deserves further investigation. We identify two promising research directions. First, developing generic part-based architectures for representation learning that extend beyond human ReID and dynamically align comparable features while discarding missing information across object pairs. Such generic part-based methods would learn to build representations under partial observability without priors on the input object type, making them truly universal solutions. \nSecond, moving beyond representation learning entirely by designing ReID models that directly process image pairs to compute similarity scores, thus avoiding the ambiguity of intermediate representation in an embedding space under partial observations. \nIn this paradigm, the network would inherently learn to compare only mutually visible parts of the input image pair.\nThis approach faces however two key limitations compared to the first one: the need to run the neural network for every possible image pair, leading to quadratic computational complexity with gallery size (versus running the network once per image in representation learning), and the challenge of quantifying the confidence in the output similarity score when comparing partially visible objects.\n\n\n## Instructions\n### Installation\nMake sure [conda](https://www.anaconda.com/distribution/) is installed.\n\n # clone this repository\n git clone https://github.com/VlSomers/bpbreid\n\n # create conda environment\n cd bpbreid/ # enter project folder\n conda create --name bpbreid python=3.10\n conda activate bpbreid\n \n # install dependencies\n # make sure `which python` and `which pip` point to the correct path\n pip install -r requirements.txt\n \n # install torch and torchvision (select the proper cuda version to suit your machine)\n conda install pytorch torchvision cudatoolkit=9.0 -c pytorch\n \n # install torchreid (don't need to re-build it if you modify the source code)\n python setup.py develop\n \n\n### Download human parsing labels\nYou can download the human parsing labels on [GDrive](https://drive.google.com/drive/folders/1IbCAbjj3XtV3_tFOsCuqBi79ZiDqNc1H?usp=sharing). \nThese labels were generated using the [PifPaf](https://github.com/openpifpaf/openpifpaf) pose estimation model and then filtered using segmentation masks from [Mask-RCNN](https://github.com/facebookresearch/detectron2).\nWe provide the labels for five datasets: **Market-1501**, **DukeMTMC-reID**, **Occluded-Duke**, **Occluded-ReID** and **P-DukeMTMC**.\nAfter downloading, unzip the file and put the `masks` folder under the corresponding dataset directory.\nFor instance, Market-1501 should look like this:\n\n Market-1501-v15.09.15\n ├── bounding_box_test\n ├── bounding_box_train\n ├── masks\n │   └── pifpaf_maskrcnn_filtering\n │   ├── bounding_box_test\n │   ├── bounding_box_train\n │   └── query\n └── query\n\nMake also sure to set `data.root` config to your dataset root directory path, i.e., all your datasets folders (`Market-1501-v15.09.15`, `DukeMTMC-reID`, `Occluded_Duke`, `P-DukeMTMC-reID`, `Occluded_REID`) should be under this path.\n\n\n### Generate human parsing labels\n\nYou can create human parsing labels for your own dataset using the following command:\n\n conda activate bpbreid\n python torchreid/scripts/get_labels --source [Dataset Path] \n\nThe labels will be saved under the source directory in the *masks* folder as per the code convention.\n\n\n### Download the pre-trained models\nWe also provide some [state-of-the-art pre-trained models](https://drive.google.com/drive/folders/1aUjpSXXVGtAh2nzV0RVsCq0tTXuDZWoH?usp=sharing) based on the HRNet-W32 backbone.\nYou can put the downloaded weights under a 'pretrained_models/' directory or specify the path to the pre-trained weights using the `model.load_weights` parameter in the `yaml` config.\nThe configuration used to obtain the pre-trained weights is also saved within the `.pth` file: make sure to set `model.load_config` to `True` so that the parameters under the `model.bpbreid` part of the configuration tree will be loaded from this file.\n\n### Inference\nYou can test the above downloaded models using the following command:\n\n conda activate bpbreid\n python torchreid/scripts/main.py --config-file configs/bpbreid/bpbreid__test.yaml\n \nFor instance, for the Market-1501 dataset:\n\n conda activate bpbreid\n python torchreid/scripts/main.py --config-file configs/bpbreid/bpbreid_market1501_test.yaml\n \nConfiguration files for other datasets are available under `configs/bpbreid/`.\nMake sure the `model.load_weights` in these `yaml` config files points to the pre-trained weights you just downloaded. \n\n### Training\nTraining configs for five datasets (Market-1501, DukeMTMC-reID, Occluded-Duke, Occluded-ReID and P-DukeMTMC) are provided in the `configs/bpbreid/` folder. \nA training procedure can be launched with:\n\n conda activate bpbreid\n python ./torchreid/scripts/main.py --config-file configs/bpbreid/bpbreid__train.yaml\n \nFor instance, for the Occluded-Duke dataset:\n\n conda activate bpbreid\n python torchreid/scripts/main.py --config-file configs/bpbreid/bpbreid_occ_duke_train.yaml\n\nMake sure to download and install the human parsing labels for your training dataset before runing this command.\n\n### Visualization tools\nThe ranking visualization tool can be activated by setting the `test.visrank` config to `True`.\nAs illustrated below, this tool displays the Top-K ranked samples as rows (K can be set via `test.visrank_topk`). The first row with blue background is the query, and the following green/red rows indicated correct/incorrect matches.\nThe attention maps for each test embedding (foreground, parts, etc) are displayed in the row.\nAn attention map has a green/red border when it is visible/unvisible. \nThe first number under each attention map indicate the visibility score and the second number indicate the distance of the embedding to the corresponding query embedding.\nThe distances under the images in the first column on the left are the global distances of that sample to the query, which is usually computed as the average of all other distances weighted by the visibility score.\nIf you need more information about the visualization tool, fell free to open an issue.\n\n \n

\n \n\n## Other works\nIf you are looking for datasets to evaluate your re-identification models, please have a look at our other works on player re-identification for team sport events:\n- [CAMELTrack: Context-Aware Multi-cue ExpLoitation for Online Multi-Object Tracking](https://github.com/TrackingLaboratory/CAMELTrack)\n- The [SoccerNet Player Re-Identification](https://github.com/SoccerNet/sn-reid) dataset\n- The [DeepSportRadar Player Re-Identification](https://github.com/DeepSportRadar/player-reidentification-challenge) dataset \n\n

\n \n \n

\n\n## Questions and suggestions\nIf you have any question/suggestion, or find any bug/issue with the code, please raise a GitHub issue in this repository, I'll be glab to help you as much as I can!\nI'll try to update the documentation regularly based on your questions. \n\n\n## Citation\nIf you use this repository for your research or wish to refer to our method [BPBReID](https://arxiv.org/abs/2211.03679), please use the following BibTeX entry:\n```\n@article{bpbreid,\n archivePrefix = {arXiv},\n arxivId = {2211.03679},\n author = {Somers, Vladimir and {De Vleeschouwer}, Christophe and Alahi, Alexandre},\n doi = {10.48550/arxiv.2211.03679},\n eprint = {2211.03679},\n isbn = {2211.03679v1},\n journal = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV23)},\n month = {nov},\n title = {{Body Part-Based Representation Learning for Occluded Person Re-Identification}},\n url = {https://arxiv.org/abs/2211.03679v1 http://arxiv.org/abs/2211.03679},\n year = {2023}\n}\n```\n\n## Acknowledgement\nThis codebase is a fork from [Torchreid](https://github.com/KaiyangZhou/deep-person-reid)\n\n\n" }, { "path": "Torchreid_original_README.rst", "content": "Torchreid\n===========\nTorchreid is a library for deep-learning person re-identification, written in `PyTorch `_.\n\nIt features:\n\n- multi-GPU training\n- support both image- and video-reid\n- end-to-end training and evaluation\n- incredibly easy preparation of reid datasets\n- multi-dataset training\n- cross-dataset evaluation\n- standard protocol used by most research papers\n- highly extensible (easy to add models, datasets, training methods, etc.)\n- implementations of state-of-the-art deep reid models\n- access to pretrained reid models\n- advanced training techniques\n- visualization tools (tensorboard, ranks, etc.)\n\n\nCode: https://github.com/KaiyangZhou/deep-person-reid.\n\nDocumentation: https://kaiyangzhou.github.io/deep-person-reid/.\n\nHow-to instructions: https://kaiyangzhou.github.io/deep-person-reid/user_guide.\n\nModel zoo: https://kaiyangzhou.github.io/deep-person-reid/MODEL_ZOO.\n\nTech report: https://arxiv.org/abs/1910.10093.\n\nYou can find some research projects that are built on top of Torchreid `here `_.\n\n\nWhat's new\n------------\n- [May 2020] Added the person attribute recognition code used in [Omni-Scale Feature Learning for Person Re-Identification (ICCV'19)](https://arxiv.org/abs/1905.00953).\n- [May 2020] ``1.2.1``: Added a simple API for feature extraction (``torchreid/utils/feature_extractor.py``). See the `documentation `_ for the instruction.\n- [Apr 2020] Code for reproducing the experiments of `deep mutual learning `_ in the `OSNet paper `__ (Supp. B) has been released at ``projects/DML``.\n- [Apr 2020] Upgraded to ``1.2.0``. The engine class has been made more model-agnostic to improve extensibility. See `Engine `_ and `ImageSoftmaxEngine `_ for more details. Credit to `Dassl.pytorch `_.\n- [Dec 2019] Our `OSNet paper `_ has been updated, with additional experiments (in section B of the supplementary) showing some useful techniques for improving OSNet's performance in practice.\n- [Nov 2019] ``ImageDataManager`` can load training data from target datasets by setting ``load_train_targets=True``, and the train-loader can be accessed with ``train_loader_t = datamanager.train_loader_t``. This feature is useful for domain adaptation research.\n\n\nInstallation\n---------------\n\nMake sure `conda `_ is installed.\n\n\n.. code-block:: bash\n\n # cd to your preferred directory and clone this repo\n git clone https://github.com/KaiyangZhou/deep-person-reid.git\n\n # create environment\n cd deep-person-reid/\n conda create --name torchreid python=3.7\n conda activate torchreid\n\n # install dependencies\n # make sure `which python` and `which pip` point to the correct path\n pip install -r requirements.txt\n\n # install torch and torchvision (select the proper cuda version to suit your machine)\n conda install pytorch torchvision cudatoolkit=9.0 -c pytorch\n\n # install torchreid (don't need to re-build it if you modify the source code)\n python setup.py develop\n\n\nGet started: 30 seconds to Torchreid\n-------------------------------------\n1. Import ``torchreid``\n\n.. code-block:: python\n \n import torchreid\n\n2. Load data manager\n\n.. code-block:: python\n \n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources='market1501',\n targets='market1501',\n height=256,\n width=128,\n batch_size_train=32,\n batch_size_test=100,\n transforms=['random_flip', 'random_crop']\n )\n\n3 Build model, optimizer and lr_scheduler\n\n.. code-block:: python\n \n model = torchreid.models.build_model(\n name='resnet50',\n num_classes=datamanager.num_train_pids,\n loss='softmax',\n pretrained=True\n )\n\n model = model.cuda()\n\n optimizer = torchreid.optim.build_optimizer(\n model,\n optim='adam',\n lr=0.0003\n )\n\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer,\n lr_scheduler='single_step',\n stepsize=20\n )\n\n4. Build engine\n\n.. code-block:: python\n \n engine = torchreid.engine.ImageSoftmaxEngine(\n datamanager,\n model,\n optimizer=optimizer,\n scheduler=scheduler,\n label_smooth=True\n )\n\n5. Run training and test\n\n.. code-block:: python\n \n engine.run(\n save_dir='log/resnet50',\n max_epoch=60,\n eval_freq=10,\n test_only=False\n )\n\n\nA unified interface\n-----------------------\nIn \"deep-person-reid/scripts/\", we provide a unified interface to train and test a model. See \"scripts/main.py\" and \"scripts/default_config.py\" for more details. The folder \"configs/\" contains some predefined configs which you can use as a starting point.\n\nBelow we provide an example to train and test `OSNet (Zhou et al. ICCV'19) `_. Assume :code:`PATH_TO_DATA` is the directory containing reid datasets. The environmental variable :code:`CUDA_VISIBLE_DEVICES` is omitted, which you need to specify if you have a pool of gpus and want to use a specific set of them.\n\nConventional setting\n^^^^^^^^^^^^^^^^^^^^^\n\nTo train OSNet on Market1501, do\n\n.. code-block:: bash\n\n python scripts/main.py \\\n --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \\\n --transforms random_flip random_erase \\\n --root $PATH_TO_DATA\n\n\nThe config file sets Market1501 as the default dataset. If you wanna use DukeMTMC-reID, do\n\n.. code-block:: bash\n\n python scripts/main.py \\\n --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \\\n -s dukemtmcreid \\\n -t dukemtmcreid \\\n --transforms random_flip random_erase \\\n --root $PATH_TO_DATA \\\n data.save_dir log/osnet_x1_0_dukemtmcreid_softmax_cosinelr\n\nThe code will automatically (download and) load the ImageNet pretrained weights. After the training is done, the model will be saved as \"log/osnet_x1_0_market1501_softmax_cosinelr/model.pth.tar-250\". Under the same folder, you can find the `tensorboard `_ file. To visualize the learning curves using tensorboard, you can run :code:`tensorboard --logdir=log/osnet_x1_0_market1501_softmax_cosinelr` in the terminal and visit :code:`http://localhost:6006/` in your web browser.\n\nEvaluation is automatically performed at the end of training. To run the test again using the trained model, do\n\n.. code-block:: bash\n\n python scripts/main.py \\\n --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad_cosine.yaml \\\n --root $PATH_TO_DATA \\\n model.load_weights log/osnet_x1_0_market1501_softmax_cosinelr/model.pth.tar-250 \\\n test.evaluate True\n\n\nCross-domain setting\n^^^^^^^^^^^^^^^^^^^^^\n\nSuppose you wanna train OSNet on DukeMTMC-reID and test its performance on Market1501, you can do\n\n.. code-block:: bash\n\n python scripts/main.py \\\n --config-file configs/im_osnet_x1_0_softmax_256x128_amsgrad.yaml \\\n -s dukemtmcreid \\\n -t market1501 \\\n --transforms random_flip color_jitter \\\n --root $PATH_TO_DATA\n\nHere we only test the cross-domain performance. However, if you also want to test the performance on the source dataset, i.e. DukeMTMC-reID, you can set :code:`-t dukemtmcreid market1501`, which will evaluate the model on the two datasets separately.\n\nDifferent from the same-domain setting, here we replace :code:`random_erase` with :code:`color_jitter`. This can improve the generalization performance on the unseen target dataset.\n\nPretrained models are available in the `Model Zoo `_.\n\n\nDatasets\n--------\n\nImage-reid datasets\n^^^^^^^^^^^^^^^^^^^^^\n- `Market1501 `_\n- `CUHK03 `_\n- `DukeMTMC-reID `_\n- `MSMT17 `_\n- `VIPeR `_\n- `GRID `_\n- `CUHK01 `_\n- `SenseReID `_\n- `QMUL-iLIDS `_\n- `PRID `_\n\nVideo-reid datasets\n^^^^^^^^^^^^^^^^^^^^^^^\n- `MARS `_\n- `iLIDS-VID `_\n- `PRID2011 `_\n- `DukeMTMC-VideoReID `_\n\n\nModels\n-------\n\nImageNet classification models\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- `ResNet `_\n- `ResNeXt `_\n- `SENet `_\n- `DenseNet `_\n- `Inception-ResNet-V2 `_\n- `Inception-V4 `_\n- `Xception `_\n- `IBN-Net `_\n\nLightweight models\n^^^^^^^^^^^^^^^^^^^\n- `NASNet `_\n- `MobileNetV2 `_\n- `ShuffleNet `_\n- `ShuffleNetV2 `_\n- `SqueezeNet `_\n\nReID-specific models\n^^^^^^^^^^^^^^^^^^^^^^\n- `MuDeep `_\n- `ResNet-mid `_\n- `HACNN `_\n- `PCB `_\n- `MLFN `_\n- `OSNet `_\n- `OSNet-AIN `_\n\n\nUseful links\n-------------\n- `OSNet-IBN1-Lite (test-only code with lite docker container) `_\n- `Deep Learning for Person Re-identification: A Survey and Outlook `_\n\n\nCitation\n---------\nIf you find this code useful to your research, please cite the following papers.\n\n.. code-block:: bash\n\n @article{torchreid,\n title={Torchreid: A Library for Deep Learning Person Re-Identification in Pytorch},\n author={Zhou, Kaiyang and Xiang, Tao},\n journal={arXiv preprint arXiv:1910.10093},\n year={2019}\n }\n \n @inproceedings{zhou2019osnet,\n title={Omni-Scale Feature Learning for Person Re-Identification},\n author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},\n booktitle={ICCV},\n year={2019}\n }\n\n @article{zhou2019learning,\n title={Learning Generalisable Omni-Scale Representations for Person Re-Identification},\n author={Zhou, Kaiyang and Yang, Yongxin and Cavallaro, Andrea and Xiang, Tao},\n journal={arXiv preprint arXiv:1910.06827},\n year={2019}\n }\n" }, { "path": "configs/bpbreid/bpbreid_dukemtmc_test.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['dukemtmcreid']\n targets: ['dukemtmcreid']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n load_weights: 'pretrained_models/bpbreid_dukemtmcreid_hrnet32_10669.pth'\n load_config: True\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\ntest:\n evaluate: True\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_dukemtmc_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['dukemtmcreid']\n targets: ['dukemtmcreid']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights: # SOTA weights for GiLt loss\n globl:\n id: 1.\n tr: 0.\n foreg:\n id: 1.\n tr: 1.\n conct:\n id: 1.\n tr: 0.\n parts:\n id: 0.\n tr: 1.\n pixls:\n ce: 0.35\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_market1501_test.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['market1501']\n targets: ['market1501']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n load_weights: 'pretrained_models/bpbreid_market1501_hrnet32_10642.pth'\n load_config: True\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\ntest:\n evaluate: True\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_market1501_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['market1501']\n targets: ['market1501']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights: # SOTA weights for GiLt loss\n globl:\n id: 1.\n tr: 0.\n foreg:\n id: 1.\n tr: 1.\n conct:\n id: 1.\n tr: 0.\n parts:\n id: 0.\n tr: 1.\n pixls:\n ce: 0.35\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_occ_duke_test.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['occluded_duke']\n targets: ['occluded_duke']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n load_weights: 'pretrained_models/bpbreid_occluded_duke_hrnet32_10670.pth'\n load_config: True\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\ntest:\n evaluate: True\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_occ_duke_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['occluded_duke']\n targets: ['occluded_duke']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'eight'\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights: # SOTA weights for GiLt loss\n globl:\n id: 1.\n tr: 0.\n foreg:\n id: 1.\n tr: 0.\n conct:\n id: 1.\n tr: 0.\n parts:\n id: 0.\n tr: 1.\n pixls:\n ce: 0.35\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_occ_reid_test.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['market1501']\n targets: ['occluded_reid']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n load_weights: 'pretrained_models/'\n load_config: True\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\ntest:\n evaluate: True\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_occ_reid_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['market1501']\n targets: ['market1501', 'occluded_reid']\n height: 384\n width: 128\n transforms: ['rc', 're', 'rf', 'cj']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'eight'\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights: # SOTA weights for GiLt loss\n globl:\n id: 1.\n tr: 0.\n foreg:\n id: 1.\n tr: 1.\n conct:\n id: 1.\n tr: 0.\n parts:\n id: 0.\n tr: 1.\n pixls:\n ce: 0.7\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_p_dukemtmc_test.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['p_dukemtmc_reid']\n targets: ['p_dukemtmc_reid']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n load_weights: 'pretrained_models/bpbreid_p_dukemtmc_hrnet32_10672.pth'\n load_config: True\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'five_v'\n\ntest:\n evaluate: True\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/bpbreid_p_dukemtmc_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['p_dukemtmc_reid']\n targets: ['p_dukemtmc_reid']\n height: 384\n width: 128\n transforms: ['rc', 're']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: True\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['bn_foreg', 'parts']\n masks:\n dir: 'pifpaf_maskrcnn_filtering'\n preprocess: 'eight'\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights: # SOTA weights for GiLt loss\n globl:\n id: 1.\n tr: 0.\n foreg:\n id: 1.\n tr: 0.\n conct:\n id: 1.\n tr: 0.\n parts:\n id: 0.\n tr: 1.\n pixls:\n ce: 0.35\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/pcb_market1501_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['market1501']\n targets: ['market1501']\n height: 384\n width: 128\n transforms: ['rc', 'rf', 're']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: False # no visibility scores used at training for PCB\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['conct'] # the holistic concatenated embedding alone is used to compute query-gallery distances\n masks:\n type: 'stripes' # PCB uses horizontal stripes\n parts_num: 6 # we use 6 horizontal stripes here\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights:\n globl:\n id: 0.\n tr: 0.\n foreg:\n id: 0.\n tr: 0.\n conct:\n id: 0.\n tr: 0.\n parts:\n id: 1. # For PCB, we use the cross-entropy (identity) loss on each of the 6 horizontal stripes\n tr: 0.\n pixls:\n ce: 0. # no body part prediction loss for PCB\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "configs/bpbreid/pcb_occ_duke_train.yaml", "content": "data:\n root: '~/datasets/reid'\n sources: ['occluded_duke']\n targets: ['occluded_duke']\n height: 384\n width: 128\n transforms: ['rc', 'rf', 're']\n\nmodel:\n name: 'bpbreid'\n bpbreid:\n mask_filtering_training: False\n mask_filtering_testing: False # no visibility scores used at training for PCB\n learnable_attention_enabled: True\n backbone: 'hrnet32'\n test_embeddings: ['conct'] # the holistic concatenated embedding alone is used to compute query-gallery distances\n masks:\n type: 'stripes' # PCB uses horizontal stripes\n parts_num: 6 # we use 6 horizontal stripes here\n\nloss:\n name: 'part_based'\n part_based:\n name: 'part_averaged_triplet_loss'\n ppl: 'cl'\n weights:\n globl:\n id: 0.\n tr: 0.\n foreg:\n id: 0.\n tr: 0.\n conct:\n id: 0.\n tr: 0.\n parts:\n id: 1. # For PCB, we use the cross-entropy (identity) loss on each of the 6 horizontal stripes\n tr: 0.\n pixls:\n ce: 0. # no body part prediction loss for PCB\n\ntrain:\n batch_size: 64\n\ntest:\n evaluate: False\n batch_size: 64\n visrank: True\n" }, { "path": "docs/AWESOME_REID.md", "content": "# Awesome-ReID\nHere is a collection of ReID-related research with links to papers and code. You are welcome to submit [PR](https://help.github.com/articles/creating-a-pull-request/)s if you find something missing.\n\n\n- [ArXiv20] Deep Learning for Person Re-identification: A Survey and Outlook [[paper](https://arxiv.org/abs/2001.04193)] [[code](https://github.com/mangye16/ReID-Survey)]\n\n- [ArXiv19] Learning Generalisable Omni-Scale Representations for Person Re-Identification [[paper](https://arxiv.org/abs/1910.06827)][[code](https://github.com/KaiyangZhou/deep-person-reid)]\n\n- [ICCV19] RGB-Infrared Cross-Modality Person Re-Identification via Joint Pixel and Feature Alignment. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Wang_RGB-Infrared_Cross-Modality_Person_Re-Identification_via_Joint_Pixel_and_Feature_Alignment_ICCV_2019_paper.pdf)] [[code](https://github.com/wangguanan/AlignGAN)]\n\n- [ICCV19] Unsupervised Graph Association for Person Re-identification. [[paper](https://github.com/yichuan9527/Unsupervised-Graph-Association-for-Person-Re-identification)] [[code](https://github.com/yichuan9527/Unsupervised-Graph-Association-for-Person-Re-identification)]\n\n- [ICCV19] Self-similarity Grouping: A Simple Unsupervised Cross Domain Adaptation Approach for Person Re-identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Fu_Self-Similarity_Grouping_A_Simple_Unsupervised_Cross_Domain_Adaptation_Approach_for_ICCV_2019_paper.pdf)] [[code](https://github.com/OasisYang/SSG)]\n\n- [ICCV19] Spectral Feature Transformation for Person Re-Identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Luo_Spectral_Feature_Transformation_for_Person_Re-Identification_ICCV_2019_paper.pdf)] [[code](https://github.com/LuckyDC/SFT_REID)]\n\n- [ICCV19] Beyond Human Parts: Dual Part-Aligned Representations for Person Re-Identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Guo_Beyond_Human_Parts_Dual_Part-Aligned_Representations_for_Person_Re-Identification_ICCV_2019_paper.pdf)] [[code](https://github.com/ggjy/P2Net.pytorch)]\n\n- [ICCV19] Co-segmentation Inspired Attention Networks for Video-based Person Re-identification. [[paper](http://openaccess.thecvf.com/content_ICCV_2019/papers/Subramaniam_Co-Segmentation_Inspired_Attention_Networks_for_Video-Based_Person_Re-Identification_ICCV_2019_paper.pdf)][[code](https://github.com/InnovArul/vidreid_cosegmentation)]\n\n- [ICCV19] Mixed High-Order Attention Network for Person Re-Identification. [[paper](https://arxiv.org/abs/1908.05819)][[code](https://github.com/chenbinghui1/MHN)]\n\n- [ICCV19] ABD-Net: Attentive but Diverse Person Re-Identification. [[paper](https://arxiv.org/abs/1908.01114)] [[code](https://github.com/TAMU-VITA/ABD-Net)]\n\n- [ICCV19] Omni-Scale Feature Learning for Person Re-Identification. [[paper](https://arxiv.org/abs/1905.00953)] [[code](https://github.com/KaiyangZhou/deep-person-reid)]\n\n- [CVPR19] Joint Discriminative and Generative Learning for Person Re-identification. [[paper](https://arxiv.org/abs/1904.07223)][[code](https://github.com/NVlabs/DG-Net)]\n- [CVPR19] Invariance Matters: Exemplar Memory for Domain Adaptive Person Re-identification. [[paper](https://arxiv.org/abs/1904.01990)][[code](https://github.com/zhunzhong07/ECN)]\n- [CVPR19] Dissecting Person Re-identification from the Viewpoint of Viewpoint. [[paper](https://arxiv.org/abs/1812.02162)][[code](https://github.com/sxzrt/Dissecting-Person-Re-ID-from-the-Viewpoint-of-Viewpoint)]\n- [CVPR19] Unsupervised Person Re-identification by Soft Multilabel Learning. [[paper](https://arxiv.org/abs/1903.06325)][[code](https://github.com/KovenYu/MAR)]\n- [CVPR19] Patch-based Discriminative Feature Learning for Unsupervised Person Re-identification. [[paper](https://kovenyu.com/publication/2019-cvpr-pedal/)][[code](https://github.com/QizeYang/PAUL)]\n\n- [AAAI19] Spatial and Temporal Mutual Promotion for Video-based Person Re-identification. [[paper](https://arxiv.org/abs/1812.10305)][[code](https://github.com/yolomax/person-reid-lib)]\n\n- [AAAI19] Spatial-Temporal Person Re-identification. [[paper](https://arxiv.org/abs/1812.03282)][[code](https://github.com/Wanggcong/Spatial-Temporal-Re-identification)]\n\n- [AAAI19] Horizontal Pyramid Matching for Person Re-identification. [[paper](https://arxiv.org/abs/1804.05275)][[code](https://github.com/OasisYang/HPM)]\n\n- [AAAI19] Backbone Can Not be Trained at Once: Rolling Back to Pre-trained Network for Person Re-identification. [[paper](https://arxiv.org/abs/1901.06140)][[code](https://github.com/youngminPIL/rollback)]\n\n- [AAAI19] A Bottom-Up Clustering Approach to Unsupervised Person Re-identification. [[paper](https://vana77.github.io/vana77.github.io/images/AAAI19.pdf)][[code](https://github.com/vana77/Bottom-up-Clustering-Person-Re-identification)]\n\n- [NIPS18] FD-GAN: Pose-guided Feature Distilling GAN for Robust Person Re-identification. [[paper](https://arxiv.org/abs/1810.02936)][[code](https://github.com/yxgeee/FD-GAN)]\n\n- [ECCV18] Generalizing A Person Retrieval Model Hetero- and Homogeneously. [[paper](http://openaccess.thecvf.com/content_ECCV_2018/papers/Zhun_Zhong_Generalizing_A_Person_ECCV_2018_paper.pdf)][[code](https://github.com/zhunzhong07/HHL)]\n\n- [ECCV18] Pose-Normalized Image Generation for Person Re-identification. [[paper](https://arxiv.org/abs/1712.02225)][[code](https://github.com/naiq/PN_GAN)]\n\n- [CVPR18] Camera Style Adaptation for Person Re-Identification. [[paper](https://arxiv.org/abs/1711.10295)][[code](https://github.com/zhunzhong07/CamStyle)]\n\n- [CVPR18] Deep Group-Shuffling Random Walk for Person Re-Identification. [[paper](https://arxiv.org/abs/1807.11178)][[code](https://github.com/YantaoShen/kpm_rw_person_reid)]\n\n- [CVPR18] End-to-End Deep Kronecker-Product Matching for Person Re-identification. [[paper](https://arxiv.org/abs/1807.11182)][[code](https://github.com/YantaoShen/kpm_rw_person_reid)]\n\n- [CVPR18] Features for Multi-Target Multi-Camera Tracking and Re-Identification. [[paper](https://arxiv.org/abs/1803.10859)][[code](https://github.com/ergysr/DeepCC)]\n\n- [CVPR18] Group Consistent Similarity Learning via Deep CRF for Person Re-Identification. [[paper](http://openaccess.thecvf.com/content_cvpr_2018/papers/Chen_Group_Consistent_Similarity_CVPR_2018_paper.pdf)][[code](https://github.com/dapengchen123/crf_affinity)]\n\n- [CVPR18] Harmonious Attention Network for Person Re-Identification. [[paper](https://arxiv.org/abs/1802.08122)][[code](https://github.com/KaiyangZhou/deep-person-reid)]\n\n- [CVPR18] Human Semantic Parsing for Person Re-Identification. [[paper](https://arxiv.org/abs/1804.00216)][[code](https://github.com/emrahbasaran/SPReID)]\n\n- [CVPR18] Multi-Level Factorisation Net for Person Re-Identification. [[paper](https://arxiv.org/abs/1803.09132)][[code](https://github.com/KaiyangZhou/deep-person-reid)]\n\n- [CVPR18] Resource Aware Person Re-identification across Multiple Resolutions. [[paper](https://arxiv.org/abs/1805.08805)][[code](https://github.com/mileyan/DARENet)]\n\n- [CVPR18] Exploit the Unknown Gradually: One-Shot Video-Based Person Re-Identification by Stepwise Learning. [[paper](https://yu-wu.net/pdf/CVPR2018_Exploit-Unknown-Gradually.pdf)][[code](https://github.com/Yu-Wu/Exploit-Unknown-Gradually)]\n\n- [ArXiv18] Revisiting Temporal Modeling for Video-based Person ReID. [[paper](https://arxiv.org/abs/1805.02104)][[code](https://github.com/jiyanggao/Video-Person-ReID)]\n" }, { "path": "docs/MODEL_ZOO.md", "content": "# Model Zoo\n\n- Results are presented in the format of **.\n- When computing model size and FLOPs, only layers that are used at test time are considered (see `torchreid.utils.compute_model_complexity`).\n- Asterisk (\\*) means the model is trained from scratch.\n- `combineall=True` means all images in the dataset are used for model training.\n\n\n## ImageNet pretrained models\n\n\n| Model | Download |\n| :--- | :---: |\n| shufflenet | [model](https://mega.nz/#!RDpUlQCY!tr_5xBEkelzDjveIYBBcGcovNCOrgfiJO9kiidz9fZM) |\n| mobilenetv2_x1_0 | [model](https://mega.nz/#!NKp2wAIA!1NH1pbNzY_M2hVk_hdsxNM1NUOWvvGPHhaNr-fASF6c) |\n| mobilenetv2_x1_4 | [model](https://mega.nz/#!RGhgEIwS!xN2s2ZdyqI6vQ3EwgmRXLEW3khr9tpXg96G9SUJugGk) |\n| mlfn | [model](https://mega.nz/#!YHxAhaxC!yu9E6zWl0x5zscSouTdbZu8gdFFytDdl-RAdD2DEfpk) |\n| osnet_x1_0 | [model](https://mega.nz/#!YK5GRARL!F90NsNB2XHjXGZFC3Lrw1GMic0oMw4fnfuDUnSrPAYM) |\n| osnet_x0_75 | [model](https://mega.nz/#!NPxilYBA!Se414Wtts__7eY6J5FIrowynvjUUG7a8Z5zUPfJN33s) |\n| osnet_x0_5 | [model](https://mega.nz/#!NO4ihQSJ!oMIRSZ0HlJF_8FKUbXT8Ei0vzH0xUYs5tWaf_KLrODg) |\n| osnet_x0_25 | [model](https://mega.nz/#!IDwQwaxT!TbQ_33gPK-ZchPFTf43UMc45rlNKWiWMqH4rTXB1T7k) |\n| osnet_ibn_x1_0 | [model](https://mega.nz/#!8Wo2kSDR!bNvgu4V0VkCQp_L2ZUDaudYKYRCkkSNdzcA1CcZGZTE) |\n| osnet_ain_x1_0 | [model](https://drive.google.com/open?id=1-CaioD9NaqbHK_kzSMW8VE4_3KcsRjEo) |\n\n\n## Same-domain ReID\n\n\n| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | market1501 | dukemtmcreid | msmt17 |\n| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| resnet50 | 23.5 | 2.7 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [87.9 (70.4)](https://mega.nz/#!FKZjVKaZ!4v_FR8pTvuHoMQIKdstJ_YCsRrtZW2hwWxc-T0JIlHE) | [78.3 (58.9)](https://mega.nz/#!JPZjCYhK!YVJbE_4vTc8DX19Rt_FB77YY4BaEA1P6Xb5sNJGep2M) | [63.2 (33.9)](https://mega.nz/#!APAxDY4Z!Iou9x8s3ATdYS2SlK2oiJbHrhvlzH7F1gE2qjM-GJGw) |\n| resnet50_fc512 | 24.6 | 4.1 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [90.8 (75.3)](https://mega.nz/#!EaZjhKyS!lBvD3vAJ4DOmElZkNa7gyPM1RE661GUd2v9kK84gSZE) | [81.0 (64.0)](https://mega.nz/#!lXYDSKZa!lumiXkY2H5Sm8gEgTWPBdWKv3ujy4zjrffjERaXkc9I) | [69.6 (38.4)](https://mega.nz/#!9PQTXIpL!iI5wgieTCn0Jm-pyg9RCu0RkH43pV3ntHhr1PeqSyT4) |\n| mlfn | 32.5 | 2.8 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [90.1 (74.3)](https://mega.nz/#!kHQ3ESLT!NoGc8eHEBZOJZM19THh3DFfRBXIPXzM-sdLmF1mvTXA) | [81.1 (63.2)](https://mega.nz/#!8PQXUCaI!mJO1vD9tI739hkNBj2QWUt0VPcZ-s89fSMMGPPP1msc) | [66.4 (37.2)](https://mega.nz/#!paIXFQCS!W3ZGkxyF1idwvQzTRDE2p0DhNDki2SBJRfp7S_Cwphk) |\n| hacnn* | 4.5 | 0.5 | softmax | (160, 64) | `random_flip`, `random_crop` | `euclidean` | [90.9 (75.6)](https://mega.nz/#!ULQXUQBK!S-8v_pR2xBD3ZpuY0I7Bqift-eX_V84gajHMDG6zUac) | [80.1 (63.2)](https://mega.nz/#!wPJTkAQR!XkKd39lsmBZMrCh3JjF6vnNafBZkouVIVdeBqQKdSzA) | [64.7 (37.2)](https://mega.nz/#!AXAziKjL!JtMwHz2UYy58gDMQLGakSmF3JOr72o8zmkqlQA-LIpQ) |\n| mobilenetv2_x1_0 | 2.2 | 0.2 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [85.6 (67.3)](https://mega.nz/#!8KYTFAIB!3dL35WQLxSoTSClDTv0kxa81k3fh5hXmAWA4_a3qiOI) | [74.2 (54.7)](https://mega.nz/#!hbRXDSCL!YYgqJ6PVUf4clgtUuK2s5FRhYJdU3yTibLscwOTNnDk) | [57.4 (29.3)](https://mega.nz/#!5SJTmCYb!ZQ8O2MN9JF4-WDAeX04Xex1KyuBYQ_o2aoMIsTgQ748) |\n| mobilenetv2_x1_4 | 4.3 | 0.4 | softmax | (256, 128) | `random_flip`, `random_crop` | `euclidean` | [87.0 (68.5)](https://mega.nz/#!4XZhEKCS!6lTuTRbHIWU5nzJzTPDGykA7sPME8_1ISGsUYFJXZWA) | [76.2 (55.8)](https://mega.nz/#!JbQVDIYQ!-7pnjIfpIDt1EoQOvpvuIEcTj3Qg8SE6o_3ZPGWrIcw) | [60.1 (31.5)](https://mega.nz/#!gOYDAQrK!sMJO7c_X4iIxoVfV_tXYdzeDJByPo5XkUjEN7Z2JTmM) |\n| osnet_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip` | `euclidean` | [94.2 (82.6)](https://mega.nz/#!hLoyTSba!fqt7GcKrHJhwe9BtuK0ozgVAQcrlMG8Pm6JsSfr5HEI) | [87.0 (70.2)](https://mega.nz/#!ETwGhQYB!h2gHN-H3J4X4WqcJXy2b0pPKl28paydkiS-PDHsEgPM) | [74.9 (43.8)](https://mega.nz/#!hWxE2aJA!NGcxu5uYH1qI6DfBTu0KFoi_NfoA0TJcBFW-g43pC0I) |\n| osnet_x0_75 | 1.3 | 0.57 | softmax | (256, 128) | `random_flip` | `euclidean` | [93.7 (81.2)](https://mega.nz/#!JO4WAaJa!nQuoqZnYfy0xu7vs2mp28AFceya-ZhrXTry837jvoDQ) | [85.8 (69.8)](https://mega.nz/#!lOgkEIoI!fQ5vuYIABIOcRxF-OK-6YxtEufWhyVkYkGB4qPoRYJ4) | [72.8 (41.4)](https://mega.nz/#!0exGXI5a!rxtzBayyRK0on0HFq9XO0UtWEBhbV86dFitljhjeWcs) |\n| osnet_x0_5 | 0.6 | 0.27 | softmax | (256, 128) | `random_flip` | `euclidean` | [92.5 (79.8)](https://mega.nz/#!QCx0RArD!hqz3Mh0Iif5d8PpQW0frxa-Tepn2a2g24aei7du4MFs) | [85.1 (67.4)](https://mega.nz/#!QTxCDIbT!eOZxj4dHl0uFnjKEB-J3YBY98blXZvppgWGA3CGa-tk) | [69.7 (37.5)](https://mega.nz/#!ETpiECDa!CCkq4JryztHqgw7spL5zDw0usJpAfEsSd5gPlkMufCc) |\n| osnet_x0_25 | 0.2 | 0.08 | softmax | (256, 128) | `random_flip` | `euclidean` | [91.2 (75.0)](https://mega.nz/#!VWxCgSqY!Q4WaQ3j9D7HMhK3jsbvMuwaZ7yBY80T2Zj5V8JAlAKU) | [82.0 (61.4)](https://mega.nz/#!5TpwnATK!UvU_Asdy_aJ9SNzuvqhEFoemxSSB8vm_Gm8Xe03jqiA) | [61.4 (29.5)](https://mega.nz/#!AWgE3SzD!DngUaNyA7VIqOd2gq10Aty_-ER0CmG0xTJLHLj6_36g) |\n\n\n## Cross-domain ReID\n\n#### Market1501 -> DukeMTMC-reID\n\n\n| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | Rank-1 | Rank-5 | Rank-10 | mAP | Download |\n| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| osnet_ibn_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 48.5 | 62.3 | 67.4 | 26.7 | [model](https://mega.nz/#!wXwGxKxK!f8EMk8hBt6AjxU3JIPGMFSMvX7j-Nt5Lp1Gpbqso1Ts) |\n| osnet_ain_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | 52.4 | 66.1 | 71.2 | 30.5 | [model](https://mega.nz/#!QLJE2CRI!FXYc3Vm6Y5Scwx0xvRwBJxId56kf06fIXNLwA_b_1FE) |\n\n\n#### DukeMTMC-reID -> Market1501\n\n\n| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | Rank-1 | Rank-5 | Rank-10 | mAP | Download |\n| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| osnet_ibn_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 57.7 | 73.7 | 80.0 | 26.1 | [model](https://mega.nz/#!FD4WEKJS!ZGgI-2IwVuX6re09xylChR03o6Dkjpi6KSebrbS0fAA) |\n| osnet_ain_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | 61.0 | 77.0 | 82.5 | 30.6 | [model](https://mega.nz/#!4PBQlCCL!9yMHu1WyyBVxqssubLAEyoEfHUiNP4Ggg5On0nCX2S4) |\n\n\n#### MSMT17 (`combineall=True`) -> Market1501 & DukeMTMC-reID\n\n\n| Model | # Param (10^6) | GFLOPs | Loss | Input | Transforms | Distance | msmt17 -> market1501 | msmt17 -> dukemtmcreid | Download |\n| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |\n| resnet50 | 23.5 | 2.7 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 46.3 (22.8) | 52.3 (32.1) | [model](https://mega.nz/#!VTpkWSbS!Y8gDnmg7u-sPwnZDhWXrtZNYOj7UYL4QzZkhDf1qWW4) |\n| osnet_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 66.6 (37.5) | 66.0 (45.3) | [model](https://mega.nz/#!MepG3QRC!Lb-C9d7rdS_YJjGSoJ5cRlzjYcP28P_1Cm5S5WSslW0) |\n| osnet_x0_75 | 1.3 | 0.57 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 63.6 (35.5) | 65.3 (44.5) | [model](https://mega.nz/#!tO4WDagL!8Tl6kdJWRXRHQb16GeUHR008tJqW3N7_3fyVMu-LcKM) |\n| osnet_x0_5 | 0.6 | 0.27 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 64.3 (34.9) | 65.2 (43.3) | [model](https://mega.nz/#!papSWQhY!IId-QfcHj7nXQ_muUubgv9_n0SsnZzarmb5mQgcMv74) |\n| osnet_x0_25 | 0.2 | 0.08 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 59.9 (31.0) | 61.5 (39.6) | [model](https://mega.nz/#!QCoE0Kpa!BITLANumgjiR68TUFteL__N_RIoDKkL0M5Bl3Q8LC3U) |\n| osnet_ibn_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `euclidean` | 66.5 (37.2) | 67.4 (45.6) | [model](https://mega.nz/#!dL4Q2K5B!ZdHQ_X_rs2T-xmggigM5YvzJhmT1orkr6aQ1_fHgunM) |\n| osnet_ain_x1_0 | 2.2 | 0.98 | softmax | (256, 128) | `random_flip`, `color_jitter` | `cosine` | 70.1 (43.3) | 71.1 (52.7) | [model](https://mega.nz/#!YTZFnSJY!wlbo_5oa2TpDAGyWCTKTX1hh4d6DvJhh_RUA2z6i_so) |\n" }, { "path": "docs/Makefile", "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line.\nSPHINXOPTS =\nSPHINXBUILD = sphinx-build\nSOURCEDIR = .\nBUILDDIR = _build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option. $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)" }, { "path": "docs/conf.py", "content": "# -*- coding: utf-8 -*-\n#\n# Configuration file for the Sphinx documentation builder.\n#\n# This file does only contain a selection of the most common options. For a\n# full list see the documentation:\n# http://www.sphinx-doc.org/en/master/config\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\nimport os\nimport sys\n\nsys.path.insert(0, os.path.abspath('..'))\n\n# -- Project information -----------------------------------------------------\n\nproject = u'torchreid'\ncopyright = u'2019, Kaiyang Zhou'\nauthor = u'Kaiyang Zhou'\n\nversion_file = '../torchreid/__init__.py'\nwith open(version_file, 'r') as f:\n exec(compile(f.read(), version_file, 'exec'))\n__version__ = locals()['__version__']\n\n# The short X.Y version\nversion = __version__\n# The full version, including alpha/beta/rc tags\nrelease = __version__\n\n# -- General configuration ---------------------------------------------------\n\n# If your documentation needs a minimal Sphinx version, state it here.\n#\n# needs_sphinx = '1.0'\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n 'sphinx.ext.autodoc',\n 'sphinxcontrib.napoleon',\n 'sphinx.ext.viewcode',\n 'sphinx.ext.githubpages',\n 'sphinx_markdown_tables',\n]\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = ['_templates']\n\n# The suffix(es) of source filenames.\n# You can specify multiple suffix as a list of string:\n#\nsource_suffix = ['.rst', '.md']\n# source_suffix = '.rst'\nsource_parsers = {'.md': 'recommonmark.parser.CommonMarkParser'}\n\n# The master toctree document.\nmaster_doc = 'index'\n\n# The language for content autogenerated by Sphinx. Refer to documentation\n# for a list of supported languages.\n#\n# This is also used if you do content translation via gettext catalogs.\n# Usually you set \"language\" from the command line for these cases.\nlanguage = None\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path.\nexclude_patterns = [u'_build', 'Thumbs.db', '.DS_Store']\n\n# The name of the Pygments (syntax highlighting) style to use.\npygments_style = None\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages. See the documentation for\n# a list of builtin themes.\n#\nhtml_theme = 'sphinx_rtd_theme'\n\n# Theme options are theme-specific and customize the look and feel of a theme\n# further. For a list of options available for each theme, see the\n# documentation.\n#\n# html_theme_options = {}\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named \"default.css\" will overwrite the builtin \"default.css\".\nhtml_static_path = ['_static']\n\n# Custom sidebar templates, must be a dictionary that maps document names\n# to template names.\n#\n# The default sidebars (for documents that don't match any pattern) are\n# defined by theme itself. Builtin themes are using these templates by\n# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',\n# 'searchbox.html']``.\n#\n# html_sidebars = {}\n\n# -- Options for HTMLHelp output ---------------------------------------------\n\n# Output file base name for HTML help builder.\nhtmlhelp_basename = 'torchreiddoc'\n\n# -- Options for LaTeX output ------------------------------------------------\n\nlatex_elements = {\n # The paper size ('letterpaper' or 'a4paper').\n #\n # 'papersize': 'letterpaper',\n\n # The font size ('10pt', '11pt' or '12pt').\n #\n # 'pointsize': '10pt',\n\n # Additional stuff for the LaTeX preamble.\n #\n # 'preamble': '',\n\n # Latex figure (float) alignment\n #\n # 'figure_align': 'htbp',\n}\n\n# Grouping the document tree into LaTeX files. List of tuples\n# (source start file, target name, title,\n# author, documentclass [howto, manual, or own class]).\nlatex_documents = [\n (\n master_doc, 'torchreid.tex', u'torchreid Documentation',\n u'Kaiyang Zhou', 'manual'\n ),\n]\n\n# -- Options for manual page output ------------------------------------------\n\n# One entry per manual page. List of tuples\n# (source start file, name, description, authors, manual section).\nman_pages = [\n (master_doc, 'torchreid', u'torchreid Documentation', [author], 1)\n]\n\n# -- Options for Texinfo output ----------------------------------------------\n\n# Grouping the document tree into Texinfo files. List of tuples\n# (source start file, target name, title, author,\n# dir menu entry, description, category)\ntexinfo_documents = [\n (\n master_doc, 'torchreid', u'torchreid Documentation', author,\n 'torchreid', 'One line description of project.', 'Miscellaneous'\n ),\n]\n\n# -- Options for Epub output -------------------------------------------------\n\n# Bibliographic Dublin Core info.\nepub_title = project\n\n# The unique identifier of the text. This can be a ISBN number\n# or the project homepage.\n#\n# epub_identifier = ''\n\n# A unique identification for the text.\n#\n# epub_uid = ''\n\n# A list of files that should not be packed into the epub file.\nepub_exclude_files = ['search.html']\n\n# -- Extension configuration -------------------------------------------------\n" }, { "path": "docs/datasets.rst", "content": ".. _datasets:\n\nDatasets\n=========\n\nHere we provide a guide on how to prepare datasets.\n\nSuppose you wanna store the reid data in a directory called \"path/to/reid-data/\", you need to specify the ``root`` as *root='path/to/reid-data/'* when initializing ``DataManager``. Below we use ``$REID`` to denote \"path/to/reid-data\".\n\nPlease refer to :ref:`torchreid_data` for details regarding the arguments.\n\n\n.. note::\n Dataset with a :math:`\\dagger` superscript means that the process is automated, so you can directly call the dataset in ``DataManager`` (which automatically downloads the dataset and organizes the data structure). However, we also provide a way below to help the manual setup in case the automation fails.\n\n\n.. note::\n The keys to use specific datasets are enclosed in the parantheses beside the datasets.\n\n\n.. note::\n You are suggested to use the provided names for dataset folders such as \"market1501\" for Market1501 and \"dukemtmcreid\" for DukeMTMC-reID when doing the manual setup, otherwise you need to modify the source code accordingly (i.e. the ``dataset_dir`` attribute).\n\n\n.. contents::\n :local:\n\n\nImage Datasets\n--------------\n\nMarket1501 :math:`^\\dagger` (``market1501``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create a directory named \"market1501\" under ``$REID``.\n- Download the dataset to \"market1501\" from http://www.liangzheng.org/Project/project_reid.html and extract the files.\n- The data structure should look like\n\n.. code-block:: none\n \n market1501/\n Market-1501-v15.09.15/\n query/\n bounding_box_train/\n bounding_box_test/\n\n- To use the extra 500K distractors (i.e. Market1501 + 500K), go to the **Market-1501+500k Dataset** section at http://www.liangzheng.org/Project/project_reid.html, download the zip file \"distractors_500k.zip\" and extract it under \"market1501/Market-1501-v15.09.15\". The argument to use these 500K distrctors is ``market1501_500k`` in ``ImageDataManager``.\n\n\nCUHK03 (``cuhk03``)\n^^^^^^^^^^^^^^^^^^^^^\n- Create a folder named \"cuhk03\" under ``$REID``.\n- Download the dataset to \"cuhk03/\" from http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html and extract \"cuhk03_release.zip\", resulting in \"cuhk03/cuhk03_release/\".\n- Download the new split (767/700) from `person-re-ranking `_. What you need are \"cuhk03_new_protocol_config_detected.mat\" and \"cuhk03_new_protocol_config_labeled.mat\". Put these two mat files under \"cuhk03/\".\n- The data structure should look like\n\n.. code-block:: none\n \n cuhk03/\n cuhk03_release/\n cuhk03_new_protocol_config_detected.mat\n cuhk03_new_protocol_config_labeled.mat\n\n\n- In the default mode, we load data using the new split (767/700). If you wanna use the original (20) splits (1367/100), please set ``cuhk03_classic_split`` to True in ``ImageDataManager``. As the CMC is computed differently from Market1501 for the 1367/100 split (see `here `_), you need to enable ``use_metric_cuhk03`` in ``ImageDataManager`` to activate the *single-gallery-shot* metric for fair comparison with some methods that adopt the old splits (*do not need to report mAP*). In addition, we support both *labeled* and *detected* modes. The default mode loads *detected* images. Enable ``cuhk03_labeled`` in ``ImageDataManager`` if you wanna train and test on *labeled* images.\n\n.. note::\n The code will extract images in \"cuhk-03.mat\" and save them under \"cuhk03/images_detected\" and \"cuhk03/images_labeled\". Also, four json files will be automatically generated, i.e. \"splits_classic_detected.json\", \"splits_classic_labeled.json\", \"splits_new_detected.json\" and \"splits_new_labeled.json\". If the parent path of ``$REID`` is changed, these json files should be manually deleted. The code can automatically generate new json files to match the new path.\n \n\nDukeMTMC-reID :math:`^\\dagger` (``dukemtmcreid``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create a directory called \"dukemtmc-reid\" under ``$REID``.\n- Download \"DukeMTMC-reID\" from http://vision.cs.duke.edu/DukeMTMC/ and extract it under \"dukemtmc-reid\".\n- The data structure should look like\n\n.. code-block:: none\n \n dukemtmc-reid/\n DukeMTMC-reID/\n query/\n bounding_box_train/\n bounding_box_test/\n ...\n\nMSMT17 (``msmt17``)\n^^^^^^^^^^^^^^^^^^^^^\n- Create a directory called \"msmt17\" under ``$REID``.\n- Download the dataset from http://www.pkuvmc.com/publications/msmt17.html to \"msmt17\" and extract the files.\n- The data structure should look like\n\n.. code-block:: none\n \n msmt17/\n MSMT17_V1/ # or MSMT17_V2\n train/\n test/\n list_train.txt\n list_query.txt\n list_gallery.txt\n list_val.txt\n\nVIPeR :math:`^\\dagger` (``viper``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- The download link is http://users.soe.ucsc.edu/~manduchi/VIPeR.v1.0.zip.\n- Organize the dataset in a folder named \"viper\" as follows\n\n.. code-block:: none\n \n viper/\n VIPeR/\n cam_a/\n cam_b/\n\nGRID :math:`^\\dagger` (``grid``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- The download link is http://personal.ie.cuhk.edu.hk/~ccloy/files/datasets/underground_reid.zip.\n- Organize the dataset in a folder named \"grid\" as follows\n\n.. code-block:: none\n \n grid/\n underground_reid/\n probe/\n gallery/\n ...\n\nCUHK01 (``cuhk01``)\n^^^^^^^^^^^^^^^^^^^^^^^^\n- Create a folder named \"cuhk01\" under ``$REID``.\n- Download \"CUHK01.zip\" from http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html and place it under \"cuhk01/\".\n- The code can automatically extract the files, or you can do it yourself.\n- The data structure should look like\n\n.. code-block:: none\n \n cuhk01/\n campus/\n\nSenseReID (``sensereid``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create \"sensereid\" under ``$REID``.\n- Download the dataset from this `link `_ and extract it to \"sensereid\".\n- Organize the data to be like\n\n.. code-block:: none\n\n sensereid/\n SenseReID/\n test_probe/\n test_gallery/\n\nQMUL-iLIDS :math:`^\\dagger` (``ilids``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create a folder named \"ilids\" under ``$REID``.\n- Download the dataset from http://www.eecs.qmul.ac.uk/~jason/data/i-LIDS_Pedestrian.tgz and organize it to look like\n\n.. code-block:: none\n \n ilids/\n i-LIDS_Pedestrian/\n Persons/\n\nPRID (``prid``)\n^^^^^^^^^^^^^^^^^^^\n- Create a directory named \"prid2011\" under ``$REID``.\n- Download the dataset from https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/PRID11/ and extract it under \"prid2011\".\n- The data structure should end up with\n\n.. code-block:: none\n\n prid2011/\n prid_2011/\n single_shot/\n multi_shot/\n\nCUHK02 (``cuhk02``)\n^^^^^^^^^^^^^^^^^^^^^\n- Create a folder named \"cuhk02\" under ``$REID``.\n- Download the data from http://www.ee.cuhk.edu.hk/~xgwang/CUHK_identification.html and put it under \"cuhk02/\".\n- Extract the file so the data structure looks like\n\n.. code-block:: none\n \n cuhk02/\n Dataset/\n P1/\n P2/\n P3/\n P4/\n P5/\n\nVideo Datasets\n--------------\n\nMARS (``mars``)\n^^^^^^^^^^^^^^^^^\n- Create \"mars/\" under ``$REID``.\n- Download the dataset from http://www.liangzheng.com.cn/Project/project_mars.html and place it in \"mars/\".\n- Extract \"bbox_train.zip\" and \"bbox_test.zip\".\n- Download the split metadata from https://github.com/liangzheng06/MARS-evaluation/tree/master/info and put \"info/\" in \"mars/\".\n- The data structure should end up with\n\n.. code-block:: none\n \n mars/\n bbox_test/\n bbox_train/\n info/\n\niLIDS-VID :math:`^\\dagger` (``ilidsvid``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create \"ilids-vid\" under ``$REID``.\n- Download the dataset from http://www.eecs.qmul.ac.uk/~xiatian/downloads_qmul_iLIDS-VID_ReID_dataset.html to \"ilids-vid\".\n- Organize the data structure to match\n\n.. code-block:: none\n \n ilids-vid/\n i-LIDS-VID/\n train-test people splits/\n\nPRID2011 (``prid2011``)\n^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create a directory named \"prid2011\" under ``$REID``.\n- Download the dataset from https://www.tugraz.at/institute/icg/research/team-bischof/lrs/downloads/PRID11/ and extract it under \"prid2011\".\n- Download the split created by *iLIDS-VID* from `this google drive `_ and put it under \"prid2011/\". Following the standard protocol, only 178 persons whose sequences are more than a threshold are used.\n- The data structure should end up with\n\n.. code-block:: none\n\n prid2011/\n splits_prid2011.json\n prid_2011/\n single_shot/\n multi_shot/\n\nDukeMTMC-VideoReID :math:`^\\dagger` (``dukemtmcvidreid``)\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n- Create \"dukemtmc-vidreid\" under ``$REID``.\n- Download \"DukeMTMC-VideoReID\" from http://vision.cs.duke.edu/DukeMTMC/ and unzip the file to \"dukemtmc-vidreid/\".\n- The data structure should look like\n\n.. code-block:: none\n \n dukemtmc-vidreid/\n DukeMTMC-VideoReID/\n train/\n query/\n gallery/\n" }, { "path": "docs/evaluation.rst", "content": "Evaluation\n==========\n\nImage ReID\n-----------\n- **Market1501**, **DukeMTMC-reID**, **CUHK03 (767/700 split)** and **MSMT17** have fixed split so keeping ``split_id=0`` is fine.\n- **CUHK03 (classic split)** has 20 fixed splits, so do ``split_id=0~19``.\n- **VIPeR** contains 632 identities each with 2 images under two camera views. Evaluation should be done for 10 random splits. Each split randomly divides 632 identities to 316 train ids (632 images) and the other 316 test ids (632 images). Note that, in each random split, there are two sub-splits, one using camera-A as query and camera-B as gallery while the other one using camera-B as query and camera-A as gallery. Thus, there are totally 20 splits generated with ``split_id`` starting from 0 to 19. Models can be trained on ``split_id=[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]`` (because ``split_id=0`` and ``split_id=1`` share the same train set, and so on and so forth.). At test time, models trained on ``split_id=0`` can be directly evaluated on ``split_id=1``, models trained on ``split_id=2`` can be directly evaluated on ``split_id=3``, and so on and so forth.\n- **CUHK01** is similar to VIPeR in the split generation.\n- **GRID** , **iLIDS** and **PRID** have 10 random splits, so evaluation should be done by varying ``split_id`` from 0 to 9.\n- **SenseReID** has no training images and is used for evaluation only.\n\n\n.. note::\n The ``split_id`` argument is defined in ``ImageDataManager`` and ``VideoDataManager``. Please refer to :ref:`torchreid_data`.\n\n\nVideo ReID\n-----------\n- **MARS** and **DukeMTMC-VideoReID** have fixed single split so using ``split_id=0`` is ok.\n- **iLIDS-VID** and **PRID2011** have 10 predefined splits so evaluation should be done by varying ``split_id`` from 0 to 9." }, { "path": "docs/index.rst", "content": ".. include:: ../README.rst\n\n\n.. toctree::\n :hidden:\n\n user_guide\n datasets\n evaluation\n\n.. toctree::\n :caption: Package Reference\n :hidden:\n\n pkg/data\n pkg/engine\n pkg/losses\n pkg/metrics\n pkg/models\n pkg/optim\n pkg/utils\n\n.. toctree::\n :caption: Resources\n :hidden:\n\n AWESOME_REID.md\n MODEL_ZOO.md\n\n\nIndices and tables\n==================\n\n* :ref:`genindex`\n* :ref:`modindex`" }, { "path": "docs/pkg/data.rst", "content": ".. _torchreid_data:\n\ntorchreid.data\n==============\n\n\nData Manager\n---------------------------\n\n.. automodule:: torchreid.data.datamanager\n :members:\n\n\nSampler\n-----------------------\n\n.. automodule:: torchreid.data.sampler\n :members:\n\n\nTransforms\n---------------------------\n\n.. automodule:: torchreid.data.transforms\n :members:\n\n\nDataset\n---------------------------\n\n.. automodule:: torchreid.data.datasets.dataset\n :members:\n\n\n.. automodule:: torchreid.data.datasets.__init__\n :members:\n\n\nImage Datasets\n------------------------------\n\n.. automodule:: torchreid.data.datasets.image.market1501\n :members:\n\n.. automodule:: torchreid.data.datasets.image.cuhk03\n :members:\n\n.. automodule:: torchreid.data.datasets.image.dukemtmcreid\n :members:\n\n.. automodule:: torchreid.data.datasets.image.msmt17\n :members:\n\n.. automodule:: torchreid.data.datasets.image.viper\n :members:\n\n.. automodule:: torchreid.data.datasets.image.grid\n :members:\n\n.. automodule:: torchreid.data.datasets.image.cuhk01\n :members:\n\n.. automodule:: torchreid.data.datasets.image.ilids\n :members:\n\n.. automodule:: torchreid.data.datasets.image.sensereid\n :members:\n\n.. automodule:: torchreid.data.datasets.image.prid\n :members:\n\n\nVideo Datasets\n------------------------------\n\n.. automodule:: torchreid.data.datasets.video.mars\n :members:\n\n.. automodule:: torchreid.data.datasets.video.ilidsvid\n :members:\n\n.. automodule:: torchreid.data.datasets.video.prid2011\n :members:\n\n.. automodule:: torchreid.data.datasets.video.dukemtmcvidreid\n :members:" }, { "path": "docs/pkg/engine.rst", "content": ".. _torchreid_engine:\n\ntorchreid.engine\n==================\n\n\nBase Engine\n------------\n\n.. autoclass:: torchreid.engine.engine.Engine\n :members:\n\n\nImage Engines\n-------------\n\n.. autoclass:: torchreid.engine.image.softmax.ImageSoftmaxEngine\n :members:\n\n\n.. autoclass:: torchreid.engine.image.triplet.ImageTripletEngine\n :members:\n\n\nVideo Engines\n-------------\n\n.. autoclass:: torchreid.engine.video.softmax.VideoSoftmaxEngine\n\n\n.. autoclass:: torchreid.engine.video.triplet.VideoTripletEngine" }, { "path": "docs/pkg/losses.rst", "content": ".. _torchreid_losses:\n\ntorchreid.losses\n=================\n\n\nSoftmax\n--------\n\n.. automodule:: torchreid.losses.cross_entropy_loss\n :members:\n\n\nTriplet\n-------\n\n.. automodule:: torchreid.losses.hard_mine_triplet_loss\n :members:" }, { "path": "docs/pkg/metrics.rst", "content": ".. _torchreid_metrics:\n\ntorchreid.metrics\n=================\n\n\nDistance\n---------\n\n.. automodule:: torchreid.metrics.distance\n :members:\n\n\nAccuracy\n--------\n\n.. automodule:: torchreid.metrics.accuracy\n :members:\n\n\nRank\n-----\n\n.. automodule:: torchreid.metrics.rank\n :members: evaluate_rank" }, { "path": "docs/pkg/models.rst", "content": ".. _torchreid_models:\n\ntorchreid.models\n=================\n\nInterface\n---------\n\n.. automodule:: torchreid.models.__init__\n :members:\n\n\nImageNet Classification Models\n-------------------------------\n\n.. autoclass:: torchreid.models.resnet.ResNet\n.. autoclass:: torchreid.models.senet.SENet\n.. autoclass:: torchreid.models.densenet.DenseNet\n.. autoclass:: torchreid.models.inceptionresnetv2.InceptionResNetV2\n.. autoclass:: torchreid.models.inceptionv4.InceptionV4\n.. autoclass:: torchreid.models.xception.Xception\n\n\nLightweight Models\n------------------\n\n.. autoclass:: torchreid.models.nasnet.NASNetAMobile\n.. autoclass:: torchreid.models.mobilenetv2.MobileNetV2\n.. autoclass:: torchreid.models.shufflenet.ShuffleNet\n.. autoclass:: torchreid.models.squeezenet.SqueezeNet\n.. autoclass:: torchreid.models.shufflenetv2.ShuffleNetV2\n\n\nReID-specific Models\n--------------------\n\n.. autoclass:: torchreid.models.mudeep.MuDeep\n.. autoclass:: torchreid.models.resnetmid.ResNetMid\n.. autoclass:: torchreid.models.hacnn.HACNN\n.. autoclass:: torchreid.models.pcb.PCB\n.. autoclass:: torchreid.models.mlfn.MLFN\n.. autoclass:: torchreid.models.osnet.OSNet\n.. autoclass:: torchreid.models.osnet_ain.OSNet" }, { "path": "docs/pkg/optim.rst", "content": ".. _torchreid_optim:\n\ntorchreid.optim\n=================\n\n\nOptimizer\n----------\n\n.. automodule:: torchreid.optim.optimizer\n :members: build_optimizer\n\n\nLR Scheduler\n-------------\n\n.. automodule:: torchreid.optim.lr_scheduler\n :members: build_lr_scheduler" }, { "path": "docs/pkg/utils.rst", "content": ".. _torchreid_utils:\n\ntorchreid.utils\n=================\n\nAverage Meter\n--------------\n\n.. automodule:: torchreid.utils.avgmeter\n :members:\n\n\nLoggers\n-------\n\n.. automodule:: torchreid.utils.loggers\n :members:\n\n\nGeneric Tools\n---------------\n.. automodule:: torchreid.utils.tools\n :members:\n\n\nReID Tools\n----------\n\n.. automodule:: torchreid.utils.reidtools\n :members:\n\n\nTorch Tools\n------------\n\n.. automodule:: torchreid.utils.torchtools\n :members:\n\n\n.. automodule:: torchreid.utils.model_complexity\n :members:\n" }, { "path": "docs/requirements.txt", "content": "sphinx\nsphinx-markdown-tables\nsphinx-rtd-theme\nsphinxcontrib-napoleon\nsphinxcontrib-websupport\nrecommonmark" }, { "path": "docs/user_guide.rst", "content": "How-to\n============\n\n.. contents::\n :local:\n\n\nPrepare datasets\n-----------------\nSee :ref:`datasets`.\n\n\nFind model keys\n-----------------\nKeys are listed under the *Public keys* section within each model class in :ref:`torchreid_models`.\n\n\nShow available models\n----------------------\n\n.. code-block:: python\n \n import torchreid\n torchreid.models.show_avai_models()\n\n\nChange the training sampler\n-----------------------------\nThe default ``train_sampler`` is \"RandomSampler\". You can give the specific sampler name as input to ``train_sampler``, e.g. ``train_sampler='RandomIdentitySampler'`` for triplet loss.\n\n\nChoose an optimizer/lr_scheduler\n----------------------------------\nPlease refer to the source code of ``build_optimizer``/``build_lr_scheduler`` in :ref:`torchreid_optim` for details.\n\n\nResume training\n----------------\nSuppose the checkpoint is saved in \"log/resnet50/model.pth.tar-30\", you can do\n\n.. code-block:: python\n \n start_epoch = torchreid.utils.resume_from_checkpoint(\n 'log/resnet50/model.pth.tar-30',\n model,\n optimizer\n )\n\n engine.run(\n save_dir='log/resnet50',\n max_epoch=60,\n start_epoch=start_epoch\n )\n\n\nCompute model complexity\n--------------------------\nWe provide a tool in ``torchreid.utils.model_complexity.py`` to automatically compute the model complexity, i.e. number of parameters and FLOPs.\n\n.. code-block:: python\n \n from torchreid import models, utils\n \n model = models.build_model(name='resnet50', num_classes=1000)\n num_params, flops = utils.compute_model_complexity(model, (1, 3, 256, 128))\n\n # show detailed complexity for each module\n utils.compute_model_complexity(model, (1, 3, 256, 128), verbose=True)\n\n # count flops for all layers including ReLU and BatchNorm\n utils.compute_model_complexity(model, (1, 3, 256, 128), verbose=True, only_conv_linear=False)\n\nNote that (1) this function only provides an estimate of the theoretical time complexity rather than the actual running time which depends on implementations and hardware; (2) the FLOPs is only counted for layers that are used at test time. This means that redundant layers such as person ID classification layer will be ignored. The inference graph depends on how you define the computations in ``forward()``.\n\n\nCombine multiple datasets\n---------------------------\nEasy. Just give whatever datasets (keys) you want to the ``sources`` argument when instantiating a data manager. For example,\n\n.. code-block:: python\n \n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources=['market1501', 'dukemtmcreid', 'cuhk03', 'msmt17'],\n height=256,\n width=128,\n batch_size=32\n )\n\nIn this example, the target datasets are Market1501, DukeMTMC-reID, CUHK03 and MSMT17 as the ``targets`` argument is not specified. Please refer to ``Engine.test()`` in :ref:`torchreid_engine` for details regarding how evaluation is performed.\n\n\nDo cross-dataset evaluation\n-----------------------------\nEasy. Just give whatever datasets (keys) you want to the argument ``targets``, like\n\n.. code-block:: python\n \n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources='market1501',\n targets='dukemtmcreid', # or targets='cuhk03' or targets=['dukemtmcreid', 'cuhk03']\n height=256,\n width=128,\n batch_size=32\n )\n\n\nCombine train, query and gallery\n---------------------------------\nThis can be easily done by setting ``combineall=True`` when instantiating a data manager. Below is an example of using Market1501,\n\n.. code-block:: python\n \n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources='market1501',\n height=256,\n width=128,\n batch_size=32,\n market1501_500k=False,\n combineall=True # it's me, here\n )\n\nMore specifically, with ``combineall=False``, you will get\n\n.. code-block:: none\n \n => Loaded Market1501\n ----------------------------------------\n subset | # ids | # images | # cameras\n ----------------------------------------\n train | 751 | 12936 | 6\n query | 750 | 3368 | 6\n gallery | 751 | 15913 | 6\n ---------------------------------------\n\nwith ``combineall=True``, you will get\n\n.. code-block:: none\n \n => Loaded Market1501\n ----------------------------------------\n subset | # ids | # images | # cameras\n ----------------------------------------\n train | 1501 | 29419 | 6\n query | 750 | 3368 | 6\n gallery | 751 | 15913 | 6\n ---------------------------------------\n\n\nOptimize layers with different learning rates\n-----------------------------------------------\nA common practice for fine-tuning pretrained models is to use a smaller learning rate for base layers and a large learning rate for randomly initialized layers (referred to as ``new_layers``). ``torchreid.optim.optimizer`` has implemented such feature. What you need to do is to set ``staged_lr=True`` and give the names of ``new_layers`` such as \"classifier\".\n\nBelow is an example of setting different learning rates for base layers and new layers in ResNet50,\n\n.. code-block:: python\n \n # New layer \"classifier\" has a learning rate of 0.01\n # The base layers have a learning rate of 0.001\n optimizer = torchreid.optim.build_optimizer(\n model,\n optim='sgd',\n lr=0.01,\n staged_lr=True,\n new_layers='classifier',\n base_lr_mult=0.1\n )\n\nPlease refer to :ref:`torchreid_optim` for more details.\n\n\nDo two-stepped transfer learning\n-------------------------------------\nTo prevent the pretrained layers from being damaged by harmful gradients back-propagated from randomly initialized layers, one can adopt the *two-stepped transfer learning strategy* presented in `Deep Transfer Learning for Person Re-identification `_. The basic idea is to pretrain the randomly initialized layers for few epochs while keeping the base layers frozen before training all layers end-to-end.\n\nThis has been implemented in ``Engine.train()`` (see :ref:`torchreid_engine`). The arguments related to this feature are ``fixbase_epoch`` and ``open_layers``. Intuitively, ``fixbase_epoch`` denotes the number of epochs to keep the base layers frozen; ``open_layers`` means which layer is open for training.\n\nFor example, say you want to pretrain the classification layer named \"classifier\" in ResNet50 for 5 epochs before training all layers, you can do\n\n.. code-block:: python\n \n engine.run(\n save_dir='log/resnet50',\n max_epoch=60,\n eval_freq=10,\n test_only=False,\n fixbase_epoch=5,\n open_layers='classifier'\n )\n # or open_layers=['fc', 'classifier'] if there is another fc layer that\n # is randomly initialized, like resnet50_fc512\n\nNote that ``fixbase_epoch`` is counted into ``max_epoch``. In the above example, the base network will be fixed for 5 epochs and then open for training for 55 epochs. Thus, if you want to freeze some layers throughout the training, what you can do is to set ``fixbase_epoch`` equal to ``max_epoch`` and put the layer names in ``open_layers`` which you want to train.\n\n\nTest a trained model\n----------------------\nYou can load a trained model using :code:`torchreid.utils.load_pretrained_weights(model, weight_path)` and set ``test_only=True`` in ``engine.run()``.\n\n\nFine-tune a model pre-trained on reid datasets\n-----------------------------------------------\nUse :code:`torchreid.utils.load_pretrained_weights(model, weight_path)` to load the pre-trained weights and then fine-tune on the dataset you want.\n\n\nVisualize learning curves with tensorboard\n--------------------------------------------\nThe ``SummaryWriter()`` for tensorboard will be automatically initialized in ``engine.run()`` when you are training your model. Therefore, you do not need to do extra jobs. After the training is done, the ``*tf.events*`` file will be saved in ``save_dir``. Then, you just call ``tensorboard --logdir=your_save_dir`` in your terminal and visit ``http://localhost:6006/`` in a web browser. See `pytorch tensorboard `_ for further information.\n\n\nVisualize ranking results\n---------------------------\nThis can be achieved by setting ``visrank`` to true in ``engine.run()``. ``visrank_topk`` determines the top-k images to be visualized (Default is ``visrank_topk=10``). Note that ``visrank`` can only be used in test mode, i.e. ``test_only=True`` in ``engine.run()``. The output will be saved under ``save_dir/visrank_DATASETNAME`` where each plot contains the top-k similar gallery images given a query. An example is shown below where red and green denote incorrect and correct matches respectively.\n\n.. image:: figures/ranking_results.jpg\n :width: 800px\n :align: center\n\n\nVisualize activation maps\n--------------------------\nTo understand where the CNN focuses on to extract features for ReID, you can visualize the activation maps as in `OSNet `_. This is implemented in ``tools/visualize_actmap.py`` (check the code for more details). An example running command is\n\n.. code-block:: shell\n \n python tools/visualize_actmap.py \\\n --root $DATA/reid \\\n -d market1501 \\\n -m osnet_x1_0 \\\n --weights PATH_TO_PRETRAINED_WEIGHTS \\\n --save-dir log/visactmap_osnet_x1_0_market1501\n\nThe output will look like (from left to right: image, activation map, overlapped image)\n\n.. image:: figures/actmap.jpg\n :width: 300px\n :align: center\n\n\n.. note::\n In order to visualize activation maps, the CNN needs to output the last convolutional feature maps at eval mode. See ``torchreid/models/osnet.py`` for example.\n\n\nUse your own dataset\n----------------------\n1. Write your own dataset class. Below is a template for image dataset. However, it can also be applied to a video dataset class, for which you simply change ``ImageDataset`` to ``VideoDataset``.\n\n.. code-block:: python\n \n from __future__ import absolute_import\n from __future__ import print_function\n from __future__ import division\n\n import sys\n import os\n import os.path as osp\n\n from torchreid.data import ImageDataset\n\n\n class NewDataset(ImageDataset):\n dataset_dir = 'new_dataset'\n\n def __init__(self, root='', **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n\n # All you need to do here is to generate three lists,\n # which are train, query and gallery.\n # Each list contains tuples of (img_path, pid, camid),\n # where\n # - img_path (str): absolute path to an image.\n # - pid (int): person ID, e.g. 0, 1.\n # - camid (int): camera ID, e.g. 0, 1.\n # Note that\n # - pid and camid should be 0-based.\n # - query and gallery should share the same pid scope (e.g.\n # pid=0 in query refers to the same person as pid=0 in gallery).\n # - train, query and gallery share the same camid scope (e.g.\n # camid=0 in train refers to the same camera as camid=0\n # in query/gallery).\n train = ...\n query = ...\n gallery = ...\n\n super(NewDataset, self).__init__(train, query, gallery, **kwargs)\n\n\n2. Register your dataset.\n\n.. code-block:: python\n \n import torchreid\n torchreid.data.register_image_dataset('new_dataset', NewDataset)\n\n\n3. Initialize a data manager with your dataset.\n\n.. code-block:: python\n \n # use your own dataset only\n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources='new_dataset'\n )\n # combine with other datasets\n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources=['new_dataset', 'dukemtmcreid']\n )\n # cross-dataset evaluation\n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources=['new_dataset', 'dukemtmcreid'],\n targets='market1501' # or targets=['market1501', 'cuhk03']\n )\n\n\n\nDesign your own Engine\n------------------------\nA new Engine should be designed if you have your own loss function. The base Engine class ``torchreid.engine.Engine`` has implemented some generic methods which you can inherit to avoid re-writing. Please refer to the source code for more details. You are suggested to see how ``ImageSoftmaxEngine`` and ``ImageTripletEngine`` are constructed (also ``VideoSoftmaxEngine`` and ``VideoTripletEngine``). All you need to implement might be just a ``forward_backward()`` function.\n\n\nUse Torchreid as a feature extractor in your projects\n-------------------------------------------------------\nWe have provided a simple API for feature extraction, which accepts input of various types such as a list of image paths or numpy arrays. More details can be found in the code at ``torchreid/utils/feature_extractor.py``. Here we show a simple example of how to extract features given a list of image paths.\n\n.. code-block:: python\n\n from torchreid.utils import FeatureExtractor\n\n extractor = FeatureExtractor(\n model_name='osnet_x1_0',\n model_path='a/b/c/model.pth.tar',\n device='cuda'\n )\n\n image_list = [\n 'a/b/c/image001.jpg',\n 'a/b/c/image002.jpg',\n 'a/b/c/image003.jpg',\n 'a/b/c/image004.jpg',\n 'a/b/c/image005.jpg'\n ]\n\n features = extractor(image_list)\n print(features.shape) # output (5, 512)\n" }, { "path": "linter.sh", "content": "echo \"Running isort\"\nisort -y -sp .\necho \"Done\"\n\necho \"Running yapf\"\nyapf -i -r -vv -e build .\necho \"Done\"\n\necho \"Running flake8\"\nflake8 .\necho \"Done\"" }, { "path": "pyproject.toml", "content": "[project]\nname = \"torchreid\"\nauthors = [{name=\"Vladimir Somers\"}, {name=\"Kaiyang Zhou\"}]\nurls = {homepage = \"https://github.com/VlSomers/bpbreid\"}\ndynamic = [\"version\", \"description\", \"license\", \"dependencies\", \"keywords\", \"readme\", \"optional-dependencies\"]\n\n[tool.setuptools.dynamic]\nversion = {attr = \"torchreid.__version__\"}\n\n[build-system]\nrequires = [\"setuptools\", \"numpy\", \"cython\"]\nbuild-backend = \"setuptools.build_meta\"\n" }, { "path": "requirements.txt", "content": "# pip install -r requirements.txt\n\n# base torchreid ----------------------------------------\nnumpy\nCython\nh5py\nPillow\nsix\nscipy\nopencv-python\nmatplotlib\ntb-nightly\nfuture\nyacs\ngdown\nflake8\nyapf\nisort\n\n# added for bpbreid\nalbumentations\npandas\ntabulate\ndeepdiff\nwandb\nmonai\ntorchmetrics==0.10.3\n\n" }, { "path": "requirements_labels.txt", "content": "openpifpaf\ndetectron2 @ git+https://github.com/facebookresearch/detectron2.git\n" }, { "path": "setup.py", "content": "import numpy as np\nimport os.path as osp\nfrom setuptools import setup, find_packages\nfrom distutils.extension import Extension\nfrom Cython.Build import cythonize\n\ndef readme():\n with open('README.md') as f:\n content = f.read()\n return content\n\n\ndef numpy_include():\n try:\n numpy_include = np.get_include()\n except AttributeError:\n numpy_include = np.get_numpy_include()\n return numpy_include\n\n\next_modules = [\n Extension(\n 'torchreid.metrics.rank_cylib.rank_cy',\n ['torchreid/metrics/rank_cylib/rank_cy.pyx'],\n include_dirs=[numpy_include()],\n )\n]\n\n\ndef get_requirements(filename='requirements.txt'):\n here = osp.dirname(osp.realpath(__file__))\n with open(osp.join(here, filename), 'r') as f:\n requires = [line.replace('\\n', '') for line in f.readlines()]\n return requires\n\nsetup(\n description='A library for deep learning person re-ID in PyTorch',\n license='MIT',\n long_description=readme(),\n packages=find_packages(),\n install_requires=get_requirements(),\n extras_require={\"labels\": get_requirements(\"requirements_labels.txt\")},\n keywords=['Person Re-Identification', 'Deep Learning', 'Computer Vision'],\n ext_modules=cythonize(ext_modules)\n)\n" }, { "path": "torchreid/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom torchreid import data, optim, utils, engine, losses, models, metrics\n\n__version__ = '1.2.4'\n__author__ = 'Kaiyang Zhou'\n__homepage__ = 'https://kaiyangzhou.github.io/'\n__description__ = 'Deep learning person re-identification in PyTorch'\n__url__ = 'https://github.com/KaiyangZhou/deep-person-reid'\n" }, { "path": "torchreid/data/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom .datasets import (\n Dataset, ImageDataset, VideoDataset, register_image_dataset,\n register_video_dataset, get_dataset_nickname\n)\nfrom .datamanager import ImageDataManager, VideoDataManager\n" }, { "path": "torchreid/data/data_augmentation/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom .random_occlusion import *\n" }, { "path": "torchreid/data/data_augmentation/random_occlusion.py", "content": "# Source: https://github.com/isarandi/synthetic-occlusion/blob/master/augmentation.py\nimport math\nimport os.path\nimport random\nimport sys\nimport xml.etree.ElementTree\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport skimage.data\nimport cv2\nimport PIL.Image\nfrom albumentations import (\n DualTransform, functional\n)\n\n\ndef main():\n \"\"\"Demo of how to use the code\"\"\"\n\n # path = 'something/something/VOCtrainval_11-May-2012/VOCdevkit/VOC2012'\n path = sys.argv[1]\n\n print('Loading occluders from Pascal VOC dataset...')\n occluders = load_occluders(pascal_voc_root_path=path)\n print('Found {} suitable objects'.format(len(occluders)))\n\n original_im = cv2.resize(skimage.data.astronaut(), (256, 256))\n fig, axarr = plt.subplots(3, 3, figsize=(7, 7))\n for ax in axarr.ravel():\n occluded_im = occlude_with_objects(original_im, occluders)\n ax.imshow(occluded_im, interpolation=\"none\")\n ax.axis('off')\n\n fig.tight_layout(h_pad=0)\n # plt.savefig('examples.jpg', dpi=150, bbox_inches='tight')\n plt.show()\n\n\ndef load_occluders(\n pascal_voc_root_path,\n classes_filter=None,\n):\n occluders = []\n structuring_element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (8, 8))\n if classes_filter is None:\n # classes_filter = [\"person\", \"bicycle\", \"boat\", \"bus\", \"car\", \"motorbike\", \"train\", \"chair\", \"dining\", \"table\", \"plant\", \"sofa\"]\n classes_filter = [\"person\", \"bicycle\", \"boat\", \"bus\", \"car\", \"motorbike\", \"train\"]\n # classes_filter = [\"person\"]\n annotation_paths = list_filepaths(os.path.join(pascal_voc_root_path, 'Annotations'))\n for annotation_path in annotation_paths:\n xml_root = xml.etree.ElementTree.parse(annotation_path).getroot()\n is_segmented = (xml_root.find('segmented').text != '0')\n\n if not is_segmented:\n continue\n\n boxes = []\n for i_obj, obj in enumerate(xml_root.findall('object')):\n is_authorized_class = (obj.find('name').text in classes_filter)\n is_difficult = (obj.find('difficult').text != '0')\n is_truncated = (obj.find('truncated').text != '0')\n if is_authorized_class and not is_difficult and not is_truncated:\n bndbox = obj.find('bndbox')\n box = [int(bndbox.find(s).text) for s in ['xmin', 'ymin', 'xmax', 'ymax']]\n boxes.append((i_obj, box))\n\n if not boxes:\n continue\n\n im_filename = xml_root.find('filename').text\n seg_filename = im_filename.replace('jpg', 'png')\n\n im_path = os.path.join(pascal_voc_root_path, 'JPEGImages', im_filename)\n seg_path = os.path.join(pascal_voc_root_path, 'SegmentationObject', seg_filename)\n\n im = np.asarray(PIL.Image.open(im_path))\n labels = np.asarray(PIL.Image.open(seg_path))\n\n for i_obj, (xmin, ymin, xmax, ymax) in boxes:\n object_mask = (labels[ymin:ymax, xmin:xmax] == i_obj + 1).astype(np.uint8) * 255\n object_image = im[ymin:ymax, xmin:xmax]\n if cv2.countNonZero(object_mask) < 500:\n # Ignore small objects\n continue\n\n # Reduce the opacity of the mask along the border for smoother blending\n eroded = cv2.erode(object_mask, structuring_element)\n object_mask[eroded < object_mask] = 192\n object_with_mask = np.concatenate([object_image, object_mask[..., np.newaxis]], axis=-1)\n\n # Downscale for efficiency\n object_with_mask = resize_by_factor(object_with_mask, 0.5)\n occluders.append(object_with_mask)\n\n return occluders\n\n\ndef occlude_with_objects(im, occluders, n=1, min_overlap=0.1, max_overlap=0.6):\n \"\"\"Returns an augmented version of `im`, containing some occluders from the Pascal VOC dataset.\"\"\"\n\n result = im.copy()\n width_height = np.asarray([im.shape[1], im.shape[0]])\n im_area = im.shape[1] * im.shape[0]\n count = np.random.randint(1, n+1)\n\n for _ in range(count):\n occluder = random.choice(occluders)\n occluder_area = occluder.shape[1] * occluder.shape[0]\n overlap = random.uniform(min_overlap, max_overlap)\n scale_factor = math.sqrt(overlap * im_area / occluder_area)\n occluder = resize_by_factor(occluder, scale_factor)\n assert (occluder.shape[1] * occluder.shape[0]) / im_area == overlap\n center = np.random.uniform([0, 0], width_height)\n paste_over(im_src=occluder, im_dst=result, center=center)\n return result\n\n\ndef paste_over(im_src, im_dst, center, is_mask=False):\n \"\"\"Pastes `im_src` onto `im_dst` at a specified position, with alpha blending, in place.\n\n Locations outside the bounds of `im_dst` are handled as expected (only a part or none of\n `im_src` becomes visible).\n\n Args:\n im_src: The RGBA image to be pasted onto `im_dst`. Its size can be arbitrary.\n im_dst: The target image.\n alpha: A float (0.0-1.0) array of the same size as `im_src` controlling the alpha blending\n at each pixel. Large values mean more visibility for `im_src`.\n center: coordinates in `im_dst` where the center of `im_src` should be placed.\n \"\"\"\n\n width_height_src = np.asarray([im_src.shape[1], im_src.shape[0]])\n width_height_dst = np.asarray([im_dst.shape[1], im_dst.shape[0]])\n\n center = np.round(center).astype(np.int32)\n raw_start_dst = center - width_height_src // 2\n raw_end_dst = raw_start_dst + width_height_src\n\n start_dst = np.clip(raw_start_dst, 0, width_height_dst)\n end_dst = np.clip(raw_end_dst, 0, width_height_dst)\n region_dst = im_dst[start_dst[1]:end_dst[1], start_dst[0]:end_dst[0]]\n\n start_src = start_dst - raw_start_dst\n end_src = width_height_src + (end_dst - raw_end_dst)\n region_src = im_src[start_src[1]:end_src[1], start_src[0]:end_src[0]]\n color_src = region_src[..., 0:3]\n alpha = region_src[..., 3:].astype(np.float32) / 255\n\n if is_mask: # if this is a segmentation mask, just apply alpha erasing\n im_dst[start_dst[1]:end_dst[1], start_dst[0]:end_dst[0]] = (1 - alpha) * region_dst\n else:\n im_dst[start_dst[1]:end_dst[1], start_dst[0]:end_dst[0]] = (\n alpha * color_src + (1 - alpha) * region_dst)\n im_area = im_src.shape[1] * im_src.shape[0]\n bbox_overlap = (color_src.shape[0] * color_src.shape[1]) / im_area\n pxls_overlap = np.count_nonzero(alpha) / im_area\n return bbox_overlap, pxls_overlap\n\n\ndef resize_by_factor(im, factor):\n \"\"\"Returns a copy of `im` resized by `factor`, using bilinear interp for up and area interp\n for downscaling.\n \"\"\"\n new_size = tuple(np.round(np.array([im.shape[1], im.shape[0]]) * factor).astype(int))\n interp = cv2.INTER_LINEAR if factor > 1.0 else cv2.INTER_AREA\n return cv2.resize(im, new_size, fx=factor, fy=factor, interpolation=interp)\n\n\ndef list_filepaths(dirpath):\n names = os.listdir(dirpath)\n paths = [os.path.join(dirpath, name) for name in names]\n return sorted(filter(os.path.isfile, paths))\n\n\nclass RandomOcclusion(DualTransform):\n def __init__(self,\n path,\n im_shape,\n always_apply=False,\n p=.5,\n n=1,\n min_overlap=0.5,\n max_overlap=0.8,\n ):\n super(RandomOcclusion, self).__init__(always_apply, p)\n print('Loading occluders from Pascal VOC dataset...')\n self.all_occluders = load_occluders(pascal_voc_root_path=path)\n self.bbox_overlaps = []\n self.pxls_overlaps = []\n self.count = 0\n self.n = n\n self.min_overlap = min_overlap\n self.max_overlap = max_overlap\n self.im_shape = im_shape\n\n def check_range(self, dimension):\n if isinstance(dimension, float) and not 0 <= dimension < 1.0:\n raise ValueError(\n \"Invalid value {}. If using floats, the value should be in the range [0.0, 1.0)\".format(dimension)\n )\n\n def apply(self, image, occluders=(), centers=(), **params):\n for occluder, center in zip(occluders, centers):\n bbox_overlap, pxls_overlap = paste_over(im_src=occluder, im_dst=image, center=center)\n self.bbox_overlaps.append(bbox_overlap)\n self.pxls_overlaps.append(pxls_overlap)\n self.count += 1\n if self.count % 10000 == 0:\n bbox_overlaps = np.array(self.bbox_overlaps)\n pxls_overlaps = np.array(self.pxls_overlaps)\n print(\"RandomOcclusion #{}: bbox_overlap=[{:.2f},{:.2f},{:.2f}], pxls_overlap=[{:.2f},{:.2f},{:.2f}]\"\n .format(self.count,\n bbox_overlaps.min(),\n bbox_overlaps.max(),\n bbox_overlaps.mean(),\n pxls_overlaps.min(),\n pxls_overlaps.max(),\n pxls_overlaps.mean()\n )\n )\n return image\n\n def apply_to_mask(self, image, occluders=(), centers=(), **params):\n for occluder, center in zip(occluders, centers):\n paste_over(im_src=occluder, im_dst=image, center=center, is_mask=True)\n return image\n\n def get_params_dependent_on_targets(self, params):\n img = params[\"image\"]\n count = np.random.randint(1, self.n + 1)\n width_height = np.asarray([img.shape[1], img.shape[0]])\n im_area = self.im_shape[1] * self.im_shape[0]\n occluders = []\n centers = []\n for _ in range(count):\n occluder = random.choice(self.all_occluders)\n occluder_area = occluder.shape[1] * occluder.shape[0]\n overlap = random.uniform(self.min_overlap, self.max_overlap)\n scale_factor = math.sqrt(overlap * im_area / occluder_area)\n occluder = resize_by_factor(occluder, scale_factor)\n # assert abs((occluder.shape[1] * occluder.shape[0]) / im_area - overlap) < 0.005\n center = np.random.uniform([0, 0], width_height)\n occluders.append(occluder)\n centers.append(center)\n\n return {\"occluders\": occluders,\n \"centers\": centers}\n\n @property\n def targets_as_params(self):\n return [\"image\"]\n\n def get_transform_init_args_names(self):\n return (\n \"max_holes\",\n \"max_height\",\n \"max_width\",\n \"min_holes\",\n \"min_height\",\n \"min_width\",\n \"fill_value\",\n \"mask_fill_value\",\n )\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "torchreid/data/datamanager.py", "content": "from __future__ import division, print_function, absolute_import\nimport torch\n\nfrom torchreid.data.masks_transforms import masks_preprocess_transforms\nfrom torchreid.data.sampler import build_train_sampler\nfrom torchreid.data.datasets import init_image_dataset, init_video_dataset, get_image_dataset\nfrom torchreid.data.transforms import build_transforms\n\n\nclass DataManager(object):\n r\"\"\"Base data manager.\n\n Args:\n sources (str or list): source dataset(s).\n targets (str or list, optional): target dataset(s). If not given,\n it equals to ``sources``.\n height (int, optional): target image height. Default is 256.\n width (int, optional): target image width. Default is 128.\n transforms (str or list of str, optional): transformations applied to model training.\n Default is 'random_flip'.\n norm_mean (list or None, optional): data mean. Default is None (use imagenet mean).\n norm_std (list or None, optional): data std. Default is None (use imagenet std).\n use_gpu (bool, optional): use gpu. Default is True.\n \"\"\"\n\n def __init__(\n self,\n config,\n sources=None,\n targets=None,\n height=256,\n width=128,\n transforms='random_flip',\n norm_mean=None,\n norm_std=None,\n use_gpu=False,\n use_masks=False,\n masks_dir='',\n ):\n self.sources = sources\n self.targets = targets\n self.height = height\n self.width = width\n self.masks_dir = masks_dir\n self.config = config\n\n if self.sources is None:\n raise ValueError('sources must not be None')\n\n if isinstance(self.sources, str):\n self.sources = [self.sources]\n\n if self.targets is None:\n self.targets = self.sources\n\n if isinstance(self.targets, str):\n self.targets = [self.targets]\n\n masks_config = get_image_dataset(self.sources[0]).get_masks_config(self.masks_dir)\n self.transform_tr, self.transform_te = build_transforms(\n self.height,\n self.width,\n config,\n transforms=transforms,\n norm_mean=norm_mean,\n norm_std=norm_std,\n remove_background_mask=masks_config[1] if masks_config is not None else False,\n masks_preprocess=config.model.bpbreid.masks.preprocess,\n softmax_weight=config.model.bpbreid.masks.softmax_weight,\n background_computation_strategy=config.model.bpbreid.masks.background_computation_strategy,\n mask_filtering_threshold=config.model.bpbreid.masks.mask_filtering_threshold,\n )\n\n self.use_gpu = (torch.cuda.is_available() and use_gpu)\n\n @property\n def num_train_pids(self):\n \"\"\"Returns the number of training person identities.\"\"\"\n return self._num_train_pids\n\n @property\n def num_train_cams(self):\n \"\"\"Returns the number of training cameras.\"\"\"\n return self._num_train_cams\n\n def fetch_test_loaders(self, name):\n \"\"\"Returns query and gallery of a test dataset, each containing\n tuples of (img_path(s), pid, camid).\n\n Args:\n name (str): dataset name.\n \"\"\"\n query_loader = self.test_dataset[name]['query']\n gallery_loader = self.test_dataset[name]['gallery']\n return query_loader, gallery_loader\n\n def preprocess_pil_img(self, img):\n \"\"\"Transforms a PIL image to torch tensor for testing.\"\"\"\n return self.transform_te(img)\n\n\nclass ImageDataManager(DataManager):\n r\"\"\"Image data manager.\n\n Args:\n root (str): root path to datasets.\n sources (str or list): source dataset(s).\n targets (str or list, optional): target dataset(s). If not given,\n it equals to ``sources``.\n height (int, optional): target image height. Default is 256.\n width (int, optional): target image width. Default is 128.\n transforms (str or list of str, optional): transformations applied to model training.\n Default is 'random_flip'.\n norm_mean (list or None, optional): data mean. Default is None (use imagenet mean).\n norm_std (list or None, optional): data std. Default is None (use imagenet std).\n use_gpu (bool, optional): use gpu. Default is True.\n split_id (int, optional): split id (*0-based*). Default is 0.\n combineall (bool, optional): combine train, query and gallery in a dataset for\n training. Default is False.\n load_train_targets (bool, optional): construct train-loader for target datasets.\n Default is False. This is useful for domain adaptation research.\n batch_size_train (int, optional): number of images in a training batch. Default is 32.\n batch_size_test (int, optional): number of images in a test batch. Default is 32.\n workers (int, optional): number of workers. Default is 4.\n num_instances (int, optional): number of instances per identity in a batch.\n Default is 4.\n train_sampler (str, optional): sampler. Default is RandomSampler.\n train_sampler_t (str, optional): sampler for target train loader. Default is RandomSampler.\n cuhk03_labeled (bool, optional): use cuhk03 labeled images.\n Default is False (defaul is to use detected images).\n cuhk03_classic_split (bool, optional): use the classic split in cuhk03.\n Default is False.\n market1501_500k (bool, optional): add 500K distractors to the gallery\n set in market1501. Default is False.\n\n Examples::\n\n datamanager = torchreid.data.ImageDataManager(\n root='path/to/reid-data',\n sources='market1501',\n height=256,\n width=128,\n batch_size_train=32,\n batch_size_test=100\n )\n\n # return train loader of source data\n train_loader = datamanager.train_loader\n\n # return test loader of target data\n test_loader = datamanager.test_loader\n\n # return train loader of target data\n train_loader_t = datamanager.train_loader_t\n \"\"\"\n data_type = 'image'\n\n def __init__(\n self,\n config,\n root='',\n sources=None,\n targets=None,\n height=256,\n width=128,\n mask_scale=8,\n transforms='random_flip',\n norm_mean=None,\n norm_std=None,\n use_gpu=True,\n split_id=0,\n combineall=False,\n load_train_targets=False,\n batch_size_train=32,\n batch_size_test=32,\n workers=4,\n num_instances=4,\n train_sampler='RandomSampler',\n train_sampler_t='RandomSampler',\n cuhk03_labeled=False,\n cuhk03_classic_split=False,\n market1501_500k=False,\n use_masks=False,\n masks_dir=None,\n ):\n\n super(ImageDataManager, self).__init__(\n sources=sources,\n targets=targets,\n height=height,\n width=width,\n transforms=transforms,\n norm_mean=norm_mean,\n norm_std=norm_std,\n use_gpu=use_gpu,\n use_masks=use_masks,\n masks_dir=masks_dir,\n config=config\n )\n\n print('=> Loading train (source) dataset')\n trainset = []\n for name in self.sources:\n trainset_ = init_image_dataset(\n name,\n config=config,\n transform_tr=self.transform_tr,\n transform_te=self.transform_te,\n mode='train',\n combineall=combineall,\n root=root,\n split_id=split_id,\n cuhk03_labeled=cuhk03_labeled,\n cuhk03_classic_split=cuhk03_classic_split,\n market1501_500k=market1501_500k,\n use_masks=use_masks,\n masks_dir=masks_dir,\n load_masks=self.config.model.bpbreid.masks.preprocess in masks_preprocess_transforms,\n )\n trainset.append(trainset_)\n trainset = sum(trainset)\n\n self._num_train_pids = trainset.num_train_pids\n self._num_train_cams = trainset.num_train_cams\n\n self.train_loader = torch.utils.data.DataLoader(\n trainset,\n sampler=build_train_sampler(\n trainset.train,\n train_sampler,\n batch_size=batch_size_train,\n num_instances=num_instances\n ),\n batch_size=batch_size_train,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=True\n )\n\n self.train_loader_t = None\n if load_train_targets:\n # check if sources and targets are identical\n assert len(set(self.sources) & set(self.targets)) == 0, \\\n 'sources={} and targets={} must not have overlap'.format(self.sources, self.targets)\n\n print('=> Loading train (target) dataset')\n trainset_t = []\n for name in self.targets:\n trainset_t_ = init_image_dataset(\n name,\n config=config,\n transform_tr=self.transform_tr,\n transform_te=self.transform_te,\n mode='train',\n combineall=False, # only use the training data\n root=root,\n split_id=split_id,\n cuhk03_labeled=cuhk03_labeled,\n cuhk03_classic_split=cuhk03_classic_split,\n market1501_500k=market1501_500k,\n use_masks=use_masks,\n masks_dir=masks_dir,\n load_masks=self.config.model.bpbreid.masks.preprocess in masks_preprocess_transforms,\n )\n trainset_t.append(trainset_t_)\n trainset_t = sum(trainset_t)\n\n self.train_loader_t = torch.utils.data.DataLoader(\n trainset_t,\n sampler=build_train_sampler(\n trainset_t.train,\n train_sampler_t,\n batch_size=batch_size_train,\n num_instances=num_instances\n ),\n batch_size=batch_size_train,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=True\n )\n\n print('=> Loading test (target) dataset')\n self.test_loader = {\n name: {\n 'query': None,\n 'gallery': None\n }\n for name in self.targets\n }\n self.test_dataset = {\n name: {\n 'query': None,\n 'gallery': None\n }\n for name in self.targets\n }\n\n for name in self.targets:\n # build query loader\n queryset = init_image_dataset(\n name,\n config=config,\n transform_tr=self.transform_tr,\n transform_te=self.transform_te,\n mode='query',\n combineall=combineall,\n root=root,\n split_id=split_id,\n cuhk03_labeled=cuhk03_labeled,\n cuhk03_classic_split=cuhk03_classic_split,\n market1501_500k=market1501_500k,\n use_masks=use_masks,\n masks_dir=masks_dir,\n load_masks=self.config.model.bpbreid.masks.preprocess in masks_preprocess_transforms,\n )\n self.test_loader[name]['query'] = torch.utils.data.DataLoader(\n queryset,\n batch_size=batch_size_test,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=False\n )\n\n # build gallery loader\n galleryset = init_image_dataset(\n name,\n config=config,\n transform_tr=self.transform_tr,\n transform_te=self.transform_te,\n mode='gallery',\n combineall=combineall,\n verbose=False,\n root=root,\n split_id=split_id,\n cuhk03_labeled=cuhk03_labeled,\n cuhk03_classic_split=cuhk03_classic_split,\n market1501_500k=market1501_500k,\n use_masks=use_masks,\n masks_dir=masks_dir,\n load_masks=self.config.model.bpbreid.masks.preprocess in masks_preprocess_transforms,\n )\n self.test_loader[name]['gallery'] = torch.utils.data.DataLoader(\n galleryset,\n batch_size=batch_size_test,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=False\n )\n\n self.test_dataset[name]['query'] = queryset.query\n self.test_dataset[name]['gallery'] = galleryset.gallery\n\n print('\\n')\n print(' **************** Summary ****************')\n print(' source : {}'.format(self.sources))\n print(' # source datasets : {}'.format(len(self.sources)))\n print(' # source ids : {}'.format(self.num_train_pids))\n print(' # source images : {}'.format(len(trainset)))\n print(' # source cameras : {}'.format(self.num_train_cams))\n if load_train_targets:\n print(\n ' # target images : {} (unlabeled)'.format(len(trainset_t))\n )\n print(' target : {}'.format(self.targets))\n print(' *****************************************')\n print('\\n')\n\n\n\nclass VideoDataManager(DataManager):\n r\"\"\"Video data manager.\n\n Args:\n root (str): root path to datasets.\n sources (str or list): source dataset(s).\n targets (str or list, optional): target dataset(s). If not given,\n it equals to ``sources``.\n height (int, optional): target image height. Default is 256.\n width (int, optional): target image width. Default is 128.\n transforms (str or list of str, optional): transformations applied to model training.\n Default is 'random_flip'.\n norm_mean (list or None, optional): data mean. Default is None (use imagenet mean).\n norm_std (list or None, optional): data std. Default is None (use imagenet std).\n use_gpu (bool, optional): use gpu. Default is True.\n split_id (int, optional): split id (*0-based*). Default is 0.\n combineall (bool, optional): combine train, query and gallery in a dataset for\n training. Default is False.\n batch_size_train (int, optional): number of tracklets in a training batch. Default is 3.\n batch_size_test (int, optional): number of tracklets in a test batch. Default is 3.\n workers (int, optional): number of workers. Default is 4.\n num_instances (int, optional): number of instances per identity in a batch.\n Default is 4.\n train_sampler (str, optional): sampler. Default is RandomSampler.\n seq_len (int, optional): how many images to sample in a tracklet. Default is 15.\n sample_method (str, optional): how to sample images in a tracklet. Default is \"evenly\".\n Choices are [\"evenly\", \"random\", \"all\"]. \"evenly\" and \"random\" will sample ``seq_len``\n images in a tracklet while \"all\" samples all images in a tracklet, where the batch size\n needs to be set to 1.\n\n Examples::\n\n datamanager = torchreid.data.VideoDataManager(\n root='path/to/reid-data',\n sources='mars',\n height=256,\n width=128,\n batch_size_train=3,\n batch_size_test=3,\n seq_len=15,\n sample_method='evenly'\n )\n\n # return train loader of source data\n train_loader = datamanager.train_loader\n\n # return test loader of target data\n test_loader = datamanager.test_loader\n\n .. note::\n The current implementation only supports image-like training. Therefore, each image in a\n sampled tracklet will undergo independent transformation functions. To achieve tracklet-aware\n training, you need to modify the transformation functions for video reid such that each function\n applies the same operation to all images in a tracklet to keep consistency.\n \"\"\"\n data_type = 'video'\n\n def __init__(\n self,\n config,\n root='',\n sources=None,\n targets=None,\n height=256,\n width=128,\n transforms='random_flip',\n norm_mean=None,\n norm_std=None,\n use_gpu=True,\n split_id=0,\n combineall=False,\n batch_size_train=3,\n batch_size_test=3,\n workers=4,\n num_instances=4,\n train_sampler='RandomSampler',\n seq_len=15,\n sample_method='evenly'\n ):\n\n super(VideoDataManager, self).__init__(\n config=config,\n sources=sources,\n targets=targets,\n height=height,\n width=width,\n transforms=transforms,\n norm_mean=norm_mean,\n norm_std=norm_std,\n use_gpu=use_gpu\n )\n\n print('=> Loading train (source) dataset')\n trainset = []\n for name in self.sources:\n trainset_ = init_video_dataset(\n name,\n transform=self.transform_tr,\n mode='train',\n combineall=combineall,\n root=root,\n split_id=split_id,\n seq_len=seq_len,\n sample_method=sample_method\n )\n trainset.append(trainset_)\n trainset = sum(trainset)\n\n self._num_train_pids = trainset.num_train_pids\n self._num_train_cams = trainset.num_train_cams\n\n train_sampler = build_train_sampler(\n trainset.train,\n train_sampler,\n batch_size=batch_size_train,\n num_instances=num_instances\n )\n\n self.train_loader = torch.utils.data.DataLoader(\n trainset,\n sampler=train_sampler,\n batch_size=batch_size_train,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=True\n )\n\n print('=> Loading test (target) dataset')\n self.test_loader = {\n name: {\n 'query': None,\n 'gallery': None\n }\n for name in self.targets\n }\n self.test_dataset = {\n name: {\n 'query': None,\n 'gallery': None\n }\n for name in self.targets\n }\n\n for name in self.targets:\n # build query loader\n queryset = init_video_dataset(\n name,\n transform=self.transform_te,\n mode='query',\n combineall=combineall,\n root=root,\n split_id=split_id,\n seq_len=seq_len,\n sample_method=sample_method\n )\n self.test_loader[name]['query'] = torch.utils.data.DataLoader(\n queryset,\n batch_size=batch_size_test,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=False\n )\n\n # build gallery loader\n galleryset = init_video_dataset(\n name,\n transform=self.transform_te,\n mode='gallery',\n combineall=combineall,\n verbose=False,\n root=root,\n split_id=split_id,\n seq_len=seq_len,\n sample_method=sample_method\n )\n self.test_loader[name]['gallery'] = torch.utils.data.DataLoader(\n galleryset,\n batch_size=batch_size_test,\n shuffle=False,\n num_workers=workers,\n pin_memory=self.use_gpu,\n drop_last=False\n )\n\n self.test_dataset[name]['query'] = queryset.query\n self.test_dataset[name]['gallery'] = galleryset.gallery\n\n print('\\n')\n print(' **************** Summary ****************')\n print(' source : {}'.format(self.sources))\n print(' # source datasets : {}'.format(len(self.sources)))\n print(' # source ids : {}'.format(self.num_train_pids))\n print(' # source tracklets : {}'.format(len(trainset)))\n print(' # source cameras : {}'.format(self.num_train_cams))\n print(' target : {}'.format(self.targets))\n print(' *****************************************')\n print('\\n')\n" }, { "path": "torchreid/data/datasets/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nimport copy\n\nfrom .image import (\n GRID, PRID, CUHK01, CUHK02, CUHK03, MSMT17, VIPeR, SenseReID, Market1501,\n DukeMTMCreID, iLIDS, OccludedDuke, OccludedReID, Partial_iLIDS, Partial_REID, PDukemtmcReid,\n P_ETHZ\n)\nfrom .video import PRID2011, Mars, DukeMTMCVidReID, iLIDSVID\nfrom .dataset import Dataset, ImageDataset, VideoDataset\n\n__image_datasets = {\n 'market1501': Market1501,\n 'cuhk03': CUHK03,\n 'dukemtmcreid': DukeMTMCreID,\n 'msmt17': MSMT17,\n 'viper': VIPeR,\n 'grid': GRID,\n 'cuhk01': CUHK01,\n 'ilids': iLIDS,\n 'sensereid': SenseReID,\n 'prid': PRID,\n 'cuhk02': CUHK02,\n 'occluded_duke': OccludedDuke,\n 'occluded_reid': OccludedReID,\n 'partial_reid': Partial_REID,\n 'partial_ilids': Partial_iLIDS,\n 'p_ETHZ': P_ETHZ,\n 'p_dukemtmc_reid': PDukemtmcReid,\n}\n\n__datasets_nicknames = {\n 'market1501': 'mk',\n 'cuhk03': 'c03',\n 'dukemtmcreid': 'du',\n 'msmt17': 'ms',\n 'viper': 'vi',\n 'grid': 'gr',\n 'cuhk01': 'c01',\n 'ilids': 'il',\n 'sensereid': 'se',\n 'prid': 'pr',\n 'cuhk02': 'c02',\n 'occluded_duke': 'od',\n 'occluded_reid': 'or',\n 'partial_reid': 'pr',\n 'partial_ilids': 'pi',\n 'p_ETHZ': 'pz',\n 'p_dukemtmc_reid': 'pd',\n}\n\n__video_datasets = {\n 'mars': Mars,\n 'ilidsvid': iLIDSVID,\n 'prid2011': PRID2011,\n 'dukemtmcvidreid': DukeMTMCVidReID\n}\n\n__datasets_cache = {}\n\n\ndef configure_dataset_class(clazz, **ext_kwargs):\n \"\"\"\n Wrapper function to provide the class with args external to torchreid\n \"\"\"\n class ClazzWrapper(clazz):\n def __init__(self, **kwargs):\n self.__name__ = clazz.__name__\n super(ClazzWrapper, self).__init__(**{**kwargs, **ext_kwargs})\n\n ClazzWrapper.__name__ = clazz.__name__\n\n return ClazzWrapper\n\n\ndef get_dataset_nickname(name):\n return __datasets_nicknames.get(name, name)\n\n\ndef get_image_dataset(name):\n avai_datasets = list(__image_datasets.keys())\n if name not in avai_datasets:\n raise ValueError(\n 'Invalid dataset name. Received \"{}\", '\n 'but expected to be one of {}'.format(name, avai_datasets)\n )\n return __image_datasets[name]\n\n\ndef init_image_dataset(name, mode='train', **kwargs):\n \"\"\"\n Initializes an image dataset.\n The copy.copy() was introduced to fix Torchreid implementing multiple times the same dataset.\n In Datamanager, each dataset was instantiated multiple times via 'init_image_dataset': one for train, one for query\n and one for gallery. Each instance had its own 'data' field containing either train, query or gallery set, based on\n the 'mode' field passed as argument, and its own transforms, to perform training time or test time data transformation.\n However, instantiating the same dataset multiple times is not efficient, as it requires to load the dataset metadata from\n disk multiple times. Moreover, other printing (such as dataset summary) are displayed multiple times.\n To fix this, we copy the dataset class but not its contained objects (such as train/query/gallery) and set a new 'mode' on each copy.\n Thanks to that hack, the data list is created only once, and only the Dataset class is instantiated multiple times\n (for each 'mode'). Therefore, each Dataset uses the same data lists in the background, switching\n between train, query and gallery based on the 'mode' field.\n \"\"\"\n if name in __datasets_cache:\n print(\"Using cached dataset {}.\".format(name))\n dataset = __datasets_cache[name]\n else:\n print(\"Creating new dataset {} and add it to the datasets cache.\".format(name))\n dataset = get_image_dataset(name)(mode=mode, **kwargs)\n __datasets_cache[name] = dataset\n mode_dataset = copy.copy(dataset)\n mode_dataset.mode = mode\n return mode_dataset\n\n\ndef init_video_dataset(name, **kwargs):\n \"\"\"Initializes a video dataset.\"\"\"\n avai_datasets = list(__video_datasets.keys())\n if name not in avai_datasets:\n raise ValueError(\n 'Invalid dataset name. Received \"{}\", '\n 'but expected to be one of {}'.format(name, avai_datasets)\n )\n return __video_datasets[name](**kwargs)\n\n\ndef register_image_dataset(name, dataset, nickname=None):\n \"\"\"Registers a new image dataset.\n\n Args:\n name (str): key corresponding to the new dataset.\n dataset (Dataset): the new dataset class.\n\n Examples::\n \n import torchreid\n import NewDataset\n torchreid.data.register_image_dataset('new_dataset', NewDataset)\n # single dataset case\n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources='new_dataset'\n )\n # multiple dataset case\n datamanager = torchreid.data.ImageDataManager(\n root='reid-data',\n sources=['new_dataset', 'dukemtmcreid']\n )\n \"\"\"\n global __image_datasets\n curr_datasets = list(__image_datasets.keys())\n if name in curr_datasets:\n raise ValueError(\n 'The given name already exists, please choose '\n 'another name excluding {}'.format(curr_datasets)\n )\n __image_datasets[name] = dataset\n __datasets_nicknames[name] = nickname if nickname is not None else name\n\n\ndef register_video_dataset(name, dataset):\n \"\"\"Registers a new video dataset.\n\n Args:\n name (str): key corresponding to the new dataset.\n dataset (Dataset): the new dataset class.\n\n Examples::\n \n import torchreid\n import NewDataset\n torchreid.data.register_video_dataset('new_dataset', NewDataset)\n # single dataset case\n datamanager = torchreid.data.VideoDataManager(\n root='reid-data',\n sources='new_dataset'\n )\n # multiple dataset case\n datamanager = torchreid.data.VideoDataManager(\n root='reid-data',\n sources=['new_dataset', 'ilidsvid']\n )\n \"\"\"\n global __video_datasets\n curr_datasets = list(__video_datasets.keys())\n if name in curr_datasets:\n raise ValueError(\n 'The given name already exists, please choose '\n 'another name excluding {}'.format(curr_datasets)\n )\n __video_datasets[name] = dataset\n" }, { "path": "torchreid/data/datasets/dataset.py", "content": "from __future__ import division, print_function, absolute_import\nimport copy\nimport os\n\nimport numpy as np\nimport os.path as osp\nimport tarfile\nimport zipfile\nimport torch\n\nfrom torchreid.utils import read_masks, read_image, download_url, mkdir_if_missing\n\n\nclass Dataset(object):\n \"\"\"An abstract class representing a Dataset.\n\n This is the base class for ``ImageDataset`` and ``VideoDataset``.\n\n Args:\n train (list): contains tuples of (img_path(s), pid, camid).\n query (list): contains tuples of (img_path(s), pid, camid).\n gallery (list): contains tuples of (img_path(s), pid, camid).\n transform: transform function.\n mode (str): 'train', 'query' or 'gallery'.\n combineall (bool): combines train, query and gallery in a\n dataset for training.\n verbose (bool): show information.\n \"\"\"\n _junk_pids = [\n ] # contains useless person IDs, e.g. background, false detections\n\n masks_base_dir = None\n eval_metric = 'default' # default to market101\n\n def gallery_filter(self, q_pid, q_camid, q_ann, g_pids, g_camids, g_anns):\n \"\"\" Remove gallery samples that have the same pid and camid as the query sample, since ReID is a cross-camera\n person retrieval task for most datasets. However, we still keep samples from the same camera but of different\n identity as distractors.\"\"\"\n remove = (g_camids == q_camid) & (g_pids == q_pid)\n return remove\n\n def infer_masks_path(self, img_path):\n masks_path = os.path.join(self.dataset_dir, self.masks_base_dir, self.masks_dir, os.path.basename(os.path.dirname(img_path)), os.path.splitext(os.path.basename(img_path))[0] + self.masks_suffix)\n return masks_path\n\n def __init__(\n self,\n train,\n query,\n gallery,\n config=None,\n transform_tr=None,\n transform_te=None,\n mode='train',\n combineall=False,\n verbose=True,\n use_masks=False,\n masks_dir=None,\n masks_base_dir=None,\n load_masks=False,\n **kwargs\n ):\n self.train = train\n self.query = query\n self.gallery = gallery\n self.transform_tr = transform_tr\n self.transform_te = transform_te\n self.cfg = config\n self.mode = mode\n self.combineall = combineall\n self.verbose = verbose\n self.use_masks = use_masks\n self.masks_dir = masks_dir\n self.load_masks = load_masks\n if masks_base_dir is not None:\n self.masks_base_dir = masks_base_dir\n\n self.num_train_pids = self.get_num_pids(self.train)\n self.num_train_cams = self.get_num_cams(self.train)\n\n if self.combineall:\n self.combine_all()\n\n if self.verbose:\n self.show_summary()\n\n def transforms(self, mode):\n \"\"\"Returns the transforms of a specific mode.\"\"\"\n if mode == 'train':\n return self.transform_tr\n elif mode == 'query':\n return self.transform_te\n elif mode == 'gallery':\n return self.transform_te\n else:\n raise ValueError(\"Invalid mode. Got {}, but expected to be \"\n \"'train', 'query' or 'gallery'\".format(mode))\n\n def data(self, mode):\n \"\"\"Returns the data of a specific mode.\n\n Args:\n mode (str): 'train', 'query' or 'gallery'.\n\n Returns:\n list: contains tuples of (img_path(s), pid, camid).\n \"\"\"\n if mode == 'train':\n return self.train\n elif mode == 'query':\n return self.query\n elif mode == 'gallery':\n return self.gallery\n else:\n raise ValueError(\"Invalid mode. Got {}, but expected to be \"\n \"'train', 'query' or 'gallery'\".format(mode))\n\n def __getitem__(self, index):\n raise NotImplementedError\n\n def __len__(self): # kept for backward compatibility\n return self.len(self.mode)\n\n def len(self, mode):\n return len(self.data(mode))\n\n def __add__(self, other):\n \"\"\"Adds two datasets together (only the train set).\"\"\"\n train = copy.deepcopy(self.train)\n\n for sample in other.train:\n sample['pid'] += self.num_train_pids\n train.append(sample)\n\n ###################################\n # Things to do beforehand:\n # 1. set verbose=False to avoid unnecessary print\n # 2. set combineall=False because combineall would have been applied\n # if it was True for a specific dataset, setting it to True will\n # create new IDs that should have been included\n ###################################\n\n\n # FIXME find better implementation for combining datasets and masks\n assert self.use_masks == other.use_masks\n\n if isinstance(self, ImageDataset):\n return ImageDataset(\n train,\n self.query,\n self.gallery,\n transform=self.transform,\n mode=self.mode,\n combineall=False,\n verbose=False,\n use_masks=self.use_masks,\n masks_base_dir=self.masks_base_dir,\n )\n else:\n return VideoDataset(\n train,\n self.query,\n self.gallery,\n transform=self.transform,\n mode=self.mode,\n combineall=False,\n verbose=False,\n seq_len=self.seq_len,\n sample_method=self.sample_method\n )\n\n def __radd__(self, other):\n \"\"\"Supports sum([dataset1, dataset2, dataset3]).\"\"\"\n if other == 0:\n return self\n else:\n return self.__add__(other)\n\n def parse_data(self, data):\n \"\"\"Parses data list and returns the number of person IDs\n and the number of camera views.\n\n Args:\n data (list): contains tuples of (img_path(s), pid, camid)\n \"\"\"\n pids = set()\n cams = set()\n for i, sample in enumerate(data):\n pids.add(sample['pid'])\n cams.add(sample['camid'])\n return len(pids), len(cams)\n\n def get_num_pids(self, data):\n \"\"\"Returns the number of training person identities.\"\"\"\n return self.parse_data(data)[0]\n\n def get_num_cams(self, data):\n \"\"\"Returns the number of training cameras.\"\"\"\n return self.parse_data(data)[1]\n\n def show_summary(self):\n \"\"\"Shows dataset statistics.\"\"\"\n pass\n\n def combine_all(self):\n \"\"\"Combines train, query and gallery in a dataset for training.\"\"\"\n combined = copy.deepcopy(self.train)\n\n # relabel pids in gallery (query shares the same scope)\n g_pids = set()\n for sample in self.gallery:\n pid = sample['pid']\n if pid in self._junk_pids:\n continue\n g_pids.add(pid)\n pid2label = {pid: i for i, pid in enumerate(g_pids)}\n\n def _combine_data(data):\n for sample in data:\n pid = sample['pid']\n if pid in self._junk_pids:\n continue\n sample['pid'] = pid2label[pid] + self.num_train_pids\n combined.append(sample)\n\n _combine_data(self.query)\n _combine_data(self.gallery)\n\n self.train = combined\n self.num_train_pids = self.get_num_pids(self.train)\n\n def download_dataset(self, dataset_dir, dataset_url):\n \"\"\"Downloads and extracts dataset.\n\n Args:\n dataset_dir (str): dataset directory.\n dataset_url (str): url to download dataset.\n \"\"\"\n if osp.exists(dataset_dir):\n return\n\n if dataset_url is None:\n raise RuntimeError(\n '{} dataset needs to be manually '\n 'prepared, please follow the '\n 'document to prepare this dataset'.format(\n self.__class__.__name__\n )\n )\n\n print('Creating directory \"{}\"'.format(dataset_dir))\n mkdir_if_missing(dataset_dir)\n fpath = osp.join(dataset_dir, osp.basename(dataset_url))\n\n print(\n 'Downloading {} dataset to \"{}\"'.format(\n self.__class__.__name__, dataset_dir\n )\n )\n download_url(dataset_url, fpath)\n\n print('Extracting \"{}\"'.format(fpath))\n try:\n tar = tarfile.open(fpath)\n tar.extractall(path=dataset_dir)\n tar.close()\n except:\n zip_ref = zipfile.ZipFile(fpath, 'r')\n zip_ref.extractall(dataset_dir)\n zip_ref.close()\n\n print('{} dataset is ready'.format(self.__class__.__name__))\n\n def check_before_run(self, required_files):\n \"\"\"Checks if required files exist before going deeper.\n\n Args:\n required_files (str or list): string file name(s).\n \"\"\"\n if isinstance(required_files, str):\n required_files = [required_files]\n\n for fpath in required_files:\n if not osp.exists(fpath):\n raise RuntimeError('\"{}\" is not found'.format(fpath))\n\n def __repr__(self):\n num_train_pids, num_train_cams = self.parse_data(self.train)\n num_query_pids, num_query_cams = self.parse_data(self.query)\n num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)\n\n msg = ' ----------------------------------------\\n' \\\n ' subset | # ids | # items | # cameras\\n' \\\n ' ----------------------------------------\\n' \\\n ' train | {:5d} | {:7d} | {:9d}\\n' \\\n ' query | {:5d} | {:7d} | {:9d}\\n' \\\n ' gallery | {:5d} | {:7d} | {:9d}\\n' \\\n ' ----------------------------------------\\n' \\\n ' items: images/tracklets for image/video dataset\\n'.format(\n num_train_pids, len(self.train), num_train_cams,\n num_query_pids, len(self.query), num_query_cams,\n num_gallery_pids, len(self.gallery), num_gallery_cams\n )\n\n return msg\n\n\nclass ImageDataset(Dataset):\n \"\"\"A base class representing ImageDataset.\n\n All other image datasets should subclass it.\n\n ``__getitem__`` returns an image given index.\n It will return ``img``, ``pid``, ``camid`` and ``img_path``\n where ``img`` has shape (channel, height, width). As a result,\n data in each batch has shape (batch_size, channel, height, width).\n \"\"\"\n\n def __init__(self, train, query, gallery, **kwargs):\n super(ImageDataset, self).__init__(train, query, gallery, **kwargs)\n\n def __getitem__(self, index): # kept for backward compatibility\n return self.getitem(index, self.mode)\n\n def getitem(self, index, mode):\n # BPBreID can work with None masks\n # list all combination: source vs target, merged/joined vs not, cross domain or not, load from disk vs fixed for BoT/PBP transform vs None,\n # need masks when available for pixel accuracy prediction\n sample = self.data(mode)[index]\n transf_args = {\"image\": read_image(sample['img_path'])}\n if self.use_masks:\n if self.load_masks and 'masks_path' in sample:\n transf_args[\"mask\"] = read_masks(sample['masks_path'])\n elif not self.load_masks:\n # hack for BoT and PCB masks that are generated in transform().\n # FIXME BoT and PCB masks should not be generated here, but later in BPBreID model with a config\n transf_args[\"mask\"] = np.ones((1, 2, 2))\n else:\n pass\n result = self.transforms(mode)(**transf_args)\n sample.update(result)\n return sample\n\n def show_summary(self):\n num_train_pids, num_train_cams = self.parse_data(self.train)\n num_query_pids, num_query_cams = self.parse_data(self.query)\n num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)\n\n print('=> Loaded {}'.format(self.__class__.__name__))\n print(' ----------------------------------------')\n print(' subset | # ids | # images | # cameras')\n print(' ----------------------------------------')\n print(\n ' train | {:5d} | {:8d} | {:9d}'.format(\n num_train_pids, len(self.train), num_train_cams\n )\n )\n print(\n ' query | {:5d} | {:8d} | {:9d}'.format(\n num_query_pids, len(self.query), num_query_cams\n )\n )\n print(\n ' gallery | {:5d} | {:8d} | {:9d}'.format(\n num_gallery_pids, len(self.gallery), num_gallery_cams\n )\n )\n print(' ----------------------------------------')\n\n\nclass VideoDataset(Dataset):\n \"\"\"A base class representing VideoDataset.\n\n All other video datasets should subclass it.\n\n ``__getitem__`` returns an image given index.\n It will return ``imgs``, ``pid`` and ``camid``\n where ``imgs`` has shape (seq_len, channel, height, width). As a result,\n data in each batch has shape (batch_size, seq_len, channel, height, width).\n \"\"\"\n\n def __init__(\n self,\n train,\n query,\n gallery,\n seq_len=15,\n sample_method='evenly',\n **kwargs\n ):\n super(VideoDataset, self).__init__(train, query, gallery, **kwargs)\n self.seq_len = seq_len\n self.sample_method = sample_method\n\n if self.transform is None:\n raise RuntimeError('transform must not be None')\n\n def getitem(self, index, mode):\n img_paths, pid, camid = self.data(mode)[index] # FIXME new format\n num_imgs = len(img_paths)\n\n if self.sample_method == 'random':\n # Randomly samples seq_len images from a tracklet of length num_imgs,\n # if num_imgs is smaller than seq_len, then replicates images\n indices = np.arange(num_imgs)\n replace = False if num_imgs >= self.seq_len else True\n indices = np.random.choice(\n indices, size=self.seq_len, replace=replace\n )\n # sort indices to keep temporal order (comment it to be order-agnostic)\n indices = np.sort(indices)\n\n elif self.sample_method == 'evenly':\n # Evenly samples seq_len images from a tracklet\n if num_imgs >= self.seq_len:\n num_imgs -= num_imgs % self.seq_len\n indices = np.arange(0, num_imgs, num_imgs / self.seq_len)\n else:\n # if num_imgs is smaller than seq_len, simply replicate the last image\n # until the seq_len requirement is satisfied\n indices = np.arange(0, num_imgs)\n num_pads = self.seq_len - num_imgs\n indices = np.concatenate(\n [\n indices,\n np.ones(num_pads).astype(np.int32) * (num_imgs-1)\n ]\n )\n assert len(indices) == self.seq_len\n\n elif self.sample_method == 'all':\n # Samples all images in a tracklet. batch_size must be set to 1\n indices = np.arange(num_imgs)\n\n else:\n raise ValueError(\n 'Unknown sample method: {}'.format(self.sample_method)\n )\n\n imgs = []\n for index in indices:\n img_path = img_paths[int(index)]\n img = read_image(img_path)\n if self.transform is not None:\n img = self.transform(img)\n img = img.unsqueeze(0) # img must be torch.Tensor\n imgs.append(img)\n imgs = torch.cat(imgs, dim=0)\n\n return imgs, pid, camid\n\n def show_summary(self):\n num_train_pids, num_train_cams = self.parse_data(self.train)\n num_query_pids, num_query_cams = self.parse_data(self.query)\n num_gallery_pids, num_gallery_cams = self.parse_data(self.gallery)\n\n print('=> Loaded {}'.format(self.__class__.__name__))\n print(' -------------------------------------------')\n print(' subset | # ids | # tracklets | # cameras')\n print(' -------------------------------------------')\n print(\n ' train | {:5d} | {:11d} | {:9d}'.format(\n num_train_pids, len(self.train), num_train_cams\n )\n )\n print(\n ' query | {:5d} | {:11d} | {:9d}'.format(\n num_query_pids, len(self.query), num_query_cams\n )\n )\n print(\n ' gallery | {:5d} | {:11d} | {:9d}'.format(\n num_gallery_pids, len(self.gallery), num_gallery_cams\n )\n )\n print(' -------------------------------------------')\n" }, { "path": "torchreid/data/datasets/image/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom .grid import GRID\nfrom .prid import PRID\nfrom .ilids import iLIDS\nfrom .viper import VIPeR\nfrom .cuhk01 import CUHK01\nfrom .cuhk02 import CUHK02\nfrom .cuhk03 import CUHK03\nfrom .msmt17 import MSMT17\nfrom .sensereid import SenseReID\nfrom .market1501 import Market1501\nfrom .dukemtmcreid import DukeMTMCreID\nfrom .occluded_dukemtmc import OccludedDuke\nfrom .occluded_reid import OccludedReID\nfrom .partial_reid import Partial_REID\nfrom .partial_ilids import Partial_iLIDS\nfrom .p_ETHZ import P_ETHZ\nfrom .p_dukemtmc_reid import PDukemtmcReid\n" }, { "path": "torchreid/data/datasets/image/cuhk01.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport numpy as np\nimport os.path as osp\nimport zipfile\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import ImageDataset\n\n\nclass CUHK01(ImageDataset):\n \"\"\"CUHK01.\n\n Reference:\n Li et al. Human Reidentification with Transferred Metric Learning. ACCV 2012.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 971.\n - images: 3884.\n - cameras: 4.\n \"\"\"\n dataset_dir = 'cuhk01'\n dataset_url = None\n eval_metric = 'default'\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.zip_path = osp.join(self.dataset_dir, 'CUHK01.zip')\n self.campus_dir = osp.join(self.dataset_dir, 'campus')\n self.split_path = osp.join(self.dataset_dir, 'splits.json')\n\n self.extract_file()\n\n required_files = [self.dataset_dir, self.campus_dir]\n self.check_before_run(required_files)\n\n self.prepare_split()\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, but expected between 0 and {}'\n .format(split_id,\n len(splits) - 1)\n )\n split = splits[split_id]\n\n train = split['train']\n query = split['query']\n gallery = split['gallery']\n\n train = [tuple(item) for item in train]\n query = [tuple(item) for item in query]\n gallery = [tuple(item) for item in gallery]\n\n super(CUHK01, self).__init__(train, query, gallery, **kwargs)\n\n def extract_file(self):\n if not osp.exists(self.campus_dir):\n print('Extracting files')\n zip_ref = zipfile.ZipFile(self.zip_path, 'r')\n zip_ref.extractall(self.dataset_dir)\n zip_ref.close()\n\n def prepare_split(self):\n \"\"\"\n Image name format: 0001001.png, where first four digits represent identity\n and last four digits represent cameras. Camera 1&2 are considered the same\n view and camera 3&4 are considered the same view.\n \"\"\"\n if not osp.exists(self.split_path):\n print('Creating 10 random splits of train ids and test ids')\n img_paths = sorted(glob.glob(osp.join(self.campus_dir, '*.png')))\n img_list = []\n pid_container = set()\n for img_path in img_paths:\n img_name = osp.basename(img_path)\n pid = int(img_name[:4]) - 1\n camid = (int(img_name[4:7]) - 1) // 2 # result is either 0 or 1\n img_list.append((img_path, pid, camid))\n pid_container.add(pid)\n\n num_pids = len(pid_container)\n num_train_pids = num_pids // 2\n\n splits = []\n for _ in range(10):\n order = np.arange(num_pids)\n np.random.shuffle(order)\n train_idxs = order[:num_train_pids]\n train_idxs = np.sort(train_idxs)\n idx2label = {\n idx: label\n for label, idx in enumerate(train_idxs)\n }\n\n train, test_a, test_b = [], [], []\n for img_path, pid, camid in img_list:\n if pid in train_idxs:\n train.append((img_path, idx2label[pid], camid))\n else:\n if camid == 0:\n test_a.append((img_path, pid, camid))\n else:\n test_b.append((img_path, pid, camid))\n\n # use cameraA as query and cameraB as gallery\n split = {\n 'train': train,\n 'query': test_a,\n 'gallery': test_b,\n 'num_train_pids': num_train_pids,\n 'num_query_pids': num_pids - num_train_pids,\n 'num_gallery_pids': num_pids - num_train_pids\n }\n splits.append(split)\n\n # use cameraB as query and cameraA as gallery\n split = {\n 'train': train,\n 'query': test_b,\n 'gallery': test_a,\n 'num_train_pids': num_train_pids,\n 'num_query_pids': num_pids - num_train_pids,\n 'num_gallery_pids': num_pids - num_train_pids\n }\n splits.append(split)\n\n print('Totally {} splits are created'.format(len(splits)))\n write_json(splits, self.split_path)\n print('Split file saved to {}'.format(self.split_path))\n" }, { "path": "torchreid/data/datasets/image/cuhk02.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport os.path as osp\n\nfrom ..dataset import ImageDataset\n\n\nclass CUHK02(ImageDataset):\n \"\"\"CUHK02.\n\n Reference:\n Li and Wang. Locally Aligned Feature Transforms across Views. CVPR 2013.\n\n URL: ``_\n \n Dataset statistics:\n - 5 camera view pairs each with two cameras\n - 971, 306, 107, 193 and 239 identities from P1 - P5\n - totally 1,816 identities\n - image format is png\n\n Protocol: Use P1 - P4 for training and P5 for evaluation.\n \"\"\"\n dataset_dir = 'cuhk02'\n cam_pairs = ['P1', 'P2', 'P3', 'P4', 'P5']\n test_cam_pair = 'P5'\n eval_metric = 'default'\n\n def __init__(self, root='', **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir, 'Dataset')\n\n required_files = [self.dataset_dir]\n self.check_before_run(required_files)\n\n train, query, gallery = self.get_data_list()\n\n super(CUHK02, self).__init__(train, query, gallery, **kwargs)\n\n def get_data_list(self):\n num_train_pids, camid = 0, 0\n train, query, gallery = [], [], []\n\n for cam_pair in self.cam_pairs:\n cam_pair_dir = osp.join(self.dataset_dir, cam_pair)\n\n cam1_dir = osp.join(cam_pair_dir, 'cam1')\n cam2_dir = osp.join(cam_pair_dir, 'cam2')\n\n impaths1 = glob.glob(osp.join(cam1_dir, '*.png'))\n impaths2 = glob.glob(osp.join(cam2_dir, '*.png'))\n\n if cam_pair == self.test_cam_pair:\n # add images to query\n for impath in impaths1:\n pid = osp.basename(impath).split('_')[0]\n pid = int(pid)\n query.append((impath, pid, camid))\n camid += 1\n\n # add images to gallery\n for impath in impaths2:\n pid = osp.basename(impath).split('_')[0]\n pid = int(pid)\n gallery.append((impath, pid, camid))\n camid += 1\n\n else:\n pids1 = [\n osp.basename(impath).split('_')[0] for impath in impaths1\n ]\n pids2 = [\n osp.basename(impath).split('_')[0] for impath in impaths2\n ]\n pids = set(pids1 + pids2)\n pid2label = {\n pid: label + num_train_pids\n for label, pid in enumerate(pids)\n }\n\n # add images to train from cam1\n for impath in impaths1:\n pid = osp.basename(impath).split('_')[0]\n pid = pid2label[pid]\n train.append((impath, pid, camid))\n camid += 1\n\n # add images to train from cam2\n for impath in impaths2:\n pid = osp.basename(impath).split('_')[0]\n pid = pid2label[pid]\n train.append((impath, pid, camid))\n camid += 1\n num_train_pids += len(pids)\n\n return train, query, gallery\n" }, { "path": "torchreid/data/datasets/image/cuhk03.py", "content": "from __future__ import division, print_function, absolute_import\nimport os.path as osp\n\nfrom torchreid.utils import read_json, write_json, mkdir_if_missing\n\nfrom ..dataset import ImageDataset\n\n\nclass CUHK03(ImageDataset):\n \"\"\"CUHK03.\n\n Reference:\n Li et al. DeepReID: Deep Filter Pairing Neural Network for Person Re-identification. CVPR 2014.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 1360.\n - images: 13164.\n - cameras: 6.\n - splits: 20 (classic).\n \"\"\"\n dataset_dir = 'cuhk03'\n dataset_url = None\n eval_metric = 'default'\n\n def __init__(\n self,\n root='',\n split_id=0,\n cuhk03_labeled=False,\n cuhk03_classic_split=False,\n **kwargs\n ):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.data_dir = osp.join(self.dataset_dir, 'cuhk03_release')\n self.raw_mat_path = osp.join(self.data_dir, 'cuhk-03.mat')\n\n self.imgs_detected_dir = osp.join(self.dataset_dir, 'images_detected')\n self.imgs_labeled_dir = osp.join(self.dataset_dir, 'images_labeled')\n\n self.split_classic_det_json_path = osp.join(\n self.dataset_dir, 'splits_classic_detected.json'\n )\n self.split_classic_lab_json_path = osp.join(\n self.dataset_dir, 'splits_classic_labeled.json'\n )\n\n self.split_new_det_json_path = osp.join(\n self.dataset_dir, 'splits_new_detected.json'\n )\n self.split_new_lab_json_path = osp.join(\n self.dataset_dir, 'splits_new_labeled.json'\n )\n\n self.split_new_det_mat_path = osp.join(\n self.dataset_dir, 'cuhk03_new_protocol_config_detected.mat'\n )\n self.split_new_lab_mat_path = osp.join(\n self.dataset_dir, 'cuhk03_new_protocol_config_labeled.mat'\n )\n\n required_files = [\n self.dataset_dir, self.data_dir, self.raw_mat_path,\n self.split_new_det_mat_path, self.split_new_lab_mat_path\n ]\n self.check_before_run(required_files)\n\n self.preprocess_split()\n\n if cuhk03_labeled:\n split_path = self.split_classic_lab_json_path if cuhk03_classic_split else self.split_new_lab_json_path\n else:\n split_path = self.split_classic_det_json_path if cuhk03_classic_split else self.split_new_det_json_path\n\n splits = read_json(split_path)\n assert split_id < len(\n splits\n ), 'Condition split_id ({}) < len(splits) ({}) is false'.format(\n split_id, len(splits)\n )\n split = splits[split_id]\n\n train = split['train']\n query = split['query']\n gallery = split['gallery']\n\n super(CUHK03, self).__init__(train, query, gallery, **kwargs)\n\n def preprocess_split(self):\n # This function is a bit complex and ugly, what it does is\n # 1. extract data from cuhk-03.mat and save as png images\n # 2. create 20 classic splits (Li et al. CVPR'14)\n # 3. create new split (Zhong et al. CVPR'17)\n if osp.exists(self.imgs_labeled_dir) \\\n and osp.exists(self.imgs_detected_dir) \\\n and osp.exists(self.split_classic_det_json_path) \\\n and osp.exists(self.split_classic_lab_json_path) \\\n and osp.exists(self.split_new_det_json_path) \\\n and osp.exists(self.split_new_lab_json_path):\n return\n\n import h5py\n from scipy.misc import imsave\n from scipy.io import loadmat\n\n mkdir_if_missing(self.imgs_detected_dir)\n mkdir_if_missing(self.imgs_labeled_dir)\n\n print(\n 'Extract image data from \"{}\" and save as png'.format(\n self.raw_mat_path\n )\n )\n mat = h5py.File(self.raw_mat_path, 'r')\n\n def _deref(ref):\n return mat[ref][:].T\n\n def _process_images(img_refs, campid, pid, save_dir):\n img_paths = [] # Note: some persons only have images for one view\n for imgid, img_ref in enumerate(img_refs):\n img = _deref(img_ref)\n if img.size == 0 or img.ndim < 3:\n continue # skip empty cell\n # images are saved with the following format, index-1 (ensure uniqueness)\n # campid: index of camera pair (1-5)\n # pid: index of person in 'campid'-th camera pair\n # viewid: index of view, {1, 2}\n # imgid: index of image, (1-10)\n viewid = 1 if imgid < 5 else 2\n img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format(\n campid + 1, pid + 1, viewid, imgid + 1\n )\n img_path = osp.join(save_dir, img_name)\n if not osp.isfile(img_path):\n imsave(img_path, img)\n img_paths.append(img_path)\n return img_paths\n\n def _extract_img(image_type):\n print('Processing {} images ...'.format(image_type))\n meta_data = []\n imgs_dir = self.imgs_detected_dir if image_type == 'detected' else self.imgs_labeled_dir\n for campid, camp_ref in enumerate(mat[image_type][0]):\n camp = _deref(camp_ref)\n num_pids = camp.shape[0]\n for pid in range(num_pids):\n img_paths = _process_images(\n camp[pid, :], campid, pid, imgs_dir\n )\n assert len(img_paths) > 0, \\\n 'campid{}-pid{} has no images'.format(campid, pid)\n meta_data.append((campid + 1, pid + 1, img_paths))\n print(\n '- done camera pair {} with {} identities'.format(\n campid + 1, num_pids\n )\n )\n return meta_data\n\n meta_detected = _extract_img('detected')\n meta_labeled = _extract_img('labeled')\n\n def _extract_classic_split(meta_data, test_split):\n train, test = [], []\n num_train_pids, num_test_pids = 0, 0\n num_train_imgs, num_test_imgs = 0, 0\n for i, (campid, pid, img_paths) in enumerate(meta_data):\n\n if [campid, pid] in test_split:\n for img_path in img_paths:\n camid = int(\n osp.basename(img_path).split('_')[2]\n ) - 1 # make it 0-based\n test.append((img_path, num_test_pids, camid))\n num_test_pids += 1\n num_test_imgs += len(img_paths)\n else:\n for img_path in img_paths:\n camid = int(\n osp.basename(img_path).split('_')[2]\n ) - 1 # make it 0-based\n train.append((img_path, num_train_pids, camid))\n num_train_pids += 1\n num_train_imgs += len(img_paths)\n return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs\n\n print('Creating classic splits (# = 20) ...')\n splits_classic_det, splits_classic_lab = [], []\n for split_ref in mat['testsets'][0]:\n test_split = _deref(split_ref).tolist()\n\n # create split for detected images\n train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \\\n _extract_classic_split(meta_detected, test_split)\n splits_classic_det.append(\n {\n 'train': train,\n 'query': test,\n 'gallery': test,\n 'num_train_pids': num_train_pids,\n 'num_train_imgs': num_train_imgs,\n 'num_query_pids': num_test_pids,\n 'num_query_imgs': num_test_imgs,\n 'num_gallery_pids': num_test_pids,\n 'num_gallery_imgs': num_test_imgs\n }\n )\n\n # create split for labeled images\n train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \\\n _extract_classic_split(meta_labeled, test_split)\n splits_classic_lab.append(\n {\n 'train': train,\n 'query': test,\n 'gallery': test,\n 'num_train_pids': num_train_pids,\n 'num_train_imgs': num_train_imgs,\n 'num_query_pids': num_test_pids,\n 'num_query_imgs': num_test_imgs,\n 'num_gallery_pids': num_test_pids,\n 'num_gallery_imgs': num_test_imgs\n }\n )\n\n write_json(splits_classic_det, self.split_classic_det_json_path)\n write_json(splits_classic_lab, self.split_classic_lab_json_path)\n\n def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel):\n tmp_set = []\n unique_pids = set()\n for idx in idxs:\n img_name = filelist[idx][0]\n camid = int(img_name.split('_')[2]) - 1 # make it 0-based\n pid = pids[idx]\n if relabel:\n pid = pid2label[pid]\n img_path = osp.join(img_dir, img_name)\n tmp_set.append((img_path, int(pid), camid))\n unique_pids.add(pid)\n return tmp_set, len(unique_pids), len(idxs)\n\n def _extract_new_split(split_dict, img_dir):\n train_idxs = split_dict['train_idx'].flatten() - 1 # index-0\n pids = split_dict['labels'].flatten()\n train_pids = set(pids[train_idxs])\n pid2label = {pid: label for label, pid in enumerate(train_pids)}\n query_idxs = split_dict['query_idx'].flatten() - 1\n gallery_idxs = split_dict['gallery_idx'].flatten() - 1\n filelist = split_dict['filelist'].flatten()\n train_info = _extract_set(\n filelist, pids, pid2label, train_idxs, img_dir, relabel=True\n )\n query_info = _extract_set(\n filelist, pids, pid2label, query_idxs, img_dir, relabel=False\n )\n gallery_info = _extract_set(\n filelist,\n pids,\n pid2label,\n gallery_idxs,\n img_dir,\n relabel=False\n )\n return train_info, query_info, gallery_info\n\n print('Creating new split for detected images (767/700) ...')\n train_info, query_info, gallery_info = _extract_new_split(\n loadmat(self.split_new_det_mat_path), self.imgs_detected_dir\n )\n split = [\n {\n 'train': train_info[0],\n 'query': query_info[0],\n 'gallery': gallery_info[0],\n 'num_train_pids': train_info[1],\n 'num_train_imgs': train_info[2],\n 'num_query_pids': query_info[1],\n 'num_query_imgs': query_info[2],\n 'num_gallery_pids': gallery_info[1],\n 'num_gallery_imgs': gallery_info[2]\n }\n ]\n write_json(split, self.split_new_det_json_path)\n\n print('Creating new split for labeled images (767/700) ...')\n train_info, query_info, gallery_info = _extract_new_split(\n loadmat(self.split_new_lab_mat_path), self.imgs_labeled_dir\n )\n split = [\n {\n 'train': train_info[0],\n 'query': query_info[0],\n 'gallery': gallery_info[0],\n 'num_train_pids': train_info[1],\n 'num_train_imgs': train_info[2],\n 'num_query_pids': query_info[1],\n 'num_query_imgs': query_info[2],\n 'num_gallery_pids': gallery_info[1],\n 'num_gallery_imgs': gallery_info[2]\n }\n ]\n write_json(split, self.split_new_lab_json_path)\n" }, { "path": "torchreid/data/datasets/image/dukemtmcreid.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport re\nimport glob\nimport os.path as osp\n\nfrom ..dataset import ImageDataset\n\n\nclass DukeMTMCreID(ImageDataset):\n \"\"\"DukeMTMC-reID.\n\n Reference:\n - Ristani et al. Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCVW 2016.\n - Zheng et al. Unlabeled Samples Generated by GAN Improve the Person Re-identification Baseline in vitro. ICCV 2017.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 1404 (train + query).\n - images:16522 (train) + 2228 (query) + 17661 (gallery).\n - cameras: 8.\n \"\"\"\n dataset_dir = 'DukeMTMC-reID'\n dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-reID.zip'\n masks_base_dir = 'masks'\n\n masks_dirs = {\n # dir_name: (parts_num, masks_stack_size, contains_background_mask)\n 'pifpaf': (36, False, '.jpg.confidence_fields.npy'),\n 'pifpaf_maskrcnn_filtering': (36, False, '.npy'),\n }\n\n @staticmethod\n def get_masks_config(masks_dir):\n if masks_dir not in DukeMTMCreID.masks_dirs:\n return None\n else:\n return DukeMTMCreID.masks_dirs[masks_dir]\n\n def __init__(self, root='', masks_dir=None, **kwargs):\n self.masks_dir = masks_dir\n if self.masks_dir in self.masks_dirs:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = self.masks_dirs[self.masks_dir]\n else:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = None, None, None\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')\n self.query_dir = osp.join(self.dataset_dir, 'query')\n self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')\n\n required_files = [\n self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir\n ]\n self.check_before_run(required_files)\n\n train = self.process_dir(self.train_dir, relabel=True)\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False)\n\n super(DukeMTMCreID, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, relabel=False):\n img_paths = glob.glob(osp.join(dir_path, '*.jpg'))\n pattern = re.compile(r'([-\\d]+)_c(\\d)')\n\n pid_container = set()\n for img_path in img_paths:\n pid, _ = map(int, pattern.search(img_path).groups())\n pid_container.add(pid)\n pid2label = {pid: label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n pid, camid = map(int, pattern.search(img_path).groups())\n assert 1 <= camid <= 8\n camid -= 1 # index starts from 0\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n\n return data\n" }, { "path": "torchreid/data/datasets/image/grid.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport os.path as osp\nfrom scipy.io import loadmat\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import ImageDataset\n\n\nclass GRID(ImageDataset):\n \"\"\"GRID.\n\n Reference:\n Loy et al. Multi-camera activity correlation analysis. CVPR 2009.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 250.\n - images: 1275.\n - cameras: 8.\n \"\"\"\n dataset_dir = 'grid'\n dataset_url = 'http://personal.ie.cuhk.edu.hk/~ccloy/files/datasets/underground_reid.zip'\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.probe_path = osp.join(\n self.dataset_dir, 'underground_reid', 'probe'\n )\n self.gallery_path = osp.join(\n self.dataset_dir, 'underground_reid', 'gallery'\n )\n self.split_mat_path = osp.join(\n self.dataset_dir, 'underground_reid', 'features_and_partitions.mat'\n )\n self.split_path = osp.join(self.dataset_dir, 'splits.json')\n\n required_files = [\n self.dataset_dir, self.probe_path, self.gallery_path,\n self.split_mat_path\n ]\n self.check_before_run(required_files)\n\n self.prepare_split()\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, '\n 'but expected between 0 and {}'.format(\n split_id,\n len(splits) - 1\n )\n )\n split = splits[split_id]\n\n train = split['train']\n query = split['query']\n gallery = split['gallery']\n\n train = [tuple(item) for item in train]\n query = [tuple(item) for item in query]\n gallery = [tuple(item) for item in gallery]\n\n super(GRID, self).__init__(train, query, gallery, **kwargs)\n\n def prepare_split(self):\n if not osp.exists(self.split_path):\n print('Creating 10 random splits')\n split_mat = loadmat(self.split_mat_path)\n trainIdxAll = split_mat['trainIdxAll'][0] # length = 10\n probe_img_paths = sorted(\n glob.glob(osp.join(self.probe_path, '*.jpeg'))\n )\n gallery_img_paths = sorted(\n glob.glob(osp.join(self.gallery_path, '*.jpeg'))\n )\n\n splits = []\n for split_idx in range(10):\n train_idxs = trainIdxAll[split_idx][0][0][2][0].tolist()\n assert len(train_idxs) == 125\n idx2label = {\n idx: label\n for label, idx in enumerate(train_idxs)\n }\n\n train, query, gallery = [], [], []\n\n # processing probe folder\n for img_path in probe_img_paths:\n img_name = osp.basename(img_path)\n img_idx = int(img_name.split('_')[0])\n camid = int(\n img_name.split('_')[1]\n ) - 1 # index starts from 0\n if img_idx in train_idxs:\n train.append((img_path, idx2label[img_idx], camid))\n else:\n query.append((img_path, img_idx, camid))\n\n # process gallery folder\n for img_path in gallery_img_paths:\n img_name = osp.basename(img_path)\n img_idx = int(img_name.split('_')[0])\n camid = int(\n img_name.split('_')[1]\n ) - 1 # index starts from 0\n if img_idx in train_idxs:\n train.append((img_path, idx2label[img_idx], camid))\n else:\n gallery.append((img_path, img_idx, camid))\n\n split = {\n 'train': train,\n 'query': query,\n 'gallery': gallery,\n 'num_train_pids': 125,\n 'num_query_pids': 125,\n 'num_gallery_pids': 900\n }\n splits.append(split)\n\n print('Totally {} splits are created'.format(len(splits)))\n write_json(splits, self.split_path)\n print('Split file saved to {}'.format(self.split_path))\n" }, { "path": "torchreid/data/datasets/image/ilids.py", "content": "from __future__ import division, print_function, absolute_import\nimport copy\nimport glob\nimport random\nimport os.path as osp\nfrom collections import defaultdict\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import ImageDataset\n\n\nclass iLIDS(ImageDataset):\n \"\"\"QMUL-iLIDS.\n\n Reference:\n Zheng et al. Associating Groups of People. BMVC 2009.\n \n Dataset statistics:\n - identities: 119.\n - images: 476.\n - cameras: 8 (not explicitly provided).\n \"\"\"\n dataset_dir = 'ilids'\n dataset_url = 'http://www.eecs.qmul.ac.uk/~jason/data/i-LIDS_Pedestrian.tgz'\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.data_dir = osp.join(self.dataset_dir, 'i-LIDS_Pedestrian/Persons')\n self.split_path = osp.join(self.dataset_dir, 'splits.json')\n\n required_files = [self.dataset_dir, self.data_dir]\n self.check_before_run(required_files)\n\n self.prepare_split()\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, but '\n 'expected between 0 and {}'.format(split_id,\n len(splits) - 1)\n )\n split = splits[split_id]\n\n train, query, gallery = self.process_split(split)\n\n super(iLIDS, self).__init__(train, query, gallery, **kwargs)\n\n def prepare_split(self):\n if not osp.exists(self.split_path):\n print('Creating splits ...')\n\n paths = glob.glob(osp.join(self.data_dir, '*.jpg'))\n img_names = [osp.basename(path) for path in paths]\n num_imgs = len(img_names)\n assert num_imgs == 476, 'There should be 476 images, but ' \\\n 'got {}, please check the data'.format(num_imgs)\n\n # store image names\n # image naming format:\n # the first four digits denote the person ID\n # the last four digits denote the sequence index\n pid_dict = defaultdict(list)\n for img_name in img_names:\n pid = int(img_name[:4])\n pid_dict[pid].append(img_name)\n pids = list(pid_dict.keys())\n num_pids = len(pids)\n assert num_pids == 119, 'There should be 119 identities, ' \\\n 'but got {}, please check the data'.format(num_pids)\n\n num_train_pids = int(num_pids * 0.5)\n\n splits = []\n for _ in range(10):\n # randomly choose num_train_pids train IDs and the rest for test IDs\n pids_copy = copy.deepcopy(pids)\n random.shuffle(pids_copy)\n train_pids = pids_copy[:num_train_pids]\n test_pids = pids_copy[num_train_pids:]\n\n train = []\n query = []\n gallery = []\n\n # for train IDs, all images are used in the train set.\n for pid in train_pids:\n img_names = pid_dict[pid]\n train.extend(img_names)\n\n # for each test ID, randomly choose two images, one for\n # query and the other one for gallery.\n for pid in test_pids:\n img_names = pid_dict[pid]\n samples = random.sample(img_names, 2)\n query.append(samples[0])\n gallery.append(samples[1])\n\n split = {'train': train, 'query': query, 'gallery': gallery}\n splits.append(split)\n\n print('Totally {} splits are created'.format(len(splits)))\n write_json(splits, self.split_path)\n print('Split file is saved to {}'.format(self.split_path))\n\n def get_pid2label(self, img_names):\n pid_container = set()\n for img_name in img_names:\n pid = int(img_name[:4])\n pid_container.add(pid)\n pid2label = {pid: label for label, pid in enumerate(pid_container)}\n return pid2label\n\n def parse_img_names(self, img_names, pid2label=None):\n data = []\n\n for img_name in img_names:\n pid = int(img_name[:4])\n if pid2label is not None:\n pid = pid2label[pid]\n camid = int(img_name[4:7]) - 1 # 0-based\n img_path = osp.join(self.data_dir, img_name)\n data.append({'img_path': img_path, 'pid': pid, 'camid': camid})\n\n return data\n\n def process_split(self, split):\n train_pid2label = self.get_pid2label(split['train'])\n train = self.parse_img_names(split['train'], train_pid2label)\n query = self.parse_img_names(split['query'])\n gallery = self.parse_img_names(split['gallery'])\n return train, query, gallery\n" }, { "path": "torchreid/data/datasets/image/market1501.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport re\nimport glob\nimport os.path as osp\nimport warnings\n\nfrom ..dataset import ImageDataset\n\n\nclass Market1501(ImageDataset):\n \"\"\"Market1501.\n\n Reference:\n Zheng et al. Scalable Person Re-identification: A Benchmark. ICCV 2015.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 1501 (+1 for background).\n - images: 12936 (train) + 3368 (query) + 15913 (gallery).\n \"\"\"\n _junk_pids = [0, -1]\n dataset_dir = 'Market-1501-v15.09.15'\n masks_base_dir = 'masks'\n dataset_url = 'http://188.138.127.15:81/Datasets/Market-1501-v15.09.15.zip'\n\n masks_dirs = {\n # dir_name: (parts_num, masks_stack_size, contains_background_mask)\n 'pifpaf': (36, False, '.jpg.confidence_fields.npy'),\n 'pifpaf_maskrcnn_filtering': (36, False, '.npy'),\n }\n\n @staticmethod\n def get_masks_config(masks_dir):\n if masks_dir not in Market1501.masks_dirs:\n return None\n else:\n return Market1501.masks_dirs[masks_dir]\n\n def __init__(self, root='', market1501_500k=False, masks_dir=None, **kwargs):\n self.masks_dir = masks_dir\n if self.masks_dir in self.masks_dirs:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = self.masks_dirs[self.masks_dir]\n else:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = None, None, None\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n self.masks_dir = masks_dir\n\n # allow alternative directory structure\n if not osp.isdir(self.dataset_dir):\n warnings.warn(\n 'The current data structure is deprecated. Please '\n 'put data folders such as \"bounding_box_train\" under '\n '\"Market-1501-v15.09.15\".'\n )\n\n self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')\n self.query_dir = osp.join(self.dataset_dir, 'query')\n self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')\n self.extra_gallery_dir = osp.join(self.dataset_dir, 'images')\n self.market1501_500k = market1501_500k\n\n required_files = [\n self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir\n ]\n if self.market1501_500k:\n required_files.append(self.extra_gallery_dir)\n self.check_before_run(required_files)\n train = self.process_dir(self.train_dir, relabel=True)\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False)\n if self.market1501_500k:\n gallery += self.process_dir(self.extra_gallery_dir, relabel=False)\n super(Market1501, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, relabel=False):\n img_paths = glob.glob(osp.join(dir_path, '*.jpg'))\n pattern = re.compile(r'([-\\d]+)_c(\\d)')\n\n pid_container = set()\n for img_path in img_paths:\n pid, _ = map(int, pattern.search(img_path).groups())\n if pid == -1:\n continue # junk images are just ignored\n pid_container.add(pid)\n pid2label = {pid: label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n pid, camid = map(int, pattern.search(img_path).groups())\n if pid == -1:\n continue # junk images are just ignored\n assert 0 <= pid <= 1501 # pid == 0 means background\n assert 1 <= camid <= 6\n camid -= 1 # index starts from 0\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n return data\n" }, { "path": "torchreid/data/datasets/image/msmt17.py", "content": "from __future__ import division, print_function, absolute_import\nimport os.path as osp\n\nfrom ..dataset import ImageDataset\n\n# Log\n# 22.01.2019\n# - add v2\n# - v1 and v2 differ in dir names\n# - note that faces in v2 are blurred\nTRAIN_DIR_KEY = 'train_dir'\nTEST_DIR_KEY = 'test_dir'\nVERSION_DICT = {\n 'MSMT17_V1': {\n TRAIN_DIR_KEY: 'train',\n TEST_DIR_KEY: 'test',\n },\n 'MSMT17_V2': {\n TRAIN_DIR_KEY: 'mask_train_v2',\n TEST_DIR_KEY: 'mask_test_v2',\n }\n}\n\n\nclass MSMT17(ImageDataset):\n \"\"\"MSMT17.\n\n Reference:\n Wei et al. Person Transfer GAN to Bridge Domain Gap for Person Re-Identification. CVPR 2018.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 4101.\n - images: 32621 (train) + 11659 (query) + 82161 (gallery).\n - cameras: 15.\n \"\"\"\n dataset_dir = 'msmt17'\n dataset_url = None\n\n masks_dirs = {\n # dir_name: (masks_stack_size, contains_background_mask)\n }\n\n @staticmethod\n def get_masks_config(masks_dir):\n if masks_dir not in MSMT17.masks_dirs:\n return None\n else:\n return MSMT17.masks_dirs[masks_dir]\n\n def __init__(self, root='', masks_dir=None, **kwargs):\n self.masks_dir = masks_dir\n if self.masks_dir in self.masks_dirs:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = self.masks_dirs[self.masks_dir]\n else:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = None, None, None\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n has_main_dir = False\n for main_dir in VERSION_DICT:\n if osp.exists(osp.join(self.dataset_dir, main_dir)):\n train_dir = VERSION_DICT[main_dir][TRAIN_DIR_KEY]\n test_dir = VERSION_DICT[main_dir][TEST_DIR_KEY]\n has_main_dir = True\n break\n assert has_main_dir, 'Dataset folder not found'\n\n self.train_dir = osp.join(self.dataset_dir, main_dir, train_dir)\n self.test_dir = osp.join(self.dataset_dir, main_dir, test_dir)\n self.list_train_path = osp.join(\n self.dataset_dir, main_dir, 'list_train.txt'\n )\n self.list_val_path = osp.join(\n self.dataset_dir, main_dir, 'list_val.txt'\n )\n self.list_query_path = osp.join(\n self.dataset_dir, main_dir, 'list_query.txt'\n )\n self.list_gallery_path = osp.join(\n self.dataset_dir, main_dir, 'list_gallery.txt'\n )\n\n required_files = [self.dataset_dir, self.train_dir, self.test_dir]\n self.check_before_run(required_files)\n\n train = self.process_dir(self.train_dir, self.list_train_path)\n val = self.process_dir(self.train_dir, self.list_val_path)\n query = self.process_dir(self.test_dir, self.list_query_path)\n gallery = self.process_dir(self.test_dir, self.list_gallery_path)\n\n # Note: to fairly compare with published methods on the conventional ReID setting,\n # do not add val images to the training set.\n if 'combineall' in kwargs and kwargs['combineall']:\n train += val\n\n super(MSMT17, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, list_path):\n with open(list_path, 'r') as txt:\n lines = txt.readlines()\n\n data = []\n\n for img_idx, img_info in enumerate(lines):\n img_path, pid = img_info.split(' ')\n pid = int(pid) # no need to relabel\n camid = int(img_path.split('_')[2]) - 1 # index starts from 0\n img_path = osp.join(dir_path, img_path)\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n\n return data\n" }, { "path": "torchreid/data/datasets/image/occluded_dukemtmc.py", "content": "from __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport os.path as osp\nimport glob\nimport re\nfrom ..dataset import ImageDataset\n\n# Sources :\n# https://github.com/hh23333/PVPM\n# https://github.com/lightas/Occluded-DukeMTMC-Dataset\n# Miao, J., Wu, Y., Liu, P., DIng, Y., & Yang, Y. (2019). \"Pose-guided feature alignment for occluded person re-identification\". ICCV 2019\n\nclass OccludedDuke(ImageDataset):\n dataset_dir = 'Occluded_Duke'\n masks_base_dir = 'masks'\n\n masks_dirs = {\n # dir_name: (parts_num, masks_stack_size, contains_background_mask)\n 'pifpaf': (36, False, '.jpg.confidence_fields.npy'),\n 'pifpaf_maskrcnn_filtering': (36, False, '.jpg.confidence_fields.npy'),\n 'isp_6_parts': (5, True, '.jpg.confidence_fields.npy', [\"p{}\".format(p) for p in range(1, 5+1)])\n }\n\n @staticmethod\n def get_masks_config(masks_dir):\n if masks_dir not in OccludedDuke.masks_dirs:\n return None\n else:\n return OccludedDuke.masks_dirs[masks_dir]\n\n def __init__(self, root='', masks_dir=None, **kwargs):\n self.masks_dir = masks_dir\n if self.masks_dir in self.masks_dirs:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = self.masks_dirs[self.masks_dir]\n else:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = None, None, None\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')\n self.query_dir = osp.join(self.dataset_dir, 'query')\n self.gallery_dir = osp.join(self.dataset_dir, 'bounding_box_test')\n\n required_files = [\n self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir\n ]\n self.check_before_run(required_files)\n\n train = self.process_dir(self.train_dir, relabel=True)\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False)\n\n super(OccludedDuke, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, relabel=False):\n img_paths = glob.glob(osp.join(dir_path, '*.jpg'))\n pattern = re.compile(r'([-\\d]+)_c(\\d)')\n\n pid_container = set()\n for img_path in img_paths:\n pid, _ = map(int, pattern.search(img_path).groups())\n pid_container.add(pid)\n pid2label = {pid: label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n pid, camid = map(int, pattern.search(img_path).groups())\n assert 1 <= camid <= 8\n camid -= 1 # index starts from 0\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n\n return data\n" }, { "path": "torchreid/data/datasets/image/occluded_reid.py", "content": "from __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport os\nimport os.path as osp\nimport glob\nimport warnings\nfrom ..dataset import ImageDataset\n\n# Sources :\n# https://github.com/hh23333/PVPM\n# Zhuo, J., Chen, Z., Lai, J., & Wang, G. (2018). Occluded Person Re-identification.\n\n\nclass OccludedReID(ImageDataset):\n dataset_dir = 'Occluded_REID'\n masks_base_dir = 'masks'\n\n masks_dirs = {\n # dir_name: (parts_num, masks_stack_size, contains_background_mask)\n 'pifpaf': (36, False, '.tif.confidence_fields.npy'),\n 'pifpaf_maskrcnn_filtering': (36, False, '.npy'),\n }\n\n @staticmethod\n def get_masks_config(masks_dir):\n if masks_dir not in OccludedReID.masks_dirs:\n return None\n else:\n return OccludedReID.masks_dirs[masks_dir]\n\n def infer_masks_path(self, img_path):\n masks_path = os.path.join(self.dataset_dir, self.masks_base_dir, self.masks_dir, os.path.basename(os.path.dirname(os.path.dirname(img_path))), os.path.splitext(os.path.basename(img_path))[0] + self.masks_suffix)\n return masks_path\n\n def __init__(self, root='', masks_dir=None, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.masks_dir = masks_dir\n if self.masks_dir in self.masks_dirs:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = self.masks_dirs[self.masks_dir]\n else:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = None, None, None\n # allow alternative directory structure\n if not osp.isdir(self.dataset_dir):\n warnings.warn(\n 'The current data structure is deprecated. Please '\n 'put data folders such as \"bounding_box_train\" under '\n '\"Market-1501-v15.09.15\".'\n )\n self.query_dir = osp.join(self.dataset_dir, 'occluded_body_images')\n self.gallery_dir = osp.join(self.dataset_dir, 'whole_body_images')\n\n train = []\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False, is_query=False)\n super(OccludedReID, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, relabel=False, is_query=True):\n img_paths = glob.glob(osp.join(dir_path, '*', '*.tif'))\n if is_query:\n camid = 0\n else:\n camid = 1\n pid_container = set()\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n pid_container.add(pid)\n pid2label = {pid:label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n return data\n" }, { "path": "torchreid/data/datasets/image/p_ETHZ.py", "content": "from __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport os.path as osp\nimport glob\nimport warnings\n\nfrom ..dataset import ImageDataset\n\n# Sources :\n# https://github.com/hh23333/PVPM\n# A. Ess, B. Leibe, K. Schindler, and L. Van Gool, \"A mobile vision system for robust multi-person tracking\" in CVPR, 2008\n\n\nclass P_ETHZ(ImageDataset):\n dataset_dir = 'P_ETHZ'\n\n def __init__(self, root='', **kwargs):\n self.root=osp.abspath(osp.expanduser(root))\n # self.dataset_dir = self.root\n data_dir = osp.join(self.root, self.dataset_dir)\n if osp.isdir(data_dir):\n self.data_dir = data_dir\n else:\n warnings.warn('The current data structure is deprecated.')\n self.query_dir=osp.join(self.data_dir, 'occluded_body_images')\n self.gallery_dir=osp.join(self.data_dir, 'whole_body_images')\n\n train = []\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False, is_query=False)\n super(P_ETHZ, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, relabel=False, is_query=True):\n img_paths = glob.glob(osp.join(dir_path, '*', '*.png'))\n if is_query:\n camid = 0\n else:\n camid = 1\n pid_container = set()\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n pid_container.add(pid)\n pid2label = {pid:label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n if relabel:\n pid = pid2label[pid]\n data.append({'img_path': img_path, 'pid': pid, 'camid': camid})\n return data\n\n" }, { "path": "torchreid/data/datasets/image/p_dukemtmc_reid.py", "content": "from __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport os\nimport os.path as osp\nimport glob\nimport warnings\n\nfrom ..dataset import ImageDataset\n\n# Sources :\n# https://github.com/hh23333/PVPM\n# Z.D. Zheng, L. Zheng, and Y. Yang, \"Unlabeled samples generated by gan improve the person re-identification baseline in vitro\" in ICCV, 2017.\n\nclass PDukemtmcReid(ImageDataset):\n dataset_dir = 'P-DukeMTMC-reID'\n masks_base_dir = 'masks'\n\n masks_dirs = {\n # dir_name: (parts_num, masks_stack_size, contains_background_mask)\n 'pifpaf': (36, False, '.jpg.confidence_fields.npy'),\n 'pifpaf_maskrcnn_filtering': (36, False, '.npy'),\n }\n\n @staticmethod\n def get_masks_config(masks_dir):\n if masks_dir not in PDukemtmcReid.masks_dirs:\n return None\n else:\n return PDukemtmcReid.masks_dirs[masks_dir]\n\n def infer_masks_path(self, img_path):\n masks_path = os.path.join(self.dataset_dir,\n self.masks_base_dir,\n self.masks_dir,\n os.path.basename(os.path.dirname(os.path.dirname(os.path.dirname(img_path)))),\n os.path.basename(os.path.dirname(os.path.dirname(img_path))),\n # FIXME ugly\n os.path.splitext(os.path.basename(img_path))[0] + self.masks_suffix)\n return masks_path\n\n def __init__(self, root='', masks_dir=None, **kwargs):\n self.masks_dir = masks_dir\n if self.masks_dir in self.masks_dirs:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = self.masks_dirs[self.masks_dir]\n else:\n self.masks_parts_numbers, self.has_background, self.masks_suffix = None, None, None\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n if not osp.isdir(self.dataset_dir):\n warnings.warn(\n 'The current data structure is deprecated. Please '\n 'put data folders such as \"bounding_box_train\" under '\n '\"Market-1501-v15.09.15\".'\n )\n self.train_dir=osp.join(self.dataset_dir, 'train')\n self.query_dir=osp.join(self.dataset_dir, 'test', 'occluded_body_images')\n self.gallery_dir=osp.join(self.dataset_dir, 'test', 'whole_body_images')\n\n train = self.process_train_dir(self.train_dir, relabel=True)\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False, is_query=False)\n super(PDukemtmcReid, self).__init__(train, query, gallery, **kwargs)\n\n def process_train_dir(self, dir_path, relabel=True):\n img_paths = glob.glob(osp.join(dir_path,'whole_body_images', '*', '*.jpg'))\n camid=1\n pid_container = set()\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n pid_container.add(pid)\n pid2label = {pid:label for label, pid in enumerate(pid_container)}\n data = []\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n img_paths = glob.glob(osp.join(dir_path,'occluded_body_images','*','*.jpg'))\n camid=0\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n return data\n\n def process_dir(self, dir_path, relabel=False, is_query=True):\n img_paths = glob.glob(osp.join(dir_path, '*', '*.jpg'))\n if is_query:\n camid = 0\n else:\n camid = 1\n pid_container = set()\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n pid_container.add(pid)\n pid2label = {pid:label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n if relabel:\n pid = pid2label[pid]\n masks_path = self.infer_masks_path(img_path)\n data.append({'img_path': img_path,\n 'pid': pid,\n 'masks_path': masks_path,\n 'camid': camid})\n return data\n" }, { "path": "torchreid/data/datasets/image/partial_ilids.py", "content": "from __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport os.path as osp\nimport glob\nimport warnings\n\nfrom ..dataset import ImageDataset\n\n# Sources :\n# https://github.com/hh23333/PVPM\n# Lingxiao He, Jian Liang, Haiqing Li, and Zhenan Sun, \"Deep spatial feature reconstruction for partial person reidentification: Alignment-free approach\", 2018\n\nclass Partial_iLIDS(ImageDataset):\n dataset_dir = 'Partial_iLIDS'\n\n def __init__(self, root='', **kwargs):\n self.root=osp.abspath(osp.expanduser(root))\n # self.dataset_dir = self.root\n data_dir = osp.join(self.root, self.dataset_dir)\n if osp.isdir(data_dir):\n self.data_dir = data_dir\n else:\n warnings.warn('The current data structure is deprecated.')\n self.query_dir = osp.join(self.data_dir, 'partial_body_images')\n self.gallery_dir = osp.join(self.data_dir, 'whole_body_images')\n\n train = []\n query = self.process_dir(self.query_dir)\n gallery = self.process_dir(self.gallery_dir, is_query=False)\n super(Partial_iLIDS, self).__init__(train, query, gallery, **kwargs)\n self.load_pose = isinstance(self.transform, tuple)\n if self.load_pose:\n if self.mode == 'query':\n self.pose_dir = osp.join(self.data_dir, 'occluded_body_pose')\n elif self.mode == 'gallery':\n self.pose_dir = osp.join(self.data_dir, 'whole_body_pose')\n else:\n self.pose_dir = ''\n\n def process_dir(self, dir_path, is_query=True):\n img_paths = glob.glob(osp.join(dir_path, '*.jpg'))\n if is_query:\n camid = 0\n else:\n camid = 1\n\n data = []\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('.')[0])\n data.append({'img_path': img_path, 'pid': pid, 'camid': camid})\n return data\n" }, { "path": "torchreid/data/datasets/image/partial_reid.py", "content": "from __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport os.path as osp\nimport glob\nimport warnings\n\nfrom ..dataset import ImageDataset\n\n# Source :\n# https://github.com/hh23333/PVPM\n# Zheng, W. S., Li, X., Xiang, T., Liao, S., Lai, J., & Gong, S. (2015). Partial person re-identification. ICCV, 2015\n\n\nclass Partial_REID(ImageDataset):\n dataset_dir = 'Partial_REID'\n\n def __init__(self, root='', **kwargs):\n self.root=osp.abspath(osp.expanduser(root))\n # self.dataset_dir = self.root\n data_dir = osp.join(self.root, self.dataset_dir)\n if osp.isdir(data_dir):\n self.data_dir = data_dir\n else:\n warnings.warn('The current data structure is deprecated.')\n self.query_dir = osp.join(self.data_dir, 'partial_body_images')\n self.gallery_dir = osp.join(self.data_dir, 'whole_body_images')\n\n train = []\n query = self.process_dir(self.query_dir, relabel=False)\n gallery = self.process_dir(self.gallery_dir, relabel=False, is_query=False)\n super(Partial_REID, self).__init__(train, query, gallery, **kwargs)\n self.load_pose = isinstance(self.transform, tuple)\n if self.load_pose:\n if self.mode == 'query':\n self.pose_dir = osp.join(self.data_dir, 'occluded_body_pose')\n elif self.mode == 'gallery':\n self.pose_dir = osp.join(self.data_dir, 'whole_body_pose')\n else:\n self.pose_dir = ''\n\n def process_dir(self, dir_path, relabel=False, is_query=True):\n img_paths = glob.glob(osp.join(dir_path, '*.jpg'))\n if is_query:\n camid = 0\n else:\n camid = 1\n pid_container = set()\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n pid_container.add(pid)\n pid2label = {pid: label for label, pid in enumerate(pid_container)}\n\n data = []\n for img_path in img_paths:\n img_name = img_path.split('/')[-1]\n pid = int(img_name.split('_')[0])\n if relabel:\n pid = pid2label[pid]\n data.append({'img_path': img_path, 'pid': pid, 'camid': camid})\n return data\n" }, { "path": "torchreid/data/datasets/image/prid.py", "content": "from __future__ import division, print_function, absolute_import\nimport random\nimport os.path as osp\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import ImageDataset\n\n\nclass PRID(ImageDataset):\n \"\"\"PRID (single-shot version of prid-2011)\n\n Reference:\n Hirzer et al. Person Re-Identification by Descriptive and Discriminative\n Classification. SCIA 2011.\n\n URL: ``_\n \n Dataset statistics:\n - Two views.\n - View A captures 385 identities.\n - View B captures 749 identities.\n - 200 identities appear in both views.\n \"\"\"\n dataset_dir = 'prid2011'\n dataset_url = None\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.cam_a_dir = osp.join(\n self.dataset_dir, 'prid_2011', 'single_shot', 'cam_a'\n )\n self.cam_b_dir = osp.join(\n self.dataset_dir, 'prid_2011', 'single_shot', 'cam_b'\n )\n self.split_path = osp.join(self.dataset_dir, 'splits_single_shot.json')\n\n required_files = [self.dataset_dir, self.cam_a_dir, self.cam_b_dir]\n self.check_before_run(required_files)\n\n self.prepare_split()\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, but expected between 0 and {}'\n .format(split_id,\n len(splits) - 1)\n )\n split = splits[split_id]\n\n train, query, gallery = self.process_split(split)\n\n super(PRID, self).__init__(train, query, gallery, **kwargs)\n\n def prepare_split(self):\n if not osp.exists(self.split_path):\n print('Creating splits ...')\n\n splits = []\n for _ in range(10):\n # randomly sample 100 IDs for train and use the rest 100 IDs for test\n # (note: there are only 200 IDs appearing in both views)\n pids = [i for i in range(1, 201)]\n train_pids = random.sample(pids, 100)\n train_pids.sort()\n test_pids = [i for i in pids if i not in train_pids]\n split = {'train': train_pids, 'test': test_pids}\n splits.append(split)\n\n print('Totally {} splits are created'.format(len(splits)))\n write_json(splits, self.split_path)\n print('Split file is saved to {}'.format(self.split_path))\n\n def process_split(self, split):\n train_pids = split['train']\n test_pids = split['test']\n\n train_pid2label = {pid: label for label, pid in enumerate(train_pids)}\n\n # train\n train = []\n for pid in train_pids:\n img_name = 'person_' + str(pid).zfill(4) + '.png'\n pid = train_pid2label[pid]\n img_a_path = osp.join(self.cam_a_dir, img_name)\n train.append((img_a_path, pid, 0))\n img_b_path = osp.join(self.cam_b_dir, img_name)\n train.append((img_b_path, pid, 1))\n\n # query and gallery\n query, gallery = [], []\n for pid in test_pids:\n img_name = 'person_' + str(pid).zfill(4) + '.png'\n img_a_path = osp.join(self.cam_a_dir, img_name)\n query.append((img_a_path, pid, 0))\n img_b_path = osp.join(self.cam_b_dir, img_name)\n gallery.append((img_b_path, pid, 1))\n for pid in range(201, 750):\n img_name = 'person_' + str(pid).zfill(4) + '.png'\n img_b_path = osp.join(self.cam_b_dir, img_name)\n gallery.append((img_b_path, pid, 1))\n\n return train, query, gallery\n" }, { "path": "torchreid/data/datasets/image/sensereid.py", "content": "from __future__ import division, print_function, absolute_import\nimport copy\nimport glob\nimport os.path as osp\n\nfrom ..dataset import ImageDataset\n\n\nclass SenseReID(ImageDataset):\n \"\"\"SenseReID.\n\n This dataset is used for test purpose only.\n\n Reference:\n Zhao et al. Spindle Net: Person Re-identification with Human Body\n Region Guided Feature Decomposition and Fusion. CVPR 2017.\n\n URL: ``_\n\n Dataset statistics:\n - query: 522 ids, 1040 images.\n - gallery: 1717 ids, 3388 images.\n \"\"\"\n dataset_dir = 'sensereid'\n dataset_url = None\n\n def __init__(self, root='', **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.query_dir = osp.join(self.dataset_dir, 'SenseReID', 'test_probe')\n self.gallery_dir = osp.join(\n self.dataset_dir, 'SenseReID', 'test_gallery'\n )\n\n required_files = [self.dataset_dir, self.query_dir, self.gallery_dir]\n self.check_before_run(required_files)\n\n query = self.process_dir(self.query_dir)\n gallery = self.process_dir(self.gallery_dir)\n\n # relabel\n g_pids = set()\n for _, pid, _ in gallery:\n g_pids.add(pid)\n pid2label = {pid: i for i, pid in enumerate(g_pids)}\n\n query = [\n (img_path, pid2label[pid], camid) for img_path, pid, camid in query\n ]\n gallery = [\n (img_path, pid2label[pid], camid)\n for img_path, pid, camid in gallery\n ]\n train = copy.deepcopy(query) + copy.deepcopy(gallery) # dummy variable\n\n super(SenseReID, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path):\n img_paths = glob.glob(osp.join(dir_path, '*.jpg'))\n data = []\n\n for img_path in img_paths:\n img_name = osp.splitext(osp.basename(img_path))[0]\n pid, camid = img_name.split('_')\n pid, camid = int(pid), int(camid)\n data.append({'img_path': img_path, 'pid': pid, 'camid': camid})\n\n return data\n" }, { "path": "torchreid/data/datasets/image/viper.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport numpy as np\nimport os.path as osp\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import ImageDataset\n\n\nclass VIPeR(ImageDataset):\n \"\"\"VIPeR.\n\n Reference:\n Gray et al. Evaluating appearance models for recognition, reacquisition, and tracking. PETS 2007.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 632.\n - images: 632 x 2 = 1264.\n - cameras: 2.\n \"\"\"\n dataset_dir = 'viper'\n dataset_url = 'http://users.soe.ucsc.edu/~manduchi/VIPeR.v1.0.zip'\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.cam_a_dir = osp.join(self.dataset_dir, 'VIPeR', 'cam_a')\n self.cam_b_dir = osp.join(self.dataset_dir, 'VIPeR', 'cam_b')\n self.split_path = osp.join(self.dataset_dir, 'splits.json')\n\n required_files = [self.dataset_dir, self.cam_a_dir, self.cam_b_dir]\n self.check_before_run(required_files)\n\n self.prepare_split()\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, '\n 'but expected between 0 and {}'.format(\n split_id,\n len(splits) - 1\n )\n )\n split = splits[split_id]\n\n train = split['train']\n query = split['query'] # query and gallery share the same images\n gallery = split['gallery']\n\n train = [tuple(item) for item in train]\n query = [tuple(item) for item in query]\n gallery = [tuple(item) for item in gallery]\n\n super(VIPeR, self).__init__(train, query, gallery, **kwargs)\n\n def prepare_split(self):\n if not osp.exists(self.split_path):\n print('Creating 10 random splits of train ids and test ids')\n\n cam_a_imgs = sorted(glob.glob(osp.join(self.cam_a_dir, '*.bmp')))\n cam_b_imgs = sorted(glob.glob(osp.join(self.cam_b_dir, '*.bmp')))\n assert len(cam_a_imgs) == len(cam_b_imgs)\n num_pids = len(cam_a_imgs)\n print('Number of identities: {}'.format(num_pids))\n num_train_pids = num_pids // 2\n \"\"\"\n In total, there will be 20 splits because each random split creates two\n sub-splits, one using cameraA as query and cameraB as gallery\n while the other using cameraB as query and cameraA as gallery.\n Therefore, results should be averaged over 20 splits (split_id=0~19).\n \n In practice, a model trained on split_id=0 can be applied to split_id=0&1\n as split_id=0&1 share the same training data (so on and so forth).\n \"\"\"\n splits = []\n for _ in range(10):\n order = np.arange(num_pids)\n np.random.shuffle(order)\n train_idxs = order[:num_train_pids]\n test_idxs = order[num_train_pids:]\n assert not bool(set(train_idxs) & set(test_idxs)), \\\n 'Error: train and test overlap'\n\n train = []\n for pid, idx in enumerate(train_idxs):\n cam_a_img = cam_a_imgs[idx]\n cam_b_img = cam_b_imgs[idx]\n train.append((cam_a_img, pid, 0))\n train.append((cam_b_img, pid, 1))\n\n test_a = []\n test_b = []\n for pid, idx in enumerate(test_idxs):\n cam_a_img = cam_a_imgs[idx]\n cam_b_img = cam_b_imgs[idx]\n test_a.append((cam_a_img, pid, 0))\n test_b.append((cam_b_img, pid, 1))\n\n # use cameraA as query and cameraB as gallery\n split = {\n 'train': train,\n 'query': test_a,\n 'gallery': test_b,\n 'num_train_pids': num_train_pids,\n 'num_query_pids': num_pids - num_train_pids,\n 'num_gallery_pids': num_pids - num_train_pids\n }\n splits.append(split)\n\n # use cameraB as query and cameraA as gallery\n split = {\n 'train': train,\n 'query': test_b,\n 'gallery': test_a,\n 'num_train_pids': num_train_pids,\n 'num_query_pids': num_pids - num_train_pids,\n 'num_gallery_pids': num_pids - num_train_pids\n }\n splits.append(split)\n\n print('Totally {} splits are created'.format(len(splits)))\n write_json(splits, self.split_path)\n print('Split file saved to {}'.format(self.split_path))\n" }, { "path": "torchreid/data/datasets/video/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom .mars import Mars\nfrom .ilidsvid import iLIDSVID\nfrom .prid2011 import PRID2011\nfrom .dukemtmcvidreid import DukeMTMCVidReID\n" }, { "path": "torchreid/data/datasets/video/dukemtmcvidreid.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport os.path as osp\nimport warnings\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import VideoDataset\n\n\nclass DukeMTMCVidReID(VideoDataset):\n \"\"\"DukeMTMCVidReID.\n\n Reference:\n - Ristani et al. Performance Measures and a Data Set for Multi-Target,\n Multi-Camera Tracking. ECCVW 2016.\n - Wu et al. Exploit the Unknown Gradually: One-Shot Video-Based Person\n Re-Identification by Stepwise Learning. CVPR 2018.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 702 (train) + 702 (test).\n - tracklets: 2196 (train) + 2636 (test).\n \"\"\"\n dataset_dir = 'dukemtmc-vidreid'\n dataset_url = 'http://vision.cs.duke.edu/DukeMTMC/data/misc/DukeMTMC-VideoReID.zip'\n\n def __init__(self, root='', min_seq_len=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.train_dir = osp.join(self.dataset_dir, 'DukeMTMC-VideoReID/train')\n self.query_dir = osp.join(self.dataset_dir, 'DukeMTMC-VideoReID/query')\n self.gallery_dir = osp.join(\n self.dataset_dir, 'DukeMTMC-VideoReID/gallery'\n )\n self.split_train_json_path = osp.join(\n self.dataset_dir, 'split_train.json'\n )\n self.split_query_json_path = osp.join(\n self.dataset_dir, 'split_query.json'\n )\n self.split_gallery_json_path = osp.join(\n self.dataset_dir, 'split_gallery.json'\n )\n self.min_seq_len = min_seq_len\n\n required_files = [\n self.dataset_dir, self.train_dir, self.query_dir, self.gallery_dir\n ]\n self.check_before_run(required_files)\n\n train = self.process_dir(\n self.train_dir, self.split_train_json_path, relabel=True\n )\n query = self.process_dir(\n self.query_dir, self.split_query_json_path, relabel=False\n )\n gallery = self.process_dir(\n self.gallery_dir, self.split_gallery_json_path, relabel=False\n )\n\n super(DukeMTMCVidReID, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dir_path, json_path, relabel):\n if osp.exists(json_path):\n split = read_json(json_path)\n return split['tracklets']\n\n print('=> Generating split json file (** this might take a while **)')\n pdirs = glob.glob(osp.join(dir_path, '*')) # avoid .DS_Store\n print(\n 'Processing \"{}\" with {} person identities'.format(\n dir_path, len(pdirs)\n )\n )\n\n pid_container = set()\n for pdir in pdirs:\n pid = int(osp.basename(pdir))\n pid_container.add(pid)\n pid2label = {pid: label for label, pid in enumerate(pid_container)}\n\n tracklets = []\n for pdir in pdirs:\n pid = int(osp.basename(pdir))\n if relabel:\n pid = pid2label[pid]\n tdirs = glob.glob(osp.join(pdir, '*'))\n for tdir in tdirs:\n raw_img_paths = glob.glob(osp.join(tdir, '*.jpg'))\n num_imgs = len(raw_img_paths)\n\n if num_imgs < self.min_seq_len:\n continue\n\n img_paths = []\n for img_idx in range(num_imgs):\n # some tracklet starts from 0002 instead of 0001\n img_idx_name = 'F' + str(img_idx + 1).zfill(4)\n res = glob.glob(\n osp.join(tdir, '*' + img_idx_name + '*.jpg')\n )\n if len(res) == 0:\n warnings.warn(\n 'Index name {} in {} is missing, skip'.format(\n img_idx_name, tdir\n )\n )\n continue\n img_paths.append(res[0])\n img_name = osp.basename(img_paths[0])\n if img_name.find('_') == -1:\n # old naming format: 0001C6F0099X30823.jpg\n camid = int(img_name[5]) - 1\n else:\n # new naming format: 0001_C6_F0099_X30823.jpg\n camid = int(img_name[6]) - 1\n img_paths = tuple(img_paths)\n tracklets.append((img_paths, pid, camid))\n\n print('Saving split to {}'.format(json_path))\n split_dict = {'tracklets': tracklets}\n write_json(split_dict, json_path)\n\n return tracklets\n" }, { "path": "torchreid/data/datasets/video/ilidsvid.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport os.path as osp\nfrom scipy.io import loadmat\n\nfrom torchreid.utils import read_json, write_json\n\nfrom ..dataset import VideoDataset\n\n\nclass iLIDSVID(VideoDataset):\n \"\"\"iLIDS-VID.\n\n Reference:\n Wang et al. Person Re-Identification by Video Ranking. ECCV 2014.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 300.\n - tracklets: 600.\n - cameras: 2.\n \"\"\"\n dataset_dir = 'ilids-vid'\n dataset_url = 'http://www.eecs.qmul.ac.uk/~xiatian/iLIDS-VID/iLIDS-VID.tar'\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.data_dir = osp.join(self.dataset_dir, 'i-LIDS-VID')\n self.split_dir = osp.join(self.dataset_dir, 'train-test people splits')\n self.split_mat_path = osp.join(\n self.split_dir, 'train_test_splits_ilidsvid.mat'\n )\n self.split_path = osp.join(self.dataset_dir, 'splits.json')\n self.cam_1_path = osp.join(\n self.dataset_dir, 'i-LIDS-VID/sequences/cam1'\n )\n self.cam_2_path = osp.join(\n self.dataset_dir, 'i-LIDS-VID/sequences/cam2'\n )\n\n required_files = [self.dataset_dir, self.data_dir, self.split_dir]\n self.check_before_run(required_files)\n\n self.prepare_split()\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, but expected between 0 and {}'\n .format(split_id,\n len(splits) - 1)\n )\n split = splits[split_id]\n train_dirs, test_dirs = split['train'], split['test']\n\n train = self.process_data(train_dirs, cam1=True, cam2=True)\n query = self.process_data(test_dirs, cam1=True, cam2=False)\n gallery = self.process_data(test_dirs, cam1=False, cam2=True)\n\n super(iLIDSVID, self).__init__(train, query, gallery, **kwargs)\n\n def prepare_split(self):\n if not osp.exists(self.split_path):\n print('Creating splits ...')\n mat_split_data = loadmat(self.split_mat_path)['ls_set']\n\n num_splits = mat_split_data.shape[0]\n num_total_ids = mat_split_data.shape[1]\n assert num_splits == 10\n assert num_total_ids == 300\n num_ids_each = num_total_ids // 2\n\n # pids in mat_split_data are indices, so we need to transform them\n # to real pids\n person_cam1_dirs = sorted(\n glob.glob(osp.join(self.cam_1_path, '*'))\n )\n person_cam2_dirs = sorted(\n glob.glob(osp.join(self.cam_2_path, '*'))\n )\n\n person_cam1_dirs = [\n osp.basename(item) for item in person_cam1_dirs\n ]\n person_cam2_dirs = [\n osp.basename(item) for item in person_cam2_dirs\n ]\n\n # make sure persons in one camera view can be found in the other camera view\n assert set(person_cam1_dirs) == set(person_cam2_dirs)\n\n splits = []\n for i_split in range(num_splits):\n # first 50% for testing and the remaining for training, following Wang et al. ECCV'14.\n train_idxs = sorted(\n list(mat_split_data[i_split, num_ids_each:])\n )\n test_idxs = sorted(\n list(mat_split_data[i_split, :num_ids_each])\n )\n\n train_idxs = [int(i) - 1 for i in train_idxs]\n test_idxs = [int(i) - 1 for i in test_idxs]\n\n # transform pids to person dir names\n train_dirs = [person_cam1_dirs[i] for i in train_idxs]\n test_dirs = [person_cam1_dirs[i] for i in test_idxs]\n\n split = {'train': train_dirs, 'test': test_dirs}\n splits.append(split)\n\n print(\n 'Totally {} splits are created, following Wang et al. ECCV\\'14'\n .format(len(splits))\n )\n print('Split file is saved to {}'.format(self.split_path))\n write_json(splits, self.split_path)\n\n def process_data(self, dirnames, cam1=True, cam2=True):\n tracklets = []\n dirname2pid = {dirname: i for i, dirname in enumerate(dirnames)}\n\n for dirname in dirnames:\n if cam1:\n person_dir = osp.join(self.cam_1_path, dirname)\n img_names = glob.glob(osp.join(person_dir, '*.png'))\n assert len(img_names) > 0\n img_names = tuple(img_names)\n pid = dirname2pid[dirname]\n tracklets.append((img_names, pid, 0))\n\n if cam2:\n person_dir = osp.join(self.cam_2_path, dirname)\n img_names = glob.glob(osp.join(person_dir, '*.png'))\n assert len(img_names) > 0\n img_names = tuple(img_names)\n pid = dirname2pid[dirname]\n tracklets.append((img_names, pid, 1))\n\n return tracklets\n" }, { "path": "torchreid/data/datasets/video/mars.py", "content": "from __future__ import division, print_function, absolute_import\nimport os.path as osp\nimport warnings\nfrom scipy.io import loadmat\n\nfrom ..dataset import VideoDataset\n\n\nclass Mars(VideoDataset):\n \"\"\"MARS.\n\n Reference:\n Zheng et al. MARS: A Video Benchmark for Large-Scale Person Re-identification. ECCV 2016.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 1261.\n - tracklets: 8298 (train) + 1980 (query) + 9330 (gallery).\n - cameras: 6.\n \"\"\"\n dataset_dir = 'mars'\n dataset_url = None\n\n def __init__(self, root='', **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.train_name_path = osp.join(\n self.dataset_dir, 'info/train_name.txt'\n )\n self.test_name_path = osp.join(self.dataset_dir, 'info/test_name.txt')\n self.track_train_info_path = osp.join(\n self.dataset_dir, 'info/tracks_train_info.mat'\n )\n self.track_test_info_path = osp.join(\n self.dataset_dir, 'info/tracks_test_info.mat'\n )\n self.query_IDX_path = osp.join(self.dataset_dir, 'info/query_IDX.mat')\n\n required_files = [\n self.dataset_dir, self.train_name_path, self.test_name_path,\n self.track_train_info_path, self.track_test_info_path,\n self.query_IDX_path\n ]\n self.check_before_run(required_files)\n\n train_names = self.get_names(self.train_name_path)\n test_names = self.get_names(self.test_name_path)\n track_train = loadmat(self.track_train_info_path\n )['track_train_info'] # numpy.ndarray (8298, 4)\n track_test = loadmat(self.track_test_info_path\n )['track_test_info'] # numpy.ndarray (12180, 4)\n query_IDX = loadmat(self.query_IDX_path\n )['query_IDX'].squeeze() # numpy.ndarray (1980,)\n query_IDX -= 1 # index from 0\n track_query = track_test[query_IDX, :]\n gallery_IDX = [\n i for i in range(track_test.shape[0]) if i not in query_IDX\n ]\n track_gallery = track_test[gallery_IDX, :]\n\n train = self.process_data(\n train_names, track_train, home_dir='bbox_train', relabel=True\n )\n query = self.process_data(\n test_names, track_query, home_dir='bbox_test', relabel=False\n )\n gallery = self.process_data(\n test_names, track_gallery, home_dir='bbox_test', relabel=False\n )\n\n super(Mars, self).__init__(train, query, gallery, **kwargs)\n\n def get_names(self, fpath):\n names = []\n with open(fpath, 'r') as f:\n for line in f:\n new_line = line.rstrip()\n names.append(new_line)\n return names\n\n def process_data(\n self, names, meta_data, home_dir=None, relabel=False, min_seq_len=0\n ):\n assert home_dir in ['bbox_train', 'bbox_test']\n num_tracklets = meta_data.shape[0]\n pid_list = list(set(meta_data[:, 2].tolist()))\n\n if relabel:\n pid2label = {pid: label for label, pid in enumerate(pid_list)}\n tracklets = []\n\n for tracklet_idx in range(num_tracklets):\n data = meta_data[tracklet_idx, ...]\n start_index, end_index, pid, camid = data\n if pid == -1:\n continue # junk images are just ignored\n assert 1 <= camid <= 6\n if relabel:\n pid = pid2label[pid]\n camid -= 1 # index starts from 0\n img_names = names[start_index - 1:end_index]\n\n # make sure image names correspond to the same person\n pnames = [img_name[:4] for img_name in img_names]\n assert len(\n set(pnames)\n ) == 1, 'Error: a single tracklet contains different person images'\n\n # make sure all images are captured under the same camera\n camnames = [img_name[5] for img_name in img_names]\n assert len(\n set(camnames)\n ) == 1, 'Error: images are captured under different cameras!'\n\n # append image names with directory information\n img_paths = [\n osp.join(self.dataset_dir, home_dir, img_name[:4], img_name)\n for img_name in img_names\n ]\n if len(img_paths) >= min_seq_len:\n img_paths = tuple(img_paths)\n tracklets.append((img_paths, pid, camid))\n\n return tracklets\n\n def combine_all(self):\n warnings.warn(\n 'Some query IDs do not appear in gallery. Therefore, combineall '\n 'does not make any difference to Mars'\n )\n" }, { "path": "torchreid/data/datasets/video/prid2011.py", "content": "from __future__ import division, print_function, absolute_import\nimport glob\nimport os.path as osp\n\nfrom torchreid.utils import read_json\n\nfrom ..dataset import VideoDataset\n\n\nclass PRID2011(VideoDataset):\n \"\"\"PRID2011.\n\n Reference:\n Hirzer et al. Person Re-Identification by Descriptive and\n Discriminative Classification. SCIA 2011.\n\n URL: ``_\n \n Dataset statistics:\n - identities: 200.\n - tracklets: 400.\n - cameras: 2.\n \"\"\"\n dataset_dir = 'prid2011'\n dataset_url = None\n\n def __init__(self, root='', split_id=0, **kwargs):\n self.root = osp.abspath(osp.expanduser(root))\n self.dataset_dir = osp.join(self.root, self.dataset_dir)\n self.download_dataset(self.dataset_dir, self.dataset_url)\n\n self.split_path = osp.join(self.dataset_dir, 'splits_prid2011.json')\n self.cam_a_dir = osp.join(\n self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_a'\n )\n self.cam_b_dir = osp.join(\n self.dataset_dir, 'prid_2011', 'multi_shot', 'cam_b'\n )\n\n required_files = [self.dataset_dir, self.cam_a_dir, self.cam_b_dir]\n self.check_before_run(required_files)\n\n splits = read_json(self.split_path)\n if split_id >= len(splits):\n raise ValueError(\n 'split_id exceeds range, received {}, but expected between 0 and {}'\n .format(split_id,\n len(splits) - 1)\n )\n split = splits[split_id]\n train_dirs, test_dirs = split['train'], split['test']\n\n train = self.process_dir(train_dirs, cam1=True, cam2=True)\n query = self.process_dir(test_dirs, cam1=True, cam2=False)\n gallery = self.process_dir(test_dirs, cam1=False, cam2=True)\n\n super(PRID2011, self).__init__(train, query, gallery, **kwargs)\n\n def process_dir(self, dirnames, cam1=True, cam2=True):\n tracklets = []\n dirname2pid = {dirname: i for i, dirname in enumerate(dirnames)}\n\n for dirname in dirnames:\n if cam1:\n person_dir = osp.join(self.cam_a_dir, dirname)\n img_names = glob.glob(osp.join(person_dir, '*.png'))\n assert len(img_names) > 0\n img_names = tuple(img_names)\n pid = dirname2pid[dirname]\n tracklets.append((img_names, pid, 0))\n\n if cam2:\n person_dir = osp.join(self.cam_b_dir, dirname)\n img_names = glob.glob(osp.join(person_dir, '*.png'))\n assert len(img_names) > 0\n img_names = tuple(img_names)\n pid = dirname2pid[dirname]\n tracklets.append((img_names, pid, 1))\n\n return tracklets\n" }, { "path": "torchreid/data/masks_transforms/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom .mask_transform import *\nfrom .pcb_transforms import *\nfrom .pifpaf_mask_transform import *\nfrom .coco_keypoints_transforms import *\nfrom torchreid.data.datasets import get_image_dataset\n\nmasks_preprocess_pifpaf = {\n 'full': CombinePifPafIntoFullBodyMask,\n 'bs_fu': AddFullBodyMaskToBaseMasks,\n 'bs_fu_bb': AddFullBodyMaskAndFullBoundingBoxToBaseMasks,\n 'mu_sc': CombinePifPafIntoMultiScaleBodyMasks,\n 'one': CombinePifPafIntoOneBodyMasks,\n 'two_v': CombinePifPafIntoTwoBodyMasks,\n 'three_v': CombinePifPafIntoThreeBodyMasks,\n 'four': CombinePifPafIntoFourBodyMasks,\n 'four_no': CombinePifPafIntoFourBodyMasksNoOverlap,\n 'four_v': CombinePifPafIntoFourVerticalParts,\n 'four_v_pif': CombinePifPafIntoFourVerticalPartsPif,\n 'five_v': CombinePifPafIntoFiveVerticalParts,\n 'five': CombinePifPafIntoFiveBodyMasks,\n 'six': CombinePifPafIntoSixBodyMasks,\n 'six_v': CombinePifPafIntoSixVerticalParts,\n 'six_no': CombinePifPafIntoSixBodyMasksSum,\n 'six_new': CombinePifPafIntoSixBodyMasksSimilarToEight,\n 'seven_v': CombinePifPafIntoSevenVerticalBodyMasks,\n 'seven_new': CombinePifPafIntoSevenBodyMasksSimilarToEight,\n 'eight': CombinePifPafIntoEightBodyMasks,\n 'eight_v': CombinePifPafIntoEightVerticalBodyMasks,\n 'ten_ms': CombinePifPafIntoTenMSBodyMasks,\n 'eleven': CombinePifPafIntoElevenBodyMasks,\n 'fourteen': CombinePifPafIntoFourteenBodyMasks,\n}\n\nmasks_preprocess_coco = {\n 'cc6': CocoToSixBodyMasks\n}\n\nmasks_preprocess_fixed = {\n 'id': IdentityMask,\n 'strp_2': PCBMasks2,\n 'strp_3': PCBMasks3,\n 'strp_4': PCBMasks4,\n 'strp_5': PCBMasks5,\n 'strp_6': PCBMasks6,\n 'strp_7': PCBMasks7,\n 'strp_8': PCBMasks8,\n}\n\nmasks_preprocess_transforms = {**masks_preprocess_pifpaf, **masks_preprocess_coco}\nmasks_preprocess_all = {**masks_preprocess_pifpaf, **masks_preprocess_fixed, **masks_preprocess_coco}\n\n\ndef compute_parts_num_and_names(cfg):\n mask_config = get_image_dataset(cfg.data.sources[0]).get_masks_config(cfg.model.bpbreid.masks.dir)\n if cfg.loss.name == 'part_based':\n if (mask_config is not None and mask_config[1]) or cfg.model.bpbreid.masks.preprocess == 'none':\n # ISP masks or no transform\n cfg.model.bpbreid.masks.parts_num = mask_config[0]\n cfg.model.bpbreid.masks.parts_names = mask_config[3] if 3 in mask_config else [\"p{}\".format(p) for p in range(1, cfg.data.parts_num+1)]\n else:\n masks_transform = masks_preprocess_all[cfg.model.bpbreid.masks.preprocess]()\n cfg.model.bpbreid.masks.parts_num = masks_transform.parts_num\n cfg.model.bpbreid.masks.parts_names = masks_transform.parts_names" }, { "path": "torchreid/data/masks_transforms/coco_keypoints_transforms.py", "content": "from torchreid.data.masks_transforms.mask_transform import MaskGroupingTransform\n\nCOCO_KEYPOINTS = [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_shoulder\", \"right_shoulder\",\n \"left_elbow\", \"right_elbow\", \"left_wrist\", \"right_wrist\", \"left_hip\", \"right_hip\", \"left_knee\",\n \"right_knee\", \"left_ankle\", \"right_ankle\"]\n\nCOCO_KEYPOINTS_MAP = {k: i for i, k in enumerate(COCO_KEYPOINTS)}\n\n\nclass CocoToSixBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\"],\n \"torso\": [\"left_shoulder\", \"right_shoulder\", \"left_hip\", \"right_hip\"],\n \"left_arm\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\"],\n \"right_arm\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\"],\n \"left_leg\": [\"left_hip\", \"left_knee\", \"left_ankle\"],\n \"right_leg\": [\"right_hip\", \"right_knee\", \"right_ankle\"]\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, COCO_KEYPOINTS_MAP)" }, { "path": "torchreid/data/masks_transforms/mask_transform.py", "content": "import torch\nfrom torch import nn\nfrom albumentations import DualTransform\nimport torch.nn.functional as F\n\n\nclass MaskTransform(DualTransform):\n def __init__(self):\n super(MaskTransform, self).__init__(always_apply=True, p=1)\n\n def apply(self, img, **params):\n return img\n\n def apply_to_bbox(self, bbox, **params):\n raise NotImplementedError(\"Method apply_to_bbox is not implemented in class \" + self.__class__.__name__)\n\n def apply_to_keypoint(self, keypoint, **params):\n raise NotImplementedError(\"Method apply_to_keypoint is not implemented in class \" + self.__class__.__name__)\n\n\nclass MaskGroupingTransform(MaskTransform):\n\n def __init__(self, parts_grouping, parts_map, combine_mode='max'):\n super().__init__()\n self.parts_grouping = parts_grouping\n self.parts_map = parts_map\n self.parts_names = list(parts_grouping.keys())\n self.parts_num = len(self.parts_names)\n self.combine_mode = combine_mode\n\n def apply_to_mask(self, masks, **params):\n parts_masks = []\n for i, part in enumerate(self.parts_names):\n if self.combine_mode == 'sum':\n parts_masks.append(masks[[self.parts_map[k] for k in self.parts_grouping[part]]].sum(dim=0).clamp(0, 1))\n else:\n parts_masks.append(masks[[self.parts_map[k] for k in self.parts_grouping[part]]].max(dim=0)[0].clamp(0, 1))\n return torch.stack(parts_masks)\n\n\nclass PermuteMasksDim(MaskTransform):\n def apply_to_mask(self, masks, **params):\n return masks.permute(2, 0, 1)\n\n\nclass ResizeMasks(MaskTransform):\n def __init__(self, height, width, mask_scale):\n super(ResizeMasks, self).__init__()\n self._size = (int(height/mask_scale), int(width/mask_scale))\n\n def apply_to_mask(self, masks, **params):\n return nn.functional.interpolate(masks.unsqueeze(0), self._size, mode='nearest').squeeze(0) # Best perf with nearest here and bilinear in parts engine\n\n\nclass RemoveBackgroundMask(MaskTransform):\n def apply_to_mask(self, masks, **params):\n return masks[:, :, 1::]\n\n\nclass AddBackgroundMask(MaskTransform):\n def __init__(self, background_computation_strategy='sum', softmax_weight=0, mask_filtering_threshold=0.3):\n super().__init__()\n self.background_computation_strategy = background_computation_strategy\n self.softmax_weight = softmax_weight\n self.mask_filtering_threshold = mask_filtering_threshold\n\n def apply_to_mask(self, masks, **params):\n if self.background_computation_strategy == 'sum':\n background_mask = 1 - masks.sum(dim=0)\n background_mask = background_mask.clamp(0, 1)\n masks = torch.cat([background_mask.unsqueeze(0), masks])\n elif self.background_computation_strategy == 'threshold':\n background_mask = masks.max(dim=0)[0] < self.mask_filtering_threshold\n masks = torch.cat([background_mask.unsqueeze(0), masks])\n elif self.background_computation_strategy == 'diff_from_max':\n background_mask = 1 - masks.max(dim=0)[0]\n background_mask = background_mask.clamp(0, 1)\n masks = torch.cat([background_mask.unsqueeze(0), masks])\n else:\n raise ValueError('Background mask combine strategy {} not supported'.format(self.background_computation_strategy))\n if self.softmax_weight > 0:\n masks = F.softmax(masks * self.softmax_weight, dim=0)\n else:\n masks = masks / masks.sum(dim=0)\n return masks\n\n\nclass IdentityMask(MaskTransform):\n parts_names = ['id']\n parts_num = 1\n def apply_to_mask(self, masks, **params):\n return torch.ones((1, masks.shape[1], masks.shape[2]))\n" }, { "path": "torchreid/data/masks_transforms/pcb_transforms.py", "content": "import numpy as np\nimport torch\n\nfrom torchreid.data.masks_transforms.mask_transform import MaskTransform\n\n\nclass PCBMasks(MaskTransform):\n def apply_to_mask(self, masks, **params):\n self._size = masks.shape[1:3]\n self.stripe_height = self._size[0] / self.parts_num\n\n self.pcb_masks = torch.zeros((self.parts_num, self._size[0], self._size[1]))\n\n stripes_range = np.round(np.arange(0, self.parts_num + 1) * self._size[0] / self.parts_num).astype(int)\n for i in range(0, stripes_range.size-1):\n self.pcb_masks[i, stripes_range[i]:stripes_range[i+1], :] = 1\n\n return self.pcb_masks\n\n\nclass PCBMasks2(PCBMasks):\n parts_num = 2\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n\nclass PCBMasks3(PCBMasks):\n parts_num = 3\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n\nclass PCBMasks4(PCBMasks):\n parts_num = 4\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n\nclass PCBMasks5(PCBMasks):\n parts_num = 5\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n\nclass PCBMasks6(PCBMasks):\n parts_num = 6\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n\nclass PCBMasks7(PCBMasks):\n parts_num = 7\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n\nclass PCBMasks8(PCBMasks):\n parts_num = 8\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n" }, { "path": "torchreid/data/masks_transforms/pifpaf_mask_transform.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport torch\nfrom torchreid.data.masks_transforms.mask_transform import MaskGroupingTransform\n\nPIFPAF_KEYPOINTS = [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_shoulder\", \"right_shoulder\",\n \"left_elbow\", \"right_elbow\", \"left_wrist\", \"right_wrist\", \"left_hip\", \"right_hip\", \"left_knee\",\n \"right_knee\", \"left_ankle\", \"right_ankle\"]\n\nPIFPAF_JOINTS = [\"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\", \"left_hip_to_right_hip\", \"left_shoulder_to_left_hip\",\n \"right_shoulder_to_right_hip\", \"left_shoulder_to_right_shoulder\", \"left_shoulder_to_left_elbow\",\n \"right_shoulder_to_right_elbow\", \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\",\n \"left_eye_to_right_eye\", \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\", \"right_ear_to_right_shoulder\"]\n\nPIFPAF_PARTS = PIFPAF_KEYPOINTS + PIFPAF_JOINTS\nPIFPAF_SINGLE_GROUPS = {k:k for k in PIFPAF_PARTS}\nPIFPAF_PARTS_MAP = {k: i for i, k in enumerate(PIFPAF_PARTS)}\n\n\nclass CombinePifPafIntoFullBodyMask(MaskGroupingTransform):\n parts_grouping = {\n \"full_body\": PIFPAF_PARTS\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass AddFullBodyMaskToBaseMasks(MaskGroupingTransform):\n parts_grouping = {**PIFPAF_SINGLE_GROUPS,\n **{\n \"full_body\": PIFPAF_PARTS\n }\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass AddFullBodyMaskAndFullBoundingBoxToBaseMasks(MaskGroupingTransform):\n parts_num = 38\n parts_names = [\"p{}\".format(p) for p in range(1, parts_num+1)]\n\n def apply_to_mask(self, masks, **params):\n full_body_mask = torch.max(masks, 0, keepdim=True)[0]\n\n full_bounding_box = torch.ones(masks.shape[1:3]).unsqueeze(0)\n\n return torch.cat([masks,\n full_body_mask,\n full_bounding_box\n ])\n\n\nclass CombinePifPafIntoMultiScaleBodyMasks(MaskGroupingTransform):\n parts_grouping = {**PIFPAF_SINGLE_GROUPS, **{\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"arms_mask\": [\"left_shoulder\", \"right_shoulder\", \"left_elbow\", \"right_elbow\", \"left_wrist\",\n \"right_wrist\", \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_shoulder\", \"right_shoulder\", \"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\", \"left_ankle\", \"right_ankle\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\", \"left_hip_to_right_hip\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n\n \"upper_body\": [\"torso_mask\", \"arms_mask\", \"head_mask\"],\n\n \"lower_body\": [\"legs_mask\", \"feet_mask\"],\n\n \"full_body_mask\": PIFPAF_PARTS\n }\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoOneBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"full\": PIFPAF_PARTS\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoTwoBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"torso_arms_head\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\", \"left_shoulder\", \"right_shoulder\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\",\n \"left_elbow\", \"right_elbow\", \"left_wrist\",\n \"right_wrist\", \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\"],\n \"legs\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\", \"left_ankle\", \"right_ankle\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\", \"left_hip_to_right_hip\"]\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoThreeBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"torso_arms_mask\": [\"left_shoulder\", \"right_shoulder\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\",\n \"left_elbow\", \"right_elbow\", \"left_wrist\",\n \"right_wrist\", \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\", \"left_ankle\", \"right_ankle\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\", \"left_hip_to_right_hip\"]\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFourBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"arms_mask\": [\"left_shoulder\", \"right_shoulder\", \"left_elbow\", \"right_elbow\", \"left_wrist\",\n \"right_wrist\", \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_shoulder\", \"right_shoulder\", \"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\", \"left_ankle\", \"right_ankle\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\", \"left_hip_to_right_hip\"]\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFourBodyMasksNoOverlap(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"arms_mask\": [\"left_elbow\", \"right_elbow\", \"left_wrist\",\n \"right_wrist\", \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_shoulder\", \"right_shoulder\", \"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"legs_mask\": [\"left_knee\", \"right_knee\", \"left_ankle\", \"right_ankle\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\", \"left_hip_to_right_hip\"]\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFourVerticalParts(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"arms_torso_mask\": [\"left_elbow\", \"right_elbow\", \"left_wrist\",\n \"right_wrist\", \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\", \"left_shoulder\", \"right_shoulder\", \"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFourVerticalPartsPif(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\"],\n \"arms_torso_mask\": [\"left_elbow\", \"right_elbow\", \"left_wrist\", \"right_wrist\", \"left_shoulder\", \"right_shoulder\", \"left_hip\", \"right_hip\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFiveVerticalParts(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"upper_arms_torso_mask\": [\"left_elbow\", \"right_elbow\",\n \"left_shoulder_to_left_elbow\", \"right_shoulder_to_right_elbow\",\n \"left_shoulder\", \"right_shoulder\",\n \"left_shoulder_to_right_shoulder\"],\n \"lower_arms_torso_mask\": [\"left_wrist\", \"right_wrist\",\n \"left_elbow_to_left_wrist\", \"right_elbow_to_right_wrist\",\n \"left_hip\", \"right_hip\",\n \"right_shoulder_to_right_hip\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFiveBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n\n \"arms_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\", \"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n\n \"legs_mask\": [\"left_hip_to_right_hip\", \"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoSixVerticalParts(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"arms_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\", \"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"upper_torso_mask\": [\"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"lower_torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\"],\n \"legs_mask\": [\"left_hip\", \"right_hip\", \"left_knee\", \"right_knee\",\n \"left_ankle_to_left_knee\", \"left_knee_to_left_hip\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoSixBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"left_leg_mask\": [\"left_knee\", \"left_ankle\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_ankle\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoSixBodyMasksSum(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"left_leg_mask\": [\"left_knee\", \"left_ankle\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_ankle\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS, 'sum')\n\n\nclass CombinePifPafIntoSixBodyMasksSimilarToEight(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"leg_mask\": [\"left_knee\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\", \"right_knee\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoEightBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"left_leg_mask\": [\"left_knee\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"left_feet_mask\": [\"left_ankle\"],\n \"right_feet_mask\": [\"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\n\nclass CombinePifPafIntoEightVerticalBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"left_leg_mask\": [\"left_knee\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"left_feet_mask\": [\"left_ankle\"],\n \"right_feet_mask\": [\"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoTenMSBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"left_leg_mask\": [\"left_knee\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"left_feet_mask\": [\"left_ankle\"],\n \"right_feet_mask\": [\"right_ankle\"],\n\n \"upper_body_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\",\n \"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\",\n \"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\",\n \"left_hip\", \"right_hip\", \"left_hip_to_right_hip\",\n \"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"lower_body_mask\": [\"left_knee\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\",\n \"right_knee\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\",\n \"left_ankle\",\n \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoSevenVerticalBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"shoulders_mask\": [\"left_shoulder\", \"right_shoulder\", \"left_shoulder_to_right_shoulder\"],\n \"elbow_mask\": [\"left_elbow\", \"right_elbow\"],\n \"wrist_mask\": [\"left_wrist\", \"right_wrist\"],\n \"hip_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\"],\n \"knee_mask\": [\"left_knee\", \"right_knee\"],\n \"ankle_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoSevenBodyMasksSimilarToEight(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_arm_mask\": [\"left_shoulder\", \"left_elbow\", \"left_wrist\", \"left_shoulder_to_left_elbow\",\n \"left_elbow_to_left_wrist\"],\n \"right_arm_mask\": [\"right_shoulder\", \"right_elbow\", \"right_wrist\", \"right_shoulder_to_right_elbow\",\n \"right_elbow_to_right_wrist\"],\n \"upper_torso_mask\": [\"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\",\n \"left_shoulder_to_right_shoulder\"],\n \"lower_torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\"],\n \"leg_mask\": [\"left_knee\", \"left_ankle_to_left_knee\",\n \"left_knee_to_left_hip\", \"left_hip_to_right_hip\", \"right_knee\", \"right_ankle_to_right_knee\",\n \"right_knee_to_right_hip\"],\n \"feet_mask\": [\"left_ankle\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\n\nclass CombinePifPafIntoElevenBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\", \"left_ear_to_left_shoulder\",\n \"right_ear_to_right_shoulder\"],\n \"left_elbow_mask\": [\"left_shoulder\", \"left_elbow\", \"left_shoulder_to_left_elbow\"],\n \"left_wrist_mask\": [\"left_wrist\", \"left_elbow_to_left_wrist\"],\n\n \"right_elbow_mask\": [\"right_shoulder\", \"right_elbow\", \"right_shoulder_to_right_elbow\"],\n \"right_wrist_mask\": [\"right_wrist\", \"right_elbow_to_right_wrist\"],\n\n \"upper_torso_mask\": [\"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\", \"left_shoulder_to_right_shoulder\"],\n\n \"lower_torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\"],\n \"left_leg_mask\": [\"left_knee\", \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_knee_to_right_hip\"],\n \"left_feet_mask\": [\"left_ankle_to_left_knee\", \"left_ankle\"],\n \"right_feet_mask\": [\"right_ankle_to_right_knee\", \"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n\n\nclass CombinePifPafIntoFourteenBodyMasks(MaskGroupingTransform):\n parts_grouping = {\n \"head_mask\": [\"nose\", \"left_eye\", \"right_eye\", \"left_ear\", \"right_ear\", \"left_eye_to_right_eye\",\n \"nose_to_left_eye\", \"nose_to_right_eye\", \"left_eye_to_left_ear\",\n \"right_eye_to_right_ear\"],\n \"neck_mask\": [\"left_ear_to_left_shoulder\", \"right_ear_to_right_shoulder\"],\n \"left_elbow_mask\": [\"left_shoulder\", \"left_elbow\", \"left_shoulder_to_left_elbow\"],\n \"left_wrist_mask\": [\"left_wrist\", \"left_elbow_to_left_wrist\"],\n \"right_elbow_mask\": [\"right_shoulder\", \"right_elbow\", \"right_shoulder_to_right_elbow\"],\n \"right_wrist_mask\": [\"right_wrist\", \"right_elbow_to_right_wrist\"],\n \"upper_torso_mask\": [\"left_shoulder_to_left_hip\", \"right_shoulder_to_right_hip\", \"left_shoulder_to_right_shoulder\"],\n \"lower_torso_mask\": [\"left_hip\", \"right_hip\", \"left_hip_to_right_hip\"],\n \"left_leg_mask\": [\"left_knee\", \"left_knee_to_left_hip\", \"left_hip_to_right_hip\"],\n \"right_leg_mask\": [\"right_knee\", \"right_knee_to_right_hip\"],\n \"left_tibia_mask\": [\"left_ankle_to_left_knee\"],\n \"right_tibia_mask\": [\"right_ankle_to_right_knee\"],\n \"left_feet_mask\": [\"left_ankle\"],\n \"right_feet_mask\": [\"right_ankle\"],\n }\n\n def __init__(self):\n super().__init__(self.parts_grouping, PIFPAF_PARTS_MAP)\n" }, { "path": "torchreid/data/sampler.py", "content": "from __future__ import division, absolute_import\nimport copy\nimport numpy as np\nimport random\nfrom collections import defaultdict\nfrom torch.utils.data.sampler import Sampler, RandomSampler, SequentialSampler\n\nAVAI_SAMPLERS = ['RandomIdentitySampler', 'SequentialSampler', 'RandomSampler']\n\n\nclass RandomIdentitySampler(Sampler):\n \"\"\"Randomly samples N identities each with K instances.\n\n Args:\n data_source (list): contains tuples of (img_path(s), pid, camid).\n batch_size (int): batch size.\n num_instances (int): number of instances per identity in a batch.\n \"\"\"\n\n def __init__(self, data_source, batch_size, num_instances):\n if batch_size < num_instances:\n raise ValueError(\n 'batch_size={} must be no less '\n 'than num_instances={}'.format(batch_size, num_instances)\n )\n\n self.data_source = data_source\n self.batch_size = batch_size\n self.num_instances = num_instances\n self.num_pids_per_batch = self.batch_size // self.num_instances\n self.index_dic = defaultdict(list)\n for index, sample in enumerate(self.data_source):\n self.index_dic[sample['pid']].append(index)\n self.pids = list(self.index_dic.keys())\n\n # estimate number of examples in an epoch\n self.length = 0\n for pid in self.pids:\n idxs = self.index_dic[pid]\n num = len(idxs)\n if num < self.num_instances:\n num = self.num_instances\n self.length += num - num % self.num_instances\n\n def __iter__(self):\n batch_idxs_dict = defaultdict(list)\n\n for pid in self.pids:\n idxs = copy.deepcopy(self.index_dic[pid])\n if len(idxs) < self.num_instances:\n idxs = np.random.choice(\n idxs, size=self.num_instances, replace=True\n )\n random.shuffle(idxs)\n batch_idxs = []\n for idx in idxs:\n batch_idxs.append(idx)\n if len(batch_idxs) == self.num_instances:\n batch_idxs_dict[pid].append(batch_idxs)\n batch_idxs = []\n\n avai_pids = copy.deepcopy(self.pids)\n final_idxs = []\n\n while len(avai_pids) >= self.num_pids_per_batch:\n selected_pids = random.sample(avai_pids, self.num_pids_per_batch)\n for pid in selected_pids:\n batch_idxs = batch_idxs_dict[pid].pop(0)\n final_idxs.extend(batch_idxs)\n if len(batch_idxs_dict[pid]) == 0:\n avai_pids.remove(pid)\n\n return iter(final_idxs)\n\n def __len__(self):\n return self.length\n\n\ndef build_train_sampler(\n data_source, train_sampler, batch_size=32, num_instances=4, **kwargs\n):\n \"\"\"Builds a training sampler.\n\n Args:\n data_source (list): contains tuples of (img_path(s), pid, camid).\n train_sampler (str): sampler name (default: ``RandomSampler``).\n batch_size (int, optional): batch size. Default is 32.\n num_instances (int, optional): number of instances per identity in a\n batch (when using ``RandomIdentitySampler``). Default is 4.\n \"\"\"\n assert train_sampler in AVAI_SAMPLERS, \\\n 'train_sampler must be one of {}, but got {}'.format(AVAI_SAMPLERS, train_sampler)\n\n if train_sampler == 'RandomIdentitySampler':\n sampler = RandomIdentitySampler(data_source, batch_size, num_instances)\n\n elif train_sampler == 'SequentialSampler':\n sampler = SequentialSampler(data_source)\n\n elif train_sampler == 'RandomSampler':\n sampler = RandomSampler(data_source)\n\n return sampler\n" }, { "path": "torchreid/data/transforms.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport cv2\nimport torch\nimport numpy as np\nfrom albumentations import (\n Resize, Compose, Normalize, ColorJitter, HorizontalFlip, CoarseDropout, RandomCrop, PadIfNeeded\n)\nfrom albumentations.pytorch import ToTensorV2\nfrom torchreid.data.masks_transforms import masks_preprocess_all, AddBackgroundMask, ResizeMasks, PermuteMasksDim, \\\n RemoveBackgroundMask\nfrom torchreid.data.data_augmentation import RandomOcclusion\n\n\nclass NpToTensor(object):\n def __call__(self, masks):\n assert isinstance(masks, np.ndarray)\n return torch.as_tensor(masks)\n\n def __repr__(self):\n return self.__class__.__name__ + '()'\n\n\ndef build_transforms(\n height,\n width,\n config,\n mask_scale=4,\n transforms='random_flip',\n norm_mean=[0.485, 0.456, 0.406],\n norm_std=[0.229, 0.224, 0.225],\n remove_background_mask=False,\n masks_preprocess = 'none',\n softmax_weight = 0,\n mask_filtering_threshold = 0.3,\n background_computation_strategy = 'threshold',\n **kwargs\n):\n \"\"\"Builds train and test transform functions.\n\n Args:\n height (int): target image height.\n width (int): target image width.\n transforms (str or list of str, optional): transformations applied to model training.\n Default is 'random_flip'.\n norm_mean (list or None, optional): normalization mean values. Default is ImageNet means.\n norm_std (list or None, optional): normalization standard deviation values. Default is\n ImageNet standard deviation values.\n \"\"\"\n if transforms is None:\n transforms = []\n\n if isinstance(transforms, str):\n transforms = [transforms]\n\n if not isinstance(transforms, list):\n raise ValueError(\n 'transforms must be a list of strings, but found to be {}'.format(\n type(transforms)\n )\n )\n\n if len(transforms) > 0:\n transforms = [t.lower() for t in transforms]\n\n if norm_mean is None or norm_std is None:\n norm_mean = [0.485, 0.456, 0.406] # imagenet mean\n norm_std = [0.229, 0.224, 0.225] # imagenet std\n normalize = Normalize(mean=norm_mean, std=norm_std)\n\n print('Building train transforms ...')\n transform_tr = []\n\n print('+ resize to {}x{}'.format(height, width))\n transform_tr += [Resize(height, width)]\n\n if 'random_occlusion' in transforms or 'ro' in transforms:\n print('+ random occlusion')\n transform_tr += [RandomOcclusion(path=config.data.ro.path,\n im_shape=[config.data.height, config.data.width],\n p=config.data.ro.p,\n n=config.data.ro.n,\n min_overlap=config.data.ro.min_overlap,\n max_overlap=config.data.ro.max_overlap,\n )]\n\n if 'random_flip' in transforms or 'rf' in transforms:\n print('+ random flip')\n transform_tr += [HorizontalFlip()]\n\n if 'random_crop' in transforms or 'rc' in transforms:\n print('+ random crop')\n pad_size = 10\n transform_tr += [PadIfNeeded(min_height=height+pad_size*2, min_width=width+pad_size*2, border_mode=cv2.BORDER_CONSTANT, value=0, mask_value=0, p=1),\n RandomCrop(height, width, p=1)]\n\n if 'color_jitter' in transforms or 'cj' in transforms:\n print('+ color jitter')\n transform_tr += [\n ColorJitter(brightness=config.data.cj.brightness,\n contrast=config.data.cj.contrast,\n saturation=config.data.cj.saturation,\n hue=config.data.cj.hue,\n always_apply=config.data.cj.always_apply,\n p=config.data.cj.p,\n )\n ]\n\n print('+ normalization (mean={}, std={})'.format(norm_mean, norm_std))\n transform_tr += [normalize]\n\n if 'random_erase' in transforms or 're' in transforms:\n print('+ random erase')\n transform_tr += [CoarseDropout(min_holes=1, max_holes=1,\n min_height=int(height*0.15), max_height=int(height*0.65),\n min_width=int(width*0.15), max_width=int(width*0.65),\n fill_value=[0.485, 0.456, 0.406], mask_fill_value=0, always_apply=False, p=0.5)]\n\n print('+ to torch tensor of range [0, 1]')\n transform_tr += [ToTensorV2()]\n\n print('Building test transforms ...')\n print('+ resize to {}x{}'.format(height, width))\n print('+ to torch tensor of range [0, 1]')\n print('+ normalization (mean={}, std={})'.format(norm_mean, norm_std))\n\n transform_te = [\n Resize(height, width),\n normalize,\n ToTensorV2()\n ]\n\n transform_tr += [PermuteMasksDim()]\n transform_te += [PermuteMasksDim()]\n\n if remove_background_mask: # ISP masks\n print('+ use remove background mask')\n # remove background before performing other transforms\n transform_tr = [RemoveBackgroundMask()] + transform_tr\n transform_te = [RemoveBackgroundMask()] + transform_te\n\n # Derive background mask from all foreground masks once other tasks have been performed\n print('+ use add background mask')\n transform_tr += [AddBackgroundMask('sum')]\n transform_te += [AddBackgroundMask('sum')]\n else: # Pifpaf confidence based masks\n if masks_preprocess != 'none':\n print('+ masks preprocess = {}'.format(masks_preprocess))\n masks_preprocess_transform = masks_preprocess_all[masks_preprocess]\n transform_tr += [masks_preprocess_transform()]\n transform_te += [masks_preprocess_transform()]\n\n print('+ use add background mask')\n transform_tr += [AddBackgroundMask(background_computation_strategy, softmax_weight, mask_filtering_threshold)]\n transform_te += [AddBackgroundMask(background_computation_strategy, softmax_weight, mask_filtering_threshold)]\n\n transform_tr += [ResizeMasks(height, width, mask_scale)]\n transform_te += [ResizeMasks(height, width, mask_scale)]\n\n transform_tr = Compose(transform_tr, is_check_shapes=False)\n transform_te = Compose(transform_te, is_check_shapes=False)\n\n return transform_tr, transform_te\n" }, { "path": "torchreid/engine/__init__.py", "content": "from __future__ import print_function, absolute_import\n\nfrom .image import ImageSoftmaxEngine, ImageTripletEngine, ImagePartBasedEngine\nfrom .video import VideoSoftmaxEngine, VideoTripletEngine\nfrom .engine import Engine\n" }, { "path": "torchreid/engine/engine.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport os.path as osp\nfrom collections import OrderedDict\n\nimport numpy as np\nimport torch\nfrom torch.nn import functional as F\nfrom torchreid import metrics\nfrom torchreid.data.datasets import get_dataset_nickname\nfrom torchreid.losses import deep_supervision\nfrom torchreid.utils import (\n re_ranking, open_all_layers, save_checkpoint,\n open_specified_layers, visualize_ranked_results, Logger, AverageMeter, perc, plot_pairs_distance_distribution\n)\nfrom torchreid.utils.torchtools import collate\n\n\nclass Engine(object):\n \"\"\"A generic base Engine class for both image- and video-reid.\n\n Args:\n datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager``\n or ``torchreid.data.VideoDataManager``.\n model (nn.Module): model instance.\n optimizer (Optimizer): an Optimizer.\n scheduler (LRScheduler, optional): if None, no learning rate decay will be performed.\n use_gpu (bool, optional): use gpu. Default is True.\n \"\"\"\n\n def __init__(self, config, datamanager, writer, engine_state, use_gpu=True, save_model_flag=False, detailed_ranking=False):\n self.config = config\n self.datamanager = datamanager\n self.train_loader = self.datamanager.train_loader\n self.test_loader = self.datamanager.test_loader\n self.use_gpu = (torch.cuda.is_available() and use_gpu)\n self.save_model_flag = save_model_flag\n self.detailed_ranking = detailed_ranking\n\n self.engine_state = engine_state\n self.writer = writer\n self.logger = Logger.current_logger()\n\n self.model = None\n self.optimizer = None\n self.scheduler = None\n\n self._models = OrderedDict()\n self._optims = OrderedDict()\n self._scheds = OrderedDict()\n\n def register_model(self, name='model', model=None, optim=None, sched=None):\n if self.__dict__.get('_models') is None:\n raise AttributeError(\n 'Cannot assign model before super().__init__() call'\n )\n\n if self.__dict__.get('_optims') is None:\n raise AttributeError(\n 'Cannot assign optim before super().__init__() call'\n )\n\n if self.__dict__.get('_scheds') is None:\n raise AttributeError(\n 'Cannot assign sched before super().__init__() call'\n )\n\n self._models[name] = model\n self._optims[name] = optim\n self._scheds[name] = sched\n\n def get_model_names(self, names=None):\n names_real = list(self._models.keys())\n if names is not None:\n if not isinstance(names, list):\n names = [names]\n for name in names:\n assert name in names_real\n return names\n else:\n return names_real\n\n def save_model(self, epoch, cmc, mAP, ssmd, save_dir, is_best=False):\n if self.save_model_flag:\n names = self.get_model_names()\n\n for name in names:\n save_checkpoint(\n {\n 'state_dict': self._models[name].state_dict(),\n 'epoch': epoch + 1,\n 'rank1': cmc,\n 'mAP': mAP,\n 'ssmd': ssmd,\n 'config': self.config,\n 'optimizer': self._optims[name].state_dict(),\n 'scheduler': self._scheds[name].state_dict()\n },\n osp.join(save_dir, self.writer.model_name + name),\n job_id=self.config.project.job_id,\n is_best=is_best\n )\n\n def set_model_mode(self, mode='train', names=None):\n assert mode in ['train', 'eval', 'test']\n names = self.get_model_names(names)\n\n for name in names:\n if mode == 'train':\n self._models[name].train()\n else:\n self._models[name].eval()\n\n def get_current_lr(self, names=None):\n names = self.get_model_names(names)\n name = names[0]\n return self._optims[name].param_groups[0]['lr']\n\n def update_lr(self, names=None):\n names = self.get_model_names(names)\n\n for name in names:\n if self._scheds[name] is not None:\n self._scheds[name].step()\n self.engine_state.update_lr(self.get_current_lr())\n\n def run(\n self,\n save_dir='log',\n fixbase_epoch=0,\n open_layers=None,\n test_only=False,\n dist_metric='euclidean',\n normalize_feature=False,\n visrank=False,\n visrank_topk=10,\n visrank_q_idx_list=[],\n visrank_count=10,\n use_metric_cuhk03=False,\n ranks=[1, 5, 10, 20],\n rerank=False,\n save_features=False\n ):\n \"\"\"A unified pipeline for training and evaluating a model.\n\n Args:\n save_dir (str): directory to save model.\n max_epoch (int): maximum epoch.\n start_epoch (int, optional): starting epoch. Default is 0.\n fixbase_epoch (int, optional): number of epochs to train ``open_layers`` (new layers)\n while keeping base layers frozen. Default is 0. ``fixbase_epoch`` is counted\n in ``max_epoch``.\n open_layers (str or list, optional): layers (attribute names) open for training.\n start_eval (int, optional): from which epoch to start evaluation. Default is 0.\n eval_freq (int, optional): evaluation frequency. Default is -1 (meaning evaluation\n is only performed at the end of training).\n test_only (bool, optional): if True, only runs evaluation on test datasets.\n Default is False.\n dist_metric (str, optional): distance metric used to compute distance matrix\n between query and gallery. Default is \"euclidean\".\n normalize_feature (bool, optional): performs L2 normalization on feature vectors before\n computing feature distance. Default is False.\n visrank (bool, optional): visualizes ranked results. Default is False. It is recommended to\n enable ``visrank`` when ``test_only`` is True. The ranked images will be saved to\n \"save_dir/visrank_dataset\", e.g. \"save_dir/visrank_market1501\".\n visrank_topk (int, optional): top-k ranked images to be visualized. Default is 10.\n use_metric_cuhk03 (bool, optional): use single-gallery-shot setting for cuhk03.\n Default is False. This should be enabled when using cuhk03 classic split.\n ranks (list, optional): cmc ranks to be computed. Default is [1, 5, 10, 20].\n rerank (bool, optional): uses person re-ranking (by Zhong et al. CVPR'17).\n Default is False. This is only enabled when test_only=True.\n save_features (bool, optional): save test query and test gallery extracted features to disk\n \"\"\"\n\n if test_only:\n self.test(\n 0,\n dist_metric=dist_metric,\n normalize_feature=normalize_feature,\n visrank=visrank,\n visrank_topk=visrank_topk,\n visrank_q_idx_list=visrank_q_idx_list,\n visrank_count=visrank_count,\n save_dir=save_dir,\n use_metric_cuhk03=use_metric_cuhk03,\n ranks=ranks,\n rerank=rerank,\n save_features=save_features\n )\n return\n\n\n self.writer.total_run_timer.start()\n self.engine_state.estimated_num_batches = len(self.train_loader)\n self.engine_state.update_lr(self.get_current_lr())\n print('=> Start training')\n self.engine_state.training_started()\n mAP = 0\n\n for epoch in range(self.engine_state.start_epoch, self.engine_state.max_epoch):\n self.writer.epoch_timer.start()\n self.engine_state.epoch_started()\n\n self.train(\n fixbase_epoch=fixbase_epoch,\n open_layers=open_layers\n )\n\n self.writer.epoch_timer.stop()\n self.engine_state.epoch_completed()\n\n if self.writer.intermediate_evaluate():\n print('=> Intermediate test')\n rank_1, mAP, ssmd = self.test(\n epoch,\n dist_metric=dist_metric,\n normalize_feature=normalize_feature,\n visrank=False,\n visrank_topk=visrank_topk,\n visrank_q_idx_list=visrank_q_idx_list,\n visrank_count=visrank_count,\n save_dir=save_dir,\n use_metric_cuhk03=use_metric_cuhk03,\n ranks=ranks,\n evalate_on_sources_only=True\n )\n self.save_model(epoch, rank_1, mAP, ssmd, save_dir)\n\n self.engine_state.training_completed()\n\n if self.engine_state.max_epoch > 0:\n print('=> Final test')\n rank_1, mAP, ssmd = self.test(\n self.engine_state.epoch,\n dist_metric=dist_metric,\n normalize_feature=normalize_feature,\n visrank=visrank,\n visrank_topk=visrank_topk,\n visrank_q_idx_list=visrank_q_idx_list,\n visrank_count=visrank_count,\n save_dir=save_dir,\n use_metric_cuhk03=use_metric_cuhk03,\n ranks=ranks,\n save_features=save_features,\n evalate_on_sources_only=False\n )\n self.save_model(self.engine_state.epoch, rank_1, mAP, ssmd, save_dir)\n\n self.writer.total_run_timer.stop()\n self.engine_state.run_completed()\n self.logger.close()\n\n return mAP\n\n def train(self, fixbase_epoch=0, open_layers=None):\n self.set_model_mode('train')\n self.logger.add_scalar('Train/lr', self.get_current_lr(), self.engine_state.epoch)\n\n self.two_stepped_transfer_learning(\n self.engine_state.epoch, fixbase_epoch, open_layers\n )\n\n self.writer.data_loading_timer.start()\n for self.batch_idx, data in enumerate(self.train_loader):\n self.writer.data_loading_timer.stop()\n self.writer.batch_timer.start()\n\n loss, loss_summary = self.forward_backward(data)\n self.writer.batch_timer.stop()\n\n self.writer.losses.update(loss_summary)\n self.writer.loss.update(loss)\n self.writer.data_loading_timer.start()\n\n self.engine_state.batch_completed()\n\n self.update_lr()\n\n def forward_backward(self, data):\n raise NotImplementedError\n\n def test(\n self,\n epoch,\n dist_metric='euclidean',\n normalize_feature=False,\n visrank=False,\n visrank_topk=10,\n visrank_q_idx_list=[],\n visrank_count=10,\n save_dir='',\n use_metric_cuhk03=False,\n ranks=[1, 5, 10, 20],\n rerank=False,\n save_features=False,\n evalate_on_sources_only=False\n ):\n \"\"\"Tests model on target datasets.\n\n .. note::\n\n This function has been called in ``run()``.\n\n .. note::\n\n The test pipeline implemented in this function suits both image- and\n video-reid. In general, a subclass of Engine only needs to re-implement\n ``extract_features()`` and ``parse_data_for_eval()`` (most of the time),\n but not a must. Please refer to the source code for more details.\n \"\"\"\n self.writer.test_timer.start()\n\n self.set_model_mode('eval')\n targets = list(self.test_loader.keys())\n if len(targets) == 0:\n raise RuntimeError(\"Test set is either empty or target dataset was not specified.\")\n cmc_avg = AverageMeter()\n mAP_avg = AverageMeter()\n ssmd_avg = AverageMeter()\n pxl_acc_avg = AverageMeter()\n # TODO: capture metrics with Pandas frame (more scalable for new metrics)\n\n cmc_per_dataset = {}\n mAP_per_dataset = {}\n ssmd_per_dataset = {}\n pxl_acc_per_dataset = {}\n for name in targets:\n is_source_dataset = name in self.datamanager.sources\n domain = 'source' if is_source_dataset else 'target'\n if is_source_dataset or not evalate_on_sources_only:\n print('##### Evaluating {} ({}) #####'.format(name, domain))\n query_loader = self.test_loader[name]['query']\n gallery_loader = self.test_loader[name]['gallery']\n cmc, mAP, ssmd, avg_pxl_pred_accuracy = self._evaluate(\n epoch,\n dataset_name=name,\n query_loader=query_loader,\n gallery_loader=gallery_loader,\n dist_metric=dist_metric,\n normalize_feature=normalize_feature,\n visrank=visrank,\n visrank_topk=visrank_topk,\n visrank_q_idx_list=visrank_q_idx_list,\n visrank_count=visrank_count,\n save_dir=save_dir,\n use_metric_cuhk03=use_metric_cuhk03,\n ranks=ranks,\n rerank=rerank,\n save_features=save_features\n )\n dataset_nickname = get_dataset_nickname(name)\n self.writer.report_performance(cmc, mAP, ssmd, avg_pxl_pred_accuracy, dataset_nickname)\n\n cmc_per_dataset[dataset_nickname] = perc(cmc)\n mAP_per_dataset[dataset_nickname] = perc(mAP)\n ssmd_per_dataset[dataset_nickname] = np.around(ssmd, 2)\n pxl_acc_per_dataset[dataset_nickname] = avg_pxl_pred_accuracy\n\n if is_source_dataset:\n cmc_avg.update(cmc)\n mAP_avg.update(mAP)\n ssmd_avg.update(ssmd)\n pxl_acc_avg.update(avg_pxl_pred_accuracy)\n else:\n print('##### Skipping {} ({}) #####'.format(name, domain))\n\n average_score_key = 'avg'\n cmc_per_dataset[average_score_key] = np.array(list(cmc_per_dataset.values())).mean(0)\n # transform dataset->cmc to cmc->dataset\n cmc_per_dataset = [{k: v[i-1] for k, v in cmc_per_dataset.items()} for i in ranks]\n mAP_per_dataset[average_score_key] = np.array(list(mAP_per_dataset.values())).mean()\n ssmd_per_dataset[average_score_key] = np.array(list(ssmd_per_dataset.values())).mean()\n pxl_acc_per_dataset[average_score_key] = np.array(list(pxl_acc_per_dataset.values())).mean()\n\n self.engine_state.test_completed()\n\n self.writer.test_timer.stop()\n\n if mAP_avg.count != 0:\n self.writer.report_performance(cmc_avg.avg, mAP_avg.avg, ssmd_avg.avg, pxl_acc_avg.avg)\n self.writer.report_global_performance(cmc_per_dataset,\n mAP_per_dataset,\n ssmd_per_dataset,\n pxl_acc_per_dataset)\n r1 = cmc_avg.avg[0] if mAP_avg.count != 0 else 0\n return r1, mAP_avg.avg, ssmd_avg.avg\n\n @torch.no_grad()\n def _evaluate(\n self,\n epoch,\n dataset_name='',\n query_loader=None,\n gallery_loader=None,\n dist_metric='euclidean',\n normalize_feature=False,\n visrank=False,\n visrank_topk=10,\n visrank_q_idx_list=[],\n visrank_count=10,\n save_dir='',\n use_metric_cuhk03=False,\n ranks=[1, 5, 10, 20],\n rerank=False,\n save_features=False\n ):\n print('Extracting features from query set:')\n qf, q_pids, q_camids, q_anns = self._feature_extraction(query_loader)\n print('Done, obtained {} tensor'.format(qf.shape))\n\n print('Extracting features from gallery set:')\n gf, g_pids, g_camids, g_anns = self._feature_extraction(gallery_loader)\n print('Done, obtained {} tensor'.format(gf.shape))\n\n print('Test batch feature extraction speed: {:.4f} sec/batch'.format(self.writer.test_batch_timer.avg))\n\n if save_features:\n features_dir = osp.join(save_dir, 'features')\n print('Saving features to : ' + features_dir)\n # TODO create if doesn't exist\n torch.save(gf, osp.join(features_dir, 'gallery_features_' + dataset_name + '.pt'))\n torch.save(qf, osp.join(features_dir, 'query_features_' + dataset_name + '.pt'))\n # save pids, camids and feature length\n\n self.writer.performance_evaluation_timer.start()\n if normalize_feature:\n print('Normalizing features with L2 norm ...')\n qf = self.normalize(qf)\n gf = self.normalize(gf)\n print('Computing distance matrix with metric={} ...'.format(dist_metric))\n distmat = metrics.compute_distance_matrix(qf, gf, dist_metric)\n distmat = distmat.numpy()\n if rerank:\n print('Applying person re-ranking ...')\n distmat_qq = metrics.compute_distance_matrix(qf, qf, dist_metric)\n distmat_gg = metrics.compute_distance_matrix(gf, gf, dist_metric)\n distmat = re_ranking(distmat, distmat_qq, distmat_gg)\n print('Computing CMC and mAP ...')\n eval_metrics = metrics.evaluate_rank(\n distmat,\n q_pids,\n g_pids,\n q_camids,\n g_camids,\n q_anns=q_anns,\n g_anns=g_anns,\n eval_metric=self.datamanager.test_loader[dataset_name]['query'].dataset.eval_metric\n )\n mAP = eval_metrics['mAP']\n cmc = eval_metrics['cmc']\n print('** Results **')\n print('mAP: {:.2%}'.format(mAP))\n print('CMC curve')\n for r in ranks:\n print('Rank-{:<3}: {:.2%}'.format(r, cmc[r - 1]))\n\n for metric in eval_metrics.keys():\n if metric != 'mAP' and metric != 'cmc':\n val, size = eval_metrics[metric]\n if val is not None:\n print('{:<20}: {:.2%} ({})'.format(metric, val, size))\n else:\n print('{:<20}: not provided'.format(metric))\n\n # TODO move below to writer\n print('Evaluate distribution of distances of pairs with same id vs different ids')\n same_ids_dist_mean, same_ids_dist_std, different_ids_dist_mean, different_ids_dist_std, ssmd = \\\n plot_pairs_distance_distribution(distmat, q_pids, g_pids,\n \"Query-gallery\") # TODO separate ssmd from plot, put plot in writer\n print(\"Positive pairs distance distribution mean: {:.3f}\".format(same_ids_dist_mean))\n print(\"Positive pairs distance distribution standard deviation: {:.3f}\".format(same_ids_dist_std))\n print(\"Negative pairs distance distribution mean: {:.3f}\".format(different_ids_dist_mean))\n print(\"Negative pairs distance distribution standard deviation: {:.3f}\".format(\n different_ids_dist_std))\n print(\"SSMD = {:.4f}\".format(ssmd))\n if visrank:\n visualize_ranked_results(\n distmat,\n self.datamanager.fetch_test_loaders(dataset_name),\n self.datamanager.data_type,\n width=self.datamanager.width,\n height=self.datamanager.height,\n save_dir=osp.join(save_dir, 'visrank_' + dataset_name),\n topk=visrank_topk\n )\n self.writer.visualize_embeddings(qf, gf, q_pids, g_pids, self.datamanager.test_loader[dataset_name],\n dataset_name, None, None, mAP, cmc[0])\n self.writer.performance_evaluation_timer.stop()\n return cmc, mAP, ssmd, 0\n\n def _feature_extraction(self, data_loader):\n f_, pids_, camids_ = [], [], []\n anns = []\n for batch_idx, data in enumerate(data_loader):\n imgs, pids, camids = self.parse_data_for_eval(data)\n if self.use_gpu:\n imgs = imgs.cuda()\n self.writer.test_batch_timer.start()\n features = self.extract_features(imgs)\n self.writer.test_batch_timer.stop()\n features = features.data.cpu()\n f_.append(features)\n pids_.extend(pids)\n camids_.extend(camids)\n anns.append(data)\n anns = collate(anns)\n f_ = torch.cat(f_, 0)\n pids_ = np.asarray(pids_)\n camids_ = np.asarray(camids_)\n return f_, pids_, camids_, anns\n\n def compute_loss(self, criterion, outputs, targets, **kwargs):\n if isinstance(outputs, (tuple, list)):\n loss = deep_supervision(criterion, outputs, targets, **kwargs)\n else:\n loss = criterion(outputs, targets, **kwargs)\n return loss\n\n def extract_features(self, input):\n return self.model(input)\n\n def parse_data_for_train(self, data):\n imgs = data['image']\n pids = data['pid']\n return imgs, pids\n\n def parse_data_for_eval(self, data):\n imgs = data['image']\n pids = data['pid']\n camids = data['camid']\n return imgs, pids, camids\n\n def two_stepped_transfer_learning(\n self, epoch, fixbase_epoch, open_layers, model=None\n ):\n \"\"\"Two-stepped transfer learning.\n\n The idea is to freeze base layers for a certain number of epochs\n and then open all layers for training.\n\n Reference: https://arxiv.org/abs/1611.05244\n \"\"\"\n model = self.model if model is None else model\n if model is None:\n return\n\n if fixbase_epoch > 0:\n if (epoch + 1) <= fixbase_epoch and open_layers is not None:\n print(\n '* Only train {} (epoch: {}/{})'.format(\n open_layers, epoch + 1, fixbase_epoch\n )\n )\n open_specified_layers(model, open_layers)\n else:\n open_all_layers(model)\n\n def normalize(self, features):\n return F.normalize(features, p=2, dim=-1)\n" }, { "path": "torchreid/engine/image/__init__.py", "content": "from __future__ import absolute_import\n\nfrom .softmax import ImageSoftmaxEngine\nfrom .triplet import ImageTripletEngine\nfrom .part_based_engine import ImagePartBasedEngine\n" }, { "path": "torchreid/engine/image/part_based_engine.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport os.path as osp\nimport torch\nimport numpy as np\nfrom tabulate import tabulate\nfrom torch import nn\nfrom tqdm import tqdm\n\nfrom ..engine import Engine\nfrom ... import metrics\nfrom ...losses.GiLt_loss import GiLtLoss\nfrom ...losses.body_part_attention_loss import BodyPartAttentionLoss\nfrom ...metrics.distance import compute_distance_matrix_using_bp_features\nfrom ...utils import plot_body_parts_pairs_distance_distribution, \\\n plot_pairs_distance_distribution, re_ranking\nfrom torchreid.utils.constants import *\nfrom ...utils.torchtools import collate\nfrom ...utils.visualization.feature_map_visualization import display_feature_maps\n\n\nclass ImagePartBasedEngine(Engine):\n r\"\"\"Training/testing engine for part-based image-reid.\n \"\"\"\n def __init__(\n self,\n datamanager,\n model,\n optimizer,\n writer,\n loss_name,\n config,\n dist_combine_strat,\n batch_size_pairwise_dist_matrix,\n engine_state,\n margin=0.3,\n scheduler=None,\n use_gpu=True,\n save_model_flag=False,\n mask_filtering_training=False,\n mask_filtering_testing=False\n ):\n super(ImagePartBasedEngine, self).__init__(config,\n datamanager,\n writer,\n engine_state,\n use_gpu=use_gpu,\n save_model_flag=save_model_flag,\n detailed_ranking=config.test.detailed_ranking)\n\n self.model = model\n self.register_model('model', model, optimizer, scheduler)\n self.optimizer = optimizer\n self.scheduler = scheduler\n self.parts_num = self.config.model.bpbreid.masks.parts_num\n self.mask_filtering_training = mask_filtering_training\n self.mask_filtering_testing = mask_filtering_testing\n self.dist_combine_strat = dist_combine_strat\n self.batch_size_pairwise_dist_matrix = batch_size_pairwise_dist_matrix\n self.losses_weights = self.config.loss.part_based.weights\n\n # Losses\n self.GiLt = GiLtLoss(self.losses_weights,\n use_visibility_scores=self.mask_filtering_training,\n triplet_margin=margin,\n loss_name=loss_name,\n writer=self.writer,\n use_gpu=self.use_gpu)\n\n self.body_part_attention_loss = BodyPartAttentionLoss(loss_type=self.config.loss.part_based.ppl, use_gpu=self.use_gpu)\n\n # Timers\n self.feature_extraction_timer = self.writer.feature_extraction_timer\n self.loss_timer = self.writer.loss_timer\n self.optimizer_timer = self.writer.optimizer_timer\n\n def forward_backward(self, data):\n imgs, target_masks, pids, imgs_path = self.parse_data_for_train(data)\n\n # feature extraction\n self.feature_extraction_timer.start()\n embeddings_dict, visibility_scores_dict, id_cls_scores_dict, pixels_cls_scores, spatial_features, masks \\\n = self.model(imgs, external_parts_masks=target_masks)\n display_feature_maps(embeddings_dict, spatial_features, masks[PARTS], imgs_path, pids)\n self.feature_extraction_timer.stop()\n\n # loss\n self.loss_timer.start()\n loss, loss_summary = self.combine_losses(visibility_scores_dict,\n embeddings_dict,\n id_cls_scores_dict,\n pids,\n pixels_cls_scores,\n target_masks,\n bpa_weight=self.losses_weights[PIXELS]['ce'])\n self.loss_timer.stop()\n\n # optimization step\n self.optimizer_timer.start()\n self.optimizer.zero_grad()\n loss.backward()\n self.optimizer.step()\n self.optimizer_timer.stop()\n\n return loss, loss_summary\n\n def combine_losses(self, visibility_scores_dict, embeddings_dict, id_cls_scores_dict, pids, pixels_cls_scores=None, target_masks=None, bpa_weight=0):\n # 1. ReID objective:\n # GiLt loss on holistic and part-based embeddings\n loss, loss_summary = self.GiLt(embeddings_dict, visibility_scores_dict, id_cls_scores_dict, pids)\n\n # 2. Part prediction objective:\n # Body part attention loss on spatial feature map\n if pixels_cls_scores is not None\\\n and target_masks is not None\\\n and bpa_weight > 0:\n # resize external masks to fit feature map size\n target_masks = nn.functional.interpolate(target_masks,\n pixels_cls_scores.shape[2::],\n mode='bilinear',\n align_corners=True)\n # compute target part index for each spatial location, i.e. each spatial location (pixel) value indicate\n # the (body) part that spatial location belong to, or 0 for background.\n pixels_cls_score_targets = target_masks.argmax(dim=1) # [N, Hf, Wf]\n # compute the classification loss for each pixel\n bpa_loss, bpa_loss_summary = self.body_part_attention_loss(pixels_cls_scores, pixels_cls_score_targets)\n loss += bpa_weight * bpa_loss\n loss_summary = {**loss_summary, **bpa_loss_summary}\n\n return loss, loss_summary\n\n def _feature_extraction(self, data_loader):\n f_, pids_, camids_, parts_visibility_, p_masks_, pxl_scores_, anns = [], [], [], [], [], [], []\n for batch_idx, data in enumerate(tqdm(data_loader, desc=f'Batches processed')):\n imgs, masks, pids, camids = self.parse_data_for_eval(data)\n if self.use_gpu:\n if masks is not None:\n masks = masks.cuda()\n imgs = imgs.cuda()\n self.writer.test_batch_timer.start()\n model_output = self.model(imgs, external_parts_masks=masks)\n features, visibility_scores, parts_masks, pixels_cls_scores = self.extract_test_embeddings(model_output)\n self.writer.test_batch_timer.stop()\n if self.mask_filtering_testing:\n parts_visibility = visibility_scores\n parts_visibility = parts_visibility.cpu()\n parts_visibility_.append(parts_visibility)\n else:\n parts_visibility_ = None\n features = features.data.cpu()\n parts_masks = parts_masks.data.cpu()\n f_.append(features)\n p_masks_.append(parts_masks)\n pxl_scores_.append(pixels_cls_scores)\n pids_.extend(pids)\n camids_.extend(camids)\n anns.append(data)\n if self.mask_filtering_testing:\n parts_visibility_ = torch.cat(parts_visibility_, 0)\n f_ = torch.cat(f_, 0)\n p_masks_ = torch.cat(p_masks_, 0)\n pxl_scores_ = torch.cat(pxl_scores_, 0) if pxl_scores_[0] is not None else None\n pids_ = np.asarray(pids_)\n camids_ = np.asarray(camids_)\n anns = collate(anns)\n return f_, pids_, camids_, parts_visibility_, p_masks_, pxl_scores_, anns\n\n @torch.no_grad()\n def _evaluate(\n self,\n epoch,\n dataset_name='',\n query_loader=None,\n gallery_loader=None,\n dist_metric='euclidean',\n normalize_feature=False,\n visrank=False,\n visrank_topk=10,\n visrank_q_idx_list=[],\n visrank_count=10,\n save_dir='',\n use_metric_cuhk03=False,\n ranks=[1, 5, 10, 20],\n rerank=False,\n save_features=False\n ):\n print('Extracting features from query set ...')\n qf, q_pids, q_camids, qf_parts_visibility, q_parts_masks, q_pxl_scores_, q_anns = self._feature_extraction(query_loader)\n print('Done, obtained {} tensor'.format(qf.shape))\n\n print('Extracting features from gallery set ...')\n gf, g_pids, g_camids, gf_parts_visibility, g_parts_masks, g_pxl_scores_, g_anns = self._feature_extraction(gallery_loader)\n print('Done, obtained {} tensor'.format(gf.shape))\n\n print('Test batch feature extraction speed: {:.4f} sec/batch'.format(self.writer.test_batch_timer.avg))\n\n if save_features:\n features_dir = osp.join(save_dir, 'features')\n print('Saving features to : ' + features_dir)\n # TODO create if doesn't exist\n torch.save(gf, osp.join(features_dir, 'gallery_features_' + dataset_name + '.pt'))\n torch.save(qf, osp.join(features_dir, 'query_features_' + dataset_name + '.pt'))\n # save pids, camids and feature length\n\n self.writer.performance_evaluation_timer.start()\n if normalize_feature:\n print('Normalizing features with L2 norm ...')\n qf = self.normalize(qf)\n gf = self.normalize(gf)\n print('Computing distance matrix with metric={} ...'.format(dist_metric))\n distmat, body_parts_distmat = compute_distance_matrix_using_bp_features(qf, gf, qf_parts_visibility,\n gf_parts_visibility,\n self.dist_combine_strat,\n self.batch_size_pairwise_dist_matrix,\n self.use_gpu, dist_metric)\n distmat = distmat.numpy()\n body_parts_distmat = body_parts_distmat.numpy()\n if rerank:\n print('Applying person re-ranking ...')\n distmat_qq, body_parts_distmat_qq = compute_distance_matrix_using_bp_features(qf, qf, qf_parts_visibility, qf_parts_visibility,\n self.dist_combine_strat, self.batch_size_pairwise_dist_matrix,\n self.use_gpu, dist_metric)\n distmat_gg, body_parts_distmat_gg = compute_distance_matrix_using_bp_features(gf, gf, gf_parts_visibility, gf_parts_visibility,\n self.dist_combine_strat, self.batch_size_pairwise_dist_matrix,\n self.use_gpu, dist_metric)\n distmat = re_ranking(distmat, distmat_qq, distmat_gg)\n\n eval_metric = self.datamanager.test_loader[dataset_name]['query'].dataset.eval_metric\n\n print('Computing CMC and mAP ...')\n eval_metrics = metrics.evaluate_rank(\n distmat,\n q_pids,\n g_pids,\n q_camids,\n g_camids,\n q_anns=q_anns,\n g_anns=g_anns,\n eval_metric=eval_metric\n )\n\n mAP = eval_metrics['mAP']\n cmc = eval_metrics['cmc']\n print('** Results **')\n print('mAP: {:.2%}'.format(mAP))\n print('CMC curve')\n for r in ranks:\n print('Rank-{:<3}: {:.2%}'.format(r, cmc[r - 1]))\n\n for metric in eval_metrics.keys():\n if metric != 'mAP' and metric != 'cmc':\n val, size = eval_metrics[metric]\n if val is not None:\n print('{:<20}: {:.2%} ({})'.format(metric, val, size))\n else:\n print('{:<20}: not provided'.format(metric))\n\n # Parts ranking\n if self.detailed_ranking:\n self.display_individual_parts_ranking_performances(body_parts_distmat, cmc, g_camids, g_pids, mAP,\n q_camids, q_pids, eval_metric)\n # TODO move below to writer\n plot_body_parts_pairs_distance_distribution(body_parts_distmat, q_pids, g_pids, \"Query-gallery\")\n print('Evaluate distribution of distances of pairs with same id vs different ids')\n same_ids_dist_mean, same_ids_dist_std, different_ids_dist_mean, different_ids_dist_std, ssmd = \\\n plot_pairs_distance_distribution(distmat, q_pids, g_pids,\n \"Query-gallery\") # TODO separate ssmd from plot, put plot in writer\n print(\"Positive pairs distance distribution mean: {:.3f}\".format(same_ids_dist_mean))\n print(\"Positive pairs distance distribution standard deviation: {:.3f}\".format(same_ids_dist_std))\n print(\"Negative pairs distance distribution mean: {:.3f}\".format(different_ids_dist_mean))\n print(\"Negative pairs distance distribution standard deviation: {:.3f}\".format(\n different_ids_dist_std))\n print(\"SSMD = {:.4f}\".format(ssmd))\n\n # if groundtruth target body masks are provided, compute part prediction accuracy\n avg_pxl_pred_accuracy = 0.0\n if 'mask' in q_anns and 'mask' in g_anns and q_pxl_scores_ is not None and g_pxl_scores_ is not None:\n q_pxl_pred_accuracy = self.compute_pixels_cls_accuracy(torch.from_numpy(q_anns['mask']), q_pxl_scores_)\n g_pxl_pred_accuracy = self.compute_pixels_cls_accuracy(torch.from_numpy(g_anns['mask']), g_pxl_scores_)\n avg_pxl_pred_accuracy = (q_pxl_pred_accuracy * len(q_parts_masks) + g_pxl_pred_accuracy * len(g_parts_masks)) /\\\n (len(q_parts_masks) + len(g_parts_masks))\n print(\"Pixel prediction accuracy for query = {:.2f}% and for gallery = {:.2f}% and on average = {:.2f}%\"\n .format(q_pxl_pred_accuracy, g_pxl_pred_accuracy, avg_pxl_pred_accuracy))\n\n if visrank:\n self.writer.visualize_rank(self.datamanager.test_loader[dataset_name], dataset_name, distmat, save_dir,\n visrank_topk, visrank_q_idx_list, visrank_count,\n body_parts_distmat, qf_parts_visibility, gf_parts_visibility, q_parts_masks,\n g_parts_masks, mAP, cmc[0])\n\n self.writer.visualize_embeddings(qf, gf, q_pids, g_pids, self.datamanager.test_loader[dataset_name],\n dataset_name,\n qf_parts_visibility, gf_parts_visibility, mAP, cmc[0])\n self.writer.performance_evaluation_timer.stop()\n return cmc, mAP, ssmd, avg_pxl_pred_accuracy\n\n def compute_pixels_cls_accuracy(self, target_masks, pixels_cls_scores):\n if pixels_cls_scores.is_cuda:\n target_masks = target_masks.cuda()\n target_masks = nn.functional.interpolate(target_masks, pixels_cls_scores.shape[2::], mode='bilinear',\n align_corners=True) # Best perf with bilinear here and nearest in resize transform\n pixels_cls_score_targets = target_masks.argmax(dim=1) # [N, Hf, Wf]\n pixels_cls_score_targets = pixels_cls_score_targets.flatten() # [N*Hf*Wf]\n pixels_cls_scores = pixels_cls_scores.permute(0, 2, 3, 1).flatten(0, 2) # [N*Hf*Wf, M]\n accuracy = metrics.accuracy(pixels_cls_scores, pixels_cls_score_targets)[0]\n return accuracy.item()\n\n def display_individual_parts_ranking_performances(self, body_parts_distmat, cmc, g_camids, g_pids, mAP, q_camids,\n q_pids, eval_metric):\n print('Parts embeddings individual rankings :')\n bp_offset = 0\n if GLOBAL in self.config.model.bpbreid.test_embeddings:\n bp_offset += 1\n if FOREGROUND in self.config.model.bpbreid.test_embeddings:\n bp_offset += 1\n table = []\n for bp in range(0, body_parts_distmat.shape[0]): # TODO DO NOT TAKE INTO ACCOUNT -1 DISTANCES!!!!\n perf_metrics = metrics.evaluate_rank(\n body_parts_distmat[bp],\n q_pids,\n g_pids,\n q_camids,\n g_camids,\n eval_metric=eval_metric\n )\n title = 'p {}'.format(bp - bp_offset)\n if bp < bp_offset:\n if bp == 0:\n if GLOBAL in self.config.model.bpbreid.test_embeddings:\n title = GLOBAL\n else:\n title = FOREGROUND\n if bp == 1:\n title = FOREGROUND\n mAP = perf_metrics['mAP']\n cmc = perf_metrics['cmc']\n table.append([title, mAP, cmc[0], cmc[4], cmc[9]])\n headers = [\"embed\", \"mAP\", \"R-1\", \"R-5\", \"R-10\"]\n print(tabulate(table, headers, tablefmt=\"fancy_grid\", floatfmt=\".3f\"))\n\n def parse_data_for_train(self, data):\n imgs = data['image']\n imgs_path = data['img_path']\n masks = data['mask'] if 'mask' in data else None\n pids = data['pid']\n\n if self.use_gpu:\n imgs = imgs.cuda()\n if masks is not None:\n masks = masks.cuda()\n pids = pids.cuda()\n\n if masks is not None:\n assert masks.shape[1] == (self.config.model.bpbreid.masks.parts_num + 1)\n\n return imgs, masks, pids, imgs_path\n\n def parse_data_for_eval(self, data):\n imgs = data['image']\n masks = data['mask'] if 'mask' in data else None\n pids = data['pid']\n camids = data['camid']\n return imgs, masks, pids, camids\n\n def extract_test_embeddings(self, model_output):\n embeddings, visibility_scores, id_cls_scores, pixels_cls_scores, spatial_features, parts_masks = model_output\n embeddings_list = []\n visibility_scores_list = []\n embeddings_masks_list = []\n\n for test_emb in self.config.model.bpbreid.test_embeddings:\n embds = embeddings[test_emb]\n embeddings_list.append(embds if len(embds.shape) == 3 else embds.unsqueeze(1))\n if test_emb in bn_correspondants:\n test_emb = bn_correspondants[test_emb]\n vis_scores = visibility_scores[test_emb]\n visibility_scores_list.append(vis_scores if len(vis_scores.shape) == 2 else vis_scores.unsqueeze(1))\n pt_masks = parts_masks[test_emb]\n embeddings_masks_list.append(pt_masks if len(pt_masks.shape) == 4 else pt_masks.unsqueeze(1))\n\n assert len(embeddings) != 0\n\n embeddings = torch.cat(embeddings_list, dim=1) # [N, P+2, D]\n visibility_scores = torch.cat(visibility_scores_list, dim=1) # [N, P+2]\n embeddings_masks = torch.cat(embeddings_masks_list, dim=1) # [N, P+2, Hf, Wf]\n\n return embeddings, visibility_scores, embeddings_masks, pixels_cls_scores\n\n" }, { "path": "torchreid/engine/image/softmax.py", "content": "from __future__ import division, print_function, absolute_import\n\nfrom torchreid import metrics\nfrom torchreid.losses import CrossEntropyLoss\n\nfrom ..engine import Engine\n\n\nclass ImageSoftmaxEngine(Engine):\n r\"\"\"Softmax-loss engine for image-reid.\n\n Args:\n datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager``\n or ``torchreid.data.VideoDataManager``.\n model (nn.Module): model instance.\n optimizer (Optimizer): an Optimizer.\n scheduler (LRScheduler, optional): if None, no learning rate decay will be performed.\n use_gpu (bool, optional): use gpu. Default is True.\n label_smooth (bool, optional): use label smoothing regularizer. Default is True.\n\n Examples::\n \n import torchreid\n datamanager = torchreid.data.ImageDataManager(\n root='path/to/reid-data',\n sources='market1501',\n height=256,\n width=128,\n combineall=False,\n batch_size=32\n )\n model = torchreid.models.build_model(\n name='resnet50',\n num_classes=datamanager.num_train_pids,\n loss='softmax'\n )\n model = model.cuda()\n optimizer = torchreid.optim.build_optimizer(\n model, optim='adam', lr=0.0003\n )\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer,\n lr_scheduler='single_step',\n stepsize=20\n )\n engine = torchreid.engine.ImageSoftmaxEngine(\n datamanager, model, optimizer, scheduler=scheduler\n )\n engine.run(\n max_epoch=60,\n save_dir='log/resnet50-softmax-market1501'\n )\n \"\"\"\n\n def __init__(\n self,\n datamanager,\n model,\n optimizer,\n writer,\n engine_state,\n scheduler=None,\n use_gpu=True,\n label_smooth=True,\n save_model_flag=False\n ):\n super(ImageSoftmaxEngine, self).__init__(datamanager, writer, engine_state, use_gpu, save_model_flag)\n\n self.model = model\n self.optimizer = optimizer\n self.scheduler = scheduler\n self.register_model('model', model, optimizer, scheduler)\n\n self.criterion = CrossEntropyLoss(\n label_smooth=label_smooth\n )\n\n def forward_backward(self, data):\n imgs, pids = self.parse_data_for_train(data)\n\n if self.use_gpu:\n imgs = imgs.cuda()\n pids = pids.cuda()\n\n outputs = self.model(imgs)\n loss = self.compute_loss(self.criterion, outputs, pids)\n\n self.optimizer.zero_grad()\n loss.backward()\n self.optimizer.step()\n\n loss_summary = {\n 'loss': loss.item(),\n 'acc': metrics.accuracy(outputs, pids)[0].item()\n }\n\n return loss, {'glb_ft': loss_summary}\n" }, { "path": "torchreid/engine/image/triplet.py", "content": "from __future__ import division, print_function, absolute_import\n\nfrom torchreid import metrics\nfrom torchreid.losses import TripletLoss, CrossEntropyLoss\n\nfrom ..engine import Engine\n\n\nclass ImageTripletEngine(Engine):\n r\"\"\"Triplet-loss engine for image-reid.\n\n Args:\n datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager``\n or ``torchreid.data.VideoDataManager``.\n model (nn.Module): model instance.\n optimizer (Optimizer): an Optimizer.\n margin (float, optional): margin for triplet loss. Default is 0.3.\n weight_t (float, optional): weight for triplet loss. Default is 1.\n weight_x (float, optional): weight for softmax loss. Default is 1.\n scheduler (LRScheduler, optional): if None, no learning rate decay will be performed.\n use_gpu (bool, optional): use gpu. Default is True.\n label_smooth (bool, optional): use label smoothing regularizer. Default is True.\n\n Examples::\n \n import torchreid\n datamanager = torchreid.data.ImageDataManager(\n root='path/to/reid-data',\n sources='market1501',\n height=256,\n width=128,\n combineall=False,\n batch_size=32,\n num_instances=4,\n train_sampler='RandomIdentitySampler' # this is important\n )\n model = torchreid.models.build_model(\n name='resnet50',\n num_classes=datamanager.num_train_pids,\n loss='triplet'\n )\n model = model.cuda()\n optimizer = torchreid.optim.build_optimizer(\n model, optim='adam', lr=0.0003\n )\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer,\n lr_scheduler='single_step',\n stepsize=20\n )\n engine = torchreid.engine.ImageTripletEngine(\n datamanager, model, optimizer, margin=0.3,\n weight_t=0.7, weight_x=1, scheduler=scheduler\n )\n engine.run(\n max_epoch=60,\n save_dir='log/resnet50-triplet-market1501'\n )\n \"\"\"\n\n def __init__(\n self,\n datamanager,\n model,\n optimizer,\n writer,\n engine_state,\n margin=0.3,\n weight_t=1,\n weight_x=1,\n scheduler=None,\n use_gpu=True,\n label_smooth=True,\n save_model_flag=False\n ):\n super(ImageTripletEngine, self).__init__(datamanager, writer, engine_state, use_gpu, save_model_flag)\n\n self.model = model\n self.optimizer = optimizer\n self.scheduler = scheduler\n self.register_model('model', model, optimizer, scheduler)\n\n self.weight_t = weight_t\n self.weight_x = weight_x\n\n self.criterion_t = TripletLoss(margin=margin)\n self.criterion_x = CrossEntropyLoss(\n label_smooth=label_smooth\n )\n\n def forward_backward(self, data):\n imgs, pids = self.parse_data_for_train(data)\n\n if self.use_gpu:\n imgs = imgs.cuda()\n pids = pids.cuda()\n\n outputs, features = self.model(imgs)\n loss_t = self.compute_loss(self.criterion_t, features, pids)\n loss_x = self.compute_loss(self.criterion_x, outputs, pids)\n loss = self.weight_t * loss_t + self.weight_x * loss_x\n\n self.optimizer.zero_grad()\n loss.backward()\n self.optimizer.step()\n\n loss_summary = {\n 'loss_t': loss_t.item(),\n 'loss_x': loss_x.item(),\n 'acc': metrics.accuracy(outputs, pids)[0].item()\n }\n\n return loss, {'glb_ft': loss_summary}\n" }, { "path": "torchreid/engine/video/__init__.py", "content": "from __future__ import absolute_import\n\nfrom .softmax import VideoSoftmaxEngine\nfrom .triplet import VideoTripletEngine\n" }, { "path": "torchreid/engine/video/softmax.py", "content": "from __future__ import division, print_function, absolute_import\nimport torch\n\nfrom torchreid.engine.image import ImageSoftmaxEngine\n\n\nclass VideoSoftmaxEngine(ImageSoftmaxEngine):\n \"\"\"Softmax-loss engine for video-reid.\n\n Args:\n datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager``\n or ``torchreid.data.VideoDataManager``.\n model (nn.Module): model instance.\n optimizer (Optimizer): an Optimizer.\n scheduler (LRScheduler, optional): if None, no learning rate decay will be performed.\n use_gpu (bool, optional): use gpu. Default is True.\n label_smooth (bool, optional): use label smoothing regularizer. Default is True.\n pooling_method (str, optional): how to pool features for a tracklet.\n Default is \"avg\" (average). Choices are [\"avg\", \"max\"].\n\n Examples::\n \n import torch\n import torchreid\n # Each batch contains batch_size*seq_len images\n datamanager = torchreid.data.VideoDataManager(\n root='path/to/reid-data',\n sources='mars',\n height=256,\n width=128,\n combineall=False,\n batch_size=8, # number of tracklets\n seq_len=15 # number of images in each tracklet\n )\n model = torchreid.models.build_model(\n name='resnet50',\n num_classes=datamanager.num_train_pids,\n loss='softmax'\n )\n model = model.cuda()\n optimizer = torchreid.optim.build_optimizer(\n model, optim='adam', lr=0.0003\n )\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer,\n lr_scheduler='single_step',\n stepsize=20\n )\n engine = torchreid.engine.VideoSoftmaxEngine(\n datamanager, model, optimizer, scheduler=scheduler,\n pooling_method='avg'\n )\n engine.run(\n max_epoch=60,\n save_dir='log/resnet50-softmax-mars'\n )\n \"\"\"\n\n def __init__(\n self,\n datamanager,\n model,\n optimizer,\n writer,\n scheduler=None,\n use_gpu=True,\n label_smooth=True,\n pooling_method='avg',\n save_model_flag=False\n ):\n super(VideoSoftmaxEngine, self).__init__(\n datamanager,\n model,\n optimizer,\n writer,\n scheduler=scheduler,\n use_gpu=use_gpu,\n label_smooth=label_smooth,\n save_model_flag=save_model_flag\n )\n self.pooling_method = pooling_method\n\n def parse_data_for_train(self, data):\n imgs = data[0]\n pids = data[1]\n if imgs.dim() == 5:\n # b: batch size\n # s: sequence length\n # c: channel depth\n # h: height\n # w: width\n b, s, c, h, w = imgs.size()\n imgs = imgs.view(b * s, c, h, w)\n pids = pids.view(b, 1).expand(b, s)\n pids = pids.contiguous().view(b * s)\n return imgs, pids\n\n def extract_features(self, input):\n # b: batch size\n # s: sequence length\n # c: channel depth\n # h: height\n # w: width\n b, s, c, h, w = input.size()\n input = input.view(b * s, c, h, w)\n features = self.model(input)\n features = features.view(b, s, -1)\n if self.pooling_method == 'avg':\n features = torch.mean(features, 1)\n else:\n features = torch.max(features, 1)[0]\n return features\n" }, { "path": "torchreid/engine/video/triplet.py", "content": "from __future__ import division, print_function, absolute_import\n\nfrom torchreid.engine.image import ImageTripletEngine\nfrom torchreid.engine.video import VideoSoftmaxEngine\n\n\nclass VideoTripletEngine(ImageTripletEngine, VideoSoftmaxEngine):\n \"\"\"Triplet-loss engine for video-reid.\n\n Args:\n datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager``\n or ``torchreid.data.VideoDataManager``.\n model (nn.Module): model instance.\n optimizer (Optimizer): an Optimizer.\n margin (float, optional): margin for triplet loss. Default is 0.3.\n weight_t (float, optional): weight for triplet loss. Default is 1.\n weight_x (float, optional): weight for softmax loss. Default is 1.\n scheduler (LRScheduler, optional): if None, no learning rate decay will be performed.\n use_gpu (bool, optional): use gpu. Default is True.\n label_smooth (bool, optional): use label smoothing regularizer. Default is True.\n pooling_method (str, optional): how to pool features for a tracklet.\n Default is \"avg\" (average). Choices are [\"avg\", \"max\"].\n\n Examples::\n\n import torch\n import torchreid\n # Each batch contains batch_size*seq_len images\n # Each identity is sampled with num_instances tracklets\n datamanager = torchreid.data.VideoDataManager(\n root='path/to/reid-data',\n sources='mars',\n height=256,\n width=128,\n combineall=False,\n num_instances=4,\n train_sampler='RandomIdentitySampler'\n batch_size=8, # number of tracklets\n seq_len=15 # number of images in each tracklet\n )\n model = torchreid.models.build_model(\n name='resnet50',\n num_classes=datamanager.num_train_pids,\n loss='triplet'\n )\n model = model.cuda()\n optimizer = torchreid.optim.build_optimizer(\n model, optim='adam', lr=0.0003\n )\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer,\n lr_scheduler='single_step',\n stepsize=20\n )\n engine = torchreid.engine.VideoTripletEngine(\n datamanager, model, optimizer, margin=0.3,\n weight_t=0.7, weight_x=1, scheduler=scheduler,\n pooling_method='avg'\n )\n engine.run(\n max_epoch=60,\n save_dir='log/resnet50-triplet-mars'\n )\n \"\"\"\n\n def __init__(\n self,\n datamanager,\n model,\n optimizer,\n writer,\n margin=0.3,\n weight_t=1,\n weight_x=1,\n scheduler=None,\n use_gpu=True,\n label_smooth=True,\n pooling_method='avg',\n save_model_flag=False\n ):\n super(VideoTripletEngine, self).__init__(\n datamanager,\n model,\n optimizer,\n writer,\n margin=margin,\n weight_t=weight_t,\n weight_x=weight_x,\n scheduler=scheduler,\n use_gpu=use_gpu,\n label_smooth=label_smooth,\n save_model_flag=save_model_flag\n )\n self.pooling_method = pooling_method\n" }, { "path": "torchreid/hyperparameter/custom_hyperparameter_optimizer.py", "content": "import os\nimport sys\nimport time\nimport os.path as osp\nimport argparse\nimport torch\nimport torch.nn as nn\n\nimport torchreid\nfrom torchreid.utils import (\n Logger, check_isfile, set_random_seed, collect_env_info,\n resume_from_checkpoint, load_pretrained_weights, compute_model_complexity, Writer\n)\n\nfrom scripts.default_config import (\n imagedata_kwargs, optimizer_kwargs, videodata_kwargs, engine_run_kwargs,\n get_default_config, lr_scheduler_kwargs\n)\n\n\ndef build_datamanager(cfg):\n if cfg.data.type == 'image':\n return torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))\n else:\n return torchreid.data.VideoDataManager(**videodata_kwargs(cfg))\n\n\ndef build_engine(cfg, datamanager, model, optimizer, scheduler, writer):\n if cfg.data.type == 'image':\n if cfg.loss.name == 'softmax':\n engine = torchreid.engine.ImageSoftmaxEngine(\n datamanager,\n model,\n optimizer=optimizer,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer\n )\n\n elif cfg.loss.name == 'triplet':\n engine = torchreid.engine.ImageTripletEngine(\n datamanager,\n model,\n optimizer=optimizer,\n margin=cfg.loss.triplet.margin,\n weight_t=cfg.loss.triplet.weight_t,\n weight_x=cfg.loss.triplet.weight_x,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer\n )\n\n elif cfg.loss.name == 'part_based':\n engine = torchreid.engine.ImagePartBasedEngine(\n datamanager,\n model,\n optimizer=optimizer,\n loss_name=cfg.loss.part_based.name,\n config=cfg,\n margin=cfg.loss.triplet.margin,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer,\n mask_filtering_training=cfg.model.bpbreid.mask_filtering_training,\n mask_filtering_testing=cfg.model.bpbreid.mask_filtering_testing,\n mask_filtering_threshold=cfg.model.bpbreid.mask_filtering_threshold,\n batch_debug_freq=cfg.train.batch_debug_freq,\n batch_size_pairwise_dist_matrix=cfg.test.batch_size_pairwise_dist_matrix\n )\n\n else:\n if cfg.loss.name == 'softmax':\n engine = torchreid.engine.VideoSoftmaxEngine(\n datamanager,\n model,\n optimizer=optimizer,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n pooling_method=cfg.video.pooling_method,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer\n )\n\n else:\n engine = torchreid.engine.VideoTripletEngine(\n datamanager,\n model,\n optimizer=optimizer,\n margin=cfg.loss.triplet.margin,\n weight_t=cfg.loss.triplet.weight_t,\n weight_x=cfg.loss.triplet.weight_x,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer\n )\n\n return engine\n\n\ndef reset_config(cfg, args):\n if args.root:\n cfg.data.root = args.root\n if args.sources:\n cfg.data.sources = args.sources\n if args.targets:\n cfg.data.targets = args.targets\n if args.transforms:\n cfg.data.transforms = args.transforms\n\n\ndef main():\n parser = argparse.ArgumentParser(\n formatter_class=argparse.ArgumentDefaultsHelpFormatter\n )\n parser.add_argument(\n '--config-file', type=str, default='', help='path to config file'\n )\n parser.add_argument(\n '-s',\n '--sources',\n type=str,\n nargs='+',\n help='source datasets (delimited by space)'\n )\n parser.add_argument(\n '-t',\n '--targets',\n type=str,\n nargs='+',\n help='target datasets (delimited by space)'\n )\n parser.add_argument(\n '--transforms', type=str, nargs='+', help='data augmentation'\n )\n parser.add_argument(\n '--root', type=str, default='', help='path to data root'\n )\n parser.add_argument(\n 'opts',\n default=None,\n nargs=argparse.REMAINDER,\n help='Modify config options using the command-line'\n )\n args = parser.parse_args()\n\n cfg = get_default_config()\n cfg.use_gpu = torch.cuda.is_available()\n if args.config_file:\n cfg.merge_from_file(args.config_file)\n cfg.project.config_file = os.path.basename(args.config_file)\n reset_config(cfg, args)\n cfg.merge_from_list(args.opts)\n\n params = []\n for pooling in cfg.tuning.model.bpbreid.pooling:\n for normalization in cfg.tuning.model.bpbreid.normalization:\n for mask_filtering_training in cfg.tuning.model.bpbreid.mask_filtering_training:\n for mask_filtering_testing in cfg.tuning.model.bpbreid.mask_filtering_testing:\n for mask_filtering_threshold in cfg.tuning.model.bpbreid.mask_filtering_threshold:\n for transforms in cfg.tuning.data.transforms:\n for pose_name in cfg.tuning.loss.part_based.name:\n params.append([pooling, normalization, mask_filtering_training, mask_filtering_testing, mask_filtering_threshold, transforms, pose_name])\n\n print(\"Will optimize #{} params combination\".format(len(params)))\n for count, param in enumerate(params):\n cfg.model.bpbreid.pooling = param[0]\n cfg.model.bpbreid.normalization = param[1]\n cfg.model.bpbreid.mask_filtering_training = param[2]\n cfg.model.bpbreid.mask_filtering_testing = param[3]\n cfg.model.bpbreid.mask_filtering_threshold = param[4]\n cfg.data.transforms = param[5]\n cfg.loss.part_based.name = param[6]\n\n if cfg.project.debug_mode:\n torch.autograd.set_detect_anomaly(True)\n\n writer = Writer(cfg)\n\n print(\"Hyper param tuning {}/{}\".format(count+1, len(param)))\n print(\"Hyper param = {}\".format(params))\n\n set_random_seed(cfg.train.seed)\n\n log_name = 'test_log' if cfg.test.evaluate else 'train_log'\n log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')\n log_name += '.txt'\n sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))\n\n print('Show configuration\\n{}\\n'.format(cfg))\n print('Collecting env info ...')\n print('** System info **\\n{}\\n'.format(collect_env_info()))\n\n if cfg.use_gpu:\n torch.backends.cudnn.benchmark = True\n\n datamanager = build_datamanager(cfg)\n\n print('Building model: {}'.format(cfg.model.name))\n model = torchreid.models.build_model(\n name=cfg.model.name,\n num_classes=datamanager.num_train_pids,\n loss=cfg.loss.name,\n pretrained=cfg.model.pretrained,\n use_gpu=cfg.use_gpu,\n pooling=cfg.model.bpbreid.pooling,\n normalization=cfg.model.bpbreid.normalization\n )\n num_params, flops = compute_model_complexity(\n model, (1, 3, cfg.data.height, cfg.data.width)\n )\n print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))\n\n if cfg.model.load_weights and check_isfile(cfg.model.load_weights):\n load_pretrained_weights(model, cfg.model.load_weights)\n\n if cfg.use_gpu:\n model = nn.DataParallel(model).cuda()\n\n optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer, **lr_scheduler_kwargs(cfg)\n )\n\n if cfg.model.resume and check_isfile(cfg.model.resume):\n cfg.train.start_epoch = resume_from_checkpoint(\n cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler\n )\n\n print(\n 'Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)\n )\n\n engine = build_engine(cfg, datamanager, model, optimizer, scheduler, writer)\n engine.run(**engine_run_kwargs(cfg))\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "torchreid/hyperparameter/hyperparameter_optimizer.py", "content": "import logging\n\nfrom clearml import Task\nfrom clearml.automation import (HyperParameterOptimizer, GridSearch, UniformIntegerParameterRange,\n DiscreteParameterRange, ParameterSet)\n\n\ndef job_complete_callback(\n job_id, # type: str\n objective_value, # type: float\n objective_iteration, # type: int\n job_parameters, # type: dict\n top_performance_job_id # type: str\n):\n print('Job completed!', job_id, objective_value, objective_iteration, job_parameters)\n if job_id == top_performance_job_id:\n print('WOOT WOOT we broke the record! Objective reached {}'.format(objective_value))\n\n\n# Connecting clearml\ntask = Task.init(project_name='Hyper-Parameter Optimization',\n task_name='Automatic Hyper-Parameter Optimization',\n task_type=Task.TaskTypes.optimizer,\n reuse_last_task_id=False)\n\n# experiment template to optimize in the hyper-parameter optimization\nargs = {\n 'template_task_id': None,\n 'run_as_service': False,\n}\nargs = task.connect(args)\n\n# Get the template task experiment that we want to optimize\nif not args['template_task_id']:\n args['template_task_id'] = '71379629c09449b8b3160d881c58657a'\n\n# Example use case:\nan_optimizer = HyperParameterOptimizer(\n # This is the experiment we want to optimize\n base_task_id=args['template_task_id'],\n # here we define the hyper-parameters to optimize\n hyper_parameters=[\n DiscreteParameterRange('General/loss/part_based/name', ['inter_parts_triplet_loss', 'part_max_triplet_loss', 'part_averaged_triplet_loss', 'intra_parts_triplet_loss']),\n DiscreteParameterRange('General/train/batch_size', values=[32, 128])\n ],\n # this is the objective metric we want to maximize/minimize\n objective_metric_title='Test/rank1',\n objective_metric_series='Test/rank1',\n # now we decide if we want to maximize it or minimize it (accuracy we maximize)\n objective_metric_sign='max',\n # let us limit the number of concurrent experiments,\n # this in turn will make sure we do dont bombard the scheduler with experiments.\n # if we have an auto-scaler connected, this, by proxy, will limit the number of machine\n max_number_of_concurrent_tasks=1,\n # this is the optimizer class (actually doing the optimization)\n # Currently, we can choose from GridSearch, RandomSearch or OptimizerBOHB (Bayesian optimization Hyper-Band)\n # more are coming soon...\n optimizer_class=GridSearch,\n # Select an execution queue to schedule the experiments for execution\n execution_queue='default',\n # Optional: Limit the execution time of a single experiment, in minutes.\n # (this is optional, and if using OptimizerBOHB, it is ignored)\n time_limit_per_job=None,\n # Check the experiments every 12 seconds is way too often, we should probably set it to 5 min,\n # assuming a single experiment is usually hours...\n pool_period_min=5,\n # set the maximum number of jobs to launch for the optimization, default (None) unlimited\n # If OptimizerBOHB is used, it defined the maximum budget in terms of full jobs\n # basically the cumulative number of iterations will not exceed total_max_jobs * max_iteration_per_job\n total_max_jobs=10,\n # set the minimum number of iterations for an experiment, before early stopping.\n # Does not apply for simple strategies such as RandomSearch or GridSearch\n min_iteration_per_job=10,\n # Set the maximum number of iterations for an experiment to execute\n # (This is optional, unless using OptimizerBOHB where this is a must)\n max_iteration_per_job=None,\n)\n\n# if we are running as a service, just enqueue ourselves into the services queue and let it run the optimization\nif args['run_as_service']:\n # if this code is executed by `clearml-agent` the function call does nothing.\n # if executed locally, the local process will be terminated, and a remote copy will be executed instead\n task.execute_remotely(queue_name='services', exit_process=True)\n\n# report every 12 seconds, this is way too often, but we are testing here J\nan_optimizer.set_report_period(5)\n# start the optimization process, callback function to be called every time an experiment is completed\n# this function returns immediately\nan_optimizer.start(job_complete_callback=job_complete_callback)\n# set the time limit for the optimization process (2 hours)\nan_optimizer.set_time_limit(in_minutes=60.0 * 24)\n# wait until process is done (notice we are controlling the optimization process in the background)\nan_optimizer.wait()\n# optimization is completed, print the top performing experiments id\ntop_exp = an_optimizer.get_top_experiments(top_k=3)\nprint([t.id for t in top_exp])\n# make sure background optimization stopped\nan_optimizer.stop()\n\nprint('We are done, good bye')\n" }, { "path": "torchreid/hyperparameter/optuna_hyperparameter_optimizer.py", "content": "# import os\n# import sys\n# import time\n# import os.path as osp\n# import argparse\n#\n# import cv2\n# import optuna\n# import torch\n# import torch.nn as nn\n# from optuna import Trial\n# from optuna.samplers import GridSampler\n#\n# import torchreid\n# from torchreid.utils import (\n# Logger, check_isfile, set_random_seed, collect_env_info,\n# resume_from_checkpoint, load_pretrained_weights, compute_model_complexity, Writer\n# )\n#\n# from scripts.default_config import (\n# imagedata_kwargs, optimizer_kwargs, videodata_kwargs, engine_run_kwargs,\n# get_default_config, lr_scheduler_kwargs\n# )\n#\n#\n# def build_datamanager(cfg):\n# if cfg.data.type == 'image':\n# return torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))\n# else:\n# return torchreid.data.VideoDataManager(**videodata_kwargs(cfg))\n#\n#\n# def build_engine(cfg, datamanager, model, optimizer, scheduler, writer):\n# if cfg.data.type == 'image':\n# if cfg.loss.name == 'softmax':\n# engine = torchreid.engine.ImageSoftmaxEngine(\n# datamanager,\n# model,\n# optimizer=optimizer,\n# scheduler=scheduler,\n# use_gpu=cfg.use_gpu,\n# label_smooth=cfg.loss.softmax.label_smooth,\n# save_model_flag=cfg.model.save_model_flag,\n# writer=writer\n# )\n#\n# elif cfg.loss.name == 'triplet':\n# engine = torchreid.engine.ImageTripletEngine(\n# datamanager,\n# model,\n# optimizer=optimizer,\n# margin=cfg.loss.triplet.margin,\n# weight_t=cfg.loss.triplet.weight_t,\n# weight_x=cfg.loss.triplet.weight_x,\n# scheduler=scheduler,\n# use_gpu=cfg.use_gpu,\n# label_smooth=cfg.loss.softmax.label_smooth,\n# save_model_flag=cfg.model.save_model_flag,\n# writer=writer\n# )\n#\n# elif cfg.loss.name == 'part_based':\n# engine = torchreid.engine.ImagePartBasedEngine(\n# datamanager,\n# model,\n# optimizer=optimizer,\n# loss_name=cfg.loss.part_based.name,\n# config=cfg,\n# margin=cfg.loss.triplet.margin,\n# scheduler=scheduler,\n# use_gpu=cfg.use_gpu,\n# save_model_flag=cfg.model.save_model_flag,\n# writer=writer,\n# mask_filtering_training=cfg.model.bpbreid.mask_filtering_training,\n# mask_filtering_testing=cfg.model.bpbreid.mask_filtering_testing,\n# mask_filtering_threshold=cfg.model.bpbreid.mask_filtering_threshold,\n# batch_debug_freq=cfg.train.batch_debug_freq,\n# batch_size_pairwise_dist_matrix=cfg.test.batch_size_pairwise_dist_matrix\n# )\n#\n# else:\n# if cfg.loss.name == 'softmax':\n# engine = torchreid.engine.VideoSoftmaxEngine(\n# datamanager,\n# model,\n# optimizer=optimizer,\n# scheduler=scheduler,\n# use_gpu=cfg.use_gpu,\n# label_smooth=cfg.loss.softmax.label_smooth,\n# pooling_method=cfg.video.pooling_method,\n# save_model_flag=cfg.model.save_model_flag,\n# writer=writer\n# )\n#\n# else:\n# engine = torchreid.engine.VideoTripletEngine(\n# datamanager,\n# model,\n# optimizer=optimizer,\n# margin=cfg.loss.triplet.margin,\n# weight_t=cfg.loss.triplet.weight_t,\n# weight_x=cfg.loss.triplet.weight_x,\n# scheduler=scheduler,\n# use_gpu=cfg.use_gpu,\n# label_smooth=cfg.loss.softmax.label_smooth,\n# save_model_flag=cfg.model.save_model_flag,\n# writer=writer\n# )\n#\n# return engine\n#\n#\n# def reset_config(cfg, args):\n# if args.root:\n# cfg.data.root = args.root\n# if args.sources:\n# cfg.data.sources = args.sources\n# if args.targets:\n# cfg.data.targets = args.targets\n# if args.transforms:\n# cfg.data.transforms = args.transforms\n#\n#\n# def merge_optuna_hyperparams(cfg, trial):\n# # cfg.data.sources\n# # cfg.data.targets\n# # cfg.sampler.num_instances\n# # cfg.train.optim\n# # cfg.train.lr\n# # cfg.train.weight_decay\n# # cfg.train.max_epoch\n# # cfg.train.start_epoch\n# # cfg.train.batch_size\n# # cfg.train.fixbase_epoch\n# # cfg.train.open_layers\n# # cfg.train.staged_lr\n# # cfg.train.new_layers\n# # cfg.train.base_lr_mult\n# # cfg.train.lr_scheduler\n# # cfg.train.stepsize\n# # cfg.train.gamma\n# # cfg.train.seed\n# # cfg.train.eval_freq\n# # cfg.train.batch_debug_freq\n# # cfg.sgd.momentum\n# # cfg.loss.triplet.margin\n# # cfg.sgd.dampening\n# # cfg.sgd.nesterov\n# # cfg.rmsprop.alpha\n# # cfg.adam.beta1\n# # cfg.adam.beta2\n#\n# # Categorical parameter\n# # optimizer = trial.suggest_categorical('optimizer', ['MomentumSGD', 'Adam'])\n#\n# # Int parameter\n# # num_layers = trial.suggest_int('num_layers', 1, 3)\n#\n# # Uniform parameter\n# # dropout_rate = trial.suggest_uniform('dropout_rate', 0.0, 1.0)\n#\n# # Loguniform parameter\n# # learning_rate = trial.suggest_loguniform('learning_rate', 1e-5, 1e-2)\n#\n# # Discrete-uniform parameter\n# # drop_path_rate = trial.suggest_discrete_uniform('drop_path_rate', 0.0, 1.0, 0.1)\n#\n# if len(cfg.tuning.model.bpbreid.pooling) != 0:\n# cfg.model.bpbreid.pooling = trial.suggest_categorical(\"model.bpbreid.pooling\", cfg.tuning.model.bpbreid.pooling)\n# if len(cfg.tuning.model.bpbreid.normalization) != 0:\n# cfg.model.bpbreid.normalization = trial.suggest_categorical(\"model.bpbreid.normalization\", cfg.tuning.model.bpbreid.normalization)\n# if len(cfg.tuning.model.bpbreid.mask_filtering_training) != 0:\n# cfg.model.bpbreid.mask_filtering_training = trial.suggest_categorical(\"model.bpbreid.mask_filtering_training\", cfg.tuning.model.bpbreid.mask_filtering_training)\n# if len(cfg.tuning.model.bpbreid.mask_filtering_testing) != 0:\n# cfg.model.bpbreid.mask_filtering_testing = trial.suggest_categorical(\"model.bpbreid.mask_filtering_testing\", cfg.tuning.model.bpbreid.mask_filtering_testing)\n# if len(cfg.tuning.model.bpbreid.mask_filtering_threshold) != 0:\n# cfg.model.bpbreid.mask_filtering_threshold = trial.suggest_categorical(\"model.bpbreid.mask_filtering_threshold\", cfg.tuning.model.bpbreid.mask_filtering_threshold)\n# if len(cfg.tuning.loss.part_based.name) != 0:\n# cfg.loss.part_based.name = trial.suggest_categorical(\"loss.part_based.name\", cfg.tuning.loss.part_based.name)\n#\n#\n# def main():\n# parser = argparse.ArgumentParser(\n# formatter_class=argparse.ArgumentDefaultsHelpFormatter\n# )\n# parser.add_argument(\n# '--config-file', type=str, default='', help='path to config file'\n# )\n# parser.add_argument(\n# '-s',\n# '--sources',\n# type=str,\n# nargs='+',\n# help='source datasets (delimited by space)'\n# )\n# parser.add_argument(\n# '-t',\n# '--targets',\n# type=str,\n# nargs='+',\n# help='target datasets (delimited by space)'\n# )\n# parser.add_argument(\n# '--transforms', type=str, nargs='+', help='data augmentation'\n# )\n# parser.add_argument(\n# '--root', type=str, default='', help='path to data root'\n# )\n# parser.add_argument(\n# 'opts',\n# default=None,\n# nargs=argparse.REMAINDER,\n# help='Modify config options using the command-line'\n# )\n# args = parser.parse_args()\n#\n# cfg = get_default_config()\n# cfg.use_gpu = torch.cuda.is_available()\n# if args.config_file:\n# cfg.merge_from_file(args.config_file)\n# cfg.project.config_file = os.path.basename(args.config_file)\n# reset_config(cfg, args)\n# cfg.merge_from_list(args.opts)\n#\n# # Create objective function with access to config from main context\n# def objective(trial: Trial):\n# # import trial hyper parameters into corresponding config field\n# merge_optuna_hyperparams(cfg, trial)\n#\n# if cfg.project.debug_mode:\n# torch.autograd.set_detect_anomaly(True)\n# writer = Writer(cfg)\n# set_random_seed(cfg.train.seed)\n# log_name = 'test_log' if cfg.test.evaluate else 'train_log'\n# log_name += time.strftime('-%Y-%m-%d-%H-%M-%S')\n# log_name += '.txt'\n# sys.stdout = Logger(osp.join(cfg.data.save_dir, log_name))\n# print('Show configuration\\n{}\\n'.format(cfg))\n# print('Collecting env info ...')\n# print('** System info **\\n{}\\n'.format(collect_env_info()))\n# if cfg.use_gpu:\n# torch.backends.cudnn.benchmark = True\n# datamanager = build_datamanager(cfg)\n# print('Building model: {}'.format(cfg.model.name))\n# model = torchreid.models.build_model(\n# name=cfg.model.name,\n# num_classes=datamanager.num_train_pids,\n# loss=cfg.loss.name,\n# pretrained=cfg.model.pretrained,\n# use_gpu=cfg.use_gpu,\n# pooling=cfg.model.bpbreid.pooling,\n# normalization=cfg.model.bpbreid.normalization\n# )\n# num_params, flops = compute_model_complexity(\n# model, (1, 3, cfg.data.height, cfg.data.width)\n# )\n# print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))\n# if cfg.model.load_weights and check_isfile(cfg.model.load_weights):\n# load_pretrained_weights(model, cfg.model.load_weights)\n# if cfg.use_gpu:\n# model = nn.DataParallel(model).cuda()\n# optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))\n# scheduler = torchreid.optim.build_lr_scheduler(\n# optimizer, **lr_scheduler_kwargs(cfg)\n# )\n# if cfg.model.resume and check_isfile(cfg.model.resume):\n# cfg.train.start_epoch = resume_from_checkpoint(\n# cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler\n# )\n# print(\n# 'Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)\n# )\n# engine = build_engine(cfg, datamanager, model, optimizer, scheduler, writer)\n# mAP = engine.run(**engine_run_kwargs(cfg))\n# return mAP\n#\n# # gridSampler = GridSampler()\n# # study = optuna.create_study(direction=\"maximize\", sampler=gridSampler)\n# # TODO ETA\n# study = optuna.create_study(direction=\"maximize\")\n# study.optimize(objective, n_trials=4, timeout=600)\n#\n# pruned_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]\n# complete_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]\n#\n# print(\"Study statistics: \")\n# print(\" Number of finished trials: \", len(study.trials))\n# print(\" Number of pruned trials: \", len(pruned_trials))\n# print(\" Number of complete trials: \", len(complete_trials))\n#\n# print(\"Best trial:\")\n# trial = study.best_trial\n#\n# print(\" Value: \", trial.value)\n#\n# print(\" Params: \")\n# for key, value in trial.params.items():\n# print(\" {}: {}\".format(key, value))\n#\n#\n# if __name__ == '__main__':\n# main()\n" }, { "path": "torchreid/losses/GiLt_loss.py", "content": "from __future__ import division, absolute_import\n\nimport torch\nimport torch.nn as nn\nfrom collections import OrderedDict\nfrom torchmetrics import Accuracy\nfrom torchreid.losses import init_part_based_triplet_loss, CrossEntropyLoss\nfrom torchreid.utils.constants import GLOBAL, FOREGROUND, CONCAT_PARTS, PARTS\n\n\nclass GiLtLoss(nn.Module):\n \"\"\" The Global-identity Local-triplet 'GiLt' loss as described in our paper:\n 'Somers V. & al, Body Part-Based Representation Learning for Occluded Person Re-Identification, WACV23'.\n Source: https://github.com/VlSomers/bpbreid\n The default weights for the GiLt strategy (as described in the paper) are provided in 'default_losses_weights': the\n identity loss is applied only on holistic embeddings and the triplet loss is applied only on part-based embeddings.\n 'tr' denotes 'triplet' for the triplet loss and 'id' denotes 'identity' for the identity cross-entropy loss.\n \"\"\"\n\n default_losses_weights = {\n GLOBAL: {'id': 1., 'tr': 0.},\n FOREGROUND: {'id': 1., 'tr': 0.},\n CONCAT_PARTS: {'id': 1., 'tr': 0.},\n PARTS: {'id': 0., 'tr': 1.}\n }\n\n def __init__(self,\n losses_weights=None,\n use_visibility_scores=False,\n triplet_margin=0.3,\n loss_name='part_averaged_triplet_loss',\n use_gpu=False,\n writer=None):\n super().__init__()\n if losses_weights is None:\n losses_weights = self.default_losses_weights\n self.pred_accuracy = Accuracy(top_k=1)\n if use_gpu:\n self.pred_accuracy = self.pred_accuracy.cuda()\n self.losses_weights = losses_weights\n self.part_triplet_loss = init_part_based_triplet_loss(loss_name, margin=triplet_margin, writer=writer)\n self.identity_loss = CrossEntropyLoss(label_smooth=True)\n self.use_visibility_scores = use_visibility_scores\n\n def forward(self, embeddings_dict, visibility_scores_dict, id_cls_scores_dict, pids):\n \"\"\"\n Keys in the input dictionaries are from {'globl', 'foreg', 'conct', 'parts'} and correspond to the different\n types of embeddings. In the documentation below, we denote the batch size by 'N' and the number of parts by 'K'.\n :param embeddings_dict: a dictionary of embeddings, where the keys are the embedding types and the values are\n Tensors of size [N, D] or [N, K*D] or [N, K, D].\n :param visibility_scores_dict: a dictionary of visibility scores, where the keys are the embedding types and the\n values are Tensors of size [N] or [N, K].\n :param id_cls_scores_dict: a dictionary of identity classification scores, where the keys are the embedding types\n and the values are Tensors of size [N, num_classes] or [N, K, num_classes]\n :param pids: A Tensor of size [N] containing the person IDs.\n :return: a tupel with the total combined loss and a dictionnary with performance information for each individual\n loss.\n \"\"\"\n loss_summary = {}\n losses = []\n # global, foreground and parts embeddings id loss\n for key in [GLOBAL, FOREGROUND, CONCAT_PARTS, PARTS]:\n loss_info = OrderedDict() if key not in loss_summary else loss_summary[key]\n ce_w = self.losses_weights[key]['id']\n if ce_w > 0:\n parts_id_loss, parts_id_accuracy = self.compute_id_cls_loss(id_cls_scores_dict[key],\n visibility_scores_dict[key], pids)\n losses.append((ce_w, parts_id_loss))\n loss_info['c'] = parts_id_loss\n loss_info['a'] = parts_id_accuracy\n\n loss_summary[key] = loss_info\n\n # global, foreground and parts embeddings triplet loss\n for key in [GLOBAL, FOREGROUND, CONCAT_PARTS, PARTS]:\n loss_info = OrderedDict() if key not in loss_summary else loss_summary[key]\n tr_w = self.losses_weights[key]['tr']\n if tr_w > 0:\n parts_triplet_loss, parts_trivial_triplets_ratio, parts_valid_triplets_ratio = \\\n self.compute_triplet_loss(embeddings_dict[key], visibility_scores_dict[key], pids)\n losses.append((tr_w, parts_triplet_loss))\n loss_info['t'] = parts_triplet_loss\n loss_info['tt'] = parts_trivial_triplets_ratio\n loss_info['vt'] = parts_valid_triplets_ratio\n\n loss_summary[key] = loss_info\n\n # weighted sum of all losses\n if len(losses) == 0:\n return torch.tensor(0., device=(pids.get_device() if pids.is_cuda else None)), loss_summary\n else:\n loss = torch.stack([weight * loss for weight, loss in losses]).sum()\n return loss, loss_summary\n\n def compute_triplet_loss(self, embeddings, visibility_scores, pids):\n if self.use_visibility_scores:\n visibility = visibility_scores if len(visibility_scores.shape) == 2 else visibility_scores.unsqueeze(1)\n else:\n visibility = None\n embeddings = embeddings if len(embeddings.shape) == 3 else embeddings.unsqueeze(1)\n triplet_loss, trivial_triplets_ratio, valid_triplets_ratio = self.part_triplet_loss(embeddings, pids,\n parts_visibility=visibility)\n return triplet_loss, trivial_triplets_ratio, valid_triplets_ratio\n\n def compute_id_cls_loss(self, id_cls_scores, visibility_scores, pids):\n if len(id_cls_scores.shape) == 3:\n M = id_cls_scores.shape[1]\n id_cls_scores = id_cls_scores.flatten(0, 1)\n pids = pids.unsqueeze(1).expand(-1, M).flatten(0, 1)\n visibility_scores = visibility_scores.flatten(0, 1)\n weights = None\n if self.use_visibility_scores and visibility_scores.dtype is torch.bool:\n id_cls_scores = id_cls_scores[visibility_scores]\n pids = pids[visibility_scores]\n elif self.use_visibility_scores and visibility_scores.dtype is not torch.bool:\n weights = visibility_scores\n cls_loss = self.identity_loss(id_cls_scores, pids, weights)\n accuracy = self.pred_accuracy(id_cls_scores, pids)\n return cls_loss, accuracy\n" }, { "path": "torchreid/losses/__init__.py", "content": "from __future__ import division, print_function, absolute_import\n\nfrom .inter_parts_triplet_loss import InterPartsTripletLoss\nfrom .part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom .part_max_triplet_loss import PartMaxTripletLoss\nfrom .part_max_min_triplet_loss import PartMaxMinTripletLoss\nfrom .part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom .part_min_triplet_loss import PartMinTripletLoss\nfrom .part_random_max_min_triplet_loss import PartRandomMaxMinTripletLoss\nfrom .part_individual_triplet_loss import PartIndividualTripletLoss\nfrom .cross_entropy_loss import CrossEntropyLoss\nfrom .hard_mine_triplet_loss import TripletLoss\n\n__body_parts_losses = {\n 'part_averaged_triplet_loss': PartAveragedTripletLoss, # Part-Averaged triplet loss described in the paper\n 'part_max_triplet_loss': PartMaxTripletLoss,\n 'part_min_triplet_loss': PartMinTripletLoss,\n 'part_max_min_triplet_loss': PartMaxMinTripletLoss,\n 'part_random_max_min_triplet_loss': PartRandomMaxMinTripletLoss,\n 'inter_parts_triplet_loss': InterPartsTripletLoss,\n 'intra_parts_triplet_loss': PartIndividualTripletLoss,\n}\n\ndef init_part_based_triplet_loss(name, **kwargs):\n \"\"\"Initializes the part based triplet loss based on the part-based distance combination strategy.\"\"\"\n avai_body_parts_losses = list(__body_parts_losses.keys())\n if name not in avai_body_parts_losses:\n raise ValueError(\n 'Invalid loss name. Received \"{}\", '\n 'but expected to be one of {}'.format(name, avai_body_parts_losses)\n )\n return __body_parts_losses[name](**kwargs)\n\n\ndef deep_supervision(criterion, xs, y):\n \"\"\"DeepSupervision\n\n Applies criterion to each element in a list.\n\n Args:\n criterion: loss function\n xs: tuple of inputs\n y: ground truth\n \"\"\"\n loss = 0.\n for x in xs:\n loss += criterion(x, y)\n loss /= len(xs)\n return loss\n" }, { "path": "torchreid/losses/body_part_attention_loss.py", "content": "from __future__ import division, absolute_import\n\nimport torch.nn as nn\nfrom collections import OrderedDict\nfrom monai.losses import FocalLoss, DiceLoss\nfrom torch.nn import CrossEntropyLoss\nfrom torchmetrics import Accuracy\nfrom torchreid.utils.constants import PIXELS\n\n\nclass BodyPartAttentionLoss(nn.Module):\n \"\"\" A body part attention loss as described in our paper\n 'Somers V. & al, Body Part-Based Representation Learning for Occluded Person Re-Identification, WACV23'.\n Source: https://github.com/VlSomers/bpbreid\n \"\"\"\n\n def __init__(self, loss_type='cl', label_smoothing=0.1, use_gpu=False):\n super().__init__()\n self.pred_accuracy = Accuracy(top_k=1)\n if use_gpu:\n self.pred_accuracy = self.pred_accuracy.cuda()\n if loss_type == 'cl':\n self.part_prediction_loss_1 = CrossEntropyLoss(label_smoothing=label_smoothing)\n elif loss_type == 'fl':\n self.part_prediction_loss = FocalLoss(to_onehot_y=True, gamma=1.0)\n elif loss_type == 'dl':\n self.part_prediction_loss = DiceLoss(to_onehot_y=True, softmax=True)\n else:\n raise ValueError(\"Loss {} for part prediction is not supported\".format(loss_type))\n\n def forward(self, pixels_cls_scores, targets):\n \"\"\" Compute loss for body part attention prediction.\n Args:\n pixels_cls_scores [N, K, H, W]\n targets [N, H, W]\n Returns:\n \"\"\"\n loss_summary = {}\n loss_summary[PIXELS] = OrderedDict()\n pixels_cls_loss, pixels_cls_accuracy = self.compute_pixels_cls_loss(pixels_cls_scores, targets)\n loss_summary[PIXELS]['c'] = pixels_cls_loss\n loss_summary[PIXELS]['a'] = pixels_cls_accuracy\n return pixels_cls_loss, loss_summary\n\n def compute_pixels_cls_loss(self, pixels_cls_scores, targets):\n if pixels_cls_scores.is_cuda:\n targets = targets.cuda()\n pixels_cls_score_targets = targets.flatten() # [N*Hf*Wf]\n pixels_cls_scores = pixels_cls_scores.permute(0, 2, 3, 1).flatten(0, 2) # [N*Hf*Wf, M]\n loss = self.part_prediction_loss_1(pixels_cls_scores, pixels_cls_score_targets)\n accuracy = self.pred_accuracy(pixels_cls_scores, pixels_cls_score_targets)\n return loss, accuracy.item()\n" }, { "path": "torchreid/losses/cross_entropy_loss.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.nn as nn\n\n\nclass CrossEntropyLoss(nn.Module):\n r\"\"\"Cross entropy loss with label smoothing regularizer.\n \n Reference:\n Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.\n\n With label smoothing, the label :math:`y` for a class is computed by\n \n .. math::\n \\begin{equation}\n (1 - \\eps) \\times y + \\frac{\\eps}{K},\n \\end{equation}\n\n where :math:`K` denotes the number of classes and :math:`\\eps` is a weight. When\n :math:`\\eps = 0`, the loss function reduces to the normal cross entropy.\n \n Args:\n num_classes (int): number of classes.\n eps (float, optional): weight. Default is 0.1.\n use_gpu (bool, optional): whether to use gpu devices. Default is True.\n label_smooth (bool, optional): whether to apply label smoothing. Default is True.\n \"\"\"\n\n def __init__(self, eps=0.1, label_smooth=True):\n super(CrossEntropyLoss, self).__init__()\n self.eps = eps if label_smooth else 0\n self.logsoftmax = nn.LogSoftmax(dim=1)\n\n def forward(self, inputs, targets, weights=None):\n \"\"\"\n Args:\n inputs (torch.Tensor): prediction matrix (before softmax) with\n shape (batch_size, num_classes).\n targets (torch.LongTensor): ground truth labels with shape (batch_size).\n Each position contains the label index.\n \"\"\"\n assert inputs.shape[0] == targets.shape[0]\n num_classes = inputs.shape[1]\n log_probs = self.logsoftmax(inputs)\n zeros = torch.zeros(log_probs.size())\n targets = zeros.scatter_(1, targets.unsqueeze(1).data.cpu(), 1)\n if inputs.is_cuda:\n targets = targets.cuda()\n targets = (1 - self.eps) * targets + self.eps / num_classes\n if weights is not None:\n result = (-targets * log_probs).sum(dim=1)\n result = result * nn.functional.normalize(weights, p=1, dim=0)\n result = result.sum()\n else:\n result = (-targets * log_probs).mean(0).sum()\n return result\n" }, { "path": "torchreid/losses/hard_mine_triplet_loss.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.nn as nn\n\n\nclass TripletLoss(nn.Module):\n \"\"\"Triplet loss with hard positive/negative mining.\n \n Reference:\n Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.\n \n Imported from ``_.\n \n Args:\n margin (float, optional): margin for triplet. Default is 0.3.\n \"\"\"\n\n def __init__(self, margin=0.3):\n super(TripletLoss, self).__init__()\n self.margin = margin\n self.ranking_loss = nn.MarginRankingLoss(margin=margin)\n\n def forward(self, inputs, targets):\n \"\"\"\n Args:\n inputs (torch.Tensor): feature matrix with shape (batch_size, feat_dim).\n targets (torch.LongTensor): ground truth labels with shape (num_classes).\n \"\"\"\n\n # Compute pairwise distance\n dist = self.compute_dist_matrix(inputs)\n\n # For each anchor, find the hardest positive and negative\n return self.compute_hard_mine_triplet_loss(dist, inputs, targets)\n\n def compute_hard_mine_triplet_loss(self, dist, inputs, targets):\n n = inputs.size(0)\n mask = targets.expand(n, n).eq(targets.expand(n, n).t())\n dist_ap, dist_an = [], []\n for i in range(n):\n dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))\n dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))\n dist_ap = torch.cat(dist_ap)\n dist_an = torch.cat(dist_an)\n # Compute ranking hinge loss\n y = torch.ones_like(dist_an)\n return self.ranking_loss(dist_an, dist_ap, y)\n\n def compute_dist_matrix(self, inputs):\n n = inputs.size(0)\n # dist(a, b) = sqrt(sum((a_i - b_i)^2)) = sqrt(sum(a_i^2) + sum(b_i^2) - 2*sum(a_i*b_i))\n dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)\n dist = dist + dist.t() # sum(a_i^2) + sum(b_i^2)\n dist.addmm_(inputs, inputs.t(), beta=1, alpha=-2) # sum(a_i^2) + sum(b_i^2) - 2*sum(a_i*b_i)\n dist = dist.clamp(min=1e-12) # for numerical stability\n dist = dist.sqrt() # sqrt(sum(a_i^2) + sum(b_i^2) - 2*sum(a_i*b_i))\n return dist\n" }, { "path": "torchreid/losses/inter_parts_triplet_loss.py", "content": "from __future__ import division, absolute_import\nfrom torchreid.losses.part_averaged_triplet_loss import PartAveragedTripletLoss\nimport torch\n\n\nclass InterPartsTripletLoss(PartAveragedTripletLoss):\n\n def __init__(self, **kwargs):\n super(InterPartsTripletLoss, self).__init__(**kwargs)\n\n def forward(self, body_parts_features, targets, n_iter=0, parts_visibility=None):\n # body_parts_features.shape = [M, N, N]\n body_parts_dist_matrices = self.compute_mixed_body_parts_dist_matrices(body_parts_features)\n return self.hard_mine_triplet_loss(body_parts_dist_matrices, targets)\n\n def compute_mixed_body_parts_dist_matrices(self, body_parts_features):\n body_parts_features = body_parts_features.flatten(start_dim=0, end_dim=1).unsqueeze(1)\n body_parts_dist_matrices = self._part_based_pairwise_distance_matrix(body_parts_features, False, self.epsilon).squeeze()\n return body_parts_dist_matrices\n\n def hard_mine_triplet_loss(self, dist, targets): # TODO extract code for mask generation into separate method\n # TODO cleanup\n nm = dist.shape[0]\n n = targets.size(0)\n m = int(nm / n)\n expanded_targets = targets.repeat(m).expand(nm, -1)\n pids_mask = expanded_targets.eq(expanded_targets.t())\n\n body_parts_targets = []\n for i in range(0, m):\n body_parts_targets.append(torch.full_like(targets, i))\n body_parts_targets = torch.cat(body_parts_targets)\n expanded_body_parts_targets = body_parts_targets.expand(nm, -1)\n body_parts_mask = expanded_body_parts_targets.eq(expanded_body_parts_targets.t())\n\n mask_p = torch.logical_and(pids_mask, body_parts_mask)\n mask_n = pids_mask == 0\n\n # Create two big mask of size BxB with B = M*N (one entry per embedding).\n # positive mask cell = true if embedding of same identity and same body part\n # to create that, compute AND between classic pids M*N mask and a new body part mask :\n # 110000 112233.expand().eq(112233.expand().t())\n # 110000\n # 001100\n # 001100\n # 000011\n # 000011\n # negative mask cell = true if embeddings from different ids\n\n dist_ap, dist_an = [], []\n for i in range(nm):\n i_pos_dist = dist[i][mask_p[i]]\n dist_ap.append(i_pos_dist.max().unsqueeze(0))\n i_neg_dist = dist[i][mask_n[i]]\n assert i_neg_dist.nelement() != 0, \"embedding %r should have at least one negative counterpart\" % i\n dist_an.append(i_neg_dist.min().unsqueeze(0))\n dist_ap = torch.cat(dist_ap)\n dist_an = torch.cat(dist_an)\n # Compute ranking hinge loss\n y = torch.ones_like(dist_an)\n return self.ranking_loss(dist_an, dist_ap, y)\n" }, { "path": "torchreid/losses/part_averaged_triplet_loss.py", "content": "from __future__ import division, absolute_import\n\nimport warnings\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchreid.utils.tensortools import masked_mean\n\n\nclass PartAveragedTripletLoss(nn.Module):\n \"\"\"Compute the part-averaged triplet loss as described in our paper:\n 'Somers V. & al, Body Part-Based Representation Learning for Occluded Person Re-Identification, WACV23'.\n Source: https://github.com/VlSomers/bpbreid\n This class provides a generic implementation of the batch-hard triplet loss for part-based models, i.e. models\n outputting multiple embeddings (part-based/local representations) per input sample/image.\n When K=1 parts are provided and the parts_visiblity scores are set to one (or not provided), this implementation is\n strictly equal to the standard batch-hard triplet loss described in:\n 'Alexander Hermans, Lucas Beyer, and Bastian Leibe. In Defense of the Triplet Loss for Person Re-Identification.'\n It is therefore valid to use this implementation for global embeddings too.\n Part-based distances are combined into a global sample-to-sample distance using a 'mean' operation.\n Other subclasses of PartAveragedTripletLoss provide different strategies to combine local distances into a global\n one.\n This implementation is optimized, using only tensors operations and no Python loops.\n \"\"\"\n\n def __init__(self, margin=0.3, epsilon=1e-16, writer=None):\n super(PartAveragedTripletLoss, self).__init__()\n self.margin = margin\n self.writer = writer\n self.batch_debug = False\n self.imgs = None\n self.masks = None\n self.epsilon = epsilon\n\n def forward(self, part_based_embeddings, labels, parts_visibility=None):\n \"\"\"\n The part averaged triplet loss is computed in three steps.\n Firstly, we compute the part-based pairwise distance matrix of size [K, N, N] for the K parts and the N \n training samples.\n Secondly we compute the (samples) pairwise distance matrix of size [N, N] by combining the part-based distances.\n The part-based distances can be combined by averaging, max, min, etc.\n Thirdly, we compute the standard batch-hard triplet loss using the pairwise distance matrix.\n Compared to a standard triplet loss implementation, some entries in the pairwise distance matrix can have a\n value of -1. These entries correspond to pairs of samples that could not be compared, because there was no\n common visible parts for instance. Such pairs should be ignored for computing the batch hard triplets.\n \n Args:\n part_based_embeddings (torch.Tensor): feature matrix with shape (batch_size, parts_num, feat_dim).\n labels (torch.LongTensor): ground truth labels with shape (num_classes).\n \"\"\"\n\n # Compute pairwise distance matrix for each part\n part_based_pairwise_dist = self._part_based_pairwise_distance_matrix(part_based_embeddings.transpose(1, 0), squared=False)\n\n if parts_visibility is not None:\n parts_visibility = parts_visibility.t()\n valid_part_based_pairwise_dist_mask = parts_visibility.unsqueeze(1) * parts_visibility.unsqueeze(2)\n if valid_part_based_pairwise_dist_mask.dtype is not torch.bool:\n valid_part_based_pairwise_dist_mask = torch.sqrt(valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist_mask = None\n\n pairwise_dist = self._combine_part_based_dist_matrices(part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels)\n\n return self._hard_mine_triplet_loss(pairwise_dist, labels, self.margin)\n\n def _combine_part_based_dist_matrices(self, part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels):\n if valid_part_based_pairwise_dist_mask is not None:\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n pairwise_dist = masked_mean(part_based_pairwise_dist, valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist = part_based_pairwise_dist\n pairwise_dist = valid_part_based_pairwise_dist.mean(0)\n\n return pairwise_dist\n\n def _part_based_pairwise_distance_matrix(self, embeddings, squared=False):\n \"\"\"\n embeddings.shape = (K, N, C)\n ||a-b||^2 = |a|^2 - 2* + |b|^2\n \"\"\"\n dot_product = torch.matmul(embeddings, embeddings.transpose(2, 1))\n square_sum = dot_product.diagonal(dim1=1, dim2=2)\n distances = square_sum.unsqueeze(2) - 2 * dot_product + square_sum.unsqueeze(1)\n distances = F.relu(distances)\n\n if not squared:\n mask = torch.eq(distances, 0).float()\n distances = distances + mask * self.epsilon # for numerical stability (infinite derivative of sqrt in 0)\n distances = torch.sqrt(distances)\n distances = distances * (1 - mask)\n\n return distances\n\n def _hard_mine_triplet_loss(self, batch_pairwise_dist, labels, margin):\n \"\"\"\n A generic implementation of the batch-hard triplet loss.\n K (part-based) distance matrix between N samples are provided in tensor 'batch_pairwise_dist' of size [K, N, N].\n The standard batch-hard triplet loss is then computed for each of the K distance matrix, yielding a total of KxN\n triplet losses.\n When a pairwise distance matrix of size [1, N, N] is provided with K=1, this function behave like a standard\n batch-hard triplet loss.\n When a pairwise distance matrix of size [K, N, N] is provided, this function will apply the batch-hard triplet\n loss strategy K times, i.e. one time for each of the K part-based distance matrix. It will then average all\n KxN triplet losses for all K parts into one loss value.\n For the part-averaged triplet loss described in the paper, all part-based distance are first averaged before\n calling this function, and a pairwise distance matrix of size [1, N, N] is provided here.\n When the triplet loss is applied individually for each part, without considering the global/combined distance\n between two training samples (as implemented by 'PartIndividualTripletLoss'), then a (part-based) pairwise\n distance matrix of size [K, N, N] is given as input.\n Compute distance matrix; i.e. for each anchor a_i with i=range(0, batch_size) :\n - find the (a_i,p_i) pair with greatest distance s.t. a_i and p_i have the same label\n - find the (a_i,n_i) pair with smallest distance s.t. a_i and n_i have different label\n - compute triplet loss for each triplet (a_i, p_i, n_i), average them\n Source :\n - https://github.com/lyakaap/NetVLAD-pytorch/blob/master/hard_triplet_loss.py\n - https://github.com/Yuol96/pytorch-triplet-loss/blob/master/model/triplet_loss.py\n Args:\n batch_pairwise_dist: pairwise distances between samples, of size (K, N, N). A value of -1 means no distance\n could be computed between the two sample, that pair should therefore not be considered for triplet\n mining.\n labels: id labels for the batch, of size (N,)\n Returns:\n triplet_loss: scalar tensor containing the batch hard triplet loss, which is the result of the average of a\n maximum of KxN triplet losses. Triplets are generated for anchors with at least one valid negative and\n one valid positive. Invalid negatives and invalid positives are marked with a -1 distance in\n batch_pairwise_dist input tensor.\n trivial_triplets_ratio: scalar between [0, 1] indicating the ratio of hard triplets that are 'trivial', i.e.\n for which the triplet loss value is 0 because the margin condition is already satisfied.\n valid_triplets_ratio: scalar between [0, 1] indicating the ratio of hard triplets that are valid. A triplet \n is invalid if the anchor could not be compared with any positive or negative sample. Two samples cannot \n be compared if they have no mutually visible parts (therefore no distance could be computed).\n \"\"\"\n max_value = torch.finfo(batch_pairwise_dist.dtype).max\n\n valid_pairwise_dist_mask = (batch_pairwise_dist != float(-1))\n\n self.writer.update_invalid_pairwise_distances_count(batch_pairwise_dist)\n\n # Get the hardest positive pairs\n # invalid positive distance were set to -1 to\n mask_anchor_positive = self._get_anchor_positive_mask(labels).unsqueeze(0)\n mask_anchor_positive = mask_anchor_positive * valid_pairwise_dist_mask\n valid_positive_dist = batch_pairwise_dist * mask_anchor_positive.float() - (~mask_anchor_positive).float()\n hardest_positive_dist, _ = torch.max(valid_positive_dist, dim=-1) # [K, N]\n\n # Get the hardest negative pairs\n mask_anchor_negative = self._get_anchor_negative_mask(labels).unsqueeze(0)\n mask_anchor_negative = mask_anchor_negative * valid_pairwise_dist_mask\n valid_negative_dist = batch_pairwise_dist * mask_anchor_negative.float() + (~mask_anchor_negative).float() * max_value\n hardest_negative_dist, _ = torch.min(valid_negative_dist, dim=-1) # [K, N]\n\n # Hardest negative/positive with dist=float.max/-1 are invalid: no valid negative/positive found for this anchor\n # Do not generate triplet for such anchor\n valid_hardest_positive_dist_mask = hardest_positive_dist != -1\n valid_hardest_negative_dist_mask = hardest_negative_dist != max_value\n valid_triplets_mask = valid_hardest_positive_dist_mask * valid_hardest_negative_dist_mask # [K, N]\n hardest_dist = torch.stack([hardest_positive_dist, hardest_negative_dist], 2) # [K, N, 2]\n valid_hardest_dist = hardest_dist[valid_triplets_mask, :] # [K*N, 2]\n\n if valid_hardest_dist.nelement() == 0:\n warnings.warn(\"CRITICAL WARNING: no valid triplets were generated for current batch\")\n return None\n\n # Build valid triplets and compute triplet loss\n if self.margin > 0:\n triplet_loss, trivial_triplets_ratio, valid_triplets_ratio = self.hard_margin_triplet_loss(margin, valid_hardest_dist,\n valid_triplets_mask)\n else:\n triplet_loss, trivial_triplets_ratio, valid_triplets_ratio = self.soft_margin_triplet_loss(0.3, valid_hardest_dist,\n valid_triplets_mask)\n\n return triplet_loss, trivial_triplets_ratio, valid_triplets_ratio\n\n def hard_margin_triplet_loss(self, margin, valid_hardest_dist, valid_triplets_mask):\n triplet_losses = F.relu(valid_hardest_dist[:, 0] - valid_hardest_dist[:, 1] + margin)\n triplet_loss = torch.mean(triplet_losses)\n trivial_triplets_ratio = (triplet_losses == 0.).sum() / triplet_losses.nelement()\n valid_triplets_ratio = valid_triplets_mask.sum() / valid_triplets_mask.nelement()\n return triplet_loss, trivial_triplets_ratio, valid_triplets_ratio\n\n def soft_margin_triplet_loss(self, margin, valid_hardest_dist, valid_triplets_mask):\n triplet_losses = F.relu(valid_hardest_dist[:, 0] - valid_hardest_dist[:, 1] + margin)\n hard_margin_triplet_loss = torch.mean(triplet_losses)\n trivial_triplets_ratio = (triplet_losses == 0.).sum() / triplet_losses.nelement()\n valid_triplets_ratio = valid_triplets_mask.sum() / valid_triplets_mask.nelement()\n\n # valid_hardest_dist[:, 0] = hardest positive dist\n # valid_hardest_dist[:, 1] = hardest negative dist\n y = valid_hardest_dist[:, 0].new().resize_as_(valid_hardest_dist[:, 0]).fill_(1)\n soft_margin_triplet_loss = F.soft_margin_loss(valid_hardest_dist[:, 1] - valid_hardest_dist[:, 0], y)\n if soft_margin_triplet_loss == float('Inf'):\n print(\"soft_margin_triplet_loss = inf\")\n return hard_margin_triplet_loss, trivial_triplets_ratio, valid_triplets_ratio\n return soft_margin_triplet_loss, trivial_triplets_ratio, valid_triplets_ratio\n\n @staticmethod\n def _get_anchor_positive_mask(labels):\n \"\"\"\n To be a valid positive pair (a,p) :\n - a and p are different embeddings\n - a and p have the same label\n \"\"\"\n indices_equal_mask = torch.eye(labels.shape[0], dtype=torch.bool, device=(labels.get_device() if labels.is_cuda else None))\n indices_not_equal_mask = ~indices_equal_mask\n\n # Check if labels[i] == labels[j]\n labels_equal_mask = torch.eq(labels.unsqueeze(0), labels.unsqueeze(1))\n\n mask_anchor_positive = indices_not_equal_mask * labels_equal_mask\n\n return mask_anchor_positive\n\n @staticmethod\n def _get_anchor_negative_mask(labels):\n \"\"\"\n To be a valid negative pair (a,n) :\n - a and n have different labels (and therefore are different embeddings)\n \"\"\"\n\n # Check if labels[i] != labels[k]\n labels_not_equal_mask = torch.ne(torch.unsqueeze(labels, 0), torch.unsqueeze(labels, 1))\n\n return labels_not_equal_mask\n" }, { "path": "torchreid/losses/part_individual_triplet_loss.py", "content": "from __future__ import division, absolute_import\n\nfrom torchreid.losses.part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom torchreid.utils.tensortools import replace_values\n\n\nclass PartIndividualTripletLoss(PartAveragedTripletLoss):\n \"\"\"A triplet loss applied individually for each part, without considering the global/combined distance\n between two training samples. If the model outputs K embeddings (for K parts), this loss will compute the\n batch-hard triplet loss K times and output the average of them. With the part-averaged triplet loss, the global\n distance between two training samples is used in the triplet loss equation: that global distance is obtained by\n combining all K part-based distance between two samples into one value ('combining' = mean, max, min, etc).\n With the part-individual triplet loss, the triplet loss is applied only on local distance individually, i.e.,\n the distance between two local parts is used in the triplet loss equation. This part-individual triplet loss is\n therefore more sensitive to occluded parts (if 'valid_part_based_pairwise_dist_mask' is not used) and to\n non-discriminative parts, i.e. parts from two different identities having similar appearance.\n 'Somers V. & al, Body Part-Based Representation Learning for Occluded Person Re-Identification, WACV23'.\n Source: https://github.com/VlSomers/bpbreid\n \"\"\"\n def __init__(self, **kwargs):\n super(PartIndividualTripletLoss, self).__init__(**kwargs)\n\n def _combine_part_based_dist_matrices(self, part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels):\n \"\"\"Do not combine part-based distance, simply return the input part-based pairwise distances, and optionally\n replace non-valid part-based distance with -1\"\"\"\n if valid_part_based_pairwise_dist_mask is not None:\n valid_part_based_pairwise_dist = replace_values(part_based_pairwise_dist, ~valid_part_based_pairwise_dist_mask, -1)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist = part_based_pairwise_dist\n\n return valid_part_based_pairwise_dist\n" }, { "path": "torchreid/losses/part_max_min_triplet_loss.py", "content": "from __future__ import division, absolute_import\n\nimport torch\n\nfrom torchreid.losses.part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom torchreid.utils.tensortools import replace_values\n\n\nclass PartMaxMinTripletLoss(PartAveragedTripletLoss):\n\n def __init__(self, **kwargs):\n super(PartMaxMinTripletLoss, self).__init__(**kwargs)\n\n def _combine_part_based_dist_matrices(self, part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels):\n if valid_part_based_pairwise_dist_mask is not None:\n valid_part_based_pairwise_dist_for_max = replace_values(part_based_pairwise_dist, ~valid_part_based_pairwise_dist_mask, -1)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist_for_max = part_based_pairwise_dist\n\n max_pairwise_dist, part_id_for_max = valid_part_based_pairwise_dist_for_max.max(0)\n\n if valid_part_based_pairwise_dist_mask is not None:\n max_value = torch.finfo(part_based_pairwise_dist.dtype).max\n valid_part_based_pairwise_dist_for_min = replace_values(part_based_pairwise_dist,\n ~valid_part_based_pairwise_dist_mask, max_value)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist_for_min = part_based_pairwise_dist\n\n min_pairwise_dist, part_id_for_min = valid_part_based_pairwise_dist_for_min.min(0)\n\n labels_equal_mask = torch.eq(labels.unsqueeze(0), labels.unsqueeze(1))\n pairwise_dist = max_pairwise_dist * labels_equal_mask + min_pairwise_dist * ~labels_equal_mask\n part_id = part_id_for_max * labels_equal_mask + part_id_for_min * ~labels_equal_mask\n\n if valid_part_based_pairwise_dist_mask is not None:\n invalid_pairwise_dist_mask = valid_part_based_pairwise_dist_mask.sum(dim=0) == 0\n pairwise_dist = replace_values(pairwise_dist, invalid_pairwise_dist_mask, -1)\n\n if part_based_pairwise_dist.shape[0] > 1:\n self.writer.used_parts_statistics(part_based_pairwise_dist.shape[0], part_id)\n\n return pairwise_dist\n" }, { "path": "torchreid/losses/part_max_triplet_loss.py", "content": "from __future__ import division, absolute_import\n\nfrom torchreid.losses.part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom torchreid.utils.tensortools import replace_values\n\n\nclass PartMaxTripletLoss(PartAveragedTripletLoss):\n\n def __init__(self, **kwargs):\n super(PartMaxTripletLoss, self).__init__(**kwargs)\n\n def _combine_part_based_dist_matrices(self, part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels):\n if valid_part_based_pairwise_dist_mask is not None:\n valid_part_based_pairwise_dist = replace_values(part_based_pairwise_dist, ~valid_part_based_pairwise_dist_mask, -1)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist = part_based_pairwise_dist\n\n pairwise_dist, part_id = valid_part_based_pairwise_dist.max(0)\n\n parts_count = part_based_pairwise_dist.shape[0]\n\n if part_based_pairwise_dist.shape[0] > 1:\n self.writer.used_parts_statistics(parts_count, part_id)\n\n return pairwise_dist\n" }, { "path": "torchreid/losses/part_min_triplet_loss.py", "content": "from __future__ import division, absolute_import\n\nimport torch\n\nfrom torchreid.losses.part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom torchreid.utils.tensortools import replace_values\n\n\nclass PartMinTripletLoss(PartAveragedTripletLoss):\n\n def __init__(self, **kwargs):\n super(PartMinTripletLoss, self).__init__(**kwargs)\n\n def _combine_part_based_dist_matrices(self, part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels):\n if valid_part_based_pairwise_dist_mask is not None:\n max_value = torch.finfo(part_based_pairwise_dist.dtype).max\n valid_part_based_pairwise_dist = replace_values(part_based_pairwise_dist, ~valid_part_based_pairwise_dist_mask, max_value)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n else:\n valid_part_based_pairwise_dist = part_based_pairwise_dist\n\n pairwise_dist, part_id = valid_part_based_pairwise_dist.min(0)\n\n if valid_part_based_pairwise_dist_mask is not None:\n invalid_pairwise_dist_mask = valid_part_based_pairwise_dist_mask.sum(dim=0) == 0\n pairwise_dist = replace_values(pairwise_dist, invalid_pairwise_dist_mask, -1)\n\n parts_count = part_based_pairwise_dist.shape[0]\n if part_based_pairwise_dist.shape[0] > 1:\n self.writer.used_parts_statistics(parts_count, part_id)\n\n return pairwise_dist\n" }, { "path": "torchreid/losses/part_random_max_min_triplet_loss.py", "content": "from __future__ import division, absolute_import\n\nimport torch\n\nfrom torchreid.losses.part_averaged_triplet_loss import PartAveragedTripletLoss\nfrom torchreid.utils.tensortools import replace_values\n\n\nclass PartRandomMaxMinTripletLoss(PartAveragedTripletLoss):\n\n def __init__(self, **kwargs):\n super(PartRandomMaxMinTripletLoss, self).__init__(**kwargs)\n\n def _combine_part_based_dist_matrices(self, part_based_pairwise_dist, valid_part_based_pairwise_dist_mask, labels):\n if valid_part_based_pairwise_dist_mask is None:\n valid_part_based_pairwise_dist_mask = torch.ones(part_based_pairwise_dist.shape, dtype=torch.bool, device=(labels.get_device() if labels.is_cuda else None))\n # put some random entries to 0 (20%)\n dropout_mask = torch.rand(size=valid_part_based_pairwise_dist_mask.shape, device=(labels.get_device() if labels.is_cuda else None)) > 0.5 # FIXME should be symmetric?\n # dropout_mask = torch.rand(size=valid_body_part_pairwise_dist_mask.shape[0:2], device=(labels.get_device() if labels.is_cuda else None)) > 0.3\n # dropout_mask = dropout_mask.unsqueeze(1) * dropout_mask.unsqueeze(2)\n\n valid_part_based_pairwise_dist_mask *= dropout_mask\n\n valid_part_based_pairwise_dist_for_max = replace_values(part_based_pairwise_dist, ~valid_part_based_pairwise_dist_mask, -1)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n max_pairwise_dist, part_id_for_max = valid_part_based_pairwise_dist_for_max.max(0)\n\n max_value = torch.finfo(part_based_pairwise_dist.dtype).max\n valid_part_based_pairwise_dist_for_min = replace_values(part_based_pairwise_dist,\n ~valid_part_based_pairwise_dist_mask, max_value)\n self.writer.update_invalid_part_based_pairwise_distances_count(valid_part_based_pairwise_dist_mask)\n min_pairwise_dist, part_id_for_min = valid_part_based_pairwise_dist_for_min.min(0)\n\n labels_equal_mask = torch.eq(labels.unsqueeze(0), labels.unsqueeze(1))\n pairwise_dist = max_pairwise_dist * labels_equal_mask + min_pairwise_dist * ~labels_equal_mask\n part_id = part_id_for_max * labels_equal_mask + part_id_for_min * ~labels_equal_mask\n\n invalid_pairwise_dist_mask = valid_part_based_pairwise_dist_mask.sum(dim=0) == 0\n pairwise_dist = replace_values(pairwise_dist, invalid_pairwise_dist_mask, -1)\n\n if part_based_pairwise_dist.shape[0] > 1:\n self.writer.used_parts_statistics(part_based_pairwise_dist.shape[0], part_id)\n\n return pairwise_dist\n" }, { "path": "torchreid/metrics/__init__.py", "content": "from __future__ import absolute_import\n\nfrom .rank import evaluate_rank\nfrom .accuracy import accuracy\nfrom .distance import compute_distance_matrix\n" }, { "path": "torchreid/metrics/accuracy.py", "content": "from __future__ import division, print_function, absolute_import\n\n\ndef accuracy(output, target, topk=(1, )):\n \"\"\"Computes the accuracy over the k top predictions for\n the specified values of k.\n\n Args:\n output (torch.Tensor): prediction matrix with shape (batch_size, num_classes).\n target (torch.LongTensor): ground truth labels with shape (batch_size).\n topk (tuple, optional): accuracy at top-k will be computed. For example,\n topk=(1, 5) means accuracy at top-1 and top-5 will be computed.\n\n Returns:\n list: accuracy at top-k.\n\n Examples::\n >>> from torchreid import metrics\n >>> metrics.accuracy(output, target)\n \"\"\"\n maxk = max(topk)\n batch_size = target.size(0)\n\n if isinstance(output, (tuple, list)):\n output = output[0]\n\n _, pred = output.topk(maxk, 1, True, True)\n pred = pred.t()\n correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n res = []\n for k in topk:\n correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n acc = correct_k.mul_(100.0 / batch_size)\n res.append(acc)\n\n return res\n" }, { "path": "torchreid/metrics/distance.py", "content": "from __future__ import division, print_function, absolute_import\nimport torch\nfrom torch.nn import functional as F\n\nfrom torchreid.utils.writer import Writer\nfrom torchreid.utils.tensortools import replace_values, masked_mean\n\n\ndef compute_distance_matrix(input1, input2, metric='euclidean'):\n \"\"\"A wrapper function for computing distance matrix.\n\n Args:\n input1 (torch.Tensor): 2-D feature matrix.\n input2 (torch.Tensor): 2-D feature matrix.\n metric (str, optional): \"euclidean\" or \"cosine\".\n Default is \"euclidean\".\n\n Returns:\n torch.Tensor: distance matrix.\n\n Examples::\n >>> from torchreid import metrics\n >>> input1 = torch.rand(10, 2048)\n >>> input2 = torch.rand(100, 2048)\n >>> distmat = metrics.compute_distance_matrix(input1, input2)\n >>> distmat.size() # (10, 100)\n \"\"\"\n # check input\n assert isinstance(input1, torch.Tensor)\n assert isinstance(input2, torch.Tensor)\n assert input1.dim() == 2, 'Expected 2-D tensor, but got {}-D'.format(\n input1.dim()\n )\n assert input2.dim() == 2, 'Expected 2-D tensor, but got {}-D'.format(\n input2.dim()\n )\n assert input1.size(1) == input2.size(1)\n\n if metric == 'euclidean':\n distmat = euclidean_squared_distance(input1, input2)\n elif metric == 'cosine':\n distmat = cosine_distance(input1, input2)\n else:\n raise ValueError(\n 'Unknown distance metric: {}. '\n 'Please choose either \"euclidean\" or \"cosine\"'.format(metric)\n )\n\n return distmat\n\n\ndef euclidean_squared_distance(input1, input2):\n \"\"\"Computes euclidean squared distance.\n\n Args:\n input1 (torch.Tensor): 2-D feature matrix.\n input2 (torch.Tensor): 2-D feature matrix.\n\n Returns:\n torch.Tensor: distance matrix.\n \"\"\"\n # dist(a, b) = sum((a_i - b_i)^2) = sum(a_i^2) + sum(b_i^2) - 2*sum(a_i*b_i)\n m, n = input1.size(0), input2.size(0)\n mat1 = torch.pow(input1, 2).sum(dim=1, keepdim=True).expand(m, n) # sum(a_i^2)\n mat2 = torch.pow(input2, 2).sum(dim=1, keepdim=True).expand(n, m).t() # sum(b_i^2)\n distmat = mat1 + mat2 # sum(a_i^2) + sum(b_i^2)\n distmat.addmm_(input1, input2.t(), beta=1, alpha=-2) # sum(a_i^2) + sum(b_i^2) - 2*sum(a_i*b_i)\n return distmat\n\n\ndef cosine_distance(input1, input2):\n \"\"\"Computes cosine distance.\n\n Args:\n input1 (torch.Tensor): 2-D feature matrix.\n input2 (torch.Tensor): 2-D feature matrix.\n\n Returns:\n torch.Tensor: distance matrix.\n \"\"\"\n input1_normed = F.normalize(input1, p=2, dim=1)\n input2_normed = F.normalize(input2, p=2, dim=1)\n distmat = 1 - torch.mm(input1_normed, input2_normed.t())\n return distmat\n\n\ndef compute_distance_matrix_using_bp_features(qf, gf, qf_parts_visibility=None, gf_parts_visibility=None, dist_combine_strat='mean', batch_size_pairwise_dist_matrix=5000, use_gpu=False, metric='euclidean'):\n \"\"\"Computes distance matrix between each pair of samples using their part-based features. 3 implementations here: without visibility scores, with boolean/binary visibility scores and with continuous [0, 1] visibility scores.\"\"\"\n # TODO keep only one generic implementation\n if qf_parts_visibility is not None and gf_parts_visibility is not None:\n if qf_parts_visibility.dtype is torch.bool and gf_parts_visibility.dtype is torch.bool:\n # boolean visibility scores\n return _compute_distance_matrix_using_bp_features_and_masks(qf, gf, qf_parts_visibility, gf_parts_visibility, dist_combine_strat, batch_size_pairwise_dist_matrix, use_gpu, metric)\n else:\n # continuous visibility scores\n return _compute_distance_matrix_using_bp_features_and_visibility_scores(qf, gf, qf_parts_visibility, gf_parts_visibility, dist_combine_strat, batch_size_pairwise_dist_matrix, use_gpu, metric)\n else:\n # no visibility scores\n return _compute_distance_matrix_using_bp_features(qf, gf, dist_combine_strat, batch_size_pairwise_dist_matrix, use_gpu, metric)\n\n\ndef _compute_distance_matrix_using_bp_features(qf, gf, dist_combine_strat, batch_size_pairwise_dist_matrix, use_gpu, metric):\n if use_gpu:\n qf = qf.cuda()\n\n pairwise_dist_, body_part_pairwise_dist_ = [], []\n for batch_gf in torch.split(gf, batch_size_pairwise_dist_matrix):\n if use_gpu:\n batch_gf = batch_gf.cuda()\n batch_body_part_pairwise_dist = _compute_body_parts_dist_matrices(qf, batch_gf, metric)\n if dist_combine_strat == 'max':\n batch_pairwise_dist, _ = batch_body_part_pairwise_dist.max(dim=0)\n elif dist_combine_strat == 'mean':\n batch_pairwise_dist = batch_body_part_pairwise_dist.mean(dim=0)\n else:\n raise ValueError('Body parts distance combination strategy \"{}\" not supported'.format(dist_combine_strat))\n batch_body_part_pairwise_dist = batch_body_part_pairwise_dist.cpu()\n body_part_pairwise_dist_.append(batch_body_part_pairwise_dist)\n\n batch_pairwise_dist = batch_pairwise_dist.cpu()\n pairwise_dist_.append(batch_pairwise_dist)\n\n pairwise_dist = torch.cat(pairwise_dist_, 1)\n body_part_pairwise_dist = torch.cat(body_part_pairwise_dist_, 2)\n\n if Writer.current_writer() is not None:\n Writer.current_writer().qg_pairwise_dist_statistics(pairwise_dist, body_part_pairwise_dist, None, None)\n\n return pairwise_dist, body_part_pairwise_dist\n\ndef _compute_distance_matrix_using_bp_features_and_masks(qf, gf, qf_parts_visibility, gf_parts_visibility, dist_combine_strat, batch_size_pairwise_dist_matrix, use_gpu, metric):\n batch_gf_list = torch.split(gf, batch_size_pairwise_dist_matrix)\n batch_gf_parts_visibility_list = torch.split(gf_parts_visibility, batch_size_pairwise_dist_matrix)\n\n qf_parts_visibility_cpu = qf_parts_visibility\n if use_gpu:\n qf = qf.cuda()\n qf_parts_visibility = qf_parts_visibility.cuda()\n\n qf_parts_visibility = qf_parts_visibility.t()\n pairwise_dist_, body_part_pairwise_dist_ = [], []\n for batch_gf, batch_gf_parts_visibility in zip(batch_gf_list, batch_gf_parts_visibility_list):\n if use_gpu:\n batch_gf = batch_gf.cuda()\n batch_gf_parts_visibility = batch_gf_parts_visibility.cuda()\n\n batch_body_part_pairwise_dist = _compute_body_parts_dist_matrices(qf, batch_gf, metric)\n assert qf_parts_visibility.dtype is torch.bool and batch_gf_parts_visibility.dtype is torch.bool\n batch_gf_parts_visibility = batch_gf_parts_visibility.t()\n valid_body_part_pairwise_dist_mask = qf_parts_visibility.unsqueeze(2) * batch_gf_parts_visibility.unsqueeze(1)\n\n if dist_combine_strat == 'max':\n valid_body_part_pairwise_dist = replace_values(batch_body_part_pairwise_dist,\n ~valid_body_part_pairwise_dist_mask, -1)\n batch_pairwise_dist, _ = valid_body_part_pairwise_dist.max(dim=0)\n elif dist_combine_strat == 'mean':\n batch_pairwise_dist = masked_mean(batch_body_part_pairwise_dist, valid_body_part_pairwise_dist_mask)\n valid_body_part_pairwise_dist = replace_values(batch_body_part_pairwise_dist,\n ~valid_body_part_pairwise_dist_mask, -1)\n else:\n raise ValueError('Body parts distance combination strategy \"{}\" not supported'.format(dist_combine_strat))\n\n valid_body_part_pairwise_dist = valid_body_part_pairwise_dist.cpu()\n body_part_pairwise_dist_.append(valid_body_part_pairwise_dist)\n batch_pairwise_dist = batch_pairwise_dist.cpu()\n pairwise_dist_.append(batch_pairwise_dist)\n pairwise_dist = torch.cat(pairwise_dist_, 1)\n body_part_pairwise_dist = torch.cat(body_part_pairwise_dist_, 2)\n\n if Writer.current_writer() is not None:\n Writer.current_writer().qg_pairwise_dist_statistics(pairwise_dist, body_part_pairwise_dist, qf_parts_visibility_cpu, gf_parts_visibility)\n\n max_value = body_part_pairwise_dist.max() + 1 # FIXME not clean with cosine dist\n valid_pairwise_dist_mask = (pairwise_dist != float(-1))\n pairwise_dist = replace_values(pairwise_dist, ~valid_pairwise_dist_mask, max_value)\n body_part_pairwise_dist = replace_values(body_part_pairwise_dist, (body_part_pairwise_dist == -1), max_value)\n\n return pairwise_dist, body_part_pairwise_dist\n\n\ndef _compute_distance_matrix_using_bp_features_and_visibility_scores(qf, gf, qf_parts_visibility, gf_parts_visibility, dist_combine_strat, batch_size_pairwise_dist_matrix, use_gpu, metric):\n batch_gf_list = torch.split(gf, batch_size_pairwise_dist_matrix)\n batch_gf_parts_visibility_list = torch.split(gf_parts_visibility, batch_size_pairwise_dist_matrix)\n\n qf_parts_visibility_cpu = qf_parts_visibility\n if use_gpu:\n qf = qf.cuda()\n qf_parts_visibility = qf_parts_visibility.cuda()\n\n qf_parts_visibility = qf_parts_visibility.t()\n pairwise_dist_, body_part_pairwise_dist_ = [], []\n for batch_gf, batch_gf_parts_visibility in zip(batch_gf_list, batch_gf_parts_visibility_list):\n if use_gpu:\n batch_gf = batch_gf.cuda()\n batch_gf_parts_visibility = batch_gf_parts_visibility.cuda()\n\n batch_body_part_pairwise_dist = _compute_body_parts_dist_matrices(qf, batch_gf, metric)\n batch_gf_parts_visibility = batch_gf_parts_visibility.t()\n valid_body_part_pairwise_dist_mask = torch.sqrt(qf_parts_visibility.unsqueeze(2) * batch_gf_parts_visibility.unsqueeze(1))\n\n batch_pairwise_dist = masked_mean(batch_body_part_pairwise_dist, valid_body_part_pairwise_dist_mask)\n valid_body_part_pairwise_dist = batch_body_part_pairwise_dist\n\n valid_body_part_pairwise_dist = valid_body_part_pairwise_dist.cpu()\n body_part_pairwise_dist_.append(valid_body_part_pairwise_dist)\n batch_pairwise_dist = batch_pairwise_dist.cpu()\n pairwise_dist_.append(batch_pairwise_dist)\n pairwise_dist = torch.cat(pairwise_dist_, 1)\n body_part_pairwise_dist = torch.cat(body_part_pairwise_dist_, 2)\n\n # TODO check if still valid:\n if Writer.current_writer() is not None:\n Writer.current_writer().qg_pairwise_dist_statistics(pairwise_dist, body_part_pairwise_dist, qf_parts_visibility_cpu, gf_parts_visibility)\n\n max_value = body_part_pairwise_dist.max() + 1\n valid_pairwise_dist_mask = (pairwise_dist != float(-1))\n pairwise_dist = replace_values(pairwise_dist, ~valid_pairwise_dist_mask, max_value)\n\n return pairwise_dist, body_part_pairwise_dist\n\n\ndef _compute_body_parts_dist_matrices(qf, gf, metric='euclidean'):\n \"\"\"\n gf, qf shapes = (N, M, C)\n ||a-b||^2 = |a|^2 - 2* + |b|^2\n \"\"\"\n if metric == 'euclidean':\n qf = qf.transpose(1, 0)\n gf = gf.transpose(1, 0)\n dot_product = torch.matmul(qf, gf.transpose(2, 1))\n qf_square_sum = qf.pow(2).sum(dim=-1)\n gf_square_sum = gf.pow(2).sum(dim=-1)\n\n distances = qf_square_sum.unsqueeze(2) - 2 * dot_product + gf_square_sum.unsqueeze(1)\n distances = F.relu(distances)\n distances = torch.sqrt(distances)\n elif metric == 'cosine':\n qf = qf.transpose(1, 0)\n gf = gf.transpose(1, 0)\n distances = 1 - torch.matmul(qf, gf.transpose(2, 1))\n else:\n raise ValueError(\n 'Unknown distance metric: {}. '\n 'Please choose either \"euclidean\" or \"cosine\"'.format(metric)\n )\n\n return distances\n" }, { "path": "torchreid/metrics/rank.py", "content": "from __future__ import division, print_function, absolute_import\nimport numpy as np\nimport warnings\nfrom collections import defaultdict\n\ntry:\n from torchreid.metrics.rank_cylib.rank_cy import evaluate_cy\n IS_CYTHON_AVAI = True\nexcept ImportError:\n IS_CYTHON_AVAI = False\n warnings.warn(\n 'Cython evaluation (very fast so highly recommended) is '\n 'unavailable, now use python evaluation.'\n )\n\n\ndef eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):\n \"\"\"Evaluation with cuhk03 metric\n Key: one image for each gallery identity is randomly sampled for each query identity.\n Random sampling is performed num_repeats times.\n \"\"\"\n num_repeats = 10\n num_q, num_g = distmat.shape\n\n if num_g < max_rank:\n max_rank = num_g\n print(\n 'Note: number of gallery samples is quite small, got {}'.\n format(num_g)\n )\n\n indices = np.argsort(distmat, axis=1)\n matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)\n\n # compute cmc curve for each query\n all_cmc = []\n all_AP = []\n num_valid_q = 0. # number of valid query\n\n for q_idx in range(num_q):\n # get query pid and camid\n q_pid = q_pids[q_idx]\n q_camid = q_camids[q_idx]\n\n # remove gallery samples that have the same pid and camid with query\n order = indices[q_idx]\n remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)\n keep = np.invert(remove)\n\n # compute cmc curve\n raw_cmc = matches[q_idx][\n keep] # binary vector, positions with value 1 are correct matches\n if not np.any(raw_cmc):\n # this condition is true when query identity does not appear in gallery\n continue\n\n kept_g_pids = g_pids[order][keep]\n g_pids_dict = defaultdict(list)\n for idx, pid in enumerate(kept_g_pids):\n g_pids_dict[pid].append(idx)\n\n cmc = 0.\n for repeat_idx in range(num_repeats):\n mask = np.zeros(len(raw_cmc), dtype=np.bool)\n for _, idxs in g_pids_dict.items():\n # randomly sample one image for each gallery person\n rnd_idx = np.random.choice(idxs)\n mask[rnd_idx] = True\n masked_raw_cmc = raw_cmc[mask]\n _cmc = masked_raw_cmc.cumsum()\n _cmc[_cmc > 1] = 1\n cmc += _cmc[:max_rank].astype(np.float32)\n\n cmc /= num_repeats\n all_cmc.append(cmc)\n # compute AP\n num_rel = raw_cmc.sum()\n tmp_cmc = raw_cmc.cumsum()\n tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]\n tmp_cmc = np.asarray(tmp_cmc) * raw_cmc\n AP = tmp_cmc.sum() / num_rel\n all_AP.append(AP)\n num_valid_q += 1.\n\n assert num_valid_q > 0, 'Error: all query identities do not appear in gallery'\n\n all_cmc = np.asarray(all_cmc).astype(np.float32)\n cmc = all_cmc.sum(0) / num_valid_q\n mAP = np.mean(all_AP)\n\n return {\n 'cmc': cmc,\n 'mAP': mAP,\n }\n\n\ndef eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank):\n \"\"\"Evaluation with market1501 metric\n Key: for each query identity, its gallery images from the same camera view are discarded.\n \"\"\"\n num_q, num_g = distmat.shape\n\n if num_g < max_rank:\n max_rank = num_g\n print(\n 'Note: number of gallery samples is quite small, got {}'.\n format(num_g)\n )\n\n indices = np.argsort(distmat, axis=1)\n matches = (g_pids[indices] == q_pids[:, np.newaxis]).astype(np.int32)\n\n # compute cmc curve for each query\n all_cmc = []\n all_AP = []\n num_valid_q = 0. # number of valid query\n\n for q_idx in range(num_q):\n # get query pid and camid\n q_pid = q_pids[q_idx]\n q_camid = q_camids[q_idx]\n\n # remove gallery samples that have the same pid and camid with query\n order = indices[q_idx]\n remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)\n keep = np.invert(remove)\n\n # compute cmc curve\n raw_cmc = matches[q_idx][\n keep] # binary vector, positions with value 1 are correct matches\n if not np.any(raw_cmc):\n # this condition is true when query identity does not appear in gallery\n continue\n\n cmc = raw_cmc.cumsum()\n cmc[cmc > 1] = 1\n\n all_cmc.append(cmc[:max_rank])\n num_valid_q += 1.\n\n # compute average precision\n # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision\n num_rel = raw_cmc.sum()\n tmp_cmc = raw_cmc.cumsum()\n tmp_cmc = [x / (i+1.) for i, x in enumerate(tmp_cmc)]\n tmp_cmc = np.asarray(tmp_cmc) * raw_cmc\n AP = tmp_cmc.sum() / num_rel\n all_AP.append(AP)\n\n assert num_valid_q > 0, 'Error: all query identities do not appear in gallery'\n\n all_cmc = np.asarray(all_cmc).astype(np.float32)\n cmc = all_cmc.sum(0) / num_valid_q\n mAP = np.mean(all_AP)\n\n return {\n 'cmc': cmc,\n 'mAP': mAP,\n }\n\n\ndef evaluate_py(\n distmat, q_pids, g_pids, q_camids, g_camids, max_rank, eval_metric, q_anns=None, g_anns=None,\n):\n if eval_metric == 'default':\n return eval_market1501(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)\n elif eval_metric == 'cuhk03':\n return eval_cuhk03(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)\n else:\n raise ValueError(\"Incorrect eval_metric value '{}'\".format(eval_metric))\n\n\ndef evaluate_rank(\n distmat,\n q_pids,\n g_pids,\n q_camids,\n g_camids,\n max_rank=50,\n eval_metric='default',\n q_anns=None,\n g_anns=None,\n use_cython=True\n):\n \"\"\"Evaluates CMC rank.\n\n Args:\n distmat (numpy.ndarray): distance matrix of shape (num_query, num_gallery).\n q_pids (numpy.ndarray): 1-D array containing person identities\n of each query instance.\n g_pids (numpy.ndarray): 1-D array containing person identities\n of each gallery instance.\n q_camids (numpy.ndarray): 1-D array containing camera views under\n which each query instance is captured.\n g_camids (numpy.ndarray): 1-D array containing camera views under\n which each gallery instance is captured.\n max_rank (int, optional): maximum CMC rank to be computed. Default is 50.\n eval_metric (str, optional): use multi-gallery-shot setting with 'default', single-gallery-shot\n setting with 'cuhk03' or action-to-replay setting with 'soccernetv3'.\n Default is 'default'.\n use_cython (bool, optional): use cython code for evaluation. Default is True.\n This is highly recommended as the cython code can speed up the cmc computation\n by more than 10x. This requires Cython to be installed.\n \"\"\"\n if use_cython and IS_CYTHON_AVAI and (eval_metric == 'default' or eval_metric == 'cuhk03'):\n return evaluate_py(\n distmat, q_pids, g_pids, q_camids, g_camids, max_rank,\n eval_metric\n )\n else:\n return evaluate_py(\n distmat, q_pids, g_pids, q_camids, g_camids, max_rank,\n eval_metric, q_anns=q_anns, g_anns=g_anns\n )\n" }, { "path": "torchreid/metrics/rank_cylib/Makefile", "content": "all:\n\t$(PYTHON) setup.py build_ext --inplace\n\trm -rf build\nclean:\n\trm -rf build\n\trm -f rank_cy.c *.so" }, { "path": "torchreid/metrics/rank_cylib/__init__.py", "content": "" }, { "path": "torchreid/metrics/rank_cylib/rank_cy.pyx", "content": "# cython: boundscheck=False, wraparound=False, nonecheck=False, cdivision=True\n\nfrom __future__ import print_function\n\nimport cython\nimport numpy as np\ncimport numpy as np\nfrom collections import defaultdict\nimport random\n\n\n\"\"\"\nCompiler directives:\nhttps://github.com/cython/cython/wiki/enhancements-compilerdirectives\n\nCython tutorial:\nhttps://cython.readthedocs.io/en/latest/src/userguide/numpy_tutorial.html\n\nCredit to https://github.com/luzai\n\"\"\"\n\n\n# Main interface\ncpdef evaluate_cy(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_metric_cuhk03=False):\n distmat = np.asarray(distmat, dtype=np.float32)\n q_pids = np.asarray(q_pids, dtype=np.int64)\n g_pids = np.asarray(g_pids, dtype=np.int64)\n q_camids = np.asarray(q_camids, dtype=np.int64)\n g_camids = np.asarray(g_camids, dtype=np.int64)\n if use_metric_cuhk03:\n return eval_cuhk03_cy(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)\n return eval_market1501_cy(distmat, q_pids, g_pids, q_camids, g_camids, max_rank)\n\n\ncpdef eval_cuhk03_cy(float[:,:] distmat, long[:] q_pids, long[:]g_pids,\n long[:]q_camids, long[:]g_camids, long max_rank):\n \n cdef long num_q = distmat.shape[0]\n cdef long num_g = distmat.shape[1]\n\n if num_g < max_rank:\n max_rank = num_g\n print('Note: number of gallery samples is quite small, got {}'.format(num_g))\n \n cdef:\n long num_repeats = 10\n long[:,:] indices = np.argsort(distmat, axis=1)\n long[:,:] matches = (np.asarray(g_pids)[np.asarray(indices)] == np.asarray(q_pids)[:, np.newaxis]).astype(np.int64)\n\n float[:,:] all_cmc = np.zeros((num_q, max_rank), dtype=np.float32)\n float[:] all_AP = np.zeros(num_q, dtype=np.float32)\n float num_valid_q = 0. # number of valid query\n\n long q_idx, q_pid, q_camid, g_idx\n long[:] order = np.zeros(num_g, dtype=np.int64)\n long keep\n\n float[:] raw_cmc = np.zeros(num_g, dtype=np.float32) # binary vector, positions with value 1 are correct matches\n float[:] masked_raw_cmc = np.zeros(num_g, dtype=np.float32)\n float[:] cmc, masked_cmc\n long num_g_real, num_g_real_masked, rank_idx, rnd_idx\n unsigned long meet_condition\n float AP\n long[:] kept_g_pids, mask\n\n float num_rel\n float[:] tmp_cmc = np.zeros(num_g, dtype=np.float32)\n float tmp_cmc_sum\n \n for q_idx in range(num_q):\n # get query pid and camid\n q_pid = q_pids[q_idx]\n q_camid = q_camids[q_idx]\n\n # remove gallery samples that have the same pid and camid with query\n for g_idx in range(num_g):\n order[g_idx] = indices[q_idx, g_idx]\n num_g_real = 0\n meet_condition = 0\n kept_g_pids = np.zeros(num_g, dtype=np.int64)\n\n for g_idx in range(num_g):\n if (g_pids[order[g_idx]] != q_pid) or (g_camids[order[g_idx]] != q_camid):\n raw_cmc[num_g_real] = matches[q_idx][g_idx]\n kept_g_pids[num_g_real] = g_pids[order[g_idx]]\n num_g_real += 1\n if matches[q_idx][g_idx] > 1e-31:\n meet_condition = 1\n \n if not meet_condition:\n # this condition is true when query identity does not appear in gallery\n continue\n\n # cuhk03-specific setting\n g_pids_dict = defaultdict(list) # overhead!\n for g_idx in range(num_g_real):\n g_pids_dict[kept_g_pids[g_idx]].append(g_idx)\n\n cmc = np.zeros(max_rank, dtype=np.float32)\n for _ in range(num_repeats):\n mask = np.zeros(num_g_real, dtype=np.int64)\n \n for _, idxs in g_pids_dict.items():\n # randomly sample one image for each gallery person\n rnd_idx = np.random.choice(idxs)\n #rnd_idx = idxs[0] # use deterministic for debugging\n mask[rnd_idx] = 1\n\n num_g_real_masked = 0\n for g_idx in range(num_g_real):\n if mask[g_idx] == 1:\n masked_raw_cmc[num_g_real_masked] = raw_cmc[g_idx]\n num_g_real_masked += 1\n\n masked_cmc = np.zeros(num_g, dtype=np.float32)\n function_cumsum(masked_raw_cmc, masked_cmc, num_g_real_masked)\n for g_idx in range(num_g_real_masked):\n if masked_cmc[g_idx] > 1:\n masked_cmc[g_idx] = 1\n\n for rank_idx in range(max_rank):\n cmc[rank_idx] += masked_cmc[rank_idx] / num_repeats\n \n for rank_idx in range(max_rank):\n all_cmc[q_idx, rank_idx] = cmc[rank_idx]\n # compute average precision\n # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision\n function_cumsum(raw_cmc, tmp_cmc, num_g_real)\n num_rel = 0\n tmp_cmc_sum = 0\n for g_idx in range(num_g_real):\n tmp_cmc_sum += (tmp_cmc[g_idx] / (g_idx + 1.)) * raw_cmc[g_idx]\n num_rel += raw_cmc[g_idx]\n all_AP[q_idx] = tmp_cmc_sum / num_rel\n num_valid_q += 1.\n\n assert num_valid_q > 0, 'Error: all query identities do not appear in gallery'\n\n # compute averaged cmc\n cdef float[:] avg_cmc = np.zeros(max_rank, dtype=np.float32)\n for rank_idx in range(max_rank):\n for q_idx in range(num_q):\n avg_cmc[rank_idx] += all_cmc[q_idx, rank_idx]\n avg_cmc[rank_idx] /= num_valid_q\n \n cdef float mAP = 0\n for q_idx in range(num_q):\n mAP += all_AP[q_idx]\n mAP /= num_valid_q\n\n return np.asarray(avg_cmc).astype(np.float32), mAP\n\n\ncpdef eval_market1501_cy(float[:,:] distmat, long[:] q_pids, long[:]g_pids,\n long[:]q_camids, long[:]g_camids, long max_rank):\n \n cdef long num_q = distmat.shape[0]\n cdef long num_g = distmat.shape[1]\n\n if num_g < max_rank:\n max_rank = num_g\n print('Note: number of gallery samples is quite small, got {}'.format(num_g))\n \n cdef:\n long[:,:] indices = np.argsort(distmat, axis=1)\n long[:,:] matches = (np.asarray(g_pids)[np.asarray(indices)] == np.asarray(q_pids)[:, np.newaxis]).astype(np.int64)\n\n float[:,:] all_cmc = np.zeros((num_q, max_rank), dtype=np.float32)\n float[:] all_AP = np.zeros(num_q, dtype=np.float32)\n float num_valid_q = 0. # number of valid query\n\n long q_idx, q_pid, q_camid, g_idx\n long[:] order = np.zeros(num_g, dtype=np.int64)\n long keep\n\n float[:] raw_cmc = np.zeros(num_g, dtype=np.float32) # binary vector, positions with value 1 are correct matches\n float[:] cmc = np.zeros(num_g, dtype=np.float32)\n long num_g_real, rank_idx\n unsigned long meet_condition\n\n float num_rel\n float[:] tmp_cmc = np.zeros(num_g, dtype=np.float32)\n float tmp_cmc_sum\n \n for q_idx in range(num_q):\n # get query pid and camid\n q_pid = q_pids[q_idx]\n q_camid = q_camids[q_idx]\n\n # remove gallery samples that have the same pid and camid with query\n for g_idx in range(num_g):\n order[g_idx] = indices[q_idx, g_idx]\n num_g_real = 0\n meet_condition = 0\n \n for g_idx in range(num_g):\n if (g_pids[order[g_idx]] != q_pid) or (g_camids[order[g_idx]] != q_camid):\n raw_cmc[num_g_real] = matches[q_idx][g_idx]\n num_g_real += 1\n if matches[q_idx][g_idx] > 1e-31:\n meet_condition = 1\n \n if not meet_condition:\n # this condition is true when query identity does not appear in gallery\n continue\n\n # compute cmc\n function_cumsum(raw_cmc, cmc, num_g_real)\n for g_idx in range(num_g_real):\n if cmc[g_idx] > 1:\n cmc[g_idx] = 1\n\n for rank_idx in range(max_rank):\n all_cmc[q_idx, rank_idx] = cmc[rank_idx]\n num_valid_q += 1.\n\n # compute average precision\n # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision\n function_cumsum(raw_cmc, tmp_cmc, num_g_real)\n num_rel = 0\n tmp_cmc_sum = 0\n for g_idx in range(num_g_real):\n tmp_cmc_sum += (tmp_cmc[g_idx] / (g_idx + 1.)) * raw_cmc[g_idx]\n num_rel += raw_cmc[g_idx]\n all_AP[q_idx] = tmp_cmc_sum / num_rel\n\n assert num_valid_q > 0, 'Error: all query identities do not appear in gallery'\n\n # compute averaged cmc\n cdef float[:] avg_cmc = np.zeros(max_rank, dtype=np.float32)\n for rank_idx in range(max_rank):\n for q_idx in range(num_q):\n avg_cmc[rank_idx] += all_cmc[q_idx, rank_idx]\n avg_cmc[rank_idx] /= num_valid_q\n \n cdef float mAP = 0\n for q_idx in range(num_q):\n mAP += all_AP[q_idx]\n mAP /= num_valid_q\n\n return np.asarray(avg_cmc).astype(np.float32), mAP\n\n\n# Compute the cumulative sum\ncdef void function_cumsum(cython.numeric[:] src, cython.numeric[:] dst, long n):\n cdef long i\n dst[0] = src[0]\n for i in range(1, n):\n dst[i] = src[i] + dst[i - 1]" }, { "path": "torchreid/metrics/rank_cylib/setup.py", "content": "import numpy as np\nfrom distutils.core import setup\nfrom distutils.extension import Extension\nfrom Cython.Build import cythonize\n\n\ndef numpy_include():\n try:\n numpy_include = np.get_include()\n except AttributeError:\n numpy_include = np.get_numpy_include()\n return numpy_include\n\n\next_modules = [\n Extension(\n 'rank_cy',\n ['rank_cy.pyx'],\n include_dirs=[numpy_include()],\n )\n]\n\nsetup(\n name='Cython-based reid evaluation code',\n ext_modules=cythonize(ext_modules)\n)\n" }, { "path": "torchreid/metrics/rank_cylib/test_cython.py", "content": "from __future__ import print_function\nimport sys\nimport numpy as np\nimport timeit\nimport os.path as osp\n\nfrom torchreid import metrics\n\nsys.path.insert(0, osp.dirname(osp.abspath(__file__)) + '/../../..')\n\"\"\"\nTest the speed of cython-based evaluation code. The speed improvements\ncan be much bigger when using the real reid data, which contains a larger\namount of query and gallery images.\n\nNote: you might encounter the following error:\n 'AssertionError: Error: all query identities do not appear in gallery'.\nThis is normal because the inputs are random numbers. Just try again.\n\"\"\"\n\nprint('*** Compare running time ***')\n\nsetup = '''\nimport sys\nimport os.path as osp\nimport numpy as np\nsys.path.insert(0, osp.dirname(osp.abspath(__file__)) + '/../../..')\nfrom torchreid import metrics\nnum_q = 30\nnum_g = 300\nmax_rank = 5\ndistmat = np.random.rand(num_q, num_g) * 20\nq_pids = np.random.randint(0, num_q, size=num_q)\ng_pids = np.random.randint(0, num_g, size=num_g)\nq_camids = np.random.randint(0, 5, size=num_q)\ng_camids = np.random.randint(0, 5, size=num_g)\n'''\n\nprint('=> Using market1501\\'s metric')\npytime = timeit.timeit(\n 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=False)',\n setup=setup,\n number=20\n)\ncytime = timeit.timeit(\n 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=True)',\n setup=setup,\n number=20\n)\nprint('Python time: {} s'.format(pytime))\nprint('Cython time: {} s'.format(cytime))\nprint('Cython is {} times faster than python\\n'.format(pytime / cytime))\n\nprint('=> Using cuhk03\\'s metric')\npytime = timeit.timeit(\n 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_metric_cuhk03=True, use_cython=False)',\n setup=setup,\n number=20\n)\ncytime = timeit.timeit(\n 'metrics.evaluate_rank(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_metric_cuhk03=True, use_cython=True)',\n setup=setup,\n number=20\n)\nprint('Python time: {} s'.format(pytime))\nprint('Cython time: {} s'.format(cytime))\nprint('Cython is {} times faster than python\\n'.format(pytime / cytime))\n\"\"\"\nprint(\"=> Check precision\")\n\nnum_q = 30\nnum_g = 300\nmax_rank = 5\ndistmat = np.random.rand(num_q, num_g) * 20\nq_pids = np.random.randint(0, num_q, size=num_q)\ng_pids = np.random.randint(0, num_g, size=num_g)\nq_camids = np.random.randint(0, 5, size=num_q)\ng_camids = np.random.randint(0, 5, size=num_g)\n\ncmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=False)\nprint(\"Python:\\nmAP = {} \\ncmc = {}\\n\".format(mAP, cmc))\ncmc, mAP = evaluate(distmat, q_pids, g_pids, q_camids, g_camids, max_rank, use_cython=True)\nprint(\"Cython:\\nmAP = {} \\ncmc = {}\\n\".format(mAP, cmc))\n\"\"\"\n" }, { "path": "torchreid/models/__init__.py", "content": "from __future__ import absolute_import\nfrom .hrnet import hrnet32\nfrom .pcb import *\nfrom .mlfn import *\nfrom .hacnn import *\nfrom .osnet import *\nfrom .pvpm import pose_resnet50_256_p4, pose_resnet50_256_p6, pose_resnet50_256_p6_pscore_reg, \\\n pose_resnet50_256_p4_pscore_reg\nfrom .resnet_fastreid import build_resnet_backbone, fastreid_resnet, fastreid_resnet_ibn, fastreid_resnet_nl, \\\n fastreid_resnet_ibn_nl\nfrom .senet import *\nfrom .mudeep import *\nfrom .nasnet import *\nfrom .resnet import *\nfrom .densenet import *\nfrom .xception import *\nfrom .osnet_ain import *\nfrom .resnetmid import *\nfrom .shufflenet import *\nfrom .squeezenet import *\nfrom .inceptionv4 import *\nfrom .mobilenetv2 import *\nfrom .resnet_ibn_a import *\nfrom .resnet_ibn_b import *\nfrom .shufflenetv2 import *\nfrom .inceptionresnetv2 import *\nfrom .bpbreid import *\n\n__model_factory = {\n # image classification models\n 'resnet18': resnet18,\n 'resnet34': resnet34,\n 'resnet50': resnet50,\n 'resnet101': resnet101,\n 'resnet152': resnet152,\n 'resnext50_32x4d': resnext50_32x4d,\n 'resnext101_32x8d': resnext101_32x8d,\n 'resnet50_fc512': resnet50_fc512,\n 'se_resnet50': se_resnet50,\n 'se_resnet50_fc512': se_resnet50_fc512,\n 'se_resnet101': se_resnet101,\n 'se_resnext50_32x4d': se_resnext50_32x4d,\n 'se_resnext101_32x4d': se_resnext101_32x4d,\n 'densenet121': densenet121,\n 'densenet169': densenet169,\n 'densenet201': densenet201,\n 'densenet161': densenet161,\n 'densenet121_fc512': densenet121_fc512,\n 'inceptionresnetv2': inceptionresnetv2,\n 'inceptionv4': inceptionv4,\n 'xception': xception,\n 'resnet50_ibn_a': resnet50_ibn_a,\n 'resnet50_ibn_b': resnet50_ibn_b,\n # lightweight models\n 'nasnsetmobile': nasnetamobile,\n 'mobilenetv2_x1_0': mobilenetv2_x1_0,\n 'mobilenetv2_x1_4': mobilenetv2_x1_4,\n 'shufflenet': shufflenet,\n 'squeezenet1_0': squeezenet1_0,\n 'squeezenet1_0_fc512': squeezenet1_0_fc512,\n 'squeezenet1_1': squeezenet1_1,\n 'shufflenet_v2_x0_5': shufflenet_v2_x0_5,\n 'shufflenet_v2_x1_0': shufflenet_v2_x1_0,\n 'shufflenet_v2_x1_5': shufflenet_v2_x1_5,\n 'shufflenet_v2_x2_0': shufflenet_v2_x2_0,\n # reid-specific models\n 'mudeep': MuDeep,\n 'resnet50mid': resnet50mid,\n 'hacnn': HACNN,\n 'pcb_p6': pcb_p6,\n 'pcb_p4': pcb_p4,\n 'mlfn': mlfn,\n 'osnet_x1_0': osnet_x1_0,\n 'osnet_x0_75': osnet_x0_75,\n 'osnet_x0_5': osnet_x0_5,\n 'osnet_x0_25': osnet_x0_25,\n 'osnet_ibn_x1_0': osnet_ibn_x1_0,\n 'osnet_ain_x1_0': osnet_ain_x1_0,\n 'pose_p4': pose_resnet50_256_p4,\n 'pose_p6': pose_resnet50_256_p6,\n 'pose_p6s': pose_resnet50_256_p6_pscore_reg,\n 'pose_p4s': pose_resnet50_256_p4_pscore_reg,\n 'hrnet32': hrnet32,\n 'bpbreid': bpbreid,\n 'fastreid_resnet': fastreid_resnet,\n 'fastreid_resnet_ibn': fastreid_resnet_ibn,\n 'fastreid_resnet_nl': fastreid_resnet_nl,\n 'fastreid_resnet_ibn_nl': fastreid_resnet_ibn_nl,\n}\n\n\ndef show_avai_models():\n \"\"\"Displays available models.\n\n Examples::\n >>> from torchreid import models\n >>> models.show_avai_models()\n \"\"\"\n print(list(__model_factory.keys()))\n\n\ndef build_model(\n name, num_classes, loss='softmax', pretrained=True, use_gpu=True, **kwargs\n):\n \"\"\"A function wrapper for building a model.\n\n Args:\n name (str): model name.\n num_classes (int): number of training identities.\n loss (str, optional): loss function to optimize the model. Currently\n supports \"softmax\" and \"triplet\". Default is \"softmax\".\n pretrained (bool, optional): whether to load ImageNet-pretrained weights.\n Default is True.\n use_gpu (bool, optional): whether to use gpu. Default is True.\n\n Returns:\n nn.Module\n\n Examples::\n >>> from torchreid import models\n >>> model = models.build_model('resnet50', 751, loss='softmax')\n \"\"\"\n avai_models = list(__model_factory.keys())\n if name not in avai_models:\n raise KeyError(\n 'Unknown model: {}. Must be one of {}'.format(name, avai_models)\n )\n return __model_factory[name](\n num_classes=num_classes,\n loss=loss,\n pretrained=pretrained,\n use_gpu=use_gpu,\n **kwargs\n )\n" }, { "path": "torchreid/models/bpbreid.py", "content": "from __future__ import division, absolute_import\n\nimport torch\nimport torch.nn.functional as F\nimport numpy as np\nfrom torch import nn\nfrom torchreid import models\nfrom torchreid.utils.constants import *\n\n__all__ = [\n 'bpbreid'\n]\n\n\nclass BPBreID(nn.Module):\n \"\"\"Posed based feature extraction network\n \"\"\"\n def __init__(self, num_classes, pretrained, loss, model_cfg, horizontal_stripes=False, **kwargs):\n super(BPBreID, self).__init__()\n\n # Init config\n self.model_cfg = model_cfg\n # number of training classes/identities\n self.num_classes = num_classes\n # number of parts K\n self.parts_num = self.model_cfg.masks.parts_num\n # whether to perform horizontal stripes pooling similar to PCB\n self.horizontal_stripes = horizontal_stripes\n # use shared weights/parameters between each part branch for the identity classifier\n self.shared_parts_id_classifier = self.model_cfg.shared_parts_id_classifier\n # at test time, perform a 'soft' or 'hard' merging of the learned attention maps with the external part masks\n self.test_use_target_segmentation = self.model_cfg.test_use_target_segmentation\n # use continuous or binary visibility scores at train time:\n self.training_binary_visibility_score = self.model_cfg.training_binary_visibility_score\n # use continuous or binary visibility scores at test time:\n self.testing_binary_visibility_score = self.model_cfg.testing_binary_visibility_score\n\n # Init backbone feature extractor\n self.backbone_appearance_feature_extractor = models.build_model(self.model_cfg.backbone,\n num_classes,\n loss=loss,\n pretrained=pretrained,\n last_stride=self.model_cfg.last_stride,\n enable_dim_reduction=(self.model_cfg.dim_reduce=='before_pooling'),\n dim_reduction_channels=self.model_cfg.dim_reduce_output,\n pretrained_path=self.model_cfg.hrnet_pretrained_path\n )\n self.spatial_feature_size = self.backbone_appearance_feature_extractor.feature_dim\n\n # Init dim reduce layers\n self.init_dim_reduce_layers(self.model_cfg.dim_reduce,\n self.spatial_feature_size,\n self.model_cfg.dim_reduce_output)\n\n # Init pooling layers\n self.global_pooling_head = nn.AdaptiveAvgPool2d(1)\n self.foreground_attention_pooling_head = GlobalAveragePoolingHead(self.dim_reduce_output)\n self.background_attention_pooling_head = GlobalAveragePoolingHead(self.dim_reduce_output)\n self.parts_attention_pooling_head = init_part_attention_pooling_head(self.model_cfg.normalization,\n self.model_cfg.pooling,\n self.dim_reduce_output)\n\n # Init parts classifier\n self.learnable_attention_enabled = self.model_cfg.learnable_attention_enabled\n self.pixel_classifier = PixelToPartClassifier(self.spatial_feature_size, self.parts_num)\n\n # Init id classifier\n self.global_identity_classifier = BNClassifier(self.dim_reduce_output, self.num_classes)\n self.background_identity_classifier = BNClassifier(self.dim_reduce_output, self.num_classes)\n self.foreground_identity_classifier = BNClassifier(self.dim_reduce_output, self.num_classes)\n self.concat_parts_identity_classifier = BNClassifier(self.parts_num * self.dim_reduce_output, self.num_classes)\n if self.shared_parts_id_classifier:\n # the same identity classifier weights are used for each part branch\n self.parts_identity_classifier = BNClassifier(self.dim_reduce_output, self.num_classes)\n else:\n # each part branch has its own identity classifier\n self.parts_identity_classifier = nn.ModuleList(\n [\n BNClassifier(self.dim_reduce_output, self.num_classes)\n for _ in range(self.parts_num)\n ]\n )\n\n def init_dim_reduce_layers(self, dim_reduce_mode, spatial_feature_size, dim_reduce_output):\n self.dim_reduce_output = dim_reduce_output\n self.after_pooling_dim_reduce = False\n self.before_pooling_dim_reduce = None\n\n if dim_reduce_mode == 'before_pooling':\n self.before_pooling_dim_reduce = BeforePoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.spatial_feature_size = dim_reduce_output\n elif dim_reduce_mode == 'after_pooling':\n self.after_pooling_dim_reduce = True\n self.global_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.foreground_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.background_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.parts_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n elif dim_reduce_mode == 'before_and_after_pooling':\n self.before_pooling_dim_reduce = BeforePoolingDimReduceLayer(spatial_feature_size, dim_reduce_output * 2)\n spatial_feature_size = dim_reduce_output * 2\n self.spatial_feature_size = spatial_feature_size\n self.after_pooling_dim_reduce = True\n self.global_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.foreground_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.background_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n self.parts_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output)\n elif dim_reduce_mode == 'after_pooling_with_dropout':\n self.after_pooling_dim_reduce = True\n self.global_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output, 0.5)\n self.foreground_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output, 0.5)\n self.background_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output, 0.5)\n self.parts_after_pooling_dim_reduce = AfterPoolingDimReduceLayer(spatial_feature_size, dim_reduce_output, 0.5)\n else:\n self.dim_reduce_output = spatial_feature_size\n\n def forward(self, images, external_parts_masks=None):\n \"\"\"\n :param images: images tensor of size [N, C, Hi, Wi], where N is the batch size, C channel depth (3 for RGB), and\n (Hi, Wi) are the image height and width.\n :param external_parts_masks: masks tensor of size [N, K+1, Hm, Wm], where N is the batch size, K is the number\n parts, and (Hm, Wm) are the masks height and width. The first index (index 0) along the parts K+1 dimension\n is the background by convention. The masks are expected to have values in the range [0, 1]. Spatial entry at\n location external_parts_masks[i, k+1, h, w] is the probability that the pixel at location (h, w) belongs to\n part k for batch sample i. The masks are NOT expected to be of the same size as the images.\n :return:\n \"\"\"\n\n # Global spatial_features\n spatial_features = self.backbone_appearance_feature_extractor(images) # [N, D, Hf, Wf]\n N, _, Hf, Wf = spatial_features.shape\n\n if self.before_pooling_dim_reduce is not None \\\n and spatial_features.shape[1] != self.dim_reduce_output: # When HRNet used as backbone, already done\n spatial_features = self.before_pooling_dim_reduce(spatial_features) # [N, dim_reduce_output, Hf, Wf]\n\n # Pixels classification and parts attention weights\n if self.horizontal_stripes:\n pixels_cls_scores = None\n feature_map_shape = (Hf, Wf)\n stripes_range = np.round(np.arange(0, self.parts_num + 1) * feature_map_shape[0] / self.parts_num).astype(int)\n pcb_masks = torch.zeros((self.parts_num, feature_map_shape[0], feature_map_shape[1]))\n for i in range(0, stripes_range.size - 1):\n pcb_masks[i, stripes_range[i]:stripes_range[i + 1], :] = 1\n pixels_parts_probabilities = pcb_masks\n pixels_parts_probabilities.requires_grad = False\n elif self.learnable_attention_enabled:\n pixels_cls_scores = self.pixel_classifier(spatial_features) # [N, K, Hf, Wf]\n pixels_parts_probabilities = F.softmax(pixels_cls_scores, dim=1)\n else:\n pixels_cls_scores = None\n assert external_parts_masks is not None\n external_parts_masks = external_parts_masks.type(spatial_features.dtype)\n pixels_parts_probabilities = nn.functional.interpolate(external_parts_masks, (Hf, Wf), mode='bilinear', align_corners=True)\n pixels_parts_probabilities.requires_grad = False\n assert pixels_parts_probabilities.max() <= 1 and pixels_parts_probabilities.min() >= 0\n\n background_masks = pixels_parts_probabilities[:, 0]\n parts_masks = pixels_parts_probabilities[:, 1:]\n\n # Explicit pixels segmentation of re-id target using external part masks\n if not self.training and self.test_use_target_segmentation == 'hard':\n assert external_parts_masks is not None\n # hard masking\n external_parts_masks = nn.functional.interpolate(external_parts_masks, (Hf, Wf), mode='bilinear',\n align_corners=True)\n target_segmentation_mask = external_parts_masks[:, 1::].max(dim=1)[0] > external_parts_masks[:, 0]\n background_masks = ~target_segmentation_mask\n parts_masks[background_masks.unsqueeze(1).expand_as(parts_masks)] = 1e-12\n\n if not self.training and self.test_use_target_segmentation == 'soft':\n assert external_parts_masks is not None\n # soft masking\n external_parts_masks = nn.functional.interpolate(external_parts_masks, (Hf, Wf), mode='bilinear',\n align_corners=True)\n parts_masks = parts_masks * external_parts_masks[:, 1::]\n\n # foreground_masks = parts_masks.sum(dim=1)\n foreground_masks = parts_masks.max(dim=1)[0]\n global_masks = torch.ones_like(foreground_masks)\n\n # Parts visibility\n if (self.training and self.training_binary_visibility_score) or (not self.training and self.testing_binary_visibility_score):\n pixels_parts_predictions = pixels_parts_probabilities.argmax(dim=1) # [N, Hf, Wf]\n pixels_parts_predictions_one_hot = F.one_hot(pixels_parts_predictions, self.parts_num + 1).permute(0, 3, 1, 2) # [N, K+1, Hf, Wf]\n parts_visibility = pixels_parts_predictions_one_hot.amax(dim=(2, 3)).to(torch.bool) # [N, K+1]\n else:\n parts_visibility = pixels_parts_probabilities.amax(dim=(2, 3)) # [N, K+1]\n background_visibility = parts_visibility[:, 0] # [N]\n foreground_visibility = parts_visibility.amax(dim=1) # [N]\n parts_visibility = parts_visibility[:, 1:] # [N, K]\n concat_parts_visibility = foreground_visibility\n global_visibility = torch.ones_like(foreground_visibility) # [N]\n\n # Global embedding\n global_embeddings = self.global_pooling_head(spatial_features).view(N, -1) # [N, D]\n\n # Foreground and background embeddings\n foreground_embeddings = self.foreground_attention_pooling_head(spatial_features, foreground_masks.unsqueeze(1)).flatten(1, 2) # [N, D]\n background_embeddings = self.background_attention_pooling_head(spatial_features, background_masks.unsqueeze(1)).flatten(1, 2) # [N, D]\n\n # Part features\n parts_embeddings = self.parts_attention_pooling_head(spatial_features, parts_masks) # [N, K, D]\n\n # Dim reduction\n if self.after_pooling_dim_reduce:\n global_embeddings = self.global_after_pooling_dim_reduce(global_embeddings) # [N, D]\n foreground_embeddings = self.foreground_after_pooling_dim_reduce(foreground_embeddings) # [N, D]\n background_embeddings = self.background_after_pooling_dim_reduce(background_embeddings) # [N, D]\n parts_embeddings = self.parts_after_pooling_dim_reduce(parts_embeddings) # [N, M, D]\n\n # Concatenated part features\n concat_parts_embeddings = parts_embeddings.flatten(1, 2) # [N, K*D]\n\n # Identity classification scores\n bn_global_embeddings, global_cls_score = self.global_identity_classifier(global_embeddings) # [N, D], [N, num_classes]\n bn_background_embeddings, background_cls_score = self.background_identity_classifier(background_embeddings) # [N, D], [N, num_classes]\n bn_foreground_embeddings, foreground_cls_score = self.foreground_identity_classifier(foreground_embeddings) # [N, D], [N, num_classes]\n bn_concat_parts_embeddings, concat_parts_cls_score = self.concat_parts_identity_classifier(concat_parts_embeddings) # [N, K*D], [N, num_classes]\n bn_parts_embeddings, parts_cls_score = self.parts_identity_classification(self.dim_reduce_output, N, parts_embeddings) # [N, K, D], [N, K, num_classes]\n\n # Outputs\n embeddings = {\n GLOBAL: global_embeddings, # [N, D]\n BACKGROUND: background_embeddings, # [N, D]\n FOREGROUND: foreground_embeddings, # [N, D]\n CONCAT_PARTS: concat_parts_embeddings, # [N, K*D]\n PARTS: parts_embeddings, # [N, K, D]\n BN_GLOBAL: bn_global_embeddings, # [N, D]\n BN_BACKGROUND: bn_background_embeddings, # [N, D]\n BN_FOREGROUND: bn_foreground_embeddings, # [N, D]\n BN_CONCAT_PARTS: bn_concat_parts_embeddings, # [N, K*D]\n BN_PARTS: bn_parts_embeddings, # [N, K, D]\n }\n\n visibility_scores = {\n GLOBAL: global_visibility, # [N]\n BACKGROUND: background_visibility, # [N]\n FOREGROUND: foreground_visibility, # [N]\n CONCAT_PARTS: concat_parts_visibility, # [N]\n PARTS: parts_visibility, # [N, K]\n }\n\n id_cls_scores = {\n GLOBAL: global_cls_score, # [N, num_classes]\n BACKGROUND: background_cls_score, # [N, num_classes]\n FOREGROUND: foreground_cls_score, # [N, num_classes]\n CONCAT_PARTS: concat_parts_cls_score, # [N, num_classes]\n PARTS: parts_cls_score, # [N, K, num_classes]\n }\n\n masks = {\n GLOBAL: global_masks, # [N, Hf, Wf]\n BACKGROUND: background_masks, # [N, Hf, Wf]\n FOREGROUND: foreground_masks, # [N, Hf, Wf]\n CONCAT_PARTS: foreground_masks, # [N, Hf, Wf]\n PARTS: parts_masks, # [N, K, Hf, Wf]\n }\n\n return embeddings, visibility_scores, id_cls_scores, pixels_cls_scores, spatial_features, masks\n\n def parts_identity_classification(self, D, N, parts_embeddings):\n if self.shared_parts_id_classifier:\n # apply the same classifier on each part embedding, classifier weights are therefore shared across parts\n parts_embeddings = parts_embeddings.flatten(0, 1) # [N*K, D]\n bn_part_embeddings, part_cls_score = self.parts_identity_classifier(parts_embeddings)\n bn_part_embeddings = bn_part_embeddings.view([N, self.parts_num, D])\n part_cls_score = part_cls_score.view([N, self.parts_num, -1])\n else:\n # apply K classifiers on each of the K part embedding, each part has therefore it's own classifier weights\n scores = []\n embeddings = []\n for i, parts_identity_classifier in enumerate(self.parts_identity_classifier):\n bn_part_embeddings, part_cls_score = parts_identity_classifier(parts_embeddings[:, i])\n scores.append(part_cls_score.unsqueeze(1))\n embeddings.append(bn_part_embeddings.unsqueeze(1))\n part_cls_score = torch.cat(scores, 1)\n bn_part_embeddings = torch.cat(embeddings, 1)\n\n return bn_part_embeddings, part_cls_score\n\n\n########################################\n# Dimensionality reduction layers #\n########################################\n\nclass BeforePoolingDimReduceLayer(nn.Module):\n def __init__(self, input_dim, output_dim):\n super(BeforePoolingDimReduceLayer, self).__init__()\n layers = []\n layers.append(\n nn.Conv2d(\n input_dim, output_dim, 1, stride=1, padding=0\n )\n )\n layers.append(nn.BatchNorm2d(output_dim))\n layers.append(nn.ReLU(inplace=True))\n\n self.layers = nn.Sequential(*layers)\n self._init_params()\n\n def forward(self, x):\n return self.layers(x)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n\nclass AfterPoolingDimReduceLayer(nn.Module):\n def __init__(self, input_dim, output_dim, dropout_p=None):\n super(AfterPoolingDimReduceLayer, self).__init__()\n # dim reduction used in ResNet and PCB\n layers = []\n layers.append(\n nn.Linear(\n input_dim, output_dim, bias=True\n )\n )\n layers.append(nn.BatchNorm1d(output_dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.opout(p=dropout_p))\n\n self.layers = nn.Sequential(*layers)\n self._init_params()\n\n def forward(self, x):\n if len(x.size()) == 3:\n N, K, _ = x.size() # [N, K, input_dim]\n x = x.flatten(0, 1)\n x = self.layers(x)\n x = x.view(N, K, -1)\n else:\n x = self.layers(x)\n return x\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n\n########################################\n# Classifiers #\n########################################\n\nclass PixelToPartClassifier(nn.Module):\n def __init__(self, dim_reduce_output, parts_num):\n super(PixelToPartClassifier, self).__init__()\n self.bn = torch.nn.BatchNorm2d(dim_reduce_output)\n self.classifier = nn.Conv2d(in_channels=dim_reduce_output, out_channels=parts_num + 1, kernel_size=1, stride=1, padding=0)\n self._init_params()\n\n def forward(self, x):\n x = self.bn(x)\n return self.classifier(x)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Conv2d):\n nn.init.normal_(m.weight, 0, 0.001) # ResNet = 0.01, Bof and ISP-reid = 0.001\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n\nclass BNClassifier(nn.Module):\n # Source: https://github.com/upgirlnana/Pytorch-Person-REID-Baseline-Bag-of-Tricks\n def __init__(self, in_dim, class_num):\n super(BNClassifier, self).__init__()\n\n self.in_dim = in_dim\n self.class_num = class_num\n\n self.bn = nn.BatchNorm1d(self.in_dim)\n self.bn.bias.requires_grad_(False) # BoF: this doesn't have a big impact on perf according to author on github\n self.classifier = nn.Linear(self.in_dim, self.class_num, bias=False)\n\n self._init_params()\n\n def forward(self, x):\n feature = self.bn(x)\n cls_score = self.classifier(feature)\n return feature, cls_score\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.001) # ResNet = 0.01, Bof and ISP-reid = 0.001\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n\n########################################\n# Pooling heads #\n########################################\n\ndef init_part_attention_pooling_head(normalization, pooling, dim_reduce_output):\n if pooling == 'gap':\n parts_attention_pooling_head = GlobalAveragePoolingHead(dim_reduce_output, normalization)\n elif pooling == 'gmp':\n parts_attention_pooling_head = GlobalMaxPoolingHead(dim_reduce_output, normalization)\n elif pooling == 'gwap':\n parts_attention_pooling_head = GlobalWeightedAveragePoolingHead(dim_reduce_output, normalization)\n else:\n raise ValueError('pooling type {} not supported'.format(pooling))\n return parts_attention_pooling_head\n\n\nclass GlobalMaskWeightedPoolingHead(nn.Module):\n def __init__(self, depth, normalization='identity'):\n super().__init__()\n if normalization == 'identity':\n self.normalization = nn.Identity()\n elif normalization == 'batch_norm_3d':\n self.normalization = torch.nn.BatchNorm3d(depth, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n elif normalization == 'batch_norm_2d':\n self.normalization = torch.nn.BatchNorm2d(depth, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n elif normalization == 'batch_norm_1d':\n self.normalization = torch.nn.BatchNorm1d(depth, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n else:\n raise ValueError('normalization type {} not supported'.format(normalization))\n\n def forward(self, features, part_masks):\n part_masks = torch.unsqueeze(part_masks, 2)\n features = torch.unsqueeze(features, 1)\n parts_features = torch.mul(part_masks, features)\n\n N, M, _, _, _ = parts_features.size()\n parts_features = parts_features.flatten(0, 1)\n parts_features = self.normalization(parts_features)\n parts_features = self.global_pooling(parts_features)\n parts_features = parts_features.view(N, M, -1)\n return parts_features\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.BatchNorm1d) or isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm3d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.001) # ResNet = 0.01, Bof and ISP-reid = 0.001\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n\nclass GlobalMaxPoolingHead(GlobalMaskWeightedPoolingHead):\n global_pooling = nn.AdaptiveMaxPool2d((1, 1))\n\n\nclass GlobalAveragePoolingHead(GlobalMaskWeightedPoolingHead):\n global_pooling = nn.AdaptiveAvgPool2d((1, 1))\n\n\nclass GlobalWeightedAveragePoolingHead(GlobalMaskWeightedPoolingHead):\n def forward(self, features, part_masks):\n part_masks = torch.unsqueeze(part_masks, 2)\n features = torch.unsqueeze(features, 1)\n parts_features = torch.mul(part_masks, features)\n\n N, M, _, _, _ = parts_features.size()\n parts_features = parts_features.flatten(0, 1)\n parts_features = self.normalization(parts_features)\n parts_features = torch.sum(parts_features, dim=(-2, -1))\n part_masks_sum = torch.sum(part_masks.flatten(0, 1), dim=(-2, -1))\n part_masks_sum = torch.clamp(part_masks_sum, min=1e-6)\n parts_features_avg = torch.div(parts_features, part_masks_sum)\n parts_features = parts_features_avg.view(N, M, -1)\n return parts_features\n\n\n########################################\n# Constructors #\n########################################\n\ndef bpbreid(num_classes, loss='part_based', pretrained=True, config=None, **kwargs):\n model = BPBreID(\n num_classes,\n pretrained,\n loss,\n config.model.bpbreid,\n **kwargs\n )\n return model\n\n\ndef pcb(num_classes, loss='part_based', pretrained=True, config=None, **kwargs):\n config.model.bpbreid.learnable_attention_enabled = False\n model = BPBreID(\n num_classes,\n pretrained,\n loss,\n config.model.bpbreid,\n horizontal_stipes=True,\n config=config,\n **kwargs\n )\n return model\n\n\ndef bot(num_classes, loss='part_based', pretrained=True, config=None, **kwargs):\n config.model.bpbreid.masks.parts_num = 1\n config.model.bpbreid.learnable_attention_enabled = False\n model = BPBreID(\n num_classes,\n pretrained,\n loss,\n config.model.bpbreid,\n horizontal_stipes=True,\n config=config,\n **kwargs\n )\n return model\n" }, { "path": "torchreid/models/compact_bilinear_pooling.py", "content": "import types\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import Function\n\n\ndef CountSketchFn_forward(h, s, output_size, x, force_cpu_scatter_add=False):\n x_size = tuple(x.size())\n\n s_view = (1,) * (len(x_size) - 1) + (x_size[-1],)\n\n out_size = x_size[:-1] + (output_size,)\n\n # Broadcast s and compute x * s\n s = s.view(s_view)\n xs = x * s\n\n # Broadcast h then compute h:\n # out[h_i] += x_i * s_i\n h = h.view(s_view).expand(x_size)\n\n if force_cpu_scatter_add:\n out = x.new(*out_size).zero_().cpu()\n return out.scatter_add_(-1, h.cpu(), xs.cpu()).cuda()\n else:\n out = x.new(*out_size).zero_()\n return out.scatter_add_(-1, h, xs)\n\n\ndef CountSketchFn_backward(h, s, x_size, grad_output):\n s_view = (1,) * (len(x_size) - 1) + (x_size[-1],)\n\n s = s.view(s_view)\n h = h.view(s_view).expand(x_size)\n\n grad_x = grad_output.gather(-1, h)\n grad_x = grad_x * s\n return grad_x\n\n\nclass CountSketchFn(Function):\n\n @staticmethod\n def forward(ctx, h, s, output_size, x, force_cpu_scatter_add=False):\n x_size = tuple(x.size())\n\n ctx.save_for_backward(h, s)\n ctx.x_size = tuple(x.size())\n\n return CountSketchFn_forward(h, s, output_size, x, force_cpu_scatter_add)\n\n @staticmethod\n def backward(ctx, grad_output):\n h, s = ctx.saved_variables\n\n grad_x = CountSketchFn_backward(h, s, ctx.x_size, grad_output)\n return None, None, None, grad_x\n\n\nclass CountSketch(nn.Module):\n r\"\"\"Compute the count sketch over an input signal.\n\n .. math::\n\n out_j = \\sum_{i : j = h_i} s_i x_i\n\n Args:\n input_size (int): Number of channels in the input array\n output_size (int): Number of channels in the output sketch\n h (array, optional): Optional array of size input_size of indices in the range [0,output_size]\n s (array, optional): Optional array of size input_size of -1 and 1.\n\n .. note::\n\n If h and s are None, they will be automatically be generated using LongTensor.random_.\n\n Shape:\n - Input: (...,input_size)\n - Output: (...,output_size)\n\n References:\n Yang Gao et al. \"Compact Bilinear Pooling\" in Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (2016).\n Akira Fukui et al. \"Multimodal Compact Bilinear Pooling for Visual Question Answering and Visual Grounding\", arXiv:1606.01847 (2016).\n \"\"\"\n\n def __init__(self, input_size, output_size, h=None, s=None):\n super(CountSketch, self).__init__()\n\n self.input_size = input_size\n self.output_size = output_size\n\n if h is None:\n h = torch.LongTensor(input_size).random_(0, output_size)\n if s is None:\n s = 2 * torch.Tensor(input_size).random_(0, 2) - 1\n\n # The Variable h being a list of indices,\n # If the type of this module is changed (e.g. float to double),\n # the variable h should remain a LongTensor\n # therefore we force float() and double() to be no-ops on the variable h.\n def identity(self):\n return self\n\n h.float = types.MethodType(identity, h)\n h.double = types.MethodType(identity, h)\n\n self.register_buffer('h', h)\n self.register_buffer('s', s)\n\n def forward(self, x):\n x_size = list(x.size())\n\n assert (x_size[-1] == self.input_size)\n\n return CountSketchFn.apply(self.h, self.s, self.output_size, x)\n\n\ndef ComplexMultiply_forward(X_re, X_im, Y_re, Y_im):\n Z_re = torch.addcmul(X_re * Y_re, -1, X_im, Y_im)\n Z_im = torch.addcmul(X_re * Y_im, 1, X_im, Y_re)\n return Z_re, Z_im\n\n\ndef ComplexMultiply_backward(X_re, X_im, Y_re, Y_im, grad_Z_re, grad_Z_im):\n grad_X_re = torch.addcmul(grad_Z_re * Y_re, 1, grad_Z_im, Y_im)\n grad_X_im = torch.addcmul(grad_Z_im * Y_re, -1, grad_Z_re, Y_im)\n grad_Y_re = torch.addcmul(grad_Z_re * X_re, 1, grad_Z_im, X_im)\n grad_Y_im = torch.addcmul(grad_Z_im * X_re, -1, grad_Z_re, X_im)\n return grad_X_re, grad_X_im, grad_Y_re, grad_Y_im\n\n\nclass ComplexMultiply(torch.autograd.Function):\n\n @staticmethod\n def forward(ctx, X_re, X_im, Y_re, Y_im):\n ctx.save_for_backward(X_re, X_im, Y_re, Y_im)\n return ComplexMultiply_forward(X_re, X_im, Y_re, Y_im)\n\n @staticmethod\n def backward(ctx, grad_Z_re, grad_Z_im):\n X_re, X_im, Y_re, Y_im = ctx.saved_tensors\n return ComplexMultiply_backward(X_re, X_im, Y_re, Y_im, grad_Z_re, grad_Z_im)\n\n\nclass CompactBilinearPoolingFn(Function):\n\n @staticmethod\n def forward(ctx, h1, s1, h2, s2, output_size, x, y, force_cpu_scatter_add=False):\n ctx.save_for_backward(h1, s1, h2, s2, x, y)\n ctx.x_size = tuple(x.size())\n ctx.y_size = tuple(y.size())\n ctx.force_cpu_scatter_add = force_cpu_scatter_add\n ctx.output_size = output_size\n\n # Compute the count sketch of each input\n px = CountSketchFn_forward(h1, s1, output_size, x, force_cpu_scatter_add)\n fx = torch.rfft(px, 1)\n re_fx = fx.select(-1, 0)\n im_fx = fx.select(-1, 1)\n del px\n py = CountSketchFn_forward(h2, s2, output_size, y, force_cpu_scatter_add)\n fy = torch.rfft(py, 1)\n re_fy = fy.select(-1, 0)\n im_fy = fy.select(-1, 1)\n del py\n\n # Convolution of the two sketch using an FFT.\n # Compute the FFT of each sketch\n\n # Complex multiplication\n re_prod, im_prod = ComplexMultiply_forward(re_fx, im_fx, re_fy, im_fy)\n\n # Back to real domain\n # The imaginary part should be zero's\n re = torch.irfft(torch.stack((re_prod, im_prod), re_prod.dim()), 1, signal_sizes=(output_size,))\n\n return re\n\n @staticmethod\n def backward(ctx, grad_output):\n h1, s1, h2, s2, x, y = ctx.saved_tensors\n\n # Recompute part of the forward pass to get the input to the complex product\n # Compute the count sketch of each input\n px = CountSketchFn_forward(h1, s1, ctx.output_size, x, ctx.force_cpu_scatter_add)\n py = CountSketchFn_forward(h2, s2, ctx.output_size, y, ctx.force_cpu_scatter_add)\n\n # Then convert the output to Fourier domain\n grad_output = grad_output.contiguous()\n grad_prod = torch.rfft(grad_output, 1)\n grad_re_prod = grad_prod.select(-1, 0)\n grad_im_prod = grad_prod.select(-1, 1)\n\n # Compute the gradient of x first then y\n\n # Gradient of x\n # Recompute fy\n fy = torch.rfft(py, 1)\n re_fy = fy.select(-1, 0)\n im_fy = fy.select(-1, 1)\n del py\n # Compute the gradient of fx, then back to temporal space\n grad_re_fx = torch.addcmul(grad_re_prod * re_fy, 1, grad_im_prod, im_fy)\n grad_im_fx = torch.addcmul(grad_im_prod * re_fy, -1, grad_re_prod, im_fy)\n grad_fx = torch.irfft(torch.stack((grad_re_fx, grad_im_fx), grad_re_fx.dim()), 1,\n signal_sizes=(ctx.output_size,))\n # Finally compute the gradient of x\n grad_x = CountSketchFn_backward(h1, s1, ctx.x_size, grad_fx)\n del re_fy, im_fy, grad_re_fx, grad_im_fx, grad_fx\n\n # Gradient of y\n # Recompute fx\n fx = torch.rfft(px, 1)\n re_fx = fx.select(-1, 0)\n im_fx = fx.select(-1, 1)\n del px\n # Compute the gradient of fy, then back to temporal space\n grad_re_fy = torch.addcmul(grad_re_prod * re_fx, 1, grad_im_prod, im_fx)\n grad_im_fy = torch.addcmul(grad_im_prod * re_fx, -1, grad_re_prod, im_fx)\n grad_fy = torch.irfft(torch.stack((grad_re_fy, grad_im_fy), grad_re_fy.dim()), 1,\n signal_sizes=(ctx.output_size,))\n # Finally compute the gradient of y\n grad_y = CountSketchFn_backward(h2, s2, ctx.y_size, grad_fy)\n del re_fx, im_fx, grad_re_fy, grad_im_fy, grad_fy\n\n return None, None, None, None, None, grad_x, grad_y, None\n\n\nclass CompactBilinearPooling(nn.Module):\n r\"\"\"Compute the compact bilinear pooling between two input array x and y\n\n .. math::\n\n out = \\Psi (x,h_1,s_1) \\ast \\Psi (y,h_2,s_2)\n\n Args:\n input_size1 (int): Number of channels in the first input array\n input_size2 (int): Number of channels in the second input array\n output_size (int): Number of channels in the output array\n h1 (array, optional): Optional array of size input_size of indices in the range [0,output_size]\n s1 (array, optional): Optional array of size input_size of -1 and 1.\n h2 (array, optional): Optional array of size input_size of indices in the range [0,output_size]\n s2 (array, optional): Optional array of size input_size of -1 and 1.\n force_cpu_scatter_add (boolean, optional): Force the scatter_add operation to run on CPU for testing purposes\n\n .. note::\n\n If h1, s1, s2, h2 are None, they will be automatically be generated using LongTensor.random_.\n\n Shape:\n - Input 1: (...,input_size1)\n - Input 2: (...,input_size2)\n - Output: (...,output_size)\n\n References:\n Yang Gao et al. \"Compact Bilinear Pooling\" in Proceedings of IEEE Conference on Computer Vision and Pattern Recognition (2016).\n Akira Fukui et al. \"Multimodal Compact Bilinear Pooling for Visual Question Answering and Visual Grounding\", arXiv:1606.01847 (2016).\n \"\"\"\n\n def __init__(self, input1_size, input2_size, output_size, h1=None, s1=None, h2=None, s2=None,\n force_cpu_scatter_add=False):\n super(CompactBilinearPooling, self).__init__()\n self.add_module('sketch1', CountSketch(input1_size, output_size, h1, s1))\n self.add_module('sketch2', CountSketch(input2_size, output_size, h2, s2))\n self.output_size = output_size\n self.force_cpu_scatter_add = force_cpu_scatter_add\n\n def forward(self, x, y=None):\n if y is None:\n y = x\n\n return CompactBilinearPoolingFn.apply(self.sketch1.h, self.sketch1.s, self.sketch2.h, self.sketch2.s,\n self.output_size, x, y, self.force_cpu_scatter_add)\n\n\n\nif __name__ == '__main__':\n\n input_size = 2048\n output_size = 16000\n mcb = CompactBilinearPooling(input_size, input_size, output_size).cuda()\n x = torch.rand(4, input_size).cuda()\n y = torch.rand(4, input_size).cuda()\n\n z = mcb(x, y)\n print(z)\n" }, { "path": "torchreid/models/densenet.py", "content": "\"\"\"\nCode source: https://github.com/pytorch/vision\n\"\"\"\nfrom __future__ import division, absolute_import\nimport re\nfrom collections import OrderedDict\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\nfrom torch.utils import model_zoo\n\n__all__ = [\n 'densenet121', 'densenet169', 'densenet201', 'densenet161',\n 'densenet121_fc512'\n]\n\nmodel_urls = {\n 'densenet121':\n 'https://download.pytorch.org/models/densenet121-a639ec97.pth',\n 'densenet169':\n 'https://download.pytorch.org/models/densenet169-b2777c0a.pth',\n 'densenet201':\n 'https://download.pytorch.org/models/densenet201-c1103571.pth',\n 'densenet161':\n 'https://download.pytorch.org/models/densenet161-8d451a50.pth',\n}\n\n\nclass _DenseLayer(nn.Sequential):\n\n def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):\n super(_DenseLayer, self).__init__()\n self.add_module('norm1', nn.BatchNorm2d(num_input_features)),\n self.add_module('relu1', nn.ReLU(inplace=True)),\n self.add_module(\n 'conv1',\n nn.Conv2d(\n num_input_features,\n bn_size * growth_rate,\n kernel_size=1,\n stride=1,\n bias=False\n )\n ),\n self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),\n self.add_module('relu2', nn.ReLU(inplace=True)),\n self.add_module(\n 'conv2',\n nn.Conv2d(\n bn_size * growth_rate,\n growth_rate,\n kernel_size=3,\n stride=1,\n padding=1,\n bias=False\n )\n ),\n self.drop_rate = drop_rate\n\n def forward(self, x):\n new_features = super(_DenseLayer, self).forward(x)\n if self.drop_rate > 0:\n new_features = F.dropout(\n new_features, p=self.drop_rate, training=self.training\n )\n return torch.cat([x, new_features], 1)\n\n\nclass _DenseBlock(nn.Sequential):\n\n def __init__(\n self, num_layers, num_input_features, bn_size, growth_rate, drop_rate\n ):\n super(_DenseBlock, self).__init__()\n for i in range(num_layers):\n layer = _DenseLayer(\n num_input_features + i*growth_rate, growth_rate, bn_size,\n drop_rate\n )\n self.add_module('denselayer%d' % (i+1), layer)\n\n\nclass _Transition(nn.Sequential):\n\n def __init__(self, num_input_features, num_output_features):\n super(_Transition, self).__init__()\n self.add_module('norm', nn.BatchNorm2d(num_input_features))\n self.add_module('relu', nn.ReLU(inplace=True))\n self.add_module(\n 'conv',\n nn.Conv2d(\n num_input_features,\n num_output_features,\n kernel_size=1,\n stride=1,\n bias=False\n )\n )\n self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))\n\n\nclass DenseNet(nn.Module):\n \"\"\"Densely connected network.\n \n Reference:\n Huang et al. Densely Connected Convolutional Networks. CVPR 2017.\n\n Public keys:\n - ``densenet121``: DenseNet121.\n - ``densenet169``: DenseNet169.\n - ``densenet201``: DenseNet201.\n - ``densenet161``: DenseNet161.\n - ``densenet121_fc512``: DenseNet121 + FC.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n growth_rate=32,\n block_config=(6, 12, 24, 16),\n num_init_features=64,\n bn_size=4,\n drop_rate=0,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n ):\n\n super(DenseNet, self).__init__()\n self.loss = loss\n\n # First convolution\n self.features = nn.Sequential(\n OrderedDict(\n [\n (\n 'conv0',\n nn.Conv2d(\n 3,\n num_init_features,\n kernel_size=7,\n stride=2,\n padding=3,\n bias=False\n )\n ),\n ('norm0', nn.BatchNorm2d(num_init_features)),\n ('relu0', nn.ReLU(inplace=True)),\n (\n 'pool0',\n nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n ),\n ]\n )\n )\n\n # Each denseblock\n num_features = num_init_features\n for i, num_layers in enumerate(block_config):\n block = _DenseBlock(\n num_layers=num_layers,\n num_input_features=num_features,\n bn_size=bn_size,\n growth_rate=growth_rate,\n drop_rate=drop_rate\n )\n self.features.add_module('denseblock%d' % (i+1), block)\n num_features = num_features + num_layers*growth_rate\n if i != len(block_config) - 1:\n trans = _Transition(\n num_input_features=num_features,\n num_output_features=num_features // 2\n )\n self.features.add_module('transition%d' % (i+1), trans)\n num_features = num_features // 2\n\n # Final batch norm\n self.features.add_module('norm5', nn.BatchNorm2d(num_features))\n\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.feature_dim = num_features\n self.fc = self._construct_fc_layer(fc_dims, num_features, dropout_p)\n\n # Linear layer\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer.\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n f = self.features(x)\n f = F.relu(f, inplace=True)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if self.fc is not None:\n v = self.fc(v)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n\n # '.'s are no longer allowed in module names, but pervious _DenseLayer\n # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.\n # They are also in the checkpoints in model_urls. This pattern is used\n # to find such keys.\n pattern = re.compile(\n r'^(.*denselayer\\d+\\.(?:norm|relu|conv))\\.((?:[12])\\.(?:weight|bias|running_mean|running_var))$'\n )\n for key in list(pretrain_dict.keys()):\n res = pattern.match(key)\n if res:\n new_key = res.group(1) + res.group(2)\n pretrain_dict[new_key] = pretrain_dict[key]\n del pretrain_dict[key]\n\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\n\"\"\"\nDense network configurations:\n--\ndensenet121: num_init_features=64, growth_rate=32, block_config=(6, 12, 24, 16)\ndensenet169: num_init_features=64, growth_rate=32, block_config=(6, 12, 32, 32)\ndensenet201: num_init_features=64, growth_rate=32, block_config=(6, 12, 48, 32)\ndensenet161: num_init_features=96, growth_rate=48, block_config=(6, 12, 36, 24)\n\"\"\"\n\n\ndef densenet121(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = DenseNet(\n num_classes=num_classes,\n loss=loss,\n num_init_features=64,\n growth_rate=32,\n block_config=(6, 12, 24, 16),\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['densenet121'])\n return model\n\n\ndef densenet169(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = DenseNet(\n num_classes=num_classes,\n loss=loss,\n num_init_features=64,\n growth_rate=32,\n block_config=(6, 12, 32, 32),\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['densenet169'])\n return model\n\n\ndef densenet201(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = DenseNet(\n num_classes=num_classes,\n loss=loss,\n num_init_features=64,\n growth_rate=32,\n block_config=(6, 12, 48, 32),\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['densenet201'])\n return model\n\n\ndef densenet161(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = DenseNet(\n num_classes=num_classes,\n loss=loss,\n num_init_features=96,\n growth_rate=48,\n block_config=(6, 12, 36, 24),\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['densenet161'])\n return model\n\n\ndef densenet121_fc512(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = DenseNet(\n num_classes=num_classes,\n loss=loss,\n num_init_features=64,\n growth_rate=32,\n block_config=(6, 12, 24, 16),\n fc_dims=[512],\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['densenet121'])\n return model\n" }, { "path": "torchreid/models/hacnn.py", "content": "from __future__ import division, absolute_import\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['HACNN']\n\n\nclass ConvBlock(nn.Module):\n \"\"\"Basic convolutional block.\n \n convolution + batch normalization + relu.\n\n Args:\n in_c (int): number of input channels.\n out_c (int): number of output channels.\n k (int or tuple): kernel size.\n s (int or tuple): stride.\n p (int or tuple): padding.\n \"\"\"\n\n def __init__(self, in_c, out_c, k, s=1, p=0):\n super(ConvBlock, self).__init__()\n self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p)\n self.bn = nn.BatchNorm2d(out_c)\n\n def forward(self, x):\n return F.relu(self.bn(self.conv(x)))\n\n\nclass InceptionA(nn.Module):\n\n def __init__(self, in_channels, out_channels):\n super(InceptionA, self).__init__()\n mid_channels = out_channels // 4\n\n self.stream1 = nn.Sequential(\n ConvBlock(in_channels, mid_channels, 1),\n ConvBlock(mid_channels, mid_channels, 3, p=1),\n )\n self.stream2 = nn.Sequential(\n ConvBlock(in_channels, mid_channels, 1),\n ConvBlock(mid_channels, mid_channels, 3, p=1),\n )\n self.stream3 = nn.Sequential(\n ConvBlock(in_channels, mid_channels, 1),\n ConvBlock(mid_channels, mid_channels, 3, p=1),\n )\n self.stream4 = nn.Sequential(\n nn.AvgPool2d(3, stride=1, padding=1),\n ConvBlock(in_channels, mid_channels, 1),\n )\n\n def forward(self, x):\n s1 = self.stream1(x)\n s2 = self.stream2(x)\n s3 = self.stream3(x)\n s4 = self.stream4(x)\n y = torch.cat([s1, s2, s3, s4], dim=1)\n return y\n\n\nclass InceptionB(nn.Module):\n\n def __init__(self, in_channels, out_channels):\n super(InceptionB, self).__init__()\n mid_channels = out_channels // 4\n\n self.stream1 = nn.Sequential(\n ConvBlock(in_channels, mid_channels, 1),\n ConvBlock(mid_channels, mid_channels, 3, s=2, p=1),\n )\n self.stream2 = nn.Sequential(\n ConvBlock(in_channels, mid_channels, 1),\n ConvBlock(mid_channels, mid_channels, 3, p=1),\n ConvBlock(mid_channels, mid_channels, 3, s=2, p=1),\n )\n self.stream3 = nn.Sequential(\n nn.MaxPool2d(3, stride=2, padding=1),\n ConvBlock(in_channels, mid_channels * 2, 1),\n )\n\n def forward(self, x):\n s1 = self.stream1(x)\n s2 = self.stream2(x)\n s3 = self.stream3(x)\n y = torch.cat([s1, s2, s3], dim=1)\n return y\n\n\nclass SpatialAttn(nn.Module):\n \"\"\"Spatial Attention (Sec. 3.1.I.1)\"\"\"\n\n def __init__(self):\n super(SpatialAttn, self).__init__()\n self.conv1 = ConvBlock(1, 1, 3, s=2, p=1)\n self.conv2 = ConvBlock(1, 1, 1)\n\n def forward(self, x):\n # global cross-channel averaging\n x = x.mean(1, keepdim=True)\n # 3-by-3 conv\n x = self.conv1(x)\n # bilinear resizing\n x = F.upsample(\n x, (x.size(2) * 2, x.size(3) * 2),\n mode='bilinear',\n align_corners=True\n )\n # scaling conv\n x = self.conv2(x)\n return x\n\n\nclass ChannelAttn(nn.Module):\n \"\"\"Channel Attention (Sec. 3.1.I.2)\"\"\"\n\n def __init__(self, in_channels, reduction_rate=16):\n super(ChannelAttn, self).__init__()\n assert in_channels % reduction_rate == 0\n self.conv1 = ConvBlock(in_channels, in_channels // reduction_rate, 1)\n self.conv2 = ConvBlock(in_channels // reduction_rate, in_channels, 1)\n\n def forward(self, x):\n # squeeze operation (global average pooling)\n x = F.avg_pool2d(x, x.size()[2:])\n # excitation operation (2 conv layers)\n x = self.conv1(x)\n x = self.conv2(x)\n return x\n\n\nclass SoftAttn(nn.Module):\n \"\"\"Soft Attention (Sec. 3.1.I)\n \n Aim: Spatial Attention + Channel Attention\n \n Output: attention maps with shape identical to input.\n \"\"\"\n\n def __init__(self, in_channels):\n super(SoftAttn, self).__init__()\n self.spatial_attn = SpatialAttn()\n self.channel_attn = ChannelAttn(in_channels)\n self.conv = ConvBlock(in_channels, in_channels, 1)\n\n def forward(self, x):\n y_spatial = self.spatial_attn(x)\n y_channel = self.channel_attn(x)\n y = y_spatial * y_channel\n y = torch.sigmoid(self.conv(y))\n return y\n\n\nclass HardAttn(nn.Module):\n \"\"\"Hard Attention (Sec. 3.1.II)\"\"\"\n\n def __init__(self, in_channels):\n super(HardAttn, self).__init__()\n self.fc = nn.Linear(in_channels, 4 * 2)\n self.init_params()\n\n def init_params(self):\n self.fc.weight.data.zero_()\n self.fc.bias.data.copy_(\n torch.tensor(\n [0, -0.75, 0, -0.25, 0, 0.25, 0, 0.75], dtype=torch.float\n )\n )\n\n def forward(self, x):\n # squeeze operation (global average pooling)\n x = F.avg_pool2d(x, x.size()[2:]).view(x.size(0), x.size(1))\n # predict transformation parameters\n theta = torch.tanh(self.fc(x))\n theta = theta.view(-1, 4, 2)\n return theta\n\n\nclass HarmAttn(nn.Module):\n \"\"\"Harmonious Attention (Sec. 3.1)\"\"\"\n\n def __init__(self, in_channels):\n super(HarmAttn, self).__init__()\n self.soft_attn = SoftAttn(in_channels)\n self.hard_attn = HardAttn(in_channels)\n\n def forward(self, x):\n y_soft_attn = self.soft_attn(x)\n theta = self.hard_attn(x)\n return y_soft_attn, theta\n\n\nclass HACNN(nn.Module):\n \"\"\"Harmonious Attention Convolutional Neural Network.\n\n Reference:\n Li et al. Harmonious Attention Network for Person Re-identification. CVPR 2018.\n\n Public keys:\n - ``hacnn``: HACNN.\n \"\"\"\n\n # Args:\n # num_classes (int): number of classes to predict\n # nchannels (list): number of channels AFTER concatenation\n # feat_dim (int): feature dimension for a single stream\n # learn_region (bool): whether to learn region features (i.e. local branch)\n\n def __init__(\n self,\n num_classes,\n loss='softmax',\n nchannels=[128, 256, 384],\n feat_dim=512,\n learn_region=True,\n use_gpu=True,\n **kwargs\n ):\n super(HACNN, self).__init__()\n self.loss = loss\n self.learn_region = learn_region\n self.use_gpu = use_gpu\n\n self.conv = ConvBlock(3, 32, 3, s=2, p=1)\n\n # Construct Inception + HarmAttn blocks\n # ============== Block 1 ==============\n self.inception1 = nn.Sequential(\n InceptionA(32, nchannels[0]),\n InceptionB(nchannels[0], nchannels[0]),\n )\n self.ha1 = HarmAttn(nchannels[0])\n\n # ============== Block 2 ==============\n self.inception2 = nn.Sequential(\n InceptionA(nchannels[0], nchannels[1]),\n InceptionB(nchannels[1], nchannels[1]),\n )\n self.ha2 = HarmAttn(nchannels[1])\n\n # ============== Block 3 ==============\n self.inception3 = nn.Sequential(\n InceptionA(nchannels[1], nchannels[2]),\n InceptionB(nchannels[2], nchannels[2]),\n )\n self.ha3 = HarmAttn(nchannels[2])\n\n self.fc_global = nn.Sequential(\n nn.Linear(nchannels[2], feat_dim),\n nn.BatchNorm1d(feat_dim),\n nn.ReLU(),\n )\n self.classifier_global = nn.Linear(feat_dim, num_classes)\n\n if self.learn_region:\n self.init_scale_factors()\n self.local_conv1 = InceptionB(32, nchannels[0])\n self.local_conv2 = InceptionB(nchannels[0], nchannels[1])\n self.local_conv3 = InceptionB(nchannels[1], nchannels[2])\n self.fc_local = nn.Sequential(\n nn.Linear(nchannels[2] * 4, feat_dim),\n nn.BatchNorm1d(feat_dim),\n nn.ReLU(),\n )\n self.classifier_local = nn.Linear(feat_dim, num_classes)\n self.feat_dim = feat_dim * 2\n else:\n self.feat_dim = feat_dim\n\n def init_scale_factors(self):\n # initialize scale factors (s_w, s_h) for four regions\n self.scale_factors = []\n self.scale_factors.append(\n torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n )\n self.scale_factors.append(\n torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n )\n self.scale_factors.append(\n torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n )\n self.scale_factors.append(\n torch.tensor([[1, 0], [0, 0.25]], dtype=torch.float)\n )\n\n def stn(self, x, theta):\n \"\"\"Performs spatial transform\n \n x: (batch, channel, height, width)\n theta: (batch, 2, 3)\n \"\"\"\n grid = F.affine_grid(theta, x.size())\n x = F.grid_sample(x, grid)\n return x\n\n def transform_theta(self, theta_i, region_idx):\n \"\"\"Transforms theta to include (s_w, s_h), resulting in (batch, 2, 3)\"\"\"\n scale_factors = self.scale_factors[region_idx]\n theta = torch.zeros(theta_i.size(0), 2, 3)\n theta[:, :, :2] = scale_factors\n theta[:, :, -1] = theta_i\n if self.use_gpu:\n theta = theta.cuda()\n return theta\n\n def forward(self, x):\n assert x.size(2) == 160 and x.size(3) == 64, \\\n 'Input size does not match, expected (160, 64) but got ({}, {})'.format(x.size(2), x.size(3))\n x = self.conv(x)\n\n # ============== Block 1 ==============\n # global branch\n x1 = self.inception1(x)\n x1_attn, x1_theta = self.ha1(x1)\n x1_out = x1 * x1_attn\n # local branch\n if self.learn_region:\n x1_local_list = []\n for region_idx in range(4):\n x1_theta_i = x1_theta[:, region_idx, :]\n x1_theta_i = self.transform_theta(x1_theta_i, region_idx)\n x1_trans_i = self.stn(x, x1_theta_i)\n x1_trans_i = F.upsample(\n x1_trans_i, (24, 28), mode='bilinear', align_corners=True\n )\n x1_local_i = self.local_conv1(x1_trans_i)\n x1_local_list.append(x1_local_i)\n\n # ============== Block 2 ==============\n # Block 2\n # global branch\n x2 = self.inception2(x1_out)\n x2_attn, x2_theta = self.ha2(x2)\n x2_out = x2 * x2_attn\n # local branch\n if self.learn_region:\n x2_local_list = []\n for region_idx in range(4):\n x2_theta_i = x2_theta[:, region_idx, :]\n x2_theta_i = self.transform_theta(x2_theta_i, region_idx)\n x2_trans_i = self.stn(x1_out, x2_theta_i)\n x2_trans_i = F.upsample(\n x2_trans_i, (12, 14), mode='bilinear', align_corners=True\n )\n x2_local_i = x2_trans_i + x1_local_list[region_idx]\n x2_local_i = self.local_conv2(x2_local_i)\n x2_local_list.append(x2_local_i)\n\n # ============== Block 3 ==============\n # Block 3\n # global branch\n x3 = self.inception3(x2_out)\n x3_attn, x3_theta = self.ha3(x3)\n x3_out = x3 * x3_attn\n # local branch\n if self.learn_region:\n x3_local_list = []\n for region_idx in range(4):\n x3_theta_i = x3_theta[:, region_idx, :]\n x3_theta_i = self.transform_theta(x3_theta_i, region_idx)\n x3_trans_i = self.stn(x2_out, x3_theta_i)\n x3_trans_i = F.upsample(\n x3_trans_i, (6, 7), mode='bilinear', align_corners=True\n )\n x3_local_i = x3_trans_i + x2_local_list[region_idx]\n x3_local_i = self.local_conv3(x3_local_i)\n x3_local_list.append(x3_local_i)\n\n # ============== Feature generation ==============\n # global branch\n x_global = F.avg_pool2d(x3_out,\n x3_out.size()[2:]\n ).view(x3_out.size(0), x3_out.size(1))\n x_global = self.fc_global(x_global)\n # local branch\n if self.learn_region:\n x_local_list = []\n for region_idx in range(4):\n x_local_i = x3_local_list[region_idx]\n x_local_i = F.avg_pool2d(x_local_i,\n x_local_i.size()[2:]\n ).view(x_local_i.size(0), -1)\n x_local_list.append(x_local_i)\n x_local = torch.cat(x_local_list, 1)\n x_local = self.fc_local(x_local)\n\n if not self.training:\n # l2 normalization before concatenation\n if self.learn_region:\n x_global = x_global / x_global.normalization(p=2, dim=1, keepdim=True)\n x_local = x_local / x_local.normalization(p=2, dim=1, keepdim=True)\n return torch.cat([x_global, x_local], 1)\n else:\n return x_global\n\n prelogits_global = self.classifier_global(x_global)\n if self.learn_region:\n prelogits_local = self.classifier_local(x_local)\n\n if self.loss == 'softmax':\n if self.learn_region:\n return (prelogits_global, prelogits_local)\n else:\n return prelogits_global\n\n elif self.loss == 'triplet':\n if self.learn_region:\n return (prelogits_global, prelogits_local), (x_global, x_local)\n else:\n return prelogits_global, x_global\n\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n" }, { "path": "torchreid/models/hrnet.py", "content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\nimport os\nfrom pathlib import Path\n\nimport torch\nimport torch.nn as nn\nimport torch._utils\nimport torch.nn.functional as F\nfrom yacs.config import CfgNode as CN\n\n# Source:\n# https://github.com/HRNet/HRNet-Image-Classification/blob/master/lib/models/cls_hrnet.py\n# https://github.com/CASIA-IVA-Lab/ISP-reID\n__all__ = [\n 'hrnet32',\n]\n\nmodel_paths = {\n 'hrnet-w32':\n 'hrnetv2_w32_imagenet_pretrained.pth',\n}\n\n\ndef get_hrnet_config():\n _C = CN()\n\n _C.MODEL = CN()\n\n _C.MODEL.EXTRA = CN(new_allowed=True)\n _C.MODEL.EXTRA.STAGE2 = CN()\n _C.MODEL.EXTRA.STAGE2.NUM_MODULES = 1\n _C.MODEL.EXTRA.STAGE2.NUM_BRANCHES = 2\n _C.MODEL.EXTRA.STAGE2.NUM_BLOCKS = [4, 4]\n _C.MODEL.EXTRA.STAGE2.NUM_CHANNELS = [32, 64]\n _C.MODEL.EXTRA.STAGE2.BLOCK = 'BASIC'\n _C.MODEL.EXTRA.STAGE2.FUSE_METHOD = 'SUM'\n\n _C.MODEL.EXTRA.STAGE3 = CN()\n _C.MODEL.EXTRA.STAGE3.NUM_MODULES = 4\n _C.MODEL.EXTRA.STAGE3.NUM_BRANCHES = 3\n _C.MODEL.EXTRA.STAGE3.NUM_BLOCKS = [4, 4, 4]\n _C.MODEL.EXTRA.STAGE3.NUM_CHANNELS = [32, 64, 128]\n _C.MODEL.EXTRA.STAGE3.BLOCK = 'BASIC'\n _C.MODEL.EXTRA.STAGE3.FUSE_METHOD = 'SUM'\n\n _C.MODEL.EXTRA.STAGE4 = CN()\n _C.MODEL.EXTRA.STAGE4.NUM_MODULES = 3\n _C.MODEL.EXTRA.STAGE4.NUM_BRANCHES = 4\n _C.MODEL.EXTRA.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]\n _C.MODEL.EXTRA.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]\n _C.MODEL.EXTRA.STAGE4.BLOCK = 'BASIC'\n _C.MODEL.EXTRA.STAGE4.FUSE_METHOD = 'SUM'\n\n return _C\n\n\nBN_MOMENTUM = 0.1\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n padding=1, bias=False)\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n padding=1, bias=False)\n self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)\n self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,\n bias=False)\n self.bn3 = nn.BatchNorm2d(planes * self.expansion,\n momentum=BN_MOMENTUM)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass HighResolutionModule(nn.Module):\n def __init__(self, num_branches, blocks, num_blocks, num_inchannels,\n num_channels, fuse_method, multi_scale_output=True):\n super(HighResolutionModule, self).__init__()\n self._check_branches(\n num_branches, blocks, num_blocks, num_inchannels, num_channels)\n\n self.num_inchannels = num_inchannels\n self.fuse_method = fuse_method\n self.num_branches = num_branches\n\n self.multi_scale_output = multi_scale_output\n\n self.branches = self._make_branches(\n num_branches, blocks, num_blocks, num_channels)\n self.fuse_layers = self._make_fuse_layers()\n self.relu = nn.ReLU(False)\n\n def _check_branches(self, num_branches, blocks, num_blocks,\n num_inchannels, num_channels):\n if num_branches != len(num_blocks):\n error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(\n num_branches, len(num_blocks))\n print(error_msg)\n raise ValueError(error_msg)\n\n if num_branches != len(num_channels):\n error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(\n num_branches, len(num_channels))\n print(error_msg)\n raise ValueError(error_msg)\n\n if num_branches != len(num_inchannels):\n error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(\n num_branches, len(num_inchannels))\n print(error_msg)\n raise ValueError(error_msg)\n\n def _make_one_branch(self, branch_index, block, num_blocks, num_channels,\n stride=1):\n downsample = None\n if stride != 1 or \\\n self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.num_inchannels[branch_index],\n num_channels[branch_index] * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n nn.BatchNorm2d(num_channels[branch_index] * block.expansion,\n momentum=BN_MOMENTUM),\n )\n\n layers = []\n layers.append(block(self.num_inchannels[branch_index],\n num_channels[branch_index], stride, downsample))\n self.num_inchannels[branch_index] = \\\n num_channels[branch_index] * block.expansion\n for i in range(1, num_blocks[branch_index]):\n layers.append(block(self.num_inchannels[branch_index],\n num_channels[branch_index]))\n\n return nn.Sequential(*layers)\n\n def _make_branches(self, num_branches, block, num_blocks, num_channels):\n branches = []\n\n for i in range(num_branches):\n branches.append(\n self._make_one_branch(i, block, num_blocks, num_channels))\n\n return nn.ModuleList(branches)\n\n def _make_fuse_layers(self):\n if self.num_branches == 1:\n return None\n\n num_branches = self.num_branches\n num_inchannels = self.num_inchannels\n fuse_layers = []\n for i in range(num_branches if self.multi_scale_output else 1):\n fuse_layer = []\n for j in range(num_branches):\n if j > i:\n fuse_layer.append(nn.Sequential(\n nn.Conv2d(num_inchannels[j],\n num_inchannels[i],\n 1,\n 1,\n 0,\n bias=False),\n nn.BatchNorm2d(num_inchannels[i],\n momentum=BN_MOMENTUM),\n nn.Upsample(scale_factor=2 ** (j - i), mode='nearest')))\n elif j == i:\n fuse_layer.append(None)\n else:\n conv3x3s = []\n for k in range(i - j):\n if k == i - j - 1:\n num_outchannels_conv3x3 = num_inchannels[i]\n conv3x3s.append(nn.Sequential(\n nn.Conv2d(num_inchannels[j],\n num_outchannels_conv3x3,\n 3, 2, 1, bias=False),\n nn.BatchNorm2d(num_outchannels_conv3x3,\n momentum=BN_MOMENTUM)))\n else:\n num_outchannels_conv3x3 = num_inchannels[j]\n conv3x3s.append(nn.Sequential(\n nn.Conv2d(num_inchannels[j],\n num_outchannels_conv3x3,\n 3, 2, 1, bias=False),\n nn.BatchNorm2d(num_outchannels_conv3x3,\n momentum=BN_MOMENTUM),\n nn.ReLU(False)))\n fuse_layer.append(nn.Sequential(*conv3x3s))\n fuse_layers.append(nn.ModuleList(fuse_layer))\n\n return nn.ModuleList(fuse_layers)\n\n def get_num_inchannels(self):\n return self.num_inchannels\n\n def forward(self, x):\n if self.num_branches == 1:\n return [self.branches[0](x[0])]\n\n for i in range(self.num_branches):\n x[i] = self.branches[i](x[i])\n\n x_fuse = []\n for i in range(len(self.fuse_layers)):\n y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])\n for j in range(1, self.num_branches):\n if i == j:\n y = y + x[j]\n else:\n y = y + self.fuse_layers[i][j](x[j])\n x_fuse.append(self.relu(y))\n\n return x_fuse\n\n\nblocks_dict = {\n 'BASIC': BasicBlock,\n 'BOTTLENECK': Bottleneck\n}\n\n\nclass ConvBlock(nn.Module):\n def __init__(self, in_c, out_c, k, s=1, p=0):\n super(ConvBlock, self).__init__()\n self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p)\n self.conv.apply(weights_init_kaiming)\n self.bn = nn.BatchNorm2d(out_c)\n\n def forward(self, x):\n return self.bn(self.conv(x))\n\n\ndef weights_init_kaiming(m):\n classname = m.__class__.__name__\n if classname.find('Linear') != -1:\n nn.init.kaiming_normal_(m.weight, a=0, mode='fan_out')\n nn.init.constant_(m.bias, 0.0)\n elif classname.find('Conv') != -1:\n nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')\n if m.bias is not None:\n nn.init.constant_(m.bias, 0.0)\n elif classname.find('BatchNorm') != -1:\n if m.affine:\n nn.init.constant_(m.weight, 1.0)\n nn.init.constant_(m.bias, 0.0)\n\n\nclass HighResolutionNet(nn.Module):\n\n def __init__(self, cfg, enable_dim_reduction, dim_reduction_channels, **kwargs):\n super(HighResolutionNet, self).__init__()\n\n self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,\n bias=False)\n self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\n self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,\n bias=False)\n self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)\n self.relu = nn.ReLU(inplace=True)\n self.layer1 = self._make_layer(Bottleneck, 64, 64, 4)\n\n self.stage2_cfg = cfg['MODEL']['EXTRA']['STAGE2']\n num_channels = self.stage2_cfg['NUM_CHANNELS']\n block = blocks_dict[self.stage2_cfg['BLOCK']]\n num_channels = [\n num_channels[i] * block.expansion for i in range(len(num_channels))]\n self.transition1 = self._make_transition_layer(\n [256], num_channels)\n self.stage2, pre_stage_channels = self._make_stage(\n self.stage2_cfg, num_channels)\n\n self.stage3_cfg = cfg['MODEL']['EXTRA']['STAGE3']\n num_channels = self.stage3_cfg['NUM_CHANNELS']\n block = blocks_dict[self.stage3_cfg['BLOCK']]\n num_channels = [\n num_channels[i] * block.expansion for i in range(len(num_channels))]\n self.transition2 = self._make_transition_layer(\n pre_stage_channels, num_channels)\n self.stage3, pre_stage_channels = self._make_stage(\n self.stage3_cfg, num_channels)\n\n self.stage4_cfg = cfg['MODEL']['EXTRA']['STAGE4']\n num_channels = self.stage4_cfg['NUM_CHANNELS']\n block = blocks_dict[self.stage4_cfg['BLOCK']]\n num_channels = [\n num_channels[i] * block.expansion for i in range(len(num_channels))]\n self.transition3 = self._make_transition_layer(\n pre_stage_channels, num_channels)\n self.stage4, pre_stage_channels = self._make_stage(\n self.stage4_cfg, num_channels, multi_scale_output=True)\n\n self.incre_modules, _, _ = self._make_head(pre_stage_channels)\n\n self.layers_out_channels = 1920\n self.dim_reduction_channels = dim_reduction_channels\n self.cls_head = nn.Sequential(\n nn.Conv2d(\n in_channels=self.layers_out_channels,\n out_channels=self.dim_reduction_channels,\n kernel_size=1,\n stride=1,\n padding=0),\n nn.BatchNorm2d(self.dim_reduction_channels, momentum=BN_MOMENTUM),\n nn.ReLU(inplace=True)\n )\n self.gap = nn.AdaptiveAvgPool2d(1)\n\n self.enable_dim_reduction = enable_dim_reduction\n if self.enable_dim_reduction:\n self.feature_dim = self.dim_reduction_channels\n else:\n self.feature_dim = self.layers_out_channels\n\n self.random_init()\n\n def _make_incre_channel_nin(self):\n head_channels = [128, 256, 512, 1024]\n incre_modules = []\n for i in range(3):\n incre_module = nn.Sequential(\n nn.Conv2d(\n in_channels=head_channels[i],\n out_channels=head_channels[i + 1],\n kernel_size=1,\n stride=1,\n padding=0),\n nn.BatchNorm2d(head_channels[i + 1], momentum=BN_MOMENTUM),\n nn.ReLU(inplace=True)\n )\n incre_modules.append(incre_module)\n incre_modules = nn.ModuleList(incre_modules)\n return incre_modules\n\n def _make_head(self, pre_stage_channels):\n head_block = Bottleneck\n head_channels = [32, 64, 128, 256]\n\n # Increasing the #channels on each resolution \n # from C, 2C, 4C, 8C to 128, 256, 512, 1024\n incre_modules = []\n for i, channels in enumerate(pre_stage_channels):\n incre_module = self._make_layer(head_block,\n channels,\n head_channels[i],\n 1,\n stride=1)\n incre_modules.append(incre_module)\n incre_modules = nn.ModuleList(incre_modules)\n\n # downsampling modules\n downsamp_modules = []\n for i in range(len(pre_stage_channels) - 1):\n in_channels = head_channels[i] * head_block.expansion\n out_channels = head_channels[i + 1] * head_block.expansion\n\n downsamp_module = nn.Sequential(\n nn.Conv2d(in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=3,\n stride=2,\n padding=1),\n nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM),\n nn.ReLU(inplace=True)\n )\n\n downsamp_modules.append(downsamp_module)\n downsamp_modules = nn.ModuleList(downsamp_modules)\n\n final_layer = nn.Sequential(\n nn.Conv2d(\n in_channels=head_channels[3] * head_block.expansion,\n out_channels=2048,\n kernel_size=1,\n stride=1,\n padding=0\n ),\n nn.BatchNorm2d(2048, momentum=BN_MOMENTUM),\n nn.ReLU(inplace=True)\n )\n\n return incre_modules, downsamp_modules, final_layer\n\n def _make_transition_layer(\n self, num_channels_pre_layer, num_channels_cur_layer):\n num_branches_cur = len(num_channels_cur_layer)\n num_branches_pre = len(num_channels_pre_layer)\n\n transition_layers = []\n for i in range(num_branches_cur):\n if i < num_branches_pre:\n if num_channels_cur_layer[i] != num_channels_pre_layer[i]:\n transition_layers.append(nn.Sequential(\n nn.Conv2d(num_channels_pre_layer[i],\n num_channels_cur_layer[i],\n 3,\n 1,\n 1,\n bias=False),\n nn.BatchNorm2d(\n num_channels_cur_layer[i], momentum=BN_MOMENTUM),\n nn.ReLU(inplace=True)))\n else:\n transition_layers.append(None)\n else:\n conv3x3s = []\n for j in range(i + 1 - num_branches_pre):\n inchannels = num_channels_pre_layer[-1]\n outchannels = num_channels_cur_layer[i] \\\n if j == i - num_branches_pre else inchannels\n conv3x3s.append(nn.Sequential(\n nn.Conv2d(\n inchannels, outchannels, 3, 2, 1, bias=False),\n nn.BatchNorm2d(outchannels, momentum=BN_MOMENTUM),\n nn.ReLU(inplace=True)))\n transition_layers.append(nn.Sequential(*conv3x3s))\n\n return nn.ModuleList(transition_layers)\n\n def _make_layer(self, block, inplanes, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(inplanes, planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),\n )\n\n layers = []\n layers.append(block(inplanes, planes, stride, downsample))\n inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(inplanes, planes))\n\n return nn.Sequential(*layers)\n\n def _make_stage(self, layer_config, num_inchannels,\n multi_scale_output=True):\n num_modules = layer_config['NUM_MODULES']\n num_branches = layer_config['NUM_BRANCHES']\n num_blocks = layer_config['NUM_BLOCKS']\n num_channels = layer_config['NUM_CHANNELS']\n block = blocks_dict[layer_config['BLOCK']]\n fuse_method = layer_config['FUSE_METHOD']\n\n modules = []\n for i in range(num_modules):\n # multi_scale_output is only used last module\n if not multi_scale_output and i == num_modules - 1:\n reset_multi_scale_output = False\n else:\n reset_multi_scale_output = True\n\n modules.append(\n HighResolutionModule(num_branches,\n block,\n num_blocks,\n num_inchannels,\n num_channels,\n fuse_method,\n reset_multi_scale_output)\n )\n num_inchannels = modules[-1].get_num_inchannels()\n\n return nn.Sequential(*modules), num_inchannels\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.conv2(x)\n x = self.bn2(x)\n x = self.relu(x)\n x = self.layer1(x)\n\n x_list = []\n for i in range(self.stage2_cfg['NUM_BRANCHES']):\n if self.transition1[i] is not None:\n x_list.append(self.transition1[i](x))\n else:\n x_list.append(x)\n y_list = self.stage2(x_list)\n\n x_list = []\n for i in range(self.stage3_cfg['NUM_BRANCHES']):\n if self.transition2[i] is not None:\n x_list.append(self.transition2[i](y_list[-1]))\n else:\n x_list.append(y_list[i])\n y_list = self.stage3(x_list)\n\n x_list = []\n for i in range(self.stage4_cfg['NUM_BRANCHES']):\n if self.transition3[i] is not None:\n x_list.append(self.transition3[i](y_list[-1]))\n else:\n x_list.append(y_list[i])\n x = self.stage4(x_list)\n\n for i in range(len(self.incre_modules)):\n x[i] = self.incre_modules[i](x[i])\n\n x0_h, x0_w = x[0].size(2), x[0].size(3)\n x1 = F.interpolate(x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=True) # torch.Size([128, 256, 64, 32])\n x2 = F.interpolate(x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=True) # torch.Size([128, 512, 64, 32])\n x3 = F.interpolate(x[3], size=(x0_h, x0_w), mode='bilinear', align_corners=True) # torch.Size([128, 1024, 64, 32])\n\n x = torch.cat([x[0], x1, x2, x3], 1) # torch.Size([b, 1920, 64, 32])\n if self.enable_dim_reduction:\n x = self.cls_head(x) # torch.Size([128, 256, 64, 32])\n return x\n\n def random_init(self):\n print('=> init weights from normal distribution')\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu')\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n def load_param(self, pretrained_path):\n if not Path(pretrained_path).exists():\n raise FileNotFoundError(f'HRNet-W32-C pretrained weights not found under \"{pretrained_path}\", please download it '\n f'first at https://github.com/HRNet/HRNet-Image-Classification or specify the correct '\n f'weights dir location with the cfg.model.bpbreid.hrnet_pretrained_path config.')\n pretrained_dict = torch.load(pretrained_path)\n print('=> loading pretrained model {}'.format(pretrained_path))\n model_dict = self.state_dict()\n pretrained_dict = {k: v for k, v in pretrained_dict.items()\n if k in model_dict.keys()}\n # for k, _ in pretrained_dict.items():\n # print(\n # '=> loading {} pretrained model {}'.format(k, pretrained_path))\n model_dict.update(pretrained_dict)\n self.load_state_dict(model_dict)\n\n\ndef init_pretrained_weights(model, pretrain_path, model_key):\n path = os.path.join(pretrain_path, model_paths[model_key])\n print('Loading pretrained ImageNet HRNet32 model at {}'.format(path))\n model.load_param(path)\n\n\ndef hrnet32(num_classes, loss='part_based', pretrained=True, enable_dim_reduction=True, dim_reduction_channels=256, pretrained_path='', **kwargs):\n cfg = get_hrnet_config()\n model = HighResolutionNet(\n cfg,\n enable_dim_reduction,\n dim_reduction_channels\n )\n if pretrained:\n init_pretrained_weights(model, pretrained_path, 'hrnet-w32')\n return model\n" }, { "path": "torchreid/models/inceptionresnetv2.py", "content": "\"\"\"\nCode imported from https://github.com/Cadene/pretrained-models.pytorch\n\"\"\"\nfrom __future__ import division, absolute_import\nimport torch\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['inceptionresnetv2']\n\npretrained_settings = {\n 'inceptionresnetv2': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/inceptionresnetv2-520b38e4.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 299, 299],\n 'input_range': [0, 1],\n 'mean': [0.5, 0.5, 0.5],\n 'std': [0.5, 0.5, 0.5],\n 'num_classes': 1000\n },\n 'imagenet+background': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/inceptionresnetv2-520b38e4.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 299, 299],\n 'input_range': [0, 1],\n 'mean': [0.5, 0.5, 0.5],\n 'std': [0.5, 0.5, 0.5],\n 'num_classes': 1001\n }\n }\n}\n\n\nclass BasicConv2d(nn.Module):\n\n def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):\n super(BasicConv2d, self).__init__()\n self.conv = nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n bias=False\n ) # verify bias false\n self.bn = nn.BatchNorm2d(\n out_planes,\n eps=0.001, # value found in tensorflow\n momentum=0.1, # default pytorch value\n affine=True\n )\n self.relu = nn.ReLU(inplace=False)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n\nclass Mixed_5b(nn.Module):\n\n def __init__(self):\n super(Mixed_5b, self).__init__()\n\n self.branch0 = BasicConv2d(192, 96, kernel_size=1, stride=1)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(192, 48, kernel_size=1, stride=1),\n BasicConv2d(48, 64, kernel_size=5, stride=1, padding=2)\n )\n\n self.branch2 = nn.Sequential(\n BasicConv2d(192, 64, kernel_size=1, stride=1),\n BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1),\n BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1)\n )\n\n self.branch3 = nn.Sequential(\n nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),\n BasicConv2d(192, 64, kernel_size=1, stride=1)\n )\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n x3 = self.branch3(x)\n out = torch.cat((x0, x1, x2, x3), 1)\n return out\n\n\nclass Block35(nn.Module):\n\n def __init__(self, scale=1.0):\n super(Block35, self).__init__()\n\n self.scale = scale\n\n self.branch0 = BasicConv2d(320, 32, kernel_size=1, stride=1)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(320, 32, kernel_size=1, stride=1),\n BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1)\n )\n\n self.branch2 = nn.Sequential(\n BasicConv2d(320, 32, kernel_size=1, stride=1),\n BasicConv2d(32, 48, kernel_size=3, stride=1, padding=1),\n BasicConv2d(48, 64, kernel_size=3, stride=1, padding=1)\n )\n\n self.conv2d = nn.Conv2d(128, 320, kernel_size=1, stride=1)\n self.relu = nn.ReLU(inplace=False)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n out = torch.cat((x0, x1, x2), 1)\n out = self.conv2d(out)\n out = out * self.scale + x\n out = self.relu(out)\n return out\n\n\nclass Mixed_6a(nn.Module):\n\n def __init__(self):\n super(Mixed_6a, self).__init__()\n\n self.branch0 = BasicConv2d(320, 384, kernel_size=3, stride=2)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(320, 256, kernel_size=1, stride=1),\n BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1),\n BasicConv2d(256, 384, kernel_size=3, stride=2)\n )\n\n self.branch2 = nn.MaxPool2d(3, stride=2)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n out = torch.cat((x0, x1, x2), 1)\n return out\n\n\nclass Block17(nn.Module):\n\n def __init__(self, scale=1.0):\n super(Block17, self).__init__()\n\n self.scale = scale\n\n self.branch0 = BasicConv2d(1088, 192, kernel_size=1, stride=1)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(1088, 128, kernel_size=1, stride=1),\n BasicConv2d(\n 128, 160, kernel_size=(1, 7), stride=1, padding=(0, 3)\n ),\n BasicConv2d(\n 160, 192, kernel_size=(7, 1), stride=1, padding=(3, 0)\n )\n )\n\n self.conv2d = nn.Conv2d(384, 1088, kernel_size=1, stride=1)\n self.relu = nn.ReLU(inplace=False)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n out = torch.cat((x0, x1), 1)\n out = self.conv2d(out)\n out = out * self.scale + x\n out = self.relu(out)\n return out\n\n\nclass Mixed_7a(nn.Module):\n\n def __init__(self):\n super(Mixed_7a, self).__init__()\n\n self.branch0 = nn.Sequential(\n BasicConv2d(1088, 256, kernel_size=1, stride=1),\n BasicConv2d(256, 384, kernel_size=3, stride=2)\n )\n\n self.branch1 = nn.Sequential(\n BasicConv2d(1088, 256, kernel_size=1, stride=1),\n BasicConv2d(256, 288, kernel_size=3, stride=2)\n )\n\n self.branch2 = nn.Sequential(\n BasicConv2d(1088, 256, kernel_size=1, stride=1),\n BasicConv2d(256, 288, kernel_size=3, stride=1, padding=1),\n BasicConv2d(288, 320, kernel_size=3, stride=2)\n )\n\n self.branch3 = nn.MaxPool2d(3, stride=2)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n x3 = self.branch3(x)\n out = torch.cat((x0, x1, x2, x3), 1)\n return out\n\n\nclass Block8(nn.Module):\n\n def __init__(self, scale=1.0, noReLU=False):\n super(Block8, self).__init__()\n\n self.scale = scale\n self.noReLU = noReLU\n\n self.branch0 = BasicConv2d(2080, 192, kernel_size=1, stride=1)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(2080, 192, kernel_size=1, stride=1),\n BasicConv2d(\n 192, 224, kernel_size=(1, 3), stride=1, padding=(0, 1)\n ),\n BasicConv2d(\n 224, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)\n )\n )\n\n self.conv2d = nn.Conv2d(448, 2080, kernel_size=1, stride=1)\n if not self.noReLU:\n self.relu = nn.ReLU(inplace=False)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n out = torch.cat((x0, x1), 1)\n out = self.conv2d(out)\n out = out * self.scale + x\n if not self.noReLU:\n out = self.relu(out)\n return out\n\n\n# ----------------\n# Model Definition\n# ----------------\nclass InceptionResNetV2(nn.Module):\n \"\"\"Inception-ResNet-V2.\n\n Reference:\n Szegedy et al. Inception-v4, Inception-ResNet and the Impact of Residual\n Connections on Learning. AAAI 2017.\n\n Public keys:\n - ``inceptionresnetv2``: Inception-ResNet-V2.\n \"\"\"\n\n def __init__(self, num_classes, loss='softmax', **kwargs):\n super(InceptionResNetV2, self).__init__()\n self.loss = loss\n\n # Modules\n self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2)\n self.conv2d_2a = BasicConv2d(32, 32, kernel_size=3, stride=1)\n self.conv2d_2b = BasicConv2d(\n 32, 64, kernel_size=3, stride=1, padding=1\n )\n self.maxpool_3a = nn.MaxPool2d(3, stride=2)\n self.conv2d_3b = BasicConv2d(64, 80, kernel_size=1, stride=1)\n self.conv2d_4a = BasicConv2d(80, 192, kernel_size=3, stride=1)\n self.maxpool_5a = nn.MaxPool2d(3, stride=2)\n self.mixed_5b = Mixed_5b()\n self.repeat = nn.Sequential(\n Block35(scale=0.17), Block35(scale=0.17), Block35(scale=0.17),\n Block35(scale=0.17), Block35(scale=0.17), Block35(scale=0.17),\n Block35(scale=0.17), Block35(scale=0.17), Block35(scale=0.17),\n Block35(scale=0.17)\n )\n self.mixed_6a = Mixed_6a()\n self.repeat_1 = nn.Sequential(\n Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),\n Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),\n Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),\n Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),\n Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),\n Block17(scale=0.10), Block17(scale=0.10), Block17(scale=0.10),\n Block17(scale=0.10), Block17(scale=0.10)\n )\n self.mixed_7a = Mixed_7a()\n self.repeat_2 = nn.Sequential(\n Block8(scale=0.20), Block8(scale=0.20), Block8(scale=0.20),\n Block8(scale=0.20), Block8(scale=0.20), Block8(scale=0.20),\n Block8(scale=0.20), Block8(scale=0.20), Block8(scale=0.20)\n )\n\n self.block8 = Block8(noReLU=True)\n self.conv2d_7b = BasicConv2d(2080, 1536, kernel_size=1, stride=1)\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.classifier = nn.Linear(1536, num_classes)\n\n def load_imagenet_weights(self):\n settings = pretrained_settings['inceptionresnetv2']['imagenet']\n pretrain_dict = model_zoo.load_url(settings['url'])\n model_dict = self.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n self.load_state_dict(model_dict)\n\n def featuremaps(self, x):\n x = self.conv2d_1a(x)\n x = self.conv2d_2a(x)\n x = self.conv2d_2b(x)\n x = self.maxpool_3a(x)\n x = self.conv2d_3b(x)\n x = self.conv2d_4a(x)\n x = self.maxpool_5a(x)\n x = self.mixed_5b(x)\n x = self.repeat(x)\n x = self.mixed_6a(x)\n x = self.repeat_1(x)\n x = self.mixed_7a(x)\n x = self.repeat_2(x)\n x = self.block8(x)\n x = self.conv2d_7b(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef inceptionresnetv2(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = InceptionResNetV2(num_classes=num_classes, loss=loss, **kwargs)\n if pretrained:\n model.load_imagenet_weights()\n return model\n" }, { "path": "torchreid/models/inceptionv4.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['inceptionv4']\n\"\"\"\nCode imported from https://github.com/Cadene/pretrained-models.pytorch\n\"\"\"\n\npretrained_settings = {\n 'inceptionv4': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/inceptionv4-8e4777a0.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 299, 299],\n 'input_range': [0, 1],\n 'mean': [0.5, 0.5, 0.5],\n 'std': [0.5, 0.5, 0.5],\n 'num_classes': 1000\n },\n 'imagenet+background': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/inceptionv4-8e4777a0.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 299, 299],\n 'input_range': [0, 1],\n 'mean': [0.5, 0.5, 0.5],\n 'std': [0.5, 0.5, 0.5],\n 'num_classes': 1001\n }\n }\n}\n\n\nclass BasicConv2d(nn.Module):\n\n def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):\n super(BasicConv2d, self).__init__()\n self.conv = nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=kernel_size,\n stride=stride,\n padding=padding,\n bias=False\n ) # verify bias false\n self.bn = nn.BatchNorm2d(\n out_planes,\n eps=0.001, # value found in tensorflow\n momentum=0.1, # default pytorch value\n affine=True\n )\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n\nclass Mixed_3a(nn.Module):\n\n def __init__(self):\n super(Mixed_3a, self).__init__()\n self.maxpool = nn.MaxPool2d(3, stride=2)\n self.conv = BasicConv2d(64, 96, kernel_size=3, stride=2)\n\n def forward(self, x):\n x0 = self.maxpool(x)\n x1 = self.conv(x)\n out = torch.cat((x0, x1), 1)\n return out\n\n\nclass Mixed_4a(nn.Module):\n\n def __init__(self):\n super(Mixed_4a, self).__init__()\n\n self.branch0 = nn.Sequential(\n BasicConv2d(160, 64, kernel_size=1, stride=1),\n BasicConv2d(64, 96, kernel_size=3, stride=1)\n )\n\n self.branch1 = nn.Sequential(\n BasicConv2d(160, 64, kernel_size=1, stride=1),\n BasicConv2d(64, 64, kernel_size=(1, 7), stride=1, padding=(0, 3)),\n BasicConv2d(64, 64, kernel_size=(7, 1), stride=1, padding=(3, 0)),\n BasicConv2d(64, 96, kernel_size=(3, 3), stride=1)\n )\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n out = torch.cat((x0, x1), 1)\n return out\n\n\nclass Mixed_5a(nn.Module):\n\n def __init__(self):\n super(Mixed_5a, self).__init__()\n self.conv = BasicConv2d(192, 192, kernel_size=3, stride=2)\n self.maxpool = nn.MaxPool2d(3, stride=2)\n\n def forward(self, x):\n x0 = self.conv(x)\n x1 = self.maxpool(x)\n out = torch.cat((x0, x1), 1)\n return out\n\n\nclass Inception_A(nn.Module):\n\n def __init__(self):\n super(Inception_A, self).__init__()\n self.branch0 = BasicConv2d(384, 96, kernel_size=1, stride=1)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(384, 64, kernel_size=1, stride=1),\n BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1)\n )\n\n self.branch2 = nn.Sequential(\n BasicConv2d(384, 64, kernel_size=1, stride=1),\n BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1),\n BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1)\n )\n\n self.branch3 = nn.Sequential(\n nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),\n BasicConv2d(384, 96, kernel_size=1, stride=1)\n )\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n x3 = self.branch3(x)\n out = torch.cat((x0, x1, x2, x3), 1)\n return out\n\n\nclass Reduction_A(nn.Module):\n\n def __init__(self):\n super(Reduction_A, self).__init__()\n self.branch0 = BasicConv2d(384, 384, kernel_size=3, stride=2)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(384, 192, kernel_size=1, stride=1),\n BasicConv2d(192, 224, kernel_size=3, stride=1, padding=1),\n BasicConv2d(224, 256, kernel_size=3, stride=2)\n )\n\n self.branch2 = nn.MaxPool2d(3, stride=2)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n out = torch.cat((x0, x1, x2), 1)\n return out\n\n\nclass Inception_B(nn.Module):\n\n def __init__(self):\n super(Inception_B, self).__init__()\n self.branch0 = BasicConv2d(1024, 384, kernel_size=1, stride=1)\n\n self.branch1 = nn.Sequential(\n BasicConv2d(1024, 192, kernel_size=1, stride=1),\n BasicConv2d(\n 192, 224, kernel_size=(1, 7), stride=1, padding=(0, 3)\n ),\n BasicConv2d(\n 224, 256, kernel_size=(7, 1), stride=1, padding=(3, 0)\n )\n )\n\n self.branch2 = nn.Sequential(\n BasicConv2d(1024, 192, kernel_size=1, stride=1),\n BasicConv2d(\n 192, 192, kernel_size=(7, 1), stride=1, padding=(3, 0)\n ),\n BasicConv2d(\n 192, 224, kernel_size=(1, 7), stride=1, padding=(0, 3)\n ),\n BasicConv2d(\n 224, 224, kernel_size=(7, 1), stride=1, padding=(3, 0)\n ),\n BasicConv2d(\n 224, 256, kernel_size=(1, 7), stride=1, padding=(0, 3)\n )\n )\n\n self.branch3 = nn.Sequential(\n nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),\n BasicConv2d(1024, 128, kernel_size=1, stride=1)\n )\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n x3 = self.branch3(x)\n out = torch.cat((x0, x1, x2, x3), 1)\n return out\n\n\nclass Reduction_B(nn.Module):\n\n def __init__(self):\n super(Reduction_B, self).__init__()\n\n self.branch0 = nn.Sequential(\n BasicConv2d(1024, 192, kernel_size=1, stride=1),\n BasicConv2d(192, 192, kernel_size=3, stride=2)\n )\n\n self.branch1 = nn.Sequential(\n BasicConv2d(1024, 256, kernel_size=1, stride=1),\n BasicConv2d(\n 256, 256, kernel_size=(1, 7), stride=1, padding=(0, 3)\n ),\n BasicConv2d(\n 256, 320, kernel_size=(7, 1), stride=1, padding=(3, 0)\n ), BasicConv2d(320, 320, kernel_size=3, stride=2)\n )\n\n self.branch2 = nn.MaxPool2d(3, stride=2)\n\n def forward(self, x):\n x0 = self.branch0(x)\n x1 = self.branch1(x)\n x2 = self.branch2(x)\n out = torch.cat((x0, x1, x2), 1)\n return out\n\n\nclass Inception_C(nn.Module):\n\n def __init__(self):\n super(Inception_C, self).__init__()\n\n self.branch0 = BasicConv2d(1536, 256, kernel_size=1, stride=1)\n\n self.branch1_0 = BasicConv2d(1536, 384, kernel_size=1, stride=1)\n self.branch1_1a = BasicConv2d(\n 384, 256, kernel_size=(1, 3), stride=1, padding=(0, 1)\n )\n self.branch1_1b = BasicConv2d(\n 384, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)\n )\n\n self.branch2_0 = BasicConv2d(1536, 384, kernel_size=1, stride=1)\n self.branch2_1 = BasicConv2d(\n 384, 448, kernel_size=(3, 1), stride=1, padding=(1, 0)\n )\n self.branch2_2 = BasicConv2d(\n 448, 512, kernel_size=(1, 3), stride=1, padding=(0, 1)\n )\n self.branch2_3a = BasicConv2d(\n 512, 256, kernel_size=(1, 3), stride=1, padding=(0, 1)\n )\n self.branch2_3b = BasicConv2d(\n 512, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)\n )\n\n self.branch3 = nn.Sequential(\n nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),\n BasicConv2d(1536, 256, kernel_size=1, stride=1)\n )\n\n def forward(self, x):\n x0 = self.branch0(x)\n\n x1_0 = self.branch1_0(x)\n x1_1a = self.branch1_1a(x1_0)\n x1_1b = self.branch1_1b(x1_0)\n x1 = torch.cat((x1_1a, x1_1b), 1)\n\n x2_0 = self.branch2_0(x)\n x2_1 = self.branch2_1(x2_0)\n x2_2 = self.branch2_2(x2_1)\n x2_3a = self.branch2_3a(x2_2)\n x2_3b = self.branch2_3b(x2_2)\n x2 = torch.cat((x2_3a, x2_3b), 1)\n\n x3 = self.branch3(x)\n\n out = torch.cat((x0, x1, x2, x3), 1)\n return out\n\n\nclass InceptionV4(nn.Module):\n \"\"\"Inception-v4.\n\n Reference:\n Szegedy et al. Inception-v4, Inception-ResNet and the Impact of Residual\n Connections on Learning. AAAI 2017.\n\n Public keys:\n - ``inceptionv4``: InceptionV4.\n \"\"\"\n\n def __init__(self, num_classes, loss, **kwargs):\n super(InceptionV4, self).__init__()\n self.loss = loss\n\n self.features = nn.Sequential(\n BasicConv2d(3, 32, kernel_size=3, stride=2),\n BasicConv2d(32, 32, kernel_size=3, stride=1),\n BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1),\n Mixed_3a(),\n Mixed_4a(),\n Mixed_5a(),\n Inception_A(),\n Inception_A(),\n Inception_A(),\n Inception_A(),\n Reduction_A(), # Mixed_6a\n Inception_B(),\n Inception_B(),\n Inception_B(),\n Inception_B(),\n Inception_B(),\n Inception_B(),\n Inception_B(),\n Reduction_B(), # Mixed_7a\n Inception_C(),\n Inception_C(),\n Inception_C()\n )\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.classifier = nn.Linear(1536, num_classes)\n\n def forward(self, x):\n f = self.features(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef inceptionv4(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = InceptionV4(num_classes, loss, **kwargs)\n if pretrained:\n model_url = pretrained_settings['inceptionv4']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n" }, { "path": "torchreid/models/mlfn.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['mlfn']\n\nmodel_urls = {\n # training epoch = 5, top1 = 51.6\n 'imagenet':\n 'https://mega.nz/#!YHxAhaxC!yu9E6zWl0x5zscSouTdbZu8gdFFytDdl-RAdD2DEfpk',\n}\n\n\nclass MLFNBlock(nn.Module):\n\n def __init__(\n self, in_channels, out_channels, stride, fsm_channels, groups=32\n ):\n super(MLFNBlock, self).__init__()\n self.groups = groups\n mid_channels = out_channels // 2\n\n # Factor Modules\n self.fm_conv1 = nn.Conv2d(in_channels, mid_channels, 1, bias=False)\n self.fm_bn1 = nn.BatchNorm2d(mid_channels)\n self.fm_conv2 = nn.Conv2d(\n mid_channels,\n mid_channels,\n 3,\n stride=stride,\n padding=1,\n bias=False,\n groups=self.groups\n )\n self.fm_bn2 = nn.BatchNorm2d(mid_channels)\n self.fm_conv3 = nn.Conv2d(mid_channels, out_channels, 1, bias=False)\n self.fm_bn3 = nn.BatchNorm2d(out_channels)\n\n # Factor Selection Module\n self.fsm = nn.Sequential(\n nn.AdaptiveAvgPool2d(1),\n nn.Conv2d(in_channels, fsm_channels[0], 1),\n nn.BatchNorm2d(fsm_channels[0]),\n nn.ReLU(inplace=True),\n nn.Conv2d(fsm_channels[0], fsm_channels[1], 1),\n nn.BatchNorm2d(fsm_channels[1]),\n nn.ReLU(inplace=True),\n nn.Conv2d(fsm_channels[1], self.groups, 1),\n nn.BatchNorm2d(self.groups),\n nn.Sigmoid(),\n )\n\n self.downsample = None\n if in_channels != out_channels or stride > 1:\n self.downsample = nn.Sequential(\n nn.Conv2d(\n in_channels, out_channels, 1, stride=stride, bias=False\n ),\n nn.BatchNorm2d(out_channels),\n )\n\n def forward(self, x):\n residual = x\n s = self.fsm(x)\n\n # reduce dimension\n x = self.fm_conv1(x)\n x = self.fm_bn1(x)\n x = F.relu(x, inplace=True)\n\n # group convolution\n x = self.fm_conv2(x)\n x = self.fm_bn2(x)\n x = F.relu(x, inplace=True)\n\n # factor selection\n b, c = x.size(0), x.size(1)\n n = c // self.groups\n ss = s.repeat(1, n, 1, 1) # from (b, g, 1, 1) to (b, g*n=c, 1, 1)\n ss = ss.view(b, n, self.groups, 1, 1)\n ss = ss.permute(0, 2, 1, 3, 4).contiguous()\n ss = ss.view(b, c, 1, 1)\n x = ss * x\n\n # recover dimension\n x = self.fm_conv3(x)\n x = self.fm_bn3(x)\n x = F.relu(x, inplace=True)\n\n if self.downsample is not None:\n residual = self.downsample(residual)\n\n return F.relu(residual + x, inplace=True), s\n\n\nclass MLFN(nn.Module):\n \"\"\"Multi-Level Factorisation Net.\n\n Reference:\n Chang et al. Multi-Level Factorisation Net for\n Person Re-Identification. CVPR 2018.\n\n Public keys:\n - ``mlfn``: MLFN (Multi-Level Factorisation Net).\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss='softmax',\n groups=32,\n channels=[64, 256, 512, 1024, 2048],\n embed_dim=1024,\n **kwargs\n ):\n super(MLFN, self).__init__()\n self.loss = loss\n self.groups = groups\n\n # first convolutional layer\n self.conv1 = nn.Conv2d(3, channels[0], 7, stride=2, padding=3)\n self.bn1 = nn.BatchNorm2d(channels[0])\n self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)\n\n # main body\n self.feature = nn.ModuleList(\n [\n # layer 1-3\n MLFNBlock(channels[0], channels[1], 1, [128, 64], self.groups),\n MLFNBlock(channels[1], channels[1], 1, [128, 64], self.groups),\n MLFNBlock(channels[1], channels[1], 1, [128, 64], self.groups),\n # layer 4-7\n MLFNBlock(\n channels[1], channels[2], 2, [256, 128], self.groups\n ),\n MLFNBlock(\n channels[2], channels[2], 1, [256, 128], self.groups\n ),\n MLFNBlock(\n channels[2], channels[2], 1, [256, 128], self.groups\n ),\n MLFNBlock(\n channels[2], channels[2], 1, [256, 128], self.groups\n ),\n # layer 8-13\n MLFNBlock(\n channels[2], channels[3], 2, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[3], channels[3], 1, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[3], channels[3], 1, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[3], channels[3], 1, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[3], channels[3], 1, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[3], channels[3], 1, [512, 128], self.groups\n ),\n # layer 14-16\n MLFNBlock(\n channels[3], channels[4], 2, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[4], channels[4], 1, [512, 128], self.groups\n ),\n MLFNBlock(\n channels[4], channels[4], 1, [512, 128], self.groups\n ),\n ]\n )\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n\n # projection functions\n self.fc_x = nn.Sequential(\n nn.Conv2d(channels[4], embed_dim, 1, bias=False),\n nn.BatchNorm2d(embed_dim),\n nn.ReLU(inplace=True),\n )\n self.fc_s = nn.Sequential(\n nn.Conv2d(self.groups * 16, embed_dim, 1, bias=False),\n nn.BatchNorm2d(embed_dim),\n nn.ReLU(inplace=True),\n )\n\n self.classifier = nn.Linear(embed_dim, num_classes)\n\n self.init_params()\n\n def init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = F.relu(x, inplace=True)\n x = self.maxpool(x)\n\n s_hat = []\n for block in self.feature:\n x, s = block(x)\n s_hat.append(s)\n s_hat = torch.cat(s_hat, 1)\n\n x = self.global_avgpool(x)\n x = self.fc_x(x)\n s_hat = self.fc_s(s_hat)\n\n v = (x+s_hat) * 0.5\n v = v.view(v.size(0), -1)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef mlfn(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = MLFN(num_classes, loss, **kwargs)\n if pretrained:\n # init_pretrained_weights(model, model_urls['imagenet'])\n import warnings\n warnings.warn(\n 'The imagenet pretrained weights need to be manually downloaded from {}'\n .format(model_urls['imagenet'])\n )\n return model\n" }, { "path": "torchreid/models/mobilenetv2.py", "content": "from __future__ import division, absolute_import\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['mobilenetv2_x1_0', 'mobilenetv2_x1_4']\n\nmodel_urls = {\n # 1.0: top-1 71.3\n 'mobilenetv2_x1_0':\n 'https://mega.nz/#!NKp2wAIA!1NH1pbNzY_M2hVk_hdsxNM1NUOWvvGPHhaNr-fASF6c',\n # 1.4: top-1 73.9\n 'mobilenetv2_x1_4':\n 'https://mega.nz/#!RGhgEIwS!xN2s2ZdyqI6vQ3EwgmRXLEW3khr9tpXg96G9SUJugGk',\n}\n\n\nclass ConvBlock(nn.Module):\n \"\"\"Basic convolutional block.\n \n convolution (bias discarded) + batch normalization + relu6.\n\n Args:\n in_c (int): number of input channels.\n out_c (int): number of output channels.\n k (int or tuple): kernel size.\n s (int or tuple): stride.\n p (int or tuple): padding.\n g (int): number of blocked connections from input channels\n to output channels (default: 1).\n \"\"\"\n\n def __init__(self, in_c, out_c, k, s=1, p=0, g=1):\n super(ConvBlock, self).__init__()\n self.conv = nn.Conv2d(\n in_c, out_c, k, stride=s, padding=p, bias=False, groups=g\n )\n self.bn = nn.BatchNorm2d(out_c)\n\n def forward(self, x):\n return F.relu6(self.bn(self.conv(x)))\n\n\nclass Bottleneck(nn.Module):\n\n def __init__(self, in_channels, out_channels, expansion_factor, stride=1):\n super(Bottleneck, self).__init__()\n mid_channels = in_channels * expansion_factor\n self.use_residual = stride == 1 and in_channels == out_channels\n self.conv1 = ConvBlock(in_channels, mid_channels, 1)\n self.dwconv2 = ConvBlock(\n mid_channels, mid_channels, 3, stride, 1, g=mid_channels\n )\n self.conv3 = nn.Sequential(\n nn.Conv2d(mid_channels, out_channels, 1, bias=False),\n nn.BatchNorm2d(out_channels),\n )\n\n def forward(self, x):\n m = self.conv1(x)\n m = self.dwconv2(m)\n m = self.conv3(m)\n if self.use_residual:\n return x + m\n else:\n return m\n\n\nclass MobileNetV2(nn.Module):\n \"\"\"MobileNetV2.\n\n Reference:\n Sandler et al. MobileNetV2: Inverted Residuals and\n Linear Bottlenecks. CVPR 2018.\n\n Public keys:\n - ``mobilenetv2_x1_0``: MobileNetV2 x1.0.\n - ``mobilenetv2_x1_4``: MobileNetV2 x1.4.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n width_mult=1,\n loss='softmax',\n fc_dims=None,\n dropout_p=None,\n **kwargs\n ):\n super(MobileNetV2, self).__init__()\n self.loss = loss\n self.in_channels = int(32 * width_mult)\n self.feature_dim = int(1280 * width_mult) if width_mult > 1 else 1280\n\n # construct layers\n self.conv1 = ConvBlock(3, self.in_channels, 3, s=2, p=1)\n self.conv2 = self._make_layer(\n Bottleneck, 1, int(16 * width_mult), 1, 1\n )\n self.conv3 = self._make_layer(\n Bottleneck, 6, int(24 * width_mult), 2, 2\n )\n self.conv4 = self._make_layer(\n Bottleneck, 6, int(32 * width_mult), 3, 2\n )\n self.conv5 = self._make_layer(\n Bottleneck, 6, int(64 * width_mult), 4, 2\n )\n self.conv6 = self._make_layer(\n Bottleneck, 6, int(96 * width_mult), 3, 1\n )\n self.conv7 = self._make_layer(\n Bottleneck, 6, int(160 * width_mult), 3, 2\n )\n self.conv8 = self._make_layer(\n Bottleneck, 6, int(320 * width_mult), 1, 1\n )\n self.conv9 = ConvBlock(self.in_channels, self.feature_dim, 1)\n\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.fc = self._construct_fc_layer(\n fc_dims, self.feature_dim, dropout_p\n )\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _make_layer(self, block, t, c, n, s):\n # t: expansion factor\n # c: output channels\n # n: number of blocks\n # s: stride for first layer\n layers = []\n layers.append(block(self.in_channels, c, t, s))\n self.in_channels = c\n for i in range(1, n):\n layers.append(block(self.in_channels, c, t))\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer.\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.conv2(x)\n x = self.conv3(x)\n x = self.conv4(x)\n x = self.conv5(x)\n x = self.conv6(x)\n x = self.conv7(x)\n x = self.conv8(x)\n x = self.conv9(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if self.fc is not None:\n v = self.fc(v)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef mobilenetv2_x1_0(num_classes, loss, pretrained=True, **kwargs):\n model = MobileNetV2(\n num_classes,\n loss=loss,\n width_mult=1,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n # init_pretrained_weights(model, model_urls['mobilenetv2_x1_0'])\n import warnings\n warnings.warn(\n 'The imagenet pretrained weights need to be manually downloaded from {}'\n .format(model_urls['mobilenetv2_x1_0'])\n )\n return model\n\n\ndef mobilenetv2_x1_4(num_classes, loss, pretrained=True, **kwargs):\n model = MobileNetV2(\n num_classes,\n loss=loss,\n width_mult=1.4,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n # init_pretrained_weights(model, model_urls['mobilenetv2_x1_4'])\n import warnings\n warnings.warn(\n 'The imagenet pretrained weights need to be manually downloaded from {}'\n .format(model_urls['mobilenetv2_x1_4'])\n )\n return model\n" }, { "path": "torchreid/models/mudeep.py", "content": "from __future__ import division, absolute_import\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['MuDeep']\n\n\nclass ConvBlock(nn.Module):\n \"\"\"Basic convolutional block.\n \n convolution + batch normalization + relu.\n\n Args:\n in_c (int): number of input channels.\n out_c (int): number of output channels.\n k (int or tuple): kernel size.\n s (int or tuple): stride.\n p (int or tuple): padding.\n \"\"\"\n\n def __init__(self, in_c, out_c, k, s, p):\n super(ConvBlock, self).__init__()\n self.conv = nn.Conv2d(in_c, out_c, k, stride=s, padding=p)\n self.bn = nn.BatchNorm2d(out_c)\n\n def forward(self, x):\n return F.relu(self.bn(self.conv(x)))\n\n\nclass ConvLayers(nn.Module):\n \"\"\"Preprocessing layers.\"\"\"\n\n def __init__(self):\n super(ConvLayers, self).__init__()\n self.conv1 = ConvBlock(3, 48, k=3, s=1, p=1)\n self.conv2 = ConvBlock(48, 96, k=3, s=1, p=1)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.conv2(x)\n x = self.maxpool(x)\n return x\n\n\nclass MultiScaleA(nn.Module):\n \"\"\"Multi-scale stream layer A (Sec.3.1)\"\"\"\n\n def __init__(self):\n super(MultiScaleA, self).__init__()\n self.stream1 = nn.Sequential(\n ConvBlock(96, 96, k=1, s=1, p=0),\n ConvBlock(96, 24, k=3, s=1, p=1),\n )\n self.stream2 = nn.Sequential(\n nn.AvgPool2d(kernel_size=3, stride=1, padding=1),\n ConvBlock(96, 24, k=1, s=1, p=0),\n )\n self.stream3 = ConvBlock(96, 24, k=1, s=1, p=0)\n self.stream4 = nn.Sequential(\n ConvBlock(96, 16, k=1, s=1, p=0),\n ConvBlock(16, 24, k=3, s=1, p=1),\n ConvBlock(24, 24, k=3, s=1, p=1),\n )\n\n def forward(self, x):\n s1 = self.stream1(x)\n s2 = self.stream2(x)\n s3 = self.stream3(x)\n s4 = self.stream4(x)\n y = torch.cat([s1, s2, s3, s4], dim=1)\n return y\n\n\nclass Reduction(nn.Module):\n \"\"\"Reduction layer (Sec.3.1)\"\"\"\n\n def __init__(self):\n super(Reduction, self).__init__()\n self.stream1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.stream2 = ConvBlock(96, 96, k=3, s=2, p=1)\n self.stream3 = nn.Sequential(\n ConvBlock(96, 48, k=1, s=1, p=0),\n ConvBlock(48, 56, k=3, s=1, p=1),\n ConvBlock(56, 64, k=3, s=2, p=1),\n )\n\n def forward(self, x):\n s1 = self.stream1(x)\n s2 = self.stream2(x)\n s3 = self.stream3(x)\n y = torch.cat([s1, s2, s3], dim=1)\n return y\n\n\nclass MultiScaleB(nn.Module):\n \"\"\"Multi-scale stream layer B (Sec.3.1)\"\"\"\n\n def __init__(self):\n super(MultiScaleB, self).__init__()\n self.stream1 = nn.Sequential(\n nn.AvgPool2d(kernel_size=3, stride=1, padding=1),\n ConvBlock(256, 256, k=1, s=1, p=0),\n )\n self.stream2 = nn.Sequential(\n ConvBlock(256, 64, k=1, s=1, p=0),\n ConvBlock(64, 128, k=(1, 3), s=1, p=(0, 1)),\n ConvBlock(128, 256, k=(3, 1), s=1, p=(1, 0)),\n )\n self.stream3 = ConvBlock(256, 256, k=1, s=1, p=0)\n self.stream4 = nn.Sequential(\n ConvBlock(256, 64, k=1, s=1, p=0),\n ConvBlock(64, 64, k=(1, 3), s=1, p=(0, 1)),\n ConvBlock(64, 128, k=(3, 1), s=1, p=(1, 0)),\n ConvBlock(128, 128, k=(1, 3), s=1, p=(0, 1)),\n ConvBlock(128, 256, k=(3, 1), s=1, p=(1, 0)),\n )\n\n def forward(self, x):\n s1 = self.stream1(x)\n s2 = self.stream2(x)\n s3 = self.stream3(x)\n s4 = self.stream4(x)\n return s1, s2, s3, s4\n\n\nclass Fusion(nn.Module):\n \"\"\"Saliency-based learning fusion layer (Sec.3.2)\"\"\"\n\n def __init__(self):\n super(Fusion, self).__init__()\n self.a1 = nn.Parameter(torch.rand(1, 256, 1, 1))\n self.a2 = nn.Parameter(torch.rand(1, 256, 1, 1))\n self.a3 = nn.Parameter(torch.rand(1, 256, 1, 1))\n self.a4 = nn.Parameter(torch.rand(1, 256, 1, 1))\n\n # We add an average pooling layer to reduce the spatial dimension\n # of feature maps, which differs from the original paper.\n self.avgpool = nn.AvgPool2d(kernel_size=4, stride=4, padding=0)\n\n def forward(self, x1, x2, x3, x4):\n s1 = self.a1.expand_as(x1) * x1\n s2 = self.a2.expand_as(x2) * x2\n s3 = self.a3.expand_as(x3) * x3\n s4 = self.a4.expand_as(x4) * x4\n y = self.avgpool(s1 + s2 + s3 + s4)\n return y\n\n\nclass MuDeep(nn.Module):\n \"\"\"Multiscale deep neural network.\n\n Reference:\n Qian et al. Multi-scale Deep Learning Architectures\n for Person Re-identification. ICCV 2017.\n\n Public keys:\n - ``mudeep``: Multiscale deep neural network.\n \"\"\"\n\n def __init__(self, num_classes, loss='softmax', **kwargs):\n super(MuDeep, self).__init__()\n self.loss = loss\n\n self.block1 = ConvLayers()\n self.block2 = MultiScaleA()\n self.block3 = Reduction()\n self.block4 = MultiScaleB()\n self.block5 = Fusion()\n\n # Due to this fully connected layer, input image has to be fixed\n # in shape, i.e. (3, 256, 128), such that the last convolutional feature\n # maps are of shape (256, 16, 8). If input shape is changed,\n # the input dimension of this layer has to be changed accordingly.\n self.fc = nn.Sequential(\n nn.Linear(256 * 16 * 8, 4096),\n nn.BatchNorm1d(4096),\n nn.ReLU(),\n )\n self.classifier = nn.Linear(4096, num_classes)\n self.feat_dim = 4096\n\n def featuremaps(self, x):\n x = self.block1(x)\n x = self.block2(x)\n x = self.block3(x)\n x = self.block4(x)\n x = self.block5(*x)\n return x\n\n def forward(self, x):\n x = self.featuremaps(x)\n x = x.view(x.size(0), -1)\n x = self.fc(x)\n y = self.classifier(x)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, x\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n" }, { "path": "torchreid/models/nasnet.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['nasnetamobile']\n\"\"\"\nNASNet Mobile\nThanks to Anastasiia (https://github.com/DagnyT) for the great help, support and motivation!\n\n\n------------------------------------------------------------------------------------\n Architecture | Top-1 Acc | Top-5 Acc | Multiply-Adds | Params (M)\n------------------------------------------------------------------------------------\n| NASNet-A (4 @ 1056) | 74.08% | 91.74% | 564 M | 5.3 |\n------------------------------------------------------------------------------------\n# References:\n - [Learning Transferable Architectures for Scalable Image Recognition]\n (https://arxiv.org/abs/1707.07012)\n\"\"\"\n\"\"\"\nCode imported from https://github.com/Cadene/pretrained-models.pytorch\n\"\"\"\n\npretrained_settings = {\n 'nasnetamobile': {\n 'imagenet': {\n # 'url': 'https://github.com/veronikayurchuk/pretrained-models.pytorch/releases/download/v1.0/nasnetmobile-7e03cead.pth.tar',\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/nasnetamobile-7e03cead.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224], # resize 256\n 'input_range': [0, 1],\n 'mean': [0.5, 0.5, 0.5],\n 'std': [0.5, 0.5, 0.5],\n 'num_classes': 1000\n },\n # 'imagenet+background': {\n # # 'url': 'http://data.lip6.fr/cadene/pretrainedmodels/nasnetalarge-a1897284.pth',\n # 'input_space': 'RGB',\n # 'input_size': [3, 224, 224], # resize 256\n # 'input_range': [0, 1],\n # 'mean': [0.5, 0.5, 0.5],\n # 'std': [0.5, 0.5, 0.5],\n # 'num_classes': 1001\n # }\n }\n}\n\n\nclass MaxPoolPad(nn.Module):\n\n def __init__(self):\n super(MaxPoolPad, self).__init__()\n self.pad = nn.ZeroPad2d((1, 0, 1, 0))\n self.pool = nn.MaxPool2d(3, stride=2, padding=1)\n\n def forward(self, x):\n x = self.pad(x)\n x = self.pool(x)\n x = x[:, :, 1:, 1:].contiguous()\n return x\n\n\nclass AvgPoolPad(nn.Module):\n\n def __init__(self, stride=2, padding=1):\n super(AvgPoolPad, self).__init__()\n self.pad = nn.ZeroPad2d((1, 0, 1, 0))\n self.pool = nn.AvgPool2d(\n 3, stride=stride, padding=padding, count_include_pad=False\n )\n\n def forward(self, x):\n x = self.pad(x)\n x = self.pool(x)\n x = x[:, :, 1:, 1:].contiguous()\n return x\n\n\nclass SeparableConv2d(nn.Module):\n\n def __init__(\n self,\n in_channels,\n out_channels,\n dw_kernel,\n dw_stride,\n dw_padding,\n bias=False\n ):\n super(SeparableConv2d, self).__init__()\n self.depthwise_conv2d = nn.Conv2d(\n in_channels,\n in_channels,\n dw_kernel,\n stride=dw_stride,\n padding=dw_padding,\n bias=bias,\n groups=in_channels\n )\n self.pointwise_conv2d = nn.Conv2d(\n in_channels, out_channels, 1, stride=1, bias=bias\n )\n\n def forward(self, x):\n x = self.depthwise_conv2d(x)\n x = self.pointwise_conv2d(x)\n return x\n\n\nclass BranchSeparables(nn.Module):\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n name=None,\n bias=False\n ):\n super(BranchSeparables, self).__init__()\n self.relu = nn.ReLU()\n self.separable_1 = SeparableConv2d(\n in_channels, in_channels, kernel_size, stride, padding, bias=bias\n )\n self.bn_sep_1 = nn.BatchNorm2d(\n in_channels, eps=0.001, momentum=0.1, affine=True\n )\n self.relu1 = nn.ReLU()\n self.separable_2 = SeparableConv2d(\n in_channels, out_channels, kernel_size, 1, padding, bias=bias\n )\n self.bn_sep_2 = nn.BatchNorm2d(\n out_channels, eps=0.001, momentum=0.1, affine=True\n )\n self.name = name\n\n def forward(self, x):\n x = self.relu(x)\n if self.name == 'specific':\n x = nn.ZeroPad2d((1, 0, 1, 0))(x)\n x = self.separable_1(x)\n if self.name == 'specific':\n x = x[:, :, 1:, 1:].contiguous()\n\n x = self.bn_sep_1(x)\n x = self.relu1(x)\n x = self.separable_2(x)\n x = self.bn_sep_2(x)\n return x\n\n\nclass BranchSeparablesStem(nn.Module):\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n bias=False\n ):\n super(BranchSeparablesStem, self).__init__()\n self.relu = nn.ReLU()\n self.separable_1 = SeparableConv2d(\n in_channels, out_channels, kernel_size, stride, padding, bias=bias\n )\n self.bn_sep_1 = nn.BatchNorm2d(\n out_channels, eps=0.001, momentum=0.1, affine=True\n )\n self.relu1 = nn.ReLU()\n self.separable_2 = SeparableConv2d(\n out_channels, out_channels, kernel_size, 1, padding, bias=bias\n )\n self.bn_sep_2 = nn.BatchNorm2d(\n out_channels, eps=0.001, momentum=0.1, affine=True\n )\n\n def forward(self, x):\n x = self.relu(x)\n x = self.separable_1(x)\n x = self.bn_sep_1(x)\n x = self.relu1(x)\n x = self.separable_2(x)\n x = self.bn_sep_2(x)\n return x\n\n\nclass BranchSeparablesReduction(BranchSeparables):\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n z_padding=1,\n bias=False\n ):\n BranchSeparables.__init__(\n self, in_channels, out_channels, kernel_size, stride, padding, bias\n )\n self.padding = nn.ZeroPad2d((z_padding, 0, z_padding, 0))\n\n def forward(self, x):\n x = self.relu(x)\n x = self.padding(x)\n x = self.separable_1(x)\n x = x[:, :, 1:, 1:].contiguous()\n x = self.bn_sep_1(x)\n x = self.relu1(x)\n x = self.separable_2(x)\n x = self.bn_sep_2(x)\n return x\n\n\nclass CellStem0(nn.Module):\n\n def __init__(self, stem_filters, num_filters=42):\n super(CellStem0, self).__init__()\n self.num_filters = num_filters\n self.stem_filters = stem_filters\n self.conv_1x1 = nn.Sequential()\n self.conv_1x1.add_module('relu', nn.ReLU())\n self.conv_1x1.add_module(\n 'conv',\n nn.Conv2d(\n self.stem_filters, self.num_filters, 1, stride=1, bias=False\n )\n )\n self.conv_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n self.num_filters, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.comb_iter_0_left = BranchSeparables(\n self.num_filters, self.num_filters, 5, 2, 2\n )\n self.comb_iter_0_right = BranchSeparablesStem(\n self.stem_filters, self.num_filters, 7, 2, 3, bias=False\n )\n\n self.comb_iter_1_left = nn.MaxPool2d(3, stride=2, padding=1)\n self.comb_iter_1_right = BranchSeparablesStem(\n self.stem_filters, self.num_filters, 7, 2, 3, bias=False\n )\n\n self.comb_iter_2_left = nn.AvgPool2d(\n 3, stride=2, padding=1, count_include_pad=False\n )\n self.comb_iter_2_right = BranchSeparablesStem(\n self.stem_filters, self.num_filters, 5, 2, 2, bias=False\n )\n\n self.comb_iter_3_right = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_4_left = BranchSeparables(\n self.num_filters, self.num_filters, 3, 1, 1, bias=False\n )\n self.comb_iter_4_right = nn.MaxPool2d(3, stride=2, padding=1)\n\n def forward(self, x):\n x1 = self.conv_1x1(x)\n\n x_comb_iter_0_left = self.comb_iter_0_left(x1)\n x_comb_iter_0_right = self.comb_iter_0_right(x)\n x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right\n\n x_comb_iter_1_left = self.comb_iter_1_left(x1)\n x_comb_iter_1_right = self.comb_iter_1_right(x)\n x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right\n\n x_comb_iter_2_left = self.comb_iter_2_left(x1)\n x_comb_iter_2_right = self.comb_iter_2_right(x)\n x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right\n\n x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)\n x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1\n\n x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)\n x_comb_iter_4_right = self.comb_iter_4_right(x1)\n x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right\n\n x_out = torch.cat(\n [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1\n )\n return x_out\n\n\nclass CellStem1(nn.Module):\n\n def __init__(self, stem_filters, num_filters):\n super(CellStem1, self).__init__()\n self.num_filters = num_filters\n self.stem_filters = stem_filters\n self.conv_1x1 = nn.Sequential()\n self.conv_1x1.add_module('relu', nn.ReLU())\n self.conv_1x1.add_module(\n 'conv',\n nn.Conv2d(\n 2 * self.num_filters,\n self.num_filters,\n 1,\n stride=1,\n bias=False\n )\n )\n self.conv_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n self.num_filters, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.relu = nn.ReLU()\n self.path_1 = nn.Sequential()\n self.path_1.add_module(\n 'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)\n )\n self.path_1.add_module(\n 'conv',\n nn.Conv2d(\n self.stem_filters,\n self.num_filters // 2,\n 1,\n stride=1,\n bias=False\n )\n )\n self.path_2 = nn.ModuleList()\n self.path_2.add_module('pad', nn.ZeroPad2d((0, 1, 0, 1)))\n self.path_2.add_module(\n 'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)\n )\n self.path_2.add_module(\n 'conv',\n nn.Conv2d(\n self.stem_filters,\n self.num_filters // 2,\n 1,\n stride=1,\n bias=False\n )\n )\n\n self.final_path_bn = nn.BatchNorm2d(\n self.num_filters, eps=0.001, momentum=0.1, affine=True\n )\n\n self.comb_iter_0_left = BranchSeparables(\n self.num_filters,\n self.num_filters,\n 5,\n 2,\n 2,\n name='specific',\n bias=False\n )\n self.comb_iter_0_right = BranchSeparables(\n self.num_filters,\n self.num_filters,\n 7,\n 2,\n 3,\n name='specific',\n bias=False\n )\n\n # self.comb_iter_1_left = nn.MaxPool2d(3, stride=2, padding=1)\n self.comb_iter_1_left = MaxPoolPad()\n self.comb_iter_1_right = BranchSeparables(\n self.num_filters,\n self.num_filters,\n 7,\n 2,\n 3,\n name='specific',\n bias=False\n )\n\n # self.comb_iter_2_left = nn.AvgPool2d(3, stride=2, padding=1, count_include_pad=False)\n self.comb_iter_2_left = AvgPoolPad()\n self.comb_iter_2_right = BranchSeparables(\n self.num_filters,\n self.num_filters,\n 5,\n 2,\n 2,\n name='specific',\n bias=False\n )\n\n self.comb_iter_3_right = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_4_left = BranchSeparables(\n self.num_filters,\n self.num_filters,\n 3,\n 1,\n 1,\n name='specific',\n bias=False\n )\n # self.comb_iter_4_right = nn.MaxPool2d(3, stride=2, padding=1)\n self.comb_iter_4_right = MaxPoolPad()\n\n def forward(self, x_conv0, x_stem_0):\n x_left = self.conv_1x1(x_stem_0)\n\n x_relu = self.relu(x_conv0)\n # path 1\n x_path1 = self.path_1(x_relu)\n # path 2\n x_path2 = self.path_2.pad(x_relu)\n x_path2 = x_path2[:, :, 1:, 1:]\n x_path2 = self.path_2.avgpool(x_path2)\n x_path2 = self.path_2.conv(x_path2)\n # final path\n x_right = self.final_path_bn(torch.cat([x_path1, x_path2], 1))\n\n x_comb_iter_0_left = self.comb_iter_0_left(x_left)\n x_comb_iter_0_right = self.comb_iter_0_right(x_right)\n x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right\n\n x_comb_iter_1_left = self.comb_iter_1_left(x_left)\n x_comb_iter_1_right = self.comb_iter_1_right(x_right)\n x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right\n\n x_comb_iter_2_left = self.comb_iter_2_left(x_left)\n x_comb_iter_2_right = self.comb_iter_2_right(x_right)\n x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right\n\n x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)\n x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1\n\n x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)\n x_comb_iter_4_right = self.comb_iter_4_right(x_left)\n x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right\n\n x_out = torch.cat(\n [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1\n )\n return x_out\n\n\nclass FirstCell(nn.Module):\n\n def __init__(\n self, in_channels_left, out_channels_left, in_channels_right,\n out_channels_right\n ):\n super(FirstCell, self).__init__()\n self.conv_1x1 = nn.Sequential()\n self.conv_1x1.add_module('relu', nn.ReLU())\n self.conv_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_right, out_channels_right, 1, stride=1, bias=False\n )\n )\n self.conv_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_right, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.relu = nn.ReLU()\n self.path_1 = nn.Sequential()\n self.path_1.add_module(\n 'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)\n )\n self.path_1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_left, out_channels_left, 1, stride=1, bias=False\n )\n )\n self.path_2 = nn.ModuleList()\n self.path_2.add_module('pad', nn.ZeroPad2d((0, 1, 0, 1)))\n self.path_2.add_module(\n 'avgpool', nn.AvgPool2d(1, stride=2, count_include_pad=False)\n )\n self.path_2.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_left, out_channels_left, 1, stride=1, bias=False\n )\n )\n\n self.final_path_bn = nn.BatchNorm2d(\n out_channels_left * 2, eps=0.001, momentum=0.1, affine=True\n )\n\n self.comb_iter_0_left = BranchSeparables(\n out_channels_right, out_channels_right, 5, 1, 2, bias=False\n )\n self.comb_iter_0_right = BranchSeparables(\n out_channels_right, out_channels_right, 3, 1, 1, bias=False\n )\n\n self.comb_iter_1_left = BranchSeparables(\n out_channels_right, out_channels_right, 5, 1, 2, bias=False\n )\n self.comb_iter_1_right = BranchSeparables(\n out_channels_right, out_channels_right, 3, 1, 1, bias=False\n )\n\n self.comb_iter_2_left = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_3_left = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n self.comb_iter_3_right = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_4_left = BranchSeparables(\n out_channels_right, out_channels_right, 3, 1, 1, bias=False\n )\n\n def forward(self, x, x_prev):\n x_relu = self.relu(x_prev)\n # path 1\n x_path1 = self.path_1(x_relu)\n # path 2\n x_path2 = self.path_2.pad(x_relu)\n x_path2 = x_path2[:, :, 1:, 1:]\n x_path2 = self.path_2.avgpool(x_path2)\n x_path2 = self.path_2.conv(x_path2)\n # final path\n x_left = self.final_path_bn(torch.cat([x_path1, x_path2], 1))\n\n x_right = self.conv_1x1(x)\n\n x_comb_iter_0_left = self.comb_iter_0_left(x_right)\n x_comb_iter_0_right = self.comb_iter_0_right(x_left)\n x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right\n\n x_comb_iter_1_left = self.comb_iter_1_left(x_left)\n x_comb_iter_1_right = self.comb_iter_1_right(x_left)\n x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right\n\n x_comb_iter_2_left = self.comb_iter_2_left(x_right)\n x_comb_iter_2 = x_comb_iter_2_left + x_left\n\n x_comb_iter_3_left = self.comb_iter_3_left(x_left)\n x_comb_iter_3_right = self.comb_iter_3_right(x_left)\n x_comb_iter_3 = x_comb_iter_3_left + x_comb_iter_3_right\n\n x_comb_iter_4_left = self.comb_iter_4_left(x_right)\n x_comb_iter_4 = x_comb_iter_4_left + x_right\n\n x_out = torch.cat(\n [\n x_left, x_comb_iter_0, x_comb_iter_1, x_comb_iter_2,\n x_comb_iter_3, x_comb_iter_4\n ], 1\n )\n return x_out\n\n\nclass NormalCell(nn.Module):\n\n def __init__(\n self, in_channels_left, out_channels_left, in_channels_right,\n out_channels_right\n ):\n super(NormalCell, self).__init__()\n self.conv_prev_1x1 = nn.Sequential()\n self.conv_prev_1x1.add_module('relu', nn.ReLU())\n self.conv_prev_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_left, out_channels_left, 1, stride=1, bias=False\n )\n )\n self.conv_prev_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_left, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.conv_1x1 = nn.Sequential()\n self.conv_1x1.add_module('relu', nn.ReLU())\n self.conv_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_right, out_channels_right, 1, stride=1, bias=False\n )\n )\n self.conv_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_right, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.comb_iter_0_left = BranchSeparables(\n out_channels_right, out_channels_right, 5, 1, 2, bias=False\n )\n self.comb_iter_0_right = BranchSeparables(\n out_channels_left, out_channels_left, 3, 1, 1, bias=False\n )\n\n self.comb_iter_1_left = BranchSeparables(\n out_channels_left, out_channels_left, 5, 1, 2, bias=False\n )\n self.comb_iter_1_right = BranchSeparables(\n out_channels_left, out_channels_left, 3, 1, 1, bias=False\n )\n\n self.comb_iter_2_left = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_3_left = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n self.comb_iter_3_right = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_4_left = BranchSeparables(\n out_channels_right, out_channels_right, 3, 1, 1, bias=False\n )\n\n def forward(self, x, x_prev):\n x_left = self.conv_prev_1x1(x_prev)\n x_right = self.conv_1x1(x)\n\n x_comb_iter_0_left = self.comb_iter_0_left(x_right)\n x_comb_iter_0_right = self.comb_iter_0_right(x_left)\n x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right\n\n x_comb_iter_1_left = self.comb_iter_1_left(x_left)\n x_comb_iter_1_right = self.comb_iter_1_right(x_left)\n x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right\n\n x_comb_iter_2_left = self.comb_iter_2_left(x_right)\n x_comb_iter_2 = x_comb_iter_2_left + x_left\n\n x_comb_iter_3_left = self.comb_iter_3_left(x_left)\n x_comb_iter_3_right = self.comb_iter_3_right(x_left)\n x_comb_iter_3 = x_comb_iter_3_left + x_comb_iter_3_right\n\n x_comb_iter_4_left = self.comb_iter_4_left(x_right)\n x_comb_iter_4 = x_comb_iter_4_left + x_right\n\n x_out = torch.cat(\n [\n x_left, x_comb_iter_0, x_comb_iter_1, x_comb_iter_2,\n x_comb_iter_3, x_comb_iter_4\n ], 1\n )\n return x_out\n\n\nclass ReductionCell0(nn.Module):\n\n def __init__(\n self, in_channels_left, out_channels_left, in_channels_right,\n out_channels_right\n ):\n super(ReductionCell0, self).__init__()\n self.conv_prev_1x1 = nn.Sequential()\n self.conv_prev_1x1.add_module('relu', nn.ReLU())\n self.conv_prev_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_left, out_channels_left, 1, stride=1, bias=False\n )\n )\n self.conv_prev_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_left, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.conv_1x1 = nn.Sequential()\n self.conv_1x1.add_module('relu', nn.ReLU())\n self.conv_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_right, out_channels_right, 1, stride=1, bias=False\n )\n )\n self.conv_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_right, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.comb_iter_0_left = BranchSeparablesReduction(\n out_channels_right, out_channels_right, 5, 2, 2, bias=False\n )\n self.comb_iter_0_right = BranchSeparablesReduction(\n out_channels_right, out_channels_right, 7, 2, 3, bias=False\n )\n\n self.comb_iter_1_left = MaxPoolPad()\n self.comb_iter_1_right = BranchSeparablesReduction(\n out_channels_right, out_channels_right, 7, 2, 3, bias=False\n )\n\n self.comb_iter_2_left = AvgPoolPad()\n self.comb_iter_2_right = BranchSeparablesReduction(\n out_channels_right, out_channels_right, 5, 2, 2, bias=False\n )\n\n self.comb_iter_3_right = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_4_left = BranchSeparablesReduction(\n out_channels_right, out_channels_right, 3, 1, 1, bias=False\n )\n self.comb_iter_4_right = MaxPoolPad()\n\n def forward(self, x, x_prev):\n x_left = self.conv_prev_1x1(x_prev)\n x_right = self.conv_1x1(x)\n\n x_comb_iter_0_left = self.comb_iter_0_left(x_right)\n x_comb_iter_0_right = self.comb_iter_0_right(x_left)\n x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right\n\n x_comb_iter_1_left = self.comb_iter_1_left(x_right)\n x_comb_iter_1_right = self.comb_iter_1_right(x_left)\n x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right\n\n x_comb_iter_2_left = self.comb_iter_2_left(x_right)\n x_comb_iter_2_right = self.comb_iter_2_right(x_left)\n x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right\n\n x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)\n x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1\n\n x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)\n x_comb_iter_4_right = self.comb_iter_4_right(x_right)\n x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right\n\n x_out = torch.cat(\n [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1\n )\n return x_out\n\n\nclass ReductionCell1(nn.Module):\n\n def __init__(\n self, in_channels_left, out_channels_left, in_channels_right,\n out_channels_right\n ):\n super(ReductionCell1, self).__init__()\n self.conv_prev_1x1 = nn.Sequential()\n self.conv_prev_1x1.add_module('relu', nn.ReLU())\n self.conv_prev_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_left, out_channels_left, 1, stride=1, bias=False\n )\n )\n self.conv_prev_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_left, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.conv_1x1 = nn.Sequential()\n self.conv_1x1.add_module('relu', nn.ReLU())\n self.conv_1x1.add_module(\n 'conv',\n nn.Conv2d(\n in_channels_right, out_channels_right, 1, stride=1, bias=False\n )\n )\n self.conv_1x1.add_module(\n 'bn',\n nn.BatchNorm2d(\n out_channels_right, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.comb_iter_0_left = BranchSeparables(\n out_channels_right,\n out_channels_right,\n 5,\n 2,\n 2,\n name='specific',\n bias=False\n )\n self.comb_iter_0_right = BranchSeparables(\n out_channels_right,\n out_channels_right,\n 7,\n 2,\n 3,\n name='specific',\n bias=False\n )\n\n # self.comb_iter_1_left = nn.MaxPool2d(3, stride=2, padding=1)\n self.comb_iter_1_left = MaxPoolPad()\n self.comb_iter_1_right = BranchSeparables(\n out_channels_right,\n out_channels_right,\n 7,\n 2,\n 3,\n name='specific',\n bias=False\n )\n\n # self.comb_iter_2_left = nn.AvgPool2d(3, stride=2, padding=1, count_include_pad=False)\n self.comb_iter_2_left = AvgPoolPad()\n self.comb_iter_2_right = BranchSeparables(\n out_channels_right,\n out_channels_right,\n 5,\n 2,\n 2,\n name='specific',\n bias=False\n )\n\n self.comb_iter_3_right = nn.AvgPool2d(\n 3, stride=1, padding=1, count_include_pad=False\n )\n\n self.comb_iter_4_left = BranchSeparables(\n out_channels_right,\n out_channels_right,\n 3,\n 1,\n 1,\n name='specific',\n bias=False\n )\n # self.comb_iter_4_right = nn.MaxPool2d(3, stride=2, padding=1)\n self.comb_iter_4_right = MaxPoolPad()\n\n def forward(self, x, x_prev):\n x_left = self.conv_prev_1x1(x_prev)\n x_right = self.conv_1x1(x)\n\n x_comb_iter_0_left = self.comb_iter_0_left(x_right)\n x_comb_iter_0_right = self.comb_iter_0_right(x_left)\n x_comb_iter_0 = x_comb_iter_0_left + x_comb_iter_0_right\n\n x_comb_iter_1_left = self.comb_iter_1_left(x_right)\n x_comb_iter_1_right = self.comb_iter_1_right(x_left)\n x_comb_iter_1 = x_comb_iter_1_left + x_comb_iter_1_right\n\n x_comb_iter_2_left = self.comb_iter_2_left(x_right)\n x_comb_iter_2_right = self.comb_iter_2_right(x_left)\n x_comb_iter_2 = x_comb_iter_2_left + x_comb_iter_2_right\n\n x_comb_iter_3_right = self.comb_iter_3_right(x_comb_iter_0)\n x_comb_iter_3 = x_comb_iter_3_right + x_comb_iter_1\n\n x_comb_iter_4_left = self.comb_iter_4_left(x_comb_iter_0)\n x_comb_iter_4_right = self.comb_iter_4_right(x_right)\n x_comb_iter_4 = x_comb_iter_4_left + x_comb_iter_4_right\n\n x_out = torch.cat(\n [x_comb_iter_1, x_comb_iter_2, x_comb_iter_3, x_comb_iter_4], 1\n )\n return x_out\n\n\nclass NASNetAMobile(nn.Module):\n \"\"\"Neural Architecture Search (NAS).\n\n Reference:\n Zoph et al. Learning Transferable Architectures\n for Scalable Image Recognition. CVPR 2018.\n\n Public keys:\n - ``nasnetamobile``: NASNet-A Mobile.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n stem_filters=32,\n penultimate_filters=1056,\n filters_multiplier=2,\n **kwargs\n ):\n super(NASNetAMobile, self).__init__()\n self.stem_filters = stem_filters\n self.penultimate_filters = penultimate_filters\n self.filters_multiplier = filters_multiplier\n self.loss = loss\n\n filters = self.penultimate_filters // 24\n # 24 is default value for the architecture\n\n self.conv0 = nn.Sequential()\n self.conv0.add_module(\n 'conv',\n nn.Conv2d(\n in_channels=3,\n out_channels=self.stem_filters,\n kernel_size=3,\n padding=0,\n stride=2,\n bias=False\n )\n )\n self.conv0.add_module(\n 'bn',\n nn.BatchNorm2d(\n self.stem_filters, eps=0.001, momentum=0.1, affine=True\n )\n )\n\n self.cell_stem_0 = CellStem0(\n self.stem_filters, num_filters=filters // (filters_multiplier**2)\n )\n self.cell_stem_1 = CellStem1(\n self.stem_filters, num_filters=filters // filters_multiplier\n )\n\n self.cell_0 = FirstCell(\n in_channels_left=filters,\n out_channels_left=filters // 2, # 1, 0.5\n in_channels_right=2 * filters,\n out_channels_right=filters\n ) # 2, 1\n self.cell_1 = NormalCell(\n in_channels_left=2 * filters,\n out_channels_left=filters, # 2, 1\n in_channels_right=6 * filters,\n out_channels_right=filters\n ) # 6, 1\n self.cell_2 = NormalCell(\n in_channels_left=6 * filters,\n out_channels_left=filters, # 6, 1\n in_channels_right=6 * filters,\n out_channels_right=filters\n ) # 6, 1\n self.cell_3 = NormalCell(\n in_channels_left=6 * filters,\n out_channels_left=filters, # 6, 1\n in_channels_right=6 * filters,\n out_channels_right=filters\n ) # 6, 1\n\n self.reduction_cell_0 = ReductionCell0(\n in_channels_left=6 * filters,\n out_channels_left=2 * filters, # 6, 2\n in_channels_right=6 * filters,\n out_channels_right=2 * filters\n ) # 6, 2\n\n self.cell_6 = FirstCell(\n in_channels_left=6 * filters,\n out_channels_left=filters, # 6, 1\n in_channels_right=8 * filters,\n out_channels_right=2 * filters\n ) # 8, 2\n self.cell_7 = NormalCell(\n in_channels_left=8 * filters,\n out_channels_left=2 * filters, # 8, 2\n in_channels_right=12 * filters,\n out_channels_right=2 * filters\n ) # 12, 2\n self.cell_8 = NormalCell(\n in_channels_left=12 * filters,\n out_channels_left=2 * filters, # 12, 2\n in_channels_right=12 * filters,\n out_channels_right=2 * filters\n ) # 12, 2\n self.cell_9 = NormalCell(\n in_channels_left=12 * filters,\n out_channels_left=2 * filters, # 12, 2\n in_channels_right=12 * filters,\n out_channels_right=2 * filters\n ) # 12, 2\n\n self.reduction_cell_1 = ReductionCell1(\n in_channels_left=12 * filters,\n out_channels_left=4 * filters, # 12, 4\n in_channels_right=12 * filters,\n out_channels_right=4 * filters\n ) # 12, 4\n\n self.cell_12 = FirstCell(\n in_channels_left=12 * filters,\n out_channels_left=2 * filters, # 12, 2\n in_channels_right=16 * filters,\n out_channels_right=4 * filters\n ) # 16, 4\n self.cell_13 = NormalCell(\n in_channels_left=16 * filters,\n out_channels_left=4 * filters, # 16, 4\n in_channels_right=24 * filters,\n out_channels_right=4 * filters\n ) # 24, 4\n self.cell_14 = NormalCell(\n in_channels_left=24 * filters,\n out_channels_left=4 * filters, # 24, 4\n in_channels_right=24 * filters,\n out_channels_right=4 * filters\n ) # 24, 4\n self.cell_15 = NormalCell(\n in_channels_left=24 * filters,\n out_channels_left=4 * filters, # 24, 4\n in_channels_right=24 * filters,\n out_channels_right=4 * filters\n ) # 24, 4\n\n self.relu = nn.ReLU()\n self.dropout = nn.Dropout()\n self.classifier = nn.Linear(24 * filters, num_classes)\n\n self._init_params()\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def features(self, input):\n x_conv0 = self.conv0(input)\n x_stem_0 = self.cell_stem_0(x_conv0)\n x_stem_1 = self.cell_stem_1(x_conv0, x_stem_0)\n\n x_cell_0 = self.cell_0(x_stem_1, x_stem_0)\n x_cell_1 = self.cell_1(x_cell_0, x_stem_1)\n x_cell_2 = self.cell_2(x_cell_1, x_cell_0)\n x_cell_3 = self.cell_3(x_cell_2, x_cell_1)\n\n x_reduction_cell_0 = self.reduction_cell_0(x_cell_3, x_cell_2)\n\n x_cell_6 = self.cell_6(x_reduction_cell_0, x_cell_3)\n x_cell_7 = self.cell_7(x_cell_6, x_reduction_cell_0)\n x_cell_8 = self.cell_8(x_cell_7, x_cell_6)\n x_cell_9 = self.cell_9(x_cell_8, x_cell_7)\n\n x_reduction_cell_1 = self.reduction_cell_1(x_cell_9, x_cell_8)\n\n x_cell_12 = self.cell_12(x_reduction_cell_1, x_cell_9)\n x_cell_13 = self.cell_13(x_cell_12, x_reduction_cell_1)\n x_cell_14 = self.cell_14(x_cell_13, x_cell_12)\n x_cell_15 = self.cell_15(x_cell_14, x_cell_13)\n\n x_cell_15 = self.relu(x_cell_15)\n x_cell_15 = F.avg_pool2d(\n x_cell_15,\n x_cell_15.size()[2:]\n ) # global average pool\n x_cell_15 = x_cell_15.view(x_cell_15.size(0), -1)\n x_cell_15 = self.dropout(x_cell_15)\n\n return x_cell_15\n\n def forward(self, input):\n v = self.features(input)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef nasnetamobile(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = NASNetAMobile(num_classes, loss, **kwargs)\n if pretrained:\n model_url = pretrained_settings['nasnetamobile']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n" }, { "path": "torchreid/models/osnet.py", "content": "from __future__ import division, absolute_import\nimport warnings\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = [\n 'osnet_x1_0', 'osnet_x0_75', 'osnet_x0_5', 'osnet_x0_25', 'osnet_ibn_x1_0'\n]\n\npretrained_urls = {\n 'osnet_x1_0':\n 'https://drive.google.com/uc?id=1LaG1EJpHrxdAxKnSCJ_i0u-nbxSAeiFY',\n 'osnet_x0_75':\n 'https://drive.google.com/uc?id=1uwA9fElHOk3ZogwbeY5GkLI6QPTX70Hq',\n 'osnet_x0_5':\n 'https://drive.google.com/uc?id=16DGLbZukvVYgINws8u8deSaOqjybZ83i',\n 'osnet_x0_25':\n 'https://drive.google.com/uc?id=1rb8UN5ZzPKRc_xvtHlyDh-cSz88YX9hs',\n 'osnet_ibn_x1_0':\n 'https://drive.google.com/uc?id=1sr90V6irlYYDd4_4ISU2iruoRG8J__6l'\n}\n\n\n##########\n# Basic layers\n##########\nclass ConvLayer(nn.Module):\n \"\"\"Convolution layer (conv + bn + relu).\"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride=1,\n padding=0,\n groups=1,\n IN=False\n ):\n super(ConvLayer, self).__init__()\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n kernel_size,\n stride=stride,\n padding=padding,\n bias=False,\n groups=groups\n )\n if IN:\n self.bn = nn.InstanceNorm2d(out_channels, affine=True)\n else:\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n\nclass Conv1x1(nn.Module):\n \"\"\"1x1 convolution + bn + relu.\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1, groups=1):\n super(Conv1x1, self).__init__()\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n 1,\n stride=stride,\n padding=0,\n bias=False,\n groups=groups\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n\nclass Conv1x1Linear(nn.Module):\n \"\"\"1x1 convolution + bn (w/o non-linearity).\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1):\n super(Conv1x1Linear, self).__init__()\n self.conv = nn.Conv2d(\n in_channels, out_channels, 1, stride=stride, padding=0, bias=False\n )\n self.bn = nn.BatchNorm2d(out_channels)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n return x\n\n\nclass Conv3x3(nn.Module):\n \"\"\"3x3 convolution + bn + relu.\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1, groups=1):\n super(Conv3x3, self).__init__()\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n 3,\n stride=stride,\n padding=1,\n bias=False,\n groups=groups\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n\nclass LightConv3x3(nn.Module):\n \"\"\"Lightweight 3x3 convolution.\n\n 1x1 (linear) + dw 3x3 (nonlinear).\n \"\"\"\n\n def __init__(self, in_channels, out_channels):\n super(LightConv3x3, self).__init__()\n self.conv1 = nn.Conv2d(\n in_channels, out_channels, 1, stride=1, padding=0, bias=False\n )\n self.conv2 = nn.Conv2d(\n out_channels,\n out_channels,\n 3,\n stride=1,\n padding=1,\n bias=False,\n groups=out_channels\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.conv2(x)\n x = self.bn(x)\n x = self.relu(x)\n return x\n\n\n##########\n# Building blocks for omni-scale feature learning\n##########\nclass ChannelGate(nn.Module):\n \"\"\"A mini-network that generates channel-wise gates conditioned on input tensor.\"\"\"\n\n def __init__(\n self,\n in_channels,\n num_gates=None,\n return_gates=False,\n gate_activation='sigmoid',\n reduction=16,\n layer_norm=False\n ):\n super(ChannelGate, self).__init__()\n if num_gates is None:\n num_gates = in_channels\n self.return_gates = return_gates\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.fc1 = nn.Conv2d(\n in_channels,\n in_channels // reduction,\n kernel_size=1,\n bias=True,\n padding=0\n )\n self.norm1 = None\n if layer_norm:\n self.norm1 = nn.LayerNorm((in_channels // reduction, 1, 1))\n self.relu = nn.ReLU(inplace=True)\n self.fc2 = nn.Conv2d(\n in_channels // reduction,\n num_gates,\n kernel_size=1,\n bias=True,\n padding=0\n )\n if gate_activation == 'sigmoid':\n self.gate_activation = nn.Sigmoid()\n elif gate_activation == 'relu':\n self.gate_activation = nn.ReLU(inplace=True)\n elif gate_activation == 'linear':\n self.gate_activation = None\n else:\n raise RuntimeError(\n \"Unknown gate activation: {}\".format(gate_activation)\n )\n\n def forward(self, x):\n input = x\n x = self.global_avgpool(x)\n x = self.fc1(x)\n if self.norm1 is not None:\n x = self.norm1(x)\n x = self.relu(x)\n x = self.fc2(x)\n if self.gate_activation is not None:\n x = self.gate_activation(x)\n if self.return_gates:\n return x\n return input * x\n\n\nclass OSBlock(nn.Module):\n \"\"\"Omni-scale feature learning block.\"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n IN=False,\n bottleneck_reduction=4,\n **kwargs\n ):\n super(OSBlock, self).__init__()\n mid_channels = out_channels // bottleneck_reduction\n self.conv1 = Conv1x1(in_channels, mid_channels)\n self.conv2a = LightConv3x3(mid_channels, mid_channels)\n self.conv2b = nn.Sequential(\n LightConv3x3(mid_channels, mid_channels),\n LightConv3x3(mid_channels, mid_channels),\n )\n self.conv2c = nn.Sequential(\n LightConv3x3(mid_channels, mid_channels),\n LightConv3x3(mid_channels, mid_channels),\n LightConv3x3(mid_channels, mid_channels),\n )\n self.conv2d = nn.Sequential(\n LightConv3x3(mid_channels, mid_channels),\n LightConv3x3(mid_channels, mid_channels),\n LightConv3x3(mid_channels, mid_channels),\n LightConv3x3(mid_channels, mid_channels),\n )\n self.gate = ChannelGate(mid_channels)\n self.conv3 = Conv1x1Linear(mid_channels, out_channels)\n self.downsample = None\n if in_channels != out_channels:\n self.downsample = Conv1x1Linear(in_channels, out_channels)\n self.IN = None\n if IN:\n self.IN = nn.InstanceNorm2d(out_channels, affine=True)\n\n def forward(self, x):\n identity = x\n x1 = self.conv1(x)\n x2a = self.conv2a(x1)\n x2b = self.conv2b(x1)\n x2c = self.conv2c(x1)\n x2d = self.conv2d(x1)\n x2 = self.gate(x2a) + self.gate(x2b) + self.gate(x2c) + self.gate(x2d)\n x3 = self.conv3(x2)\n if self.downsample is not None:\n identity = self.downsample(identity)\n out = x3 + identity\n if self.IN is not None:\n out = self.IN(out)\n return F.relu(out)\n\n\n##########\n# Network architecture\n##########\nclass OSNet(nn.Module):\n \"\"\"Omni-Scale Network.\n \n Reference:\n - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.\n - Zhou et al. Learning Generalisable Omni-Scale Representations\n for Person Re-Identification. arXiv preprint, 2019.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n blocks,\n layers,\n channels,\n feature_dim=512,\n loss='softmax',\n IN=False,\n **kwargs\n ):\n super(OSNet, self).__init__()\n num_blocks = len(blocks)\n assert num_blocks == len(layers)\n assert num_blocks == len(channels) - 1\n self.loss = loss\n\n # convolutional backbone\n self.conv1 = ConvLayer(3, channels[0], 7, stride=2, padding=3, IN=IN)\n self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)\n self.conv2 = self._make_layer(\n blocks[0],\n layers[0],\n channels[0],\n channels[1],\n reduce_spatial_size=True,\n IN=IN\n )\n self.conv3 = self._make_layer(\n blocks[1],\n layers[1],\n channels[1],\n channels[2],\n reduce_spatial_size=True\n )\n self.conv4 = self._make_layer(\n blocks[2],\n layers[2],\n channels[2],\n channels[3],\n reduce_spatial_size=False\n )\n self.conv5 = Conv1x1(channels[3], channels[3])\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n # fully connected layer\n self.fc = self._construct_fc_layer(\n feature_dim, channels[3], dropout_p=None\n )\n # identity classification layer\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _make_layer(\n self,\n block,\n layer,\n in_channels,\n out_channels,\n reduce_spatial_size,\n IN=False\n ):\n layers = []\n\n layers.append(block(in_channels, out_channels, IN=IN))\n for i in range(1, layer):\n layers.append(block(out_channels, out_channels, IN=IN))\n\n if reduce_spatial_size:\n layers.append(\n nn.Sequential(\n Conv1x1(out_channels, out_channels),\n nn.AvgPool2d(2, stride=2)\n )\n )\n\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n if fc_dims is None or fc_dims < 0:\n self.feature_dim = input_dim\n return None\n\n if isinstance(fc_dims, int):\n fc_dims = [fc_dims]\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.maxpool(x)\n x = self.conv2(x)\n x = self.conv3(x)\n x = self.conv4(x)\n x = self.conv5(x)\n return x\n\n def forward(self, x, return_featuremaps=False):\n x = self.featuremaps(x)\n if return_featuremaps:\n return x\n v = self.global_avgpool(x)\n v = v.view(v.size(0), -1)\n if self.fc is not None:\n v = self.fc(v)\n if not self.training:\n return v\n y = self.classifier(v)\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, key=''):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n import os\n import errno\n import gdown\n from collections import OrderedDict\n\n def _get_torch_home():\n ENV_TORCH_HOME = 'TORCH_HOME'\n ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'\n DEFAULT_CACHE_DIR = '~/.cache'\n torch_home = os.path.expanduser(\n os.getenv(\n ENV_TORCH_HOME,\n os.path.join(\n os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'\n )\n )\n )\n return torch_home\n\n torch_home = _get_torch_home()\n model_dir = os.path.join(torch_home, 'checkpoints')\n try:\n os.makedirs(model_dir)\n except OSError as e:\n if e.errno == errno.EEXIST:\n # Directory already exists, ignore.\n pass\n else:\n # Unexpected OSError, re-raise.\n raise\n filename = key + '_imagenet.pth'\n cached_file = os.path.join(model_dir, filename)\n\n if not os.path.exists(cached_file):\n gdown.download(pretrained_urls[key], cached_file, quiet=False)\n\n state_dict = torch.load(cached_file)\n model_dict = model.state_dict()\n new_state_dict = OrderedDict()\n matched_layers, discarded_layers = [], []\n\n for k, v in state_dict.items():\n if k.startswith('module.'):\n k = k[7:] # discard module.\n\n if k in model_dict and model_dict[k].size() == v.size():\n new_state_dict[k] = v\n matched_layers.append(k)\n else:\n discarded_layers.append(k)\n\n model_dict.update(new_state_dict)\n model.load_state_dict(model_dict)\n\n if len(matched_layers) == 0:\n warnings.warn(\n 'The pretrained weights from \"{}\" cannot be loaded, '\n 'please check the key names manually '\n '(** ignored and continue **)'.format(cached_file)\n )\n else:\n print(\n 'Successfully loaded imagenet pretrained weights from \"{}\"'.\n format(cached_file)\n )\n if len(discarded_layers) > 0:\n print(\n '** The following layers are discarded '\n 'due to unmatched keys or layer size: {}'.\n format(discarded_layers)\n )\n\n\n##########\n# Instantiation\n##########\ndef osnet_x1_0(num_classes=1000, pretrained=True, loss='softmax', **kwargs):\n # standard size (width x1.0)\n model = OSNet(\n num_classes,\n blocks=[OSBlock, OSBlock, OSBlock],\n layers=[2, 2, 2],\n channels=[64, 256, 384, 512],\n loss=loss,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, key='osnet_x1_0')\n return model\n\n\ndef osnet_x0_75(num_classes=1000, pretrained=True, loss='softmax', **kwargs):\n # medium size (width x0.75)\n model = OSNet(\n num_classes,\n blocks=[OSBlock, OSBlock, OSBlock],\n layers=[2, 2, 2],\n channels=[48, 192, 288, 384],\n loss=loss,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, key='osnet_x0_75')\n return model\n\n\ndef osnet_x0_5(num_classes=1000, pretrained=True, loss='softmax', **kwargs):\n # tiny size (width x0.5)\n model = OSNet(\n num_classes,\n blocks=[OSBlock, OSBlock, OSBlock],\n layers=[2, 2, 2],\n channels=[32, 128, 192, 256],\n loss=loss,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, key='osnet_x0_5')\n return model\n\n\ndef osnet_x0_25(num_classes=1000, pretrained=True, loss='softmax', **kwargs):\n # very tiny size (width x0.25)\n model = OSNet(\n num_classes,\n blocks=[OSBlock, OSBlock, OSBlock],\n layers=[2, 2, 2],\n channels=[16, 64, 96, 128],\n loss=loss,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, key='osnet_x0_25')\n return model\n\n\ndef osnet_ibn_x1_0(\n num_classes=1000, pretrained=True, loss='softmax', **kwargs\n):\n # standard size (width x1.0) + IBN layer\n # Ref: Pan et al. Two at Once: Enhancing Learning and Generalization Capacities via IBN-Net. ECCV, 2018.\n model = OSNet(\n num_classes,\n blocks=[OSBlock, OSBlock, OSBlock],\n layers=[2, 2, 2],\n channels=[64, 256, 384, 512],\n loss=loss,\n IN=True,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, key='osnet_ibn_x1_0')\n return model\n" }, { "path": "torchreid/models/osnet_ain.py", "content": "from __future__ import division, absolute_import\nimport warnings\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['osnet_ain_x1_0']\n\npretrained_urls = {\n 'osnet_ain_x1_0':\n 'https://drive.google.com/uc?id=1-CaioD9NaqbHK_kzSMW8VE4_3KcsRjEo'\n}\n\n\n##########\n# Basic layers\n##########\nclass ConvLayer(nn.Module):\n \"\"\"Convolution layer (conv + bn + relu).\"\"\"\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride=1,\n padding=0,\n groups=1,\n IN=False\n ):\n super(ConvLayer, self).__init__()\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n kernel_size,\n stride=stride,\n padding=padding,\n bias=False,\n groups=groups\n )\n if IN:\n self.bn = nn.InstanceNorm2d(out_channels, affine=True)\n else:\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU()\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n return self.relu(x)\n\n\nclass Conv1x1(nn.Module):\n \"\"\"1x1 convolution + bn + relu.\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1, groups=1):\n super(Conv1x1, self).__init__()\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n 1,\n stride=stride,\n padding=0,\n bias=False,\n groups=groups\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU()\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n return self.relu(x)\n\n\nclass Conv1x1Linear(nn.Module):\n \"\"\"1x1 convolution + bn (w/o non-linearity).\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1, bn=True):\n super(Conv1x1Linear, self).__init__()\n self.conv = nn.Conv2d(\n in_channels, out_channels, 1, stride=stride, padding=0, bias=False\n )\n self.bn = None\n if bn:\n self.bn = nn.BatchNorm2d(out_channels)\n\n def forward(self, x):\n x = self.conv(x)\n if self.bn is not None:\n x = self.bn(x)\n return x\n\n\nclass Conv3x3(nn.Module):\n \"\"\"3x3 convolution + bn + relu.\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1, groups=1):\n super(Conv3x3, self).__init__()\n self.conv = nn.Conv2d(\n in_channels,\n out_channels,\n 3,\n stride=stride,\n padding=1,\n bias=False,\n groups=groups\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU()\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n return self.relu(x)\n\n\nclass LightConv3x3(nn.Module):\n \"\"\"Lightweight 3x3 convolution.\n\n 1x1 (linear) + dw 3x3 (nonlinear).\n \"\"\"\n\n def __init__(self, in_channels, out_channels):\n super(LightConv3x3, self).__init__()\n self.conv1 = nn.Conv2d(\n in_channels, out_channels, 1, stride=1, padding=0, bias=False\n )\n self.conv2 = nn.Conv2d(\n out_channels,\n out_channels,\n 3,\n stride=1,\n padding=1,\n bias=False,\n groups=out_channels\n )\n self.bn = nn.BatchNorm2d(out_channels)\n self.relu = nn.ReLU()\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.conv2(x)\n x = self.bn(x)\n return self.relu(x)\n\n\nclass LightConvStream(nn.Module):\n \"\"\"Lightweight convolution stream.\"\"\"\n\n def __init__(self, in_channels, out_channels, depth):\n super(LightConvStream, self).__init__()\n assert depth >= 1, 'depth must be equal to or larger than 1, but got {}'.format(\n depth\n )\n layers = []\n layers += [LightConv3x3(in_channels, out_channels)]\n for i in range(depth - 1):\n layers += [LightConv3x3(out_channels, out_channels)]\n self.layers = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.layers(x)\n\n\n##########\n# Building blocks for omni-scale feature learning\n##########\nclass ChannelGate(nn.Module):\n \"\"\"A mini-network that generates channel-wise gates conditioned on input tensor.\"\"\"\n\n def __init__(\n self,\n in_channels,\n num_gates=None,\n return_gates=False,\n gate_activation='sigmoid',\n reduction=16,\n layer_norm=False\n ):\n super(ChannelGate, self).__init__()\n if num_gates is None:\n num_gates = in_channels\n self.return_gates = return_gates\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.fc1 = nn.Conv2d(\n in_channels,\n in_channels // reduction,\n kernel_size=1,\n bias=True,\n padding=0\n )\n self.norm1 = None\n if layer_norm:\n self.norm1 = nn.LayerNorm((in_channels // reduction, 1, 1))\n self.relu = nn.ReLU()\n self.fc2 = nn.Conv2d(\n in_channels // reduction,\n num_gates,\n kernel_size=1,\n bias=True,\n padding=0\n )\n if gate_activation == 'sigmoid':\n self.gate_activation = nn.Sigmoid()\n elif gate_activation == 'relu':\n self.gate_activation = nn.ReLU()\n elif gate_activation == 'linear':\n self.gate_activation = None\n else:\n raise RuntimeError(\n \"Unknown gate activation: {}\".format(gate_activation)\n )\n\n def forward(self, x):\n input = x\n x = self.global_avgpool(x)\n x = self.fc1(x)\n if self.norm1 is not None:\n x = self.norm1(x)\n x = self.relu(x)\n x = self.fc2(x)\n if self.gate_activation is not None:\n x = self.gate_activation(x)\n if self.return_gates:\n return x\n return input * x\n\n\nclass OSBlock(nn.Module):\n \"\"\"Omni-scale feature learning block.\"\"\"\n\n def __init__(self, in_channels, out_channels, reduction=4, T=4, **kwargs):\n super(OSBlock, self).__init__()\n assert T >= 1\n assert out_channels >= reduction and out_channels % reduction == 0\n mid_channels = out_channels // reduction\n\n self.conv1 = Conv1x1(in_channels, mid_channels)\n self.conv2 = nn.ModuleList()\n for t in range(1, T + 1):\n self.conv2 += [LightConvStream(mid_channels, mid_channels, t)]\n self.gate = ChannelGate(mid_channels)\n self.conv3 = Conv1x1Linear(mid_channels, out_channels)\n self.downsample = None\n if in_channels != out_channels:\n self.downsample = Conv1x1Linear(in_channels, out_channels)\n\n def forward(self, x):\n identity = x\n x1 = self.conv1(x)\n x2 = 0\n for conv2_t in self.conv2:\n x2_t = conv2_t(x1)\n x2 = x2 + self.gate(x2_t)\n x3 = self.conv3(x2)\n if self.downsample is not None:\n identity = self.downsample(identity)\n out = x3 + identity\n return F.relu(out)\n\n\nclass OSBlockINin(nn.Module):\n \"\"\"Omni-scale feature learning block with instance normalization.\"\"\"\n\n def __init__(self, in_channels, out_channels, reduction=4, T=4, **kwargs):\n super(OSBlockINin, self).__init__()\n assert T >= 1\n assert out_channels >= reduction and out_channels % reduction == 0\n mid_channels = out_channels // reduction\n\n self.conv1 = Conv1x1(in_channels, mid_channels)\n self.conv2 = nn.ModuleList()\n for t in range(1, T + 1):\n self.conv2 += [LightConvStream(mid_channels, mid_channels, t)]\n self.gate = ChannelGate(mid_channels)\n self.conv3 = Conv1x1Linear(mid_channels, out_channels, bn=False)\n self.downsample = None\n if in_channels != out_channels:\n self.downsample = Conv1x1Linear(in_channels, out_channels)\n self.IN = nn.InstanceNorm2d(out_channels, affine=True)\n\n def forward(self, x):\n identity = x\n x1 = self.conv1(x)\n x2 = 0\n for conv2_t in self.conv2:\n x2_t = conv2_t(x1)\n x2 = x2 + self.gate(x2_t)\n x3 = self.conv3(x2)\n x3 = self.IN(x3) # IN inside residual\n if self.downsample is not None:\n identity = self.downsample(identity)\n out = x3 + identity\n return F.relu(out)\n\n\n##########\n# Network architecture\n##########\nclass OSNet(nn.Module):\n \"\"\"Omni-Scale Network.\n \n Reference:\n - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019.\n - Zhou et al. Learning Generalisable Omni-Scale Representations\n for Person Re-Identification. arXiv preprint, 2019.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n blocks,\n layers,\n channels,\n feature_dim=512,\n loss='softmax',\n conv1_IN=False,\n **kwargs\n ):\n super(OSNet, self).__init__()\n num_blocks = len(blocks)\n assert num_blocks == len(layers)\n assert num_blocks == len(channels) - 1\n self.loss = loss\n self.feature_dim = feature_dim\n\n # convolutional backbone\n self.conv1 = ConvLayer(\n 3, channels[0], 7, stride=2, padding=3, IN=conv1_IN\n )\n self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)\n self.conv2 = self._make_layer(\n blocks[0], layers[0], channels[0], channels[1]\n )\n self.pool2 = nn.Sequential(\n Conv1x1(channels[1], channels[1]), nn.AvgPool2d(2, stride=2)\n )\n self.conv3 = self._make_layer(\n blocks[1], layers[1], channels[1], channels[2]\n )\n self.pool3 = nn.Sequential(\n Conv1x1(channels[2], channels[2]), nn.AvgPool2d(2, stride=2)\n )\n self.conv4 = self._make_layer(\n blocks[2], layers[2], channels[2], channels[3]\n )\n self.conv5 = Conv1x1(channels[3], channels[3])\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n # fully connected layer\n self.fc = self._construct_fc_layer(\n self.feature_dim, channels[3], dropout_p=None\n )\n # identity classification layer\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _make_layer(self, blocks, layer, in_channels, out_channels):\n layers = []\n layers += [blocks[0](in_channels, out_channels)]\n for i in range(1, len(blocks)):\n layers += [blocks[i](out_channels, out_channels)]\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n if fc_dims is None or fc_dims < 0:\n self.feature_dim = input_dim\n return None\n\n if isinstance(fc_dims, int):\n fc_dims = [fc_dims]\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU())\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.InstanceNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.maxpool(x)\n x = self.conv2(x)\n x = self.pool2(x)\n x = self.conv3(x)\n x = self.pool3(x)\n x = self.conv4(x)\n x = self.conv5(x)\n return x\n\n def forward(self, x, return_featuremaps=False):\n x = self.featuremaps(x)\n if return_featuremaps:\n return x\n v = self.global_avgpool(x)\n v = v.view(v.size(0), -1)\n if self.fc is not None:\n v = self.fc(v)\n if not self.training:\n return v\n y = self.classifier(v)\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, key=''):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n import os\n import errno\n import gdown\n from collections import OrderedDict\n\n def _get_torch_home():\n ENV_TORCH_HOME = 'TORCH_HOME'\n ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'\n DEFAULT_CACHE_DIR = '~/.cache'\n torch_home = os.path.expanduser(\n os.getenv(\n ENV_TORCH_HOME,\n os.path.join(\n os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'\n )\n )\n )\n return torch_home\n\n torch_home = _get_torch_home()\n model_dir = os.path.join(torch_home, 'checkpoints')\n try:\n os.makedirs(model_dir)\n except OSError as e:\n if e.errno == errno.EEXIST:\n # Directory already exists, ignore.\n pass\n else:\n # Unexpected OSError, re-raise.\n raise\n filename = key + '_imagenet.pth'\n cached_file = os.path.join(model_dir, filename)\n\n if not os.path.exists(cached_file):\n gdown.download(pretrained_urls[key], cached_file, quiet=False)\n\n state_dict = torch.load(cached_file)\n model_dict = model.state_dict()\n new_state_dict = OrderedDict()\n matched_layers, discarded_layers = [], []\n\n for k, v in state_dict.items():\n if k.startswith('module.'):\n k = k[7:] # discard module.\n\n if k in model_dict and model_dict[k].size() == v.size():\n new_state_dict[k] = v\n matched_layers.append(k)\n else:\n discarded_layers.append(k)\n\n model_dict.update(new_state_dict)\n model.load_state_dict(model_dict)\n\n if len(matched_layers) == 0:\n warnings.warn(\n 'The pretrained weights from \"{}\" cannot be loaded, '\n 'please check the key names manually '\n '(** ignored and continue **)'.format(cached_file)\n )\n else:\n print(\n 'Successfully loaded imagenet pretrained weights from \"{}\"'.\n format(cached_file)\n )\n if len(discarded_layers) > 0:\n print(\n '** The following layers are discarded '\n 'due to unmatched keys or layer size: {}'.\n format(discarded_layers)\n )\n\n\n##########\n# Instantiation\n##########\ndef osnet_ain_x1_0(\n num_classes=1000, pretrained=True, loss='softmax', **kwargs\n):\n model = OSNet(\n num_classes,\n blocks=[\n [OSBlockINin, OSBlockINin], [OSBlock, OSBlockINin],\n [OSBlockINin, OSBlock]\n ],\n layers=[2, 2, 2],\n channels=[64, 256, 384, 512],\n loss=loss,\n conv1_IN=True,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, key='osnet_ain_x1_0')\n return model\n" }, { "path": "torchreid/models/pcb.py", "content": "from __future__ import division, absolute_import\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['pcb_p6', 'pcb_p4']\n\nmodel_urls = {\n 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm2d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = nn.BatchNorm2d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(\n planes,\n planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(\n planes, planes * self.expansion, kernel_size=1, bias=False\n )\n self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass DimReduceLayer(nn.Module):\n\n def __init__(self, in_channels, out_channels, nonlinear):\n super(DimReduceLayer, self).__init__()\n layers = []\n layers.append(\n nn.Conv2d(\n in_channels, out_channels, 1, stride=1, padding=0, bias=False\n )\n )\n layers.append(nn.BatchNorm2d(out_channels))\n\n if nonlinear == 'relu':\n layers.append(nn.ReLU(inplace=True))\n elif nonlinear == 'leakyrelu':\n layers.append(nn.LeakyReLU(0.1))\n\n self.layers = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.layers(x)\n\n\nclass PCB(nn.Module):\n \"\"\"Part-based Convolutional Baseline.\n\n Reference:\n Sun et al. Beyond Part Models: Person Retrieval with Refined\n Part Pooling (and A Strong Convolutional Baseline). ECCV 2018.\n\n Public keys:\n - ``pcb_p4``: PCB with 4-part strips.\n - ``pcb_p6``: PCB with 6-part strips.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n block,\n layers,\n parts=6,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n ):\n self.inplanes = 64\n super(PCB, self).__init__()\n self.loss = loss\n self.parts = parts\n self.feature_dim = 512 * block.expansion\n\n # backbone network\n self.conv1 = nn.Conv2d(\n 3, 64, kernel_size=7, stride=2, padding=3, bias=False\n )\n self.bn1 = nn.BatchNorm2d(64)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=1)\n\n # pcb layers\n self.parts_avgpool = nn.AdaptiveAvgPool2d((self.parts, 1))\n self.dropout = nn.Dropout(p=0.5)\n self.conv5 = DimReduceLayer(\n 512 * block.expansion, reduced_dim, nonlinear=nonlinear\n )\n self.feature_dim = reduced_dim\n self.classifier = nn.ModuleList(\n [\n nn.Linear(self.feature_dim, num_classes)\n for _ in range(self.parts)\n ]\n )\n\n self._init_params()\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(\n self.inplanes,\n planes * block.expansion,\n kernel_size=1,\n stride=stride,\n bias=False\n ),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes))\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v_g = self.parts_avgpool(f)\n\n if not self.training:\n v_g = F.normalize(v_g, p=2, dim=1)\n return v_g.view(v_g.size(0), -1)\n\n v_g = self.dropout(v_g)\n v_h = self.conv5(v_g)\n\n y = []\n for i in range(self.parts):\n v_h_i = v_h[:, :, i, :]\n v_h_i = v_h_i.view(v_h_i.size(0), -1)\n y_i = self.classifier[i](v_h_i)\n y.append(y_i)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n v_g = F.normalize(v_g, p=2, dim=1)\n return y, v_g.view(v_g.size(0), -1)\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef pcb_p6(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = PCB(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=1,\n parts=6,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\ndef pcb_p4(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = PCB(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=1,\n parts=4,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n" }, { "path": "torchreid/models/pvpm.py", "content": "from __future__ import absolute_import\nfrom __future__ import division\n\n# Source: https://github.com/hh23333/PVPM\n\n__all__ = ['pcb_p6', 'pcb_p4', 'pose_resnet50_256_p4', 'pose_resnet50_256_p6',\n 'pose_resnet50_256_p6_pscore_reg', 'pose_resnet50_256_p4_pscore_reg']\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nimport torchvision\nimport torch.utils.model_zoo as model_zoo\nfrom .osnet import ConvLayer, Conv1x1, Conv1x1Linear, Conv3x3, LightConv3x3, OSBlock\n\nmodel_urls = {\n 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,\n padding=1, bias=False)\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm2d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = nn.BatchNorm2d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n padding=1, bias=False)\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass DimReduceLayer(nn.Module):\n\n def __init__(self, in_channels, out_channels, nonlinear):\n super(DimReduceLayer, self).__init__()\n layers = []\n layers.append(nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False))\n layers.append(nn.BatchNorm2d(out_channels))\n\n if nonlinear == 'relu':\n layers.append(nn.ReLU(inplace=True))\n elif nonlinear == 'leakyrelu':\n layers.append(nn.LeakyReLU(0.1))\n\n self.layers = nn.Sequential(*layers)\n\n def forward(self, x):\n return self.layers(x)\n\n\nclass PCB(nn.Module):\n \"\"\"Part-based Convolutional Baseline.\n\n Reference:\n Sun et al. Beyond Part Models: Person Retrieval with Refined\n Part Pooling (and A Strong Convolutional Baseline). ECCV 2018.\n\n Public keys:\n - ``pcb_p4``: PCB with 4-part strips.\n - ``pcb_p6``: PCB with 6-part strips.\n \"\"\"\n\n def __init__(self, num_classes, loss, block, layers,\n parts=6,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs):\n self.inplanes = 64\n super(PCB, self).__init__()\n self.loss = loss\n self.parts = parts\n self.feature_dim = 512 * block.expansion\n\n # backbone network\n self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\n self.bn1 = nn.BatchNorm2d(64)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=1)\n\n # pcb layers\n self.parts_avgpool = nn.AdaptiveAvgPool2d((self.parts, 1))\n self.dropout = nn.Dropout(p=0.5)\n self.em = nn.ModuleList( # TODO: before = DimReduceLayer\n [self._construct_em_layer(reduced_dim, 512 * block.expansion) for _ in range(self.parts)])\n self.feature_dim = reduced_dim\n self.classifier = nn.ModuleList(\n [nn.Linear(self.feature_dim, num_classes, bias=False) for _ in range(self.parts)])\n\n self._init_params()\n\n def _construct_em_layer(self, fc_dims, input_dim, dropout_p=0.5): # TODO new\n \"\"\"\n Construct fully connected layer\n\n - fc_dims (list or tuple): dimensions of fc layers, if None,\n no fc layers are constructed\n - input_dim (int): input dimension\n - dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n layers = []\n\n # layers.append(nn.Linear(input_dim, fc_dims))\n layers.append(nn.Conv2d(input_dim, fc_dims, 1, stride=1, padding=0))\n layers.append(nn.BatchNorm2d(fc_dims))\n layers.append(nn.ReLU(inplace=True))\n # layers.append(nn.Dropout(p=dropout_p))\n\n # self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.inplanes, planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes))\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.001) # TODO Learning rate of pre-trained layers: 0.1 x base learning rate ?\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n # vis_featmat_Kmeans(f) # TODO\n # vis_featmat_DBSCAN(f) # TODO\n v_g = self.parts_avgpool(f) # nn.AdaptiveAvgPool2d((self.parts, 1))\n\n if not self.training:\n v_g = F.normalize(v_g, p=2, dim=1)\n return v_g.view(v_g.size(0), -1)\n\n # v_g = self.dropout(v_g)\n # v_h = self.conv5(v_g)\n\n y = []\n v = []\n # v_g.shape = [n, 2048, 6, 1] ?\n for i in range(self.parts): # TODO new\n v_g_i = v_g[:, :, i, :].view(v_g.size(0), -1, 1, 1)\n v_g_i = self.em[i](v_g_i) # fully connected layer, Conv2d-BatchNorm2d-ReLU\n v_h_i = v_g_i.view(v_g_i.size(0), -1)\n y_i = self.classifier[i](v_h_i)\n y.append(y_i)\n v.append(v_g_i)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n v_g = F.normalize(v_g, p=2, dim=1) # TODO\n return y, v_g.view(v_g.size(0), -1)\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n\n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size()}\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef pcb_p6(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = PCB(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=1,\n parts=6,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\ndef pcb_p4(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = PCB(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=1,\n parts=4,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\nclass Conv1x1_att(nn.Module):\n \"\"\"1x1 convolution + bn + relu.\"\"\"\n\n def __init__(self, in_channels, out_channels, stride=1, groups=1):\n super(Conv1x1_att, self).__init__()\n self.conv = nn.Conv2d(in_channels, out_channels, 1, stride=stride, padding=0,\n bias=False, groups=groups)\n self.bn = nn.BatchNorm2d(out_channels)\n self.activation = nn.Sigmoid()\n\n def forward(self, x):\n x = self.conv(x)\n x = self.bn(x)\n x = self.activation(x)\n return x\n\n\nclass score_embedding(nn.Module):\n \"\"\"1x1 convolution + bn + relu.\"\"\"\n\n def __init__(self, in_channels, out_channels):\n super(score_embedding, self).__init__()\n self.pool = nn.AdaptiveAvgPool2d((1, 1))\n self.reg = nn.Linear(in_channels, out_channels, bias=False)\n self.bn = nn.BatchNorm1d(out_channels)\n self.activation = nn.Sigmoid()\n\n def forward(self, x):\n x = self.pool(x)\n x = x.view(x.size(0), -1)\n x = self.reg(x)\n x = self.bn(x)\n x = self.activation(x)\n return x\n\n\nclass Pose_Subnet(nn.Module): # TODO\n '''\n PVP and PGA\n '''\n def __init__(self, blocks, in_channels, channels, att_num=1, IN=False, matching_score_reg=False):\n super(Pose_Subnet, self).__init__()\n num_blocks = len(blocks)\n self.conv1 = ConvLayer(in_channels, channels[0], 7, stride=1, padding=3, IN=IN)\n self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)\n self.conv2 = self._make_layer(blocks[0], 1, channels[0], channels[1], reduce_spatial_size=True)\n self.conv3 = self._make_layer(blocks[1], 1, channels[1], channels[2], reduce_spatial_size=False)\n self.conv4 = Conv3x3(channels[2], channels[2])\n # PGA\n self.conv_out = Conv1x1_att(channels[2], att_num)\n # PVP\n self.matching_score_reg = matching_score_reg\n if self.matching_score_reg:\n self.conv_score = score_embedding(channels[2], att_num)\n\n self._init_params()\n\n def _make_layer(self, block, layer, in_channels, out_channels, reduce_spatial_size, IN=False):\n layers = []\n layers.append(block(in_channels, out_channels, IN=IN, gate_reduction=4))\n for i in range(1, layer):\n layers.append(block(out_channels, out_channels, IN=IN, gate_reduction=4))\n\n if reduce_spatial_size:\n layers.append(\n nn.Sequential(\n Conv1x1(out_channels, out_channels),\n nn.AvgPool2d(2, stride=2)\n )\n )\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.maxpool(x)\n x = self.conv2(x)\n x = self.conv3(x)\n x_ = self.conv4(x)\n x = self.conv_out(x_)\n _, max_index = x.max(dim=1, keepdim=True)\n onehot_index = torch.zeros_like(x).scatter_(1, max_index, 1)\n if self.matching_score_reg:\n score = self.conv_score(x_)\n return x, score, onehot_index\n else:\n return x, onehot_index\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n\nclass pose_guide_att_Resnet(PCB):\n def __init__(self, num_classes, loss, block, layers, last_stride=2, parts=4, reduced_dim=None,\n nonlinear='relu', pose_inchannel=56, part_score_reg=False, **kwargs):\n super(pose_guide_att_Resnet, self).__init__(num_classes, loss, block, layers,\n last_stride=last_stride, parts=parts,\n reduced_dim=reduced_dim,\n nonlinear=nonlinear, **kwargs)\n self.part_score_reg = part_score_reg\n self.pose_subnet = Pose_Subnet(blocks=[OSBlock, OSBlock], in_channels=pose_inchannel,\n channels=[32, 32, 32], att_num=parts, matching_score_reg=part_score_reg)\n self.pose_pool = nn.AdaptiveAvgPool2d((1, 1))\n self.parts_avgpool = nn.ModuleList([nn.AdaptiveAvgPool2d((1, 1)) for _ in range(self.parts)]) # TODO why not use the same?????\n\n def forward(self, x, pose_map):\n f = self.featuremaps(x)\n if self.part_score_reg:\n pose_att, part_score, onehot_index = self.pose_subnet(pose_map) # TODO\n else:\n pose_att, onehot_index = self.pose_subnet(pose_map)\n pose_att = pose_att * onehot_index\n pose_att_pool = self.pose_pool(pose_att) # AdaptiveAvgPool2d -> (1,1)\n v_g = []\n for i in range(self.parts):\n v_g_i = f * pose_att[:, i, :, :].unsqueeze(1) / (pose_att_pool[:, i, :, :].unsqueeze(1) + 1e-6) # TODO why divide by average?\n v_g_i = self.parts_avgpool[i](v_g_i)\n v_g.append(v_g_i)\n\n if not self.training:\n v_g = torch.cat(v_g, dim=2)\n v_g = F.normalize(v_g, p=2, dim=1) # TODO apply normalize myself?\n if self.part_score_reg:\n return v_g.squeeze(), part_score\n else:\n return v_g.view(v_g.size(0), -1)\n y = []\n v = []\n for i in range(self.parts): # add final fc layer\n v_g_i = self.em[i](v_g[i])\n v_h_i = v_g_i.view(v_g_i.size(0), -1)\n y_i = self.classifier[i](v_h_i)\n y.append(y_i)\n v.append(v_g_i)\n\n if self.loss == 'softmax':\n if self.training:\n if self.part_score_reg:\n return y, pose_att, part_score, v_g\n else:\n return y, pose_att\n else:\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef pose_resnet50_256_p4(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = pose_guide_att_Resnet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n parts=4,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\ndef pose_resnet50_256_p6(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = pose_guide_att_Resnet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n parts=6,\n reduced_dim=256,\n nonlinear='relu',\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\ndef pose_resnet50_256_p6_pscore_reg(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = pose_guide_att_Resnet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n parts=6,\n reduced_dim=256,\n nonlinear='relu',\n part_score_reg=True,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\ndef pose_resnet50_256_p4_pscore_reg(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = pose_guide_att_Resnet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n parts=4,\n reduced_dim=256,\n nonlinear='relu',\n part_score_reg=True,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n" }, { "path": "torchreid/models/resnet.py", "content": "\"\"\"\nCode source: https://github.com/pytorch/vision\n\"\"\"\nfrom __future__ import division, absolute_import\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\n\n__all__ = [\n 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152',\n 'resnext50_32x4d', 'resnext101_32x8d', 'resnet50_fc512'\n]\n\nmodel_urls = {\n 'resnet18':\n 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n 'resnet34':\n 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n 'resnet50':\n 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101':\n 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152':\n 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n 'resnext50_32x4d':\n 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',\n 'resnext101_32x8d':\n 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=3,\n stride=stride,\n padding=dilation,\n groups=groups,\n bias=False,\n dilation=dilation\n )\n\n\ndef conv1x1(in_planes, out_planes, stride=1):\n \"\"\"1x1 convolution\"\"\"\n return nn.Conv2d(\n in_planes, out_planes, kernel_size=1, stride=stride, bias=False\n )\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(\n self,\n inplanes,\n planes,\n stride=1,\n downsample=None,\n groups=1,\n base_width=64,\n dilation=1,\n norm_layer=None\n ):\n super(BasicBlock, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n if groups != 1 or base_width != 64:\n raise ValueError(\n 'BasicBlock only supports groups=1 and base_width=64'\n )\n if dilation > 1:\n raise NotImplementedError(\n \"Dilation > 1 not supported in BasicBlock\"\n )\n # Both self.conv1 and self.downsample layers downsample the input when stride != 1\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = norm_layer(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = norm_layer(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(\n self,\n inplanes,\n planes,\n stride=1,\n downsample=None,\n groups=1,\n base_width=64,\n dilation=1,\n norm_layer=None\n ):\n super(Bottleneck, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n width = int(planes * (base_width/64.)) * groups\n # Both self.conv2 and self.downsample layers downsample the input when stride != 1\n self.conv1 = conv1x1(inplanes, width)\n self.bn1 = norm_layer(width)\n self.conv2 = conv3x3(width, width, stride, groups, dilation)\n self.bn2 = norm_layer(width)\n self.conv3 = conv1x1(width, planes * self.expansion)\n self.bn3 = norm_layer(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n \"\"\"Residual network.\n \n Reference:\n - He et al. Deep Residual Learning for Image Recognition. CVPR 2016.\n - Xie et al. Aggregated Residual Transformations for Deep Neural Networks. CVPR 2017.\n\n Public keys:\n - ``resnet18``: ResNet18.\n - ``resnet34``: ResNet34.\n - ``resnet50``: ResNet50.\n - ``resnet101``: ResNet101.\n - ``resnet152``: ResNet152.\n - ``resnext50_32x4d``: ResNeXt50.\n - ``resnext101_32x8d``: ResNeXt101.\n - ``resnet50_fc512``: ResNet50 + FC.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n block,\n layers,\n zero_init_residual=False,\n groups=1,\n width_per_group=64,\n replace_stride_with_dilation=None,\n norm_layer=None,\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n ):\n super(ResNet, self).__init__()\n if norm_layer is None:\n norm_layer = nn.BatchNorm2d\n self._norm_layer = norm_layer\n self.loss = loss\n self.feature_dim = 512 * block.expansion\n self.inplanes = 64\n self.dilation = 1\n if replace_stride_with_dilation is None:\n # each element in the tuple indicates if we should replace\n # the 2x2 stride with a dilated convolution instead\n replace_stride_with_dilation = [False, False, False]\n if len(replace_stride_with_dilation) != 3:\n raise ValueError(\n \"replace_stride_with_dilation should be None \"\n \"or a 3-element tuple, got {}\".\n format(replace_stride_with_dilation)\n )\n self.groups = groups\n self.base_width = width_per_group\n self.conv1 = nn.Conv2d(\n 3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False\n )\n self.bn1 = norm_layer(self.inplanes)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(\n block,\n 128,\n layers[1],\n stride=2,\n dilate=replace_stride_with_dilation[0]\n )\n self.layer3 = self._make_layer(\n block,\n 256,\n layers[2],\n stride=2,\n dilate=replace_stride_with_dilation[1]\n )\n self.layer4 = self._make_layer(\n block,\n 512,\n layers[3],\n stride=last_stride,\n dilate=replace_stride_with_dilation[2]\n )\n self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))\n self.fc = self._construct_fc_layer(\n fc_dims, 512 * block.expansion, dropout_p\n )\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n # Zero-initialize the last BN in each residual branch,\n # so that the residual branch starts with zeros, and each residual block behaves like an identity.\n # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677\n if zero_init_residual:\n for m in self.modules():\n if isinstance(m, Bottleneck):\n nn.init.constant_(m.bn3.weight, 0)\n elif isinstance(m, BasicBlock):\n nn.init.constant_(m.bn2.weight, 0)\n\n def _make_layer(self, block, planes, blocks, stride=1, dilate=False):\n norm_layer = self._norm_layer\n downsample = None\n previous_dilation = self.dilation\n if dilate:\n self.dilation *= stride\n stride = 1\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n conv1x1(self.inplanes, planes * block.expansion, stride),\n norm_layer(planes * block.expansion),\n )\n\n layers = []\n layers.append(\n block(\n self.inplanes, planes, stride, downsample, self.groups,\n self.base_width, previous_dilation, norm_layer\n )\n )\n self.inplanes = planes * block.expansion\n for _ in range(1, blocks):\n layers.append(\n block(\n self.inplanes,\n planes,\n groups=self.groups,\n base_width=self.base_width,\n dilation=self.dilation,\n norm_layer=norm_layer\n )\n )\n\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n # x torch.Size([1, 3, 256, 128])\n x = self.conv1(x) # torch.Size([1, 64, 128, 64])\n x = self.bn1(x) # torch.Size([1, 64, 128, 64])\n x = self.relu(x) # torch.Size([1, 64, 128, 64])\n x = self.maxpool(x) # torch.Size([1, 64, 64, 32])\n x = self.layer1(x) # torch.Size([1, 256, 64, 32]) - 3 blocks\n x = self.layer2(x) # torch.Size([1, 512, 32, 16]) - 4 blocks\n x = self.layer3(x) # torch.Size([1, 1024, 16, 8]) - 6 blocks\n x = self.layer4(x) # torch.Size([1, 2048, 8, 4]) - 3 blocks\n return x\n\n def forward(self, x):\n f = self.featuremaps(x) # torch.Size([1, 2048, 8, 4])\n\n if self.loss == 'part_based':\n return f\n\n v = self.global_avgpool(f) # torch.Size([1, 2048, 1, 1])\n v = v.view(v.size(0), -1) # torch.Size([1, 2048])\n\n if self.fc is not None:\n v = self.fc(v) # torch.Size([1, 512])\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\n\"\"\"ResNet\"\"\"\n\n\ndef resnet18(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=BasicBlock,\n layers=[2, 2, 2, 2],\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet18'])\n return model\n\n\ndef resnet34(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=BasicBlock,\n layers=[3, 4, 6, 3],\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet34'])\n return model\n\n\ndef resnet50(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n\n\ndef resnet101(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 23, 3],\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet101'])\n return model\n\n\ndef resnet152(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 8, 36, 3],\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet152'])\n return model\n\n\n\"\"\"ResNeXt\"\"\"\n\n\ndef resnext50_32x4d(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n groups=32,\n width_per_group=4,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnext50_32x4d'])\n return model\n\n\ndef resnext101_32x8d(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 23, 3],\n last_stride=2,\n fc_dims=None,\n dropout_p=None,\n groups=32,\n width_per_group=8,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnext101_32x8d'])\n return model\n\n\n\"\"\"\nResNet + FC\n\"\"\"\n\n\ndef resnet50_fc512(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNet(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=1,\n fc_dims=[512],\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n" }, { "path": "torchreid/models/resnet_fastreid.py", "content": "# encoding: utf-8\n\"\"\"\n@author: liaoxingyu\n@contact: sherlockliao01@gmail.com\n\"\"\"\n\nimport math\n\nimport torch\nfrom torch import nn\n\n# from fastreid.layers import (\n# IBN,\n# SELayer,\n# Non_local,\n# get_norm,\n# )\n# from fastreid.utils.checkpoint import get_missing_parameters_message, get_unexpected_parameters_message\n# from .build import BACKBONE_REGISTRY\n# from fastreid.utils import comm\n\n\nmodel_urls = {\n '18x': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n '34x': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n '50x': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n '101x': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'ibn_18x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet18_ibn_a-2f571257.pth',\n 'ibn_34x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet34_ibn_a-94bc1577.pth',\n 'ibn_50x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet50_ibn_a-d9d0bb7b.pth',\n 'ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/resnet101_ibn_a-59ea0ac6.pth',\n 'se_ibn_101x': 'https://github.com/XingangPan/IBN-Net/releases/download/v1.0/se_resnet101_ibn_a-fabed4e2.pth',\n}\n\n\ndef get_norm(norm, out_channels, **kwargs):\n \"\"\"\n Args:\n norm (str or callable): either one of BN, GhostBN, FrozenBN, GN or SyncBN;\n or a callable that takes a channel number and returns\n the normalization layer as a nn.Module\n out_channels: number of channels for normalization layer\n\n Returns:\n nn.Module or None: the normalization layer\n \"\"\"\n if isinstance(norm, str):\n if len(norm) == 0:\n return None\n norm = {\n \"BN\": BatchNorm,\n # \"syncBN\": SyncBatchNorm,\n # \"GhostBN\": GhostBatchNorm,\n # \"FrozenBN\": FrozenBatchNorm,\n \"GN\": lambda channels, **args: nn.GroupNorm(32, channels),\n }[norm]\n return norm(out_channels, **kwargs)\n\n\nclass Non_local(nn.Module):\n def __init__(self, in_channels, bn_norm, reduc_ratio=2):\n super(Non_local, self).__init__()\n\n self.in_channels = in_channels\n self.inter_channels = reduc_ratio // reduc_ratio\n\n self.g = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n self.W = nn.Sequential(\n nn.Conv2d(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0),\n get_norm(bn_norm, self.in_channels),\n )\n nn.init.constant_(self.W[1].weight, 0.0)\n nn.init.constant_(self.W[1].bias, 0.0)\n\n self.theta = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n self.phi = nn.Conv2d(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n def forward(self, x):\n \"\"\"\n :param x: (b, t, h, w)\n :return x: (b, t, h, w)\n \"\"\"\n batch_size = x.size(0)\n g_x = self.g(x).view(batch_size, self.inter_channels, -1)\n g_x = g_x.permute(0, 2, 1)\n\n theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)\n theta_x = theta_x.permute(0, 2, 1)\n phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)\n f = torch.matmul(theta_x, phi_x)\n N = f.size(-1)\n f_div_C = f / N\n\n y = torch.matmul(f_div_C, g_x)\n y = y.permute(0, 2, 1).contiguous()\n y = y.view(batch_size, self.inter_channels, *x.size()[2:])\n W_y = self.W(y)\n z = W_y + x\n return z\n\n\nclass IBN(nn.Module):\n def __init__(self, planes, bn_norm, **kwargs):\n super(IBN, self).__init__()\n half1 = int(planes / 2)\n self.half = half1\n half2 = planes - half1\n self.IN = nn.InstanceNorm2d(half1, affine=True)\n self.BN = get_norm(bn_norm, half2, **kwargs)\n\n def forward(self, x):\n split = torch.split(x, self.half, 1)\n out1 = self.IN(split[0].contiguous())\n out2 = self.BN(split[1].contiguous())\n out = torch.cat((out1, out2), 1)\n return out\n\n\nclass BatchNorm(nn.BatchNorm2d):\n def __init__(self, num_features, eps=1e-05, momentum=0.1, weight_freeze=False, bias_freeze=False, weight_init=1.0,\n bias_init=0.0, **kwargs):\n super().__init__(num_features, eps=eps, momentum=momentum)\n if weight_init is not None: nn.init.constant_(self.weight, weight_init)\n if bias_init is not None: nn.init.constant_(self.bias, bias_init)\n self.weight.requires_grad_(not weight_freeze)\n self.bias.requires_grad_(not bias_freeze)\n\n\nclass SELayer(nn.Module):\n def __init__(self, channel, reduction=16):\n super(SELayer, self).__init__()\n self.avg_pool = nn.AdaptiveAvgPool2d(1)\n self.fc = nn.Sequential(\n nn.Linear(channel, int(channel / reduction), bias=False),\n nn.ReLU(inplace=True),\n nn.Linear(int(channel / reduction), channel, bias=False),\n nn.Sigmoid()\n )\n\n def forward(self, x):\n b, c, _, _ = x.size()\n y = self.avg_pool(x).view(b, c)\n y = self.fc(y).view(b, c, 1, 1)\n return x * y.expand_as(x)\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, bn_norm, with_ibn=False, with_se=False,\n stride=1, downsample=None, reduction=16):\n super(BasicBlock, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n if with_ibn:\n self.bn1 = IBN(planes, bn_norm)\n else:\n self.bn1 = get_norm(bn_norm, planes)\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)\n self.bn2 = get_norm(bn_norm, planes)\n self.relu = nn.ReLU(inplace=True)\n if with_se:\n self.se = SELayer(planes, reduction)\n else:\n self.se = nn.Identity()\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.se(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, bn_norm, with_ibn=False, with_se=False,\n stride=1, downsample=None, reduction=16):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n if with_ibn:\n self.bn1 = IBN(planes, bn_norm)\n else:\n self.bn1 = get_norm(bn_norm, planes)\n self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,\n padding=1, bias=False)\n self.bn2 = get_norm(bn_norm, planes)\n self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)\n self.bn3 = get_norm(bn_norm, planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n if with_se:\n self.se = SELayer(planes * self.expansion, reduction)\n else:\n self.se = nn.Identity()\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n out = self.se(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n def __init__(self, last_stride, bn_norm, with_ibn, with_se, with_nl, block, layers, non_layers):\n self.inplanes = 64\n super().__init__()\n self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,\n bias=False)\n self.bn1 = get_norm(bn_norm, 64)\n self.relu = nn.ReLU(inplace=True)\n # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)\n self.layer1 = self._make_layer(block, 64, layers[0], 1, bn_norm, with_ibn, with_se)\n self.layer2 = self._make_layer(block, 128, layers[1], 2, bn_norm, with_ibn, with_se)\n self.layer3 = self._make_layer(block, 256, layers[2], 2, bn_norm, with_ibn, with_se)\n self.layer4 = self._make_layer(block, 512, layers[3], last_stride, bn_norm, with_se=with_se)\n\n self.random_init()\n\n # fmt: off\n if with_nl: self._build_nonlocal(layers, non_layers, bn_norm)\n else: self.NL_1_idx = self.NL_2_idx = self.NL_3_idx = self.NL_4_idx = []\n # fmt: on\n self.feature_dim = 2048\n\n def _make_layer(self, block, planes, blocks, stride=1, bn_norm=\"BN\", with_ibn=False, with_se=False):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(self.inplanes, planes * block.expansion,\n kernel_size=1, stride=stride, bias=False),\n get_norm(bn_norm, planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, bn_norm, with_ibn, with_se, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes, bn_norm, with_ibn, with_se))\n\n return nn.Sequential(*layers)\n\n def _build_nonlocal(self, layers, non_layers, bn_norm):\n self.NL_1 = nn.ModuleList(\n [Non_local(256, bn_norm) for _ in range(non_layers[0])])\n self.NL_1_idx = sorted([layers[0] - (i + 1) for i in range(non_layers[0])])\n self.NL_2 = nn.ModuleList(\n [Non_local(512, bn_norm) for _ in range(non_layers[1])])\n self.NL_2_idx = sorted([layers[1] - (i + 1) for i in range(non_layers[1])])\n self.NL_3 = nn.ModuleList(\n [Non_local(1024, bn_norm) for _ in range(non_layers[2])])\n self.NL_3_idx = sorted([layers[2] - (i + 1) for i in range(non_layers[2])])\n self.NL_4 = nn.ModuleList(\n [Non_local(2048, bn_norm) for _ in range(non_layers[3])])\n self.NL_4_idx = sorted([layers[3] - (i + 1) for i in range(non_layers[3])])\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n\n # layer 1\n NL1_counter = 0\n if len(self.NL_1_idx) == 0:\n self.NL_1_idx = [-1]\n for i in range(len(self.layer1)):\n x = self.layer1[i](x)\n if i == self.NL_1_idx[NL1_counter]:\n _, C, H, W = x.shape\n x = self.NL_1[NL1_counter](x)\n NL1_counter += 1\n # layer 2\n NL2_counter = 0\n if len(self.NL_2_idx) == 0:\n self.NL_2_idx = [-1]\n for i in range(len(self.layer2)):\n x = self.layer2[i](x)\n if i == self.NL_2_idx[NL2_counter]:\n _, C, H, W = x.shape\n x = self.NL_2[NL2_counter](x)\n NL2_counter += 1\n\n # layer 3\n NL3_counter = 0\n if len(self.NL_3_idx) == 0:\n self.NL_3_idx = [-1]\n for i in range(len(self.layer3)):\n x = self.layer3[i](x)\n if i == self.NL_3_idx[NL3_counter]:\n _, C, H, W = x.shape\n x = self.NL_3[NL3_counter](x)\n NL3_counter += 1\n\n # layer 4\n NL4_counter = 0\n if len(self.NL_4_idx) == 0:\n self.NL_4_idx = [-1]\n for i in range(len(self.layer4)):\n x = self.layer4[i](x)\n if i == self.NL_4_idx[NL4_counter]:\n _, C, H, W = x.shape\n x = self.NL_4[NL4_counter](x)\n NL4_counter += 1\n\n return x\n\n def random_init(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n nn.init.normal_(m.weight, 0, math.sqrt(2. / n))\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n\ndef init_pretrained_weights(key):\n \"\"\"Initializes model with pretrained weights.\n\n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n import os\n import errno\n import gdown\n\n def _get_torch_home():\n ENV_TORCH_HOME = 'TORCH_HOME'\n ENV_XDG_CACHE_HOME = 'XDG_CACHE_HOME'\n DEFAULT_CACHE_DIR = '~/.cache'\n torch_home = os.path.expanduser(\n os.getenv(\n ENV_TORCH_HOME,\n os.path.join(\n os.getenv(ENV_XDG_CACHE_HOME, DEFAULT_CACHE_DIR), 'torch'\n )\n )\n )\n return torch_home\n\n torch_home = _get_torch_home()\n model_dir = os.path.join(torch_home, 'checkpoints')\n try:\n os.makedirs(model_dir)\n except OSError as e:\n if e.errno == errno.EEXIST:\n # Directory already exists, ignore.\n pass\n else:\n # Unexpected OSError, re-raise.\n raise\n\n filename = model_urls[key].split('/')[-1]\n\n cached_file = os.path.join(model_dir, filename)\n\n if not os.path.exists(cached_file):\n print(f\"Pretrain model don't exist, downloading from {model_urls[key]}\")\n gdown.download(model_urls[key], cached_file, quiet=False)\n\n print(f\"Loading pretrained model from {cached_file}\")\n state_dict = torch.load(cached_file, map_location=torch.device('cpu'))\n\n return state_dict\n\ndef fastreid_resnet(pretrained=True, **kwargs):\n return build_resnet_backbone(pretrained)\n\ndef fastreid_resnet_ibn(pretrained=True, **kwargs):\n return build_resnet_backbone(pretrained, with_ibn=True)\n\ndef fastreid_resnet_nl(pretrained=True, **kwargs):\n return build_resnet_backbone(pretrained, with_nl=True)\n\ndef fastreid_resnet_ibn_nl(pretrained=True, **kwargs):\n return build_resnet_backbone(pretrained, with_ibn=True, with_nl=True)\n\ndef build_resnet_backbone(pretrained=True, with_ibn=False, with_nl=False, last_stride=1, **kwargs):\n \"\"\"\n Create a ResNet instance from config.\n Returns:\n ResNet: a :class:`ResNet` instance.\n \"\"\"\n\n # fmt: off\n pretrain = pretrained\n pretrain_path = False\n bn_norm = \"BN\"\n # with_ibn = cfg.MODEL.BACKBONE.WITH_IBN\n with_se = False\n # with_nl = cfg.MODEL.BACKBONE.WITH_NL\n depth = \"50x\"\n # fmt: on\n\n num_blocks_per_stage = {\n '18x': [2, 2, 2, 2],\n '34x': [3, 4, 6, 3],\n '50x': [3, 4, 6, 3],\n '101x': [3, 4, 23, 3],\n }[depth]\n\n nl_layers_per_stage = {\n '18x': [0, 0, 0, 0],\n '34x': [0, 0, 0, 0],\n '50x': [0, 2, 3, 0],\n '101x': [0, 2, 9, 0]\n }[depth]\n\n block = {\n '18x': BasicBlock,\n '34x': BasicBlock,\n '50x': Bottleneck,\n '101x': Bottleneck\n }[depth]\n\n model = ResNet(last_stride, bn_norm, with_ibn, with_se, with_nl, block,\n num_blocks_per_stage, nl_layers_per_stage)\n if pretrain:\n # Load pretrain path if specifically\n if pretrain_path:\n try:\n state_dict = torch.load(pretrain_path, map_location=torch.device('cpu'))\n print(f\"Loading pretrained model from {pretrain_path}\")\n except FileNotFoundError as e:\n print(f'{pretrain_path} is not found! Please check this path.')\n raise e\n except KeyError as e:\n print(\"State dict keys error! Please check the state dict.\")\n raise e\n else:\n key = depth\n if with_ibn: key = 'ibn_' + key\n if with_se: key = 'se_' + key\n\n state_dict = init_pretrained_weights(key)\n\n incompatible = model.load_state_dict(state_dict, strict=False)\n if incompatible.missing_keys:\n print(\n \"incompatible.missing_keys\".format(incompatible.missing_keys)\n )\n if incompatible.unexpected_keys:\n print(\n \"incompatible.unexpected_keys\".format(incompatible.unexpected_keys)\n )\n\n return model\n" }, { "path": "torchreid/models/resnet_ibn_a.py", "content": "\"\"\"\nCredit to https://github.com/XingangPan/IBN-Net.\n\"\"\"\nfrom __future__ import division, absolute_import\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['resnet50_ibn_a']\n\nmodel_urls = {\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"3x3 convolution with padding\"\n return nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm2d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = nn.BatchNorm2d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass IBN(nn.Module):\n\n def __init__(self, planes):\n super(IBN, self).__init__()\n half1 = int(planes / 2)\n self.half = half1\n half2 = planes - half1\n self.IN = nn.InstanceNorm2d(half1, affine=True)\n self.BN = nn.BatchNorm2d(half2)\n\n def forward(self, x):\n split = torch.split(x, self.half, 1)\n out1 = self.IN(split[0].contiguous())\n out2 = self.BN(split[1].contiguous())\n out = torch.cat((out1, out2), 1)\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, ibn=False, stride=1, downsample=None):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n if ibn:\n self.bn1 = IBN(planes)\n else:\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(\n planes,\n planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(\n planes, planes * self.expansion, kernel_size=1, bias=False\n )\n self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n \"\"\"Residual network + IBN layer.\n \n Reference:\n - He et al. Deep Residual Learning for Image Recognition. CVPR 2016.\n - Pan et al. Two at Once: Enhancing Learning and Generalization\n Capacities via IBN-Net. ECCV 2018.\n \"\"\"\n\n def __init__(\n self,\n block,\n layers,\n num_classes=1000,\n loss='softmax',\n fc_dims=None,\n dropout_p=None,\n **kwargs\n ):\n scale = 64\n self.inplanes = scale\n super(ResNet, self).__init__()\n self.loss = loss\n self.feature_dim = scale * 8 * block.expansion\n\n self.conv1 = nn.Conv2d(\n 3, scale, kernel_size=7, stride=2, padding=3, bias=False\n )\n self.bn1 = nn.BatchNorm2d(scale)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, scale, layers[0])\n self.layer2 = self._make_layer(block, scale * 2, layers[1], stride=2)\n self.layer3 = self._make_layer(block, scale * 4, layers[2], stride=2)\n self.layer4 = self._make_layer(block, scale * 8, layers[3], stride=2)\n self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n self.fc = self._construct_fc_layer(\n fc_dims, scale * 8 * block.expansion, dropout_p\n )\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n m.weight.data.normal_(0, math.sqrt(2. / n))\n elif isinstance(m, nn.BatchNorm2d):\n m.weight.data.fill_(1)\n m.bias.data.zero_()\n elif isinstance(m, nn.InstanceNorm2d):\n m.weight.data.fill_(1)\n m.bias.data.zero_()\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(\n self.inplanes,\n planes * block.expansion,\n kernel_size=1,\n stride=stride,\n bias=False\n ),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n ibn = True\n if planes == 512:\n ibn = False\n layers.append(block(self.inplanes, planes, ibn, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes, ibn))\n\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.avgpool(f)\n v = v.view(v.size(0), -1)\n if self.fc is not None:\n v = self.fc(v)\n if not self.training:\n return v\n y = self.classifier(v)\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef resnet50_ibn_a(num_classes, loss='softmax', pretrained=False, **kwargs):\n model = ResNet(\n Bottleneck, [3, 4, 6, 3], num_classes=num_classes, loss=loss, **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n" }, { "path": "torchreid/models/resnet_ibn_b.py", "content": "\"\"\"\nCredit to https://github.com/XingangPan/IBN-Net.\n\"\"\"\nfrom __future__ import division, absolute_import\nimport math\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['resnet50_ibn_b']\n\nmodel_urls = {\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"3x3 convolution with padding\"\n return nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm2d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = nn.BatchNorm2d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None, IN=False):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(\n planes,\n planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(\n planes, planes * self.expansion, kernel_size=1, bias=False\n )\n self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n self.IN = None\n if IN:\n self.IN = nn.InstanceNorm2d(planes * 4, affine=True)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n if self.IN is not None:\n out = self.IN(out)\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n \"\"\"Residual network + IBN layer.\n \n Reference:\n - He et al. Deep Residual Learning for Image Recognition. CVPR 2016.\n - Pan et al. Two at Once: Enhancing Learning and Generalization\n Capacities via IBN-Net. ECCV 2018.\n \"\"\"\n\n def __init__(\n self,\n block,\n layers,\n num_classes=1000,\n loss='softmax',\n fc_dims=None,\n dropout_p=None,\n **kwargs\n ):\n scale = 64\n self.inplanes = scale\n super(ResNet, self).__init__()\n self.loss = loss\n self.feature_dim = scale * 8 * block.expansion\n\n self.conv1 = nn.Conv2d(\n 3, scale, kernel_size=7, stride=2, padding=3, bias=False\n )\n self.bn1 = nn.InstanceNorm2d(scale, affine=True)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(\n block, scale, layers[0], stride=1, IN=True\n )\n self.layer2 = self._make_layer(\n block, scale * 2, layers[1], stride=2, IN=True\n )\n self.layer3 = self._make_layer(block, scale * 4, layers[2], stride=2)\n self.layer4 = self._make_layer(block, scale * 8, layers[3], stride=2)\n self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n self.fc = self._construct_fc_layer(\n fc_dims, scale * 8 * block.expansion, dropout_p\n )\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n m.weight.data.normal_(0, math.sqrt(2. / n))\n elif isinstance(m, nn.BatchNorm2d):\n m.weight.data.fill_(1)\n m.bias.data.zero_()\n elif isinstance(m, nn.InstanceNorm2d):\n m.weight.data.fill_(1)\n m.bias.data.zero_()\n\n def _make_layer(self, block, planes, blocks, stride=1, IN=False):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(\n self.inplanes,\n planes * block.expansion,\n kernel_size=1,\n stride=stride,\n bias=False\n ),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks - 1):\n layers.append(block(self.inplanes, planes))\n layers.append(block(self.inplanes, planes, IN=IN))\n\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.avgpool(f)\n v = v.view(v.size(0), -1)\n if self.fc is not None:\n v = self.fc(v)\n if not self.training:\n return v\n y = self.classifier(v)\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef resnet50_ibn_b(num_classes, loss='softmax', pretrained=False, **kwargs):\n model = ResNet(\n Bottleneck, [3, 4, 6, 3], num_classes=num_classes, loss=loss, **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n" }, { "path": "torchreid/models/resnetmid.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\n\n__all__ = ['resnet50mid']\n\nmodel_urls = {\n 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n}\n\n\ndef conv3x3(in_planes, out_planes, stride=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv2d(\n in_planes,\n out_planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x3(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm2d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x3(planes, planes)\n self.bn2 = nn.BatchNorm2d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(Bottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(\n planes,\n planes,\n kernel_size=3,\n stride=stride,\n padding=1,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(\n planes, planes * self.expansion, kernel_size=1, bias=False\n )\n self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out += residual\n out = self.relu(out)\n\n return out\n\n\nclass ResNetMid(nn.Module):\n \"\"\"Residual network + mid-level features.\n \n Reference:\n Yu et al. The Devil is in the Middle: Exploiting Mid-level Representations for\n Cross-Domain Instance Matching. arXiv:1711.08106.\n\n Public keys:\n - ``resnet50mid``: ResNet50 + mid-level feature fusion.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n block,\n layers,\n last_stride=2,\n fc_dims=None,\n **kwargs\n ):\n self.inplanes = 64\n super(ResNetMid, self).__init__()\n self.loss = loss\n self.feature_dim = 512 * block.expansion\n\n # backbone network\n self.conv1 = nn.Conv2d(\n 3, 64, kernel_size=7, stride=2, padding=3, bias=False\n )\n self.bn1 = nn.BatchNorm2d(64)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(\n block, 512, layers[3], stride=last_stride\n )\n\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n assert fc_dims is not None\n self.fc_fusion = self._construct_fc_layer(\n fc_dims, 512 * block.expansion * 2\n )\n self.feature_dim += 512 * block.expansion\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(\n self.inplanes,\n planes * block.expansion,\n kernel_size=1,\n stride=stride,\n bias=False\n ),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample))\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes))\n\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x4a = self.layer4[0](x)\n x4b = self.layer4[1](x4a)\n x4c = self.layer4[2](x4b)\n return x4a, x4b, x4c\n\n def forward(self, x):\n x4a, x4b, x4c = self.featuremaps(x)\n\n v4a = self.global_avgpool(x4a)\n v4b = self.global_avgpool(x4b)\n v4c = self.global_avgpool(x4c)\n v4ab = torch.cat([v4a, v4b], 1)\n v4ab = v4ab.view(v4ab.size(0), -1)\n v4ab = self.fc_fusion(v4ab)\n v4c = v4c.view(v4c.size(0), -1)\n v = torch.cat([v4ab, v4c], 1)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\n\"\"\"\nResidual network configurations:\n--\nresnet18: block=BasicBlock, layers=[2, 2, 2, 2]\nresnet34: block=BasicBlock, layers=[3, 4, 6, 3]\nresnet50: block=Bottleneck, layers=[3, 4, 6, 3]\nresnet101: block=Bottleneck, layers=[3, 4, 23, 3]\nresnet152: block=Bottleneck, layers=[3, 8, 36, 3]\n\"\"\"\n\n\ndef resnet50mid(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ResNetMid(\n num_classes=num_classes,\n loss=loss,\n block=Bottleneck,\n layers=[3, 4, 6, 3],\n last_stride=2,\n fc_dims=[1024],\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['resnet50'])\n return model\n" }, { "path": "torchreid/models/senet.py", "content": "from __future__ import division, absolute_import\nimport math\nfrom collections import OrderedDict\nimport torch.nn as nn\nfrom torch.utils import model_zoo\n\n__all__ = [\n 'senet154', 'se_resnet50', 'se_resnet101', 'se_resnet152',\n 'se_resnext50_32x4d', 'se_resnext101_32x4d', 'se_resnet50_fc512'\n]\n\"\"\"\nCode imported from https://github.com/Cadene/pretrained-models.pytorch\n\"\"\"\n\npretrained_settings = {\n 'senet154': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/senet154-c7b49a05.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224],\n 'input_range': [0, 1],\n 'mean': [0.485, 0.456, 0.406],\n 'std': [0.229, 0.224, 0.225],\n 'num_classes': 1000\n }\n },\n 'se_resnet50': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet50-ce0d4300.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224],\n 'input_range': [0, 1],\n 'mean': [0.485, 0.456, 0.406],\n 'std': [0.229, 0.224, 0.225],\n 'num_classes': 1000\n }\n },\n 'se_resnet101': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet101-7e38fcc6.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224],\n 'input_range': [0, 1],\n 'mean': [0.485, 0.456, 0.406],\n 'std': [0.229, 0.224, 0.225],\n 'num_classes': 1000\n }\n },\n 'se_resnet152': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnet152-d17c99b7.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224],\n 'input_range': [0, 1],\n 'mean': [0.485, 0.456, 0.406],\n 'std': [0.229, 0.224, 0.225],\n 'num_classes': 1000\n }\n },\n 'se_resnext50_32x4d': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnext50_32x4d-a260b3a4.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224],\n 'input_range': [0, 1],\n 'mean': [0.485, 0.456, 0.406],\n 'std': [0.229, 0.224, 0.225],\n 'num_classes': 1000\n }\n },\n 'se_resnext101_32x4d': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/se_resnext101_32x4d-3b2fe3d8.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 224, 224],\n 'input_range': [0, 1],\n 'mean': [0.485, 0.456, 0.406],\n 'std': [0.229, 0.224, 0.225],\n 'num_classes': 1000\n }\n },\n}\n\n\nclass SEModule(nn.Module):\n\n def __init__(self, channels, reduction):\n super(SEModule, self).__init__()\n self.avg_pool = nn.AdaptiveAvgPool2d(1)\n self.fc1 = nn.Conv2d(\n channels, channels // reduction, kernel_size=1, padding=0\n )\n self.relu = nn.ReLU(inplace=True)\n self.fc2 = nn.Conv2d(\n channels // reduction, channels, kernel_size=1, padding=0\n )\n self.sigmoid = nn.Sigmoid()\n\n def forward(self, x):\n module_input = x\n x = self.avg_pool(x)\n x = self.fc1(x)\n x = self.relu(x)\n x = self.fc2(x)\n x = self.sigmoid(x)\n return module_input * x\n\n\nclass Bottleneck(nn.Module):\n \"\"\"\n Base class for bottlenecks that implements `forward()` method.\n \"\"\"\n\n def forward(self, x):\n residual = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n residual = self.downsample(x)\n\n out = self.se_module(out) + residual\n out = self.relu(out)\n\n return out\n\n\nclass SEBottleneck(Bottleneck):\n \"\"\"\n Bottleneck for SENet154.\n \"\"\"\n expansion = 4\n\n def __init__(\n self, inplanes, planes, groups, reduction, stride=1, downsample=None\n ):\n super(SEBottleneck, self).__init__()\n self.conv1 = nn.Conv2d(inplanes, planes * 2, kernel_size=1, bias=False)\n self.bn1 = nn.BatchNorm2d(planes * 2)\n self.conv2 = nn.Conv2d(\n planes * 2,\n planes * 4,\n kernel_size=3,\n stride=stride,\n padding=1,\n groups=groups,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(planes * 4)\n self.conv3 = nn.Conv2d(\n planes * 4, planes * 4, kernel_size=1, bias=False\n )\n self.bn3 = nn.BatchNorm2d(planes * 4)\n self.relu = nn.ReLU(inplace=True)\n self.se_module = SEModule(planes * 4, reduction=reduction)\n self.downsample = downsample\n self.stride = stride\n\n\nclass SEResNetBottleneck(Bottleneck):\n \"\"\"\n ResNet bottleneck with a Squeeze-and-Excitation module. It follows Caffe\n implementation and uses `stride=stride` in `conv1` and not in `conv2`\n (the latter is used in the torchvision implementation of ResNet).\n \"\"\"\n expansion = 4\n\n def __init__(\n self, inplanes, planes, groups, reduction, stride=1, downsample=None\n ):\n super(SEResNetBottleneck, self).__init__()\n self.conv1 = nn.Conv2d(\n inplanes, planes, kernel_size=1, bias=False, stride=stride\n )\n self.bn1 = nn.BatchNorm2d(planes)\n self.conv2 = nn.Conv2d(\n planes,\n planes,\n kernel_size=3,\n padding=1,\n groups=groups,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(planes)\n self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)\n self.bn3 = nn.BatchNorm2d(planes * 4)\n self.relu = nn.ReLU(inplace=True)\n self.se_module = SEModule(planes * 4, reduction=reduction)\n self.downsample = downsample\n self.stride = stride\n\n\nclass SEResNeXtBottleneck(Bottleneck):\n \"\"\"ResNeXt bottleneck type C with a Squeeze-and-Excitation module\"\"\"\n expansion = 4\n\n def __init__(\n self,\n inplanes,\n planes,\n groups,\n reduction,\n stride=1,\n downsample=None,\n base_width=4\n ):\n super(SEResNeXtBottleneck, self).__init__()\n width = int(math.floor(planes * (base_width/64.)) * groups)\n self.conv1 = nn.Conv2d(\n inplanes, width, kernel_size=1, bias=False, stride=1\n )\n self.bn1 = nn.BatchNorm2d(width)\n self.conv2 = nn.Conv2d(\n width,\n width,\n kernel_size=3,\n stride=stride,\n padding=1,\n groups=groups,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(width)\n self.conv3 = nn.Conv2d(width, planes * 4, kernel_size=1, bias=False)\n self.bn3 = nn.BatchNorm2d(planes * 4)\n self.relu = nn.ReLU(inplace=True)\n self.se_module = SEModule(planes * 4, reduction=reduction)\n self.downsample = downsample\n self.stride = stride\n\n\nclass SENet(nn.Module):\n \"\"\"Squeeze-and-excitation network.\n \n Reference:\n Hu et al. Squeeze-and-Excitation Networks. CVPR 2018.\n\n Public keys:\n - ``senet154``: SENet154.\n - ``se_resnet50``: ResNet50 + SE.\n - ``se_resnet101``: ResNet101 + SE.\n - ``se_resnet152``: ResNet152 + SE.\n - ``se_resnext50_32x4d``: ResNeXt50 (groups=32, width=4) + SE.\n - ``se_resnext101_32x4d``: ResNeXt101 (groups=32, width=4) + SE.\n - ``se_resnet50_fc512``: (ResNet50 + SE) + FC.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n block,\n layers,\n groups,\n reduction,\n dropout_p=0.2,\n inplanes=128,\n input_3x3=True,\n downsample_kernel_size=3,\n downsample_padding=1,\n last_stride=2,\n fc_dims=None,\n **kwargs\n ):\n \"\"\"\n Parameters\n ----------\n block (nn.Module): Bottleneck class.\n - For SENet154: SEBottleneck\n - For SE-ResNet models: SEResNetBottleneck\n - For SE-ResNeXt models: SEResNeXtBottleneck\n layers (list of ints): Number of residual blocks for 4 layers of the\n network (layer1...layer4).\n groups (int): Number of groups for the 3x3 convolution in each\n bottleneck block.\n - For SENet154: 64\n - For SE-ResNet models: 1\n - For SE-ResNeXt models: 32\n reduction (int): Reduction ratio for Squeeze-and-Excitation modules.\n - For all models: 16\n dropout_p (float or None): Drop probability for the Dropout layer.\n If `None` the Dropout layer is not used.\n - For SENet154: 0.2\n - For SE-ResNet models: None\n - For SE-ResNeXt models: None\n inplanes (int): Number of input channels for layer1.\n - For SENet154: 128\n - For SE-ResNet models: 64\n - For SE-ResNeXt models: 64\n input_3x3 (bool): If `True`, use three 3x3 convolutions instead of\n a single 7x7 convolution in layer0.\n - For SENet154: True\n - For SE-ResNet models: False\n - For SE-ResNeXt models: False\n downsample_kernel_size (int): Kernel size for downsampling convolutions\n in layer2, layer3 and layer4.\n - For SENet154: 3\n - For SE-ResNet models: 1\n - For SE-ResNeXt models: 1\n downsample_padding (int): Padding for downsampling convolutions in\n layer2, layer3 and layer4.\n - For SENet154: 1\n - For SE-ResNet models: 0\n - For SE-ResNeXt models: 0\n num_classes (int): Number of outputs in `classifier` layer.\n \"\"\"\n super(SENet, self).__init__()\n self.inplanes = inplanes\n self.loss = loss\n\n if input_3x3:\n layer0_modules = [\n (\n 'conv1',\n nn.Conv2d(3, 64, 3, stride=2, padding=1, bias=False)\n ),\n ('bn1', nn.BatchNorm2d(64)),\n ('relu1', nn.ReLU(inplace=True)),\n (\n 'conv2',\n nn.Conv2d(64, 64, 3, stride=1, padding=1, bias=False)\n ),\n ('bn2', nn.BatchNorm2d(64)),\n ('relu2', nn.ReLU(inplace=True)),\n (\n 'conv3',\n nn.Conv2d(\n 64, inplanes, 3, stride=1, padding=1, bias=False\n )\n ),\n ('bn3', nn.BatchNorm2d(inplanes)),\n ('relu3', nn.ReLU(inplace=True)),\n ]\n else:\n layer0_modules = [\n (\n 'conv1',\n nn.Conv2d(\n 3,\n inplanes,\n kernel_size=7,\n stride=2,\n padding=3,\n bias=False\n )\n ),\n ('bn1', nn.BatchNorm2d(inplanes)),\n ('relu1', nn.ReLU(inplace=True)),\n ]\n # To preserve compatibility with Caffe weights `ceil_mode=True`\n # is used instead of `padding=1`.\n layer0_modules.append(\n ('pool', nn.MaxPool2d(3, stride=2, ceil_mode=True))\n )\n self.layer0 = nn.Sequential(OrderedDict(layer0_modules))\n self.layer1 = self._make_layer(\n block,\n planes=64,\n blocks=layers[0],\n groups=groups,\n reduction=reduction,\n downsample_kernel_size=1,\n downsample_padding=0\n )\n self.layer2 = self._make_layer(\n block,\n planes=128,\n blocks=layers[1],\n stride=2,\n groups=groups,\n reduction=reduction,\n downsample_kernel_size=downsample_kernel_size,\n downsample_padding=downsample_padding\n )\n self.layer3 = self._make_layer(\n block,\n planes=256,\n blocks=layers[2],\n stride=2,\n groups=groups,\n reduction=reduction,\n downsample_kernel_size=downsample_kernel_size,\n downsample_padding=downsample_padding\n )\n self.layer4 = self._make_layer(\n block,\n planes=512,\n blocks=layers[3],\n stride=last_stride,\n groups=groups,\n reduction=reduction,\n downsample_kernel_size=downsample_kernel_size,\n downsample_padding=downsample_padding\n )\n\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.fc = self._construct_fc_layer(\n fc_dims, 512 * block.expansion, dropout_p\n )\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n def _make_layer(\n self,\n block,\n planes,\n blocks,\n groups,\n reduction,\n stride=1,\n downsample_kernel_size=1,\n downsample_padding=0\n ):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n nn.Conv2d(\n self.inplanes,\n planes * block.expansion,\n kernel_size=downsample_kernel_size,\n stride=stride,\n padding=downsample_padding,\n bias=False\n ),\n nn.BatchNorm2d(planes * block.expansion),\n )\n\n layers = []\n layers.append(\n block(\n self.inplanes, planes, groups, reduction, stride, downsample\n )\n )\n self.inplanes = planes * block.expansion\n for i in range(1, blocks):\n layers.append(block(self.inplanes, planes, groups, reduction))\n\n return nn.Sequential(*layers)\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"\n Construct fully connected layer\n\n - fc_dims (list or tuple): dimensions of fc layers, if None,\n no fc layers are constructed\n - input_dim (int): input dimension\n - dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def featuremaps(self, x):\n x = self.layer0(x)\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if self.fc is not None:\n v = self.fc(v)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef senet154(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEBottleneck,\n layers=[3, 8, 36, 3],\n groups=64,\n reduction=16,\n dropout_p=0.2,\n last_stride=2,\n fc_dims=None,\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['senet154']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n\n\ndef se_resnet50(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEResNetBottleneck,\n layers=[3, 4, 6, 3],\n groups=1,\n reduction=16,\n dropout_p=None,\n inplanes=64,\n input_3x3=False,\n downsample_kernel_size=1,\n downsample_padding=0,\n last_stride=2,\n fc_dims=None,\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['se_resnet50']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n\n\ndef se_resnet50_fc512(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEResNetBottleneck,\n layers=[3, 4, 6, 3],\n groups=1,\n reduction=16,\n dropout_p=None,\n inplanes=64,\n input_3x3=False,\n downsample_kernel_size=1,\n downsample_padding=0,\n last_stride=1,\n fc_dims=[512],\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['se_resnet50']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n\n\ndef se_resnet101(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEResNetBottleneck,\n layers=[3, 4, 23, 3],\n groups=1,\n reduction=16,\n dropout_p=None,\n inplanes=64,\n input_3x3=False,\n downsample_kernel_size=1,\n downsample_padding=0,\n last_stride=2,\n fc_dims=None,\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['se_resnet101']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n\n\ndef se_resnet152(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEResNetBottleneck,\n layers=[3, 8, 36, 3],\n groups=1,\n reduction=16,\n dropout_p=None,\n inplanes=64,\n input_3x3=False,\n downsample_kernel_size=1,\n downsample_padding=0,\n last_stride=2,\n fc_dims=None,\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['se_resnet152']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n\n\ndef se_resnext50_32x4d(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEResNeXtBottleneck,\n layers=[3, 4, 6, 3],\n groups=32,\n reduction=16,\n dropout_p=None,\n inplanes=64,\n input_3x3=False,\n downsample_kernel_size=1,\n downsample_padding=0,\n last_stride=2,\n fc_dims=None,\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['se_resnext50_32x4d']['imagenet']['url'\n ]\n init_pretrained_weights(model, model_url)\n return model\n\n\ndef se_resnext101_32x4d(\n num_classes, loss='softmax', pretrained=True, **kwargs\n):\n model = SENet(\n num_classes=num_classes,\n loss=loss,\n block=SEResNeXtBottleneck,\n layers=[3, 4, 23, 3],\n groups=32,\n reduction=16,\n dropout_p=None,\n inplanes=64,\n input_3x3=False,\n downsample_kernel_size=1,\n downsample_padding=0,\n last_stride=2,\n fc_dims=None,\n **kwargs\n )\n if pretrained:\n model_url = pretrained_settings['se_resnext101_32x4d']['imagenet'][\n 'url']\n init_pretrained_weights(model, model_url)\n return model\n" }, { "path": "torchreid/models/shufflenet.py", "content": "from __future__ import division, absolute_import\nimport torch\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\nfrom torch.nn import functional as F\n\n__all__ = ['shufflenet']\n\nmodel_urls = {\n # training epoch = 90, top1 = 61.8\n 'imagenet':\n 'https://mega.nz/#!RDpUlQCY!tr_5xBEkelzDjveIYBBcGcovNCOrgfiJO9kiidz9fZM',\n}\n\n\nclass ChannelShuffle(nn.Module):\n\n def __init__(self, num_groups):\n super(ChannelShuffle, self).__init__()\n self.g = num_groups\n\n def forward(self, x):\n b, c, h, w = x.size()\n n = c // self.g\n # reshape\n x = x.view(b, self.g, n, h, w)\n # transpose\n x = x.permute(0, 2, 1, 3, 4).contiguous()\n # flatten\n x = x.view(b, c, h, w)\n return x\n\n\nclass Bottleneck(nn.Module):\n\n def __init__(\n self,\n in_channels,\n out_channels,\n stride,\n num_groups,\n group_conv1x1=True\n ):\n super(Bottleneck, self).__init__()\n assert stride in [1, 2], 'Warning: stride must be either 1 or 2'\n self.stride = stride\n mid_channels = out_channels // 4\n if stride == 2:\n out_channels -= in_channels\n # group conv is not applied to first conv1x1 at stage 2\n num_groups_conv1x1 = num_groups if group_conv1x1 else 1\n self.conv1 = nn.Conv2d(\n in_channels,\n mid_channels,\n 1,\n groups=num_groups_conv1x1,\n bias=False\n )\n self.bn1 = nn.BatchNorm2d(mid_channels)\n self.shuffle1 = ChannelShuffle(num_groups)\n self.conv2 = nn.Conv2d(\n mid_channels,\n mid_channels,\n 3,\n stride=stride,\n padding=1,\n groups=mid_channels,\n bias=False\n )\n self.bn2 = nn.BatchNorm2d(mid_channels)\n self.conv3 = nn.Conv2d(\n mid_channels, out_channels, 1, groups=num_groups, bias=False\n )\n self.bn3 = nn.BatchNorm2d(out_channels)\n if stride == 2:\n self.shortcut = nn.AvgPool2d(3, stride=2, padding=1)\n\n def forward(self, x):\n out = F.relu(self.bn1(self.conv1(x)))\n out = self.shuffle1(out)\n out = self.bn2(self.conv2(out))\n out = self.bn3(self.conv3(out))\n if self.stride == 2:\n res = self.shortcut(x)\n out = F.relu(torch.cat([res, out], 1))\n else:\n out = F.relu(x + out)\n return out\n\n\n# configuration of (num_groups: #out_channels) based on Table 1 in the paper\ncfg = {\n 1: [144, 288, 576],\n 2: [200, 400, 800],\n 3: [240, 480, 960],\n 4: [272, 544, 1088],\n 8: [384, 768, 1536],\n}\n\n\nclass ShuffleNet(nn.Module):\n \"\"\"ShuffleNet.\n\n Reference:\n Zhang et al. ShuffleNet: An Extremely Efficient Convolutional Neural\n Network for Mobile Devices. CVPR 2018.\n\n Public keys:\n - ``shufflenet``: ShuffleNet (groups=3).\n \"\"\"\n\n def __init__(self, num_classes, loss='softmax', num_groups=3, **kwargs):\n super(ShuffleNet, self).__init__()\n self.loss = loss\n\n self.conv1 = nn.Sequential(\n nn.Conv2d(3, 24, 3, stride=2, padding=1, bias=False),\n nn.BatchNorm2d(24),\n nn.ReLU(),\n nn.MaxPool2d(3, stride=2, padding=1),\n )\n\n self.stage2 = nn.Sequential(\n Bottleneck(\n 24, cfg[num_groups][0], 2, num_groups, group_conv1x1=False\n ),\n Bottleneck(cfg[num_groups][0], cfg[num_groups][0], 1, num_groups),\n Bottleneck(cfg[num_groups][0], cfg[num_groups][0], 1, num_groups),\n Bottleneck(cfg[num_groups][0], cfg[num_groups][0], 1, num_groups),\n )\n\n self.stage3 = nn.Sequential(\n Bottleneck(cfg[num_groups][0], cfg[num_groups][1], 2, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n Bottleneck(cfg[num_groups][1], cfg[num_groups][1], 1, num_groups),\n )\n\n self.stage4 = nn.Sequential(\n Bottleneck(cfg[num_groups][1], cfg[num_groups][2], 2, num_groups),\n Bottleneck(cfg[num_groups][2], cfg[num_groups][2], 1, num_groups),\n Bottleneck(cfg[num_groups][2], cfg[num_groups][2], 1, num_groups),\n Bottleneck(cfg[num_groups][2], cfg[num_groups][2], 1, num_groups),\n )\n\n self.classifier = nn.Linear(cfg[num_groups][2], num_classes)\n self.feat_dim = cfg[num_groups][2]\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.stage2(x)\n x = self.stage3(x)\n x = self.stage4(x)\n x = F.avg_pool2d(x, x.size()[2:]).view(x.size(0), -1)\n\n if not self.training:\n return x\n\n y = self.classifier(x)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, x\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef shufflenet(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ShuffleNet(num_classes, loss, **kwargs)\n if pretrained:\n # init_pretrained_weights(model, model_urls['imagenet'])\n import warnings\n warnings.warn(\n 'The imagenet pretrained weights need to be manually downloaded from {}'\n .format(model_urls['imagenet'])\n )\n return model\n" }, { "path": "torchreid/models/shufflenetv2.py", "content": "\"\"\"\nCode source: https://github.com/pytorch/vision\n\"\"\"\nfrom __future__ import division, absolute_import\nimport torch\nimport torch.utils.model_zoo as model_zoo\nfrom torch import nn\n\n__all__ = [\n 'shufflenet_v2_x0_5', 'shufflenet_v2_x1_0', 'shufflenet_v2_x1_5',\n 'shufflenet_v2_x2_0'\n]\n\nmodel_urls = {\n 'shufflenetv2_x0.5':\n 'https://download.pytorch.org/models/shufflenetv2_x0.5-f707e7126e.pth',\n 'shufflenetv2_x1.0':\n 'https://download.pytorch.org/models/shufflenetv2_x1-5666bf0f80.pth',\n 'shufflenetv2_x1.5': None,\n 'shufflenetv2_x2.0': None,\n}\n\n\ndef channel_shuffle(x, groups):\n batchsize, num_channels, height, width = x.data.size()\n channels_per_group = num_channels // groups\n\n # reshape\n x = x.view(batchsize, groups, channels_per_group, height, width)\n\n x = torch.transpose(x, 1, 2).contiguous()\n\n # flatten\n x = x.view(batchsize, -1, height, width)\n\n return x\n\n\nclass InvertedResidual(nn.Module):\n\n def __init__(self, inp, oup, stride):\n super(InvertedResidual, self).__init__()\n\n if not (1 <= stride <= 3):\n raise ValueError('illegal stride value')\n self.stride = stride\n\n branch_features = oup // 2\n assert (self.stride != 1) or (inp == branch_features << 1)\n\n if self.stride > 1:\n self.branch1 = nn.Sequential(\n self.depthwise_conv(\n inp, inp, kernel_size=3, stride=self.stride, padding=1\n ),\n nn.BatchNorm2d(inp),\n nn.Conv2d(\n inp,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False\n ),\n nn.BatchNorm2d(branch_features),\n nn.ReLU(inplace=True),\n )\n\n self.branch2 = nn.Sequential(\n nn.Conv2d(\n inp if (self.stride > 1) else branch_features,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False\n ),\n nn.BatchNorm2d(branch_features),\n nn.ReLU(inplace=True),\n self.depthwise_conv(\n branch_features,\n branch_features,\n kernel_size=3,\n stride=self.stride,\n padding=1\n ),\n nn.BatchNorm2d(branch_features),\n nn.Conv2d(\n branch_features,\n branch_features,\n kernel_size=1,\n stride=1,\n padding=0,\n bias=False\n ),\n nn.BatchNorm2d(branch_features),\n nn.ReLU(inplace=True),\n )\n\n @staticmethod\n def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):\n return nn.Conv2d(\n i, o, kernel_size, stride, padding, bias=bias, groups=i\n )\n\n def forward(self, x):\n if self.stride == 1:\n x1, x2 = x.chunk(2, dim=1)\n out = torch.cat((x1, self.branch2(x2)), dim=1)\n else:\n out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)\n\n out = channel_shuffle(out, 2)\n\n return out\n\n\nclass ShuffleNetV2(nn.Module):\n \"\"\"ShuffleNetV2.\n \n Reference:\n Ma et al. ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design. ECCV 2018.\n\n Public keys:\n - ``shufflenet_v2_x0_5``: ShuffleNetV2 x0.5.\n - ``shufflenet_v2_x1_0``: ShuffleNetV2 x1.0.\n - ``shufflenet_v2_x1_5``: ShuffleNetV2 x1.5.\n - ``shufflenet_v2_x2_0``: ShuffleNetV2 x2.0.\n \"\"\"\n\n def __init__(\n self, num_classes, loss, stages_repeats, stages_out_channels, **kwargs\n ):\n super(ShuffleNetV2, self).__init__()\n self.loss = loss\n\n if len(stages_repeats) != 3:\n raise ValueError(\n 'expected stages_repeats as list of 3 positive ints'\n )\n if len(stages_out_channels) != 5:\n raise ValueError(\n 'expected stages_out_channels as list of 5 positive ints'\n )\n self._stage_out_channels = stages_out_channels\n\n input_channels = 3\n output_channels = self._stage_out_channels[0]\n self.conv1 = nn.Sequential(\n nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False),\n nn.BatchNorm2d(output_channels),\n nn.ReLU(inplace=True),\n )\n input_channels = output_channels\n\n self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n\n stage_names = ['stage{}'.format(i) for i in [2, 3, 4]]\n for name, repeats, output_channels in zip(\n stage_names, stages_repeats, self._stage_out_channels[1:]\n ):\n seq = [InvertedResidual(input_channels, output_channels, 2)]\n for i in range(repeats - 1):\n seq.append(\n InvertedResidual(output_channels, output_channels, 1)\n )\n setattr(self, name, nn.Sequential(*seq))\n input_channels = output_channels\n\n output_channels = self._stage_out_channels[-1]\n self.conv5 = nn.Sequential(\n nn.Conv2d(input_channels, output_channels, 1, 1, 0, bias=False),\n nn.BatchNorm2d(output_channels),\n nn.ReLU(inplace=True),\n )\n self.global_avgpool = nn.AdaptiveAvgPool2d((1, 1))\n\n self.classifier = nn.Linear(output_channels, num_classes)\n\n def featuremaps(self, x):\n x = self.conv1(x)\n x = self.maxpool(x)\n x = self.stage2(x)\n x = self.stage3(x)\n x = self.stage4(x)\n x = self.conv5(x)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError(\"Unsupported loss: {}\".format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n if model_url is None:\n import warnings\n warnings.warn(\n 'ImageNet pretrained weights are unavailable for this model'\n )\n return\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef shufflenet_v2_x0_5(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ShuffleNetV2(\n num_classes, loss, [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['shufflenetv2_x0.5'])\n return model\n\n\ndef shufflenet_v2_x1_0(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ShuffleNetV2(\n num_classes, loss, [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['shufflenetv2_x1.0'])\n return model\n\n\ndef shufflenet_v2_x1_5(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ShuffleNetV2(\n num_classes, loss, [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['shufflenetv2_x1.5'])\n return model\n\n\ndef shufflenet_v2_x2_0(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = ShuffleNetV2(\n num_classes, loss, [4, 8, 4], [24, 244, 488, 976, 2048], **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['shufflenetv2_x2.0'])\n return model\n" }, { "path": "torchreid/models/squeezenet.py", "content": "\"\"\"\nCode source: https://github.com/pytorch/vision\n\"\"\"\nfrom __future__ import division, absolute_import\nimport torch\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['squeezenet1_0', 'squeezenet1_1', 'squeezenet1_0_fc512']\n\nmodel_urls = {\n 'squeezenet1_0':\n 'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth',\n 'squeezenet1_1':\n 'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth',\n}\n\n\nclass Fire(nn.Module):\n\n def __init__(\n self, inplanes, squeeze_planes, expand1x1_planes, expand3x3_planes\n ):\n super(Fire, self).__init__()\n self.inplanes = inplanes\n self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)\n self.squeeze_activation = nn.ReLU(inplace=True)\n self.expand1x1 = nn.Conv2d(\n squeeze_planes, expand1x1_planes, kernel_size=1\n )\n self.expand1x1_activation = nn.ReLU(inplace=True)\n self.expand3x3 = nn.Conv2d(\n squeeze_planes, expand3x3_planes, kernel_size=3, padding=1\n )\n self.expand3x3_activation = nn.ReLU(inplace=True)\n\n def forward(self, x):\n x = self.squeeze_activation(self.squeeze(x))\n return torch.cat(\n [\n self.expand1x1_activation(self.expand1x1(x)),\n self.expand3x3_activation(self.expand3x3(x))\n ], 1\n )\n\n\nclass SqueezeNet(nn.Module):\n \"\"\"SqueezeNet.\n\n Reference:\n Iandola et al. SqueezeNet: AlexNet-level accuracy with 50x fewer parameters\n and< 0.5 MB model size. arXiv:1602.07360.\n\n Public keys:\n - ``squeezenet1_0``: SqueezeNet (version=1.0).\n - ``squeezenet1_1``: SqueezeNet (version=1.1).\n - ``squeezenet1_0_fc512``: SqueezeNet (version=1.0) + FC.\n \"\"\"\n\n def __init__(\n self,\n num_classes,\n loss,\n version=1.0,\n fc_dims=None,\n dropout_p=None,\n **kwargs\n ):\n super(SqueezeNet, self).__init__()\n self.loss = loss\n self.feature_dim = 512\n\n if version not in [1.0, 1.1]:\n raise ValueError(\n 'Unsupported SqueezeNet version {version}:'\n '1.0 or 1.1 expected'.format(version=version)\n )\n\n if version == 1.0:\n self.features = nn.Sequential(\n nn.Conv2d(3, 96, kernel_size=7, stride=2),\n nn.ReLU(inplace=True),\n nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),\n Fire(96, 16, 64, 64),\n Fire(128, 16, 64, 64),\n Fire(128, 32, 128, 128),\n nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),\n Fire(256, 32, 128, 128),\n Fire(256, 48, 192, 192),\n Fire(384, 48, 192, 192),\n Fire(384, 64, 256, 256),\n nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),\n Fire(512, 64, 256, 256),\n )\n else:\n self.features = nn.Sequential(\n nn.Conv2d(3, 64, kernel_size=3, stride=2),\n nn.ReLU(inplace=True),\n nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),\n Fire(64, 16, 64, 64),\n Fire(128, 16, 64, 64),\n nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),\n Fire(128, 32, 128, 128),\n Fire(256, 32, 128, 128),\n nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),\n Fire(256, 48, 192, 192),\n Fire(384, 48, 192, 192),\n Fire(384, 64, 256, 256),\n Fire(512, 64, 256, 256),\n )\n\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.fc = self._construct_fc_layer(fc_dims, 512, dropout_p)\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def forward(self, x):\n f = self.features(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if self.fc is not None:\n v = self.fc(v)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initializes model with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url, map_location=None)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef squeezenet1_0(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SqueezeNet(\n num_classes, loss, version=1.0, fc_dims=None, dropout_p=None, **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['squeezenet1_0'])\n return model\n\n\ndef squeezenet1_0_fc512(\n num_classes, loss='softmax', pretrained=True, **kwargs\n):\n model = SqueezeNet(\n num_classes,\n loss,\n version=1.0,\n fc_dims=[512],\n dropout_p=None,\n **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['squeezenet1_0'])\n return model\n\n\ndef squeezenet1_1(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = SqueezeNet(\n num_classes, loss, version=1.1, fc_dims=None, dropout_p=None, **kwargs\n )\n if pretrained:\n init_pretrained_weights(model, model_urls['squeezenet1_1'])\n return model\n" }, { "path": "torchreid/models/xception.py", "content": "from __future__ import division, absolute_import\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.model_zoo as model_zoo\n\n__all__ = ['xception']\n\npretrained_settings = {\n 'xception': {\n 'imagenet': {\n 'url':\n 'http://data.lip6.fr/cadene/pretrainedmodels/xception-43020ad28.pth',\n 'input_space': 'RGB',\n 'input_size': [3, 299, 299],\n 'input_range': [0, 1],\n 'mean': [0.5, 0.5, 0.5],\n 'std': [0.5, 0.5, 0.5],\n 'num_classes': 1000,\n 'scale':\n 0.8975 # The resize parameter of the validation transform should be 333, and make sure to center crop at 299x299\n }\n }\n}\n\n\nclass SeparableConv2d(nn.Module):\n\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size=1,\n stride=1,\n padding=0,\n dilation=1,\n bias=False\n ):\n super(SeparableConv2d, self).__init__()\n\n self.conv1 = nn.Conv2d(\n in_channels,\n in_channels,\n kernel_size,\n stride,\n padding,\n dilation,\n groups=in_channels,\n bias=bias\n )\n self.pointwise = nn.Conv2d(\n in_channels, out_channels, 1, 1, 0, 1, 1, bias=bias\n )\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.pointwise(x)\n return x\n\n\nclass Block(nn.Module):\n\n def __init__(\n self,\n in_filters,\n out_filters,\n reps,\n strides=1,\n start_with_relu=True,\n grow_first=True\n ):\n super(Block, self).__init__()\n\n if out_filters != in_filters or strides != 1:\n self.skip = nn.Conv2d(\n in_filters, out_filters, 1, stride=strides, bias=False\n )\n self.skipbn = nn.BatchNorm2d(out_filters)\n else:\n self.skip = None\n\n self.relu = nn.ReLU(inplace=True)\n rep = []\n\n filters = in_filters\n if grow_first:\n rep.append(self.relu)\n rep.append(\n SeparableConv2d(\n in_filters,\n out_filters,\n 3,\n stride=1,\n padding=1,\n bias=False\n )\n )\n rep.append(nn.BatchNorm2d(out_filters))\n filters = out_filters\n\n for i in range(reps - 1):\n rep.append(self.relu)\n rep.append(\n SeparableConv2d(\n filters, filters, 3, stride=1, padding=1, bias=False\n )\n )\n rep.append(nn.BatchNorm2d(filters))\n\n if not grow_first:\n rep.append(self.relu)\n rep.append(\n SeparableConv2d(\n in_filters,\n out_filters,\n 3,\n stride=1,\n padding=1,\n bias=False\n )\n )\n rep.append(nn.BatchNorm2d(out_filters))\n\n if not start_with_relu:\n rep = rep[1:]\n else:\n rep[0] = nn.ReLU(inplace=False)\n\n if strides != 1:\n rep.append(nn.MaxPool2d(3, strides, 1))\n self.rep = nn.Sequential(*rep)\n\n def forward(self, inp):\n x = self.rep(inp)\n\n if self.skip is not None:\n skip = self.skip(inp)\n skip = self.skipbn(skip)\n else:\n skip = inp\n\n x += skip\n return x\n\n\nclass Xception(nn.Module):\n \"\"\"Xception.\n \n Reference:\n Chollet. Xception: Deep Learning with Depthwise\n Separable Convolutions. CVPR 2017.\n\n Public keys:\n - ``xception``: Xception.\n \"\"\"\n\n def __init__(\n self, num_classes, loss, fc_dims=None, dropout_p=None, **kwargs\n ):\n super(Xception, self).__init__()\n self.loss = loss\n\n self.conv1 = nn.Conv2d(3, 32, 3, 2, 0, bias=False)\n self.bn1 = nn.BatchNorm2d(32)\n\n self.conv2 = nn.Conv2d(32, 64, 3, bias=False)\n self.bn2 = nn.BatchNorm2d(64)\n\n self.block1 = Block(\n 64, 128, 2, 2, start_with_relu=False, grow_first=True\n )\n self.block2 = Block(\n 128, 256, 2, 2, start_with_relu=True, grow_first=True\n )\n self.block3 = Block(\n 256, 728, 2, 2, start_with_relu=True, grow_first=True\n )\n\n self.block4 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n self.block5 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n self.block6 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n self.block7 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n\n self.block8 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n self.block9 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n self.block10 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n self.block11 = Block(\n 728, 728, 3, 1, start_with_relu=True, grow_first=True\n )\n\n self.block12 = Block(\n 728, 1024, 2, 2, start_with_relu=True, grow_first=False\n )\n\n self.conv3 = SeparableConv2d(1024, 1536, 3, 1, 1)\n self.bn3 = nn.BatchNorm2d(1536)\n\n self.conv4 = SeparableConv2d(1536, 2048, 3, 1, 1)\n self.bn4 = nn.BatchNorm2d(2048)\n\n self.global_avgpool = nn.AdaptiveAvgPool2d(1)\n self.feature_dim = 2048\n self.fc = self._construct_fc_layer(fc_dims, 2048, dropout_p)\n self.classifier = nn.Linear(self.feature_dim, num_classes)\n\n self._init_params()\n\n def _construct_fc_layer(self, fc_dims, input_dim, dropout_p=None):\n \"\"\"Constructs fully connected layer.\n\n Args:\n fc_dims (list or tuple): dimensions of fc layers, if None, no fc layers are constructed\n input_dim (int): input dimension\n dropout_p (float): dropout probability, if None, dropout is unused\n \"\"\"\n if fc_dims is None:\n self.feature_dim = input_dim\n return None\n\n assert isinstance(\n fc_dims, (list, tuple)\n ), 'fc_dims must be either list or tuple, but got {}'.format(\n type(fc_dims)\n )\n\n layers = []\n for dim in fc_dims:\n layers.append(nn.Linear(input_dim, dim))\n layers.append(nn.BatchNorm1d(dim))\n layers.append(nn.ReLU(inplace=True))\n if dropout_p is not None:\n layers.append(nn.Dropout(p=dropout_p))\n input_dim = dim\n\n self.feature_dim = fc_dims[-1]\n\n return nn.Sequential(*layers)\n\n def _init_params(self):\n for m in self.modules():\n if isinstance(m, nn.Conv2d):\n nn.init.kaiming_normal_(\n m.weight, mode='fan_out', nonlinearity='relu'\n )\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm2d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.Linear):\n nn.init.normal_(m.weight, 0, 0.01)\n if m.bias is not None:\n nn.init.constant_(m.bias, 0)\n\n def featuremaps(self, input):\n x = self.conv1(input)\n x = self.bn1(x)\n x = F.relu(x, inplace=True)\n\n x = self.conv2(x)\n x = self.bn2(x)\n x = F.relu(x, inplace=True)\n\n x = self.block1(x)\n x = self.block2(x)\n x = self.block3(x)\n x = self.block4(x)\n x = self.block5(x)\n x = self.block6(x)\n x = self.block7(x)\n x = self.block8(x)\n x = self.block9(x)\n x = self.block10(x)\n x = self.block11(x)\n x = self.block12(x)\n\n x = self.conv3(x)\n x = self.bn3(x)\n x = F.relu(x, inplace=True)\n\n x = self.conv4(x)\n x = self.bn4(x)\n x = F.relu(x, inplace=True)\n return x\n\n def forward(self, x):\n f = self.featuremaps(x)\n v = self.global_avgpool(f)\n v = v.view(v.size(0), -1)\n\n if self.fc is not None:\n v = self.fc(v)\n\n if not self.training:\n return v\n\n y = self.classifier(v)\n\n if self.loss == 'softmax':\n return y\n elif self.loss == 'triplet':\n return y, v\n else:\n raise KeyError('Unsupported loss: {}'.format(self.loss))\n\n\ndef init_pretrained_weights(model, model_url):\n \"\"\"Initialize models with pretrained weights.\n \n Layers that don't match with pretrained layers in name or size are kept unchanged.\n \"\"\"\n pretrain_dict = model_zoo.load_url(model_url)\n model_dict = model.state_dict()\n pretrain_dict = {\n k: v\n for k, v in pretrain_dict.items()\n if k in model_dict and model_dict[k].size() == v.size()\n }\n model_dict.update(pretrain_dict)\n model.load_state_dict(model_dict)\n\n\ndef xception(num_classes, loss='softmax', pretrained=True, **kwargs):\n model = Xception(num_classes, loss, fc_dims=None, dropout_p=None, **kwargs)\n if pretrained:\n model_url = pretrained_settings['xception']['imagenet']['url']\n init_pretrained_weights(model, model_url)\n return model\n" }, { "path": "torchreid/optim/__init__.py", "content": "from __future__ import absolute_import\n\nfrom .optimizer import build_optimizer\nfrom .lr_scheduler import build_lr_scheduler\n" }, { "path": "torchreid/optim/lr_scheduler.py", "content": "from __future__ import print_function, absolute_import\nfrom bisect import bisect_right\nimport torch\n\nAVAI_SCH = ['single_step', 'multi_step', 'warmup_multi_step', 'cosine']\n\n\ndef build_lr_scheduler(\n optimizer, lr_scheduler='single_step', stepsize=1, gamma=0.1, max_epoch=1\n):\n \"\"\"A function wrapper for building a learning rate scheduler.\n\n Args:\n optimizer (Optimizer): an Optimizer.\n lr_scheduler (str, optional): learning rate scheduler method. Default is single_step.\n stepsize (int or list, optional): step size to decay learning rate. When ``lr_scheduler``\n is \"single_step\", ``stepsize`` should be an integer. When ``lr_scheduler`` is\n \"multi_step\", ``stepsize`` is a list. Default is 1.\n gamma (float, optional): decay rate. Default is 0.1.\n max_epoch (int, optional): maximum epoch (for cosine annealing). Default is 1.\n\n Examples::\n >>> # Decay learning rate by every 20 epochs.\n >>> scheduler = torchreid.optim.build_lr_scheduler(\n >>> optimizer, lr_scheduler='single_step', stepsize=20\n >>> )\n >>> # Decay learning rate at 30, 50 and 55 epochs.\n >>> scheduler = torchreid.optim.build_lr_scheduler(\n >>> optimizer, lr_scheduler='multi_step', stepsize=[30, 50, 55]\n >>> )\n \"\"\"\n if lr_scheduler not in AVAI_SCH:\n raise ValueError(\n 'Unsupported scheduler: {}. Must be one of {}'.format(\n lr_scheduler, AVAI_SCH\n )\n )\n\n if lr_scheduler == 'single_step':\n if isinstance(stepsize, list):\n stepsize = stepsize[-1]\n\n if not isinstance(stepsize, int):\n raise TypeError(\n 'For single_step lr_scheduler, stepsize must '\n 'be an integer, but got {}'.format(type(stepsize))\n )\n\n scheduler = torch.optim.lr_scheduler.StepLR(\n optimizer, step_size=stepsize, gamma=gamma\n )\n\n elif lr_scheduler == 'multi_step':\n if not isinstance(stepsize, list):\n raise TypeError(\n 'For multi_step lr_scheduler, stepsize must '\n 'be a list, but got {}'.format(type(stepsize))\n )\n\n scheduler = torch.optim.lr_scheduler.MultiStepLR(\n optimizer, milestones=stepsize, gamma=gamma\n )\n\n elif lr_scheduler == 'warmup_multi_step':\n if not isinstance(stepsize, list):\n raise TypeError(\n 'For warmup_multi_step lr_scheduler, stepsize must '\n 'be a list, but got {}'.format(type(stepsize))\n )\n\n scheduler = WarmupMultiStepLR(\n optimizer,\n milestones=stepsize,\n gamma=gamma,\n warmup_factor=0.01,\n warmup_iters=10,\n warmup_method=\"linear\"\n )\n\n elif lr_scheduler == 'cosine':\n scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n optimizer, float(max_epoch)\n )\n\n return scheduler\n\n\nclass WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):\n \"\"\"Source: https://github.com/michuanhaohao/reid-strong-baseline\"\"\"\n def __init__(\n self,\n optimizer,\n milestones,\n gamma=0.1,\n warmup_factor=1.0 / 3,\n warmup_iters=500,\n warmup_method=\"linear\",\n last_epoch=-1,\n ):\n if not list(milestones) == sorted(milestones):\n raise ValueError(\n \"Milestones should be a list of\" \" increasing integers. Got {}\",\n milestones,\n )\n\n if warmup_method not in (\"constant\", \"linear\"):\n raise ValueError(\n \"Only 'constant' or 'linear' warmup_method accepted\"\n \"got {}\".format(warmup_method)\n )\n self.milestones = milestones\n self.gamma = gamma\n self.warmup_factor = warmup_factor\n self.warmup_iters = warmup_iters\n self.warmup_method = warmup_method\n super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch)\n\n def get_lr(self):\n warmup_factor = 1\n if self.last_epoch < self.warmup_iters:\n if self.warmup_method == \"constant\":\n warmup_factor = self.warmup_factor\n elif self.warmup_method == \"linear\":\n alpha = self.last_epoch / self.warmup_iters\n warmup_factor = self.warmup_factor * (1 - alpha) + alpha\n return [\n base_lr\n * warmup_factor\n * self.gamma ** bisect_right(self.milestones, self.last_epoch)\n for base_lr in self.base_lrs\n ]\n" }, { "path": "torchreid/optim/optimizer.py", "content": "from __future__ import print_function, absolute_import\nimport warnings\nimport torch\nimport torch.nn as nn\n\nfrom .radam import RAdam\n\nAVAI_OPTIMS = ['adam', 'amsgrad', 'sgd', 'rmsprop', 'radam']\n\n\ndef build_optimizer(\n model,\n optim='adam',\n lr=0.0003,\n weight_decay=5e-04,\n momentum=0.9,\n sgd_dampening=0,\n sgd_nesterov=False,\n rmsprop_alpha=0.99,\n adam_beta1=0.9,\n adam_beta2=0.99,\n staged_lr=False,\n new_layers='',\n base_lr_mult=0.1\n):\n \"\"\"A function wrapper for building an optimizer.\n\n Args:\n model (nn.Module): model.\n optim (str, optional): optimizer. Default is \"adam\".\n lr (float, optional): learning rate. Default is 0.0003.\n weight_decay (float, optional): weight decay (L2 penalty). Default is 5e-04.\n momentum (float, optional): momentum factor in sgd. Default is 0.9,\n sgd_dampening (float, optional): dampening for momentum. Default is 0.\n sgd_nesterov (bool, optional): enables Nesterov momentum. Default is False.\n rmsprop_alpha (float, optional): smoothing constant for rmsprop. Default is 0.99.\n adam_beta1 (float, optional): beta-1 value in adam. Default is 0.9.\n adam_beta2 (float, optional): beta-2 value in adam. Default is 0.99,\n staged_lr (bool, optional): uses different learning rates for base and new layers. Base\n layers are pretrained layers while new layers are randomly initialized, e.g. the\n identity classification layer. Enabling ``staged_lr`` can allow the base layers to\n be trained with a smaller learning rate determined by ``base_lr_mult``, while the new\n layers will take the ``lr``. Default is False.\n new_layers (str or list): attribute names in ``model``. Default is empty.\n base_lr_mult (float, optional): learning rate multiplier for base layers. Default is 0.1.\n\n Examples::\n >>> # A normal optimizer can be built by\n >>> optimizer = torchreid.optim.build_optimizer(model, optim='sgd', lr=0.01)\n >>> # If you want to use a smaller learning rate for pretrained layers\n >>> # and the attribute name for the randomly initialized layer is 'classifier',\n >>> # you can do\n >>> optimizer = torchreid.optim.build_optimizer(\n >>> model, optim='sgd', lr=0.01, staged_lr=True,\n >>> new_layers='classifier', base_lr_mult=0.1\n >>> )\n >>> # Now the `classifier` has learning rate 0.01 but the base layers\n >>> # have learning rate 0.01 * 0.1.\n >>> # new_layers can also take multiple attribute names. Say the new layers\n >>> # are 'fc' and 'classifier', you can do\n >>> optimizer = torchreid.optim.build_optimizer(\n >>> model, optim='sgd', lr=0.01, staged_lr=True,\n >>> new_layers=['fc', 'classifier'], base_lr_mult=0.1\n >>> )\n \"\"\"\n if optim not in AVAI_OPTIMS:\n raise ValueError(\n 'Unsupported optim: {}. Must be one of {}'.format(\n optim, AVAI_OPTIMS\n )\n )\n\n if not isinstance(model, nn.Module):\n raise TypeError(\n 'model given to build_optimizer must be an instance of nn.Module'\n )\n\n if staged_lr:\n if isinstance(new_layers, str):\n if new_layers is None:\n warnings.warn(\n 'new_layers is empty, therefore, staged_lr is useless'\n )\n new_layers = [new_layers]\n\n if isinstance(model, nn.DataParallel):\n model = model.module\n\n base_params = []\n base_layers = []\n new_params = []\n\n for name, module in model.named_children():\n if name in new_layers:\n new_params += [p for p in module.parameters()]\n else:\n base_params += [p for p in module.parameters()]\n base_layers.append(name)\n\n param_groups = [\n {\n 'params': base_params,\n 'lr': lr * base_lr_mult\n },\n {\n 'params': new_params\n },\n ]\n\n else:\n param_groups = model.parameters()\n\n if optim == 'adam':\n optimizer = torch.optim.Adam(\n param_groups,\n lr=lr,\n weight_decay=weight_decay,\n betas=(adam_beta1, adam_beta2),\n )\n\n elif optim == 'amsgrad':\n optimizer = torch.optim.Adam(\n param_groups,\n lr=lr,\n weight_decay=weight_decay,\n betas=(adam_beta1, adam_beta2),\n amsgrad=True,\n )\n\n elif optim == 'sgd':\n optimizer = torch.optim.SGD(\n param_groups,\n lr=lr,\n momentum=momentum,\n weight_decay=weight_decay,\n dampening=sgd_dampening,\n nesterov=sgd_nesterov,\n )\n\n elif optim == 'rmsprop':\n optimizer = torch.optim.RMSprop(\n param_groups,\n lr=lr,\n momentum=momentum,\n weight_decay=weight_decay,\n alpha=rmsprop_alpha,\n )\n\n elif optim == 'radam':\n optimizer = RAdam(\n param_groups,\n lr=lr,\n weight_decay=weight_decay,\n betas=(adam_beta1, adam_beta2)\n )\n\n return optimizer\n" }, { "path": "torchreid/optim/radam.py", "content": "\"\"\"\nImported from: https://github.com/LiyuanLucasLiu/RAdam\n\nPaper: https://arxiv.org/abs/1908.03265\n\n@article{liu2019radam,\n title={On the Variance of the Adaptive Learning Rate and Beyond},\n author={Liu, Liyuan and Jiang, Haoming and He, Pengcheng and Chen, Weizhu and Liu, Xiaodong and Gao, Jianfeng and Han, Jiawei},\n journal={arXiv preprint arXiv:1908.03265},\n year={2019}\n}\n\"\"\"\nfrom __future__ import print_function, absolute_import\nimport math\nimport torch\nfrom torch.optim.optimizer import Optimizer\n\n\nclass RAdam(Optimizer):\n\n def __init__(\n self,\n params,\n lr=1e-3,\n betas=(0.9, 0.999),\n eps=1e-8,\n weight_decay=0,\n degenerated_to_sgd=True\n ):\n if not 0.0 <= lr:\n raise ValueError(\"Invalid learning rate: {}\".format(lr))\n if not 0.0 <= eps:\n raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n if not 0.0 <= betas[0] < 1.0:\n raise ValueError(\n \"Invalid beta parameter at index 0: {}\".format(betas[0])\n )\n if not 0.0 <= betas[1] < 1.0:\n raise ValueError(\n \"Invalid beta parameter at index 1: {}\".format(betas[1])\n )\n\n self.degenerated_to_sgd = degenerated_to_sgd\n defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)\n self.buffer = [[None, None, None] for ind in range(10)]\n super(RAdam, self).__init__(params, defaults)\n\n def __setstate__(self, state):\n super(RAdam, self).__setstate__(state)\n\n def step(self, closure=None):\n\n loss = None\n if closure is not None:\n loss = closure()\n\n for group in self.param_groups:\n\n for p in group['params']:\n if p.grad is None:\n continue\n grad = p.grad.data.float()\n if grad.is_sparse:\n raise RuntimeError(\n 'RAdam does not support sparse gradients'\n )\n\n p_data_fp32 = p.data.float()\n\n state = self.state[p]\n\n if len(state) == 0:\n state['step'] = 0\n state['exp_avg'] = torch.zeros_like(p_data_fp32)\n state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n else:\n state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n state['exp_avg_sq'] = state['exp_avg_sq'].type_as(\n p_data_fp32\n )\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n beta1, beta2 = group['betas']\n\n exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n state['step'] += 1\n buffered = self.buffer[int(state['step'] % 10)]\n if state['step'] == buffered[0]:\n N_sma, step_size = buffered[1], buffered[2]\n else:\n buffered[0] = state['step']\n beta2_t = beta2**state['step']\n N_sma_max = 2 / (1-beta2) - 1\n N_sma = N_sma_max - 2 * state['step'\n ] * beta2_t / (1-beta2_t)\n buffered[1] = N_sma\n\n # more conservative since it's an approximated value\n if N_sma >= 5:\n step_size = math.sqrt(\n (1-beta2_t) * (N_sma-4) / (N_sma_max-4) *\n (N_sma-2) / N_sma * N_sma_max / (N_sma_max-2)\n ) / (1 - beta1**state['step'])\n elif self.degenerated_to_sgd:\n step_size = 1.0 / (1 - beta1**state['step'])\n else:\n step_size = -1\n buffered[2] = step_size\n\n # more conservative since it's an approximated value\n if N_sma >= 5:\n if group['weight_decay'] != 0:\n p_data_fp32.add_(\n -group['weight_decay'] * group['lr'], p_data_fp32\n )\n denom = exp_avg_sq.sqrt().add_(group['eps'])\n p_data_fp32.addcdiv_(\n -step_size * group['lr'], exp_avg, denom\n )\n p.data.copy_(p_data_fp32)\n elif step_size > 0:\n if group['weight_decay'] != 0:\n p_data_fp32.add_(\n -group['weight_decay'] * group['lr'], p_data_fp32\n )\n p_data_fp32.add_(-step_size * group['lr'], exp_avg)\n p.data.copy_(p_data_fp32)\n\n return loss\n\n\nclass PlainRAdam(Optimizer):\n\n def __init__(\n self,\n params,\n lr=1e-3,\n betas=(0.9, 0.999),\n eps=1e-8,\n weight_decay=0,\n degenerated_to_sgd=True\n ):\n if not 0.0 <= lr:\n raise ValueError(\"Invalid learning rate: {}\".format(lr))\n if not 0.0 <= eps:\n raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n if not 0.0 <= betas[0] < 1.0:\n raise ValueError(\n \"Invalid beta parameter at index 0: {}\".format(betas[0])\n )\n if not 0.0 <= betas[1] < 1.0:\n raise ValueError(\n \"Invalid beta parameter at index 1: {}\".format(betas[1])\n )\n\n self.degenerated_to_sgd = degenerated_to_sgd\n defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)\n\n super(PlainRAdam, self).__init__(params, defaults)\n\n def __setstate__(self, state):\n super(PlainRAdam, self).__setstate__(state)\n\n def step(self, closure=None):\n\n loss = None\n if closure is not None:\n loss = closure()\n\n for group in self.param_groups:\n\n for p in group['params']:\n if p.grad is None:\n continue\n grad = p.grad.data.float()\n if grad.is_sparse:\n raise RuntimeError(\n 'RAdam does not support sparse gradients'\n )\n\n p_data_fp32 = p.data.float()\n\n state = self.state[p]\n\n if len(state) == 0:\n state['step'] = 0\n state['exp_avg'] = torch.zeros_like(p_data_fp32)\n state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n else:\n state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n state['exp_avg_sq'] = state['exp_avg_sq'].type_as(\n p_data_fp32\n )\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n beta1, beta2 = group['betas']\n\n exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n state['step'] += 1\n beta2_t = beta2**state['step']\n N_sma_max = 2 / (1-beta2) - 1\n N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1-beta2_t)\n\n # more conservative since it's an approximated value\n if N_sma >= 5:\n if group['weight_decay'] != 0:\n p_data_fp32.add_(\n -group['weight_decay'] * group['lr'], p_data_fp32\n )\n step_size = group['lr'] * math.sqrt(\n (1-beta2_t) * (N_sma-4) / (N_sma_max-4) *\n (N_sma-2) / N_sma * N_sma_max / (N_sma_max-2)\n ) / (1 - beta1**state['step'])\n denom = exp_avg_sq.sqrt().add_(group['eps'])\n p_data_fp32.addcdiv_(-step_size, exp_avg, denom)\n p.data.copy_(p_data_fp32)\n elif self.degenerated_to_sgd:\n if group['weight_decay'] != 0:\n p_data_fp32.add_(\n -group['weight_decay'] * group['lr'], p_data_fp32\n )\n step_size = group['lr'] / (1 - beta1**state['step'])\n p_data_fp32.add_(-step_size, exp_avg)\n p.data.copy_(p_data_fp32)\n\n return loss\n\n\nclass AdamW(Optimizer):\n\n def __init__(\n self,\n params,\n lr=1e-3,\n betas=(0.9, 0.999),\n eps=1e-8,\n weight_decay=0,\n warmup=0\n ):\n if not 0.0 <= lr:\n raise ValueError(\"Invalid learning rate: {}\".format(lr))\n if not 0.0 <= eps:\n raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n if not 0.0 <= betas[0] < 1.0:\n raise ValueError(\n \"Invalid beta parameter at index 0: {}\".format(betas[0])\n )\n if not 0.0 <= betas[1] < 1.0:\n raise ValueError(\n \"Invalid beta parameter at index 1: {}\".format(betas[1])\n )\n\n defaults = dict(\n lr=lr,\n betas=betas,\n eps=eps,\n weight_decay=weight_decay,\n warmup=warmup\n )\n super(AdamW, self).__init__(params, defaults)\n\n def __setstate__(self, state):\n super(AdamW, self).__setstate__(state)\n\n def step(self, closure=None):\n loss = None\n if closure is not None:\n loss = closure()\n\n for group in self.param_groups:\n\n for p in group['params']:\n if p.grad is None:\n continue\n grad = p.grad.data.float()\n if grad.is_sparse:\n raise RuntimeError(\n 'Adam does not support sparse gradients, please consider SparseAdam instead'\n )\n\n p_data_fp32 = p.data.float()\n\n state = self.state[p]\n\n if len(state) == 0:\n state['step'] = 0\n state['exp_avg'] = torch.zeros_like(p_data_fp32)\n state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)\n else:\n state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)\n state['exp_avg_sq'] = state['exp_avg_sq'].type_as(\n p_data_fp32\n )\n\n exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n beta1, beta2 = group['betas']\n\n state['step'] += 1\n\n exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n exp_avg.mul_(beta1).add_(1 - beta1, grad)\n\n denom = exp_avg_sq.sqrt().add_(group['eps'])\n bias_correction1 = 1 - beta1**state['step']\n bias_correction2 = 1 - beta2**state['step']\n\n if group['warmup'] > state['step']:\n scheduled_lr = 1e-8 + state['step'] * group['lr'] / group[\n 'warmup']\n else:\n scheduled_lr = group['lr']\n\n step_size = scheduled_lr * math.sqrt(\n bias_correction2\n ) / bias_correction1\n\n if group['weight_decay'] != 0:\n p_data_fp32.add_(\n -group['weight_decay'] * scheduled_lr, p_data_fp32\n )\n\n p_data_fp32.addcdiv_(-step_size, exp_avg, denom)\n\n p.data.copy_(p_data_fp32)\n\n return loss\n" }, { "path": "torchreid/scripts/__init__.py", "content": "" }, { "path": "torchreid/scripts/default_config.py", "content": "import random\nimport uuid\nfrom datetime import datetime\nfrom yacs.config import CfgNode as CN\nfrom torchreid.utils.constants import *\nfrom deepdiff import DeepDiff\nimport re\nimport pprint\n\n\ndef get_default_config():\n cfg = CN()\n\n # project\n cfg.project = CN()\n cfg.project.name = \"BPBreID\" # will be used as WanDB project name\n cfg.project.experiment_name = \"\"\n cfg.project.diff_config = \"\"\n cfg.project.notes = \"\"\n cfg.project.tags = []\n cfg.project.config_file = \"\"\n cfg.project.debug_mode = False\n cfg.project.logger = CN() # Choose experiment manager client to use or simply use disk dump / matplotlib\n cfg.project.logger.use_clearml = False\n cfg.project.logger.use_neptune = False\n cfg.project.logger.use_tensorboard = False\n cfg.project.logger.use_wandb = False\n cfg.project.logger.matplotlib_show = False\n cfg.project.logger.save_disk = True # save images to disk\n cfg.project.job_id = random.randint(0, 1_000_000_000)\n cfg.project.experiment_id = str(uuid.uuid4())\n cfg.project.start_time = datetime.now().strftime(\"%Y_%m_%d_%H_%M_%S_%MS\")\n\n # model\n cfg.model = CN()\n cfg.model.name = 'bpbreid'\n cfg.model.pretrained = True # automatically load pretrained model weights if available (For example HRNet\n # pretrained weights on ImageNet)\n cfg.model.load_weights = '' # path to model weights, for doing inference with a model that was saved on disk with 'save_model_flag'\n cfg.model.load_config = False # load config saved with model weights and overwrite current config\n cfg.model.resume = '' # path to checkpoint for resume training\n cfg.model.save_model_flag = False # path to checkpoint for resume training\n # configs for our part-based model BPBreID\n cfg.model.bpbreid = CN()\n cfg.model.bpbreid.pooling = 'gwap' # ['gap', 'gmp', 'gwap', 'gwap2']\n cfg.model.bpbreid.normalization = 'identity' # ['identity', 'batch_norm_2d'] - obsolete, always use identity\n cfg.model.bpbreid.mask_filtering_training = False # use visibility scores at training - do not have an influence on testing performance yet, to be improved\n cfg.model.bpbreid.mask_filtering_testing = True # use visibility scores at testing - do have a big influence on testing performance when activated\n cfg.model.bpbreid.last_stride = 1 # last stride of the resnet backbone - 1 for better performance\n cfg.model.bpbreid.dim_reduce = 'after_pooling' # where to apply feature dimensionality reduction (before or after global pooling) ['none', 'before_pooling', 'after_pooling', 'before_and_after_pooling', 'after_pooling_with_dropout']\n cfg.model.bpbreid.dim_reduce_output = 512 # reduce feature dimension to this value when above config is not 'none'\n cfg.model.bpbreid.backbone = 'resnet50' # ['resnet50', 'hrnet32', 'fastreid_resnet_ibn_nl']\n cfg.model.bpbreid.learnable_attention_enabled = True # use learnable attention mechanism to pool part features, otherwise, use fixed attention weights from external (pifpaf) heatmaps/masks\n cfg.model.bpbreid.test_embeddings = ['bn_foreg', 'parts'] # embeddings to use at inference among ['globl', 'foreg', 'backg', 'conct', 'parts']: append 'bn_' suffix to use batch normed embeddings\n cfg.model.bpbreid.test_use_target_segmentation = 'none' # ['soft', 'hard', 'none'] - use external part mask to further refine the attention weights at inference\n cfg.model.bpbreid.training_binary_visibility_score = True # use binary visibility score (0 or 1) instead of continuous visibility score (0 to 1) at training\n cfg.model.bpbreid.testing_binary_visibility_score = True # use binary visibility score (0 or 1) instead of continuous visibility score (0 to 1) at testing\n cfg.model.bpbreid.shared_parts_id_classifier = False # if each part branch uses share weights for the identity classifier. Used only when the identity loss is used on part-based embeddings.\n cfg.model.bpbreid.hrnet_pretrained_path = \"pretrained_models/\" # path to pretrained weights for HRNet backbone, download on our Google Drive or on https://github.com/HRNet/HRNet-Image-Classification\n # number of horizontal stripes desired. When BPBreID is used, this variable will be automatically filled depending\n # on \"data.masks.preprocess\"\n cfg.model.bpbreid.masks = CN()\n cfg.model.bpbreid.masks.type = 'disk' # when 'disk' is used, load part masks from storage in 'cfg.model.bpbreid.masks.dir' folder\n # when 'stripes' is used, divide the image in 'cfg.model.bpbreid.masks.parts_num' horizontal stripes in a PCB style.\n # 'stripes' with parts_num=1 can be used to emulate the global method Bag of Tricks (BoT)\n cfg.model.bpbreid.masks.parts_num = 1 # number of part-based embedding to extract. When PCB is used, change this parameter to the number of stripes required\n cfg.model.bpbreid.masks.dir = 'pifpaf_maskrcnn_filtering' # masks will be loaded from 'dataset_path/masks/' directory\n cfg.model.bpbreid.masks.preprocess = 'eight' # how to group the 36 pifpaf parts into smaller human semantic groups ['eight', 'five', 'four', 'two', ...], more combination available inside 'torchreid/data/masks_transforms/__init__.masks_preprocess_pifpaf'\n cfg.model.bpbreid.masks.softmax_weight = 15\n cfg.model.bpbreid.masks.background_computation_strategy = 'threshold' # threshold, diff_from_max\n cfg.model.bpbreid.masks.mask_filtering_threshold = 0.5\n\n # data\n cfg.data = CN()\n cfg.data.type = 'image'\n cfg.data.root = 'reid-data'\n cfg.data.sources = ['market1501']\n cfg.data.targets = ['market1501']\n cfg.data.workers = 4 # number of data loading workers, set to 0 to enable breakpoint debugging in dataloader code\n cfg.data.split_id = 0 # split index\n cfg.data.height = 256 # image height\n cfg.data.width = 128 # image width\n cfg.data.combineall = False # combine train, query and gallery for training\n cfg.data.transforms = ['rc', 're'] # data augmentation from ['rf', 'rc', 're', 'cj'] = ['random flip', 'random crop', 'random erasing', 'color jitter']\n cfg.data.ro = CN() # parameters for random occlusion data augmentation with Pascal VOC, to be improved, not maintained\n cfg.data.ro.path = \"\"\n cfg.data.ro.p = 0.5\n cfg.data.ro.n = 1\n cfg.data.ro.min_overlap = 0.5\n cfg.data.ro.max_overlap = 0.8\n cfg.data.cj = CN() # parameters for color jitter data augmentation\n cfg.data.cj.brightness = 0.2\n cfg.data.cj.contrast = 0.15\n cfg.data.cj.saturation = 0.\n cfg.data.cj.hue = 0.\n cfg.data.cj.always_apply = False\n cfg.data.cj.p = 0.5\n cfg.data.norm_mean = [0.485, 0.456, 0.406] # default is imagenet mean\n cfg.data.norm_std = [0.229, 0.224, 0.225] # default is imagenet std\n cfg.data.save_dir = 'logs' # save figures, images, logs, etc. in this folder\n cfg.data.load_train_targets = False\n\n # specific datasets\n cfg.market1501 = CN()\n cfg.market1501.use_500k_distractors = False # add 500k distractors to the gallery set for market1501\n cfg.cuhk03 = CN()\n cfg.cuhk03.labeled_images = False # use labeled images, if False, use detected images\n cfg.cuhk03.classic_split = False # use classic split by Li et al. CVPR14\n cfg.cuhk03.use_metric_cuhk03 = False # use cuhk03's metric for evaluation\n\n # sampler\n cfg.sampler = CN()\n cfg.sampler.train_sampler = 'RandomIdentitySampler' # sampler for source train loader\n cfg.sampler.train_sampler_t = 'RandomIdentitySampler' # sampler for target train loader\n cfg.sampler.num_instances = 4 # number of instances per identity for RandomIdentitySampler\n\n # video reid setting\n cfg.video = CN()\n cfg.video.seq_len = 15 # number of images to sample in a tracklet\n cfg.video.sample_method = 'evenly' # how to sample images from a tracklet 'random'/'evenly'/'all'\n cfg.video.pooling_method = 'avg' # how to pool features over a tracklet\n\n # train\n cfg.train = CN()\n cfg.train.optim = 'adam'\n cfg.train.lr = 0.00035\n cfg.train.weight_decay = 5e-4\n cfg.train.max_epoch = 120\n cfg.train.start_epoch = 0\n cfg.train.batch_size = 64\n cfg.train.fixbase_epoch = 0 # number of epochs to fix base layers\n cfg.train.open_layers = [\n 'classifier'\n ] # layers for training while keeping others frozen\n cfg.train.staged_lr = False # set different lr to different layers\n cfg.train.new_layers = ['classifier'] # newly added layers with default lr\n cfg.train.base_lr_mult = 0.1 # learning rate multiplier for base layers\n cfg.train.lr_scheduler = 'warmup_multi_step'\n cfg.train.stepsize = [40, 70] # stepsize to decay learning rate\n cfg.train.gamma = 0.1 # learning rate decay multiplier\n cfg.train.seed = 1 # random seed\n cfg.train.eval_freq = -1 # evaluation frequency (-1 means to only test after training)\n cfg.train.batch_debug_freq = 0\n cfg.train.batch_log_freq = 0\n\n # optimizer\n cfg.sgd = CN()\n cfg.sgd.momentum = 0.9 # momentum factor for sgd and rmsprop\n cfg.sgd.dampening = 0. # dampening for momentum\n cfg.sgd.nesterov = False # Nesterov momentum\n cfg.rmsprop = CN()\n cfg.rmsprop.alpha = 0.99 # smoothing constant\n cfg.adam = CN()\n cfg.adam.beta1 = 0.9 # exponential decay rate for first moment\n cfg.adam.beta2 = 0.999 # exponential decay rate for second moment\n\n # loss\n cfg.loss = CN()\n cfg.loss.name = 'part_based' # use part based engine to train bpbreid with GiLt loss\n cfg.loss.part_based = CN()\n cfg.loss.part_based.name = 'part_averaged_triplet_loss' # ['inter_parts_triplet_loss', 'intra_parts_triplet_loss', 'part_max_triplet_loss', 'part_averaged_triplet_loss', 'part_min_triplet_loss', 'part_max_min_triplet_loss', 'part_random_max_min_triplet_loss']\n cfg.loss.part_based.ppl = \"cl\" # body part prediction loss: ['cl', 'fl', 'dl'] = [cross entropy loss with label smoothing, focal loss, dice loss]\n cfg.loss.part_based.weights = CN() # weights to apply for the different losses and different types of embeddings, for more details, have a look at 'torchreid/losses/GiLt_loss.py'\n cfg.loss.part_based.weights[GLOBAL] = CN()\n cfg.loss.part_based.weights[GLOBAL].id = 1.\n cfg.loss.part_based.weights[GLOBAL].tr = 0.\n cfg.loss.part_based.weights[FOREGROUND] = CN()\n cfg.loss.part_based.weights[FOREGROUND].id = 1.\n cfg.loss.part_based.weights[FOREGROUND].tr = 0.\n cfg.loss.part_based.weights[CONCAT_PARTS] = CN()\n cfg.loss.part_based.weights[CONCAT_PARTS].id = 1.\n cfg.loss.part_based.weights[CONCAT_PARTS].tr = 0.\n cfg.loss.part_based.weights[PARTS] = CN()\n cfg.loss.part_based.weights[PARTS].id = 0.\n cfg.loss.part_based.weights[PARTS].tr = 1.\n cfg.loss.part_based.weights[PIXELS] = CN()\n cfg.loss.part_based.weights[PIXELS].ce = 0.35\n cfg.loss.softmax = CN()\n cfg.loss.softmax.label_smooth = True # use label smoothing regularizer\n cfg.loss.triplet = CN()\n cfg.loss.triplet.margin = 0.3 # distance margin\n cfg.loss.triplet.weight_t = 1. # weight to balance hard triplet loss\n cfg.loss.triplet.weight_x = 0. # weight to balance cross entropy loss\n\n # test\n cfg.test = CN()\n cfg.test.batch_size = 128\n cfg.test.batch_size_pairwise_dist_matrix = 500 # query to gallery distance matrix is computed on the GPU by batch of gallery samples with this size.\n # To avoid out of memory issue, we don't compute it for all gallery samples at the same time, but we compute it\n # in batches of 'batch_size_pairwise_dist_matrix' gallery samples.\n cfg.test.dist_metric = 'euclidean' # distance metric, ['euclidean', 'cosine']\n cfg.test.normalize_feature = True # normalize feature vectors before computing distance\n cfg.test.ranks = [1, 5, 10, 20] # cmc ranks\n cfg.test.evaluate = False # test only\n cfg.test.start_eval = 0 # start to evaluate after a specific epoch\n cfg.test.rerank = False # use person re-ranking\n cfg.test.visrank = False # visualize ranked results (only available when cfg.test.evaluate=True)\n cfg.test.visrank_topk = 10 # top-k ranks to visualize\n cfg.test.visrank_count = 10 # number of top-k ranks to plot\n cfg.test.visrank_q_idx_list = [0, 1, 2, 3, 4, 5] # list of ids of queries for which we want to plot topk rank. If len(visrank_q_idx_list) < visrank_count, remaining ids will be random\n cfg.test.vis_feature_maps = False\n cfg.test.visrank_per_body_part = False\n cfg.test.vis_embedding_projection = False\n cfg.test.save_features = False # save test set extracted features to disk\n cfg.test.detailed_ranking = True # display ranking performance for each part individually\n cfg.test.part_based = CN()\n cfg.test.part_based.dist_combine_strat = \"mean\" # ['mean', 'max'] local part based distances are combined into a global distance using this strategy\n\n # inference\n cfg.inference = CN()\n cfg.inference.enabled = False\n cfg.inference.input_folder = \"\"\n\n return cfg\n\nkeys_to_ignore_in_diff = {\n \"cfg.project\",\n \"cfg.model.save_model_flag\",\n \"cfg.model.bpbreid.backbone\",\n \"cfg.model.bpbreid.learnable_attention_enabled\",\n \"cfg.model.bpbreid.masks.parts_num\",\n \"cfg.model.bpbreid.masks.dir\",\n \"cfg.data.type\",\n \"cfg.data.root\",\n \"cfg.data.sources\",\n \"cfg.data.targets\",\n \"cfg.data.workers\",\n \"cfg.data.split_id\",\n \"cfg.data.combineall\",\n \"cfg.data.save_dir\",\n \"cfg.train.eval_freq\",\n \"cfg.train.batch_debug_freq\",\n \"cfg.train.batch_log_freq\",\n \"cfg.test.batch_size\",\n \"cfg.test.batch_size_pairwise_dist_matrix\",\n \"cfg.test.dist_metric\",\n \"cfg.test.ranks\",\n \"cfg.test.evaluate\",\n \"cfg.test.start_eval\",\n \"cfg.test.rerank\",\n \"cfg.test.visrank\",\n \"cfg.test.visrank_topk\",\n \"cfg.test.visrank_count\",\n \"cfg.test.visrank_q_idx_list\",\n \"cfg.test.vis_feature_maps\",\n \"cfg.test.visrank_per_body_part\",\n \"cfg.test.vis_embedding_projection\",\n \"cfg.test.save_features\",\n \"cfg.test.detailed_ranking\",\n \"cfg.train.open_layers\",\n \"cfg.model.load_weights\",\n}\n\ndef imagedata_kwargs(cfg):\n return {\n 'config': cfg,\n 'root': cfg.data.root,\n 'sources': cfg.data.sources,\n 'targets': cfg.data.targets,\n 'height': cfg.data.height,\n 'width': cfg.data.width,\n 'transforms': cfg.data.transforms,\n 'norm_mean': cfg.data.norm_mean,\n 'norm_std': cfg.data.norm_std,\n 'use_gpu': cfg.use_gpu,\n 'split_id': cfg.data.split_id,\n 'combineall': cfg.data.combineall,\n 'load_train_targets': cfg.data.load_train_targets,\n 'batch_size_train': cfg.train.batch_size,\n 'batch_size_test': cfg.test.batch_size,\n 'workers': cfg.data.workers,\n 'num_instances': cfg.sampler.num_instances,\n 'train_sampler': cfg.sampler.train_sampler,\n 'train_sampler_t': cfg.sampler.train_sampler_t,\n # image\n 'cuhk03_labeled': cfg.cuhk03.labeled_images,\n 'cuhk03_classic_split': cfg.cuhk03.classic_split,\n 'market1501_500k': cfg.market1501.use_500k_distractors,\n 'use_masks': cfg.loss.name == 'part_based',\n 'masks_dir': cfg.model.bpbreid.masks.dir,\n }\n\n\ndef videodata_kwargs(cfg):\n return {\n 'root': cfg.data.root,\n 'sources': cfg.data.sources,\n 'targets': cfg.data.targets,\n 'height': cfg.data.height,\n 'width': cfg.data.width,\n 'transforms': cfg.data.transforms,\n 'norm_mean': cfg.data.norm_mean,\n 'norm_std': cfg.data.norm_std,\n 'use_gpu': cfg.use_gpu,\n 'split_id': cfg.data.split_id,\n 'combineall': cfg.data.combineall,\n 'batch_size_train': cfg.train.batch_size,\n 'batch_size_test': cfg.test.batch_size,\n 'workers': cfg.data.workers,\n 'num_instances': cfg.sampler.num_instances,\n 'train_sampler': cfg.sampler.train_sampler,\n # video\n 'seq_len': cfg.video.seq_len,\n 'sample_method': cfg.video.sample_method\n }\n\n\ndef optimizer_kwargs(cfg):\n return {\n 'optim': cfg.train.optim,\n 'lr': cfg.train.lr,\n 'weight_decay': cfg.train.weight_decay,\n 'momentum': cfg.sgd.momentum,\n 'sgd_dampening': cfg.sgd.dampening,\n 'sgd_nesterov': cfg.sgd.nesterov,\n 'rmsprop_alpha': cfg.rmsprop.alpha,\n 'adam_beta1': cfg.adam.beta1,\n 'adam_beta2': cfg.adam.beta2,\n 'staged_lr': cfg.train.staged_lr,\n 'new_layers': cfg.train.new_layers,\n 'base_lr_mult': cfg.train.base_lr_mult\n }\n\n\ndef lr_scheduler_kwargs(cfg):\n return {\n 'lr_scheduler': cfg.train.lr_scheduler,\n 'stepsize': cfg.train.stepsize,\n 'gamma': cfg.train.gamma,\n 'max_epoch': cfg.train.max_epoch\n }\n\n\ndef engine_run_kwargs(cfg):\n return {\n 'save_dir': cfg.data.save_dir,\n 'fixbase_epoch': cfg.train.fixbase_epoch,\n 'open_layers': cfg.train.open_layers,\n 'test_only': cfg.test.evaluate,\n 'dist_metric': cfg.test.dist_metric,\n 'normalize_feature': cfg.test.normalize_feature,\n 'visrank': cfg.test.visrank,\n 'visrank_topk': cfg.test.visrank_topk,\n 'visrank_q_idx_list': cfg.test.visrank_q_idx_list,\n 'visrank_count': cfg.test.visrank_count,\n 'use_metric_cuhk03': cfg.cuhk03.use_metric_cuhk03,\n 'ranks': cfg.test.ranks,\n 'rerank': cfg.test.rerank,\n 'save_features': cfg.test.save_features\n }\n\n\ndef display_config_diff(cfg, default_cfg_copy):\n def iterdict(d):\n for k, v in d.items():\n if isinstance(v, dict):\n iterdict(v)\n else:\n if type(v) == list:\n v = str(v)\n d.update({k: v})\n return d\n\n ddiff = DeepDiff(iterdict(default_cfg_copy), iterdict(cfg.clone()), ignore_order=True)\n cfg_diff = {}\n if 'values_changed' in ddiff:\n for k, v in ddiff['values_changed'].items():\n reformatted_key = \"cfg.\" + k.replace(\"root['\", \"\").replace(\"']['\", \".\").replace(\"']\", \"\")\n if \"[\" in reformatted_key:\n reformatted_key = reformatted_key.split(\"[\")[0]\n reformatted_key_split = reformatted_key.split(\".\")\n ignore_key = False\n for i in range(2, len(reformatted_key_split) + 1):\n prefix = \".\".join(reformatted_key_split[0:i])\n if prefix in keys_to_ignore_in_diff:\n ignore_key = True\n break\n if not ignore_key:\n key = re.findall(r\"\\['([A-Za-z0-9_]+)'\\]\", k)[-1]\n cfg_diff[key] = v['new_value']\n print(\"Diff from default config :\")\n pprint.pprint(cfg_diff)\n if len(str(cfg_diff)) < 128:\n cfg.project.diff_config = str(cfg_diff)\n else:\n cfg.project.diff_config = str(cfg_diff)[0:124] + \"...\"" }, { "path": "torchreid/scripts/get_labels.py", "content": "import argparse\nimport glob\nimport os\nfrom pathlib import Path\nfrom typing import List\n\nimport cv2\nimport detectron2.data.transforms as T\nimport numpy as np\nimport openpifpaf\nimport torch\nimport tqdm\nfrom detectron2.checkpoint import DetectionCheckpointer\nfrom detectron2.config import CfgNode, get_cfg\nfrom detectron2.model_zoo import get_checkpoint_url, get_config_file\nfrom detectron2.modeling import build_model\nfrom detectron2.structures import Instances\nfrom torch.utils.data import DataLoader, Dataset\n\n\ndef build_config_maskrcnn(model_config_name):\n cfg = get_cfg()\n cfg.merge_from_file(cfg_filename=get_config_file(model_config_name))\n cfg.MODEL.WEIGHTS = get_checkpoint_url(model_config_name)\n cfg.MODEL.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n return cfg\n\n\ndef compare_arrays(array1, array2):\n \"\"\"\n Compare two arrays and calculate Mean Absolute Error (MAE) and percentage difference.\n\n Args:\n array1 (np.ndarray): First array.\n array2 (np.ndarray): Second array.\n\n Returns:\n mae (float): Mean Absolute Error (MAE) between the arrays.\n mae_percentage (float): Percentage difference between the arrays.\n \"\"\"\n\n def calculate_mae(array1, array2):\n # Calculate Mean Absolute Error (MAE)\n mae = np.mean(np.abs(array1 - array2))\n mae_percentage = (mae / np.max(array1)) * 100\n return mae, mae_percentage\n\n print(f\"Average percentage difference: {calculate_mae(array1, array2)[1]}%\")\n\n\ndef get_image_paths(source, path_format=False):\n \"\"\"\n Get the paths of all image files in a directory.\n\n Args:\n source (str): Directory path.\n path_format (bool, optional): Return paths as Path objects if True, otherwise as strings. Default is False.\n\n Returns:\n image_paths (List[str or Path]): List of image file paths.\n \"\"\"\n image_paths = glob.glob(f\"{source}/**/*.[jJ][pP][gG]\", recursive=True) + \\\n glob.glob(f\"{source}/**/*.[pP][nN][gG]\", recursive=True) + \\\n glob.glob(f\"{source}/**/*.[jJ][pP][eE][gG]\", recursive=True) + \\\n glob.glob(f\"{source}/**/*.[tT][iI][fF]\", recursive=True) + \\\n glob.glob(f\"{source}/**/*.[tT][iI][fF][fF]\", recursive=True)\n if path_format:\n image_paths = [Path(path_str) for path_str in image_paths]\n return image_paths\n\n\ndef format_path(img_path, dataset_dir):\n \"\"\"\n Formats the given image path based on the dataset directory.\n\n Args:\n img_path (str): The path of the image file.\n dataset_dir (str): The directory path of the dataset.\n\n Returns:\n str: The formatted path of the image file.\n \"\"\"\n if \"occluded_reid\" in dataset_dir.lower() or \"occluded-reid\" in dataset_dir.lower():\n return os.path.join(os.path.basename(os.path.dirname(os.path.dirname(img_path))), os.path.basename(img_path))\n elif \"p-dukemtmc_reid\" in dataset_dir.lower() or \"p-dukemtmc-reid\" in dataset_dir.lower():\n return os.path.join(os.path.basename(os.path.dirname(os.path.dirname(os.path.dirname(img_path)))),\n os.path.basename(os.path.dirname(os.path.dirname(img_path))), os.path.basename(img_path))\n return os.path.relpath(img_path, dataset_dir)\n\n\ndef get_label_paths(is_mask, img_paths, dataset_dir):\n \"\"\"\n Get the paths of label files corresponding to the image paths.\n\n Args:\n is_mask (bool): Indicates if the label is a mask or not.\n img_paths (List[str]): List of image file paths.\n dataset_dir (str): Directory path of the dataset.\n\n Returns:\n relative_paths (List[str]): List of relative paths of the image files.\n file_paths (List[str]): List of label file paths.\n \"\"\"\n relative_paths, file_paths = [], []\n for img_name in img_paths:\n relative_path = format_path(img_name, dataset_dir)\n if not is_mask:\n file_path = os.path.join(dataset_dir, \"masks\", \"pifpaf\", relative_path + \".confidence_fields.npy\")\n else:\n file_path = os.path.join(dataset_dir, \"masks\", \"pifpaf_maskrcnn_filtering\", relative_path + \".npy\")\n relative_paths.append(relative_path)\n file_paths.append(file_path)\n return relative_paths, file_paths\n\n\ndef skip_existing(is_mask, imagery, dataset_dir):\n \"\"\"\n Filter out image paths for which label files already exist.\n\n Args:\n is_mask (bool): Indicates if the label is a mask or not.\n imagery (List[str]): List of image file paths.\n dataset_dir (str): Directory path of the dataset.\n\n Returns:\n new_imagery (List[str]): List of image file paths for which label files do not exist.\n \"\"\"\n relative_paths, file_paths = get_label_paths(is_mask=is_mask, img_paths=imagery, dataset_dir=dataset_dir)\n new_imagery = []\n for index, file_path in enumerate(file_paths):\n if not os.path.exists(file_path):\n new_imagery.append(imagery[index])\n return new_imagery\n\n\ndef save_files(files, files_path, verbose=True):\n \"\"\"\n Save files to specified paths.\n\n Args:\n files (List[object]): List of files to be saved.\n files_path (List[str]): List of paths where files will be saved.\n verbose (bool, optional): Print progress if True. Default is True.\n \"\"\"\n for file, file_path in zip(files, files_path):\n os.makedirs(os.path.dirname(file_path), exist_ok=True)\n np.save(file_path, file)\n if verbose:\n print(f\"Processed {os.path.basename(file_path)}\")\n\n\nclass ImageDataset(Dataset):\n \"\"\"\n Custom dataset class for loading images.\n\n Args:\n imagery (List[Path]): List of image file paths.\n\n Returns:\n (str, np.ndarray): Tuple containing the image file path and the loaded image.\n \"\"\"\n\n def __init__(self, imagery: List[Path]):\n self.imagery = imagery\n\n def __getitem__(self, index):\n return self.imagery[index], cv2.imread(str(self.imagery[index]))\n\n def __len__(self):\n return len(self.imagery)\n\n\nclass BatchPifPaf:\n def __init__(self, model_name: str = \"shufflenetv2k16\", batch_size: int = None, workers: int = None):\n \"\"\"\n Initializes a BatchPifPaf object.\n\n Args:\n model_name (str): Name of the OpenPifPaf model to use.\n batch_size (int): Batch size for inference.\n workers (int): Number of workers for data loading.\n \"\"\"\n models = [\n 'resnet50',\n 'shufflenetv2k16',\n 'shufflenetv2k30',\n ]\n assert model_name in models, f\"Model name must be one of {models}\"\n\n print(f\"* OpenPifPaf model -> {model_name}\")\n # Define the OpenPifPaf model\n self.model = openpifpaf.Predictor(checkpoint=model_name, visualize_image=True, visualize_processed_image=True)\n self.batch_size = batch_size if batch_size else self.model.batch_size\n self.workers = workers if workers else self.model.loader_workers if self.model.loader_workers is not None else 0\n self.__collate = openpifpaf.datasets.collate_images_anns_meta\n\n def __call__(self, imagery: List[Path] or List[str], dataset_dir: List[Path] or List[str],\n is_overwrite: bool = False, verbose: bool = False):\n \"\"\"\n Perform batch processing on the given imagery using the OpenPifPaf model.\n\n Args:\n imagery (List[Path] or List[str]): List of image paths or image file names.\n dataset_dir (List[Path] or List[str]): List of dataset directories.\n is_overwrite (bool, optional): Whether to overwrite existing files. Defaults to False.\n verbose (bool, optional): Whether to print verbose information. Defaults to False.\n\n Yields:\n torch.Tensor: Predictions for each image as a NumPy array.\n \"\"\"\n\n assert len(imagery) > 0, \"No images found in imagery.\"\n\n if not is_overwrite:\n imagery = skip_existing(False, imagery, dataset_dir)\n\n dataset = openpifpaf.datasets.ImageList(\n imagery,\n preprocess=self.model.preprocess,\n with_raw_image=True\n )\n loader = DataLoader(\n dataset,\n self.batch_size,\n shuffle=False,\n pin_memory=self.model.device.type != 'cpu',\n num_workers=self.workers,\n collate_fn=self.__collate,\n )\n\n total_batches = len(loader)\n progress_bar = tqdm.tqdm(total=total_batches, desc=\"Processing\", unit=\"batch\")\n\n with torch.no_grad():\n for batch in loader:\n if len(batch) == 3:\n processed_image_batch, gt_anns_batch, meta_batch = batch\n elif len(batch) == 4:\n image_batch, processed_image_batch, gt_anns_batch, meta_batch = batch\n\n # Specify the file path where you want to save the .npy file\n relative_paths, file_paths = get_label_paths(False, [d[\"file_name\"] for d in meta_batch], dataset_dir)\n\n # Obtain the confidence values (pifpaf_conf) for the processed image batch\n pifpaf_conf: torch.Tensor = self.__get_pifpaf_conf(processed_image_batch)\n\n # Save the NumPy array to the .npy file\n save_files(pifpaf_conf.numpy(), file_paths, verbose)\n\n progress_bar.update(1)\n\n progress_bar.close()\n\n def __get_pifpaf_conf(self, processed_image_batch: Instances):\n \"\"\"\n Get the confidence scores from the processed image batch.\n\n Args:\n processed_image_batch (Instances): Processed image batch containing pose estimation fields.\n\n Returns:\n torch.Tensor: Confidence scores for keypoints and connections.\n \"\"\"\n # Retrieve the pose estimation fields from the model processor\n fields_batch = self.model.processor.fields_batch(self.model.model, processed_image_batch,\n device=self.model.device)\n\n # Extract the pif (keypoint) and paf (connection) fields from the batch\n pif, paf = zip(*fields_batch)\n\n # Extract the confidence scores for keypoints (index 1 in each field)\n pif_confidence_scores = torch.stack(pif)[:, :, 1]\n paf_confidence_scores = torch.stack(paf)[:, :, 1]\n\n # Concatenate the confidence scores for keypoints and connections\n pifpaf_confidence_scores = torch.cat((pif_confidence_scores, paf_confidence_scores), dim=1)\n\n # Return the concatenated confidence scores\n return pifpaf_confidence_scores\n\n\nclass BatchMask:\n def __init__(self, cfg: CfgNode or str, batch_size: int = None, workers: int = None):\n \"\"\"\n Initialize the BatchMask class for performing batched instance segmentation using a Mask R-CNN model.\n\n Args:\n cfg (CfgNode or str): Configuration options for the Mask R-CNN model.\n batch_size (int, optional): Batch size for processing images. Defaults to None.\n workers (int, optional): Number of worker processes for data loading. Defaults to None.\n \"\"\"\n # Clone the provided configuration or get a default configuration\n self.cfg = build_config_maskrcnn(cfg) if isinstance(cfg, str) else cfg.clone()\n print(f\"* MaskRCNN model -> {cfg if isinstance(cfg, str) else self.cfg.MODEL.WEIGHTS}\")\n\n # Set the batch size for processing images, defaulting to 32 if not provided\n self.batch_size = batch_size if batch_size else 32\n\n # Set the number of worker processes for data loading, defaulting to the number of CPU cores\n self.workers = workers if workers is not None else 0\n\n # Build the Mask R-CNN model\n self.model = build_model(self.cfg)\n\n # Set the model to evaluation mode\n self.model.eval()\n\n # Load the pre-trained weights for the model\n checkpointer = DetectionCheckpointer(self.model)\n checkpointer.load(self.cfg.MODEL.WEIGHTS)\n\n # Define the augmentation transform for resizing images\n self.aug = T.ResizeShortestEdge(\n [self.cfg.INPUT.MIN_SIZE_TEST, self.cfg.INPUT.MIN_SIZE_TEST],\n self.cfg.INPUT.MAX_SIZE_TEST\n )\n\n # Set the input image format to RGB or BGR based on the configuration\n self.input_format = self.cfg.INPUT.FORMAT\n assert self.input_format in [\"RGB\", \"BGR\"], self.input_format\n\n def __collate(self, batch):\n \"\"\"\n Collates a batch of images and their paths for use in data loading.\n\n Args:\n batch (list): A list of tuples containing image paths and corresponding images.\n\n Returns:\n tuple: A tuple containing two lists: the paths of the images and the processed data.\n\n \"\"\"\n paths, data = [], []\n for path, image in batch:\n if self.input_format == \"RGB\":\n # Convert image format from RGB to BGR if required by the model\n image = image[:, :, ::-1]\n height, width = image.shape[:2]\n\n # Apply augmentation and transformation to the image\n image = self.aug.get_transform(image).apply_image(image)\n image = image.astype(\"float32\").transpose(2, 0, 1)\n image = torch.as_tensor(image)\n data.append({\"image\": image, \"height\": height, \"width\": width})\n paths.append(path)\n return paths, data\n\n def __call__(self, imagery: List[Path] or List[str], dataset_dir: List[Path] or List[str],\n is_overwrite: bool = False, verbose: bool = False):\n \"\"\"\n Perform the batch processing of imagery to generate and save mask files.\n\n Args:\n imagery (List[Path] or List[str]): A list of image paths or image filenames.\n dataset_dir (List[Path] or List[str]): A list of dataset directories.\n is_overwrite (bool, optional): Whether to overwrite existing mask files. Defaults to False.\n verbose (bool, optional): Whether to print verbose information. Defaults to False.\n\n \"\"\"\n assert len(imagery) > 0, \"No images found in imagery.\"\n\n if not is_overwrite:\n # Skip existing images if overwrite is disabled\n imagery = skip_existing(True, imagery, dataset_dir)\n\n # Create an instance of the ImageDataset class\n dataset = ImageDataset(imagery)\n\n # Create a data loader for batch processing\n loader = DataLoader(\n dataset,\n self.batch_size,\n shuffle=False,\n num_workers=self.workers,\n collate_fn=self.__collate,\n pin_memory=True\n )\n\n total_batches = len(loader)\n progress_bar = tqdm.tqdm(total=total_batches, desc=\"Processing\", unit=\"batch\")\n\n with torch.no_grad():\n for paths, batch in loader:\n # Get the paths and file paths for saving the mask files\n relative_paths, pifpaf_file_paths = get_label_paths(is_mask=False, img_paths=paths,\n dataset_dir=dataset_dir)\n\n assert all(os.path.exists(path) for path in\n pifpaf_file_paths), \"Some PiPaf Label File ('.confidence_fields.npy') does not exist!\"\n\n # Filter the predictions using the mask files\n pifpaf_filtered: List[np.ndarray] = self.__filter_pifpaf_with_mask(batch, pifpaf_file_paths)\n\n # Get the file paths for saving the mask files\n _, mask_file_paths = get_label_paths(is_mask=True, img_paths=paths, dataset_dir=dataset_dir)\n\n # Save the filtered mask files\n save_files(pifpaf_filtered, mask_file_paths, verbose)\n\n progress_bar.update(1)\n\n progress_bar.close()\n\n def __filter_pifpaf_with_mask(self, batch,\n pifpaf_file_paths: List[Path] or List[str]):\n \"\"\"\n Filter PifPaf predictions using segmentation masks.\n\n Args:\n paths (List[Path] or List[str]): List of image paths or filenames.\n batch: Batch data containing images.\n pifpaf_file_paths (List[Path] or List[str]): List of PifPaf label file paths.\n\n Returns:\n List[np.ndarray]: Filtered PifPaf arrays.\n\n \"\"\"\n\n # Order the bounding boxes by distance from the center of the image(default)\n def order_bbox(image_size, bbox_list, only_horizontal=False, only_vertical=False):\n distances = []\n image_height, image_width = image_size\n center_x = image_width // 2\n center_y = image_height // 2\n\n for i, bbox in enumerate(bbox_list):\n x1, y1, x2, y2 = bbox\n bbox_center_x = (x1 + x2) // 2\n bbox_center_y = (y1 + y2) // 2\n distance = bbox_center_x if only_horizontal else bbox_center_y if only_vertical else np.sqrt(\n (bbox_center_x - center_x) ** 2 + (bbox_center_y - center_y) ** 2)\n distances.append((i, distance))\n distances = sorted(distances, key=lambda x: x[1])\n return distances\n\n # Filter segmentations masks based on class and distance from the center of the image\n def filter_masks(results):\n image_size = results[0][\"instances\"].image_size\n pred_boxes, scores, pred_classes, pred_masks = results[0][\"instances\"].get_fields().values()\n if len(pred_masks) == 0:\n raise Exception(\"Error: Pifpaf model did not return any masks!\")\n\n # Filter out all masks that are not person\n filtered_boxes, filtered_masks = zip(\n *[(box.cpu().numpy(), mask.cpu().numpy()) for box, mask, cls in\n zip(pred_boxes, pred_masks, pred_classes) if cls == 0])\n\n # Order the masks by bbox distance to the center of the image\n distances = order_bbox(image_size, filtered_boxes)\n filtered_masks = [filtered_masks[i] for i, _ in distances]\n\n return filtered_masks\n\n # Filter PifPaf array using segmentation mask\n def filter_pifpaf_with_mask(pifpaf_array, mask, is_resize_pifpaf=False, interpolation=cv2.INTER_CUBIC):\n if is_resize_pifpaf:\n # Resize the PifPaf array to match the size of the mask\n pifpaf_resized = np.transpose(pifpaf_array, (1, 2, 0))\n pifpaf_resized = cv2.resize(pifpaf_resized, dsize=(mask.shape[1], mask.shape[0]),\n interpolation=interpolation)\n pifpaf_resized = np.transpose(pifpaf_resized, (2, 0, 1))\n\n # Filter the PifPaf array using the segmentation mask\n filtered_pifpaf = mask * pifpaf_resized\n filtered_pifpaf = np.array(\n [cv2.resize(slice, (9, 17), interpolation=cv2.INTER_CUBIC) for slice in filtered_pifpaf])\n\n return filtered_pifpaf\n # Resize the mask to match the size of the PifPaf array\n mask_resized = cv2.resize(mask.astype(np.uint8), (pifpaf_array.shape[2], pifpaf_array.shape[1]))\n filtered_pifpaf = mask_resized * pifpaf_array\n return filtered_pifpaf\n\n # Get the masks from the PifPaf predictions\n masks = filter_masks(self.model(batch))\n\n # Load the PifPaf label arrays\n pifpaf_labels = [np.load(pifpaf_file_path) for pifpaf_file_path in pifpaf_file_paths]\n\n # Filter the PifPaf arrays using the masks\n pifpaf_filtered = [filter_pifpaf_with_mask(pifpaf_label, mask) for pifpaf_label, mask in\n zip(pifpaf_labels, masks)]\n\n return pifpaf_filtered\n\n\ndef main():\n # Parse command line arguments\n parser = argparse.ArgumentParser(\n formatter_class=argparse.ArgumentDefaultsHelpFormatter\n )\n parser.add_argument('-s', '--source', type=str, required=True,\n help='Source dataset containing image files')\n parser.add_argument('--maskrcnn-cfg-file', type=str, default=\"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml\",\n help='Configuration file for the Mask R-CNN model')\n parser.add_argument('--pifpaf-model-name', type=str, default=\"shufflenetv2k16\",\n help='Name of the PifPaf model')\n parser.add_argument('-b', '--batch-size', type=int,\n help='Batch size for processing images')\n parser.add_argument('--num-workers', type=int,\n help='Number of worker processes for data loading')\n args = parser.parse_args()\n\n # Get image paths\n img_paths = get_image_paths(args.source)\n\n # Perform PifPaf processing\n pifpaf_model = BatchPifPaf(model_name=args.pifpaf_model_name,\n batch_size=args.batch_size,\n workers=args.num_workers)\n pifpaf_model(imagery=img_paths, dataset_dir=args.source, is_overwrite=False)\n\n # Perform Mask R-CNN processing\n mask_model = BatchMask(cfg=args.maskrcnn_cfg_file,\n batch_size=args.batch_size,\n workers=args.num_workers)\n mask_model(imagery=img_paths, dataset_dir=args.source, is_overwrite=False)\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "torchreid/scripts/main.py", "content": "import os\nimport argparse\nimport torch\nimport torch.nn as nn\nimport torchreid\nfrom torchreid.tools.extract_part_based_features import extract_reid_features\nfrom torchreid.data.masks_transforms import compute_parts_num_and_names\nfrom torchreid.utils import (\n Logger, check_isfile, set_random_seed, collect_env_info,\n resume_from_checkpoint, load_pretrained_weights, compute_model_complexity, Writer, load_checkpoint\n)\n\nfrom torchreid.scripts.default_config import (\n imagedata_kwargs, optimizer_kwargs, videodata_kwargs, engine_run_kwargs,\n get_default_config, lr_scheduler_kwargs, display_config_diff\n)\nfrom torchreid.utils.engine_state import EngineState\n\n\ndef build_datamanager(cfg):\n if cfg.data.type == 'image':\n return torchreid.data.ImageDataManager(**imagedata_kwargs(cfg))\n else:\n return torchreid.data.VideoDataManager(**videodata_kwargs(cfg))\n\n\ndef build_engine(cfg, datamanager, model, optimizer, scheduler, writer, engine_state):\n if cfg.data.type == 'image':\n if cfg.loss.name == 'softmax':\n engine = torchreid.engine.ImageSoftmaxEngine(\n datamanager,\n model,\n optimizer=optimizer,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer,\n engine_state=engine_state\n )\n\n elif cfg.loss.name == 'triplet':\n engine = torchreid.engine.ImageTripletEngine(\n datamanager,\n model,\n optimizer=optimizer,\n margin=cfg.loss.triplet.margin,\n weight_t=cfg.loss.triplet.weight_t,\n weight_x=cfg.loss.triplet.weight_x,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer,\n engine_state=engine_state\n )\n\n elif cfg.loss.name == 'part_based':\n engine = torchreid.engine.ImagePartBasedEngine(\n datamanager,\n model,\n optimizer=optimizer,\n loss_name=cfg.loss.part_based.name,\n config=cfg,\n margin=cfg.loss.triplet.margin,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer,\n engine_state=engine_state,\n dist_combine_strat=cfg.test.part_based.dist_combine_strat,\n batch_size_pairwise_dist_matrix=cfg.test.batch_size_pairwise_dist_matrix,\n mask_filtering_training=cfg.model.bpbreid.mask_filtering_training,\n mask_filtering_testing=cfg.model.bpbreid.mask_filtering_testing\n )\n\n else:\n if cfg.loss.name == 'softmax':\n engine = torchreid.engine.VideoSoftmaxEngine(\n datamanager,\n model,\n optimizer=optimizer,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n pooling_method=cfg.video.pooling_method,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer,\n engine_state=engine_state\n )\n\n else:\n engine = torchreid.engine.VideoTripletEngine(\n datamanager,\n model,\n optimizer=optimizer,\n margin=cfg.loss.triplet.margin,\n weight_t=cfg.loss.triplet.weight_t,\n weight_x=cfg.loss.triplet.weight_x,\n scheduler=scheduler,\n use_gpu=cfg.use_gpu,\n label_smooth=cfg.loss.softmax.label_smooth,\n save_model_flag=cfg.model.save_model_flag,\n writer=writer,\n engine_state=engine_state\n )\n\n return engine\n\n\ndef reset_config(cfg, args):\n if args.root:\n cfg.data.root = args.root\n if args.save_dir:\n cfg.data.save_dir = args.save_dir\n if args.inference_enabled:\n cfg.inference.enabled = args.inference_enabled\n if args.sources:\n cfg.data.sources = args.sources\n if args.targets:\n cfg.data.targets = args.targets\n if args.transforms:\n cfg.data.transforms = args.transforms\n if args.job_id:\n cfg.project.job_id = args.job_id\n\n\ndef main():\n parser = argparse.ArgumentParser(\n formatter_class=argparse.ArgumentDefaultsHelpFormatter\n )\n parser.add_argument(\n '--config-file', type=str, default='', help='path to config file'\n )\n parser.add_argument(\n '-s',\n '--sources',\n type=str,\n nargs='+',\n help='source datasets (delimited by space)'\n )\n parser.add_argument(\n '-t',\n '--targets',\n type=str,\n nargs='+',\n help='target datasets (delimited by space)'\n )\n parser.add_argument(\n '--transforms', type=str, nargs='+', help='data augmentation'\n )\n parser.add_argument(\n '--root', type=str, default='', help='path to data root'\n )\n parser.add_argument(\n '--save_dir', type=str, default='', help='path to output root dir'\n )\n parser.add_argument(\n 'opts',\n default=None,\n nargs=argparse.REMAINDER,\n help='Modify config options using the command-line'\n )\n parser.add_argument(\n '--job-id',\n type=int,\n default=None,\n help='Slurm job id'\n )\n parser.add_argument(\n '--inference-enabled',\n type=bool,\n default=False,\n )\n args = parser.parse_args()\n\n cfg = build_config(args, args.config_file)\n\n engine, model = build_torchreid_model_engine(cfg)\n print('Starting experiment {} with job id {} and creation date {}'.format(cfg.project.experiment_id,\n cfg.project.job_id,\n cfg.project.start_time))\n engine.run(**engine_run_kwargs(cfg))\n print(\n 'End of experiment {} with job id {} and creation date {}'.format(cfg.project.experiment_id, cfg.project.job_id,\n cfg.project.start_time))\n if cfg.inference.enabled:\n print(\"Starting inference on external data\")\n extract_reid_features(cfg, cfg.inference.input_folder, cfg.data.save_dir, model)\n\n\ndef build_config(args=None, config_file=None, config=None):\n cfg = get_default_config()\n default_cfg_copy = cfg.clone()\n cfg.use_gpu = torch.cuda.is_available()\n if config:\n cfg.merge_from_other_cfg(config)\n if config_file:\n cfg.merge_from_file(config_file)\n cfg.project.config_file = os.path.basename(config_file)\n if args is not None:\n reset_config(cfg, args)\n cfg.merge_from_list(args.opts)\n # set parts information (number of parts K and each part name),\n # depending on the original loaded masks size or the transformation applied:\n compute_parts_num_and_names(cfg)\n if cfg.model.load_weights and check_isfile(cfg.model.load_weights) and cfg.model.load_config:\n checkpoint = load_checkpoint(cfg.model.load_weights)\n if 'config' in checkpoint:\n print('Overwriting current config with config loaded from {}'.format(cfg.model.load_weights))\n bpbreid_config = checkpoint['config'].model.bpbreid\n if checkpoint['config'].data.sources[0] != cfg.data.targets[0]:\n print('WARNING: the train dataset of the loaded model is different from the target dataset in the '\n 'current config.')\n bpbreid_config.pop('hrnet_pretrained_path', None)\n bpbreid_config.masks.pop('dir', None)\n cfg.model.bpbreid.merge_from_other_cfg(bpbreid_config)\n else:\n print('Could not load config from file {}'.format(cfg.model.load_weights))\n display_config_diff(cfg, default_cfg_copy)\n cfg.data.save_dir = os.path.join(cfg.data.save_dir, str(cfg.project.job_id))\n os.makedirs(cfg.data.save_dir)\n return cfg\n\n\ndef build_torchreid_model_engine(cfg):\n if cfg.project.debug_mode:\n torch.autograd.set_detect_anomaly(True)\n logger = Logger(cfg)\n writer = Writer(cfg)\n set_random_seed(cfg.train.seed)\n print('Show configuration\\n{}\\n'.format(cfg))\n print('Collecting env info ...')\n print('** System info **\\n{}\\n'.format(collect_env_info()))\n if cfg.use_gpu:\n torch.backends.cudnn.benchmark = True\n datamanager = build_datamanager(cfg)\n engine_state = EngineState(cfg.train.start_epoch, cfg.train.max_epoch)\n writer.init_engine_state(engine_state, cfg.model.bpbreid.masks.parts_num)\n print('Building model: {}'.format(cfg.model.name))\n model = torchreid.models.build_model(\n name=cfg.model.name,\n num_classes=datamanager.num_train_pids,\n loss=cfg.loss.name,\n pretrained=cfg.model.pretrained,\n use_gpu=cfg.use_gpu,\n config=cfg\n )\n logger.add_model(model)\n num_params, flops = compute_model_complexity(\n model, cfg\n )\n print('Model complexity: params={:,} flops={:,}'.format(num_params, flops))\n if cfg.model.load_weights and check_isfile(cfg.model.load_weights):\n load_pretrained_weights(model, cfg.model.load_weights)\n if cfg.use_gpu:\n model = nn.DataParallel(model).cuda()\n optimizer = torchreid.optim.build_optimizer(model, **optimizer_kwargs(cfg))\n scheduler = torchreid.optim.build_lr_scheduler(\n optimizer, **lr_scheduler_kwargs(cfg)\n )\n if cfg.model.resume and check_isfile(cfg.model.resume):\n cfg.train.start_epoch = resume_from_checkpoint(\n cfg.model.resume, model, optimizer=optimizer, scheduler=scheduler\n )\n print(\n 'Building {}-engine for {}-reid'.format(cfg.loss.name, cfg.data.type)\n )\n engine = build_engine(cfg, datamanager, model, optimizer, scheduler, writer, engine_state)\n return engine, model\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "torchreid/tools/__init__.py", "content": "" }, { "path": "torchreid/tools/compute_mean_std.py", "content": "\"\"\"\nCompute channel-wise mean and standard deviation of a dataset.\n\nUsage:\n$ python compute_mean_std.py DATASET_ROOT DATASET_KEY\n\n- The first argument points to the root path where you put the datasets.\n- The second argument means the specific dataset key.\n\nFor instance, your datasets are put under $DATA and you wanna\ncompute the statistics of Market1501, do\n$ python compute_mean_std.py $DATA market1501\n\"\"\"\nimport argparse\n\nimport torchreid\n\n\ndef main():\n parser = argparse.ArgumentParser()\n parser.add_argument('root', type=str)\n parser.add_argument('sources', type=str)\n args = parser.parse_args()\n\n datamanager = torchreid.data.ImageDataManager(\n root=args.root,\n sources=args.sources,\n targets=None,\n height=256,\n width=128,\n batch_size_train=100,\n batch_size_test=100,\n transforms=None,\n norm_mean=[0., 0., 0.],\n norm_std=[1., 1., 1.],\n train_sampler='SequentialSampler'\n )\n train_loader = datamanager.train_loader\n\n print('Computing mean and std ...')\n mean = 0.\n std = 0.\n n_samples = 0.\n for data in train_loader:\n data = data[0]\n batch_size = data.size(0)\n data = data.view(batch_size, data.size(1), -1)\n mean += data.mean(2).sum(0)\n std += data.std(2).sum(0)\n n_samples += batch_size\n\n mean /= n_samples\n std /= n_samples\n print('Mean: {}'.format(mean))\n print('Std: {}'.format(std))\n\n\nif __name__ == '__main__':\n main()\n" }, { "path": "torchreid/tools/extract_part_based_features.py", "content": "import torch\nimport tqdm\nimport glob\nimport os\nimport numpy as np\nfrom torchreid.scripts.default_config import get_default_config, display_config_diff\nfrom torchreid.tools.feature_extractor import FeatureExtractor\n\n\ndef extract_part_based_features(extractor, image_list, batch_size=400):\n\n def chunks(lst, n):\n \"\"\"Yield successive n-sized chunks from lst.\"\"\"\n for i in range(0, len(lst), n):\n yield lst[i:i + n]\n\n all_embeddings = []\n all_visibility_scores = []\n all_masks = []\n\n images_chunks = chunks(image_list, batch_size)\n for chunk in tqdm.tqdm(images_chunks):\n embeddings, visibility_scores, masks = extractor(chunk)\n\n embeddings = embeddings.cpu().detach()\n visibility_scores = visibility_scores.cpu().detach()\n masks = masks.cpu().detach()\n\n all_embeddings.append(embeddings)\n all_visibility_scores.append(visibility_scores)\n all_masks.append(masks)\n\n all_embeddings = torch.cat(all_embeddings, 0).numpy()\n all_visibility_scores = torch.cat(all_visibility_scores, 0).numpy()\n all_masks = torch.cat(all_masks, 0).numpy()\n\n return {\n \"parts_embeddings\": all_embeddings,\n \"parts_visibility_scores\": all_visibility_scores,\n \"parts_masks\": all_masks,\n }\n\n\ndef extract_det_idx(img_path):\n return int(os.path.basename(img_path).split(\"_\")[0])\n\n\ndef extract_reid_features(cfg, base_folder, out_path, model=None, model_path=None, num_classes=None):\n extractor = FeatureExtractor(\n cfg,\n model_path=model_path,\n device='cuda' if torch.cuda.is_available() else 'cpu',\n num_classes=num_classes,\n model=model\n )\n\n print(\"Looking for video folders with images crops in {}\".format(base_folder))\n folder_list = glob.glob(base_folder + '/*')\n for folder in folder_list:\n image_list = glob.glob(os.path.join(folder, \"*.png\"))\n image_list.sort(key=extract_det_idx)\n print(\"{} images to process for folder {}\".format(len(image_list), folder))\n results = extract_part_based_features(extractor, image_list, batch_size=50)\n\n # dump to disk\n video_name = os.path.splitext(os.path.basename(folder))[0]\n parts_embeddings_filename = os.path.join(out_path, \"embeddings_\" + video_name + \".npy\")\n parts_visibility_scores_filanme = os.path.join(out_path, \"visibility_scores_\" + video_name + \".npy\")\n parts_masks_filename = os.path.join(out_path, \"masks_\" + video_name + \".npy\")\n\n os.makedirs(os.path.dirname(parts_embeddings_filename), exist_ok=True)\n os.makedirs(os.path.dirname(parts_visibility_scores_filanme), exist_ok=True)\n os.makedirs(os.path.dirname(parts_masks_filename), exist_ok=True)\n\n np.save(parts_embeddings_filename, results['parts_embeddings'])\n np.save(parts_visibility_scores_filanme, results['parts_visibility_scores'])\n np.save(parts_masks_filename, results['parts_masks'])\n\n print(\"features saved to {}\".format(out_path))\n" }, { "path": "torchreid/tools/feature_extractor.py", "content": "from __future__ import absolute_import\nimport numpy as np\nimport torch\nimport torchvision.transforms as T\nfrom PIL import Image\n\nfrom torchreid.utils import (\n check_isfile, load_pretrained_weights, compute_model_complexity\n)\nfrom torchreid.models import build_model\nfrom torchreid.data.transforms import build_transforms\n\n\nclass FeatureExtractor(object):\n \"\"\"A simple API for feature extraction.\n\n FeatureExtractor can be used like a python function, which\n accepts input of the following types:\n - a list of strings (image paths)\n - a list of numpy.ndarray each with shape (H, W, C)\n - a single string (image path)\n - a single numpy.ndarray with shape (H, W, C)\n - a torch.Tensor with shape (B, C, H, W) or (C, H, W)\n\n Returned is a torch tensor with shape (B, D) where D is the\n feature dimension.\n\n Args:\n model_name (str): model name.\n model_path (str): path to model weights.\n image_size (sequence or int): image height and width.\n pixel_mean (list): pixel mean for normalization.\n pixel_std (list): pixel std for normalization.\n pixel_norm (bool): whether to normalize pixels.\n device (str): 'cpu' or 'cuda' (could be specific gpu devices).\n verbose (bool): show model details.\n\n Examples::\n\n from torchreid.utils import FeatureExtractor\n\n extractor = FeatureExtractor(\n model_name='osnet_x1_0',\n model_path='a/b/c/model.pth.tar',\n device='cuda'\n )\n\n image_list = [\n 'a/b/c/image001.jpg',\n 'a/b/c/image002.jpg',\n 'a/b/c/image003.jpg',\n 'a/b/c/image004.jpg',\n 'a/b/c/image005.jpg'\n ]\n\n features = extractor(image_list)\n print(features.shape) # output (5, 512)\n \"\"\"\n\n def __init__(\n self,\n cfg,\n model_path='',\n image_size=(256, 128),\n pixel_mean=[0.485, 0.456, 0.406],\n pixel_std=[0.229, 0.224, 0.225],\n pixel_norm=True,\n device='cuda',\n num_classes=1,\n verbose=True,\n model=None\n ):\n # Build model\n if model is None:\n print(\"building model on device {}\".format(device))\n model = build_model(\n name=cfg.model.name,\n loss=cfg.loss.name,\n pretrained=cfg.model.pretrained,\n num_classes=num_classes,\n use_gpu=device.startswith('cuda'),\n pooling=cfg.model.bpbreid.pooling,\n normalization=cfg.model.bpbreid.normalization,\n last_stride=cfg.model.bpbreid.last_stride,\n config=cfg\n )\n\n model.eval()\n\n if verbose:\n num_params, flops = compute_model_complexity(\n model, cfg\n )\n print('Model: {}'.format(cfg.model.name))\n print('- params: {:,}'.format(num_params))\n print('- flops: {:,}'.format(flops))\n\n if model_path and check_isfile(model_path):\n load_pretrained_weights(model, model_path)\n\n # Build transform functions\n _, preprocess = build_transforms(image_size[0],\n image_size[1],\n cfg,\n transforms=None,\n norm_mean=pixel_mean,\n norm_std=pixel_std,\n masks_preprocess=cfg.model.bpbreid.masks.preprocess,\n softmax_weight=cfg.model.bpbreid.masks.softmax_weight,\n background_computation_strategy=cfg.model.bpbreid.masks.background_computation_strategy,\n mask_filtering_threshold=cfg.model.bpbreid.masks.mask_filtering_threshold,\n )\n\n to_pil = T.ToPILImage()\n\n device = torch.device(device)\n model.to(device)\n\n # Class attributes\n self.model = model\n self.preprocess = preprocess\n self.to_pil = to_pil\n self.device = device\n\n def __call__(self, input, external_parts_masks=None):\n if isinstance(input, list):\n images = []\n masks = []\n\n for i, element in enumerate(input):\n if isinstance(element, str):\n image = Image.open(element).convert('RGB')\n\n elif isinstance(element, np.ndarray):\n image = self.to_pil(element)\n\n else:\n raise TypeError(\n 'Type of each element must belong to [str | numpy.ndarray]'\n )\n\n transf_args = {}\n if external_parts_masks is not None:\n transf_args['mask'] = external_parts_masks[i].transpose(1, 2, 0)\n transf_args['image'] = np.array(image)\n result = self.preprocess(**transf_args)\n images.append(result['image'])\n if external_parts_masks is not None:\n masks.append(result['mask'])\n\n images = torch.stack(images, dim=0)\n images = images.to(self.device)\n if external_parts_masks is not None:\n masks = torch.stack(masks, dim=0)\n masks = masks.to(self.device)\n\n elif isinstance(input, str):\n image = Image.open(input).convert('RGB')\n transf_args = {}\n if external_parts_masks is not None:\n transf_args['mask'] = external_parts_masks.transpose(1, 2, 0)\n transf_args['image'] = np.array(image)\n result = self.preprocess(**transf_args)\n images = result['image'].unsqueeze(0).to(self.device)\n masks = result['mask'].unsqueeze(0).to(self.device)\n\n elif isinstance(input, np.ndarray):\n image = input\n transf_args = {}\n if external_parts_masks is not None:\n transf_args['mask'] = external_parts_masks.transpose(1, 2, 0)\n transf_args['image'] = np.array(image)\n result = self.preprocess(**transf_args)\n images = result['image'].unsqueeze(0).to(self.device)\n masks = result['mask'].unsqueeze(0).to(self.device)\n\n elif isinstance(input, torch.Tensor):\n if input.dim() == 3:\n input = input.unsqueeze(0)\n external_parts_masks = external_parts_masks.unsqueeze(0)\n images = input.to(self.device)\n\n masks = external_parts_masks.to(self.device)\n\n else:\n raise NotImplementedError\n\n with torch.no_grad():\n features = self.model(images, external_parts_masks=masks)\n\n return features\n" }, { "path": "torchreid/utils/__init__.py", "content": "from __future__ import absolute_import\n\nfrom .tools import *\nfrom .rerank import re_ranking\nfrom torchreid.utils.logging.deprecated_loggers import *\nfrom .avgmeter import *\nfrom .torchtools import *\nfrom torchreid.utils.logging.logger import *\nfrom .visualization import *\nfrom .reidtools import *\nfrom .writer import *\nfrom .model_complexity import compute_model_complexity\n" }, { "path": "torchreid/utils/avgmeter.py", "content": "from __future__ import division, absolute_import\n\nimport datetime\nimport warnings\nfrom collections import defaultdict, OrderedDict\nimport time\nimport torch\nimport numpy as np\n\n__all__ = ['AverageMeter', 'MetricMeter', 'TimeMeter', 'TorchTimeMeter', 'EpochMetricsMeter']\n\nfrom torchreid.utils.engine_state import EngineStateListener\n\n\nclass AverageMeter(object):\n \"\"\"Computes and stores the average and current value.\n\n Examples::\n >>> # Initialize a meter to record loss\n >>> losses = AverageMeter()\n >>> # Update meter after every minibatch update\n >>> losses.update(loss_value, batch_size)\n \"\"\"\n\n def __init__(self):\n self.reset()\n\n def reset(self):\n self.val = 0\n self.avg = 0\n self.sum = 0\n self.count = 0\n\n def update(self, val, n=1):\n self.val = val\n self.sum += val * n\n self.count += n\n self.avg = self.sum / self.count\n\n\nclass BatchMeter(object):\n def __init__(self, epoch_count, batch_count):\n self.epoch_count = epoch_count\n self.batch_count = batch_count\n self.reset()\n\n def reset(self):\n self.last_val = None\n self.values = np.zeros((self.epoch_count, self.batch_count))\n\n def update(self, epoch, batch, val):\n self.last_val = val\n self.values[epoch, batch] = val\n\n def total_for_epoch(self, epoch):\n self.values[epoch].sum()\n\n def avg_for_epoch(self, epoch):\n self.values[epoch].mean()\n\n def batch_avg(self):\n self.values.mean()\n\n def epoch_avg(self):\n self.values.mean(axis=1).mean()\n\n def total(self):\n self.values.sum()\n\n\nclass SingleMeter(EngineStateListener):\n def __init__(self, engine_state):\n self.engine_state = engine_state\n self.engine_state.add_listener(self)\n self.reset()\n self.is_empty = True\n\n def reset(self):\n self.total = 0\n self.val = 0\n\n def update(self, val, total):\n if torch.is_tensor(val):\n val = val.item()\n if torch.is_tensor(total):\n total = total.item()\n\n self.val = val\n self.total = total\n self.is_empty = False\n\n def ratio(self):\n if self.total == 0:\n return 0\n return self.val / self.total\n\n\nclass EpochMeter(EngineStateListener):\n # With RandomSample, number of batches might change from one epoch to another\n def __init__(self, engine_state):\n self.engine_state = engine_state\n self.engine_state.add_listener(self)\n self.min = np.zeros(self.engine_state.max_epoch)\n self.mean = np.zeros(self.engine_state.max_epoch)\n self.max = np.zeros(self.engine_state.max_epoch)\n self.batch_size = np.zeros(self.engine_state.max_epoch)\n self.sum = np.zeros(self.engine_state.max_epoch)\n self.total = np.zeros(self.engine_state.max_epoch)\n self.epoch_sum = []\n self.epoch_total = []\n self.is_empty = True\n\n def update(self, val, total=1.):\n self.is_empty = False\n if torch.is_tensor(val):\n val = val.item()\n if torch.is_tensor(total):\n total = total.item()\n\n self.epoch_sum.append(val)\n self.epoch_total.append(total)\n\n # Listeners\n def epoch_completed(self):\n if not self.is_empty:\n self.epoch_sum = np.array(self.epoch_sum)\n self.epoch_total = np.array(self.epoch_total)\n ratio = self.epoch_sum / self.epoch_total\n\n self.min[self.engine_state.epoch] = ratio.min()\n self.mean[self.engine_state.epoch] = ratio.mean()\n self.max[self.engine_state.epoch] = ratio.max()\n self.batch_size[self.engine_state.epoch] = len(self.epoch_sum)\n self.sum[self.engine_state.epoch] = self.epoch_sum.sum()\n self.total[self.engine_state.epoch] = self.epoch_total.sum()\n\n self.epoch_sum = []\n self.epoch_total = []\n self.is_empty = True\n\n # Utils\n def last_val(self):\n return self.epoch_sum[-1]\n\n def epoch_ratio(self, epoch):\n return self.mean[epoch]\n\n def total_ratio(self):\n return self.mean.mean()\n\n\nclass EpochArrayMeter(EngineStateListener):\n # With RandomSample, number of batches might change from one epoch to another\n def __init__(self, engine_state, array_size):\n self.engine_state = engine_state\n self.engine_state.add_listener(self)\n self.min = np.zeros((self.engine_state.max_epoch, array_size))\n self.mean = np.zeros((self.engine_state.max_epoch, array_size))\n self.max = np.zeros((self.engine_state.max_epoch, array_size))\n self.batch_size = np.zeros((self.engine_state.max_epoch, array_size))\n self.sum = np.zeros((self.engine_state.max_epoch, array_size))\n self.total = np.zeros((self.engine_state.max_epoch, array_size))\n self.epoch_sum = []\n self.epoch_total = []\n self.is_empty = True\n\n def update(self, val, total):\n self.is_empty = False\n\n if torch.is_tensor(val):\n if val.is_cuda:\n val = val.cpu()\n val = val.numpy()\n if torch.is_tensor(total):\n if val.is_cuda:\n val = val.cpu()\n val = val.numpy()\n\n self.epoch_sum.append(val)\n self.epoch_total.append(total)\n\n # Listeners\n def epoch_completed(self):\n if not self.is_empty:\n self.epoch_sum = np.array(self.epoch_sum)\n self.epoch_total = np.array(self.epoch_total)\n ratio = self.epoch_sum / self.epoch_total\n\n self.min[self.engine_state.epoch] = ratio.min(axis=0)\n self.mean[self.engine_state.epoch] = ratio.mean(axis=0)\n self.max[self.engine_state.epoch] = ratio.max(axis=0)\n self.batch_size[self.engine_state.epoch] = self.epoch_sum.shape[0]\n self.sum[self.engine_state.epoch] = self.epoch_sum.sum(axis=0)\n self.total[self.engine_state.epoch] = self.epoch_total.sum(axis=0)\n\n self.epoch_sum = []\n self.epoch_total = []\n\n # Utils\n def epoch_ratio(self, epoch):\n return self.mean[epoch]\n\n def total_ratio(self):\n return self.mean.mean(axis=0)\n\n\nclass TimeMeter(AverageMeter):\n \"\"\"Computes and stores the average time and current time value.\n \"\"\"\n\n def __init__(self, name):\n super().__init__()\n self.name = name\n self.tic = None\n\n def _format_time(self, time):\n return str(datetime.timedelta(milliseconds=round(time)))\n\n def total_time(self):\n return self._format_time(self.sum)\n\n def average_time(self):\n return self._format_time(self.avg)\n\n def start(self):\n self.tic = self._current_time_ms()\n\n def stop(self):\n if self.tic is not None:\n self.update(self._current_time_ms() - self.tic)\n self.tic = None\n return self.val\n else:\n warnings.warn(\"{0}.start() should be called before {0}.stop()\".format(self.__class__.__name__, RuntimeWarning))\n return 0\n\n @staticmethod\n def _current_time_ms():\n return time.time() * 1000\n\n\nclass TorchTimeMeter(TimeMeter):\n \"\"\"Computes and stores the average time and current time value.\n \"\"\"\n\n def __init__(self, name, plot=True):\n super().__init__(name)\n self.start_event = None\n self.end_event = None\n self.cuda = torch.cuda.is_available()\n self.plot = plot\n\n def start(self):\n if self.cuda:\n self._start_cuda()\n else:\n super().start()\n\n def _start_cuda(self):\n self.start_event = torch.cuda.Event(enable_timing=True)\n self.end_event = torch.cuda.Event(enable_timing=True)\n self.start_event.record()\n\n def stop(self):\n if self.cuda:\n return self._stop_cuda()\n else:\n return super().stop()\n\n def _stop_cuda(self):\n if self.start_event is not None:\n self.end_event.record()\n torch.cuda.synchronize() # TODO Check if slows down computation\n self.update(self.start_event.elapsed_time(self.end_event))\n self.start_event = None\n self.end_event = None\n return self.val\n else:\n warnings.warn(\"{0}.start() should be called before {0}.stop()\".format(self.__class__.__name__),\n RuntimeWarning)\n return 0\n\n\nclass EpochMetricsMeter(object):\n \"\"\"A collection of metrics.\n\n Source: https://github.com/KaiyangZhou/Dassl.pytorch\n\n Examples::\n >>> # 1. Create an instance of MetricMeter\n >>> metric = EpochMetricsMeter()\n >>> # 2. Update using a dictionary as input\n >>> input_dict = {'loss_1': value_1, 'loss_2': value_2}\n >>> metric.update(input_dict)\n >>> # 3. Convert to string and print\n >>> print(str(metric.summary(epoch)))\n \"\"\"\n\n def __init__(self, engine_state, delimiter='\\t'):\n self.engine_state = engine_state\n self.meters = {}\n self.delimiter = delimiter\n\n def update(self, input_dict):\n if input_dict is None:\n return\n\n if not isinstance(input_dict, dict):\n raise TypeError(\n 'Input to MetricMeter.update() must be a dictionary'\n )\n\n for k, v in input_dict.items():\n if isinstance(v, torch.Tensor):\n v = v.item()\n if k not in self.meters.keys():\n self.meters[k] = EpochMeter(self.engine_state)\n self.meters[k].update(v)\n\n def summary(self, epoch):\n output_str = []\n for name, meter in self.meters.items():\n output_str.append(\n '{} {:.3f} [{:.2f}, {:.2f}]'.format(name, meter.mean[epoch], meter.min[epoch], meter.max[epoch])\n )\n return self.delimiter.join(output_str)\n\nclass LossEpochMetricsMeter(object):\n def __init__(self, engine_state, delimiter='\\t'):\n self.engine_state = engine_state\n self.meters = OrderedDict()\n self.delimiter = delimiter\n\n def update(self, input_dict):\n if input_dict is None:\n return\n\n if not isinstance(input_dict, dict):\n raise TypeError(\n 'Input to MetricMeter.update() must be a dictionary'\n )\n\n for k1, v1 in input_dict.items():\n if k1 not in self.meters.keys():\n self.meters[k1] = OrderedDict()\n for k2, v2 in v1.items():\n if isinstance(v2, torch.Tensor):\n v2 = v2.item()\n if k2 not in self.meters[k1].keys():\n self.meters[k1][k2] = EpochMeter(self.engine_state)\n self.meters[k1][k2].update(v2)\n\n def summary(self, epoch):\n final_str = \"\"\n for name, dict in self.meters.items():\n if dict:\n output_str = [\"\\n\\t\" + name + \": \"]\n for key, meter in dict.items():\n output_str.append(\n '{} {:.3f} [{:.2f}, {:.2f}]'.format(key, meter.mean[epoch], meter.min[epoch], meter.max[epoch])\n )\n final_str += self.delimiter.join(output_str)\n return final_str\n\n\nclass MetricMeter(object):\n \"\"\"A collection of metrics.\n\n Source: https://github.com/KaiyangZhou/Dassl.pytorch\n\n Examples::\n >>> # 1. Create an instance of MetricMeter\n >>> metric = MetricMeter()\n >>> # 2. Update using a dictionary as input\n >>> input_dict = {'loss_1': value_1, 'loss_2': value_2}\n >>> metric.update(input_dict)\n >>> # 3. Convert to string and print\n >>> print(str(metric))\n \"\"\"\n\n def __init__(self, delimiter='\\t'):\n self.meters = defaultdict(AverageMeter)\n self.delimiter = delimiter\n\n def update(self, input_dict):\n if input_dict is None:\n return\n\n if not isinstance(input_dict, dict):\n raise TypeError(\n 'Input to MetricMeter.update() must be a dictionary'\n )\n\n for k, v in input_dict.items():\n if isinstance(v, torch.Tensor):\n v = v.item()\n self.meters[k].update(v)\n\n def __str__(self):\n output_str = []\n for name, meter in self.meters.items():\n output_str.append(\n '{} {:.4f} ({:.4f})'.format(name, meter.val, meter.avg)\n )\n return self.delimiter.join(output_str)\n" }, { "path": "torchreid/utils/constants.py", "content": "GLOBAL = 'globl'\nFOREGROUND = 'foreg'\nBACKGROUND = 'backg'\nCONCAT_PARTS = 'conct'\nPARTS = 'parts'\nBN_GLOBAL = 'bn_globl'\nBN_FOREGROUND = 'bn_foreg'\nBN_BACKGROUND = 'bn_backg'\nBN_CONCAT_PARTS = 'bn_conct'\nBN_PARTS = 'bn_parts'\nPIXELS = 'pixls'\n\nbn_correspondants = {\n BN_BACKGROUND: BACKGROUND,\n BN_GLOBAL: GLOBAL,\n BN_FOREGROUND: FOREGROUND,\n BN_CONCAT_PARTS: CONCAT_PARTS,\n BN_PARTS: PARTS,\n}\n\ndef get_test_embeddings_names(parts_names, test_embeddings):\n test_embeddings_names = []\n if GLOBAL in test_embeddings or BN_GLOBAL in test_embeddings:\n test_embeddings_names.append('global')\n # test_embeddings_names.append(('global', 'gb'))\n if FOREGROUND in test_embeddings or BN_FOREGROUND in test_embeddings:\n test_embeddings_names.append('foreground')\n # test_embeddings_names.append(('foreground', 'fg'))\n if CONCAT_PARTS in test_embeddings or BN_CONCAT_PARTS in test_embeddings:\n test_embeddings_names.append('concatenated')\n # test_embeddings_names.append(('concatenated', 'cc'))\n if PARTS in test_embeddings or BN_PARTS in test_embeddings:\n test_embeddings_names = test_embeddings_names + parts_names\n return test_embeddings_names" }, { "path": "torchreid/utils/distribution.py", "content": "import math\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom matplotlib.ticker import PercentFormatter\n\nfrom torchreid.utils import AverageMeter, Logger\nfrom torchreid.utils.engine_state import EngineState\n\n\ndef plot_body_parts_pairs_distance_distribution(body_part_pairwise_dist, q_pids, g_pids, tag):\n m = body_part_pairwise_dist.shape[0]\n cols = round(math.sqrt(m))\n rows = cols\n while rows * cols < m:\n rows += 1\n fig = plt.figure(figsize=(rows*5, cols*3))\n ssmd_meter = AverageMeter()\n for i in range(0, m):\n ax = fig.add_subplot(rows, cols, i+1)\n pos_p_mean, pos_p_std, neg_p_mean, neg_p_std, ssmd = compute_distance_distribution(\n ax, body_part_pairwise_dist[i], q_pids, g_pids, \"Bp {} pairs distance distribution\".format(i))\n ssmd_meter.update(ssmd)\n fig.tight_layout()\n Logger.current_logger().add_figure(\"{} body part pairs distance distribution\".format(tag), fig, EngineState.current_engine_state().epoch)\n return ssmd_meter.avg\n\n\ndef plot_pairs_distance_distribution(distmat, q_pids, g_pids, tag):\n fig, ax = plt.subplots()\n result = compute_distance_distribution(ax, distmat, q_pids, g_pids, \"{} pairs distance distribution\".format(tag))\n Logger.current_logger().add_figure(\"{} pairs distance distribution\".format(tag), fig, EngineState.current_engine_state().epoch)\n return result\n\n\ndef compute_distance_distribution(ax, distmat, q_pids, g_pids, title):\n pos_p = distmat[np.expand_dims(q_pids, axis=1) == np.expand_dims(g_pids, axis=0)]\n neg_p = distmat[np.expand_dims(q_pids, axis=1) != np.expand_dims(g_pids, axis=0)]\n\n pos_p_mean, pos_p_std, neg_p_mean, neg_p_std, ssmd = compute_ssmd(neg_p, pos_p)\n\n # plot_distributions(ax, neg_p, pos_p, pos_p_mean, pos_p_std, neg_p_mean, neg_p_std)\n # ax.set_title(title + \" - SSMD = {:.4f} \".format(ssmd))\n\n return pos_p_mean, pos_p_std, neg_p_mean, neg_p_std, ssmd\n\n\ndef compute_ssmd(neg_p, pos_p):\n pos_p_mean = np.mean(pos_p)\n pos_p_std = np.std(pos_p)\n neg_p_mean = np.mean(neg_p)\n neg_p_std = np.std(neg_p)\n ssmd = abs(pos_p_mean - neg_p_mean) / (pos_p_std ** 2 + neg_p_std ** 2)\n\n return pos_p_mean, pos_p_std, neg_p_mean, neg_p_std, ssmd\n\n\ndef plot_distributions(ax, neg_p, pos_p, pos_p_mean, pos_p_std, neg_p_mean, neg_p_std):\n bins = 100\n ax.hist(pos_p, weights=np.ones_like(pos_p)/len(pos_p), bins=bins, label='{} positive pairs : $\\mu={:10.3f}$, $\\sigma={:10.3f}$'.format(len(pos_p), pos_p_mean, pos_p_std), density=False, alpha=0.4, color='green')\n ax.hist(neg_p, weights=np.ones_like(neg_p)/len(neg_p), bins=bins, label='{} negative pairs : $\\mu={:10.3f}$, $\\sigma={:10.3f}$'.format(len(neg_p), neg_p_mean, neg_p_std), density=False, alpha=0.4, color='red')\n ax.axvline(x=pos_p_mean, linestyle='--', color='darkgreen')\n ax.axvline(x=neg_p_mean, linestyle='--', color='darkred')\n ax.set_xlabel(\"pairs distance\")\n ax.set_ylabel(\"pairs count\")\n ax.legend()\n ax.yaxis.set_major_formatter(PercentFormatter(xmax=1))\n" }, { "path": "torchreid/utils/engine_state.py", "content": "import queue\nfrom heapq import heappush\n\n\nclass EngineStateListener:\n def batch_completed(self):\n pass\n\n def epoch_started(self):\n pass\n\n def epoch_completed(self):\n pass\n\n def training_started(self):\n pass\n\n def training_completed(self):\n pass\n\n def test_completed(self):\n pass\n\n def run_completed(self):\n pass\n\n\nclass EngineState(EngineStateListener):\n __main_engine_state = None # type: Optional[EngineState]\n\n @classmethod\n def current_engine_state(cls):\n # type: () -> EngineState\n return cls.__main_engine_state\n\n def __init__(self, start_epoch, max_epoch):\n self.start_epoch = start_epoch\n self.max_epoch = max_epoch\n self.epoch = start_epoch\n self.batch = 0\n self.global_step = 0\n self.estimated_num_batches = 0\n self.lr = 0\n self.is_training = False\n self.listeners = []\n self.last_listeners = []\n\n EngineState.__main_engine_state = self\n\n def add_listener(self, listener, last=False):\n # FIXME ugly\n if last:\n self.last_listeners.append(listener)\n else:\n self.listeners.append(listener)\n\n def batch_completed(self):\n for listener in self.listeners + self.last_listeners:\n listener.batch_completed()\n self.batch += 1\n self.global_step += 1\n\n def epoch_started(self):\n for listener in self.listeners + self.last_listeners:\n listener.epoch_started()\n self.batch = 0\n\n def epoch_completed(self):\n for listener in self.listeners + self.last_listeners:\n listener.epoch_completed()\n if self.epoch != self.max_epoch - 1:\n self.epoch += 1\n\n def training_started(self):\n for listener in self.listeners + self.last_listeners:\n listener.training_started()\n self.is_training = True\n\n def training_completed(self):\n for listener in self.listeners + self.last_listeners:\n listener.training_completed()\n self.is_training = False\n\n def test_completed(self):\n for listener in self.listeners + self.last_listeners:\n listener.test_completed()\n\n def run_completed(self):\n for listener in self.listeners + self.last_listeners:\n listener.run_completed()\n\n def update_lr(self, lr):\n self.lr = lr" }, { "path": "torchreid/utils/imagetools.py", "content": "import numpy as np\nfrom scipy import signal\n\n\ndef gkern(kernlen=21, std=None):\n \"\"\"Returns a 2D Gaussian kernel array.\"\"\"\n if std is None:\n std = kernlen / 4\n gkern1d = signal.gaussian(kernlen, std=std).reshape(kernlen, 1)\n gkern2d = np.outer(gkern1d, gkern1d)\n return gkern2d\n\n\ndef build_gaussian_heatmaps(kp_xyc, w, h, gaussian=None):\n gaussian_heatmaps = np.zeros((len(kp_xyc), h, w))\n for i, kp in enumerate(kp_xyc):\n # do not use invisible keypoints\n if kp[2] == 0:\n continue\n\n kpx, kpy = kp[:2].astype(int)\n\n if gaussian is None:\n g_scale = 6\n g_radius = int(w / g_scale)\n gaussian = gkern(g_radius * 2 + 1)\n else:\n g_radius = gaussian.shape[0] // 2\n\n rt, rb = min(g_radius, kpy), min(g_radius, h - 1 - kpy)\n rl, rr = min(g_radius, kpx), min(g_radius, w - 1 - kpx)\n\n gaussian_heatmaps[i, kpy - rt:kpy + rb + 1, kpx - rl:kpx + rr + 1] = gaussian[\n g_radius - rt:g_radius + rb + 1,\n g_radius - rl:g_radius + rr + 1]\n return gaussian_heatmaps\n" }, { "path": "torchreid/utils/logging/__init__.py", "content": "from .logger import Logger\n" }, { "path": "torchreid/utils/logging/deprecated_loggers.py", "content": "from __future__ import absolute_import\nimport os\nimport sys\nimport os.path as osp\n\nfrom torchreid.utils.tools import mkdir_if_missing\n\n__all__ = ['StdoutLogger', 'RankLogger']\n\n\nclass StdoutLogger(object):\n \"\"\"Writes console output to external text file.\n\n Imported from ``_\n\n Args:\n fpath (str): directory to save logging file.\n\n Examples::\n >>> import sys\n >>> import os\n >>> import os.path as osp\n >>> from torchreid.utils import Logger\n >>> save_dir = 'log/resnet50-softmax-market1501'\n >>> log_name = 'train.log'\n >>> sys.stdout = Logger(osp.join(args.save_dir, log_name))\n \"\"\"\n\n def __init__(self, fpath=None):\n self.console = sys.stdout\n self.file = None\n if fpath is not None:\n mkdir_if_missing(osp.dirname(fpath))\n self.file = open(fpath, 'w')\n\n def __del__(self):\n self.close()\n\n def __enter__(self):\n pass\n\n def __exit__(self, *args):\n self.close()\n\n def write(self, msg):\n self.console.write(msg)\n if self.file is not None:\n self.file.write(msg)\n\n def flush(self):\n self.console.flush()\n if self.file is not None:\n self.file.flush()\n os.fsync(self.file.fileno())\n\n def close(self):\n self.console.close()\n if self.file is not None:\n self.file.close()\n\n\nclass RankLogger(object):\n \"\"\"Records the rank1 matching accuracy obtained for each\n test dataset at specified evaluation steps and provides a function\n to show the summarized results, which are convenient for analysis.\n\n Args:\n sources (str or list): source dataset name(s).\n targets (str or list): target dataset name(s).\n\n Examples::\n >>> from torchreid.utils import RankLogger\n >>> s = 'market1501'\n >>> t = 'market1501'\n >>> ranklogger = RankLogger(s, t)\n >>> ranklogger.write(t, 10, 0.5)\n >>> ranklogger.write(t, 20, 0.7)\n >>> ranklogger.write(t, 30, 0.9)\n >>> ranklogger.show_summary()\n >>> # You will see:\n >>> # => Show performance summary\n >>> # market1501 (source)\n >>> # - epoch 10 rank1 50.0%\n >>> # - epoch 20 rank1 70.0%\n >>> # - epoch 30 rank1 90.0%\n >>> # If there are multiple test datasets\n >>> t = ['market1501', 'dukemtmcreid']\n >>> ranklogger = RankLogger(s, t)\n >>> ranklogger.write(t[0], 10, 0.5)\n >>> ranklogger.write(t[0], 20, 0.7)\n >>> ranklogger.write(t[0], 30, 0.9)\n >>> ranklogger.write(t[1], 10, 0.1)\n >>> ranklogger.write(t[1], 20, 0.2)\n >>> ranklogger.write(t[1], 30, 0.3)\n >>> ranklogger.show_summary()\n >>> # You can see:\n >>> # => Show performance summary\n >>> # market1501 (source)\n >>> # - epoch 10 rank1 50.0%\n >>> # - epoch 20 rank1 70.0%\n >>> # - epoch 30 rank1 90.0%\n >>> # dukemtmcreid (target)\n >>> # - epoch 10 rank1 10.0%\n >>> # - epoch 20 rank1 20.0%\n >>> # - epoch 30 rank1 30.0%\n \"\"\"\n\n def __init__(self, sources, targets):\n self.sources = sources\n self.targets = targets\n\n if isinstance(self.sources, str):\n self.sources = [self.sources]\n\n if isinstance(self.targets, str):\n self.targets = [self.targets]\n\n self.logger = {\n name: {\n 'epoch': [],\n 'rank1': []\n }\n for name in self.targets\n }\n\n def write(self, name, epoch, rank1):\n \"\"\"Writes result.\n\n Args:\n name (str): dataset name.\n epoch (int): current epoch.\n rank1 (float): rank1 result.\n \"\"\"\n self.logger[name]['epoch'].append(epoch)\n self.logger[name]['rank1'].append(rank1)\n\n def show_summary(self):\n \"\"\"Shows saved results.\"\"\"\n print('=> Show performance summary')\n for name in self.targets:\n from_where = 'source' if name in self.sources else 'target'\n print('{} ({})'.format(name, from_where))\n for epoch, rank1 in zip(\n self.logger[name]['epoch'], self.logger[name]['rank1']\n ):\n print('- epoch {}\\t rank1 {:.1%}'.format(epoch, rank1))\n" }, { "path": "torchreid/utils/logging/logger.py", "content": "import os\nimport cv2\nimport wandb\nimport matplotlib.pyplot as plt\nfrom typing import Optional\nfrom pandas.io.json._normalize import nested_to_record\nfrom torch.utils.tensorboard import SummaryWriter\n\n\nclass Logger:\n \"\"\" A class to encapsulate external loggers such as Tensorboard, Allegro ClearML, Neptune, Weight and Biases, \n Comet, ...\n \"\"\"\n __main_logger = None # type: Optional[Logger]\n\n @classmethod\n def current_logger(cls):\n # type: () -> Logger\n return cls.__main_logger\n\n def __init__(self, cfg):\n # self.cfg = cfg\n # self.model_name = cfg.project.start_time + cfg.project.experiment_id\n\n # configs\n self.save_disk = cfg.project.logger.save_disk\n self.save_dir = cfg.data.save_dir\n self.matplotlib_show = cfg.project.logger.matplotlib_show\n\n # init external loggers\n self.tensorboard_logger = None\n if cfg.project.logger.use_tensorboard:\n self.tensorboard_folder = os.path.join(cfg.data.save_dir, 'tensorboard')\n self.tensorboard_logger = SummaryWriter(self.tensorboard_folder)\n\n\n self.use_wandb = cfg.project.logger.use_wandb\n if self.use_wandb:\n # os.environ[\"WANDB_SILENT\"] = \"true\"\n # wandb.init(config=cfg, sync_tensorboard=True, project=cfg.project.name, dir=cfg.data.save_dir, reinit=False)\n if cfg.project.logger.use_tensorboard:\n wandb.tensorboard.patch(pytorch=True, save=True, root_logdir=self.tensorboard_folder)\n wandb.init(config=cfg,\n project=cfg.project.name,\n dir=cfg.data.save_dir,\n reinit=False,\n name=str(cfg.project.job_id),\n notes=cfg.project.notes,\n tags=cfg.project.tags\n )\n # wandb.tensorboard.patch(save=True, tensorboardX=False)\n\n Logger.__main_logger = self\n\n def add_model(self, model):\n if self.use_wandb and wandb.run is not None:\n wandb.watch(model)\n\n def add_text(self, tag, value):\n if self.use_wandb and wandb.run is not None:\n wandb.log({tag: value})\n\n def add_scalar(self, tag, scalar_value, step):\n if self.tensorboard_logger is not None:\n self.tensorboard_logger.add_scalar(tag, scalar_value, step)\n if self.use_wandb and wandb.run is not None:\n wandb.log({tag: scalar_value})\n\n def add_figure(self, tag, figure, step):\n if self.matplotlib_show:\n figure.show()\n plt.waitforbuttonpress()\n if self.tensorboard_logger is not None:\n self.tensorboard_logger.add_figure(tag, figure, step)\n if self.use_wandb and wandb.run is not None:\n wandb.log({tag: wandb.Image(\n figure)}) # FIXME cannot give \"figure\" directly: Invalid value of type 'builtins.str' received for the 'color' property of scatter.marker Received value: 'none'\n if self.save_disk:\n figure_path = os.path.join(self.save_dir, 'figures', tag + '.png')\n os.makedirs(os.path.dirname(figure_path), exist_ok=True)\n plt.savefig(figure_path)\n plt.close(figure)\n\n def add_image(self, group, name, image, step):\n \"\"\"Input image must be in RGB format\"\"\"\n # if self.tensorboard_logger is not None:\n # self.tensorboard_logger.add_figure(tag, figure, self.global_step())\n if self.use_wandb and wandb.run is not None:\n wandb.log({group + name: wandb.Image(image)})\n if self.save_disk:\n image_path = os.path.join(self.save_dir, 'images', f\"{group}_{name}.jpg\")\n os.makedirs(os.path.dirname(image_path), exist_ok=True)\n image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n cv2.imwrite(image_path, image)\n\n def add_embeddings(self, tag, embeddings, labels, imgs, step):\n if self.tensorboard_logger is not None:\n self.tensorboard_logger.add_embedding(embeddings,\n metadata=labels,\n label_img=imgs,\n global_step=step,\n tag=tag,\n metadata_header=None)\n\n def close(self):\n if self.tensorboard_logger is not None:\n self.tensorboard_logger.close()\n if self.use_wandb and wandb.run is not None:\n wandb.finish()\n" }, { "path": "torchreid/utils/model_complexity.py", "content": "from __future__ import division, print_function, absolute_import\nimport math\nimport numpy as np\nfrom itertools import repeat\nfrom collections import namedtuple, defaultdict\nimport torch\n\n__all__ = ['compute_model_complexity']\n\nfrom torchreid.models.bpbreid import BPBreID\n\n\"\"\"\nUtility\n\"\"\"\n\n\ndef _ntuple(n):\n\n def parse(x):\n if isinstance(x, int):\n return tuple(repeat(x, n))\n return x\n\n return parse\n\n\n_single = _ntuple(1)\n_pair = _ntuple(2)\n_triple = _ntuple(3)\n\"\"\"\nConvolution\n\"\"\"\n\n\ndef hook_convNd(m, x, y):\n k = torch.prod(torch.Tensor(m.kernel_size)).item()\n cin = m.in_channels\n flops_per_ele = k * cin # + (k*cin-1)\n if m.bias is not None:\n flops_per_ele += 1\n flops = flops_per_ele * y.numel() / m.groups\n return int(flops)\n\n\n\"\"\"\nPooling\n\"\"\"\n\n\ndef hook_maxpool1d(m, x, y):\n flops_per_ele = m.kernel_size - 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_maxpool2d(m, x, y):\n k = _pair(m.kernel_size)\n k = torch.prod(torch.Tensor(k)).item()\n # ops: compare\n flops_per_ele = k - 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_maxpool3d(m, x, y):\n k = _triple(m.kernel_size)\n k = torch.prod(torch.Tensor(k)).item()\n flops_per_ele = k - 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_avgpool1d(m, x, y):\n flops_per_ele = m.kernel_size\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_avgpool2d(m, x, y):\n k = _pair(m.kernel_size)\n k = torch.prod(torch.Tensor(k)).item()\n flops_per_ele = k\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_avgpool3d(m, x, y):\n k = _triple(m.kernel_size)\n k = torch.prod(torch.Tensor(k)).item()\n flops_per_ele = k\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_adapmaxpool1d(m, x, y):\n x = x[0]\n out_size = m.output_size\n k = math.ceil(x.size(2) / out_size)\n flops_per_ele = k - 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_adapmaxpool2d(m, x, y):\n x = x[0]\n out_size = _pair(m.output_size)\n k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)\n k = torch.prod(torch.ceil(k)).item()\n flops_per_ele = k - 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_adapmaxpool3d(m, x, y):\n x = x[0]\n out_size = _triple(m.output_size)\n k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)\n k = torch.prod(torch.ceil(k)).item()\n flops_per_ele = k - 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_adapavgpool1d(m, x, y):\n x = x[0]\n out_size = m.output_size\n k = math.ceil(x.size(2) / out_size)\n flops_per_ele = k\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_adapavgpool2d(m, x, y):\n x = x[0]\n out_size = _pair(m.output_size)\n k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)\n k = torch.prod(torch.ceil(k)).item()\n flops_per_ele = k\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\ndef hook_adapavgpool3d(m, x, y):\n x = x[0]\n out_size = _triple(m.output_size)\n k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)\n k = torch.prod(torch.ceil(k)).item()\n flops_per_ele = k\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\n\"\"\"\nNon-linear activations\n\"\"\"\n\n\ndef hook_relu(m, x, y):\n # eq: max(0, x)\n num_ele = y.numel()\n return int(num_ele)\n\n\ndef hook_leakyrelu(m, x, y):\n # eq: max(0, x) + negative_slope*min(0, x)\n num_ele = y.numel()\n flops = 3 * num_ele\n return int(flops)\n\n\n\"\"\"\nNormalization\n\"\"\"\n\n\ndef hook_batchnormNd(m, x, y):\n num_ele = y.numel()\n flops = 2 * num_ele # mean and std\n if m.affine:\n flops += 2 * num_ele # gamma and beta\n return int(flops)\n\n\ndef hook_instancenormNd(m, x, y):\n return hook_batchnormNd(m, x, y)\n\n\ndef hook_groupnorm(m, x, y):\n return hook_batchnormNd(m, x, y)\n\n\ndef hook_layernorm(m, x, y):\n num_ele = y.numel()\n flops = 2 * num_ele # mean and std\n if m.elementwise_affine:\n flops += 2 * num_ele # gamma and beta\n return int(flops)\n\n\n\"\"\"\nLinear\n\"\"\"\n\n\ndef hook_linear(m, x, y):\n flops_per_ele = m.in_features # + (m.in_features-1)\n if m.bias is not None:\n flops_per_ele += 1\n flops = flops_per_ele * y.numel()\n return int(flops)\n\n\n__generic_flops_counter = {\n # Convolution\n 'Conv1d': hook_convNd,\n 'Conv2d': hook_convNd,\n 'Conv3d': hook_convNd,\n # Pooling\n 'MaxPool1d': hook_maxpool1d,\n 'MaxPool2d': hook_maxpool2d,\n 'MaxPool3d': hook_maxpool3d,\n 'AvgPool1d': hook_avgpool1d,\n 'AvgPool2d': hook_avgpool2d,\n 'AvgPool3d': hook_avgpool3d,\n 'AdaptiveMaxPool1d': hook_adapmaxpool1d,\n 'AdaptiveMaxPool2d': hook_adapmaxpool2d,\n 'AdaptiveMaxPool3d': hook_adapmaxpool3d,\n 'AdaptiveAvgPool1d': hook_adapavgpool1d,\n 'AdaptiveAvgPool2d': hook_adapavgpool2d,\n 'AdaptiveAvgPool3d': hook_adapavgpool3d,\n # Non-linear activations\n 'ReLU': hook_relu,\n 'ReLU6': hook_relu,\n 'LeakyReLU': hook_leakyrelu,\n # Normalization\n 'BatchNorm1d': hook_batchnormNd,\n 'BatchNorm2d': hook_batchnormNd,\n 'BatchNorm3d': hook_batchnormNd,\n 'InstanceNorm1d': hook_instancenormNd,\n 'InstanceNorm2d': hook_instancenormNd,\n 'InstanceNorm3d': hook_instancenormNd,\n 'GroupNorm': hook_groupnorm,\n 'LayerNorm': hook_layernorm,\n # Linear\n 'Linear': hook_linear,\n}\n\n__conv_linear_flops_counter = {\n # Convolution\n 'Conv1d': hook_convNd,\n 'Conv2d': hook_convNd,\n 'Conv3d': hook_convNd,\n # Linear\n 'Linear': hook_linear,\n}\n\n\ndef _get_flops_counter(only_conv_linear):\n if only_conv_linear:\n return __conv_linear_flops_counter\n return __generic_flops_counter\n\n\ndef compute_model_complexity(\n model, cfg, verbose=False, only_conv_linear=True\n):\n \"\"\"Returns number of parameters and FLOPs.\n\n .. note::\n (1) this function only provides an estimate of the theoretical time complexity\n rather than the actual running time which depends on implementations and hardware,\n and (2) the FLOPs is only counted for layers that are used at test time. This means\n that redundant layers such as person ID classification layer will be ignored as it\n is discarded when doing feature extraction. Note that the inference graph depends on\n how you construct the computations in ``forward()``.\n\n Args:\n model (nn.Module): network model.\n verbose (bool, optional): shows detailed complexity of\n each module. Default is False.\n only_conv_linear (bool, optional): only considers convolution\n and linear layers when counting flops. Default is True.\n If set to False, flops of all layers will be counted.\n\n Examples::\n >>> from torchreid import models, utils\n >>> model = models.build_model(name='resnet50', num_classes=1000)\n >>> num_params, flops = utils.compute_model_complexity(model, (1, 3, 256, 128), verbose=True)\n \"\"\"\n registered_handles = []\n layer_list = []\n layer = namedtuple('layer', ['class_name', 'params', 'flops'])\n\n def _add_hooks(m):\n\n def _has_submodule(m):\n return len(list(m.children())) > 0\n\n def _hook(m, x, y):\n params = sum(p.numel() for p in m.parameters())\n class_name = str(m.__class__.__name__)\n flops_counter = _get_flops_counter(only_conv_linear)\n if class_name in flops_counter:\n flops = flops_counter[class_name](m, x, y)\n else:\n flops = 0\n layer_list.append(\n layer(class_name=class_name, params=params, flops=flops)\n )\n\n # only consider the very basic nn layer\n if _has_submodule(m):\n return\n\n handle = m.register_forward_hook(_hook)\n registered_handles.append(handle)\n\n default_train_mode = model.training\n\n model.eval().apply(_add_hooks)\n input_size = (1, 3, cfg.data.height, cfg.data.width)\n input_img = torch.rand(input_size)\n if next(model.parameters()).is_cuda:\n input_img = input_img.cuda()\n if isinstance(model, BPBreID):\n model(input_img, torch.ones(1, cfg.model.bpbreid.masks.parts_num+1, 16, 8)) # forward\n else:\n model(input_img) # forward\n for handle in registered_handles:\n handle.remove()\n\n model.train(default_train_mode)\n\n if verbose:\n per_module_params = defaultdict(list)\n per_module_flops = defaultdict(list)\n\n total_params, total_flops = 0, 0\n\n for layer in layer_list:\n total_params += layer.params\n total_flops += layer.flops\n if verbose:\n per_module_params[layer.class_name].append(layer.params)\n per_module_flops[layer.class_name].append(layer.flops)\n\n if verbose:\n num_udscore = 55\n print(' {}'.format('-' * num_udscore))\n print(' Model complexity with input size {}'.format(input_size))\n print(' {}'.format('-' * num_udscore))\n for class_name in per_module_params:\n params = int(np.sum(per_module_params[class_name]))\n flops = int(np.sum(per_module_flops[class_name]))\n print(\n ' {} (params={:,}, flops={:,})'.format(\n class_name, params, flops\n )\n )\n print(' {}'.format('-' * num_udscore))\n print(\n ' Total (params={:,}, flops={:,})'.format(\n total_params, total_flops\n )\n )\n print(' {}'.format('-' * num_udscore))\n\n return total_params, total_flops\n" }, { "path": "torchreid/utils/reidtools.py", "content": "from __future__ import print_function, absolute_import\nimport numpy as np\nimport shutil\nimport os.path as osp\nimport cv2\nfrom torchreid.utils import Logger\nfrom .engine_state import EngineState\nfrom .tools import mkdir_if_missing\n\n__all__ = ['visualize_ranked_results']\n\nGRID_SPACING = 10\nQUERY_EXTRA_SPACING = 90\nBW = 5 # border width\nGREEN = (0, 255, 0)\nRED = (0, 0, 255)\n\n\ndef visualize_ranked_results(\n distmat, dataset, data_type, width=128, height=256, save_dir='', topk=10\n):\n \"\"\"Visualizes ranked results.\n\n Supports both image-reid and video-reid.\n\n For image-reid, ranks will be plotted in a single figure. For video-reid, ranks will be\n saved in folders each containing a tracklet.\n\n Args:\n distmat (numpy.ndarray): distance matrix of shape (num_query, num_gallery).\n dataset (tuple): a 2-tuple containing (query, gallery), each of which contains\n tuples of (img_path(s), pid, camid).\n data_type (str): \"image\" or \"video\".\n width (int, optional): resized image width. Default is 128.\n height (int, optional): resized image height. Default is 256.\n save_dir (str): directory to save output images.\n topk (int, optional): denoting top-k images in the rank list to be visualized.\n Default is 10.\n \"\"\"\n num_q, num_g = distmat.shape\n mkdir_if_missing(save_dir)\n\n print('# query: {}\\n# gallery {}'.format(num_q, num_g))\n print('Visualizing top-{} ranks ...'.format(topk))\n\n query, gallery = dataset\n assert num_q == len(query)\n assert num_g == len(gallery)\n\n indices = np.argsort(distmat, axis=1)\n\n def _cp_img_to(src, dst, rank, prefix, matched=False):\n \"\"\"\n Args:\n src: image path or tuple (for vidreid)\n dst: target directory\n rank: int, denoting ranked position, starting from 1\n prefix: string\n matched: bool\n \"\"\"\n if isinstance(src, (tuple, list)):\n if prefix == 'gallery':\n suffix = 'TRUE' if matched else 'FALSE'\n dst = osp.join(\n dst, prefix + '_top' + str(rank).zfill(3)\n ) + '_' + suffix\n else:\n dst = osp.join(dst, prefix + '_top' + str(rank).zfill(3))\n mkdir_if_missing(dst)\n for img_path in src:\n shutil.copy(img_path, dst)\n else:\n dst = osp.join(\n dst, prefix + '_top' + str(rank).zfill(3) + '_name_' +\n osp.basename(src)\n )\n shutil.copy(src, dst)\n\n for q_idx in range(num_q):\n qpid, qcamid, qimg_path = query[q_idx]['pid'], query[q_idx]['camid'], query[q_idx]['img_path']\n qimg_path_name = qimg_path[0] if isinstance(\n qimg_path, (tuple, list)\n ) else qimg_path\n\n if data_type == 'image':\n qimg = cv2.imread(qimg_path)\n qimg = cv2.resize(qimg, (width, height))\n qimg = cv2.copyMakeBorder(\n qimg, BW, BW, BW, BW, cv2.BORDER_CONSTANT, value=(0, 0, 0)\n )\n # resize twice to ensure that the border width is consistent across images\n qimg = cv2.resize(qimg, (width, height))\n num_cols = topk + 1\n grid_img = 255 * np.ones(\n (\n height,\n num_cols*width + topk*GRID_SPACING + QUERY_EXTRA_SPACING, 3\n ),\n dtype=np.uint8\n )\n grid_img[:, :width, :] = qimg\n else:\n qdir = osp.join(\n save_dir, osp.basename(osp.splitext(qimg_path_name)[0])\n )\n mkdir_if_missing(qdir)\n _cp_img_to(qimg_path, qdir, rank=0, prefix='query')\n\n rank_idx = 1\n for g_idx in indices[q_idx, :]:\n gpid, gcamid, gimg_path = gallery[g_idx]['pid'], gallery[g_idx]['camid'], gallery[g_idx]['img_path']\n invalid = (qpid == gpid) & (qcamid == gcamid)\n\n if not invalid:\n matched = gpid == qpid\n if data_type == 'image':\n border_color = GREEN if matched else RED\n gimg = cv2.imread(gimg_path)\n gimg = cv2.resize(gimg, (width, height))\n gimg = cv2.copyMakeBorder(\n gimg,\n BW,\n BW,\n BW,\n BW,\n cv2.BORDER_CONSTANT,\n value=border_color\n )\n gimg = cv2.resize(gimg, (width, height))\n start = rank_idx*width + rank_idx*GRID_SPACING + QUERY_EXTRA_SPACING\n end = (\n rank_idx+1\n ) * width + rank_idx*GRID_SPACING + QUERY_EXTRA_SPACING\n grid_img[:, start:end, :] = gimg\n else:\n _cp_img_to(\n gimg_path,\n qdir,\n rank=rank_idx,\n prefix='gallery',\n matched=matched\n )\n\n rank_idx += 1\n if rank_idx > topk:\n break\n\n if data_type == 'image':\n imname = osp.basename(osp.splitext(qimg_path_name)[0])\n # cv2.imwrite(osp.join(save_dir, imname + '.jpg'), grid_img)\n Logger.current_logger().add_image(\"Ranking grid\", imname + '.jpg', cv2.cvtColor(grid_img, cv2.COLOR_BGR2RGB),\n EngineState.current_engine_state().epoch)\n\n if (q_idx+1) % 100 == 0:\n print('- done {}/{}'.format(q_idx + 1, num_q))\n\n print('Done. Images have been saved to \"{}\" ...'.format(save_dir))\n" }, { "path": "torchreid/utils/rerank.py", "content": "#!/usr/bin/env python2/python3\n# -*- coding: utf-8 -*-\n\"\"\"\nSource: https://github.com/zhunzhong07/person-re-ranking\n\nCreated on Mon Jun 26 14:46:56 2017\n@author: luohao\nModified by Houjing Huang, 2017-12-22.\n- This version accepts distance matrix instead of raw features.\n- The difference of `/` division between python 2 and 3 is handled.\n- numpy.float16 is replaced by numpy.float32 for numerical precision.\n\nCVPR2017 paper:Zhong Z, Zheng L, Cao D, et al. Re-ranking Person Re-identification with k-reciprocal Encoding[J]. 2017.\nurl:http://openaccess.thecvf.com/content_cvpr_2017/papers/Zhong_Re-Ranking_Person_Re-Identification_CVPR_2017_paper.pdf\nMatlab version: https://github.com/zhunzhong07/person-re-ranking\n\nAPI\nq_g_dist: query-gallery distance matrix, numpy array, shape [num_query, num_gallery]\nq_q_dist: query-query distance matrix, numpy array, shape [num_query, num_query]\ng_g_dist: gallery-gallery distance matrix, numpy array, shape [num_gallery, num_gallery]\nk1, k2, lambda_value: parameters, the original paper is (k1=20, k2=6, lambda_value=0.3)\nReturns:\n final_dist: re-ranked distance, numpy array, shape [num_query, num_gallery]\n\"\"\"\nfrom __future__ import division, print_function, absolute_import\nimport numpy as np\n\n__all__ = ['re_ranking']\n\n\ndef re_ranking(q_g_dist, q_q_dist, g_g_dist, k1=20, k2=6, lambda_value=0.3):\n\n # The following naming, e.g. gallery_num, is different from outer scope.\n # Don't care about it.\n\n original_dist = np.concatenate(\n [\n np.concatenate([q_q_dist, q_g_dist], axis=1),\n np.concatenate([q_g_dist.T, g_g_dist], axis=1)\n ],\n axis=0\n )\n original_dist = np.power(original_dist, 2).astype(np.float32)\n original_dist = np.transpose(\n 1. * original_dist / np.max(original_dist, axis=0)\n )\n V = np.zeros_like(original_dist).astype(np.float32)\n initial_rank = np.argsort(original_dist).astype(np.int32)\n\n query_num = q_g_dist.shape[0]\n gallery_num = q_g_dist.shape[0] + q_g_dist.shape[1]\n all_num = gallery_num\n\n for i in range(all_num):\n # k-reciprocal neighbors\n forward_k_neigh_index = initial_rank[i, :k1 + 1]\n backward_k_neigh_index = initial_rank[forward_k_neigh_index, :k1 + 1]\n fi = np.where(backward_k_neigh_index == i)[0]\n k_reciprocal_index = forward_k_neigh_index[fi]\n k_reciprocal_expansion_index = k_reciprocal_index\n for j in range(len(k_reciprocal_index)):\n candidate = k_reciprocal_index[j]\n candidate_forward_k_neigh_index = initial_rank[\n candidate, :int(np.around(k1 / 2.)) + 1]\n candidate_backward_k_neigh_index = initial_rank[\n candidate_forward_k_neigh_index, :int(np.around(k1 / 2.)) + 1]\n fi_candidate = np.where(\n candidate_backward_k_neigh_index == candidate\n )[0]\n candidate_k_reciprocal_index = candidate_forward_k_neigh_index[\n fi_candidate]\n if len(\n np.\n intersect1d(candidate_k_reciprocal_index, k_reciprocal_index)\n ) > 2. / 3 * len(candidate_k_reciprocal_index):\n k_reciprocal_expansion_index = np.append(\n k_reciprocal_expansion_index, candidate_k_reciprocal_index\n )\n\n k_reciprocal_expansion_index = np.unique(k_reciprocal_expansion_index)\n weight = np.exp(-original_dist[i, k_reciprocal_expansion_index])\n V[i, k_reciprocal_expansion_index] = 1. * weight / np.sum(weight)\n original_dist = original_dist[:query_num, ]\n if k2 != 1:\n V_qe = np.zeros_like(V, dtype=np.float32)\n for i in range(all_num):\n V_qe[i, :] = np.mean(V[initial_rank[i, :k2], :], axis=0)\n V = V_qe\n del V_qe\n del initial_rank\n invIndex = []\n for i in range(gallery_num):\n invIndex.append(np.where(V[:, i] != 0)[0])\n\n jaccard_dist = np.zeros_like(original_dist, dtype=np.float32)\n\n for i in range(query_num):\n temp_min = np.zeros(shape=[1, gallery_num], dtype=np.float32)\n indNonZero = np.where(V[i, :] != 0)[0]\n indImages = []\n indImages = [invIndex[ind] for ind in indNonZero]\n for j in range(len(indNonZero)):\n temp_min[0, indImages[j]] = temp_min[0, indImages[j]] + np.minimum(\n V[i, indNonZero[j]], V[indImages[j], indNonZero[j]]\n )\n jaccard_dist[i] = 1 - temp_min / (2.-temp_min)\n\n final_dist = jaccard_dist * (1-lambda_value) + original_dist*lambda_value\n del original_dist\n del V\n del jaccard_dist\n final_dist = final_dist[:query_num, query_num:]\n return final_dist\n" }, { "path": "torchreid/utils/tensortools.py", "content": "\n\ndef replace_values(input, mask, value):\n # TODO test perfs\n output = input * (~mask) + mask * value\n # input[mask] = value\n # output = input\n # output = torch.where(mask, input, torch.tensor(value, dtype=input.dtype, device=(input.get_device() if input.is_cuda else None)))\n return output\n\n\ndef masked_mean(input, mask):\n # TODO CHECK ON RANKING GRID IF IT WORK WITH CONTINUOUS VISIBILITY\n \"\"\" output -1 where mean couldn't be computed \"\"\"\n valid_input = input * mask\n mean_weights = mask.sum(0)\n mean_weights = mean_weights + (mean_weights == 0) # to avoid division by 0\n pairwise_dist = valid_input.sum(0) / mean_weights\n invalid_pairs = (mask.sum(dim=0) == 0)\n valid_pairwise_dist = replace_values(pairwise_dist, invalid_pairs, -1)\n return valid_pairwise_dist\n" }, { "path": "torchreid/utils/tools.py", "content": "from __future__ import division, print_function, absolute_import\nimport os\nimport sys\nimport json\nimport time\nimport errno\nimport cv2\nimport numpy as np\nimport random\nimport os.path as osp\nimport warnings\nimport PIL\nimport torch\nfrom PIL import Image\nfrom torchreid.utils.constants import bn_correspondants\n\n__all__ = [\n 'mkdir_if_missing', 'check_isfile', 'read_json', 'write_json',\n 'set_random_seed', 'download_url', 'read_image', 'read_masks', 'collect_env_info', 'perc'\n]\n\n\ndef mkdir_if_missing(dirname):\n \"\"\"Creates dirname if it is missing.\"\"\"\n if not osp.exists(dirname):\n try:\n os.makedirs(dirname)\n except OSError as e:\n if e.errno != errno.EEXIST:\n raise\n\n\ndef check_isfile(fpath):\n \"\"\"Checks if the given path is a file.\n\n Args:\n fpath (str): file path.\n\n Returns:\n bool\n \"\"\"\n isfile = osp.isfile(fpath)\n if not isfile:\n warnings.warn('No file found at \"{}\"'.format(fpath))\n return isfile\n\n\ndef read_json(fpath):\n \"\"\"Reads json file from a path.\"\"\"\n with open(fpath, 'r') as f:\n obj = json.load(f)\n return obj\n\n\ndef write_json(obj, fpath):\n \"\"\"Writes to a json file.\"\"\"\n mkdir_if_missing(osp.dirname(fpath))\n with open(fpath, 'w') as f:\n json.dump(obj, f, indent=4, separators=(',', ': '))\n\n\ndef set_random_seed(seed):\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n\n\ndef download_url(url, dst):\n \"\"\"Downloads file from a url to a destination.\n\n Args:\n url (str): url to download file.\n dst (str): destination path.\n \"\"\"\n from six.moves import urllib\n print('* url=\"{}\"'.format(url))\n print('* destination=\"{}\"'.format(dst))\n\n def _reporthook(count, block_size, total_size):\n global start_time\n if count == 0:\n start_time = time.time()\n return\n duration = time.time() - start_time\n progress_size = int(count * block_size)\n speed = int(progress_size / (1024*duration))\n percent = int(count * block_size * 100 / total_size)\n sys.stdout.write(\n '\\r...%d%%, %d MB, %d KB/s, %d seconds passed' %\n (percent, progress_size / (1024*1024), speed, duration)\n )\n sys.stdout.flush()\n\n urllib.request.urlretrieve(url, dst, _reporthook)\n sys.stdout.write('\\n')\n\n\ndef read_image(path):\n \"\"\"Reads image from path using ``PIL.Image``.\n\n Args:\n path (str): path to an image.\n\n Returns:\n PIL image\n \"\"\"\n got_img = False\n if not osp.exists(path):\n raise IOError('\"{}\" does not exist'.format(path))\n while not got_img:\n try:\n img = cv2.imread(path)\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n got_img = True\n except IOError:\n print(\n 'IOError incurred when reading \"{}\". Will redo. Don\\'t worry. Just chill.'\n .format(path)\n )\n return img\n\n\ndef read_masks(masks_path):\n \"\"\"Reads part-based masks information from path.\n\n Args:\n path (str): path to an image. Part-based masks information is stored in a .npy file with image name as prefix\n\n Returns:\n Numpy array of size N x H x W where N is the number of part-based masks\n \"\"\"\n\n got_masks = False\n if not osp.exists(masks_path):\n raise IOError('Masks file\"{}\" does not exist'.format(masks_path))\n while not got_masks:\n try:\n masks = np.load(masks_path)\n masks = np.transpose(masks, (1, 2, 0))\n got_masks = True\n except IOError:\n print(\n 'IOError incurred when reading \"{}\". Will redo. Don\\'t worry. Just chill.'\n .format(masks_path)\n )\n return masks\n\n\ndef collect_env_info():\n \"\"\"Returns env info as a string.\n\n Code source: github.com/facebookresearch/maskrcnn-benchmark\n \"\"\"\n from torch.utils.collect_env import get_pretty_env_info\n env_str = get_pretty_env_info()\n env_str += '\\n Pillow ({})'.format(PIL.__version__)\n return env_str\n\n\ndef perc(val, decimals=2):\n return np.around(val*100, decimals)\n\ndef extract_test_embeddings(model_output, test_embeddings):\n embeddings, visibility_scores, id_cls_scores, pixels_cls_scores, spatial_features, parts_masks = model_output\n embeddings_list = []\n visibility_scores_list = []\n embeddings_masks_list = []\n\n for test_emb in test_embeddings:\n embds = embeddings[test_emb]\n embeddings_list.append(embds if len(embds.shape) == 3 else embds.unsqueeze(1))\n if test_emb in bn_correspondants:\n test_emb = bn_correspondants[test_emb]\n vis_scores = visibility_scores[test_emb]\n visibility_scores_list.append(vis_scores if len(vis_scores.shape) == 2 else vis_scores.unsqueeze(1))\n pt_masks = parts_masks[test_emb]\n embeddings_masks_list.append(pt_masks if len(pt_masks.shape) == 4 else pt_masks.unsqueeze(1))\n\n assert len(embeddings) != 0\n\n embeddings = torch.cat(embeddings_list, dim=1) # [N, P+2, D]\n visibility_scores = torch.cat(visibility_scores_list, dim=1) # [N, P+2]\n embeddings_masks = torch.cat(embeddings_masks_list, dim=1) # [N, P+2, Hf, Wf]\n\n return embeddings, visibility_scores, embeddings_masks, pixels_cls_scores\n\n" }, { "path": "torchreid/utils/torch_receptive_field/__init__.py", "content": "from .receptive_field import receptive_field\nfrom .receptive_field import receptive_field_for_unit\n" }, { "path": "torchreid/utils/torch_receptive_field/receptive_field.py", "content": "import torch\nimport torch.nn as nn\nfrom torch.autograd import Variable\n\nfrom collections import OrderedDict\nimport numpy as np\n\n# Source: https://github.com/Fangyh09/pytorch-receptive-field\n\ndef check_same(stride):\n if isinstance(stride, (list, tuple)):\n assert len(stride) == 2 and stride[0] == stride[1]\n stride = stride[0]\n return stride\n\ndef receptive_field(model, input_size, batch_size=-1, device=\"cuda\"):\n '''\n :parameter\n 'input_size': tuple of (Channel, Height, Width)\n\n :return OrderedDict of `Layername`->OrderedDict of receptive field stats {'j':,'r':,'start':,'conv_stage':,'output_shape':,}\n 'j' for \"jump\" denotes how many pixels do the receptive fields of spatially neighboring units in the feature tensor\n do not overlap in one direction.\n i.e. shift one unit in this feature map == how many pixels shift in the input image in one direction.\n 'r' for \"receptive_field\" is the spatial range of the receptive field in one direction.\n 'start' denotes the center of the receptive field for the first unit (start) in on direction of the feature tensor.\n Convention is to use half a pixel as the center for a range. center for `slice(0,5)` is 2.5.\n '''\n def register_hook(module):\n\n def hook(module, input, output):\n class_name = str(module.__class__).split(\".\")[-1].split(\"'\")[0]\n module_idx = len(receptive_field)\n m_key = \"%i\" % module_idx\n p_key = \"%i\" % (module_idx - 1)\n receptive_field[m_key] = OrderedDict()\n\n if not receptive_field[\"0\"][\"conv_stage\"]:\n print(\"Enter in deconv_stage\")\n receptive_field[m_key][\"j\"] = 0\n receptive_field[m_key][\"r\"] = 0\n receptive_field[m_key][\"start\"] = 0\n else:\n p_j = receptive_field[p_key][\"j\"]\n p_r = receptive_field[p_key][\"r\"]\n p_start = receptive_field[p_key][\"start\"]\n \n if class_name == \"Conv2d\" or class_name == \"MaxPool2d\":\n kernel_size = module.kernel_size\n stride = module.stride\n padding = module.padding\n dilation = module.dilation\n \n kernel_size, stride, padding, dilation = map(check_same, [kernel_size, stride, padding, dilation])\n receptive_field[m_key][\"j\"] = p_j * stride\n receptive_field[m_key][\"r\"] = p_r + ((kernel_size - 1) * dilation) * p_j\n receptive_field[m_key][\"start\"] = p_start + ((kernel_size - 1) / 2 - padding) * p_j\n elif class_name == \"BatchNorm2d\" or class_name == \"ReLU\" or class_name == \"Bottleneck\":\n receptive_field[m_key][\"j\"] = p_j\n receptive_field[m_key][\"r\"] = p_r\n receptive_field[m_key][\"start\"] = p_start\n elif class_name == \"ConvTranspose2d\":\n receptive_field[\"0\"][\"conv_stage\"] = False\n receptive_field[m_key][\"j\"] = 0\n receptive_field[m_key][\"r\"] = 0\n receptive_field[m_key][\"start\"] = 0\n else:\n raise ValueError(\"module not ok\")\n pass\n receptive_field[m_key][\"input_shape\"] = list(input[0].size()) # only one\n receptive_field[m_key][\"input_shape\"][0] = batch_size\n if isinstance(output, (list, tuple)):\n # list/tuple\n receptive_field[m_key][\"output_shape\"] = [\n [-1] + list(o.size())[1:] for o in output\n ]\n else:\n # tensor\n receptive_field[m_key][\"output_shape\"] = list(output.size())\n receptive_field[m_key][\"output_shape\"][0] = batch_size\n\n if (\n not isinstance(module, nn.Sequential)\n and not isinstance(module, nn.ModuleList)\n and not (module == model)\n ):\n hooks.append(module.register_forward_hook(hook))\n\n device = device.lower()\n assert device in [\n \"cuda\",\n \"cpu\",\n ], \"Input device is not valid, please specify 'cuda' or 'cpu'\"\n\n if device == \"cuda\" and torch.cuda.is_available():\n dtype = torch.cuda.FloatTensor\n else:\n dtype = torch.FloatTensor\n\n # check if there are multiple inputs to the network\n if isinstance(input_size[0], (list, tuple)):\n x = [Variable(torch.rand(2, *in_size)).type(dtype) for in_size in input_size]\n else:\n x = Variable(torch.rand(2, *input_size)).type(dtype)\n\n # create properties\n receptive_field = OrderedDict()\n receptive_field[\"0\"] = OrderedDict()\n receptive_field[\"0\"][\"j\"] = 1.0\n receptive_field[\"0\"][\"r\"] = 1.0\n receptive_field[\"0\"][\"start\"] = 0.5\n receptive_field[\"0\"][\"conv_stage\"] = True\n receptive_field[\"0\"][\"output_shape\"] = list(x.size())\n receptive_field[\"0\"][\"output_shape\"][0] = batch_size\n hooks = []\n\n # register hook\n model.apply(register_hook)\n\n # make a forward pass\n model(x)\n\n # remove these hooks\n for h in hooks:\n h.remove()\n\n print(\"------------------------------------------------------------------------------\")\n line_new = \"{:>20} {:>10} {:>10} {:>10} {:>15} \".format(\"Layer (type)\", \"map size\", \"start\", \"jump\", \"receptive_field\")\n print(line_new)\n print(\"==============================================================================\")\n total_params = 0\n total_output = 0\n trainable_params = 0\n for layer in receptive_field:\n # input_shape, output_shape, trainable, nb_params\n assert \"start\" in receptive_field[layer], layer\n assert len(receptive_field[layer][\"output_shape\"]) == 4\n line_new = \"{:7} {:12} {:>10} {:>10} {:>10} {:>15} \".format(\n \"\",\n layer,\n str(receptive_field[layer][\"output_shape\"][2:]),\n str(receptive_field[layer][\"start\"]),\n str(receptive_field[layer][\"j\"]),\n format(str(receptive_field[layer][\"r\"]))\n )\n print(line_new)\n\n print(\"==============================================================================\")\n # add input_shape\n receptive_field[\"input_size\"] = input_size\n return receptive_field\n\n\ndef receptive_field_for_unit(receptive_field_dict, layer, unit_position):\n \"\"\"Utility function to calculate the receptive field for a specific unit in a layer\n using the dictionary calculated above\n :parameter\n 'layer': layer name, should be a key in the result dictionary\n 'unit_position': spatial coordinate of the unit (H, W)\n\n ```\n alexnet = models.alexnet()\n model = alexnet.features.to('cuda')\n receptive_field_dict = receptive_field(model, (3, 224, 224))\n receptive_field_for_unit(receptive_field_dict, \"8\", (6,6))\n ```\n Out: [(62.0, 161.0), (62.0, 161.0)]\n \"\"\"\n input_shape = receptive_field_dict[\"input_size\"]\n if layer in receptive_field_dict:\n rf_stats = receptive_field_dict[layer]\n assert len(unit_position) == 2\n feat_map_lim = rf_stats['output_shape'][2:]\n if np.any([unit_position[idx] < 0 or\n unit_position[idx] >= feat_map_lim[idx]\n for idx in range(2)]):\n raise Exception(\"Unit position outside spatial extent of the feature tensor ((H, W) = (%d, %d)) \" % tuple(feat_map_lim))\n # X, Y = tuple(unit_position)\n rf_range = [(rf_stats['start'] + idx * rf_stats['j'] - rf_stats['r'] / 2,\n rf_stats['start'] + idx * rf_stats['j'] + rf_stats['r'] / 2) for idx in unit_position]\n if len(input_shape) == 2:\n limit = input_shape\n else: # input shape is (channel, H, W)\n limit = input_shape[1:3]\n rf_range = [(max(0, rf_range[axis][0]), min(limit[axis], rf_range[axis][1])) for axis in range(2)]\n print(\"Receptive field size for layer %s, unit_position %s, is \\n %s\" % (layer, unit_position, rf_range))\n return rf_range\n else:\n raise KeyError(\"Layer name incorrect, or not included in the model.\")\n" }, { "path": "torchreid/utils/torchtools.py", "content": "from __future__ import division, print_function, absolute_import\n\nimport itertools\nimport pickle\nimport shutil\nimport os.path as osp\nimport warnings\nfrom functools import partial\nfrom collections import OrderedDict\nimport torch\nimport torch.nn as nn\n\nfrom .tools import mkdir_if_missing\nimport collections\nimport re\n\n__all__ = [\n 'save_checkpoint', 'load_checkpoint', 'resume_from_checkpoint',\n 'open_all_layers', 'open_specified_layers', 'count_num_param',\n 'load_pretrained_weights'\n]\n\n\ndef save_checkpoint(\n state, save_dir, job_id=None, is_best=False, remove_module_from_keys=False\n):\n r\"\"\"Saves checkpoint.\n\n Args:\n state (dict): dictionary.\n save_dir (str): directory to save checkpoint.\n is_best (bool, optional): if True, this checkpoint will be copied and named\n ``model-best.pth.tar``. Default is False.\n remove_module_from_keys (bool, optional): whether to remove \"module.\"\n from layer names. Default is False.\n\n Examples::\n >>> state = {\n >>> 'state_dict': model.state_dict(),\n >>> 'epoch': 10,\n >>> 'rank1': 0.5,\n >>> 'optimizer': optimizer.state_dict()\n >>> }\n >>> save_checkpoint(state, 'log/my_model')\n \"\"\"\n mkdir_if_missing(save_dir)\n if remove_module_from_keys:\n # remove 'module.' in state_dict's keys\n state_dict = state['state_dict']\n new_state_dict = OrderedDict()\n for k, v in state_dict.items():\n if k.startswith('module.'):\n k = k[7:]\n new_state_dict[k] = v\n state['state_dict'] = new_state_dict\n # save\n epoch = state['epoch']\n fpath = osp.join(save_dir, 'job-{}_{}_model.pth.tar'.format(job_id, str(epoch)))\n torch.save(state, fpath)\n print('Checkpoint saved to \"{}\"'.format(fpath))\n if is_best:\n shutil.copy(fpath, osp.join(osp.dirname(fpath), 'model-best.pth.tar'))\n\n\ndef load_checkpoint(fpath):\n r\"\"\"Loads checkpoint.\n\n ``UnicodeDecodeError`` can be well handled, which means\n python2-saved files can be read from python3.\n\n Args:\n fpath (str): path to checkpoint.\n\n Returns:\n dict\n\n Examples::\n >>> from torchreid.utils import load_checkpoint\n >>> fpath = 'log/my_model/model.pth.tar-10'\n >>> checkpoint = load_checkpoint(fpath)\n \"\"\"\n if fpath is None:\n raise ValueError('File path is None')\n if not osp.exists(fpath):\n raise FileNotFoundError('File is not found at \"{}\"'.format(fpath))\n map_location = None if torch.cuda.is_available() else 'cpu'\n try:\n checkpoint = torch.load(fpath, map_location=map_location)\n except UnicodeDecodeError:\n pickle.load = partial(pickle.load, encoding=\"latin1\")\n pickle.Unpickler = partial(pickle.Unpickler, encoding=\"latin1\")\n checkpoint = torch.load(\n fpath, pickle_module=pickle, map_location=map_location\n )\n except Exception:\n print('Unable to load checkpoint from \"{}\"'.format(fpath))\n raise\n return checkpoint\n\n\ndef resume_from_checkpoint(fpath, model, optimizer=None, scheduler=None):\n r\"\"\"Resumes training from a checkpoint.\n\n This will load (1) model weights and (2) ``state_dict``\n of optimizer if ``optimizer`` is not None.\n\n Args:\n fpath (str): path to checkpoint.\n model (nn.Module): model.\n optimizer (Optimizer, optional): an Optimizer.\n scheduler (LRScheduler, optional): an LRScheduler.\n\n Returns:\n int: start_epoch.\n\n Examples::\n >>> from torchreid.utils import resume_from_checkpoint\n >>> fpath = 'log/my_model/model.pth.tar-10'\n >>> start_epoch = resume_from_checkpoint(\n >>> fpath, model, optimizer, scheduler\n >>> )\n \"\"\"\n print('Loading checkpoint from \"{}\"'.format(fpath))\n checkpoint = load_checkpoint(fpath)\n model.load_state_dict(checkpoint['state_dict'])\n print('Loaded model weights')\n if optimizer is not None and 'optimizer' in checkpoint.keys():\n optimizer.load_state_dict(checkpoint['optimizer'])\n print('Loaded optimizer')\n if scheduler is not None and 'scheduler' in checkpoint.keys():\n scheduler.load_state_dict(checkpoint['scheduler'])\n print('Loaded scheduler')\n start_epoch = checkpoint['epoch']\n print('Last epoch = {}'.format(start_epoch))\n if 'rank1' in checkpoint.keys():\n print('Last rank1 = {:.1%}'.format(checkpoint['rank1']))\n return start_epoch\n\n\ndef adjust_learning_rate(\n optimizer,\n base_lr,\n epoch,\n stepsize=20,\n gamma=0.1,\n linear_decay=False,\n final_lr=0,\n max_epoch=100\n):\n r\"\"\"Adjusts learning rate.\n\n Deprecated.\n \"\"\"\n if linear_decay:\n # linearly decay learning rate from base_lr to final_lr\n frac_done = epoch / max_epoch\n lr = frac_done*final_lr + (1.-frac_done) * base_lr\n else:\n # decay learning rate by gamma for every stepsize\n lr = base_lr * (gamma**(epoch // stepsize))\n\n for param_group in optimizer.param_groups:\n param_group['lr'] = lr\n\n\ndef set_bn_to_eval(m):\n r\"\"\"Sets BatchNorm layers to eval mode.\"\"\"\n # 1. no update for running mean and var\n # 2. scale and shift parameters are still trainable\n classname = m.__class__.__name__\n if classname.find('BatchNorm') != -1:\n m.eval()\n\n\ndef open_all_layers(model):\n r\"\"\"Opens all layers in model for training.\n\n Examples::\n >>> from torchreid.utils import open_all_layers\n >>> open_all_layers(model)\n \"\"\"\n model.train()\n for p in model.parameters():\n p.requires_grad = True\n\n\ndef open_specified_layers(model, open_layers):\n r\"\"\"Opens specified layers in model for training while keeping\n other layers frozen.\n\n Args:\n model (nn.Module): neural net model.\n open_layers (str or list): layers open for training.\n\n Examples::\n >>> from torchreid.utils import open_specified_layers\n >>> # Only model.classifier will be updated.\n >>> open_layers = 'classifier'\n >>> open_specified_layers(model, open_layers)\n >>> # Only model.fc and model.classifier will be updated.\n >>> open_layers = ['fc', 'classifier']\n >>> open_specified_layers(model, open_layers)\n \"\"\"\n if isinstance(model, nn.DataParallel):\n model = model.module\n\n if isinstance(open_layers, str):\n open_layers = [open_layers]\n\n for layer in open_layers:\n assert hasattr(\n model, layer\n ), '\"{}\" is not an attribute of the model, please provide the correct name'.format(\n layer\n )\n\n for name, module in model.named_children():\n if name in open_layers:\n module.train()\n for p in module.parameters():\n p.requires_grad = True\n else:\n module.eval()\n for p in module.parameters():\n p.requires_grad = False\n\n\ndef count_num_param(model):\n r\"\"\"Counts number of parameters in a model while ignoring ``self.classifier``.\n\n Args:\n model (nn.Module): network model.\n\n Examples::\n >>> from torchreid.utils import count_num_param\n >>> model_size = count_num_param(model)\n\n .. warning::\n \n This method is deprecated in favor of\n ``torchreid.utils.compute_model_complexity``.\n \"\"\"\n warnings.warn(\n 'This method is deprecated and will be removed in the future.'\n )\n\n num_param = sum(p.numel() for p in model.parameters())\n\n if isinstance(model, nn.DataParallel):\n model = model.module\n\n if hasattr(model,\n 'classifier') and isinstance(model.classifier, nn.Module):\n # we ignore the classifier because it is unused at test time\n num_param -= sum(p.numel() for p in model.classifier.parameters())\n\n return num_param\n\n\ndef load_pretrained_weights(model, weight_path):\n r\"\"\"Loads pretrained weights to model.\n\n Features::\n - Incompatible layers (unmatched in name or size) will be ignored.\n - Can automatically deal with keys containing \"module.\".\n\n Args:\n model (nn.Module): network model.\n weight_path (str): path to pretrained weights.\n\n Examples::\n >>> from torchreid.utils import load_pretrained_weights\n >>> weight_path = 'log/my_model/model-best.pth.tar'\n >>> load_pretrained_weights(model, weight_path)\n \"\"\"\n checkpoint = load_checkpoint(weight_path)\n if 'state_dict' in checkpoint:\n state_dict = checkpoint['state_dict']\n else:\n state_dict = checkpoint\n\n model_dict = model.state_dict()\n new_state_dict = OrderedDict()\n matched_layers, discarded_layers = [], []\n\n for k, v in state_dict.items():\n if k.startswith('module.'):\n k = k[7:] # discard module.\n\n if k in model_dict and model_dict[k].size() == v.size():\n new_state_dict[k] = v\n matched_layers.append(k)\n else:\n discarded_layers.append(k)\n\n model_dict.update(new_state_dict)\n model.load_state_dict(model_dict)\n\n if len(matched_layers) == 0:\n warnings.warn(\n 'The pretrained weights \"{}\" cannot be loaded, '\n 'please check the key names manually '\n '(** ignored and continue **)'.format(weight_path)\n )\n else:\n print(\n 'Successfully loaded pretrained weights from \"{}\"'.\n format(weight_path)\n )\n if len(discarded_layers) > 0:\n print(\n '** The following layers are discarded '\n 'due to unmatched keys or layer size: {}'.\n format(discarded_layers)\n )\n\n\n# Copied from torch.utils.data._utils.collate.default_collate\nnp_str_obj_array_pattern = re.compile(r'[SaUO]')\ndefault_collate_err_msg_format = (\n \"default_collate: batch must contain tensors, numpy arrays, numbers, \"\n \"dicts or lists; found {}\")\ndef collate(batch):\n r\"\"\"\n Function that takes in a batch of data and puts the elements within the batch\n into a tensor with an additional outer dimension - batch size. The exact output type can be\n a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a\n Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type.\n This is used as the default function for collation when\n `batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`.\n\n Here is the general input type (based on the type of the element within the batch) to output type mapping:\n * :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size)\n * NumPy Arrays -> :class:`torch.Tensor`\n * `float` -> :class:`torch.Tensor`\n * `int` -> :class:`torch.Tensor`\n * `str` -> `str` (unchanged)\n * `bytes` -> `bytes` (unchanged)\n * `Mapping[K, V_i]` -> `Mapping[K, default_collate([V_1, V_2, ...])]`\n * `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]`\n * `Sequence[V1_i, V2_i, ...]` -> `Sequence[default_collate([V1_1, V1_2, ...]), default_collate([V2_1, V2_2, ...]), ...]`\n\n Args:\n batch: a single batch to be collated\n\n Examples:\n >>> # Example with a batch of `int`s:\n >>> default_collate([0, 1, 2, 3])\n tensor([0, 1, 2, 3])\n >>> # Example with a batch of `str`s:\n >>> default_collate(['a', 'b', 'c'])\n ['a', 'b', 'c']\n >>> # Example with `Map` inside the batch:\n >>> default_collate([{'A': 0, 'B': 1}, {'A': 100, 'B': 100}])\n {'A': tensor([ 0, 100]), 'B': tensor([ 1, 100])}\n >>> # Example with `NamedTuple` inside the batch:\n >>> Point = namedtuple('Point', ['x', 'y'])\n >>> default_collate([Point(0, 0), Point(1, 1)])\n Point(x=tensor([0, 1]), y=tensor([0, 1]))\n >>> # Example with `Tuple` inside the batch:\n >>> default_collate([(0, 1), (2, 3)])\n [tensor([0, 2]), tensor([1, 3])]\n >>> # Example with `List` inside the batch:\n >>> default_collate([[0, 1], [2, 3]])\n [tensor([0, 2]), tensor([1, 3])]\n \"\"\"\n elem = batch[0]\n elem_type = type(elem)\n if isinstance(elem, torch.Tensor):\n out = None\n if torch.utils.data.get_worker_info() is not None:\n # If we're in a background process, concatenate directly into a\n # shared memory tensor to avoid an extra copy\n numel = sum(x.numel() for x in batch)\n storage = elem.storage()._new_shared(numel)\n out = elem.new(storage).resize_(len(batch), *list(elem.size()))\n return torch.cat(batch, dim=0, out=out).numpy() # changed to cat from original code\n elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \\\n and elem_type.__name__ != 'string_':\n if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':\n # array of string classes and object\n if np_str_obj_array_pattern.search(elem.dtype.str) is not None:\n raise TypeError(default_collate_err_msg_format.format(elem.dtype))\n\n return collate([b for b in batch])\n elif elem.shape == (): # scalars\n return torch.as_tensor(batch).numpy()\n elif isinstance(elem, float):\n return torch.tensor(batch, dtype=torch.float64).numpy()\n elif isinstance(elem, int):\n return torch.tensor(batch).numpy()\n elif isinstance(elem, str):\n return batch\n elif isinstance(elem, collections.abc.Mapping):\n try:\n return elem_type({key: collate([d[key] for d in batch]) for key in elem})\n except TypeError:\n # The mapping type may not support `__init__(iterable)`.\n return {key: collate([d[key] for d in batch]) for key in elem}\n elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple\n return elem_type(*(collate(samples) for samples in zip(*batch)))\n elif isinstance(elem, collections.abc.Sequence):\n # check to make sure that the elements in batch have consistent size\n return list(itertools.chain.from_iterable(batch))\n # it = iter(batch)\n # elem_size = len(next(it))\n # # if not all(len(elem) == elem_size for elem in it):\n # # raise RuntimeError('each element in list of batch should be of equal size')\n # transposed = list(zip(*batch)) # It may be accessed twice, so we use a list.\n #\n # if isinstance(elem, tuple):\n # return [collate(samples) for samples in transposed] # Backwards compatibility.\n # else:\n # try:\n # return elem_type([collate(samples) for samples in transposed])\n # except TypeError:\n # # The sequence type may not support `__init__(iterable)` (e.g., `range`).\n # return [collate(samples) for samples in transposed]\n\n raise TypeError(default_collate_err_msg_format.format(elem_type))\n" }, { "path": "torchreid/utils/visualization/__init__.py", "content": "from .display_batch_triplets import show_triplet_grid\nfrom .visualize_query_gallery_rankings import visualize_ranking_grid" }, { "path": "torchreid/utils/visualization/display_batch_triplets.py", "content": "import cv2\nimport matplotlib.pyplot as plt\n\n__all__ = ['show_triplet']\n\n# try:\n# import matplotlib.cm\n# CMAP_JET = copy.copy(matplotlib.cm.get_cmap('jet'))\n# CMAP_JET.set_bad('white', alpha=0.5)\n# except ImportError:\n# CMAP_JET = None\nfrom torchreid.utils import Logger\nfrom torchreid.utils.engine_state import EngineState\n\nred = [1, 0, 0]\ngreen = [0, 1, 0]\nblack = [0, 0, 0]\nimg_size = (128, 256)\n\ndef show_triplet_grid(triplets):\n\n fig11 = plt.figure(figsize=(40, 50), constrained_layout=False)\n outer_grid = fig11.add_gridspec(4, 5)\n\n count = 0\n for a in range(4):\n for b in range(5):\n print(\"grid {}-{}\".format(a, b))\n # gridspec inside gridspec\n inner_grid = outer_grid[a, b].subgridspec(1, 3)\n axs = inner_grid.subplots() # Create all subplots for the inner grid.\n triplet = triplets[count]\n pos, anc, neg, pos_dist, neg_dist = triplet\n ax1, ax2, ax3 = axs[0], axs[1], axs[2]\n show_instance(ax1, pos, pos_dist, green)\n show_instance(ax2, anc, 0, black)\n show_instance(ax3, neg, neg_dist, red)\n count += 1\n\n # show only the outside spines\n # for ax in fig11.get_axes():\n # ax.spines['top'].set_visible(ax.is_first_row())\n # ax.spines['bottom'].set_visible(ax.is_last_row())\n # ax.spines['left'].set_visible(ax.is_first_col())\n # ax.spines['right'].set_visible(ax.is_last_col())\n Logger.current_logger().add_figure(\"Batch triplets\", fig11, EngineState.current_engine_state().epoch)\n # plt.show()\n # plt.waitforbuttonpress()\n\ndef show_triplet(anc, pos, neg, pos_dist, neg_dist):\n # instance = (image, masks, id, body_part_id, body_part_name)\n f, axarr = plt.subplots(1, 3)\n ax1, ax2, ax3 = axarr[0], axarr[1], axarr[2]\n show_instance(ax1, pos, pos_dist, green)\n show_instance(ax2, anc, 0, black)\n show_instance(ax3, neg, neg_dist, red)\n f.matplotlib_show()\n plt.waitforbuttonpress()\n\n\ndef add_border(img, color):\n # border widths\n top, bottom, left, right = [5] * 4\n return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)\n\n\ndef show_instance(ax, instance, dist, color):\n mask_idx = instance[2]\n body_part = instance[3]\n img = instance[0]\n mask = instance[1]\n\n # img = overlay_mask_1(img, mask)\n img = cv2.resize(img, img_size)\n # img = add_border(img, color)\n\n ax.imshow(img)\n mask = cv2.resize(mask, dsize=(img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC)\n # mask = add_border(mask, color)\n ax.imshow(mask, cmap='jet', vmin=0, vmax=1, alpha=0.5)\n\n ax.tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, right=False,\n labelleft=False)\n ax.set_title(\"Id = {}\\n{}\".format(mask_idx, body_part))\n ax.set_xlabel('Dist = {}'.format(dist))\n for axis in ['top', 'bottom', 'left', 'right']:\n ax.spines[axis].set_color(color)\n ax.spines[axis].set_linewidth(4)" }, { "path": "torchreid/utils/visualization/embeddings_projection.py", "content": "import cv2\nimport torch\nimport numpy as np\n\nfrom torchreid.utils import Logger\nfrom torchreid.utils.engine_state import EngineState\n\n\ndef visualize_embeddings(qf, gf, q_pids, g_pids, test_loader, dataset_name, qf_parts_visibility, gf_parts_visibility, mAP, rank1):\n query_dataset = test_loader['query'].dataset\n gallery_dataset = test_loader['gallery'].dataset\n # TODO 1000 identities and 5 samples per identity\n sample_size = 1000\n q_embeddings, q_imgs, q_meta, q_idx_list = extract_samples(qf, query_dataset, sample_size)\n g_embeddings, g_imgs, g_meta, g_idx_list = extract_samples(gf, gallery_dataset, sample_size)\n\n embeddings = torch.cat([q_embeddings, g_embeddings], 0)\n imgs = torch.cat([q_imgs, g_imgs], 0)\n meta = q_meta + g_meta\n\n logger = Logger.current_logger()\n for body_part_idx in range(0, embeddings.shape[1]):\n logger.add_embeddings(\"{} query-gallery embeddings projection for {} with mAP {} and rank-1 {}\".format(dataset_name, body_part_idx, mAP, rank1), embeddings[:, body_part_idx], meta, imgs, EngineState.current_engine_state().epoch)\n\n\ndef extract_samples(features, dataset, sample_size):\n sample_size = min(sample_size, len(dataset))\n remaining_idx = np.arange(0, len(dataset))\n\n idx_list = np.random.choice(remaining_idx, replace=False, size=sample_size)\n\n embeddings = []\n meta = []\n imgs = []\n for idx in idx_list:\n _, pid, camid, img_path, masks = dataset[idx]\n embeddings.append(features[idx, :, :])\n img = cv2.imread(img_path)\n img = cv2.resize(img, (64, 64))\n img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n img = img / 255\n img = torch.from_numpy(img)\n\n imgs.append(img)\n meta.append(str(pid))\n\n embeddings = torch.stack(embeddings)\n imgs = torch.stack(imgs)\n imgs = imgs.permute(0, 3, 1, 2)\n\n return embeddings, imgs, meta, idx_list\n" }, { "path": "torchreid/utils/visualization/feature_map_visualization.py", "content": "import math\nimport os.path as osp\nimport json\n\nimport cv2\nimport numpy as np\nfrom PIL import Image\nfrom sklearn.decomposition import PCA\nimport time\nimport matplotlib.pyplot as plt\nimport torch\n\n\n# from reid import datasets, models\n# from reid.utils.data import transforms as T\n# from reid.utils.data.preprocessor import Preprocessor\n# from reid.utils.osutils import set_paths\n# from reid.utils.serialization import load_checkpoint\n# from reid.models.CompactBilinearPooling_dsybaik import CompactBilinearPooling\n\n# Source: https://github.com/yuminsuh/part_bilinear_reid/blob/master/vis_featmap.ipynb\nfrom torchreid.utils import Writer, Logger\nfrom torchreid.utils.constants import PARTS\n\n\ndef flatten(maps):\n flattened = np.transpose(maps, (1, 0, 2, 3)).reshape(maps.shape[1], int(maps.size / maps.shape[1]))\n return flattened\n\n\ndef organize(flattened, num_map, feat_dim, h, w):\n maps = flattened.reshape(feat_dim, num_map, h, w)\n maps = np.transpose(maps, (1, 0, 2, 3))\n return maps\n\n\ndef feat_to_color(maps):\n num_img, dim_feat, h, w = maps.shape\n maps_flatten = flatten(maps)\n maps_flatten_reduced = PCA(n_components=3).fit_transform(maps_flatten.transpose())\n maps_reduced = organize(maps_flatten_reduced.transpose(), num_img, 3, h, w)\n return maps_reduced\n\n\ndef normalize_01(m):\n max_value = m.max()\n min_value = m.min()\n r = max_value - min_value\n n = (m - min_value) / r\n return n\n\n\ndef mapwise_normalize(map1, feats):\n num_sample, num_channel, h, w = map1.shape # FIXME TOO MANY VALUES TO UNPACK WITH INDEITITy MASK\n # mode 1\n # norms = np.linalg.norm(np.sqrt(np.sum(map1**2, axis=(2,3))), axis=1)\n # normalized = map1\n # mode 2\n norms = np.sqrt(np.linalg.norm(feats, axis=1))\n normalized = map1 / np.reshape(norms, (num_sample, 1, 1, 1))\n return normalized\n\n\n# TODO output big grid of n identities horizontally and m samples per identity vertically\n# TODO test without norm stuff\n# TODO refactor into beautiful code\n# TODO cannot apply PCA batch by batch, should be global\ndef visualize_pca_multi(maps_all, feats, pids, tag):\n maps_all_reduced = []\n for maps in maps_all[1:]:\n if len(maps.shape) == 4:\n maps_all_reduced.append(feat_to_color(mapwise_normalize(maps, feats)))\n num_person = maps_all_reduced[0].shape[0]\n\n num_samples = len(pids)\n num_samples_per_id = 0\n for i, pid in enumerate(pids):\n if pid != pids[0]:\n num_samples_per_id = i\n break\n\n\n\n n_rows = num_samples_per_id\n n_cols = math.ceil(num_samples / num_samples_per_id)\n fig = plt.figure(figsize=(n_cols*4, n_rows*2), constrained_layout=False)\n outer_grid = fig.add_gridspec(n_rows, n_cols)\n\n # count = 0\n # for row in range(4):\n # for col in range(5):\n # print(\"grid {}-{}\".format(row, col))\n # # gridspec inside gridspec\n # inner_grid = outer_grid[row, col].subgridspec(1, 3)\n # axs = inner_grid.subplots() # Create all subplots for the inner grid.\n # triplet = triplets[count]\n # pos, anc, neg, pos_dist, neg_dist = triplet\n # ax1, ax2, ax3 = axs[0], axs[1], axs[2]\n # show_instance(ax1, pos, pos_dist, green)\n # show_instance(ax2, anc, 0, black)\n # show_instance(ax3, neg, neg_dist, red)\n # count += 1\n\n # for person_idx in range(min(num_person, num_vis)):\n person_idx = 0\n for row in range(n_rows):\n for col in range(n_cols):\n if person_idx < num_person:\n # print(person_idx)\n # fig, ax = plt.subplots(1, len(maps_all))\n # fig.set_size_inches((10,10))\n inner_grid = outer_grid[row, col].subgridspec(1, 3)\n axs = inner_grid.subplots() # Create all subplots for the inner grid.\n for idx, maps_reduced in enumerate([maps_all[0]] + maps_all_reduced):\n axs[idx].imshow(normalize_01(np.transpose(maps_reduced[person_idx, ::-1, :, :], (1, 2, 0))))\n axs[idx].set_axis_off()\n person_idx += 1\n # plt.show()\n\n # plt.tight_layout()\n plt.tight_layout(pad=1.30, h_pad=1.6, w_pad=1.6)\n Logger.current_logger().add_figure(\"features_part_maps_{}\".format(tag), fig, False)\n plt.close(fig)\n\n# fig.savefig(savepath.format(save_id, person_idx)) \n\n\ndef display_feature_maps(embeddings_dict, spatial_features, body_part_masks, imgs_path, pids):\n # TODO call at test time: display top 10 ranking for 5 queries? Take 10 random pids with 1 samples in query and\n # find 9 corresponding samples with same pid in gallery\n\n # TODO put config\n # TODO fix with new model output\n\n writer = Writer.current_writer()\n def extract_images(imgs_path):\n imgs = []\n for img_path in imgs_path:\n img = cv2.imread(img_path)\n img = cv2.resize(img, (128, 256)) # TODO size = config\n img = img / 255\n imgs.append(img)\n\n imgs = np.asarray(imgs)\n imgs = np.transpose(imgs, axes=[0, 3, 1, 2])\n return imgs\n\n if writer.engine_state.epoch == writer.engine_state.max_epoch-1 and writer.engine_state.batch < 10 and writer.cfg.test.vis_feature_maps: # TODO move away\n # TODO\n body_parts_features = embeddings_dict[PARTS]\n body_part_masks = torch.unsqueeze(body_part_masks, 2) # [N, M, 1, Hf, Wf]\n spatial_features = torch.unsqueeze(spatial_features, 1) # [N, 1, D, Hf, Wf]\n images_np = extract_images(imgs_path)\n if len(body_parts_features.shape) == 3:\n body_parts_features = body_parts_features.flatten(1, 2)\n body_parts_features_np = body_parts_features.squeeze().detach().cpu().numpy().copy()\n body_part_masks_np = body_part_masks.squeeze().detach().cpu().numpy().copy()\n spatial_features_np = spatial_features.squeeze().detach().cpu().numpy().copy()\n visualize_pca_multi([images_np, spatial_features_np, body_part_masks_np], body_parts_features_np, pids, \"e_{}_b_{}\".format(writer.engine_state.epoch, writer.engine_state.batch))\n\n# def sel_random_target(target):\n# labels = np.unique([t[1] for t in target])\n# labels = np.delete(labels, labels == 0)\n# sel_labels = np.random.choice(labels, 200, replace=False)\n# sel_target = [t for t in target if t[1] in sel_labels]\n# _, sel_target = map(list, zip(*sorted(zip([t[1] for t in sel_target], sel_target))))\n# return sel_target\n#\n# \"\"\" Settings \"\"\"\n# exp_dir = \"logs/market1501/d2_b250\"\n# batch_size = 60\n# args = json.load(open(osp.join(exp_dir, \"args.json\"), \"r\"))\n# set_paths('paths')\n# np.random.seed(0)\n# # np.random.seed(int(time.time()))\n# # savepath = 'vis_examples/{}_{}.png'\n#\n# \"\"\" Load dataset \"\"\"\n# dataset = datasets.create(args['dataset'], \"data/{}\".format(args['dataset']))\n# test_transformer = T.Compose([\n# T.RectScale(args['height'], args['width']),\n# T.CenterCrop((args['crop_height'], args['crop_width'])),\n# T.ToTensor(),\n# T.RGB_to_BGR(),\n# T.NormalizeBy(255),\n# ])\n# target = sel_random_target(list(set(dataset.query) | set(dataset.gallery)))\n# test_loader = DataLoader(\n# Preprocessor(target,\n# root=dataset.images_dir, transform=test_transformer),\n# batch_size=batch_size, num_workers=args['workers'],\n# shuffle=False, pin_memory=True)\n#\n# \"\"\" Load model \"\"\"\n# model = models.create(args['arch'], features=args['features'],\n# dilation=args['dilation'], initialize=False).cuda()\n# model_weight = osp.join(exp_dir, 'epoch_750.pth.tar')\n# checkpoint = load_checkpoint(model_weight)\n# model.app_feat_extractor.load_state_dict(checkpoint['app_state_dict'])\n# model.part_feat_extractor.load_state_dict(checkpoint['part_state_dict'])\n# model.eval()\n#\n# \"\"\" Extract feature maps \"\"\"\n# num_test = len(target)\n# feat1, feat2, feat_out, h, w = 512, 128, 512, 20, 10\n# pool = CompactBilinearPooling(feat1, feat2, feat_out, sum_pool=True)\n#\n# app_feats = np.zeros((num_test, feat1, h, w))\n# part_feats = np.zeros((num_test, feat2, h, w))\n# bilinear_feats = np.zeros((num_test, feat_out))\n# target_imgs = np.zeros((num_test, 3, args['crop_height'], args['crop_width']))\n# for i, (imgs, fnames, pids, _) in enumerate(test_loader):\n# app_feat = model.app_feat_extractor(imgs.cuda())\n# part_feat = model.part_feat_extractor(imgs.cuda())\n# bilinear_feat = pool(app_feat, part_feat)\n#\n# i_start = i * batch_size\n# i_end = min((i + 1) * batch_size, num_test)\n# app_feats[i_start:i_end, :, :, :] = app_feat.detach().cpu().numpy().copy()\n# part_feats[i_start:i_end, :, :, :] = part_feat.detach().cpu().numpy().copy()\n# bilinear_feats[i_start:i_end, :] = bilinear_feat.detach().cpu().numpy().copy()\n# target_imgs[i_start:i_end, :, :, :] = imgs.detach().cpu().numpy().copy()\n#\n# \"\"\" Visualize maps \"\"\"\n# num_vis = 200\n# visualize_pca_multi([target_imgs[:num_vis], app_feats[:num_vis], part_feats[:num_vis]], bilinear_feats[:num_vis],\n# num_vis=num_vis, save_id=\"\")" }, { "path": "torchreid/utils/visualization/visualize_query_gallery_rankings.py", "content": "import ntpath\nimport random\n\nimport cv2\nimport matplotlib\nimport numpy as np\n\nfrom torchreid.utils import Logger, perc\nfrom torchreid.utils.engine_state import EngineState\n\nGRID_SPACING_V = 100\nGRID_SPACING_H = 100\nQUERY_EXTRA_SPACING = 30\nTOP_MARGIN = 350\nLEFT_MARGIN = 150\nRIGHT_MARGIN = 500\nBOTTOM_MARGIN = 300\nROW_BACKGROUND_LEFT_MARGIN = 75\nROW_BACKGROUND_RIGHT_MARGIN = 75\nLEFT_TEXT_OFFSET = 10\nBW = 12 # border width\nGREEN = (0, 255, 0)\nRED = (0, 0, 255)\nBLUE = (255, 0, 0)\nYELLOW = (255, 255, 0)\nTEXT_FONT = cv2.FONT_HERSHEY_SIMPLEX\nTEXT_COLOR = (0, 0, 0)\nTEXT_LINE_TYPE = cv2.LINE_AA\nWIDTH = 128\nHEIGHT = 256\ncmap = matplotlib.cm.get_cmap('hsv')\n\n# TODO document and make code easier to read and adapt, i.e. less intricate\ndef visualize_ranking_grid(distmat, body_parts_distmat, test_loader, dataset_name, qf_parts_visibility, gf_parts_visibility, q_parts_masks, g_parts_masks, mAP, rank1, save_dir, topk, visrank_q_idx_list, visrank_count, config=None, bp_idx=None):\n num_q, num_g = distmat.shape\n query_dataset = test_loader['query'].dataset\n gallery_dataset = test_loader['gallery'].dataset\n assert num_q == len(query_dataset)\n assert num_g == len(gallery_dataset)\n indices = np.argsort(distmat, axis=1)\n\n mask_filtering_flag = qf_parts_visibility is not None or gf_parts_visibility is not None\n if qf_parts_visibility is None:\n qf_parts_visibility = np.ones((num_q, body_parts_distmat.shape[0]), dtype=bool)\n\n if gf_parts_visibility is None:\n gf_parts_visibility = np.ones((num_g, body_parts_distmat.shape[0]), dtype=bool)\n\n n_missing = visrank_count - len(visrank_q_idx_list)\n if n_missing > 0:\n q_idx_list = visrank_q_idx_list\n remaining_idx = np.arange(0, num_q)\n q_idx_list = np.append(q_idx_list, np.random.choice(remaining_idx, replace=False, size=n_missing))\n elif n_missing < 0:\n q_idx_list = np.array(visrank_q_idx_list[:visrank_count])\n else:\n q_idx_list = np.array(visrank_q_idx_list)\n\n q_idx_list = q_idx_list.astype(int)\n print(\"visualize_ranking_grid for dataset {}, bp {} and ids {}\".format(dataset_name, bp_idx, q_idx_list))\n for q_idx in q_idx_list:\n if q_idx >= len(query_dataset):\n # FIXME this happen when using multiple target dataset with 'visrank_q_idx_list' provided for another dataset\n new_q_idx = random.randint(0, len(query_dataset)-1)\n print(\"Invalid query index {}, using random index {} instead\".format(q_idx, new_q_idx))\n q_idx = new_q_idx\n query = query_dataset[q_idx]\n qpid, qcamid, qimg_path = query['pid'], query['camid'], query['img_path']\n qmasks = q_parts_masks[q_idx]\n if bp_idx is not None:\n qmasks = qmasks[bp_idx:bp_idx+1]\n query_sample = (q_idx, qpid, qcamid, qimg_path, qmasks, qf_parts_visibility[q_idx, :])\n gallery_topk_samples = []\n rank_idx = 1\n for g_idx in indices[q_idx, :]:\n gallery = gallery_dataset[g_idx]\n gpid, gcamid, gimg_path = gallery['pid'], gallery['camid'], gallery['img_path']\n invalid = test_loader['query'].dataset.gallery_filter(np.array(qpid),\n np.array(qcamid),\n None,\n np.array(gpid),\n np.array(gcamid),\n None).item()\n invalid = invalid or distmat[q_idx, g_idx] < 0\n\n if not invalid:\n # matched = gpid == qpid\n gmasks = g_parts_masks[g_idx]\n if bp_idx is not None:\n gmasks = gmasks[bp_idx:bp_idx+1]\n gallery_sample = (g_idx, gpid, gcamid, gimg_path, gmasks, gf_parts_visibility[g_idx, :], qpid == gpid,\n distmat[q_idx, g_idx],\n body_parts_distmat[:, q_idx, g_idx])\n gallery_topk_samples.append(gallery_sample)\n rank_idx += 1\n if rank_idx > topk:\n break\n if len(gallery_topk_samples) > 0:\n show_ranking_grid(query_sample, gallery_topk_samples, mAP, rank1, dataset_name, config, mask_filtering_flag, bp_idx)\n else:\n print(\"Skip ranking plot of query id {} ({}), no valid gallery available\".format(q_idx, qimg_path))\n\n\ndef show_ranking_grid(query_sample, gallery_topk_samples, mAP, rank1, dataset_name, config, mask_filtering_flag, bp_idx=None):\n qidx, qpid, qcamid, qimg_path, qmasks, qf_parts_visibility = query_sample\n\n topk = len(gallery_topk_samples)\n bp_num = len(qf_parts_visibility)\n\n num_cols = bp_num + 1\n num_rows = topk + 1\n grid_img = 255 * np.ones(\n (\n num_rows * HEIGHT + (num_rows + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN + BOTTOM_MARGIN,\n num_cols * WIDTH + (num_cols + 1) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN + RIGHT_MARGIN,\n 3\n ),\n dtype=np.uint8\n )\n\n samples = [query_sample] + gallery_topk_samples\n\n insert_background_line(grid_img, BLUE, 0, HEIGHT, 120, 0)\n insert_background_line(grid_img, BLUE, len(samples), HEIGHT, 0, -75)\n\n\n pos = (int(grid_img.shape[1]/2), 0)\n filtering_str = \"body part filtering with threshold {}\".format(config.model.bpbreid.masks.mask_filtering_threshold) if config.model.bpbreid.mask_filtering_testing else \"no body part filtering\"\n align_top_text(grid_img, \"Ranking for dataset {}, {}, pid {}, mAP {:.2f}%, rank1 {:.2f}%, loss {}, {}\".format(dataset_name, config.project.job_id, qpid, mAP * 100, rank1 * 100, config.loss.part_based.name, filtering_str), pos, 3.5, 7, 120)\n\n for row, sample in enumerate(samples):\n display_sample_on_row(grid_img, sample, row, (WIDTH, HEIGHT), mask_filtering_flag, qf_parts_visibility)\n\n for col in range(1, num_cols):\n parts_visibility_count = 0\n row = topk+1\n bp_idx = col - 1\n distances = []\n for i, sample in enumerate(samples):\n if i == 0:\n idx, pid, camid, img_path, masks, parts_visibility = sample\n else:\n idx, pid, camid, img_path, masks, parts_visibility, matched, dist_to_query, body_parts_dist_to_query = sample\n distances.append(body_parts_dist_to_query[bp_idx])\n parts_visibility_count += parts_visibility[bp_idx]\n distances = np.asarray(distances)\n min = distances.min()\n max = distances.max()\n mean = distances.mean()\n pos = (col * WIDTH + int(WIDTH / 2) + (col + 1) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN,\n (row) * HEIGHT + int(HEIGHT / 2) + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN)\n\n align_top_multi_text(grid_img, \"Bp={}/{}\\nMin={:.1f}\\nMean={:.1f}\\nMax={:.1f}\".format(\n parts_visibility_count, topk + 1, min, mean, max), pos, 1, 2, 60)\n\n if bp_idx is not None:\n filename = \"_{}_{}_qidx_{}_qpid_{}_{}_part_{}.jpg\".format(config.project.job_id, dataset_name, qidx, qpid, ntpath.basename(qimg_path), bp_idx)\n else:\n filename = \"_{}_{}_qidx_{}_qpid_{}_{}.jpg\".format(config.project.job_id, dataset_name, qidx, qpid, ntpath.basename(qimg_path))\n # path = os.path.join(save_dir, filename)\n # Path(os.path.dirname(path)).mkdir(parents=True, exist_ok=True)\n # cv2.imwrite(path, grid_img)\n Logger.current_logger().add_image(\"Ranking grid\", filename, cv2.cvtColor(grid_img, cv2.COLOR_BGR2RGB), EngineState.current_engine_state().epoch)\n\n\ndef insert_background_line(grid_img, match_color, row, height, padding_top=0, padding_bottom=0):\n alpha = 0.1\n color = (255 * (1-alpha) + match_color[0] * alpha,\n 255 * (1-alpha) + match_color[1] * alpha,\n 255 * (1-alpha) + match_color[2] * alpha)\n hs = row * height + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN - int(GRID_SPACING_V/2) + 15 - padding_top\n he = (row + 1) * height + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN + int(GRID_SPACING_V/2) + 15 + padding_bottom\n ws = ROW_BACKGROUND_LEFT_MARGIN\n we = grid_img.shape[1] - ROW_BACKGROUND_RIGHT_MARGIN\n grid_img[hs:he, ws:we, :] = color\n\n\ndef display_sample_on_row(grid_img, sample, row, img_shape, mask_filtering_flag, q_parts_visibility):\n if row == 0:\n idx, pid, camid, img_path, masks, parts_visibility = sample\n matched, dist_to_query, body_parts_dist_to_query = None, None, None\n else:\n idx, pid, camid, img_path, masks, parts_visibility, matched, dist_to_query, body_parts_dist_to_query = sample\n\n masks = masks.numpy()\n width, height = img_shape\n bp_num = masks.shape[0]\n img = cv2.imread(img_path)\n img = cv2.resize(img, (width, height))\n\n for col in range(0, bp_num + 1):\n bp_idx = col - 1\n if row == 0 and col == 0:\n img_to_insert = img\n img_to_insert = make_border(img_to_insert, BLUE, BW)\n pos = ((bp_num + 1) * width + (bp_num + 2) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN,\n row * height + int(height / 2) + (row + 1) * GRID_SPACING_V + TOP_MARGIN)\n align_left_multitext(grid_img, \"*Id = {}*\\n\"\n \"Visible = {}/{}\".format(\n pid, parts_visibility.sum(), bp_num), pos, 1.1, 2, 15)\n elif col == 0:\n match_color = GREEN if matched else RED\n insert_background_line(grid_img, match_color, row, height)\n img_to_insert = make_border(img, match_color, BW)\n pos = (LEFT_MARGIN + GRID_SPACING_H,\n row * height + int(height / 2) + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN)\n align_right_text(grid_img, str(row), pos, 3, 6, 30)\n pos = (LEFT_MARGIN + GRID_SPACING_H + int(width / 2),\n (row + 1) * height + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN)\n g_to_q_vis_score = np.sqrt(q_parts_visibility * parts_visibility).sum() / bp_num\n align_top_text(grid_img, \"{}% | {:.2f}\".format(int(perc(g_to_q_vis_score, 0)), dist_to_query), pos, 1.2, 2, 10)\n\n pos = ((bp_num + 1) * width + (bp_num + 2) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN,\n row * height + int(height / 2) + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN)\n if len(parts_visibility) == 1 or parts_visibility.sum() == 0:\n valid_body_parts_dist = body_parts_dist_to_query\n else:\n valid_body_parts_dist = body_parts_dist_to_query[parts_visibility > 0]\n\n align_left_multitext(grid_img, \"*Id = {}*\\n\"\n \"Idx = {}\\n\"\n \"Cam id = {}\\n\"\n \"Name = {}\\n\"\n \"Bp Visibles = {}/{}\\n\"\n \"[{:.2f}; {:.2f}; {:.2f}]\\n\"\n \"[{:.2f}; {:.2f}; {:.2f}]\".format(\n pid, idx, camid, ntpath.basename(img_path), (parts_visibility > 0).sum(), bp_num,\n body_parts_dist_to_query.min(), body_parts_dist_to_query.mean(), body_parts_dist_to_query.max(),\n valid_body_parts_dist.min(), valid_body_parts_dist.mean(), valid_body_parts_dist.max()), pos, 1, 2, 15, match_color)\n else:\n if row == 0:\n pos = (col * width + int(width / 2) + (col + 1) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN,\n TOP_MARGIN + GRID_SPACING_V)\n align_bottom_text(grid_img, str(bp_idx), pos, 2, 5, 35)\n pos = (col * width + int(width / 2) + (col + 1) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN,\n (row + 1) * height + (row + 1) * GRID_SPACING_V + TOP_MARGIN)\n align_top_text(grid_img, \"{}%\".format(int(perc(parts_visibility[bp_idx], 0))), pos, 0.9, 2, 10)\n if row != 0:\n pos = (col * width + int(width / 2) + (col + 1) * GRID_SPACING_H + QUERY_EXTRA_SPACING + LEFT_MARGIN,\n (row + 1) * height + (row + 1) * GRID_SPACING_V + QUERY_EXTRA_SPACING + TOP_MARGIN)\n thickness = 3 if body_parts_dist_to_query.argmax() == bp_idx or body_parts_dist_to_query.argmin() == bp_idx else 2\n align_top_text(grid_img, \"{}% | {:.2f}\".format(int(perc(parts_visibility[bp_idx], 0)), body_parts_dist_to_query[bp_idx]), pos, 0.9, thickness, 10)\n mask = masks[bp_idx, :, :]\n img_with_mask_overlay = mask_overlay(img, mask, interpolation=cv2.INTER_CUBIC)\n if mask_filtering_flag:\n # match_color = GREEN if parts_visibility[bp_idx] else RED\n match_color = cmap(parts_visibility[bp_idx].item()/3, bytes=True)[0:-1] # divided by three because hsv colormap goes from red to green inside [0, 0.333]\n img_to_insert = make_border(img_with_mask_overlay, (int(match_color[2]), int(match_color[1]), int(match_color[0])), BW)\n else:\n img_to_insert = img_with_mask_overlay\n\n insert_img_into_grid(grid_img, img_to_insert, row, col)\n\n\ndef mask_overlay(img, mask, clip=True, interpolation=cv2.INTER_NEAREST):\n width, height = img.shape[1], img.shape[0]\n mask = cv2.resize(mask, dsize=(width, height), interpolation=interpolation)\n if clip:\n mask = np.clip(mask, 0, 1)\n mask = (mask * 255).astype(np.uint8)\n else:\n mask = np.interp(mask, (mask.min(), mask.max()), (0, 255)).astype(np.uint8)\n mask_color = cv2.applyColorMap(mask, cv2.COLORMAP_JET)\n masked_img = cv2.addWeighted(img, 0.5, mask_color.astype(img.dtype), 0.5, 0)\n return masked_img\n\n\ndef align_top_text(img, text, pos, fontScale=1.0, thickness=1, padding=4):\n textsize = cv2.getTextSize(text, TEXT_FONT, fontScale, thickness)[0]\n textX = int(pos[0] - (textsize[0] / 2))\n textY = pos[1] + textsize[1] + padding\n cv2.putText(img, text, (textX, textY), TEXT_FONT, fontScale=fontScale, color=TEXT_COLOR, thickness=thickness,\n lineType=TEXT_LINE_TYPE)\n\n\ndef align_top_multi_text(img, text, pos, fontScale=1.0, thickness=1, padding=4, text_color=(0, 0, 0)):\n v_padding = 20\n text_lines = text.split('\\n')\n text_line_height = cv2.getTextSize(text_lines[0], TEXT_FONT, fontScale, thickness)[0][1]\n text_height = len(text_lines) * text_line_height + (len(text_lines)-1) * v_padding\n textY = int(pos[1] - text_height + text_line_height) + padding\n\n for i, text_line in enumerate(text_lines):\n bold_marker = \"*\"\n bold = text_line.startswith(bold_marker) and text_line.endswith(bold_marker)\n line_thickness = thickness+1 if bold else thickness\n if bold:\n text_line = text_line[len(bold_marker):len(text_line)-len(bold_marker)]\n textsize = cv2.getTextSize(text_line, TEXT_FONT, fontScale, thickness)[0]\n text_line_pos = (int(pos[0] - (textsize[0] / 2)), textY + (text_line_height + v_padding) * i)\n text_color = text_color if i == 0 else TEXT_COLOR\n cv2.putText(img, text_line, text_line_pos, TEXT_FONT, fontScale=fontScale, color=TEXT_COLOR, thickness=line_thickness,\n lineType=TEXT_LINE_TYPE)\n\n\ndef align_bottom_text(img, text, pos, fontScale=1.0, thickness=1, padding=4):\n textsize = cv2.getTextSize(text, TEXT_FONT, fontScale, thickness)[0]\n textX = int(pos[0] - (textsize[0] / 2))\n textY = pos[1] - padding\n cv2.putText(img, text, (textX, textY), TEXT_FONT, fontScale=fontScale, color=TEXT_COLOR, thickness=thickness,\n lineType=TEXT_LINE_TYPE)\n\n\ndef align_right_text(img, text, pos, fontScale=1.0, thickness=1, padding=4):\n textsize = cv2.getTextSize(text, TEXT_FONT, fontScale, thickness)[0]\n textX = pos[0] - textsize[0] - padding\n textY = int(pos[1] + (textsize[1] / 2))\n cv2.putText(img, text, (textX, textY), TEXT_FONT, fontScale=fontScale, color=TEXT_COLOR, thickness=thickness,\n lineType=TEXT_LINE_TYPE)\n\n\ndef align_left_multitext(img, text, pos, fontScale=1.0, thickness=1, padding=4, text_color=(0, 0, 0)):\n v_padding = 20\n text_lines = text.split('\\n')\n text_line_height = cv2.getTextSize(text_lines[0], TEXT_FONT, fontScale, thickness)[0][1]\n text_height = len(text_lines) * text_line_height + (len(text_lines)-1) * v_padding\n textX = pos[0] + padding\n textY = int(pos[1] - (text_height / 2) + text_line_height)\n\n for i, text_line in enumerate(text_lines):\n bold_marker = \"*\"\n bold = text_line.startswith(bold_marker) and text_line.endswith(bold_marker)\n line_thickness = thickness+1 if bold else thickness\n if bold:\n text_line = text_line[len(bold_marker):len(text_line)-len(bold_marker)]\n pos = (textX, textY + (text_line_height + v_padding) * i)\n text_color = text_color if i == 0 else TEXT_COLOR\n cv2.putText(img, text_line, pos, TEXT_FONT, fontScale=fontScale, color=text_color, thickness=line_thickness,\n lineType=TEXT_LINE_TYPE)\n\n\ndef centered_text(img, text, pos, fontScale=1, thickness=1):\n textsize = cv2.getTextSize(text, TEXT_FONT, fontScale, thickness)[0]\n textX = int(pos[0] - (textsize[0] / 2))\n textY = int(pos[1] + (textsize[1] / 2))\n cv2.putText(img, text, (textX, textY), TEXT_FONT, fontScale=fontScale, color=TEXT_COLOR, thickness=thickness,\n lineType=TEXT_LINE_TYPE)\n\n\ndef insert_img_into_grid(grid_img, img, row, col):\n extra_spacing_h = QUERY_EXTRA_SPACING if row > 0 else 0\n extra_spacing_w = QUERY_EXTRA_SPACING if col > 0 else 0\n width, height = img.shape[1], img.shape[0]\n hs = row * height + (row + 1) * GRID_SPACING_V + extra_spacing_h + TOP_MARGIN\n he = (row + 1) * height + (row + 1) * GRID_SPACING_V + extra_spacing_h + TOP_MARGIN\n ws = col * width + (col + 1) * GRID_SPACING_H + extra_spacing_w + LEFT_MARGIN\n we = (col + 1) * width + (col + 1) * GRID_SPACING_H + extra_spacing_w + LEFT_MARGIN\n grid_img[hs:he, ws:we, :] = img\n\n\ndef make_border(img, border_color, bw):\n img_b = cv2.copyMakeBorder(\n img,\n bw, bw, bw, bw,\n cv2.BORDER_CONSTANT,\n value=border_color\n )\n img_b = cv2.resize(img_b, (img.shape[1], img.shape[0]))\n return img_b\n\n#####################################\n# Matplotlib version - too slow #\n#####################################\n\n# GRID_SPACING = 20\n# QUERY_EXTRA_SPACING = 60\n# BW = 12 # border width\n# GREEN = (0, 255, 0)\n# RED = (0, 0, 255)\n# BLUE = (255, 0, 0)\n# YELLOW = (255,255,0)\n# FONT = cv2.FONT_HERSHEY_SIMPLEX\n# TEXT_COLOR = (0, 0, 0)\n# # width = 128\n# # height = 256\n#\n#\n# def mask_overlay(img, mask, clip=True):\n# width, height = img.shape[1], img.shape[0]\n# mask = cv2.resize(mask, dsize=(width, height), interpolation=cv2.INTER_CUBIC)\n# if clip:\n# mask = np.clip(mask, 0, 1)\n# mask = (mask*255).astype(np.uint8)\n# else:\n# mask = np.interp(mask, (mask.min(), mask.max()), (0, 255)).astype(np.uint8)\n# mask_color = cv2.applyColorMap(mask, cv2.COLORMAP_JET)\n# mask_color = cv2.cvtColor(mask_color, cv2.COLOR_BGR2RGB)\n# masked_img = cv2.addWeighted(img, 0.5, mask_color, 0.5, 0)\n# return masked_img\n#\n#\n# def show_ranking_grid(query_sample, gallery_topk_samples, config, osp=None):\n# width = 128\n# height = 256\n# samples = [query_sample] + gallery_topk_samples\n#\n# print('start {}'.format(time.time()))\n#\n# plt.close('all')\n# fig = plt.figure(figsize=(100, 66), constrained_layout=True)\n# outer_grid = fig.add_gridspec(len(samples), 1)\n# # outer_grid = plt.GridSpec(len(samples), 1, wspace=1, hspace=1)\n#\n# for row, sample in enumerate(samples):\n# print('row {} {}'.format(row, time.time()))\n# display_sample_on_row(outer_grid[row, 0], sample, row, (width, height))\n#\n# # plt.savefig('/Users/vladimirsomers/Downloads/test_ranking_viz_matplotlib/test_grid_viz_plt_{}.pdf'.format(int(time.time())), format='pdf')\n# print('savefig {}'.format(time.time()))\n# plt.savefig('/Users/vladimirsomers/Downloads/test_ranking_viz_matplotlib/test_grid_viz_plt_{}.jpg'.format(int(time.time())), format='jpg')\n# print('end {}'.format(time.time()))\n# plt.close('all')\n# # plt.show()\n# # plt.waitforbuttonpress()\n#\n#\n# def display_sample_on_row(subplot, sample, row, img_shape):\n# if row == 0:\n# pid, camid, img_path, masks_path, parts_visibility = sample\n# matched, dist_to_query, body_parts_dist_to_query = None, None, None\n# else:\n# pid, camid, img_path, masks_path, parts_visibility, matched, dist_to_query, body_parts_dist_to_query = sample\n#\n# width, height = img_shape\n# masks = read_masks(masks_path)\n# bp_num = masks.shape[0]\n# img = cv2.imread(img_path)\n# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n# img = cv2.resize(img, (width, height))\n# cols = bp_num+1\n#\n# inner_grid = subplot.subgridspec(1, cols)\n# axs = inner_grid.subplots()\n# # plt.subplots_adjust(right=0.8)\n#\n# for col in range(0, cols):\n# if row == 0 and col == 0:\n# img_to_insert = img\n# elif col == 0:\n# border_color = GREEN if matched else RED\n# img_to_insert = make_border(img, border_color, BW)\n# else:\n# bp_idx = col - 1\n# border_color = GREEN if parts_visibility[bp_idx] else RED\n# mask = masks[bp_idx, :, :]\n# img_with_mask_overlay = mask_overlay(img, mask)\n# img_to_insert = make_border(img_with_mask_overlay, border_color, BW)\n#\n# ax = axs[col]\n# ax.imshow(img_to_insert)\n# ax.tick_params(axis='both', which='both', bottom=False, top=False, labelbottom=False, left=False, right=False,\n# labelleft=False)\n# # ax.set_title(\"Id = {}\\n{}\".format(mask_idx, body_part))\n# # ax.set_xlabel('Dist = {}'.format(dist))\n# # for axis in ['top', 'bottom', 'left', 'right']:\n# # ax.spines[axis].set_color(color)\n# # ax.spines[axis].set_linewidth(4)\n#\n#\n# def insert_img_into_grid(grid_img, img, row, col):\n# extra_spacing_h = QUERY_EXTRA_SPACING if row > 0 else GRID_SPACING\n# extra_spacing_w = QUERY_EXTRA_SPACING if col > 0 else GRID_SPACING\n# width, height = img.shape[1], img.shape[0]\n# hs = (row) * height + row * GRID_SPACING + extra_spacing_h\n# he = (row + 1) * height + row * GRID_SPACING + extra_spacing_h\n# ws = (col) * width + col * GRID_SPACING + extra_spacing_w\n# we = (col + 1) * width + col * GRID_SPACING + extra_spacing_w\n# grid_img[hs:he, ws:we, :] = img\n#\n#\n# def make_border(img, border_color, bw):\n# img_b = cv2.copyMakeBorder(\n# img,\n# bw, bw, bw, bw,\n# cv2.BORDER_CONSTANT,\n# value=border_color\n# )\n# img_b = cv2.resize(img_b, (img.shape[1], img.shape[0]))\n# return img_b\n" }, { "path": "torchreid/utils/writer.py", "content": "import datetime\nimport os\nfrom typing import Optional\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom tabulate import tabulate\nfrom . import Logger, visualize_ranking_grid\nfrom .avgmeter import TorchTimeMeter, SingleMeter, \\\n EpochMeter, EpochArrayMeter, LossEpochMetricsMeter\nfrom .distribution import plot_body_parts_pairs_distance_distribution, plot_pairs_distance_distribution\nfrom .engine_state import EngineStateListener\nfrom .tools import perc\nfrom .visualization.embeddings_projection import visualize_embeddings\n\n\nclass Writer(EngineStateListener):\n # TODO Integrate this with Pytorch Lightning\n \"\"\" A class to encapsulate external loggers and writers such as Tensorboard and Allegro ClearML\n \"\"\"\n __main_writer = None # type: Optional[Writer]\n\n @classmethod\n def current_writer(cls):\n # type: () -> Writer\n return cls.__main_writer\n\n def __init__(self, cfg):\n self.cfg = cfg\n self.model_name = cfg.project.start_time + cfg.project.experiment_id\n self.logger = Logger.current_logger()\n\n # running state\n self.is_training = True\n self.batch_debug_freq = cfg.train.batch_debug_freq\n\n # configs\n self.start_eval = cfg.test.start_eval\n self.eval_freq = cfg.train.eval_freq\n self.max_epoch = cfg.train.max_epoch\n\n # init time meters\n self.total_run_timer = TorchTimeMeter(\"total run time\", False)\n self.test_timer = TorchTimeMeter(\"multi_target_test\", False)\n self.epoch_timer = TorchTimeMeter(\"epoch\", False)\n self.batch_timer = TorchTimeMeter(\"batch\")\n self.data_loading_timer = TorchTimeMeter(\"data_loading\", False)\n self.test_batch_timer = TorchTimeMeter(\"test_batch\")\n self.performance_evaluation_timer = TorchTimeMeter(\"performance_evaluation\", False)\n self.feature_extraction_timer = TorchTimeMeter(\"feature_extraction\")\n self.loss_timer = TorchTimeMeter(\"loss_computation\")\n self.optimizer_timer = TorchTimeMeter(\"optimizer_step\")\n\n Writer.__main_writer = self\n\n def init_engine_state(self, engine_state, parts_num):\n self.engine_state = engine_state\n\n # init meters\n self.invalid_pairwise_distances_count = EpochMeter(self.engine_state)\n self.uncomparable_body_parts_pairs_count = EpochMeter(self.engine_state)\n self.invalid_pairs_count_at_test_time = SingleMeter(self.engine_state)\n self.uncomparable_queries_at_test_time = SingleMeter(self.engine_state)\n self.used_body_parts_in_max = EpochArrayMeter(self.engine_state, parts_num)\n self.losses = LossEpochMetricsMeter(engine_state)\n self.loss = EpochMeter(engine_state)\n\n # writer should be the last listener to be called\n self.engine_state.add_listener(self, True)\n\n ########################\n # TRAINING STATS #\n ########################\n\n def report_performance(self, cmc, mAP, ssmd, pxl_acc_avg, name=\"\"):\n self.logger.add_scalar('r1 {}'.format(name), perc(cmc[0]), self.engine_state.epoch)\n self.logger.add_scalar('r5 {}'.format(name), perc(cmc[1]), self.engine_state.epoch)\n self.logger.add_scalar('r10 {}'.format(name), perc(cmc[2]), self.engine_state.epoch)\n self.logger.add_scalar('r20 {}'.format(name), perc(cmc[3]), self.engine_state.epoch)\n self.logger.add_scalar('mAP {}'.format(name), perc(mAP), self.engine_state.epoch)\n self.logger.add_scalar('ssmd {}'.format(name), ssmd, self.engine_state.epoch)\n self.logger.add_scalar('pxl_acc {}'.format(name), pxl_acc_avg, self.engine_state.epoch)\n\n def report_global_performance(self,\n cmc_per_dataset,\n mAP_per_dataset,\n ssmd_per_dataset,\n pxl_acc_per_dataset):\n self.logger.add_text('r1_global', str(cmc_per_dataset[0]))\n self.logger.add_text('r5_global', str(cmc_per_dataset[1]))\n self.logger.add_text('r10_global', str(cmc_per_dataset[2]))\n self.logger.add_text('r20_global', str(cmc_per_dataset[3]))\n self.logger.add_text('mAP_global', str(mAP_per_dataset))\n self.logger.add_text('ssmd_global', str(ssmd_per_dataset))\n self.logger.add_text('pxl_acc_global', str(pxl_acc_per_dataset))\n\n def intermediate_evaluate(self):\n return (self.engine_state.epoch + 1) >= self.start_eval and self.eval_freq > 0 and (\n self.engine_state.epoch + 1) % self.eval_freq == 0 and (self.engine_state.epoch + 1) != self.max_epoch\n\n def update_invalid_pairwise_distances_count(self, batch_pairwise_dist):\n self.invalid_pairwise_distances_count.update((batch_pairwise_dist == float(-1)).sum(), batch_pairwise_dist.nelement())\n\n def update_invalid_part_based_pairwise_distances_count(self, valid_body_part_pairwise_dist_mask):\n self.uncomparable_body_parts_pairs_count.update((valid_body_part_pairwise_dist_mask.nelement() - valid_body_part_pairwise_dist_mask.sum()),\n valid_body_part_pairwise_dist_mask.nelement())\n\n def used_parts_statistics(self, M, body_part_id):\n # count apparition of each body part id, remove diagonal as we don't consider pairs with same id\n used_body_parts_count = torch.bincount(body_part_id.flatten(), minlength=M) - torch.bincount(\n body_part_id.diag(), minlength=M)\n used_body_parts_count = used_body_parts_count / 2 # body parts are counted two times since matrix is symmetric\n self.used_body_parts_in_max.update(used_body_parts_count, np.ones(len(used_body_parts_count))*used_body_parts_count.sum().item())\n\n # def plot_batch_distance_distribution(self, batch_pairwise_dist, labels): # TODO report each epoch and not each batch for wandb performance issue\n # if self.batch_debug_freq > 0 and (self.engine_state.global_step + 1) % self.batch_debug_freq == 0:\n # batch_pairwise_dist = batch_pairwise_dist.detach().cpu().numpy()\n # labels = labels.detach().cpu().numpy()\n # if len(batch_pairwise_dist.shape) == 3:\n # ssmd = plot_body_parts_pairs_distance_distribution(\n # batch_pairwise_dist, labels, labels, \"Training batch\")\n # self.logger.add_scalar(\"SSMD/training batch ssmd\", ssmd, self.engine_state.global_step)\n # else:\n # pos_p_mean, pos_p_std, neg_p_mean, neg_p_std, ssmd = plot_pairs_distance_distribution(\n # batch_pairwise_dist, labels, labels, \"Training batch\")\n # self.logger.add_scalar(\"SSMD/training batch ssmd\", ssmd, self.engine_state.global_step)\n\n # TODO batch plot\n # def report_body_parts_mean_distances(self, distances): # TODO report each epoch and not each batch for wandb performance issue\n # mean_distance_per_body_part = distances.mean(dim=(1, 2))\n # for bp_id, count in enumerate(mean_distance_per_body_part):\n # self.logger.add_scalar(\"Body parts mean distances/bp_{}\".format(bp_id), count, self.engine_state.global_step)\n\n def visualize_triplets(self, images, masks, mask, dist):\n if self.batch_debug_freq > 0 and (self.engine_state.global_step + 1) % self.batch_debug_freq == 0:\n pass\n # TODO\n # np_mask = mask.clone().detach().cpu().numpy()\n # np_dist = dist.clone().detach().cpu().numpy()\n # dist_ap, dist_an = [], []\n # triplets = []\n # for i in range(20):\n # print(\"Computing triplet {}\".format(i))\n # pos_d = np_dist[i]*(np_mask[i])\n # neg_d = np_dist[i]*(np_mask[i] == 0)\n # pos_idx = pos_d.argmax()\n # neg_idx = neg_d.argmax()\n #\n # # instance = (image, masks, id, body_part_id, body_part_name)\n # pos_img = cv2.imread(images[pos_idx])\n # pos_img = cv2.cvtColor(pos_img, cv2.COLOR_BGR2RGB)\n # anc_img = cv2.imread(images[i])\n # anc_img = cv2.cvtColor(anc_img, cv2.COLOR_BGR2RGB)\n # neg_img = cv2.imread(images[neg_idx])\n # neg_img = cv2.cvtColor(neg_img, cv2.COLOR_BGR2RGB)\n #\n # pos = (pos_img, masks[pos_idx][bp], pos_idx, bp)\n # anc = (anc_img, masks[i][bp], i, bp)\n # neg = (neg_img, masks[neg_idx][bp], neg_idx, bp)\n #\n # # pos, anc, neg, pos_dist, neg_dist = triplet\n # triplet = [pos, anc, neg, pos_d[pos_idx], neg_d[neg_idx]]\n #\n # triplets.append(triplet)\n # # triplets = np.repeat(np.array(triplets), 5, axis=0)\n # show_triplet_grid(triplets, self, bp)\n\n ########################\n # TESTING STATS #\n ########################\n\n def qg_pairwise_dist_statistics(self, pairwise_dist, body_part_pairwise_dist, qf_parts_visibility, gf_parts_visibility):\n valid_pairwise_dist_mask = (pairwise_dist != float(-1))\n invalid_pairs_count = (~valid_pairwise_dist_mask).sum()\n self.invalid_pairs_count_at_test_time.update(invalid_pairs_count, valid_pairwise_dist_mask.nelement())\n uncomparable_queries_count = (~valid_pairwise_dist_mask.max(dim=1)[0]).sum()\n self.uncomparable_queries_at_test_time.update(uncomparable_queries_count, valid_pairwise_dist_mask.shape[0])\n\n part_pairwise_dist_numpy = body_part_pairwise_dist.numpy()\n self.qg_body_part_distances_boxplot(part_pairwise_dist_numpy)\n self.qg_body_part_pairs_availability_barplot(part_pairwise_dist_numpy)\n if qf_parts_visibility is not None and gf_parts_visibility is not None:\n qf_parts_visibility = qf_parts_visibility.numpy()\n gf_parts_visibility = gf_parts_visibility.numpy()\n self.qg_body_part_availability_barplot(qf_parts_visibility, gf_parts_visibility)\n self.qg_distribution_of_body_part_availability_histogram(qf_parts_visibility, gf_parts_visibility)\n\n def qg_body_part_distances_boxplot(self, body_part_pairwise_dist):\n histogram = body_part_pairwise_dist.reshape(body_part_pairwise_dist.shape[:-2] + (-1,)).transpose()\n idx_to_keep = np.random.choice(histogram.shape[0], min(2000, histogram.shape[0]), replace=False)\n idx_to_keep = np.concatenate([idx_to_keep, np.argmax(histogram, axis=0), np.argmin(histogram, axis=0)])\n sampled_histogram = histogram[idx_to_keep]\n valid_distances_histogram = [sampled_histogram[sampled_histogram[:, i] != -1, i] for i in range(0, sampled_histogram.shape[1])]\n fig, ax = plt.subplots(figsize=(24, 4))\n ax.boxplot(valid_distances_histogram, notch=True, widths=0.35, labels=range(0, body_part_pairwise_dist.shape[0]))\n ax.set_ylabel('Distance')\n ax.set_xlabel('Body part index')\n ax.set_title('Distance distribution of query-gallery body part pairs')\n fig.tight_layout()\n self.logger.add_figure(\"Query-gallery body part distances boxplot\", fig, self.engine_state.global_step)\n\n def qg_body_part_pairs_availability_barplot(self, body_part_pairwise_dist):\n body_part_pairs_availability = (body_part_pairwise_dist != -1).mean(axis=(1, 2))\n x_labels = range(0, len(body_part_pairs_availability))\n x = np.arange(len(x_labels)) # the label locations\n width = 0.7 # the width of the bars\n fig, ax = plt.subplots(figsize=(24, 4))\n rects = ax.bar(x, body_part_pairs_availability, width)\n # Add some text for x_labels, title and custom x-axis tick x_labels, etc.\n ax.set_ylabel('Availability')\n ax.set_xlabel('Body part index')\n ax.set_title('Query-gallery body parts pairs availability')\n ax.set_yticks(np.arange(0, 1.2, 0.1))\n ax.yaxis.get_major_ticks()[-1].set_visible(False)\n ax.set_xticks(x)\n ax.set_xticklabels(x_labels)\n\n def autolabel(rects):\n \"\"\"Attach a text label above each bar in *rects*, displaying its height.\"\"\"\n for rect in rects:\n height = rect.get_height()\n ax.annotate('{}%'.format(int(height*100)),\n xy=(rect.get_x() + rect.get_width() / 2, height),\n xytext=(0, 3), # 3 points vertical offset\n textcoords=\"offset points\",\n ha='center', va='bottom')\n\n autolabel(rects)\n\n fig.tight_layout()\n self.logger.add_figure(\"Query-gallery body part pairs availability barplot\", fig, self.engine_state.global_step)\n\n def qg_body_part_availability_barplot(self, qf_parts_visibility, gf_parts_visibility):\n qf_mask_availability = qf_parts_visibility.mean(axis=0)\n gf_mask_availability = gf_parts_visibility.mean(axis=0)\n x_labels = range(0, len(qf_mask_availability))\n x = np.arange(len(x_labels)) # the label locations\n width = 0.35 # the width of the bars\n fig, ax = plt.subplots(figsize=(24, 4))\n rects1 = ax.bar(x - width / 2, qf_mask_availability, width, label='Query')\n rects2 = ax.bar(x + width / 2, gf_mask_availability, width, label='Gallery')\n # Add some text for x_labels, title and custom x-axis tick x_labels, etc.\n ax.set_ylabel('Availability')\n ax.set_xlabel('Body part index')\n ax.set_title('Body parts availability for {} query and {} gallery samples'.format(qf_parts_visibility.shape[0], gf_parts_visibility.shape[0]))\n ax.set_yticks(np.arange(0, 1.2, 0.1))\n ax.yaxis.get_major_ticks()[-1].set_visible(False)\n ax.set_xticks(x)\n ax.set_xticklabels(x_labels)\n ax.legend()\n\n def autolabel(rects):\n \"\"\"Attach a text label above each bar in *rects*, displaying its height.\"\"\"\n for rect in rects:\n height = rect.get_height()\n ax.annotate('{}%'.format(int(height*100)),\n xy=(rect.get_x() + rect.get_width() / 2, height),\n xytext=(0, 3), # 3 points vertical offset\n textcoords=\"offset points\",\n ha='center', va='bottom')\n\n # autolabel(rects1)\n # autolabel(rects2)\n fig.tight_layout()\n self.logger.add_figure(\"Query-gallery body part availability barplot\", fig, self.engine_state.global_step)\n\n def qg_distribution_of_body_part_availability_histogram(self, qf_parts_visibility, gf_parts_visibility):\n qf_mask_availability = qf_parts_visibility.sum(axis=1)\n gf_mask_availability = gf_parts_visibility.sum(axis=1)\n\n x = np.arange(gf_parts_visibility.shape[1]+2)\n qf_mask_availability_distribution = np.histogram(qf_mask_availability, bins=x)[0]/len(qf_parts_visibility)\n gf_mask_availability_distribution = np.histogram(gf_mask_availability, bins=x)[0]/len(gf_parts_visibility)\n\n x_labels = np.arange(gf_parts_visibility.shape[1]+1)\n width = 0.35 # the width of the bars\n fig, ax = plt.subplots(figsize=(24, 4))\n ax.bar(x_labels - width / 2, qf_mask_availability_distribution, width, label='Query')\n ax.bar(x_labels + width / 2, gf_mask_availability_distribution, width, label='Gallery')\n # Add some text for x_labels, title and custom x-axis tick x_labels, etc.\n ax.set_ylabel('Samples count')\n ax.set_xlabel('Amount of body parts available')\n ax.set_title('Body parts availability distribution for {} query and {} gallery samples'.format(qf_parts_visibility.shape[0], gf_parts_visibility.shape[0]))\n # ax.set_yticks(np.arange(0, 1.2, 0.1))\n ax.yaxis.get_major_ticks()[-1].set_visible(False)\n ax.set_xticks(x_labels)\n ax.set_xticklabels(x_labels)\n ax.legend()\n\n fig.tight_layout()\n self.logger.add_figure(\"Query-gallery distribution of body part availability histogram\", fig, self.engine_state.global_step)\n\n def visualize_embeddings(self, qf, gf, q_pids, g_pids, test_loader, dataset_name, qf_parts_visibility, gf_parts_visibility, mAP, rank1):\n if self.cfg.test.vis_embedding_projection and not self.engine_state.is_training:\n print(\"Visualizing embeddings projection\")\n visualize_embeddings(qf, gf, q_pids, g_pids, test_loader, dataset_name, qf_parts_visibility, gf_parts_visibility, mAP, rank1)\n\n def visualize_rank(self, test_loader, dataset_name, distmat, save_dir, visrank_topk, visrank_q_idx_list, visrank_count,\n body_parts_distmat, qf_parts_visibility, gf_parts_visibility, q_parts_masks, g_parts_masks, mAP, rank1):\n if self.cfg.test.visrank:\n save_dir = os.path.join(save_dir, 'vis_bp_rank_' + dataset_name)\n visualize_ranking_grid(distmat, body_parts_distmat, test_loader, dataset_name,\n qf_parts_visibility, gf_parts_visibility, q_parts_masks, g_parts_masks, mAP, rank1, save_dir, visrank_topk, visrank_q_idx_list, visrank_count, config=self.cfg)\n\n if self.cfg.test.visrank_per_body_part:\n for bp in range(0, body_parts_distmat.shape[0]):\n qf_part_visibility = None\n if qf_parts_visibility is not None:\n qf_part_visibility = qf_parts_visibility[:, bp:bp+1]\n gf_part_visibility = None\n if gf_parts_visibility is not None:\n gf_part_visibility = gf_parts_visibility[:, bp:bp+1]\n visualize_ranking_grid(body_parts_distmat[bp], body_parts_distmat[bp:bp+1], test_loader, dataset_name,\n qf_part_visibility, gf_part_visibility, q_parts_masks, g_parts_masks, mAP, rank1, save_dir, visrank_topk, visrank_q_idx_list, visrank_count, config=self.cfg, bp_idx=bp)\n\n ########################\n # RUNNING EVENTS #\n ########################\n\n def training_started(self):\n self.report_performance([0, 0, 0, 0], 0, 0, 0)\n\n def epoch_started(self):\n self.logger.add_scalar('Other/epoch', self.engine_state.epoch, self.engine_state.global_step)\n self.logger.add_scalar('Other/batch', self.engine_state.batch, self.engine_state.global_step)\n self.logger.add_scalar('Other/iteration', self.engine_state.global_step, self.engine_state.global_step)\n\n def epoch_completed(self):\n eta_seconds = (self.max_epoch - (self.engine_state.epoch + 1)) * self.epoch_timer.avg / 1000\n eta_str = str(datetime.timedelta(seconds=int(eta_seconds)))\n\n print(\n 'epoch: [{0}/{1} e][{2} b]\\t'\n 'eta {eta}\\t'\n 'lr {lr:.8f}\\t'\n 'loss {loss:.3f}\\t'\n '{losses}'.format(\n self.engine_state.epoch + 1,\n self.max_epoch,\n self.engine_state.batch,\n eta=eta_str,\n lr=self.engine_state.lr,\n loss=self.loss.epoch_ratio(self.engine_state.epoch),\n losses=self.losses.summary(self.engine_state.epoch)\n )\n )\n\n for name, dict in self.losses.meters.items():\n for key, meter in dict.items():\n self.logger.add_scalar('Loss/' + name + \"_\" + key + '_avg', meter.mean[self.engine_state.epoch], self.engine_state.epoch)\n\n # self.logger.add_scalar(\"Training/Trivial triplets in batch\", self.zero_losses_count.epoch_ratio(self.engine_state.epoch), self.engine_state.epoch)\n\n if not self.used_body_parts_in_max.is_empty:\n for bp_id, bp_ratio in enumerate(self.used_body_parts_in_max.epoch_ratio(self.engine_state.epoch)):\n self.logger.add_scalar(\"Used body parts in training/bp {}\".format(bp_id), bp_ratio, self.engine_state.epoch)\n\n def training_completed(self):\n print(\"Training completed\")\n # TODO fix metrics affected by loss refactor\n print(\"Average image pairs that couldn't be compared within one batch: {}%\".format(perc(self.invalid_pairwise_distances_count.total_ratio())))\n print(\"Average body part pairs that couldn't be compared within one batch: {}%\".format(perc(self.uncomparable_body_parts_pairs_count.total_ratio())))\n # print(\"Average valid triplets during one epoch: {}%\".format(perc(self.valid_triplets_mask_count.total_ratio())))\n self.display_used_body_parts()\n\n def test_completed(self):\n print(\"Test completed\")\n if not self.invalid_pairs_count_at_test_time.is_empty:\n print(\"Amount of pairs query-gallery that couldn't be compared: {}%\".format(perc(self.invalid_pairs_count_at_test_time.ratio(), 3)))\n if not self.uncomparable_queries_at_test_time.is_empty:\n print(\"Amount of queries that couldn't be compared to any gallery sample: {}%\".format(perc(self.uncomparable_queries_at_test_time.ratio(), 3)))\n\n def run_completed(self):\n timer_meters = [\n self.total_run_timer,\n self.epoch_timer,\n self.batch_timer,\n self.data_loading_timer,\n self.feature_extraction_timer,\n self.loss_timer,\n self.optimizer_timer,\n self.performance_evaluation_timer,\n self.test_timer,\n self.test_batch_timer,\n ]\n table = []\n for time_meter in timer_meters:\n table.append([time_meter.name, time_meter.average_time(), time_meter.total_time(), time_meter.count])\n headers = [\"Time metric name\", \"Average\", \"Total\", \"Count\"]\n\n print(tabulate(table, headers, tablefmt=\"fancy_grid\"))\n\n ########################\n # UTILS #\n ########################\n\n def display_used_body_parts(self):\n if self.used_body_parts_in_max.is_empty:\n return\n\n # plot histogram\n body_parts_used_for_training = self.used_body_parts_in_max.total_ratio()\n x_labels = range(0, len(body_parts_used_for_training))\n x = np.arange(len(x_labels)) # the label locations\n width = 0.7 # the width of the bars\n fig, ax = plt.subplots(figsize=(24, 4))\n rects = ax.bar(x, body_parts_used_for_training, width)\n # Add some text for x_labels, title and custom x-axis tick x_labels, etc.\n ax.set_ylabel('Selection percentage')\n ax.set_xlabel('Body part index')\n ax.set_title('Body parts used for training')\n ax.set_yticks(np.arange(0, 1.2, 0.1))\n ax.yaxis.get_major_ticks()[-1].set_visible(False)\n ax.set_xticks(x)\n ax.set_xticklabels(x_labels)\n\n def autolabel(rects):\n \"\"\"Attach a text label above each bar in *rects*, displaying its height.\"\"\"\n for rect in rects:\n height = rect.get_height()\n ax.annotate('{}%'.format(np.around(height*100, 2)),\n xy=(rect.get_x() + rect.get_width() / 2, height),\n xytext=(0, 3), # 3 points vertical offset\n textcoords=\"offset points\",\n ha='center', va='bottom')\n\n autolabel(rects)\n\n fig.tight_layout()\n self.logger.add_figure(\"Body parts used for training\", fig, self.engine_state.global_step)\n" } ]