add test code

.gitignore

@@ -0,0 +1,141 @@
+log/
+output/
+.vscode/
+workspace/
+tmp_occupy_memory_saves/
+run*.sh
+py-thin-plate-spline
+wandb/
+pretrain/
+Pytorch-Correlation-extension/
+result
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

LICENSE

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

dataset/range_transform.py

@@ -0,0 +1,44 @@
+import torchvision.transforms as transforms
+import util.functional as F
+import numpy as np
+from skimage import color
+
+im_mean = (124, 116, 104)
+
+im_normalization = transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406],
+                    std=[0.229, 0.224, 0.225]
+                )
+
+inv_im_trans = transforms.Normalize(
+                mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
+                std=[1/0.229, 1/0.224, 1/0.225])
+
+# tensor l[-1, 1]  ab[-1,    1]
+# numpy  l[0 100]  ab[-127 128]
+# transforms.Normalize: x_new = (x-mean) / std
+inv_lll2rgb_trans = transforms.Normalize(
+                mean=[-1, 0, 0],
+                std=[1/50., 1/110., 1/110.])
+
+im_rgb2lab_normalization = transforms.Normalize(
+                mean=[50, 0, 0],
+                std=[50, 110, 110])
+
+class ToTensor(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, inputs):
+        return F.to_mytensor(inputs)
+
+class RGB2Lab(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, inputs):
+        # default return float64
+        # return color.rgb2lab(inputs)
+
+        # return float32
+        return np.float32(color.rgb2lab(inputs))

dataset/reseed.py

@@ -0,0 +1,6 @@
+import torch
+import random
+
+def reseed(seed):
+    random.seed(seed)
+    torch.manual_seed(seed)

dataset/static_dataset.py

@@ -0,0 +1,179 @@
+import os
+from os import path
+
+import torch
+from torch.utils.data.dataset import Dataset
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+from PIL import Image
+import numpy as np
+
+from dataset.range_transform import im_normalization, im_mean
+from dataset.tps import random_tps_warp
+from dataset.reseed import reseed
+
+
+class StaticTransformDataset(Dataset):
+    """
+    Generate pseudo VOS data by applying random transforms on static images.
+    Single-object only.
+
+    Method 0 - FSS style (class/1.jpg class/1.png)
+    Method 1 - Others style (XXX.jpg XXX.png)
+    """
+    def __init__(self, parameters, num_frames=3, max_num_obj=1):
+        self.num_frames = num_frames
+        self.max_num_obj = max_num_obj
+
+        self.im_list = []
+        for parameter in parameters:
+            root, method, multiplier = parameter
+            if method == 0:
+                # Get images
+                classes = os.listdir(root)
+                for c in classes:
+                    imgs = os.listdir(path.join(root, c))
+                    jpg_list = [im for im in imgs if 'jpg' in im[-3:].lower()]
+
+                    joint_list = [path.join(root, c, im) for im in jpg_list]
+                    self.im_list.extend(joint_list * multiplier)
+
+            elif method == 1:
+                self.im_list.extend([path.join(root, im) for im in os.listdir(root) if '.jpg' in im] * multiplier)
+
+        print(f'{len(self.im_list)} images found.')
+
+        # These set of transform is the same for im/gt pairs, but different among the 3 sampled frames
+        self.pair_im_lone_transform = transforms.Compose([
+            transforms.ColorJitter(0.1, 0.05, 0.05, 0), # No hue change here as that's not realistic
+        ])
+
+        self.pair_im_dual_transform = transforms.Compose([
+            transforms.RandomAffine(degrees=20, scale=(0.9,1.1), shear=10, interpolation=InterpolationMode.BICUBIC, fill=im_mean),
+            transforms.Resize(384, InterpolationMode.BICUBIC),
+            transforms.RandomCrop((384, 384), pad_if_needed=True, fill=im_mean),
+        ])
+
+        self.pair_gt_dual_transform = transforms.Compose([
+            transforms.RandomAffine(degrees=20, scale=(0.9,1.1), shear=10, interpolation=InterpolationMode.BICUBIC, fill=0),
+            transforms.Resize(384, InterpolationMode.NEAREST),
+            transforms.RandomCrop((384, 384), pad_if_needed=True, fill=0),
+        ])
+
+
+        # These transform are the same for all pairs in the sampled sequence
+        self.all_im_lone_transform = transforms.Compose([
+            transforms.ColorJitter(0.1, 0.05, 0.05, 0.05),
+            transforms.RandomGrayscale(0.05),
+        ])
+
+        self.all_im_dual_transform = transforms.Compose([
+            transforms.RandomAffine(degrees=0, scale=(0.8, 1.5), fill=im_mean),
+            transforms.RandomHorizontalFlip(),
+        ])
+
+        self.all_gt_dual_transform = transforms.Compose([
+            transforms.RandomAffine(degrees=0, scale=(0.8, 1.5), fill=0),
+            transforms.RandomHorizontalFlip(),
+        ])
+
+        # Final transform without randomness
+        self.final_im_transform = transforms.Compose([
+            transforms.ToTensor(),
+            im_normalization,
+        ])
+
+        self.final_gt_transform = transforms.Compose([
+            transforms.ToTensor(),
+        ])
+
+    def _get_sample(self, idx):
+        im = Image.open(self.im_list[idx]).convert('RGB')
+        gt = Image.open(self.im_list[idx][:-3]+'png').convert('L')
+
+        sequence_seed = np.random.randint(2147483647)
+
+        images = []
+        masks = []
+        for _ in range(self.num_frames):
+            reseed(sequence_seed)
+            this_im = self.all_im_dual_transform(im)
+            this_im = self.all_im_lone_transform(this_im)
+            reseed(sequence_seed)
+            this_gt = self.all_gt_dual_transform(gt)
+
+            pairwise_seed = np.random.randint(2147483647)
+            reseed(pairwise_seed)
+            this_im = self.pair_im_dual_transform(this_im)
+            this_im = self.pair_im_lone_transform(this_im)
+            reseed(pairwise_seed)
+            this_gt = self.pair_gt_dual_transform(this_gt)
+
+            # Use TPS only some of the times
+            # Not because TPS is bad -- just that it is too slow and I need to speed up data loading
+            if np.random.rand() < 0.33:
+                this_im, this_gt = random_tps_warp(this_im, this_gt, scale=0.02)
+
+            this_im = self.final_im_transform(this_im)
+            this_gt = self.final_gt_transform(this_gt)
+
+            images.append(this_im)
+            masks.append(this_gt)
+
+        images = torch.stack(images, 0)
+        masks = torch.stack(masks, 0)
+
+        return images, masks.numpy()
+
+    def __getitem__(self, idx):
+        additional_objects = np.random.randint(self.max_num_obj)
+        indices = [idx, *np.random.randint(self.__len__(), size=additional_objects)]
+
+        merged_images = None
+        merged_masks = np.zeros((self.num_frames, 384, 384), dtype=np.int)
+
+        for i, list_id in enumerate(indices):
+            images, masks = self._get_sample(list_id)
+            if merged_images is None:
+                merged_images = images
+            else:
+                merged_images = merged_images*(1-masks) + images*masks
+            merged_masks[masks[:,0]>0.5] = (i+1)
+
+        masks = merged_masks
+
+        labels = np.unique(masks[0])
+        # Remove background
+        labels = labels[labels!=0]
+        target_objects = labels.tolist()
+
+        # Generate one-hot ground-truth
+        cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int)
+        first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int)
+        for i, l in enumerate(target_objects):
+            this_mask = (masks==l)
+            cls_gt[this_mask] = i+1
+            first_frame_gt[0,i] = (this_mask[0])
+        cls_gt = np.expand_dims(cls_gt, 1)
+
+        info = {}
+        info['name'] = self.im_list[idx]
+        info['num_objects'] = max(1, len(target_objects))
+
+        # 1 if object exist, 0 otherwise
+        selector = [1 if i < info['num_objects'] else 0 for i in range(self.max_num_obj)]
+        selector = torch.FloatTensor(selector)
+
+        data = {
+            'rgb': merged_images,
+            'first_frame_gt': first_frame_gt,
+            'cls_gt': cls_gt,
+            'selector': selector,
+            'info': info
+        }
+
+        return data
+
+
+    def __len__(self):
+        return len(self.im_list)

dataset/tps.py

@@ -0,0 +1,37 @@
+import numpy as np
+from PIL import Image
+import cv2
+import thinplate as tps
+
+cv2.setNumThreads(0)
+
+def pick_random_points(h, w, n_samples):
+    y_idx = np.random.choice(np.arange(h), size=n_samples, replace=False)
+    x_idx = np.random.choice(np.arange(w), size=n_samples, replace=False)
+    return y_idx/h, x_idx/w
+
+
+def warp_dual_cv(img, mask, c_src, c_dst):
+    dshape = img.shape
+    theta = tps.tps_theta_from_points(c_src, c_dst, reduced=True)
+    grid = tps.tps_grid(theta, c_dst, dshape)
+    mapx, mapy = tps.tps_grid_to_remap(grid, img.shape)
+    return cv2.remap(img, mapx, mapy, cv2.INTER_LINEAR), cv2.remap(mask, mapx, mapy, cv2.INTER_NEAREST)
+
+
+def random_tps_warp(img, mask, scale, n_ctrl_pts=12):
+    """
+    Apply a random TPS warp of the input image and mask
+    Uses randomness from numpy
+    """
+    img = np.asarray(img)
+    mask = np.asarray(mask)
+
+    h, w = mask.shape
+    points = pick_random_points(h, w, n_ctrl_pts)
+    c_src = np.stack(points, 1)
+    c_dst = c_src + np.random.normal(scale=scale, size=c_src.shape)
+    warp_im, warp_gt = warp_dual_cv(img, mask, c_src, c_dst)
+
+    return Image.fromarray(warp_im), Image.fromarray(warp_gt)
+

dataset/util.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+def all_to_onehot(masks, labels):
+    if len(masks.shape) == 3:
+        Ms = np.zeros((len(labels), masks.shape[0], masks.shape[1], masks.shape[2]), dtype=np.uint8)
+    else:
+        Ms = np.zeros((len(labels), masks.shape[0], masks.shape[1]), dtype=np.uint8)
+
+    for ni, l in enumerate(labels):
+        Ms[ni] = (masks == l).astype(np.uint8)
+        
+    return Ms

dataset/vos_dataset.py

@@ -0,0 +1,210 @@
+import os
+from os import path, replace
+
+import torch
+from torch.utils.data.dataset import Dataset
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+from PIL import Image
+import numpy as np
+
+from dataset.range_transform import im_normalization, im_mean, im_rgb2lab_normalization, ToTensor, RGB2Lab
+from dataset.reseed import reseed
+
+import util.functional as F
+
+class VOSDataset_221128_TransColorization_batch(Dataset):
+    """
+    Works for DAVIS/YouTubeVOS/BL30K training
+    For each sequence:
+    - Pick three frames
+    - Pick two objects
+    - Apply some random transforms that are the same for all frames
+    - Apply random transform to each of the frame
+    - The distance between frames is controlled
+    """
+    def __init__(self, im_root, gt_root, max_jump, is_bl, subset=None, num_frames=3, max_num_obj=2, finetune=False):
+        self.im_root = im_root
+        self.gt_root = gt_root
+        self.max_jump = max_jump
+        self.is_bl = is_bl
+        self.num_frames = num_frames
+        self.max_num_obj = max_num_obj
+
+        self.videos = []
+        self.frames = {}
+        vid_list = sorted(os.listdir(self.im_root))
+        # Pre-filtering
+        for vid in vid_list:
+            if subset is not None:
+                if vid not in subset:
+                    continue
+            frames = sorted(os.listdir(os.path.join(self.im_root, vid)))
+            if len(frames) < num_frames:
+                continue
+            self.frames[vid] = frames
+            self.videos.append(vid)
+
+        print('%d out of %d videos accepted in %s.' % (len(self.videos), len(vid_list), im_root))
+
+        # These set of transform is the same for im/gt pairs, but different among the 3 sampled frames
+        self.pair_im_lone_transform = transforms.Compose([
+            transforms.ColorJitter(0.01, 0.01, 0.01, 0),
+        ])
+
+        self.pair_im_dual_transform = transforms.Compose([
+            transforms.RandomAffine(degrees=0 if finetune or self.is_bl else 15, shear=0 if finetune or self.is_bl else 10, interpolation=InterpolationMode.BILINEAR, fill=im_mean),
+        ])
+
+        self.pair_gt_dual_transform = transforms.Compose([
+            transforms.RandomAffine(degrees=0 if finetune or self.is_bl else 15, shear=0 if finetune or self.is_bl else 10, interpolation=InterpolationMode.NEAREST, fill=0),
+        ])
+
+        # These transform are the same for all pairs in the sampled sequence
+        self.all_im_lone_transform = transforms.Compose([
+            transforms.ColorJitter(0.1, 0.03, 0.03, 0),
+            # transforms.RandomGrayscale(0.05),
+        ])
+
+        patchsz = 448 # 224
+        self.all_im_dual_transform = transforms.Compose([
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomResizedCrop((patchsz, patchsz), scale=(0.36,1.00), interpolation=InterpolationMode.BILINEAR)
+        ])
+
+        self.all_gt_dual_transform = transforms.Compose([
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomResizedCrop((patchsz, patchsz), scale=(0.36,1.00), interpolation=InterpolationMode.NEAREST)
+        ])
+
+        # Final transform without randomness
+        self.final_im_transform = transforms.Compose([
+            RGB2Lab(),
+            ToTensor(),
+            im_rgb2lab_normalization,
+        ])
+
+    def __getitem__(self, idx):
+        video = self.videos[idx]
+        info = {}
+        info['name'] = video
+
+        vid_im_path = path.join(self.im_root, video)
+        vid_gt_path = path.join(self.gt_root, video)
+        frames = self.frames[video]
+
+        trials = 0
+        while trials < 5:
+            info['frames'] = [] # Appended with actual frames
+
+            num_frames = self.num_frames
+            length = len(frames)
+            this_max_jump = min(len(frames), self.max_jump)
+
+            # iterative sampling
+            frames_idx = [np.random.randint(length)]
+            acceptable_set = set(range(max(0, frames_idx[-1]-this_max_jump), min(length, frames_idx[-1]+this_max_jump+1))).difference(set(frames_idx))
+            while(len(frames_idx) < num_frames):
+                idx = np.random.choice(list(acceptable_set))
+                frames_idx.append(idx)
+                new_set = set(range(max(0, frames_idx[-1]-this_max_jump), min(length, frames_idx[-1]+this_max_jump+1)))
+                acceptable_set = acceptable_set.union(new_set).difference(set(frames_idx))
+
+            frames_idx = sorted(frames_idx)
+            if np.random.rand() < 0.5:
+                # Reverse time
+                frames_idx = frames_idx[::-1]
+
+            sequence_seed = np.random.randint(2147483647)
+            images = []
+            masks = []
+            target_objects = []
+            for f_idx in frames_idx:
+                jpg_name = frames[f_idx]
+                png_name = jpg_name.replace('.jpg', '.png')
+                info['frames'].append(jpg_name)
+
+                reseed(sequence_seed)
+                this_im = Image.open(path.join(vid_im_path, jpg_name)).convert('RGB')
+                this_im = self.all_im_dual_transform(this_im)
+                this_im = self.all_im_lone_transform(this_im)
+
+                reseed(sequence_seed)
+                this_gt = Image.open(path.join(vid_gt_path, png_name)).convert('P')
+                this_gt = self.all_gt_dual_transform(this_gt)
+
+                pairwise_seed = np.random.randint(2147483647)
+                reseed(pairwise_seed)
+                this_im = self.pair_im_dual_transform(this_im)
+                this_im = self.pair_im_lone_transform(this_im)
+
+                reseed(pairwise_seed)
+                this_gt = self.pair_gt_dual_transform(this_gt)
+
+                this_im = self.final_im_transform(this_im)
+                # print('1', torch.max(this_im[:1,:,:]), torch.min(this_im[:1,:,:]))
+                # print('2', torch.max(this_im[1:3,:,:]), torch.min(this_im[1:3,:,:]))
+                # print('3', torch.max(this_im), torch.min(this_im));assert 1==0
+                # print(this_im.size());assert 1==0
+
+                this_gt = np.array(this_gt)
+
+                this_im_l = this_im[:1,:,:]
+                this_im_ab = this_im[1:3,:,:]
+                # print(this_im_l.size(), this_im_ab.size());assert 1==0
+
+                # images.append(this_im_l)
+                # masks.append(this_im_ab)
+
+                this_im_lll = this_im_l.repeat(3,1,1)
+                images.append(this_im_lll)
+                masks.append(this_im_ab)
+
+            images = torch.stack(images, 0)
+            # print(images.size());assert 1==0
+
+            # target_objects = labels.tolist()
+            break
+
+        first_frame_gt = masks[0].unsqueeze(0)
+        # print(first_frame_gt.size());assert 1==0
+
+        info['num_objects'] = 2
+
+        masks = np.stack(masks, 0)
+        # print(np.shape(masks));assert 1==0
+
+
+        cls_gt = masks
+
+        # # Generate one-hot ground-truth
+        # cls_gt = np.zeros((self.num_frames, 384, 384), dtype=np.int)
+        # first_frame_gt = np.zeros((1, self.max_num_obj, 384, 384), dtype=np.int)
+        # for i, l in enumerate(target_objects):
+        #     this_mask = (masks==l)
+        #     cls_gt[this_mask] = i+1
+        #     first_frame_gt[0,i] = (this_mask[0])
+        # cls_gt = np.expand_dims(cls_gt, 1)
+
+        # 1 if object exist, 0 otherwise
+        selector = [1 if i < info['num_objects'] else 0 for i in range(self.max_num_obj)]
+
+        # print(info['num_objects'], self.max_num_obj, selector);assert 1==0
+
+        selector = torch.FloatTensor(selector)
+        
+        # print(images.size(), np.shape(first_frame_gt), np.shape(cls_gt));assert 1==0
+        ### torch.Size([8, 3, 384, 384]) torch.Size([1, 2, 384, 384]) (8, 2, 384, 384)
+
+        data = {
+            'rgb': images,
+            'first_frame_gt': first_frame_gt,
+            'cls_gt': cls_gt,
+            'selector': selector,
+            'info': info,
+        }
+
+        return data
+
+    def __len__(self):
+        return len(self.videos)

inference/data/mask_mapper.py

@@ -0,0 +1,67 @@
+import numpy as np
+import torch
+
+from dataset.util import all_to_onehot
+
+
+class MaskMapper:
+    """
+    This class is used to convert a indexed-mask to a one-hot representation.
+    It also takes care of remapping non-continuous indices
+    It has two modes:
+        1. Default. Only masks with new indices are supposed to go into the remapper.
+        This is also the case for YouTubeVOS.
+        i.e., regions with index 0 are not "background", but "don't care".
+
+        2. Exhaustive. Regions with index 0 are considered "background".
+        Every single pixel is considered to be "labeled".
+    """
+    def __init__(self):
+        self.labels = []
+        self.remappings = {}
+
+        # if coherent, no mapping is required
+        self.coherent = True
+
+    def convert_mask(self, mask, exhaustive=False):
+        # mask is in index representation, H*W numpy array
+        labels = np.unique(mask).astype(np.uint8)
+        labels = labels[labels!=0].tolist()
+
+        new_labels = list(set(labels) - set(self.labels))
+        # print('new_labels', new_labels) # [255]
+        if not exhaustive:
+            assert len(new_labels) == len(labels), 'Old labels found in non-exhaustive mode'
+
+        # add new remappings
+        for i, l in enumerate(new_labels):
+            self.remappings[l] = i+len(self.labels)+1
+            if self.coherent and i+len(self.labels)+1 != l:
+                self.coherent = False
+
+        if exhaustive:
+            new_mapped_labels = range(1, len(self.labels)+len(new_labels)+1)
+        else:
+            if self.coherent:
+                new_mapped_labels = new_labels
+            else:
+                new_mapped_labels = range(len(self.labels)+1, len(self.labels)+len(new_labels)+1)
+                # print(list(new_mapped_labels));assert 1==0 # [1]
+
+        self.labels.extend(new_labels)
+        # print(self.labels);assert 1==0 # [255]
+        mask = torch.from_numpy(all_to_onehot(mask, self.labels)).float()
+
+        # mask num_objects*H*W; new_mapped_labels: [num_objects]
+        return mask, new_mapped_labels
+
+
+    def remap_index_mask(self, mask):
+        # mask is in index representation, H*W numpy array
+        if self.coherent:
+            return mask
+
+        new_mask = np.zeros_like(mask)
+        for l, i in self.remappings.items():
+            new_mask[mask==i] = l
+        return new_mask

inference/data/test_datasets.py

@@ -0,0 +1,29 @@
+import os
+from os import path
+import json
+
+from inference.data.video_reader import VideoReader_221128_TransColorization
+
+class DAVISTestDataset_221128_TransColorization_batch:
+    def __init__(self, data_root, imset='2017/val.txt', size=-1):
+        self.image_dir = data_root
+        self.mask_dir = imset
+        self.size_dir = data_root
+        self.size = size
+
+        self.vid_list =  [clip_name for clip_name in sorted(os.listdir(data_root)) if clip_name != '.DS_Store']
+
+        # print(lst, len(lst), self.vid_list, self.vid_list_DAVIS2016, path.join(data_root, 'ImageSets', imset));assert 1==0
+
+    def get_datasets(self):
+        for video in self.vid_list:
+            # print(self.image_dir, video, path.join(self.image_dir, video));assert 1==0
+            yield VideoReader_221128_TransColorization(video, 
+                path.join(self.image_dir, video), 
+                path.join(self.mask_dir, video),
+                size=self.size,
+                size_dir=path.join(self.size_dir, video),
+            )
+
+    def __len__(self):
+        return len(self.vid_list)

inference/data/video_reader.py

@@ -0,0 +1,107 @@
+import os
+from os import path
+
+from torch.utils.data.dataset import Dataset
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+import torch.nn.functional as Ff
+from PIL import Image
+import numpy as np
+
+from dataset.range_transform import im_normalization, im_rgb2lab_normalization, ToTensor, RGB2Lab
+
+class VideoReader_221128_TransColorization(Dataset):
+    """
+    This class is used to read a video, one frame at a time
+    """
+    def __init__(self, vid_name, image_dir, mask_dir, size=-1, to_save=None, use_all_mask=False, size_dir=None):
+        """
+        image_dir - points to a directory of jpg images
+        mask_dir - points to a directory of png masks
+        size - resize min. side to size. Does nothing if <0.
+        to_save - optionally contains a list of file names without extensions 
+            where the segmentation mask is required
+        use_all_mask - when true, read all available mask in mask_dir.
+            Default false. Set to true for YouTubeVOS validation.
+        """
+        self.vid_name = vid_name
+        self.image_dir = image_dir
+        self.mask_dir = mask_dir
+        self.to_save = to_save
+        self.use_all_mask = use_all_mask
+        # print('use_all_mask', use_all_mask);assert 1==0
+        if size_dir is None:
+            self.size_dir = self.image_dir
+        else:
+            self.size_dir = size_dir
+
+        self.frames = [img for img in sorted(os.listdir(self.image_dir)) if img.endswith('.jpg') or img.endswith('.png')]
+        self.palette = Image.open(path.join(mask_dir, sorted(os.listdir(mask_dir))[0])).getpalette()
+        self.first_gt_path = path.join(self.mask_dir, sorted(os.listdir(self.mask_dir))[0])
+        self.suffix = self.first_gt_path.split('.')[-1]
+
+        if size < 0:
+            self.im_transform = transforms.Compose([
+                RGB2Lab(),
+                ToTensor(),
+                im_rgb2lab_normalization,
+            ])
+        else:
+            self.im_transform = transforms.Compose([
+                transforms.ToTensor(),
+                im_normalization,
+                transforms.Resize(size, interpolation=InterpolationMode.BILINEAR),
+            ])
+        self.size = size
+
+
+    def __getitem__(self, idx):
+        frame = self.frames[idx]
+        info = {}
+        data = {}
+        info['frame'] = frame
+        info['vid_name'] = self.vid_name
+        info['save'] = (self.to_save is None) or (frame[:-4] in self.to_save)
+
+        im_path = path.join(self.image_dir, frame)
+        img = Image.open(im_path).convert('RGB')
+
+        if self.image_dir == self.size_dir:
+            shape = np.array(img).shape[:2]
+        else:
+            size_path = path.join(self.size_dir, frame)
+            size_im = Image.open(size_path).convert('RGB')
+            shape = np.array(size_im).shape[:2]
+
+        gt_path = path.join(self.mask_dir, sorted(os.listdir(self.mask_dir))[idx]) if idx < len(os.listdir(self.mask_dir)) else None 
+
+        img = self.im_transform(img)
+        img_l = img[:1,:,:]
+        img_lll = img_l.repeat(3,1,1)
+
+        load_mask = self.use_all_mask or (gt_path == self.first_gt_path)
+        if load_mask and path.exists(gt_path):
+            mask = Image.open(gt_path).convert('RGB')
+            mask = self.im_transform(mask)
+            mask_ab = mask[1:3,:,:]
+            data['mask'] = mask_ab
+
+        info['shape'] = shape
+        info['need_resize'] = not (self.size < 0)
+        data['rgb'] = img_lll
+        data['info'] = info
+
+        return data
+
+    def resize_mask(self, mask):
+        # mask transform is applied AFTER mapper, so we need to post-process it in eval.py
+        h, w = mask.shape[-2:]
+        min_hw = min(h, w)
+        return Ff.interpolate(mask, (int(h/min_hw*self.size), int(w/min_hw*self.size)), 
+                    mode='nearest')
+
+    def get_palette(self):
+        return self.palette
+
+    def __len__(self):
+        return len(self.frames)

inference/inference_core.py

@@ -0,0 +1,111 @@
+from inference.memory_manager import MemoryManager
+from model.network import ColorMNet
+from model.aggregate import aggregate
+
+from util.tensor_util import pad_divide_by, unpad
+import torch
+
+class InferenceCore:
+    def __init__(self, network:ColorMNet, config):
+        self.config = config
+        self.network = network
+        self.mem_every = config['mem_every']
+        self.deep_update_every = config['deep_update_every']
+        self.enable_long_term = config['enable_long_term']
+
+        # if deep_update_every < 0, synchronize deep update with memory frame
+        self.deep_update_sync = (self.deep_update_every < 0)
+
+        self.clear_memory()
+        self.all_labels = None
+
+        self.last_ti_key = None
+        self.last_ti_value = None
+
+    def clear_memory(self):
+        self.curr_ti = -1
+        self.last_mem_ti = 0
+        if not self.deep_update_sync:
+            self.last_deep_update_ti = -self.deep_update_every
+        self.memory = MemoryManager(config=self.config)
+
+    def update_config(self, config):
+        self.mem_every = config['mem_every']
+        self.deep_update_every = config['deep_update_every']
+        self.enable_long_term = config['enable_long_term']
+
+        # if deep_update_every < 0, synchronize deep update with memory frame
+        self.deep_update_sync = (self.deep_update_every < 0)
+        self.memory.update_config(config)
+
+    def set_all_labels(self, all_labels):
+        # self.all_labels = [l.item() for l in all_labels]
+        self.all_labels = all_labels
+
+    def step(self, image, mask=None, valid_labels=None, end=False):
+        # image: 3*H*W
+        # mask: num_objects*H*W or None
+        self.curr_ti += 1
+        divide_by = 112 # 16
+        image, self.pad = pad_divide_by(image, divide_by)
+        image = image.unsqueeze(0) # add the batch dimension
+
+        is_mem_frame = ((self.curr_ti-self.last_mem_ti >= self.mem_every) or (mask is not None)) and (not end)
+        need_segment = (self.curr_ti > 0) and ((valid_labels is None) or (len(self.all_labels) != len(valid_labels)))
+        is_deep_update = (
+            (self.deep_update_sync and is_mem_frame) or  # synchronized
+            (not self.deep_update_sync and self.curr_ti-self.last_deep_update_ti >= self.deep_update_every) # no-sync
+        ) and (not end)
+        is_normal_update = (not self.deep_update_sync or not is_deep_update) and (not end)
+
+        key, shrinkage, selection, f16, f8, f4 = self.network.encode_key(image, 
+                                                    need_ek=(self.enable_long_term or need_segment), 
+                                                    need_sk=is_mem_frame)
+        multi_scale_features = (f16, f8, f4)
+
+        # segment the current frame is needed
+        if need_segment:
+            memory_readout = self.memory.match_memory(key, selection).unsqueeze(0)
+
+            # short term memory 
+            batch, num_objects, value_dim, h, w = self.last_ti_value.shape
+            last_ti_value = self.last_ti_value.flatten(start_dim=1, end_dim=2)
+            memory_value_short, _ = self.network.short_term_attn(key, self.last_ti_key, last_ti_value, None, key.shape[-2:])
+            memory_value_short = memory_value_short.permute(1, 2, 0).view(batch, num_objects, value_dim, h, w)
+            memory_readout += memory_value_short
+
+            hidden, _, pred_prob_with_bg = self.network.segment(multi_scale_features, memory_readout, 
+                                    self.memory.get_hidden(), h_out=is_normal_update, strip_bg=False)
+            # remove batch dim
+            pred_prob_with_bg = pred_prob_with_bg[0]
+            pred_prob_no_bg = pred_prob_with_bg
+            if is_normal_update:
+                self.memory.set_hidden(hidden)
+        else:
+            pred_prob_no_bg = pred_prob_with_bg = None
+
+        # use the input mask if any
+        if mask is not None:
+            mask, _ = pad_divide_by(mask, divide_by)
+
+            pred_prob_with_bg = mask
+
+            self.memory.create_hidden_state(2, key)
+
+        # save as memory if needed
+        if is_mem_frame:
+            value, hidden = self.network.encode_value(image, f16, self.memory.get_hidden(), 
+                                    pred_prob_with_bg.unsqueeze(0), is_deep_update=is_deep_update)
+
+            self.memory.add_memory(key, shrinkage, value, self.all_labels, 
+                                    selection=selection if self.enable_long_term else None)
+            self.last_mem_ti = self.curr_ti
+
+            self.last_ti_key = key
+            self.last_ti_value = value
+
+            if is_deep_update:
+                self.memory.set_hidden(hidden)
+                self.last_deep_update_ti = self.curr_ti
+                
+        return unpad(pred_prob_with_bg, self.pad)

inference/interact/fbrs/LICENSE

@@ -0,0 +1,373 @@
+Mozilla Public License Version 2.0
+==================================
+
+1. Definitions
+--------------
+
+1.1. "Contributor"
+    means each individual or legal entity that creates, contributes to
+    the creation of, or owns Covered Software.
+
+1.2. "Contributor Version"
+    means the combination of the Contributions of others (if any) used
+    by a Contributor and that particular Contributor's Contribution.
+
+1.3. "Contribution"
+    means Covered Software of a particular Contributor.
+
+1.4. "Covered Software"
+    means Source Code Form to which the initial Contributor has attached
+    the notice in Exhibit A, the Executable Form of such Source Code
+    Form, and Modifications of such Source Code Form, in each case
+    including portions thereof.
+
+1.5. "Incompatible With Secondary Licenses"
+    means
+
+    (a) that the initial Contributor has attached the notice described
+        in Exhibit B to the Covered Software; or
+
+    (b) that the Covered Software was made available under the terms of
+        version 1.1 or earlier of the License, but not also under the
+        terms of a Secondary License.
+
+1.6. "Executable Form"
+    means any form of the work other than Source Code Form.
+
+1.7. "Larger Work"
+    means a work that combines Covered Software with other material, in
+    a separate file or files, that is not Covered Software.
+
+1.8. "License"
+    means this document.
+
+1.9. "Licensable"
+    means having the right to grant, to the maximum extent possible,
+    whether at the time of the initial grant or subsequently, any and
+    all of the rights conveyed by this License.
+
+1.10. "Modifications"
+    means any of the following:
+
+    (a) any file in Source Code Form that results from an addition to,
+        deletion from, or modification of the contents of Covered
+        Software; or
+
+    (b) any new file in Source Code Form that contains any Covered
+        Software.
+
+1.11. "Patent Claims" of a Contributor
+    means any patent claim(s), including without limitation, method,
+    process, and apparatus claims, in any patent Licensable by such
+    Contributor that would be infringed, but for the grant of the
+    License, by the making, using, selling, offering for sale, having
+    made, import, or transfer of either its Contributions or its
+    Contributor Version.
+
+1.12. "Secondary License"
+    means either the GNU General Public License, Version 2.0, the GNU
+    Lesser General Public License, Version 2.1, the GNU Affero General
+    Public License, Version 3.0, or any later versions of those
+    licenses.
+
+1.13. "Source Code Form"
+    means the form of the work preferred for making modifications.
+
+1.14. "You" (or "Your")
+    means an individual or a legal entity exercising rights under this
+    License. For legal entities, "You" includes any entity that
+    controls, is controlled by, or is under common control with You. For
+    purposes of this definition, "control" means (a) the power, direct
+    or indirect, to cause the direction or management of such entity,
+    whether by contract or otherwise, or (b) ownership of more than
+    fifty percent (50%) of the outstanding shares or beneficial
+    ownership of such entity.
+
+2. License Grants and Conditions
+--------------------------------
+
+2.1. Grants
+
+Each Contributor hereby grants You a world-wide, royalty-free,
+non-exclusive license:
+
+(a) under intellectual property rights (other than patent or trademark)
+    Licensable by such Contributor to use, reproduce, make available,
+    modify, display, perform, distribute, and otherwise exploit its
+    Contributions, either on an unmodified basis, with Modifications, or
+    as part of a Larger Work; and
+
+(b) under Patent Claims of such Contributor to make, use, sell, offer
+    for sale, have made, import, and otherwise transfer either its
+    Contributions or its Contributor Version.
+
+2.2. Effective Date
+
+The licenses granted in Section 2.1 with respect to any Contribution
+become effective for each Contribution on the date the Contributor first
+distributes such Contribution.
+
+2.3. Limitations on Grant Scope
+
+The licenses granted in this Section 2 are the only rights granted under
+this License. No additional rights or licenses will be implied from the
+distribution or licensing of Covered Software under this License.
+Notwithstanding Section 2.1(b) above, no patent license is granted by a
+Contributor:
+
+(a) for any code that a Contributor has removed from Covered Software;
+    or
+
+(b) for infringements caused by: (i) Your and any other third party's
+    modifications of Covered Software, or (ii) the combination of its
+    Contributions with other software (except as part of its Contributor
+    Version); or
+
+(c) under Patent Claims infringed by Covered Software in the absence of
+    its Contributions.
+
+This License does not grant any rights in the trademarks, service marks,
+or logos of any Contributor (except as may be necessary to comply with
+the notice requirements in Section 3.4).
+
+2.4. Subsequent Licenses
+
+No Contributor makes additional grants as a result of Your choice to
+distribute the Covered Software under a subsequent version of this
+License (see Section 10.2) or under the terms of a Secondary License (if
+permitted under the terms of Section 3.3).
+
+2.5. Representation
+
+Each Contributor represents that the Contributor believes its
+Contributions are its original creation(s) or it has sufficient rights
+to grant the rights to its Contributions conveyed by this License.
+
+2.6. Fair Use
+
+This License is not intended to limit any rights You have under
+applicable copyright doctrines of fair use, fair dealing, or other
+equivalents.
+
+2.7. Conditions
+
+Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted
+in Section 2.1.
+
+3. Responsibilities
+-------------------
+
+3.1. Distribution of Source Form
+
+All distribution of Covered Software in Source Code Form, including any
+Modifications that You create or to which You contribute, must be under
+the terms of this License. You must inform recipients that the Source
+Code Form of the Covered Software is governed by the terms of this
+License, and how they can obtain a copy of this License. You may not
+attempt to alter or restrict the recipients' rights in the Source Code
+Form.
+
+3.2. Distribution of Executable Form
+
+If You distribute Covered Software in Executable Form then:
+
+(a) such Covered Software must also be made available in Source Code
+    Form, as described in Section 3.1, and You must inform recipients of
+    the Executable Form how they can obtain a copy of such Source Code
+    Form by reasonable means in a timely manner, at a charge no more
+    than the cost of distribution to the recipient; and
+
+(b) You may distribute such Executable Form under the terms of this
+    License, or sublicense it under different terms, provided that the
+    license for the Executable Form does not attempt to limit or alter
+    the recipients' rights in the Source Code Form under this License.
+
+3.3. Distribution of a Larger Work
+
+You may create and distribute a Larger Work under terms of Your choice,
+provided that You also comply with the requirements of this License for
+the Covered Software. If the Larger Work is a combination of Covered
+Software with a work governed by one or more Secondary Licenses, and the
+Covered Software is not Incompatible With Secondary Licenses, this
+License permits You to additionally distribute such Covered Software
+under the terms of such Secondary License(s), so that the recipient of
+the Larger Work may, at their option, further distribute the Covered
+Software under the terms of either this License or such Secondary
+License(s).
+
+3.4. Notices
+
+You may not remove or alter the substance of any license notices
+(including copyright notices, patent notices, disclaimers of warranty,
+or limitations of liability) contained within the Source Code Form of
+the Covered Software, except that You may alter any license notices to
+the extent required to remedy known factual inaccuracies.
+
+3.5. Application of Additional Terms
+
+You may choose to offer, and to charge a fee for, warranty, support,
+indemnity or liability obligations to one or more recipients of Covered
+Software. However, You may do so only on Your own behalf, and not on
+behalf of any Contributor. You must make it absolutely clear that any
+such warranty, support, indemnity, or liability obligation is offered by
+You alone, and You hereby agree to indemnify every Contributor for any
+liability incurred by such Contributor as a result of warranty, support,
+indemnity or liability terms You offer. You may include additional
+disclaimers of warranty and limitations of liability specific to any
+jurisdiction.
+
+4. Inability to Comply Due to Statute or Regulation
+---------------------------------------------------
+
+If it is impossible for You to comply with any of the terms of this
+License with respect to some or all of the Covered Software due to
+statute, judicial order, or regulation then You must: (a) comply with
+the terms of this License to the maximum extent possible; and (b)
+describe the limitations and the code they affect. Such description must
+be placed in a text file included with all distributions of the Covered
+Software under this License. Except to the extent prohibited by statute
+or regulation, such description must be sufficiently detailed for a
+recipient of ordinary skill to be able to understand it.
+
+5. Termination
+--------------
+
+5.1. The rights granted under this License will terminate automatically
+if You fail to comply with any of its terms. However, if You become
+compliant, then the rights granted under this License from a particular
+Contributor are reinstated (a) provisionally, unless and until such
+Contributor explicitly and finally terminates Your grants, and (b) on an
+ongoing basis, if such Contributor fails to notify You of the
+non-compliance by some reasonable means prior to 60 days after You have
+come back into compliance. Moreover, Your grants from a particular
+Contributor are reinstated on an ongoing basis if such Contributor
+notifies You of the non-compliance by some reasonable means, this is the
+first time You have received notice of non-compliance with this License
+from such Contributor, and You become compliant prior to 30 days after
+Your receipt of the notice.
+
+5.2. If You initiate litigation against any entity by asserting a patent
+infringement claim (excluding declaratory judgment actions,
+counter-claims, and cross-claims) alleging that a Contributor Version
+directly or indirectly infringes any patent, then the rights granted to
+You by any and all Contributors for the Covered Software under Section
+2.1 of this License shall terminate.
+
+5.3. In the event of termination under Sections 5.1 or 5.2 above, all
+end user license agreements (excluding distributors and resellers) which
+have been validly granted by You or Your distributors under this License
+prior to termination shall survive termination.
+
+************************************************************************
+*                                                                      *
+*  6. Disclaimer of Warranty                                           *
+*  -------------------------                                           *
+*                                                                      *
+*  Covered Software is provided under this License on an "as is"       *
+*  basis, without warranty of any kind, either expressed, implied, or  *
+*  statutory, including, without limitation, warranties that the       *
+*  Covered Software is free of defects, merchantable, fit for a        *
+*  particular purpose or non-infringing. The entire risk as to the     *
+*  quality and performance of the Covered Software is with You.        *
+*  Should any Covered Software prove defective in any respect, You     *
+*  (not any Contributor) assume the cost of any necessary servicing,   *
+*  repair, or correction. This disclaimer of warranty constitutes an   *
+*  essential part of this License. No use of any Covered Software is   *
+*  authorized under this License except under this disclaimer.         *
+*                                                                      *
+************************************************************************
+
+************************************************************************
+*                                                                      *
+*  7. Limitation of Liability                                          *
+*  --------------------------                                          *
+*                                                                      *
+*  Under no circumstances and under no legal theory, whether tort      *
+*  (including negligence), contract, or otherwise, shall any           *
+*  Contributor, or anyone who distributes Covered Software as          *
+*  permitted above, be liable to You for any direct, indirect,         *
+*  special, incidental, or consequential damages of any character      *
+*  including, without limitation, damages for lost profits, loss of    *
+*  goodwill, work stoppage, computer failure or malfunction, or any    *
+*  and all other commercial damages or losses, even if such party      *
+*  shall have been informed of the possibility of such damages. This   *
+*  limitation of liability shall not apply to liability for death or   *
+*  personal injury resulting from such party's negligence to the       *
+*  extent applicable law prohibits such limitation. Some               *
+*  jurisdictions do not allow the exclusion or limitation of           *
+*  incidental or consequential damages, so this exclusion and          *
+*  limitation may not apply to You.                                    *
+*                                                                      *
+************************************************************************
+
+8. Litigation
+-------------
+
+Any litigation relating to this License may be brought only in the
+courts of a jurisdiction where the defendant maintains its principal
+place of business and such litigation shall be governed by laws of that
+jurisdiction, without reference to its conflict-of-law provisions.
+Nothing in this Section shall prevent a party's ability to bring
+cross-claims or counter-claims.
+
+9. Miscellaneous
+----------------
+
+This License represents the complete agreement concerning the subject
+matter hereof. If any provision of this License is held to be
+unenforceable, such provision shall be reformed only to the extent
+necessary to make it enforceable. Any law or regulation which provides
+that the language of a contract shall be construed against the drafter
+shall not be used to construe this License against a Contributor.
+
+10. Versions of the License
+---------------------------
+
+10.1. New Versions
+
+Mozilla Foundation is the license steward. Except as provided in Section
+10.3, no one other than the license steward has the right to modify or
+publish new versions of this License. Each version will be given a
+distinguishing version number.
+
+10.2. Effect of New Versions
+
+You may distribute the Covered Software under the terms of the version
+of the License under which You originally received the Covered Software,
+or under the terms of any subsequent version published by the license
+steward.
+
+10.3. Modified Versions
+
+If you create software not governed by this License, and you want to
+create a new license for such software, you may create and use a
+modified version of this License if you rename the license and remove
+any references to the name of the license steward (except to note that
+such modified license differs from this License).
+
+10.4. Distributing Source Code Form that is Incompatible With Secondary
+Licenses
+
+If You choose to distribute Source Code Form that is Incompatible With
+Secondary Licenses under the terms of this version of the License, the
+notice described in Exhibit B of this License must be attached.
+
+Exhibit A - Source Code Form License Notice
+-------------------------------------------
+
+  This Source Code Form is subject to the terms of the Mozilla Public
+  License, v. 2.0. If a copy of the MPL was not distributed with this
+  file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+If it is not possible or desirable to put the notice in a particular
+file, then You may include the notice in a location (such as a LICENSE
+file in a relevant directory) where a recipient would be likely to look
+for such a notice.
+
+You may add additional accurate notices of copyright ownership.
+
+Exhibit B - "Incompatible With Secondary Licenses" Notice
+---------------------------------------------------------
+
+  This Source Code Form is "Incompatible With Secondary Licenses", as
+  defined by the Mozilla Public License, v. 2.0.

inference/interact/fbrs/controller.py

@@ -0,0 +1,103 @@
+import torch
+
+from ..fbrs.inference import clicker
+from ..fbrs.inference.predictors import get_predictor
+
+
+class InteractiveController:
+    def __init__(self, net, device, predictor_params, prob_thresh=0.5):
+        self.net = net.to(device)
+        self.prob_thresh = prob_thresh
+        self.clicker = clicker.Clicker()
+        self.states = []
+        self.probs_history = []
+        self.object_count = 0
+        self._result_mask = None
+
+        self.image = None
+        self.predictor = None
+        self.device = device
+        self.predictor_params = predictor_params
+        self.reset_predictor()
+
+    def set_image(self, image):
+        self.image = image
+        self._result_mask = torch.zeros(image.shape[-2:], dtype=torch.uint8)
+        self.object_count = 0
+        self.reset_last_object()
+
+    def add_click(self, x, y, is_positive):
+        self.states.append({
+            'clicker': self.clicker.get_state(),
+            'predictor': self.predictor.get_states()
+        })
+
+        click = clicker.Click(is_positive=is_positive, coords=(y, x))
+        self.clicker.add_click(click)
+        pred = self.predictor.get_prediction(self.clicker)
+        torch.cuda.empty_cache()
+
+        if self.probs_history:
+            self.probs_history.append((self.probs_history[-1][0], pred))
+        else:
+            self.probs_history.append((torch.zeros_like(pred), pred))
+
+    def undo_click(self):
+        if not self.states:
+            return
+
+        prev_state = self.states.pop()
+        self.clicker.set_state(prev_state['clicker'])
+        self.predictor.set_states(prev_state['predictor'])
+        self.probs_history.pop()
+
+    def partially_finish_object(self):
+        object_prob = self.current_object_prob
+        if object_prob is None:
+            return
+
+        self.probs_history.append((object_prob, torch.zeros_like(object_prob)))
+        self.states.append(self.states[-1])
+
+        self.clicker.reset_clicks()
+        self.reset_predictor()
+
+    def finish_object(self):
+        object_prob = self.current_object_prob
+        if object_prob is None:
+            return
+
+        self.object_count += 1
+        object_mask = object_prob > self.prob_thresh
+        self._result_mask[object_mask] = self.object_count
+        self.reset_last_object()
+
+    def reset_last_object(self):
+        self.states = []
+        self.probs_history = []
+        self.clicker.reset_clicks()
+        self.reset_predictor()
+
+    def reset_predictor(self, predictor_params=None):
+        if predictor_params is not None:
+            self.predictor_params = predictor_params
+        self.predictor = get_predictor(self.net, device=self.device,
+                                       **self.predictor_params)
+        if self.image is not None:
+            self.predictor.set_input_image(self.image)
+
+    @property
+    def current_object_prob(self):
+        if self.probs_history:
+            current_prob_total, current_prob_additive = self.probs_history[-1]
+            return torch.maximum(current_prob_total, current_prob_additive)
+        else:
+            return None
+
+    @property
+    def is_incomplete_mask(self):
+        return len(self.probs_history) > 0
+
+    @property
+    def result_mask(self):
+        return self._result_mask.clone()

inference/interact/fbrs/inference/clicker.py

@@ -0,0 +1,103 @@
+from collections import namedtuple
+
+import numpy as np
+from copy import deepcopy
+from scipy.ndimage import distance_transform_edt
+
+Click = namedtuple('Click', ['is_positive', 'coords'])
+
+
+class Clicker(object):
+    def __init__(self, gt_mask=None, init_clicks=None, ignore_label=-1):
+        if gt_mask is not None:
+            self.gt_mask = gt_mask == 1
+            self.not_ignore_mask = gt_mask != ignore_label
+        else:
+            self.gt_mask = None
+
+        self.reset_clicks()
+
+        if init_clicks is not None:
+            for click in init_clicks:
+                self.add_click(click)
+
+    def make_next_click(self, pred_mask):
+        assert self.gt_mask is not None
+        click = self._get_click(pred_mask)
+        self.add_click(click)
+
+    def get_clicks(self, clicks_limit=None):
+        return self.clicks_list[:clicks_limit]
+
+    def _get_click(self, pred_mask, padding=True):
+        fn_mask = np.logical_and(np.logical_and(self.gt_mask, np.logical_not(pred_mask)), self.not_ignore_mask)
+        fp_mask = np.logical_and(np.logical_and(np.logical_not(self.gt_mask), pred_mask), self.not_ignore_mask)
+
+        if padding:
+            fn_mask = np.pad(fn_mask, ((1, 1), (1, 1)), 'constant')
+            fp_mask = np.pad(fp_mask, ((1, 1), (1, 1)), 'constant')
+
+        fn_mask_dt = distance_transform_edt(fn_mask)
+        fp_mask_dt = distance_transform_edt(fp_mask)
+
+        if padding:
+            fn_mask_dt = fn_mask_dt[1:-1, 1:-1]
+            fp_mask_dt = fp_mask_dt[1:-1, 1:-1]
+
+        fn_mask_dt = fn_mask_dt * self.not_clicked_map
+        fp_mask_dt = fp_mask_dt * self.not_clicked_map
+
+        fn_max_dist = np.max(fn_mask_dt)
+        fp_max_dist = np.max(fp_mask_dt)
+
+        is_positive = fn_max_dist > fp_max_dist
+        if is_positive:
+            coords_y, coords_x = np.where(fn_mask_dt == fn_max_dist)  # coords is [y, x]
+        else:
+            coords_y, coords_x = np.where(fp_mask_dt == fp_max_dist)  # coords is [y, x]
+
+        return Click(is_positive=is_positive, coords=(coords_y[0], coords_x[0]))
+
+    def add_click(self, click):
+        coords = click.coords
+
+        if click.is_positive:
+            self.num_pos_clicks += 1
+        else:
+            self.num_neg_clicks += 1
+
+        self.clicks_list.append(click)
+        if self.gt_mask is not None:
+            self.not_clicked_map[coords[0], coords[1]] = False
+
+    def _remove_last_click(self):
+        click = self.clicks_list.pop()
+        coords = click.coords
+
+        if click.is_positive:
+            self.num_pos_clicks -= 1
+        else:
+            self.num_neg_clicks -= 1
+
+        if self.gt_mask is not None:
+            self.not_clicked_map[coords[0], coords[1]] = True
+
+    def reset_clicks(self):
+        if self.gt_mask is not None:
+            self.not_clicked_map = np.ones_like(self.gt_mask, dtype=np.bool)
+
+        self.num_pos_clicks = 0
+        self.num_neg_clicks = 0
+
+        self.clicks_list = []
+
+    def get_state(self):
+        return deepcopy(self.clicks_list)
+
+    def set_state(self, state):
+        self.reset_clicks()
+        for click in state:
+            self.add_click(click)
+
+    def __len__(self):
+        return len(self.clicks_list)

inference/interact/fbrs/inference/evaluation.py

@@ -0,0 +1,56 @@
+from time import time
+
+import numpy as np
+import torch
+
+from ..inference import utils
+from ..inference.clicker import Clicker
+
+try:
+    get_ipython()
+    from tqdm import tqdm_notebook as tqdm
+except NameError:
+    from tqdm import tqdm
+
+
+def evaluate_dataset(dataset, predictor, oracle_eval=False, **kwargs):
+    all_ious = []
+
+    start_time = time()
+    for index in tqdm(range(len(dataset)), leave=False):
+        sample = dataset.get_sample(index)
+        item = dataset[index]
+
+        if oracle_eval:
+            gt_mask = torch.tensor(sample['instances_mask'], dtype=torch.float32)
+            gt_mask = gt_mask.unsqueeze(0).unsqueeze(0)
+            predictor.opt_functor.mask_loss.set_gt_mask(gt_mask)
+        _, sample_ious, _ = evaluate_sample(item['images'], sample['instances_mask'], predictor, **kwargs)
+        all_ious.append(sample_ious)
+    end_time = time()
+    elapsed_time = end_time - start_time
+
+    return all_ious, elapsed_time
+
+
+def evaluate_sample(image_nd, instances_mask, predictor, max_iou_thr,
+                    pred_thr=0.49, max_clicks=20):
+    clicker = Clicker(gt_mask=instances_mask)
+    pred_mask = np.zeros_like(instances_mask)
+    ious_list = []
+
+    with torch.no_grad():
+        predictor.set_input_image(image_nd)
+
+        for click_number in range(max_clicks):
+            clicker.make_next_click(pred_mask)
+            pred_probs = predictor.get_prediction(clicker)
+            pred_mask = pred_probs > pred_thr
+
+            iou = utils.get_iou(instances_mask, pred_mask)
+            ious_list.append(iou)
+
+            if iou >= max_iou_thr:
+                break
+
+        return clicker.clicks_list, np.array(ious_list, dtype=np.float32), pred_probs

inference/interact/fbrs/inference/predictors/__init__.py

@@ -0,0 +1,95 @@
+from .base import BasePredictor
+from .brs import InputBRSPredictor, FeatureBRSPredictor, HRNetFeatureBRSPredictor
+from .brs_functors import InputOptimizer, ScaleBiasOptimizer
+from ..transforms import ZoomIn
+from ...model.is_hrnet_model import DistMapsHRNetModel
+
+
+def get_predictor(net, brs_mode, device,
+                  prob_thresh=0.49,
+                  with_flip=True,
+                  zoom_in_params=dict(),
+                  predictor_params=None,
+                  brs_opt_func_params=None,
+                  lbfgs_params=None):
+    lbfgs_params_ = {
+        'm': 20,
+        'factr': 0,
+        'pgtol': 1e-8,
+        'maxfun': 20,
+    }
+
+    predictor_params_ = {
+        'optimize_after_n_clicks': 1
+    }
+
+    if zoom_in_params is not None:
+        zoom_in = ZoomIn(**zoom_in_params)
+    else:
+        zoom_in = None
+
+    if lbfgs_params is not None:
+        lbfgs_params_.update(lbfgs_params)
+    lbfgs_params_['maxiter'] = 2 * lbfgs_params_['maxfun']
+
+    if brs_opt_func_params is None:
+        brs_opt_func_params = dict()
+
+    if brs_mode == 'NoBRS':
+        if predictor_params is not None:
+            predictor_params_.update(predictor_params)
+        predictor = BasePredictor(net, device, zoom_in=zoom_in, with_flip=with_flip, **predictor_params_)
+    elif brs_mode.startswith('f-BRS'):
+        predictor_params_.update({
+            'net_clicks_limit': 8,
+        })
+        if predictor_params is not None:
+            predictor_params_.update(predictor_params)
+
+        insertion_mode = {
+            'f-BRS-A': 'after_c4',
+            'f-BRS-B': 'after_aspp',
+            'f-BRS-C': 'after_deeplab'
+        }[brs_mode]
+
+        opt_functor = ScaleBiasOptimizer(prob_thresh=prob_thresh,
+                                         with_flip=with_flip,
+                                         optimizer_params=lbfgs_params_,
+                                         **brs_opt_func_params)
+
+        if isinstance(net, DistMapsHRNetModel):
+            FeaturePredictor = HRNetFeatureBRSPredictor
+            insertion_mode = {'after_c4': 'A', 'after_aspp': 'A', 'after_deeplab': 'C'}[insertion_mode]
+        else:
+            FeaturePredictor = FeatureBRSPredictor
+
+        predictor = FeaturePredictor(net, device,
+                                     opt_functor=opt_functor,
+                                     with_flip=with_flip,
+                                     insertion_mode=insertion_mode,
+                                     zoom_in=zoom_in,
+                                     **predictor_params_)
+    elif brs_mode == 'RGB-BRS' or brs_mode == 'DistMap-BRS':
+        use_dmaps = brs_mode == 'DistMap-BRS'
+
+        predictor_params_.update({
+            'net_clicks_limit': 5,
+        })
+        if predictor_params is not None:
+            predictor_params_.update(predictor_params)
+
+        opt_functor = InputOptimizer(prob_thresh=prob_thresh,
+                                     with_flip=with_flip,
+                                     optimizer_params=lbfgs_params_,
+                                     **brs_opt_func_params)
+
+        predictor = InputBRSPredictor(net, device,
+                                      optimize_target='dmaps' if use_dmaps else 'rgb',
+                                      opt_functor=opt_functor,
+                                      with_flip=with_flip,
+                                      zoom_in=zoom_in,
+                                      **predictor_params_)
+    else:
+        raise NotImplementedError
+
+    return predictor

inference/interact/fbrs/inference/predictors/base.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn.functional as F
+
+from ..transforms import AddHorizontalFlip, SigmoidForPred, LimitLongestSide
+
+
+class BasePredictor(object):
+    def __init__(self, net, device,
+                 net_clicks_limit=None,
+                 with_flip=False,
+                 zoom_in=None,
+                 max_size=None,
+                 **kwargs):
+        self.net = net
+        self.with_flip = with_flip
+        self.net_clicks_limit = net_clicks_limit
+        self.original_image = None
+        self.device = device
+        self.zoom_in = zoom_in
+
+        self.transforms = [zoom_in] if zoom_in is not None else []
+        if max_size is not None:
+            self.transforms.append(LimitLongestSide(max_size=max_size))
+        self.transforms.append(SigmoidForPred())
+        if with_flip:
+            self.transforms.append(AddHorizontalFlip())
+
+    def set_input_image(self, image_nd):
+        for transform in self.transforms:
+            transform.reset()
+        self.original_image = image_nd.to(self.device)
+        if len(self.original_image.shape) == 3:
+            self.original_image = self.original_image.unsqueeze(0)
+
+    def get_prediction(self, clicker):
+        clicks_list = clicker.get_clicks()
+
+        image_nd, clicks_lists, is_image_changed = self.apply_transforms(
+            self.original_image, [clicks_list]
+        )
+
+        pred_logits = self._get_prediction(image_nd, clicks_lists, is_image_changed)
+        prediction = F.interpolate(pred_logits, mode='bilinear', align_corners=True,
+                                   size=image_nd.size()[2:])
+
+        for t in reversed(self.transforms):
+            prediction = t.inv_transform(prediction)
+
+        if self.zoom_in is not None and self.zoom_in.check_possible_recalculation():
+            print('zooming')
+            return self.get_prediction(clicker)
+
+        # return prediction.cpu().numpy()[0, 0]
+        return prediction
+
+    def _get_prediction(self, image_nd, clicks_lists, is_image_changed):
+        points_nd = self.get_points_nd(clicks_lists)
+        return self.net(image_nd, points_nd)['instances']
+
+    def _get_transform_states(self):
+        return [x.get_state() for x in self.transforms]
+
+    def _set_transform_states(self, states):
+        assert len(states) == len(self.transforms)
+        for state, transform in zip(states, self.transforms):
+            transform.set_state(state)
+
+    def apply_transforms(self, image_nd, clicks_lists):
+        is_image_changed = False
+        for t in self.transforms:
+            image_nd, clicks_lists = t.transform(image_nd, clicks_lists)
+            is_image_changed |= t.image_changed
+
+        return image_nd, clicks_lists, is_image_changed
+
+    def get_points_nd(self, clicks_lists):
+        total_clicks = []
+        num_pos_clicks = [sum(x.is_positive for x in clicks_list) for clicks_list in clicks_lists]
+        num_neg_clicks = [len(clicks_list) - num_pos for clicks_list, num_pos in zip(clicks_lists, num_pos_clicks)]
+        num_max_points = max(num_pos_clicks + num_neg_clicks)
+        if self.net_clicks_limit is not None:
+            num_max_points = min(self.net_clicks_limit, num_max_points)
+        num_max_points = max(1, num_max_points)
+
+        for clicks_list in clicks_lists:
+            clicks_list = clicks_list[:self.net_clicks_limit]
+            pos_clicks = [click.coords for click in clicks_list if click.is_positive]
+            pos_clicks = pos_clicks + (num_max_points - len(pos_clicks)) * [(-1, -1)]
+
+            neg_clicks = [click.coords for click in clicks_list if not click.is_positive]
+            neg_clicks = neg_clicks + (num_max_points - len(neg_clicks)) * [(-1, -1)]
+            total_clicks.append(pos_clicks + neg_clicks)
+
+        return torch.tensor(total_clicks, device=self.device)
+
+    def get_states(self):
+        return {'transform_states': self._get_transform_states()}
+
+    def set_states(self, states):
+        self._set_transform_states(states['transform_states'])

inference/interact/fbrs/inference/predictors/brs.py

@@ -0,0 +1,280 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from scipy.optimize import fmin_l_bfgs_b
+
+from .base import BasePredictor
+from ...model.is_hrnet_model import DistMapsHRNetModel
+
+
+class BRSBasePredictor(BasePredictor):
+    def __init__(self, model, device, opt_functor, optimize_after_n_clicks=1, **kwargs):
+        super().__init__(model, device, **kwargs)
+        self.optimize_after_n_clicks = optimize_after_n_clicks
+        self.opt_functor = opt_functor
+
+        self.opt_data = None
+        self.input_data = None
+
+    def set_input_image(self, image_nd):
+        super().set_input_image(image_nd)
+        self.opt_data = None
+        self.input_data = None
+
+    def _get_clicks_maps_nd(self, clicks_lists, image_shape, radius=1):
+        pos_clicks_map = np.zeros((len(clicks_lists), 1) + image_shape, dtype=np.float32)
+        neg_clicks_map = np.zeros((len(clicks_lists), 1) + image_shape, dtype=np.float32)
+
+        for list_indx, clicks_list in enumerate(clicks_lists):
+            for click in clicks_list:
+                y, x = click.coords
+                y, x = int(round(y)), int(round(x))
+                y1, x1 = y - radius, x - radius
+                y2, x2 = y + radius + 1, x + radius + 1
+
+                if click.is_positive:
+                    pos_clicks_map[list_indx, 0, y1:y2, x1:x2] = True
+                else:
+                    neg_clicks_map[list_indx, 0, y1:y2, x1:x2] = True
+
+        with torch.no_grad():
+            pos_clicks_map = torch.from_numpy(pos_clicks_map).to(self.device)
+            neg_clicks_map = torch.from_numpy(neg_clicks_map).to(self.device)
+
+        return pos_clicks_map, neg_clicks_map
+
+    def get_states(self):
+        return {'transform_states': self._get_transform_states(), 'opt_data': self.opt_data}
+
+    def set_states(self, states):
+        self._set_transform_states(states['transform_states'])
+        self.opt_data = states['opt_data']
+
+
+class FeatureBRSPredictor(BRSBasePredictor):
+    def __init__(self, model, device, opt_functor, insertion_mode='after_deeplab', **kwargs):
+        super().__init__(model, device, opt_functor=opt_functor, **kwargs)
+        self.insertion_mode = insertion_mode
+        self._c1_features = None
+
+        if self.insertion_mode == 'after_deeplab':
+            self.num_channels = model.feature_extractor.ch
+        elif self.insertion_mode == 'after_c4':
+            self.num_channels = model.feature_extractor.aspp_in_channels
+        elif self.insertion_mode == 'after_aspp':
+            self.num_channels = model.feature_extractor.ch + 32
+        else:
+            raise NotImplementedError
+
+    def _get_prediction(self, image_nd, clicks_lists, is_image_changed):
+        points_nd = self.get_points_nd(clicks_lists)
+        pos_mask, neg_mask = self._get_clicks_maps_nd(clicks_lists, image_nd.shape[2:])
+
+        num_clicks = len(clicks_lists[0])
+        bs = image_nd.shape[0] // 2 if self.with_flip else image_nd.shape[0]
+
+        if self.opt_data is None or self.opt_data.shape[0] // (2 * self.num_channels) != bs:
+            self.opt_data = np.zeros((bs * 2 * self.num_channels), dtype=np.float32)
+
+        if num_clicks <= self.net_clicks_limit or is_image_changed or self.input_data is None:
+            self.input_data = self._get_head_input(image_nd, points_nd)
+
+        def get_prediction_logits(scale, bias):
+            scale = scale.view(bs, -1, 1, 1)
+            bias = bias.view(bs, -1, 1, 1)
+            if self.with_flip:
+                scale = scale.repeat(2, 1, 1, 1)
+                bias = bias.repeat(2, 1, 1, 1)
+
+            scaled_backbone_features = self.input_data * scale
+            scaled_backbone_features = scaled_backbone_features + bias
+            if self.insertion_mode == 'after_c4':
+                x = self.net.feature_extractor.aspp(scaled_backbone_features)
+                x = F.interpolate(x, mode='bilinear', size=self._c1_features.size()[2:],
+                                  align_corners=True)
+                x = torch.cat((x, self._c1_features), dim=1)
+                scaled_backbone_features = self.net.feature_extractor.head(x)
+            elif self.insertion_mode == 'after_aspp':
+                scaled_backbone_features = self.net.feature_extractor.head(scaled_backbone_features)
+
+            pred_logits = self.net.head(scaled_backbone_features)
+            pred_logits = F.interpolate(pred_logits, size=image_nd.size()[2:], mode='bilinear',
+                                        align_corners=True)
+            return pred_logits
+
+        self.opt_functor.init_click(get_prediction_logits, pos_mask, neg_mask, self.device)
+        if num_clicks > self.optimize_after_n_clicks:
+            opt_result = fmin_l_bfgs_b(func=self.opt_functor, x0=self.opt_data,
+                                       **self.opt_functor.optimizer_params)
+            self.opt_data = opt_result[0]
+
+        with torch.no_grad():
+            if self.opt_functor.best_prediction is not None:
+                opt_pred_logits = self.opt_functor.best_prediction
+            else:
+                opt_data_nd = torch.from_numpy(self.opt_data).to(self.device)
+                opt_vars, _ = self.opt_functor.unpack_opt_params(opt_data_nd)
+                opt_pred_logits = get_prediction_logits(*opt_vars)
+
+        return opt_pred_logits
+
+    def _get_head_input(self, image_nd, points):
+        with torch.no_grad():
+            coord_features = self.net.dist_maps(image_nd, points)
+            x = self.net.rgb_conv(torch.cat((image_nd, coord_features), dim=1))
+            if self.insertion_mode == 'after_c4' or self.insertion_mode == 'after_aspp':
+                c1, _, c3, c4 = self.net.feature_extractor.backbone(x)
+                c1 = self.net.feature_extractor.skip_project(c1)
+
+                if self.insertion_mode == 'after_aspp':
+                    x = self.net.feature_extractor.aspp(c4)
+                    x = F.interpolate(x, size=c1.size()[2:], mode='bilinear', align_corners=True)
+                    x = torch.cat((x, c1), dim=1)
+                    backbone_features = x
+                else:
+                    backbone_features = c4
+                    self._c1_features = c1
+            else:
+                backbone_features = self.net.feature_extractor(x)[0]
+
+        return backbone_features
+
+
+class HRNetFeatureBRSPredictor(BRSBasePredictor):
+    def __init__(self, model, device, opt_functor, insertion_mode='A', **kwargs):
+        super().__init__(model, device, opt_functor=opt_functor, **kwargs)
+        self.insertion_mode = insertion_mode
+        self._c1_features = None
+
+        if self.insertion_mode == 'A':
+            self.num_channels = sum(k * model.feature_extractor.width for k in [1, 2, 4, 8])
+        elif self.insertion_mode == 'C':
+            self.num_channels = 2 * model.feature_extractor.ocr_width
+        else:
+            raise NotImplementedError
+
+    def _get_prediction(self, image_nd, clicks_lists, is_image_changed):
+        points_nd = self.get_points_nd(clicks_lists)
+        pos_mask, neg_mask = self._get_clicks_maps_nd(clicks_lists, image_nd.shape[2:])
+        num_clicks = len(clicks_lists[0])
+        bs = image_nd.shape[0] // 2 if self.with_flip else image_nd.shape[0]
+
+        if self.opt_data is None or self.opt_data.shape[0] // (2 * self.num_channels) != bs:
+            self.opt_data = np.zeros((bs * 2 * self.num_channels), dtype=np.float32)
+
+        if num_clicks <= self.net_clicks_limit or is_image_changed or self.input_data is None:
+            self.input_data = self._get_head_input(image_nd, points_nd)
+
+        def get_prediction_logits(scale, bias):
+            scale = scale.view(bs, -1, 1, 1)
+            bias = bias.view(bs, -1, 1, 1)
+            if self.with_flip:
+                scale = scale.repeat(2, 1, 1, 1)
+                bias = bias.repeat(2, 1, 1, 1)
+
+            scaled_backbone_features = self.input_data * scale
+            scaled_backbone_features = scaled_backbone_features + bias
+            if self.insertion_mode == 'A':
+                out_aux = self.net.feature_extractor.aux_head(scaled_backbone_features)
+                feats = self.net.feature_extractor.conv3x3_ocr(scaled_backbone_features)
+
+                context = self.net.feature_extractor.ocr_gather_head(feats, out_aux)
+                feats = self.net.feature_extractor.ocr_distri_head(feats, context)
+                pred_logits = self.net.feature_extractor.cls_head(feats)
+            elif self.insertion_mode == 'C':
+                pred_logits = self.net.feature_extractor.cls_head(scaled_backbone_features)
+            else:
+                raise NotImplementedError
+
+            pred_logits = F.interpolate(pred_logits, size=image_nd.size()[2:], mode='bilinear',
+                                        align_corners=True)
+            return pred_logits
+
+        self.opt_functor.init_click(get_prediction_logits, pos_mask, neg_mask, self.device)
+        if num_clicks > self.optimize_after_n_clicks:
+            opt_result = fmin_l_bfgs_b(func=self.opt_functor, x0=self.opt_data,
+                                       **self.opt_functor.optimizer_params)
+            self.opt_data = opt_result[0]
+
+        with torch.no_grad():
+            if self.opt_functor.best_prediction is not None:
+                opt_pred_logits = self.opt_functor.best_prediction
+            else:
+                opt_data_nd = torch.from_numpy(self.opt_data).to(self.device)
+                opt_vars, _ = self.opt_functor.unpack_opt_params(opt_data_nd)
+                opt_pred_logits = get_prediction_logits(*opt_vars)
+
+        return opt_pred_logits
+
+    def _get_head_input(self, image_nd, points):
+        with torch.no_grad():
+            coord_features = self.net.dist_maps(image_nd, points)
+            x = self.net.rgb_conv(torch.cat((image_nd, coord_features), dim=1))
+            feats = self.net.feature_extractor.compute_hrnet_feats(x)
+            if self.insertion_mode == 'A':
+                backbone_features = feats
+            elif self.insertion_mode == 'C':
+                out_aux = self.net.feature_extractor.aux_head(feats)
+                feats = self.net.feature_extractor.conv3x3_ocr(feats)
+
+                context = self.net.feature_extractor.ocr_gather_head(feats, out_aux)
+                backbone_features = self.net.feature_extractor.ocr_distri_head(feats, context)
+            else:
+                raise NotImplementedError
+
+        return backbone_features
+
+
+class InputBRSPredictor(BRSBasePredictor):
+    def __init__(self, model, device, opt_functor, optimize_target='rgb', **kwargs):
+        super().__init__(model, device, opt_functor=opt_functor, **kwargs)
+        self.optimize_target = optimize_target
+
+    def _get_prediction(self, image_nd, clicks_lists, is_image_changed):
+        points_nd = self.get_points_nd(clicks_lists)
+        pos_mask, neg_mask = self._get_clicks_maps_nd(clicks_lists, image_nd.shape[2:])
+        num_clicks = len(clicks_lists[0])
+
+        if self.opt_data is None or is_image_changed:
+            opt_channels = 2 if self.optimize_target == 'dmaps' else 3
+            bs = image_nd.shape[0] // 2 if self.with_flip else image_nd.shape[0]
+            self.opt_data = torch.zeros((bs, opt_channels, image_nd.shape[2], image_nd.shape[3]),
+                                        device=self.device, dtype=torch.float32)
+
+        def get_prediction_logits(opt_bias):
+            input_image = image_nd
+            if self.optimize_target == 'rgb':
+                input_image = input_image + opt_bias
+            dmaps = self.net.dist_maps(input_image, points_nd)
+            if self.optimize_target == 'dmaps':
+                dmaps = dmaps + opt_bias
+
+            x = self.net.rgb_conv(torch.cat((input_image, dmaps), dim=1))
+            if self.optimize_target == 'all':
+                x = x + opt_bias
+
+            if isinstance(self.net, DistMapsHRNetModel):
+                pred_logits = self.net.feature_extractor(x)[0]
+            else:
+                backbone_features = self.net.feature_extractor(x)
+                pred_logits = self.net.head(backbone_features[0])
+            pred_logits = F.interpolate(pred_logits, size=image_nd.size()[2:], mode='bilinear', align_corners=True)
+
+            return pred_logits
+
+        self.opt_functor.init_click(get_prediction_logits, pos_mask, neg_mask, self.device,
+                                    shape=self.opt_data.shape)
+        if num_clicks > self.optimize_after_n_clicks:
+            opt_result = fmin_l_bfgs_b(func=self.opt_functor, x0=self.opt_data.cpu().numpy().ravel(),
+                                       **self.opt_functor.optimizer_params)
+
+            self.opt_data = torch.from_numpy(opt_result[0]).view(self.opt_data.shape).to(self.device)
+
+        with torch.no_grad():
+            if self.opt_functor.best_prediction is not None:
+                opt_pred_logits = self.opt_functor.best_prediction
+            else:
+                opt_vars, _ = self.opt_functor.unpack_opt_params(self.opt_data)
+                opt_pred_logits = get_prediction_logits(*opt_vars)
+
+        return opt_pred_logits

inference/interact/fbrs/inference/predictors/brs_functors.py

@@ -0,0 +1,109 @@
+import torch
+import numpy as np
+
+from ...model.metrics import _compute_iou
+from .brs_losses import BRSMaskLoss
+
+
+class BaseOptimizer:
+    def __init__(self, optimizer_params,
+                 prob_thresh=0.49,
+                 reg_weight=1e-3,
+                 min_iou_diff=0.01,
+                 brs_loss=BRSMaskLoss(),
+                 with_flip=False,
+                 flip_average=False,
+                 **kwargs):
+        self.brs_loss = brs_loss
+        self.optimizer_params = optimizer_params
+        self.prob_thresh = prob_thresh
+        self.reg_weight = reg_weight
+        self.min_iou_diff = min_iou_diff
+        self.with_flip = with_flip
+        self.flip_average = flip_average
+
+        self.best_prediction = None
+        self._get_prediction_logits = None
+        self._opt_shape = None
+        self._best_loss = None
+        self._click_masks = None
+        self._last_mask = None
+        self.device = None
+
+    def init_click(self, get_prediction_logits, pos_mask, neg_mask, device, shape=None):
+        self.best_prediction = None
+        self._get_prediction_logits = get_prediction_logits
+        self._click_masks = (pos_mask, neg_mask)
+        self._opt_shape = shape
+        self._last_mask = None
+        self.device = device
+
+    def __call__(self, x):
+        opt_params = torch.from_numpy(x).float().to(self.device)
+        opt_params.requires_grad_(True)
+
+        with torch.enable_grad():
+            opt_vars, reg_loss = self.unpack_opt_params(opt_params)
+            result_before_sigmoid = self._get_prediction_logits(*opt_vars)
+            result = torch.sigmoid(result_before_sigmoid)
+
+            pos_mask, neg_mask = self._click_masks
+            if self.with_flip and self.flip_average:
+                result, result_flipped = torch.chunk(result, 2, dim=0)
+                result = 0.5 * (result + torch.flip(result_flipped, dims=[3]))
+                pos_mask, neg_mask = pos_mask[:result.shape[0]], neg_mask[:result.shape[0]]
+
+            loss, f_max_pos, f_max_neg = self.brs_loss(result, pos_mask, neg_mask)
+            loss = loss + reg_loss
+
+        f_val = loss.detach().cpu().numpy()
+        if self.best_prediction is None or f_val < self._best_loss:
+            self.best_prediction = result_before_sigmoid.detach()
+            self._best_loss = f_val
+
+        if f_max_pos < (1 - self.prob_thresh) and f_max_neg < self.prob_thresh:
+            return [f_val, np.zeros_like(x)]
+
+        current_mask = result > self.prob_thresh
+        if self._last_mask is not None and self.min_iou_diff > 0:
+            diff_iou = _compute_iou(current_mask, self._last_mask)
+            if len(diff_iou) > 0 and diff_iou.mean() > 1 - self.min_iou_diff:
+                return [f_val, np.zeros_like(x)]
+        self._last_mask = current_mask
+
+        loss.backward()
+        f_grad = opt_params.grad.cpu().numpy().ravel().astype(np.float)
+
+        return [f_val, f_grad]
+
+    def unpack_opt_params(self, opt_params):
+        raise NotImplementedError
+
+
+class InputOptimizer(BaseOptimizer):
+    def unpack_opt_params(self, opt_params):
+        opt_params = opt_params.view(self._opt_shape)
+        if self.with_flip:
+            opt_params_flipped = torch.flip(opt_params, dims=[3])
+            opt_params = torch.cat([opt_params, opt_params_flipped], dim=0)
+        reg_loss = self.reg_weight * torch.sum(opt_params**2)
+
+        return (opt_params,), reg_loss
+
+
+class ScaleBiasOptimizer(BaseOptimizer):
+    def __init__(self, *args, scale_act=None, reg_bias_weight=10.0, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.scale_act = scale_act
+        self.reg_bias_weight = reg_bias_weight
+
+    def unpack_opt_params(self, opt_params):
+        scale, bias = torch.chunk(opt_params, 2, dim=0)
+        reg_loss = self.reg_weight * (torch.sum(scale**2) + self.reg_bias_weight * torch.sum(bias**2))
+
+        if self.scale_act == 'tanh':
+            scale = torch.tanh(scale)
+        elif self.scale_act == 'sin':
+            scale = torch.sin(scale)
+
+        return (1 + scale, bias), reg_loss

inference/interact/fbrs/inference/predictors/brs_losses.py

@@ -0,0 +1,58 @@
+import torch
+
+from ...model.losses import SigmoidBinaryCrossEntropyLoss
+
+
+class BRSMaskLoss(torch.nn.Module):
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self._eps = eps
+
+    def forward(self, result, pos_mask, neg_mask):
+        pos_diff = (1 - result) * pos_mask
+        pos_target = torch.sum(pos_diff ** 2)
+        pos_target = pos_target / (torch.sum(pos_mask) + self._eps)
+
+        neg_diff = result * neg_mask
+        neg_target = torch.sum(neg_diff ** 2)
+        neg_target = neg_target / (torch.sum(neg_mask) + self._eps)
+        
+        loss = pos_target + neg_target
+
+        with torch.no_grad():
+            f_max_pos = torch.max(torch.abs(pos_diff)).item()
+            f_max_neg = torch.max(torch.abs(neg_diff)).item()
+
+        return loss, f_max_pos, f_max_neg
+
+
+class OracleMaskLoss(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gt_mask = None
+        self.loss = SigmoidBinaryCrossEntropyLoss(from_sigmoid=True)
+        self.predictor = None
+        self.history = []
+
+    def set_gt_mask(self, gt_mask):
+        self.gt_mask = gt_mask
+        self.history = []
+
+    def forward(self, result, pos_mask, neg_mask):
+        gt_mask = self.gt_mask.to(result.device)
+        if self.predictor.object_roi is not None:
+            r1, r2, c1, c2 = self.predictor.object_roi[:4]
+            gt_mask = gt_mask[:, :, r1:r2 + 1, c1:c2 + 1]
+            gt_mask = torch.nn.functional.interpolate(gt_mask, result.size()[2:],  mode='bilinear', align_corners=True)
+
+        if result.shape[0] == 2:
+            gt_mask_flipped = torch.flip(gt_mask, dims=[3])
+            gt_mask = torch.cat([gt_mask, gt_mask_flipped], dim=0)
+
+        loss = self.loss(result, gt_mask)
+        self.history.append(loss.detach().cpu().numpy()[0])
+
+        if len(self.history) > 5 and abs(self.history[-5] - self.history[-1]) < 1e-5:
+            return 0, 0, 0
+
+        return loss, 1.0, 1.0

inference/interact/fbrs/inference/transforms/__init__.py

@@ -0,0 +1,5 @@
+from .base import SigmoidForPred
+from .flip import AddHorizontalFlip
+from .zoom_in import ZoomIn
+from .limit_longest_side import LimitLongestSide
+from .crops import Crops

inference/interact/fbrs/inference/transforms/base.py

@@ -0,0 +1,38 @@
+import torch
+
+
+class BaseTransform(object):
+    def __init__(self):
+        self.image_changed = False
+
+    def transform(self, image_nd, clicks_lists):
+        raise NotImplementedError
+
+    def inv_transform(self, prob_map):
+        raise NotImplementedError
+
+    def reset(self):
+        raise NotImplementedError
+
+    def get_state(self):
+        raise NotImplementedError
+
+    def set_state(self, state):
+        raise NotImplementedError
+
+
+class SigmoidForPred(BaseTransform):
+    def transform(self, image_nd, clicks_lists):
+        return image_nd, clicks_lists
+
+    def inv_transform(self, prob_map):
+        return torch.sigmoid(prob_map)
+
+    def reset(self):
+        pass
+
+    def get_state(self):
+        return None
+
+    def set_state(self, state):
+        pass

inference/interact/fbrs/inference/transforms/crops.py

@@ -0,0 +1,97 @@
+import math
+
+import torch
+import numpy as np
+
+from ...inference.clicker import Click
+from .base import BaseTransform
+
+
+class Crops(BaseTransform):
+    def __init__(self, crop_size=(320, 480), min_overlap=0.2):
+        super().__init__()
+        self.crop_height, self.crop_width = crop_size
+        self.min_overlap = min_overlap
+
+        self.x_offsets = None
+        self.y_offsets = None
+        self._counts = None
+
+    def transform(self, image_nd, clicks_lists):
+        assert image_nd.shape[0] == 1 and len(clicks_lists) == 1
+        image_height, image_width = image_nd.shape[2:4]
+        self._counts = None
+
+        if image_height < self.crop_height or image_width < self.crop_width:
+            return image_nd, clicks_lists
+
+        self.x_offsets = get_offsets(image_width, self.crop_width, self.min_overlap)
+        self.y_offsets = get_offsets(image_height, self.crop_height, self.min_overlap)
+        self._counts = np.zeros((image_height, image_width))
+
+        image_crops = []
+        for dy in self.y_offsets:
+            for dx in self.x_offsets:
+                self._counts[dy:dy + self.crop_height, dx:dx + self.crop_width] += 1
+                image_crop = image_nd[:, :, dy:dy + self.crop_height, dx:dx + self.crop_width]
+                image_crops.append(image_crop)
+        image_crops = torch.cat(image_crops, dim=0)
+        self._counts = torch.tensor(self._counts, device=image_nd.device, dtype=torch.float32)
+
+        clicks_list = clicks_lists[0]
+        clicks_lists = []
+        for dy in self.y_offsets:
+            for dx in self.x_offsets:
+                crop_clicks = [Click(is_positive=x.is_positive, coords=(x.coords[0] - dy, x.coords[1] - dx))
+                               for x in clicks_list]
+                clicks_lists.append(crop_clicks)
+
+        return image_crops, clicks_lists
+
+    def inv_transform(self, prob_map):
+        if self._counts is None:
+            return prob_map
+
+        new_prob_map = torch.zeros((1, 1, *self._counts.shape),
+                                   dtype=prob_map.dtype, device=prob_map.device)
+
+        crop_indx = 0
+        for dy in self.y_offsets:
+            for dx in self.x_offsets:
+                new_prob_map[0, 0, dy:dy + self.crop_height, dx:dx + self.crop_width] += prob_map[crop_indx, 0]
+                crop_indx += 1
+        new_prob_map = torch.div(new_prob_map, self._counts)
+
+        return new_prob_map
+
+    def get_state(self):
+        return self.x_offsets, self.y_offsets, self._counts
+
+    def set_state(self, state):
+        self.x_offsets, self.y_offsets, self._counts = state
+
+    def reset(self):
+        self.x_offsets = None
+        self.y_offsets = None
+        self._counts = None
+
+
+def get_offsets(length, crop_size, min_overlap_ratio=0.2):
+    if length == crop_size:
+        return [0]
+
+    N = (length / crop_size - min_overlap_ratio) / (1 - min_overlap_ratio)
+    N = math.ceil(N)
+
+    overlap_ratio = (N - length / crop_size) / (N - 1)
+    overlap_width = int(crop_size * overlap_ratio)
+
+    offsets = [0]
+    for i in range(1, N):
+        new_offset = offsets[-1] + crop_size - overlap_width
+        if new_offset + crop_size > length:
+            new_offset = length - crop_size
+
+        offsets.append(new_offset)
+
+    return offsets

inference/interact/fbrs/inference/transforms/flip.py

@@ -0,0 +1,37 @@
+import torch
+
+from ..clicker import Click
+from .base import BaseTransform
+
+
+class AddHorizontalFlip(BaseTransform):
+    def transform(self, image_nd, clicks_lists):
+        assert len(image_nd.shape) == 4
+        image_nd = torch.cat([image_nd, torch.flip(image_nd, dims=[3])], dim=0)
+
+        image_width = image_nd.shape[3]
+        clicks_lists_flipped = []
+        for clicks_list in clicks_lists:
+            clicks_list_flipped = [Click(is_positive=click.is_positive,
+                                         coords=(click.coords[0], image_width - click.coords[1] - 1))
+                                   for click in clicks_list]
+            clicks_lists_flipped.append(clicks_list_flipped)
+        clicks_lists = clicks_lists + clicks_lists_flipped
+
+        return image_nd, clicks_lists
+
+    def inv_transform(self, prob_map):
+        assert len(prob_map.shape) == 4 and prob_map.shape[0] % 2 == 0
+        num_maps = prob_map.shape[0] // 2
+        prob_map, prob_map_flipped = prob_map[:num_maps], prob_map[num_maps:]
+
+        return 0.5 * (prob_map + torch.flip(prob_map_flipped, dims=[3]))
+
+    def get_state(self):
+        return None
+
+    def set_state(self, state):
+        pass
+
+    def reset(self):
+        pass

inference/interact/fbrs/inference/transforms/limit_longest_side.py

@@ -0,0 +1,22 @@
+from .zoom_in import ZoomIn, get_roi_image_nd
+
+
+class LimitLongestSide(ZoomIn):
+    def __init__(self, max_size=800):
+        super().__init__(target_size=max_size, skip_clicks=0)
+
+    def transform(self, image_nd, clicks_lists):
+        assert image_nd.shape[0] == 1 and len(clicks_lists) == 1
+        image_max_size = max(image_nd.shape[2:4])
+        self.image_changed = False
+
+        if image_max_size <= self.target_size:
+            return image_nd, clicks_lists
+        self._input_image = image_nd
+
+        self._object_roi = (0, image_nd.shape[2] - 1, 0, image_nd.shape[3] - 1)
+        self._roi_image = get_roi_image_nd(image_nd, self._object_roi, self.target_size)
+        self.image_changed = True
+
+        tclicks_lists = [self._transform_clicks(clicks_lists[0])]
+        return self._roi_image, tclicks_lists

inference/interact/fbrs/inference/transforms/zoom_in.py

@@ -0,0 +1,171 @@
+import torch
+
+from ..clicker import Click
+from ...utils.misc import get_bbox_iou, get_bbox_from_mask, expand_bbox, clamp_bbox
+from .base import BaseTransform
+
+
+class ZoomIn(BaseTransform):
+    def __init__(self,
+                 target_size=400,
+                 skip_clicks=1,
+                 expansion_ratio=1.4,
+                 min_crop_size=200,
+                 recompute_thresh_iou=0.5,
+                 prob_thresh=0.50):
+        super().__init__()
+        self.target_size = target_size
+        self.min_crop_size = min_crop_size
+        self.skip_clicks = skip_clicks
+        self.expansion_ratio = expansion_ratio
+        self.recompute_thresh_iou = recompute_thresh_iou
+        self.prob_thresh = prob_thresh
+
+        self._input_image_shape = None
+        self._prev_probs = None
+        self._object_roi = None
+        self._roi_image = None
+
+    def transform(self, image_nd, clicks_lists):
+        assert image_nd.shape[0] == 1 and len(clicks_lists) == 1
+        self.image_changed = False
+
+        clicks_list = clicks_lists[0]
+        if len(clicks_list) <= self.skip_clicks:
+            return image_nd, clicks_lists
+
+        self._input_image_shape = image_nd.shape
+
+        current_object_roi = None
+        if self._prev_probs is not None:
+            current_pred_mask = (self._prev_probs > self.prob_thresh)[0, 0]
+            if current_pred_mask.sum() > 0:
+                current_object_roi = get_object_roi(current_pred_mask, clicks_list,
+                                                    self.expansion_ratio, self.min_crop_size)
+
+        if current_object_roi is None:
+            return image_nd, clicks_lists
+
+        update_object_roi = False
+        if self._object_roi is None:
+            update_object_roi = True
+        elif not check_object_roi(self._object_roi, clicks_list):
+            update_object_roi = True
+        elif get_bbox_iou(current_object_roi, self._object_roi) < self.recompute_thresh_iou:
+            update_object_roi = True
+
+        if update_object_roi:
+            self._object_roi = current_object_roi
+            self._roi_image = get_roi_image_nd(image_nd, self._object_roi, self.target_size)
+            self.image_changed = True
+
+        tclicks_lists = [self._transform_clicks(clicks_list)]
+        return self._roi_image.to(image_nd.device), tclicks_lists
+
+    def inv_transform(self, prob_map):
+        if self._object_roi is None:
+            self._prev_probs = prob_map.cpu().numpy()
+            return prob_map
+
+        assert prob_map.shape[0] == 1
+        rmin, rmax, cmin, cmax = self._object_roi
+        prob_map = torch.nn.functional.interpolate(prob_map, size=(rmax - rmin + 1, cmax - cmin + 1),
+                                                   mode='bilinear', align_corners=True)
+
+        if self._prev_probs is not None:
+            new_prob_map = torch.zeros(*self._prev_probs.shape, device=prob_map.device, dtype=prob_map.dtype)
+            new_prob_map[:, :, rmin:rmax + 1, cmin:cmax + 1] = prob_map
+        else:
+            new_prob_map = prob_map
+
+        self._prev_probs = new_prob_map.cpu().numpy()
+
+        return new_prob_map
+
+    def check_possible_recalculation(self):
+        if self._prev_probs is None or self._object_roi is not None or self.skip_clicks > 0:
+            return False
+
+        pred_mask = (self._prev_probs > self.prob_thresh)[0, 0]
+        if pred_mask.sum() > 0:
+            possible_object_roi = get_object_roi(pred_mask, [],
+                                                 self.expansion_ratio, self.min_crop_size)
+            image_roi = (0, self._input_image_shape[2] - 1, 0, self._input_image_shape[3] - 1)
+            if get_bbox_iou(possible_object_roi, image_roi) < 0.50:
+                return True
+        return False
+
+    def get_state(self):
+        roi_image = self._roi_image.cpu() if self._roi_image is not None else None
+        return self._input_image_shape, self._object_roi, self._prev_probs, roi_image, self.image_changed
+
+    def set_state(self, state):
+        self._input_image_shape, self._object_roi, self._prev_probs, self._roi_image, self.image_changed = state
+
+    def reset(self):
+        self._input_image_shape = None
+        self._object_roi = None
+        self._prev_probs = None
+        self._roi_image = None
+        self.image_changed = False
+
+    def _transform_clicks(self, clicks_list):
+        if self._object_roi is None:
+            return clicks_list
+
+        rmin, rmax, cmin, cmax = self._object_roi
+        crop_height, crop_width = self._roi_image.shape[2:]
+
+        transformed_clicks = []
+        for click in clicks_list:
+            new_r = crop_height * (click.coords[0] - rmin) / (rmax - rmin + 1)
+            new_c = crop_width * (click.coords[1] - cmin) / (cmax - cmin + 1)
+            transformed_clicks.append(Click(is_positive=click.is_positive, coords=(new_r, new_c)))
+        return transformed_clicks
+
+
+def get_object_roi(pred_mask, clicks_list, expansion_ratio, min_crop_size):
+    pred_mask = pred_mask.copy()
+
+    for click in clicks_list:
+        if click.is_positive:
+            pred_mask[int(click.coords[0]), int(click.coords[1])] = 1
+
+    bbox = get_bbox_from_mask(pred_mask)
+    bbox = expand_bbox(bbox, expansion_ratio, min_crop_size)
+    h, w = pred_mask.shape[0], pred_mask.shape[1]
+    bbox = clamp_bbox(bbox, 0, h - 1, 0, w - 1)
+
+    return bbox
+
+
+def get_roi_image_nd(image_nd, object_roi, target_size):
+    rmin, rmax, cmin, cmax = object_roi
+
+    height = rmax - rmin + 1
+    width = cmax - cmin + 1
+
+    if isinstance(target_size, tuple):
+        new_height, new_width = target_size
+    else:
+        scale = target_size / max(height, width)
+        new_height = int(round(height * scale))
+        new_width = int(round(width * scale))
+
+    with torch.no_grad():
+        roi_image_nd = image_nd[:, :, rmin:rmax + 1, cmin:cmax + 1]
+        roi_image_nd = torch.nn.functional.interpolate(roi_image_nd, size=(new_height, new_width),
+                                                       mode='bilinear', align_corners=True)
+
+    return roi_image_nd
+
+
+def check_object_roi(object_roi, clicks_list):
+    for click in clicks_list:
+        if click.is_positive:
+            if click.coords[0] < object_roi[0] or click.coords[0] >= object_roi[1]:
+                return False
+            if click.coords[1] < object_roi[2] or click.coords[1] >= object_roi[3]:
+                return False
+
+    return True

inference/interact/fbrs/inference/utils.py

@@ -0,0 +1,177 @@
+from datetime import timedelta
+from pathlib import Path
+
+import torch
+import numpy as np
+
+from ..model.is_deeplab_model import get_deeplab_model
+from ..model.is_hrnet_model import get_hrnet_model
+
+
+def get_time_metrics(all_ious, elapsed_time):
+    n_images = len(all_ious)
+    n_clicks = sum(map(len, all_ious))
+
+    mean_spc = elapsed_time / n_clicks
+    mean_spi = elapsed_time / n_images
+
+    return mean_spc, mean_spi
+
+
+def load_is_model(checkpoint, device, backbone='auto', **kwargs):
+    if isinstance(checkpoint, (str, Path)):
+        state_dict = torch.load(checkpoint, map_location='cpu')
+    else:
+        state_dict = checkpoint
+
+    if backbone == 'auto':
+        for k in state_dict.keys():
+            if 'feature_extractor.stage2.0.branches' in k:
+                return load_hrnet_is_model(state_dict, device, backbone, **kwargs)
+        return load_deeplab_is_model(state_dict, device, backbone, **kwargs)
+    elif 'resnet' in backbone:
+        return load_deeplab_is_model(state_dict, device, backbone, **kwargs)
+    elif 'hrnet' in backbone:
+        return load_hrnet_is_model(state_dict, device, backbone, **kwargs)
+    else:
+        raise NotImplementedError('Unknown backbone')
+
+
+def load_hrnet_is_model(state_dict, device, backbone='auto', width=48, ocr_width=256,
+                        small=False, cpu_dist_maps=False, norm_radius=260):
+    if backbone == 'auto':
+        num_fe_weights = len([x for x in state_dict.keys() if 'feature_extractor.' in x])
+        small = num_fe_weights < 1800
+
+        ocr_f_down = [v for k, v in state_dict.items() if 'object_context_block.f_down.1.0.bias' in k]
+        assert len(ocr_f_down) == 1
+        ocr_width = ocr_f_down[0].shape[0]
+
+        s2_conv1_w = [v for k, v in state_dict.items() if 'stage2.0.branches.0.0.conv1.weight' in k]
+        assert  len(s2_conv1_w) == 1
+        width = s2_conv1_w[0].shape[0]
+
+    model = get_hrnet_model(width=width, ocr_width=ocr_width, small=small,
+                            with_aux_output=False, cpu_dist_maps=cpu_dist_maps,
+                            norm_radius=norm_radius)
+
+    model.load_state_dict(state_dict, strict=False)
+    for param in model.parameters():
+        param.requires_grad = False
+    model.to(device)
+    model.eval()
+
+    return model
+
+
+def load_deeplab_is_model(state_dict, device, backbone='auto', deeplab_ch=128, aspp_dropout=0.2,
+                          cpu_dist_maps=False, norm_radius=260):
+    if backbone == 'auto':
+        num_backbone_params = len([x for x in state_dict.keys()
+                                   if 'feature_extractor.backbone' in x and not('num_batches_tracked' in x)])
+
+        if num_backbone_params <= 181:
+            backbone = 'resnet34'
+        elif num_backbone_params <= 276:
+            backbone = 'resnet50'
+        elif num_backbone_params <= 531:
+            backbone = 'resnet101'
+        else:
+            raise NotImplementedError('Unknown backbone')
+
+        if 'aspp_dropout' in state_dict:
+            aspp_dropout = float(state_dict['aspp_dropout'].cpu().numpy())
+        else:
+            aspp_project_weight = [v for k, v in state_dict.items() if 'aspp.project.0.weight' in k][0]
+            deeplab_ch = aspp_project_weight.size(0)
+            if deeplab_ch == 256:
+                aspp_dropout = 0.5
+
+    model = get_deeplab_model(backbone=backbone, deeplab_ch=deeplab_ch,
+                              aspp_dropout=aspp_dropout, cpu_dist_maps=cpu_dist_maps,
+                              norm_radius=norm_radius)
+
+    model.load_state_dict(state_dict, strict=False)
+    for param in model.parameters():
+        param.requires_grad = False
+    model.to(device)
+    model.eval()
+
+    return model
+
+
+def get_iou(gt_mask, pred_mask, ignore_label=-1):
+    ignore_gt_mask_inv = gt_mask != ignore_label
+    obj_gt_mask = gt_mask == 1
+
+    intersection = np.logical_and(np.logical_and(pred_mask, obj_gt_mask), ignore_gt_mask_inv).sum()
+    union = np.logical_and(np.logical_or(pred_mask, obj_gt_mask), ignore_gt_mask_inv).sum()
+
+    return intersection / union
+
+
+def compute_noc_metric(all_ious, iou_thrs, max_clicks=20):
+    def _get_noc(iou_arr, iou_thr):
+        vals = iou_arr >= iou_thr
+        return np.argmax(vals) + 1 if np.any(vals) else max_clicks
+
+    noc_list = []
+    over_max_list = []
+    for iou_thr in iou_thrs:
+        scores_arr = np.array([_get_noc(iou_arr, iou_thr)
+                               for iou_arr in all_ious], dtype=np.int)
+
+        score = scores_arr.mean()
+        over_max = (scores_arr == max_clicks).sum()
+
+        noc_list.append(score)
+        over_max_list.append(over_max)
+
+    return noc_list, over_max_list
+
+
+def find_checkpoint(weights_folder, checkpoint_name):
+    weights_folder = Path(weights_folder)
+    if ':' in checkpoint_name:
+        model_name, checkpoint_name = checkpoint_name.split(':')
+        models_candidates = [x for x in weights_folder.glob(f'{model_name}*') if x.is_dir()]
+        assert len(models_candidates) == 1
+        model_folder = models_candidates[0]
+    else:
+        model_folder = weights_folder
+
+    if checkpoint_name.endswith('.pth'):
+        if Path(checkpoint_name).exists():
+            checkpoint_path = checkpoint_name
+        else:
+            checkpoint_path = weights_folder / checkpoint_name
+    else:
+        model_checkpoints = list(model_folder.rglob(f'{checkpoint_name}*.pth'))
+        assert len(model_checkpoints) == 1
+        checkpoint_path = model_checkpoints[0]
+
+    return str(checkpoint_path)
+
+
+def get_results_table(noc_list, over_max_list, brs_type, dataset_name, mean_spc, elapsed_time,
+                      n_clicks=20, model_name=None):
+    table_header = (f'|{"BRS Type":^13}|{"Dataset":^11}|'
+                    f'{"NoC@80%":^9}|{"NoC@85%":^9}|{"NoC@90%":^9}|'
+                    f'{">="+str(n_clicks)+"@85%":^9}|{">="+str(n_clicks)+"@90%":^9}|'
+                    f'{"SPC,s":^7}|{"Time":^9}|')
+    row_width = len(table_header)
+
+    header = f'Eval results for model: {model_name}\n' if model_name is not None else ''
+    header += '-' * row_width + '\n'
+    header += table_header + '\n' + '-' * row_width
+
+    eval_time = str(timedelta(seconds=int(elapsed_time)))
+    table_row = f'|{brs_type:^13}|{dataset_name:^11}|'
+    table_row += f'{noc_list[0]:^9.2f}|'
+    table_row += f'{noc_list[1]:^9.2f}|' if len(noc_list) > 1 else f'{"?":^9}|'
+    table_row += f'{noc_list[2]:^9.2f}|' if len(noc_list) > 2 else f'{"?":^9}|'
+    table_row += f'{over_max_list[1]:^9}|' if len(noc_list) > 1 else f'{"?":^9}|'
+    table_row += f'{over_max_list[2]:^9}|' if len(noc_list) > 2 else f'{"?":^9}|'
+    table_row += f'{mean_spc:^7.3f}|{eval_time:^9}|'
+
+    return header, table_row

inference/interact/fbrs/model/initializer.py

@@ -0,0 +1,105 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class Initializer(object):
+    def __init__(self, local_init=True, gamma=None):
+        self.local_init = local_init
+        self.gamma = gamma
+
+    def __call__(self, m):
+        if getattr(m, '__initialized', False):
+            return
+
+        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d,
+                          nn.InstanceNorm1d, nn.InstanceNorm2d, nn.InstanceNorm3d,
+                          nn.GroupNorm, nn.SyncBatchNorm)) or 'BatchNorm' in m.__class__.__name__:
+            if m.weight is not None:
+                self._init_gamma(m.weight.data)
+            if m.bias is not None:
+                self._init_beta(m.bias.data)
+        else:
+            if getattr(m, 'weight', None) is not None:
+                self._init_weight(m.weight.data)
+            if getattr(m, 'bias', None) is not None:
+                self._init_bias(m.bias.data)
+
+        if self.local_init:
+            object.__setattr__(m, '__initialized', True)
+
+    def _init_weight(self, data):
+        nn.init.uniform_(data, -0.07, 0.07)
+
+    def _init_bias(self, data):
+        nn.init.constant_(data, 0)
+
+    def _init_gamma(self, data):
+        if self.gamma is None:
+            nn.init.constant_(data, 1.0)
+        else:
+            nn.init.normal_(data, 1.0, self.gamma)
+
+    def _init_beta(self, data):
+        nn.init.constant_(data, 0)
+
+
+class Bilinear(Initializer):
+    def __init__(self, scale, groups, in_channels, **kwargs):
+        super().__init__(**kwargs)
+        self.scale = scale
+        self.groups = groups
+        self.in_channels = in_channels
+
+    def _init_weight(self, data):
+        """Reset the weight and bias."""
+        bilinear_kernel = self.get_bilinear_kernel(self.scale)
+        weight = torch.zeros_like(data)
+        for i in range(self.in_channels):
+            if self.groups == 1:
+                j = i
+            else:
+                j = 0
+            weight[i, j] = bilinear_kernel
+        data[:] = weight
+
+    @staticmethod
+    def get_bilinear_kernel(scale):
+        """Generate a bilinear upsampling kernel."""
+        kernel_size = 2 * scale - scale % 2
+        scale = (kernel_size + 1) // 2
+        center = scale - 0.5 * (1 + kernel_size % 2)
+
+        og = np.ogrid[:kernel_size, :kernel_size]
+        kernel = (1 - np.abs(og[0] - center) / scale) * (1 - np.abs(og[1] - center) / scale)
+
+        return torch.tensor(kernel, dtype=torch.float32)
+
+
+class XavierGluon(Initializer):
+    def __init__(self, rnd_type='uniform', factor_type='avg', magnitude=3, **kwargs):
+        super().__init__(**kwargs)
+
+        self.rnd_type = rnd_type
+        self.factor_type = factor_type
+        self.magnitude = float(magnitude)
+
+    def _init_weight(self, arr):
+        fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(arr)
+
+        if self.factor_type == 'avg':
+            factor = (fan_in + fan_out) / 2.0
+        elif self.factor_type == 'in':
+            factor = fan_in
+        elif self.factor_type == 'out':
+            factor = fan_out
+        else:
+            raise ValueError('Incorrect factor type')
+        scale = np.sqrt(self.magnitude / factor)
+
+        if self.rnd_type == 'uniform':
+            nn.init.uniform_(arr, -scale, scale)
+        elif self.rnd_type == 'gaussian':
+            nn.init.normal_(arr, 0, scale)
+        else:
+            raise ValueError('Unknown random type')

inference/interact/fbrs/model/is_deeplab_model.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+
+from .ops import DistMaps
+from .modeling.deeplab_v3 import DeepLabV3Plus
+from .modeling.basic_blocks import SepConvHead
+
+
+def get_deeplab_model(backbone='resnet50', deeplab_ch=256, aspp_dropout=0.5,
+                      norm_layer=nn.BatchNorm2d, backbone_norm_layer=None,
+                      use_rgb_conv=True, cpu_dist_maps=False,
+                      norm_radius=260):
+    model = DistMapsModel(
+        feature_extractor=DeepLabV3Plus(backbone=backbone,
+                                        ch=deeplab_ch,
+                                        project_dropout=aspp_dropout,
+                                        norm_layer=norm_layer,
+                                        backbone_norm_layer=backbone_norm_layer),
+        head=SepConvHead(1, in_channels=deeplab_ch, mid_channels=deeplab_ch // 2,
+                         num_layers=2, norm_layer=norm_layer),
+        use_rgb_conv=use_rgb_conv,
+        norm_layer=norm_layer,
+        norm_radius=norm_radius,
+        cpu_dist_maps=cpu_dist_maps
+    )
+
+    return model
+
+
+class DistMapsModel(nn.Module):
+    def __init__(self, feature_extractor, head, norm_layer=nn.BatchNorm2d, use_rgb_conv=True,
+                 cpu_dist_maps=False, norm_radius=260):
+        super(DistMapsModel, self).__init__()
+
+        if use_rgb_conv:
+            self.rgb_conv = nn.Sequential(
+                nn.Conv2d(in_channels=5, out_channels=8, kernel_size=1),
+                nn.LeakyReLU(negative_slope=0.2),
+                norm_layer(8),
+                nn.Conv2d(in_channels=8, out_channels=3, kernel_size=1),
+            )
+        else:
+            self.rgb_conv = None
+
+        self.dist_maps = DistMaps(norm_radius=norm_radius, spatial_scale=1.0,
+                                  cpu_mode=cpu_dist_maps)
+        self.feature_extractor = feature_extractor
+        self.head = head
+
+    def forward(self, image, points):
+        coord_features = self.dist_maps(image, points)
+
+        if self.rgb_conv is not None:
+            x = self.rgb_conv(torch.cat((image, coord_features), dim=1))
+        else:
+            c1, c2 = torch.chunk(coord_features, 2, dim=1)
+            c3 = torch.ones_like(c1)
+            coord_features = torch.cat((c1, c2, c3), dim=1)
+            x = 0.8 * image * coord_features + 0.2 * image
+
+        backbone_features = self.feature_extractor(x)
+        instance_out = self.head(backbone_features[0])
+        instance_out = nn.functional.interpolate(instance_out, size=image.size()[2:],
+                                                 mode='bilinear', align_corners=True)
+
+        return {'instances': instance_out}
+
+    def load_weights(self, path_to_weights):
+        current_state_dict = self.state_dict()
+        new_state_dict = torch.load(path_to_weights, map_location='cpu')
+        current_state_dict.update(new_state_dict)
+        self.load_state_dict(current_state_dict)
+
+    def get_trainable_params(self):
+        backbone_params = nn.ParameterList()
+        other_params = nn.ParameterList()
+
+        for name, param in self.named_parameters():
+            if param.requires_grad:
+                if 'backbone' in name:
+                    backbone_params.append(param)
+                else:
+                    other_params.append(param)
+        return backbone_params, other_params
+
+

inference/interact/fbrs/model/is_hrnet_model.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+
+from .ops import DistMaps
+from .modeling.hrnet_ocr import HighResolutionNet
+
+
+def get_hrnet_model(width=48, ocr_width=256, small=False, norm_radius=260,
+                    use_rgb_conv=True, with_aux_output=False, cpu_dist_maps=False,
+                    norm_layer=nn.BatchNorm2d):
+    model = DistMapsHRNetModel(
+        feature_extractor=HighResolutionNet(width=width, ocr_width=ocr_width, small=small,
+                                            num_classes=1, norm_layer=norm_layer),
+        use_rgb_conv=use_rgb_conv,
+        with_aux_output=with_aux_output,
+        norm_layer=norm_layer,
+        norm_radius=norm_radius,
+        cpu_dist_maps=cpu_dist_maps
+    )
+
+    return model
+
+
+class DistMapsHRNetModel(nn.Module):
+    def __init__(self, feature_extractor, use_rgb_conv=True, with_aux_output=False,
+                 norm_layer=nn.BatchNorm2d, norm_radius=260, cpu_dist_maps=False):
+        super(DistMapsHRNetModel, self).__init__()
+        self.with_aux_output = with_aux_output
+
+        if use_rgb_conv:
+            self.rgb_conv = nn.Sequential(
+                nn.Conv2d(in_channels=5, out_channels=8, kernel_size=1),
+                nn.LeakyReLU(negative_slope=0.2),
+                norm_layer(8),
+                nn.Conv2d(in_channels=8, out_channels=3, kernel_size=1),
+            )
+        else:
+            self.rgb_conv = None
+
+        self.dist_maps = DistMaps(norm_radius=norm_radius, spatial_scale=1.0, cpu_mode=cpu_dist_maps)
+        self.feature_extractor = feature_extractor
+
+    def forward(self, image, points):
+        coord_features = self.dist_maps(image, points)
+
+        if self.rgb_conv is not None:
+            x = self.rgb_conv(torch.cat((image, coord_features), dim=1))
+        else:
+            c1, c2 = torch.chunk(coord_features, 2, dim=1)
+            c3 = torch.ones_like(c1)
+            coord_features = torch.cat((c1, c2, c3), dim=1)
+            x = 0.8 * image * coord_features + 0.2 * image
+
+        feature_extractor_out = self.feature_extractor(x)
+        instance_out = feature_extractor_out[0]
+        instance_out = nn.functional.interpolate(instance_out, size=image.size()[2:],
+                                                 mode='bilinear', align_corners=True)
+        outputs = {'instances': instance_out}
+        if self.with_aux_output:
+            instance_aux_out = feature_extractor_out[1]
+            instance_aux_out = nn.functional.interpolate(instance_aux_out, size=image.size()[2:],
+                                                         mode='bilinear', align_corners=True)
+            outputs['instances_aux'] = instance_aux_out
+
+        return outputs
+
+    def load_weights(self, path_to_weights):
+        current_state_dict = self.state_dict()
+        new_state_dict = torch.load(path_to_weights)
+        current_state_dict.update(new_state_dict)
+        self.load_state_dict(current_state_dict)
+
+    def get_trainable_params(self):
+        backbone_params = nn.ParameterList()
+        other_params = nn.ParameterList()
+        other_params_keys = []
+        nonbackbone_keywords = ['rgb_conv', 'aux_head', 'cls_head', 'conv3x3_ocr', 'ocr_distri_head']
+
+        for name, param in self.named_parameters():
+            if param.requires_grad:
+                if any(x in name for x in nonbackbone_keywords):
+                    other_params.append(param)
+                    other_params_keys.append(name)
+                else:
+                    backbone_params.append(param)
+        print('Nonbackbone params:', sorted(other_params_keys))
+        return backbone_params, other_params

inference/interact/fbrs/model/losses.py

@@ -0,0 +1,134 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..utils import misc
+
+
+class NormalizedFocalLossSigmoid(nn.Module):
+    def __init__(self, axis=-1, alpha=0.25, gamma=2,
+                 from_logits=False, batch_axis=0,
+                 weight=None, size_average=True, detach_delimeter=True,
+                 eps=1e-12, scale=1.0,
+                 ignore_label=-1):
+        super(NormalizedFocalLossSigmoid, self).__init__()
+        self._axis = axis
+        self._alpha = alpha
+        self._gamma = gamma
+        self._ignore_label = ignore_label
+        self._weight = weight if weight is not None else 1.0
+        self._batch_axis = batch_axis
+
+        self._scale = scale
+        self._from_logits = from_logits
+        self._eps = eps
+        self._size_average = size_average
+        self._detach_delimeter = detach_delimeter
+        self._k_sum = 0
+
+    def forward(self, pred, label, sample_weight=None):
+        one_hot = label > 0
+        sample_weight = label != self._ignore_label
+
+        if not self._from_logits:
+            pred = torch.sigmoid(pred)
+
+        alpha = torch.where(one_hot, self._alpha * sample_weight, (1 - self._alpha) * sample_weight)
+        pt = torch.where(one_hot, pred, 1 - pred)
+        pt = torch.where(sample_weight, pt, torch.ones_like(pt))
+
+        beta = (1 - pt) ** self._gamma
+
+        sw_sum = torch.sum(sample_weight, dim=(-2, -1), keepdim=True)
+        beta_sum = torch.sum(beta, dim=(-2, -1), keepdim=True)
+        mult = sw_sum / (beta_sum + self._eps)
+        if self._detach_delimeter:
+            mult = mult.detach()
+        beta = beta * mult
+
+        ignore_area = torch.sum(label == self._ignore_label, dim=tuple(range(1, label.dim()))).cpu().numpy()
+        sample_mult = torch.mean(mult, dim=tuple(range(1, mult.dim()))).cpu().numpy()
+        if np.any(ignore_area == 0):
+            self._k_sum = 0.9 * self._k_sum + 0.1 * sample_mult[ignore_area == 0].mean()
+
+        loss = -alpha * beta * torch.log(torch.min(pt + self._eps, torch.ones(1, dtype=torch.float).to(pt.device)))
+        loss = self._weight * (loss * sample_weight)
+
+        if self._size_average:
+            bsum = torch.sum(sample_weight, dim=misc.get_dims_with_exclusion(sample_weight.dim(), self._batch_axis))
+            loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis)) / (bsum + self._eps)
+        else:
+            loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis))
+
+        return self._scale * loss
+
+    def log_states(self, sw, name, global_step):
+        sw.add_scalar(tag=name + '_k', value=self._k_sum, global_step=global_step)
+
+
+class FocalLoss(nn.Module):
+    def __init__(self, axis=-1, alpha=0.25, gamma=2,
+                 from_logits=False, batch_axis=0,
+                 weight=None, num_class=None,
+                 eps=1e-9, size_average=True, scale=1.0):
+        super(FocalLoss, self).__init__()
+        self._axis = axis
+        self._alpha = alpha
+        self._gamma = gamma
+        self._weight = weight if weight is not None else 1.0
+        self._batch_axis = batch_axis
+
+        self._scale = scale
+        self._num_class = num_class
+        self._from_logits = from_logits
+        self._eps = eps
+        self._size_average = size_average
+
+    def forward(self, pred, label, sample_weight=None):
+        if not self._from_logits:
+            pred = F.sigmoid(pred)
+
+        one_hot = label > 0
+        pt = torch.where(one_hot, pred, 1 - pred)
+
+        t = label != -1
+        alpha = torch.where(one_hot, self._alpha * t, (1 - self._alpha) * t)
+        beta = (1 - pt) ** self._gamma
+
+        loss = -alpha * beta * torch.log(torch.min(pt + self._eps, torch.ones(1, dtype=torch.float).to(pt.device)))
+        sample_weight = label != -1
+
+        loss = self._weight * (loss * sample_weight)
+
+        if self._size_average:
+            tsum = torch.sum(label == 1, dim=misc.get_dims_with_exclusion(label.dim(), self._batch_axis))
+            loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis)) / (tsum + self._eps)
+        else:
+            loss = torch.sum(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis))
+
+        return self._scale * loss
+
+
+class SigmoidBinaryCrossEntropyLoss(nn.Module):
+    def __init__(self, from_sigmoid=False, weight=None, batch_axis=0, ignore_label=-1):
+        super(SigmoidBinaryCrossEntropyLoss, self).__init__()
+        self._from_sigmoid = from_sigmoid
+        self._ignore_label = ignore_label
+        self._weight = weight if weight is not None else 1.0
+        self._batch_axis = batch_axis
+
+    def forward(self, pred, label):
+        label = label.view(pred.size())
+        sample_weight = label != self._ignore_label
+        label = torch.where(sample_weight, label, torch.zeros_like(label))
+
+        if not self._from_sigmoid:
+            loss = torch.relu(pred) - pred * label + F.softplus(-torch.abs(pred))
+        else:
+            eps = 1e-12
+            loss = -(torch.log(pred + eps) * label
+                     + torch.log(1. - pred + eps) * (1. - label))
+
+        loss = self._weight * (loss * sample_weight)
+        return torch.mean(loss, dim=misc.get_dims_with_exclusion(loss.dim(), self._batch_axis))

inference/interact/fbrs/model/metrics.py

@@ -0,0 +1,101 @@
+import torch
+import numpy as np
+
+from ..utils import misc
+
+
+class TrainMetric(object):
+    def __init__(self, pred_outputs, gt_outputs):
+        self.pred_outputs = pred_outputs
+        self.gt_outputs = gt_outputs
+
+    def update(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def get_epoch_value(self):
+        raise NotImplementedError
+
+    def reset_epoch_stats(self):
+        raise NotImplementedError
+
+    def log_states(self, sw, tag_prefix, global_step):
+        pass
+
+    @property
+    def name(self):
+        return type(self).__name__
+
+
+class AdaptiveIoU(TrainMetric):
+    def __init__(self, init_thresh=0.4, thresh_step=0.025, thresh_beta=0.99, iou_beta=0.9,
+                 ignore_label=-1, from_logits=True,
+                 pred_output='instances', gt_output='instances'):
+        super().__init__(pred_outputs=(pred_output,), gt_outputs=(gt_output,))
+        self._ignore_label = ignore_label
+        self._from_logits = from_logits
+        self._iou_thresh = init_thresh
+        self._thresh_step = thresh_step
+        self._thresh_beta = thresh_beta
+        self._iou_beta = iou_beta
+        self._ema_iou = 0.0
+        self._epoch_iou_sum = 0.0
+        self._epoch_batch_count = 0
+
+    def update(self, pred, gt):
+        gt_mask = gt > 0
+        if self._from_logits:
+            pred = torch.sigmoid(pred)
+
+        gt_mask_area = torch.sum(gt_mask, dim=(1, 2)).detach().cpu().numpy()
+        if np.all(gt_mask_area == 0):
+            return
+
+        ignore_mask = gt == self._ignore_label
+        max_iou = _compute_iou(pred > self._iou_thresh, gt_mask, ignore_mask).mean()
+        best_thresh = self._iou_thresh
+        for t in [best_thresh - self._thresh_step, best_thresh + self._thresh_step]:
+            temp_iou = _compute_iou(pred > t, gt_mask, ignore_mask).mean()
+            if temp_iou > max_iou:
+                max_iou = temp_iou
+                best_thresh = t
+
+        self._iou_thresh = self._thresh_beta * self._iou_thresh + (1 - self._thresh_beta) * best_thresh
+        self._ema_iou = self._iou_beta * self._ema_iou + (1 - self._iou_beta) * max_iou
+        self._epoch_iou_sum += max_iou
+        self._epoch_batch_count += 1
+
+    def get_epoch_value(self):
+        if self._epoch_batch_count > 0:
+            return self._epoch_iou_sum / self._epoch_batch_count
+        else:
+            return 0.0
+
+    def reset_epoch_stats(self):
+        self._epoch_iou_sum = 0.0
+        self._epoch_batch_count = 0
+
+    def log_states(self, sw, tag_prefix, global_step):
+        sw.add_scalar(tag=tag_prefix + '_ema_iou', value=self._ema_iou, global_step=global_step)
+        sw.add_scalar(tag=tag_prefix + '_iou_thresh', value=self._iou_thresh, global_step=global_step)
+
+    @property
+    def iou_thresh(self):
+        return self._iou_thresh
+
+
+def _compute_iou(pred_mask, gt_mask, ignore_mask=None, keep_ignore=False):
+    if ignore_mask is not None:
+        pred_mask = torch.where(ignore_mask, torch.zeros_like(pred_mask), pred_mask)
+
+    reduction_dims = misc.get_dims_with_exclusion(gt_mask.dim(), 0)
+    union = torch.mean((pred_mask | gt_mask).float(), dim=reduction_dims).detach().cpu().numpy()
+    intersection = torch.mean((pred_mask & gt_mask).float(), dim=reduction_dims).detach().cpu().numpy()
+    nonzero = union > 0
+
+    iou = intersection[nonzero] / union[nonzero]
+    if not keep_ignore:
+        return iou
+    else:
+        result = np.full_like(intersection, -1)
+        result[nonzero] = iou
+        return result

inference/interact/fbrs/model/modeling/basic_blocks.py

@@ -0,0 +1,71 @@
+import torch.nn as nn
+
+from ...model import ops
+
+
+class ConvHead(nn.Module):
+    def __init__(self, out_channels, in_channels=32, num_layers=1,
+                 kernel_size=3, padding=1,
+                 norm_layer=nn.BatchNorm2d):
+        super(ConvHead, self).__init__()
+        convhead = []
+
+        for i in range(num_layers):
+            convhead.extend([
+                nn.Conv2d(in_channels, in_channels, kernel_size, padding=padding),
+                nn.ReLU(),
+                norm_layer(in_channels) if norm_layer is not None else nn.Identity()
+            ])
+        convhead.append(nn.Conv2d(in_channels, out_channels, 1, padding=0))
+
+        self.convhead = nn.Sequential(*convhead)
+
+    def forward(self, *inputs):
+        return self.convhead(inputs[0])
+
+
+class SepConvHead(nn.Module):
+    def __init__(self, num_outputs, in_channels, mid_channels, num_layers=1,
+                 kernel_size=3, padding=1, dropout_ratio=0.0, dropout_indx=0,
+                 norm_layer=nn.BatchNorm2d):
+        super(SepConvHead, self).__init__()
+
+        sepconvhead = []
+
+        for i in range(num_layers):
+            sepconvhead.append(
+                SeparableConv2d(in_channels=in_channels if i == 0 else mid_channels,
+                                out_channels=mid_channels,
+                                dw_kernel=kernel_size, dw_padding=padding,
+                                norm_layer=norm_layer, activation='relu')
+            )
+            if dropout_ratio > 0 and dropout_indx == i:
+                sepconvhead.append(nn.Dropout(dropout_ratio))
+
+        sepconvhead.append(
+            nn.Conv2d(in_channels=mid_channels, out_channels=num_outputs, kernel_size=1, padding=0)
+        )
+
+        self.layers = nn.Sequential(*sepconvhead)
+
+    def forward(self, *inputs):
+        x = inputs[0]
+
+        return self.layers(x)
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, dw_kernel, dw_padding, dw_stride=1,
+                 activation=None, use_bias=False, norm_layer=None):
+        super(SeparableConv2d, self).__init__()
+        _activation = ops.select_activation_function(activation)
+        self.body = nn.Sequential(
+            nn.Conv2d(in_channels, in_channels, kernel_size=dw_kernel, stride=dw_stride,
+                      padding=dw_padding, bias=use_bias, groups=in_channels),
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=use_bias),
+            norm_layer(out_channels) if norm_layer is not None else nn.Identity(),
+            _activation()
+        )
+
+    def forward(self, x):
+        return self.body(x)

inference/interact/fbrs/model/modeling/deeplab_v3.py

@@ -0,0 +1,176 @@
+from contextlib import ExitStack
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from .basic_blocks import SeparableConv2d
+from .resnet import ResNetBackbone
+from ...model import ops
+
+
+class DeepLabV3Plus(nn.Module):
+    def __init__(self, backbone='resnet50', norm_layer=nn.BatchNorm2d,
+                 backbone_norm_layer=None,
+                 ch=256,
+                 project_dropout=0.5,
+                 inference_mode=False,
+                 **kwargs):
+        super(DeepLabV3Plus, self).__init__()
+        if backbone_norm_layer is None:
+            backbone_norm_layer = norm_layer
+
+        self.backbone_name = backbone
+        self.norm_layer = norm_layer
+        self.backbone_norm_layer = backbone_norm_layer
+        self.inference_mode = False
+        self.ch = ch
+        self.aspp_in_channels = 2048
+        self.skip_project_in_channels = 256  # layer 1 out_channels
+
+        self._kwargs = kwargs
+        if backbone == 'resnet34':
+            self.aspp_in_channels = 512
+            self.skip_project_in_channels = 64
+
+        self.backbone = ResNetBackbone(backbone=self.backbone_name, pretrained_base=False,
+                                       norm_layer=self.backbone_norm_layer, **kwargs)
+
+        self.head = _DeepLabHead(in_channels=ch + 32, mid_channels=ch, out_channels=ch,
+                                 norm_layer=self.norm_layer)
+        self.skip_project = _SkipProject(self.skip_project_in_channels, 32, norm_layer=self.norm_layer)
+        self.aspp = _ASPP(in_channels=self.aspp_in_channels,
+                          atrous_rates=[12, 24, 36],
+                          out_channels=ch,
+                          project_dropout=project_dropout,
+                          norm_layer=self.norm_layer)
+
+        if inference_mode:
+            self.set_prediction_mode()
+
+    def load_pretrained_weights(self):
+        pretrained = ResNetBackbone(backbone=self.backbone_name, pretrained_base=True,
+                                    norm_layer=self.backbone_norm_layer, **self._kwargs)
+        backbone_state_dict = self.backbone.state_dict()
+        pretrained_state_dict = pretrained.state_dict()
+
+        backbone_state_dict.update(pretrained_state_dict)
+        self.backbone.load_state_dict(backbone_state_dict)
+
+        if self.inference_mode:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+    def set_prediction_mode(self):
+        self.inference_mode = True
+        self.eval()
+
+    def forward(self, x):
+        with ExitStack() as stack:
+            if self.inference_mode:
+                stack.enter_context(torch.no_grad())
+
+            c1, _, c3, c4 = self.backbone(x)
+            c1 = self.skip_project(c1)
+
+            x = self.aspp(c4)
+            x = F.interpolate(x, c1.size()[2:], mode='bilinear', align_corners=True)
+            x = torch.cat((x, c1), dim=1)
+            x = self.head(x)
+
+        return x,
+
+
+class _SkipProject(nn.Module):
+    def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d):
+        super(_SkipProject, self).__init__()
+        _activation = ops.select_activation_function("relu")
+
+        self.skip_project = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
+            norm_layer(out_channels),
+            _activation()
+        )
+
+    def forward(self, x):
+        return self.skip_project(x)
+
+
+class _DeepLabHead(nn.Module):
+    def __init__(self, out_channels, in_channels, mid_channels=256, norm_layer=nn.BatchNorm2d):
+        super(_DeepLabHead, self).__init__()
+
+        self.block = nn.Sequential(
+            SeparableConv2d(in_channels=in_channels, out_channels=mid_channels, dw_kernel=3,
+                            dw_padding=1, activation='relu', norm_layer=norm_layer),
+            SeparableConv2d(in_channels=mid_channels, out_channels=mid_channels, dw_kernel=3,
+                            dw_padding=1, activation='relu', norm_layer=norm_layer),
+            nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=1)
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class _ASPP(nn.Module):
+    def __init__(self, in_channels, atrous_rates, out_channels=256,
+                 project_dropout=0.5, norm_layer=nn.BatchNorm2d):
+        super(_ASPP, self).__init__()
+
+        b0 = nn.Sequential(
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU()
+        )
+
+        rate1, rate2, rate3 = tuple(atrous_rates)
+        b1 = _ASPPConv(in_channels, out_channels, rate1, norm_layer)
+        b2 = _ASPPConv(in_channels, out_channels, rate2, norm_layer)
+        b3 = _ASPPConv(in_channels, out_channels, rate3, norm_layer)
+        b4 = _AsppPooling(in_channels, out_channels, norm_layer=norm_layer)
+
+        self.concurent = nn.ModuleList([b0, b1, b2, b3, b4])
+
+        project = [
+            nn.Conv2d(in_channels=5*out_channels, out_channels=out_channels,
+                      kernel_size=1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU()
+        ]
+        if project_dropout > 0:
+            project.append(nn.Dropout(project_dropout))
+        self.project = nn.Sequential(*project)
+
+    def forward(self, x):
+        x = torch.cat([block(x) for block in self.concurent], dim=1)
+
+        return self.project(x)
+
+
+class _AsppPooling(nn.Module):
+    def __init__(self, in_channels, out_channels, norm_layer):
+        super(_AsppPooling, self).__init__()
+
+        self.gap = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=1, bias=False),
+            norm_layer(out_channels),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        pool = self.gap(x)
+        return F.interpolate(pool, x.size()[2:], mode='bilinear', align_corners=True)
+
+
+def _ASPPConv(in_channels, out_channels, atrous_rate, norm_layer):
+    block = nn.Sequential(
+        nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                  kernel_size=3, padding=atrous_rate,
+                  dilation=atrous_rate, bias=False),
+        norm_layer(out_channels),
+        nn.ReLU()
+    )
+
+    return block

inference/interact/fbrs/model/modeling/hrnet_ocr.py

@@ -0,0 +1,399 @@
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch._utils
+import torch.nn.functional as F
+from .ocr import SpatialOCR_Module, SpatialGather_Module
+from .resnetv1b import BasicBlockV1b, BottleneckV1b
+
+relu_inplace = True
+
+
+class HighResolutionModule(nn.Module):
+    def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
+                 num_channels, fuse_method,multi_scale_output=True,
+                 norm_layer=nn.BatchNorm2d, align_corners=True):
+        super(HighResolutionModule, self).__init__()
+        self._check_branches(num_branches, num_blocks, num_inchannels, num_channels)
+
+        self.num_inchannels = num_inchannels
+        self.fuse_method = fuse_method
+        self.num_branches = num_branches
+        self.norm_layer = norm_layer
+        self.align_corners = align_corners
+
+        self.multi_scale_output = multi_scale_output
+
+        self.branches = self._make_branches(
+            num_branches, blocks, num_blocks, num_channels)
+        self.fuse_layers = self._make_fuse_layers()
+        self.relu = nn.ReLU(inplace=relu_inplace)
+
+    def _check_branches(self, num_branches, num_blocks, num_inchannels, num_channels):
+        if num_branches != len(num_blocks):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
+                num_branches, len(num_blocks))
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_channels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
+                num_branches, len(num_channels))
+            raise ValueError(error_msg)
+
+        if num_branches != len(num_inchannels):
+            error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
+                num_branches, len(num_inchannels))
+            raise ValueError(error_msg)
+
+    def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
+                         stride=1):
+        downsample = None
+        if stride != 1 or \
+                self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.num_inchannels[branch_index],
+                          num_channels[branch_index] * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                self.norm_layer(num_channels[branch_index] * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.num_inchannels[branch_index],
+                            num_channels[branch_index], stride,
+                            downsample=downsample, norm_layer=self.norm_layer))
+        self.num_inchannels[branch_index] = \
+            num_channels[branch_index] * block.expansion
+        for i in range(1, num_blocks[branch_index]):
+            layers.append(block(self.num_inchannels[branch_index],
+                                num_channels[branch_index],
+                                norm_layer=self.norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _make_branches(self, num_branches, block, num_blocks, num_channels):
+        branches = []
+
+        for i in range(num_branches):
+            branches.append(
+                self._make_one_branch(i, block, num_blocks, num_channels))
+
+        return nn.ModuleList(branches)
+
+    def _make_fuse_layers(self):
+        if self.num_branches == 1:
+            return None
+
+        num_branches = self.num_branches
+        num_inchannels = self.num_inchannels
+        fuse_layers = []
+        for i in range(num_branches if self.multi_scale_output else 1):
+            fuse_layer = []
+            for j in range(num_branches):
+                if j > i:
+                    fuse_layer.append(nn.Sequential(
+                        nn.Conv2d(in_channels=num_inchannels[j],
+                                  out_channels=num_inchannels[i],
+                                  kernel_size=1,
+                                  bias=False),
+                        self.norm_layer(num_inchannels[i])))
+                elif j == i:
+                    fuse_layer.append(None)
+                else:
+                    conv3x3s = []
+                    for k in range(i - j):
+                        if k == i - j - 1:
+                            num_outchannels_conv3x3 = num_inchannels[i]
+                            conv3x3s.append(nn.Sequential(
+                                nn.Conv2d(num_inchannels[j],
+                                          num_outchannels_conv3x3,
+                                          kernel_size=3, stride=2, padding=1, bias=False),
+                                self.norm_layer(num_outchannels_conv3x3)))
+                        else:
+                            num_outchannels_conv3x3 = num_inchannels[j]
+                            conv3x3s.append(nn.Sequential(
+                                nn.Conv2d(num_inchannels[j],
+                                          num_outchannels_conv3x3,
+                                          kernel_size=3, stride=2, padding=1, bias=False),
+                                self.norm_layer(num_outchannels_conv3x3),
+                                nn.ReLU(inplace=relu_inplace)))
+                    fuse_layer.append(nn.Sequential(*conv3x3s))
+            fuse_layers.append(nn.ModuleList(fuse_layer))
+
+        return nn.ModuleList(fuse_layers)
+
+    def get_num_inchannels(self):
+        return self.num_inchannels
+
+    def forward(self, x):
+        if self.num_branches == 1:
+            return [self.branches[0](x[0])]
+
+        for i in range(self.num_branches):
+            x[i] = self.branches[i](x[i])
+
+        x_fuse = []
+        for i in range(len(self.fuse_layers)):
+            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
+            for j in range(1, self.num_branches):
+                if i == j:
+                    y = y + x[j]
+                elif j > i:
+                    width_output = x[i].shape[-1]
+                    height_output = x[i].shape[-2]
+                    y = y + F.interpolate(
+                        self.fuse_layers[i][j](x[j]),
+                        size=[height_output, width_output],
+                        mode='bilinear', align_corners=self.align_corners)
+                else:
+                    y = y + self.fuse_layers[i][j](x[j])
+            x_fuse.append(self.relu(y))
+
+        return x_fuse
+
+
+class HighResolutionNet(nn.Module):
+    def __init__(self, width, num_classes, ocr_width=256, small=False,
+                 norm_layer=nn.BatchNorm2d, align_corners=True):
+        super(HighResolutionNet, self).__init__()
+        self.norm_layer = norm_layer
+        self.width = width
+        self.ocr_width = ocr_width
+        self.align_corners = align_corners
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = norm_layer(64)
+        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn2 = norm_layer(64)
+        self.relu = nn.ReLU(inplace=relu_inplace)
+
+        num_blocks = 2 if small else 4
+
+        stage1_num_channels = 64
+        self.layer1 = self._make_layer(BottleneckV1b, 64, stage1_num_channels, blocks=num_blocks)
+        stage1_out_channel = BottleneckV1b.expansion * stage1_num_channels
+
+        self.stage2_num_branches = 2
+        num_channels = [width, 2 * width]
+        num_inchannels = [
+            num_channels[i] * BasicBlockV1b.expansion for i in range(len(num_channels))]
+        self.transition1 = self._make_transition_layer(
+            [stage1_out_channel], num_inchannels)
+        self.stage2, pre_stage_channels = self._make_stage(
+            BasicBlockV1b, num_inchannels=num_inchannels, num_modules=1, num_branches=self.stage2_num_branches,
+            num_blocks=2 * [num_blocks], num_channels=num_channels)
+
+        self.stage3_num_branches = 3
+        num_channels = [width, 2 * width, 4 * width]
+        num_inchannels = [
+            num_channels[i] * BasicBlockV1b.expansion for i in range(len(num_channels))]
+        self.transition2 = self._make_transition_layer(
+            pre_stage_channels, num_inchannels)
+        self.stage3, pre_stage_channels = self._make_stage(
+            BasicBlockV1b, num_inchannels=num_inchannels,
+            num_modules=3 if small else 4, num_branches=self.stage3_num_branches,
+            num_blocks=3 * [num_blocks], num_channels=num_channels)
+
+        self.stage4_num_branches = 4
+        num_channels = [width, 2 * width, 4 * width, 8 * width]
+        num_inchannels = [
+            num_channels[i] * BasicBlockV1b.expansion for i in range(len(num_channels))]
+        self.transition3 = self._make_transition_layer(
+            pre_stage_channels, num_inchannels)
+        self.stage4, pre_stage_channels = self._make_stage(
+            BasicBlockV1b, num_inchannels=num_inchannels, num_modules=2 if small else 3,
+            num_branches=self.stage4_num_branches,
+            num_blocks=4 * [num_blocks], num_channels=num_channels)
+
+        last_inp_channels = np.int(np.sum(pre_stage_channels))
+        ocr_mid_channels = 2 * ocr_width
+        ocr_key_channels = ocr_width
+
+        self.conv3x3_ocr = nn.Sequential(
+            nn.Conv2d(last_inp_channels, ocr_mid_channels,
+                      kernel_size=3, stride=1, padding=1),
+            norm_layer(ocr_mid_channels),
+            nn.ReLU(inplace=relu_inplace),
+        )
+        self.ocr_gather_head = SpatialGather_Module(num_classes)
+
+        self.ocr_distri_head = SpatialOCR_Module(in_channels=ocr_mid_channels,
+                                                 key_channels=ocr_key_channels,
+                                                 out_channels=ocr_mid_channels,
+                                                 scale=1,
+                                                 dropout=0.05,
+                                                 norm_layer=norm_layer,
+                                                 align_corners=align_corners)
+        self.cls_head = nn.Conv2d(
+            ocr_mid_channels, num_classes, kernel_size=1, stride=1, padding=0, bias=True)
+
+        self.aux_head = nn.Sequential(
+            nn.Conv2d(last_inp_channels, last_inp_channels,
+                      kernel_size=1, stride=1, padding=0),
+            norm_layer(last_inp_channels),
+            nn.ReLU(inplace=relu_inplace),
+            nn.Conv2d(last_inp_channels, num_classes,
+                      kernel_size=1, stride=1, padding=0, bias=True)
+        )
+
+    def _make_transition_layer(
+            self, num_channels_pre_layer, num_channels_cur_layer):
+        num_branches_cur = len(num_channels_cur_layer)
+        num_branches_pre = len(num_channels_pre_layer)
+
+        transition_layers = []
+        for i in range(num_branches_cur):
+            if i < num_branches_pre:
+                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
+                    transition_layers.append(nn.Sequential(
+                        nn.Conv2d(num_channels_pre_layer[i],
+                                  num_channels_cur_layer[i],
+                                  kernel_size=3,
+                                  stride=1,
+                                  padding=1,
+                                  bias=False),
+                        self.norm_layer(num_channels_cur_layer[i]),
+                        nn.ReLU(inplace=relu_inplace)))
+                else:
+                    transition_layers.append(None)
+            else:
+                conv3x3s = []
+                for j in range(i + 1 - num_branches_pre):
+                    inchannels = num_channels_pre_layer[-1]
+                    outchannels = num_channels_cur_layer[i] \
+                        if j == i - num_branches_pre else inchannels
+                    conv3x3s.append(nn.Sequential(
+                        nn.Conv2d(inchannels, outchannels,
+                                  kernel_size=3, stride=2, padding=1, bias=False),
+                        self.norm_layer(outchannels),
+                        nn.ReLU(inplace=relu_inplace)))
+                transition_layers.append(nn.Sequential(*conv3x3s))
+
+        return nn.ModuleList(transition_layers)
+
+    def _make_layer(self, block, inplanes, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                self.norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(inplanes, planes, stride,
+                            downsample=downsample, norm_layer=self.norm_layer))
+        inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(inplanes, planes, norm_layer=self.norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _make_stage(self, block, num_inchannels,
+                    num_modules, num_branches, num_blocks, num_channels,
+                    fuse_method='SUM',
+                    multi_scale_output=True):
+        modules = []
+        for i in range(num_modules):
+            # multi_scale_output is only used last module
+            if not multi_scale_output and i == num_modules - 1:
+                reset_multi_scale_output = False
+            else:
+                reset_multi_scale_output = True
+            modules.append(
+                HighResolutionModule(num_branches,
+                                     block,
+                                     num_blocks,
+                                     num_inchannels,
+                                     num_channels,
+                                     fuse_method,
+                                     reset_multi_scale_output,
+                                     norm_layer=self.norm_layer,
+                                     align_corners=self.align_corners)
+            )
+            num_inchannels = modules[-1].get_num_inchannels()
+
+        return nn.Sequential(*modules), num_inchannels
+
+    def forward(self, x):
+        feats = self.compute_hrnet_feats(x)
+        out_aux = self.aux_head(feats)
+        feats = self.conv3x3_ocr(feats)
+
+        context = self.ocr_gather_head(feats, out_aux)
+        feats = self.ocr_distri_head(feats, context)
+        out = self.cls_head(feats)
+
+        return [out, out_aux]
+
+    def compute_hrnet_feats(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+        x = self.layer1(x)
+
+        x_list = []
+        for i in range(self.stage2_num_branches):
+            if self.transition1[i] is not None:
+                x_list.append(self.transition1[i](x))
+            else:
+                x_list.append(x)
+        y_list = self.stage2(x_list)
+
+        x_list = []
+        for i in range(self.stage3_num_branches):
+            if self.transition2[i] is not None:
+                if i < self.stage2_num_branches:
+                    x_list.append(self.transition2[i](y_list[i]))
+                else:
+                    x_list.append(self.transition2[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        y_list = self.stage3(x_list)
+
+        x_list = []
+        for i in range(self.stage4_num_branches):
+            if self.transition3[i] is not None:
+                if i < self.stage3_num_branches:
+                    x_list.append(self.transition3[i](y_list[i]))
+                else:
+                    x_list.append(self.transition3[i](y_list[-1]))
+            else:
+                x_list.append(y_list[i])
+        x = self.stage4(x_list)
+
+        # Upsampling
+        x0_h, x0_w = x[0].size(2), x[0].size(3)
+        x1 = F.interpolate(x[1], size=(x0_h, x0_w),
+                           mode='bilinear', align_corners=self.align_corners)
+        x2 = F.interpolate(x[2], size=(x0_h, x0_w),
+                           mode='bilinear', align_corners=self.align_corners)
+        x3 = F.interpolate(x[3], size=(x0_h, x0_w),
+                           mode='bilinear', align_corners=self.align_corners)
+
+        return torch.cat([x[0], x1, x2, x3], 1)
+
+    def load_pretrained_weights(self, pretrained_path=''):
+        model_dict = self.state_dict()
+
+        if not os.path.exists(pretrained_path):
+            print(f'\nFile "{pretrained_path}" does not exist.')
+            print('You need to specify the correct path to the pre-trained weights.\n'
+                  'You can download the weights for HRNet from the repository:\n'
+                  'https://github.com/HRNet/HRNet-Image-Classification')
+            exit(1)
+        pretrained_dict = torch.load(pretrained_path, map_location={'cuda:0': 'cpu'})
+        pretrained_dict = {k.replace('last_layer', 'aux_head').replace('model.', ''): v for k, v in
+                           pretrained_dict.items()}
+
+        print('model_dict-pretrained_dict:', sorted(list(set(model_dict) - set(pretrained_dict))))
+        print('pretrained_dict-model_dict:', sorted(list(set(pretrained_dict) - set(model_dict))))
+
+        pretrained_dict = {k: v for k, v in pretrained_dict.items()
+                           if k in model_dict.keys()}
+
+        model_dict.update(pretrained_dict)
+        self.load_state_dict(model_dict)

inference/interact/fbrs/model/modeling/ocr.py

@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+import torch._utils
+import torch.nn.functional as F
+
+
+class SpatialGather_Module(nn.Module):
+    """
+        Aggregate the context features according to the initial
+        predicted probability distribution.
+        Employ the soft-weighted method to aggregate the context.
+    """
+
+    def __init__(self, cls_num=0, scale=1):
+        super(SpatialGather_Module, self).__init__()
+        self.cls_num = cls_num
+        self.scale = scale
+
+    def forward(self, feats, probs):
+        batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3)
+        probs = probs.view(batch_size, c, -1)
+        feats = feats.view(batch_size, feats.size(1), -1)
+        feats = feats.permute(0, 2, 1)  # batch x hw x c
+        probs = F.softmax(self.scale * probs, dim=2)  # batch x k x hw
+        ocr_context = torch.matmul(probs, feats) \
+            .permute(0, 2, 1).unsqueeze(3)  # batch x k x c
+        return ocr_context
+
+
+class SpatialOCR_Module(nn.Module):
+    """
+    Implementation of the OCR module:
+    We aggregate the global object representation to update the representation for each pixel.
+    """
+
+    def __init__(self,
+                 in_channels,
+                 key_channels,
+                 out_channels,
+                 scale=1,
+                 dropout=0.1,
+                 norm_layer=nn.BatchNorm2d,
+                 align_corners=True):
+        super(SpatialOCR_Module, self).__init__()
+        self.object_context_block = ObjectAttentionBlock2D(in_channels, key_channels, scale,
+                                                           norm_layer, align_corners)
+        _in_channels = 2 * in_channels
+
+        self.conv_bn_dropout = nn.Sequential(
+            nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(out_channels), nn.ReLU(inplace=True)),
+            nn.Dropout2d(dropout)
+        )
+
+    def forward(self, feats, proxy_feats):
+        context = self.object_context_block(feats, proxy_feats)
+
+        output = self.conv_bn_dropout(torch.cat([context, feats], 1))
+
+        return output
+
+
+class ObjectAttentionBlock2D(nn.Module):
+    '''
+    The basic implementation for object context block
+    Input:
+        N X C X H X W
+    Parameters:
+        in_channels       : the dimension of the input feature map
+        key_channels      : the dimension after the key/query transform
+        scale             : choose the scale to downsample the input feature maps (save memory cost)
+        bn_type           : specify the bn type
+    Return:
+        N X C X H X W
+    '''
+
+    def __init__(self,
+                 in_channels,
+                 key_channels,
+                 scale=1,
+                 norm_layer=nn.BatchNorm2d,
+                 align_corners=True):
+        super(ObjectAttentionBlock2D, self).__init__()
+        self.scale = scale
+        self.in_channels = in_channels
+        self.key_channels = key_channels
+        self.align_corners = align_corners
+
+        self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
+        self.f_pixel = nn.Sequential(
+            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True)),
+            nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True))
+        )
+        self.f_object = nn.Sequential(
+            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True)),
+            nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True))
+        )
+        self.f_down = nn.Sequential(
+            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(self.key_channels), nn.ReLU(inplace=True))
+        )
+        self.f_up = nn.Sequential(
+            nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels,
+                      kernel_size=1, stride=1, padding=0, bias=False),
+            nn.Sequential(norm_layer(self.in_channels), nn.ReLU(inplace=True))
+        )
+
+    def forward(self, x, proxy):
+        batch_size, h, w = x.size(0), x.size(2), x.size(3)
+        if self.scale > 1:
+            x = self.pool(x)
+
+        query = self.f_pixel(x).view(batch_size, self.key_channels, -1)
+        query = query.permute(0, 2, 1)
+        key = self.f_object(proxy).view(batch_size, self.key_channels, -1)
+        value = self.f_down(proxy).view(batch_size, self.key_channels, -1)
+        value = value.permute(0, 2, 1)
+
+        sim_map = torch.matmul(query, key)
+        sim_map = (self.key_channels ** -.5) * sim_map
+        sim_map = F.softmax(sim_map, dim=-1)
+
+        # add bg context ...
+        context = torch.matmul(sim_map, value)
+        context = context.permute(0, 2, 1).contiguous()
+        context = context.view(batch_size, self.key_channels, *x.size()[2:])
+        context = self.f_up(context)
+        if self.scale > 1:
+            context = F.interpolate(input=context, size=(h, w),
+                                    mode='bilinear', align_corners=self.align_corners)
+
+        return context

inference/interact/fbrs/model/modeling/resnet.py

@@ -0,0 +1,39 @@
+import torch
+from .resnetv1b import resnet34_v1b, resnet50_v1s, resnet101_v1s, resnet152_v1s
+
+
+class ResNetBackbone(torch.nn.Module):
+    def __init__(self, backbone='resnet50', pretrained_base=True, dilated=True, **kwargs):
+        super(ResNetBackbone, self).__init__()
+
+        if backbone == 'resnet34':
+            pretrained = resnet34_v1b(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        elif backbone == 'resnet50':
+            pretrained = resnet50_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        elif backbone == 'resnet101':
+            pretrained = resnet101_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        elif backbone == 'resnet152':
+            pretrained = resnet152_v1s(pretrained=pretrained_base, dilated=dilated, **kwargs)
+        else:
+            raise RuntimeError(f'unknown backbone: {backbone}')
+
+        self.conv1 = pretrained.conv1
+        self.bn1 = pretrained.bn1
+        self.relu = pretrained.relu
+        self.maxpool = pretrained.maxpool
+        self.layer1 = pretrained.layer1
+        self.layer2 = pretrained.layer2
+        self.layer3 = pretrained.layer3
+        self.layer4 = pretrained.layer4
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        c1 = self.layer1(x)
+        c2 = self.layer2(c1)
+        c3 = self.layer3(c2)
+        c4 = self.layer4(c3)
+
+        return c1, c2, c3, c4

inference/interact/fbrs/model/modeling/resnetv1b.py

@@ -0,0 +1,276 @@
+import torch
+import torch.nn as nn
+GLUON_RESNET_TORCH_HUB = 'rwightman/pytorch-pretrained-gluonresnet'
+
+
+class BasicBlockV1b(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,
+                 previous_dilation=1, norm_layer=nn.BatchNorm2d):
+        super(BasicBlockV1b, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride,
+                               padding=dilation, dilation=dilation, bias=False)
+        self.bn1 = norm_layer(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1,
+                               padding=previous_dilation, dilation=previous_dilation, bias=False)
+        self.bn2 = norm_layer(planes)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out = out + residual
+        out = self.relu(out)
+
+        return out
+
+
+class BottleneckV1b(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None,
+                 previous_dilation=1, norm_layer=nn.BatchNorm2d):
+        super(BottleneckV1b, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = norm_layer(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=dilation, dilation=dilation, bias=False)
+        self.bn2 = norm_layer(planes)
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = norm_layer(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out = out + residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNetV1b(nn.Module):
+    """ Pre-trained ResNetV1b Model, which produces the strides of 8 featuremaps at conv5.
+
+    Parameters
+    ----------
+    block : Block
+        Class for the residual block. Options are BasicBlockV1, BottleneckV1.
+    layers : list of int
+        Numbers of layers in each block
+    classes : int, default 1000
+        Number of classification classes.
+    dilated : bool, default False
+        Applying dilation strategy to pretrained ResNet yielding a stride-8 model,
+        typically used in Semantic Segmentation.
+    norm_layer : object
+        Normalization layer used (default: :class:`nn.BatchNorm2d`)
+    deep_stem : bool, default False
+        Whether to replace the 7x7 conv1 with 3 3x3 convolution layers.
+    avg_down : bool, default False
+        Whether to use average pooling for projection skip connection between stages/downsample.
+    final_drop : float, default 0.0
+        Dropout ratio before the final classification layer.
+
+    Reference:
+        - He, Kaiming, et al. "Deep residual learning for image recognition."
+        Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.
+
+        - Yu, Fisher, and Vladlen Koltun. "Multi-scale context aggregation by dilated convolutions."
+    """
+    def __init__(self, block, layers, classes=1000, dilated=True, deep_stem=False, stem_width=32,
+                 avg_down=False, final_drop=0.0, norm_layer=nn.BatchNorm2d):
+        self.inplanes = stem_width*2 if deep_stem else 64
+        super(ResNetV1b, self).__init__()
+        if not deep_stem:
+            self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        else:
+            self.conv1 = nn.Sequential(
+                nn.Conv2d(3, stem_width, kernel_size=3, stride=2, padding=1, bias=False),
+                norm_layer(stem_width),
+                nn.ReLU(True),
+                nn.Conv2d(stem_width, stem_width, kernel_size=3, stride=1, padding=1, bias=False),
+                norm_layer(stem_width),
+                nn.ReLU(True),
+                nn.Conv2d(stem_width, 2*stem_width, kernel_size=3, stride=1, padding=1, bias=False)
+            )
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(True)
+        self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0], avg_down=avg_down,
+                                       norm_layer=norm_layer)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, avg_down=avg_down,
+                                       norm_layer=norm_layer)
+        if dilated:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=1, dilation=2,
+                                           avg_down=avg_down, norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=4,
+                                           avg_down=avg_down, norm_layer=norm_layer)
+        else:
+            self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                           avg_down=avg_down, norm_layer=norm_layer)
+            self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                           avg_down=avg_down, norm_layer=norm_layer)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.drop = None
+        if final_drop > 0.0:
+            self.drop = nn.Dropout(final_drop)
+        self.fc = nn.Linear(512 * block.expansion, classes)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1,
+                    avg_down=False, norm_layer=nn.BatchNorm2d):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = []
+            if avg_down:
+                if dilation == 1:
+                    downsample.append(
+                        nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True, count_include_pad=False)
+                    )
+                else:
+                    downsample.append(
+                        nn.AvgPool2d(kernel_size=1, stride=1, ceil_mode=True, count_include_pad=False)
+                    )
+                downsample.extend([
+                    nn.Conv2d(self.inplanes, out_channels=planes * block.expansion,
+                              kernel_size=1, stride=1, bias=False),
+                    norm_layer(planes * block.expansion)
+                ])
+                downsample = nn.Sequential(*downsample)
+            else:
+                downsample = nn.Sequential(
+                    nn.Conv2d(self.inplanes, out_channels=planes * block.expansion,
+                              kernel_size=1, stride=stride, bias=False),
+                    norm_layer(planes * block.expansion)
+                )
+
+        layers = []
+        if dilation in (1, 2):
+            layers.append(block(self.inplanes, planes, stride, dilation=1, downsample=downsample,
+                                previous_dilation=dilation, norm_layer=norm_layer))
+        elif dilation == 4:
+            layers.append(block(self.inplanes, planes, stride, dilation=2, downsample=downsample,
+                                previous_dilation=dilation, norm_layer=norm_layer))
+        else:
+            raise RuntimeError("=> unknown dilation size: {}".format(dilation))
+
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation,
+                                previous_dilation=dilation, norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        if self.drop is not None:
+            x = self.drop(x)
+        x = self.fc(x)
+
+        return x
+
+
+def _safe_state_dict_filtering(orig_dict, model_dict_keys):
+    filtered_orig_dict = {}
+    for k, v in orig_dict.items():
+        if k in model_dict_keys:
+            filtered_orig_dict[k] = v
+        else:
+            print(f"[ERROR] Failed to load <{k}> in backbone")
+    return filtered_orig_dict
+
+
+def resnet34_v1b(pretrained=False, **kwargs):
+    model = ResNetV1b(BasicBlockV1b, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model_dict = model.state_dict()
+        filtered_orig_dict = _safe_state_dict_filtering(
+            torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet34_v1b', pretrained=True).state_dict(),
+            model_dict.keys()
+        )
+        model_dict.update(filtered_orig_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet50_v1s(pretrained=False, **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 4, 6, 3], deep_stem=True, stem_width=64, **kwargs)
+    if pretrained:
+        model_dict = model.state_dict()
+        filtered_orig_dict = _safe_state_dict_filtering(
+            torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet50_v1s', pretrained=True).state_dict(),
+            model_dict.keys()
+        )
+        model_dict.update(filtered_orig_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet101_v1s(pretrained=False, **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 4, 23, 3], deep_stem=True, stem_width=64, **kwargs)
+    if pretrained:
+        model_dict = model.state_dict()
+        filtered_orig_dict = _safe_state_dict_filtering(
+            torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet101_v1s', pretrained=True).state_dict(),
+            model_dict.keys()
+        )
+        model_dict.update(filtered_orig_dict)
+        model.load_state_dict(model_dict)
+    return model
+
+
+def resnet152_v1s(pretrained=False, **kwargs):
+    model = ResNetV1b(BottleneckV1b, [3, 8, 36, 3], deep_stem=True, stem_width=64, **kwargs)
+    if pretrained:
+        model_dict = model.state_dict()
+        filtered_orig_dict = _safe_state_dict_filtering(
+            torch.hub.load(GLUON_RESNET_TORCH_HUB, 'gluon_resnet152_v1s', pretrained=True).state_dict(),
+            model_dict.keys()
+        )
+        model_dict.update(filtered_orig_dict)
+        model.load_state_dict(model_dict)
+    return model

inference/interact/fbrs/model/ops.py

@@ -0,0 +1,83 @@
+import torch
+from torch import nn as nn
+import numpy as np
+
+from . import initializer as initializer
+from ..utils.cython import get_dist_maps
+
+
+def select_activation_function(activation):
+    if isinstance(activation, str):
+        if activation.lower() == 'relu':
+            return nn.ReLU
+        elif activation.lower() == 'softplus':
+            return nn.Softplus
+        else:
+            raise ValueError(f"Unknown activation type {activation}")
+    elif isinstance(activation, nn.Module):
+        return activation
+    else:
+        raise ValueError(f"Unknown activation type {activation}")
+
+
+class BilinearConvTranspose2d(nn.ConvTranspose2d):
+    def __init__(self, in_channels, out_channels, scale, groups=1):
+        kernel_size = 2 * scale - scale % 2
+        self.scale = scale
+
+        super().__init__(
+            in_channels, out_channels,
+            kernel_size=kernel_size,
+            stride=scale,
+            padding=1,
+            groups=groups,
+            bias=False)
+
+        self.apply(initializer.Bilinear(scale=scale, in_channels=in_channels, groups=groups))
+
+
+class DistMaps(nn.Module):
+    def __init__(self, norm_radius, spatial_scale=1.0, cpu_mode=False):
+        super(DistMaps, self).__init__()
+        self.spatial_scale = spatial_scale
+        self.norm_radius = norm_radius
+        self.cpu_mode = cpu_mode
+
+    def get_coord_features(self, points, batchsize, rows, cols):
+        if self.cpu_mode:
+            coords = []
+            for i in range(batchsize):
+                norm_delimeter = self.spatial_scale * self.norm_radius
+                coords.append(get_dist_maps(points[i].cpu().float().numpy(), rows, cols,
+                                            norm_delimeter))
+            coords = torch.from_numpy(np.stack(coords, axis=0)).to(points.device).float()
+        else:
+            num_points = points.shape[1] // 2
+            points = points.view(-1, 2)
+            invalid_points = torch.max(points, dim=1, keepdim=False)[0] < 0
+            row_array = torch.arange(start=0, end=rows, step=1, dtype=torch.float32, device=points.device)
+            col_array = torch.arange(start=0, end=cols, step=1, dtype=torch.float32, device=points.device)
+
+            coord_rows, coord_cols = torch.meshgrid(row_array, col_array)
+            coords = torch.stack((coord_rows, coord_cols), dim=0).unsqueeze(0).repeat(points.size(0), 1, 1, 1)
+
+            add_xy = (points * self.spatial_scale).view(points.size(0), points.size(1), 1, 1)
+            coords.add_(-add_xy)
+            coords.div_(self.norm_radius * self.spatial_scale)
+            coords.mul_(coords)
+
+            coords[:, 0] += coords[:, 1]
+            coords = coords[:, :1]
+
+            coords[invalid_points, :, :, :] = 1e6
+
+            coords = coords.view(-1, num_points, 1, rows, cols)
+            coords = coords.min(dim=1)[0]  # -> (bs * num_masks * 2) x 1 x h x w
+            coords = coords.view(-1, 2, rows, cols)
+
+        coords.sqrt_().mul_(2).tanh_()
+
+        return coords
+
+    def forward(self, x, coords):
+        return self.get_coord_features(coords, x.shape[0], x.shape[2], x.shape[3])

inference/interact/fbrs/model/syncbn/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Tamaki Kojima
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

inference/interact/fbrs/model/syncbn/README.md

@@ -0,0 +1,127 @@
+# pytorch-syncbn
+
+Tamaki Kojima(tamakoji@gmail.com)
+
+## Announcement
+
+**Pytorch 1.0 support**
+
+## Overview
+This is alternative implementation of "Synchronized Multi-GPU Batch Normalization" which computes global stats across gpus instead of locally computed. SyncBN are getting important for those input image is large, and must use multi-gpu to increase the minibatch-size for the training.
+
+The code was inspired by [Pytorch-Encoding](https://github.com/zhanghang1989/PyTorch-Encoding) and [Inplace-ABN](https://github.com/mapillary/inplace_abn)
+
+## Remarks
+- Unlike [Pytorch-Encoding](https://github.com/zhanghang1989/PyTorch-Encoding), you don't need custom `nn.DataParallel`
+- Unlike [Inplace-ABN](https://github.com/mapillary/inplace_abn), you can just replace your `nn.BatchNorm2d` to this module implementation, since it will not mark for inplace operation
+- You can plug into arbitrary module written in PyTorch to enable Synchronized BatchNorm
+- Backward computation is rewritten and tested against behavior of `nn.BatchNorm2d`
+
+## Requirements
+For PyTorch, please refer to https://pytorch.org/
+
+NOTE : The code is tested only with PyTorch v1.0.0, CUDA10/CuDNN7.4.2 on ubuntu18.04
+
+It utilize Pytorch JIT mechanism to compile seamlessly, using ninja. Please install ninja-build before use.
+
+```
+sudo apt-get install ninja-build
+```
+
+Also install all dependencies for python. For pip, run:
+
+
+```
+pip install -U -r requirements.txt
+```
+
+## Build
+
+There is no need to build. just run and JIT will take care.
+JIT and cpp extensions are supported after PyTorch0.4, however it is highly recommended to use PyTorch > 1.0 due to huge design changes.
+
+## Usage
+
+Please refer to [`test.py`](./test.py) for testing the difference between `nn.BatchNorm2d` and `modules.nn.BatchNorm2d`
+
+```
+import torch
+from modules import nn as NN
+num_gpu = torch.cuda.device_count()
+model = nn.Sequential(
+    nn.Conv2d(3, 3, 1, 1, bias=False),
+    NN.BatchNorm2d(3),
+    nn.ReLU(inplace=True),
+    nn.Conv2d(3, 3, 1, 1, bias=False),
+    NN.BatchNorm2d(3),
+).cuda()
+model = nn.DataParallel(model, device_ids=range(num_gpu))
+x = torch.rand(num_gpu, 3, 2, 2).cuda()
+z = model(x)
+```
+
+## Math
+
+### Forward
+1. compute <img src="https://latex.codecogs.com/gif.latex?\sum{x_i},\sum{x_i^2}"/> in each gpu
+2. gather all <img src="https://latex.codecogs.com/gif.latex?\sum{x_i},\sum{x_i^2}"/> from workers to master and compute <img src="https://latex.codecogs.com/gif.latex?\mu,\sigma"/> where
+
+    <img src="https://latex.codecogs.com/gif.latex?\mu=\frac{\sum{x_i}}{N}"/>
+
+    and
+
+    <img src="https://latex.codecogs.com/gif.latex?\sigma^2=\frac{\sum{x_i^2}-\mu\sum{x_i}}{N}"/></a>
+
+    and then above global stats to be shared to all gpus, update running_mean and running_var by moving average using global stats.
+
+3. forward batchnorm using global stats by
+
+    <img src="https://latex.codecogs.com/gif.latex?\hat{x_i}=\frac{x_i-\mu}{\sqrt{\sigma^2&plus;\epsilon}}"/>
+
+    and then
+
+    <img src="https://latex.codecogs.com/gif.latex?y_i=\gamma\cdot\hat{x_i}&plus;\beta"/>
+
+    where <img src="https://latex.codecogs.com/gif.latex?\gamma"/> is weight parameter and <img src="https://latex.codecogs.com/gif.latex?\beta"/> is bias parameter.
+
+4. save <img src="https://latex.codecogs.com/gif.latex?x,&space;\gamma\&space;\beta,&space;\mu,&space;\sigma^2"/> for backward
+
+### Backward
+
+1. Restore saved <img src="https://latex.codecogs.com/gif.latex?x,&space;\gamma\&space;\beta,&space;\mu,&space;\sigma^2"/>
+
+2. Compute below sums on each gpu
+
+    <img src="https://latex.codecogs.com/gif.latex?\sum_{i=1}^{N_j}(\frac{dJ}{dy_i})"/>
+
+    and
+
+    <img src="https://latex.codecogs.com/gif.latex?\sum_{i=1}^{N_j}(\frac{dJ}{dy_i}\cdot\hat{x_i})"/>
+
+    where <img src="https://latex.codecogs.com/gif.latex?j\in[0,1,....,num\_gpu]"/>
+
+    then gather them at master node to sum up global, and normalize with N where N is total number of elements for each channels. Global sums are then shared among all gpus.
+
+3. compute gradients using global stats
+
+    <img src="https://latex.codecogs.com/gif.latex?\frac{dJ}{dx_i},&space;\frac{dJ}{d\gamma},&space;\frac{dJ}{d\beta}&space;"/>
+
+    where
+
+    <img src="https://latex.codecogs.com/gif.latex?\frac{dJ}{d\gamma}=\sum_{i=1}^{N}(\frac{dJ}{dy_i}\cdot\hat{x_i})"/>
+
+    and
+
+    <img src="https://latex.codecogs.com/gif.latex?\frac{dJ}{d\beta}=\sum_{i=1}^{N}(\frac{dJ}{dy_i})"/>
+
+    and finally,
+
+    <img src="https://latex.codecogs.com/gif.latex?\frac{dJ}{dx_i}=\frac{dJ}{d\hat{x_i}}\frac{d\hat{x_i}}{dx_i}+\frac{dJ}{d\mu_i}\frac{d\mu_i}{dx_i}+\frac{dJ}{d\sigma^2_i}\frac{d\sigma^2_i}{dx_i}"/>  
+
+    <img src="https://latex.codecogs.com/gif.latex?=\frac{1}{N\sqrt{(\sigma^2+\epsilon)}}(N\frac{dJ}{d\hat{x_i}}-\sum_{j=1}^{N}(\frac{dJ}{d\hat{x_j}})-\hat{x_i}\sum_{j=1}^{N}(\frac{dJ}{d\hat{x_j}}\hat{x_j}))"/>  
+
+    <img src="https://latex.codecogs.com/gif.latex?=\frac{\gamma}{N\sqrt{(\sigma^2+\epsilon)}}(N\frac{dJ}{dy_i}-\sum_{j=1}^{N}(\frac{dJ}{dy_j})-\hat{x_i}\sum_{j=1}^{N}(\frac{dJ}{dy_j}\hat{x_j}))"/>  
+
+   Note that in the implementation, normalization with N is performed at step (2) and above equation and implementation is not exactly the same, but mathematically is same.
+
+   You can go deeper on above explanation at [Kevin Zakka's Blog](https://kevinzakka.github.io/2016/09/14/batch_normalization/)