a

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.

U2Net_/model/__init__.py

@@ -0,0 +1,2 @@
+from .u2net import U2NET
+from .u2net import U2NETP

U2Net_/model/u2net.py

@@ -0,0 +1,525 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):#UNet07DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-6 ###
+class RSU6(nn.Module):#UNet06DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):#UNet05DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):#UNet04DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):#UNet04FRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(U2NET,self).__init__()
+
+        self.stage1 = RSU7(in_ch,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        #stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #-------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,d1)
+
+        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
+
+### U^2-Net small ###
+class U2NETP(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(U2NETP,self).__init__()
+
+        self.stage1 = RSU7(in_ch,16,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,16,64)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(64,16,64)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(64,16,64)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(64,16,64)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(64,16,64)
+
+        # decoder
+        self.stage5d = RSU4F(128,16,64)
+        self.stage4d = RSU4(128,16,64)
+        self.stage3d = RSU5(128,16,64)
+        self.stage2d = RSU6(128,16,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(64,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        #stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #decoder
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,d1)
+
+        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)

U2Net_/model/u2net_refactor.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+
+import math
+
+__all__ = ['U2NET_full', 'U2NET_lite']
+
+
+def _upsample_like(x, size):
+    return nn.Upsample(size=size, mode='bilinear', align_corners=False)(x)
+
+
+def _size_map(x, height):
+    # {height: size} for Upsample
+    size = list(x.shape[-2:])
+    sizes = {}
+    for h in range(1, height):
+        sizes[h] = size
+        size = [math.ceil(w / 2) for w in size]
+    return sizes
+
+
+class REBNCONV(nn.Module):
+    def __init__(self, in_ch=3, out_ch=3, dilate=1):
+        super(REBNCONV, self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch, out_ch, 3, padding=1 * dilate, dilation=1 * dilate)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.relu_s1(self.bn_s1(self.conv_s1(x)))
+
+
+class RSU(nn.Module):
+    def __init__(self, name, height, in_ch, mid_ch, out_ch, dilated=False):
+        super(RSU, self).__init__()
+        self.name = name
+        self.height = height
+        self.dilated = dilated
+        self._make_layers(height, in_ch, mid_ch, out_ch, dilated)
+
+    def forward(self, x):
+        sizes = _size_map(x, self.height)
+        x = self.rebnconvin(x)
+
+        # U-Net like symmetric encoder-decoder structure
+        def unet(x, height=1):
+            if height < self.height:
+                x1 = getattr(self, f'rebnconv{height}')(x)
+                if not self.dilated and height < self.height - 1:
+                    x2 = unet(getattr(self, 'downsample')(x1), height + 1)
+                else:
+                    x2 = unet(x1, height + 1)
+
+                x = getattr(self, f'rebnconv{height}d')(torch.cat((x2, x1), 1))
+                return _upsample_like(x, sizes[height - 1]) if not self.dilated and height > 1 else x
+            else:
+                return getattr(self, f'rebnconv{height}')(x)
+
+        return x + unet(x)
+
+    def _make_layers(self, height, in_ch, mid_ch, out_ch, dilated=False):
+        self.add_module('rebnconvin', REBNCONV(in_ch, out_ch))
+        self.add_module('downsample', nn.MaxPool2d(2, stride=2, ceil_mode=True))
+
+        self.add_module(f'rebnconv1', REBNCONV(out_ch, mid_ch))
+        self.add_module(f'rebnconv1d', REBNCONV(mid_ch * 2, out_ch))
+
+        for i in range(2, height):
+            dilate = 1 if not dilated else 2 ** (i - 1)
+            self.add_module(f'rebnconv{i}', REBNCONV(mid_ch, mid_ch, dilate=dilate))
+            self.add_module(f'rebnconv{i}d', REBNCONV(mid_ch * 2, mid_ch, dilate=dilate))
+
+        dilate = 2 if not dilated else 2 ** (height - 1)
+        self.add_module(f'rebnconv{height}', REBNCONV(mid_ch, mid_ch, dilate=dilate))
+
+
+class U2NET(nn.Module):
+    def __init__(self, cfgs, out_ch):
+        super(U2NET, self).__init__()
+        self.out_ch = out_ch
+        self._make_layers(cfgs)
+
+    def forward(self, x):
+        sizes = _size_map(x, self.height)
+        maps = []  # storage for maps
+
+        # side saliency map
+        def unet(x, height=1):
+            if height < 6:
+                x1 = getattr(self, f'stage{height}')(x)
+                x2 = unet(getattr(self, 'downsample')(x1), height + 1)
+                x = getattr(self, f'stage{height}d')(torch.cat((x2, x1), 1))
+                side(x, height)
+                return _upsample_like(x, sizes[height - 1]) if height > 1 else x
+            else:
+                x = getattr(self, f'stage{height}')(x)
+                side(x, height)
+                return _upsample_like(x, sizes[height - 1])
+
+        def side(x, h):
+            # side output saliency map (before sigmoid)
+            x = getattr(self, f'side{h}')(x)
+            x = _upsample_like(x, sizes[1])
+            maps.append(x)
+
+        def fuse():
+            # fuse saliency probability maps
+            maps.reverse()
+            x = torch.cat(maps, 1)
+            x = getattr(self, 'outconv')(x)
+            maps.insert(0, x)
+            return [torch.sigmoid(x) for x in maps]
+
+        unet(x)
+        maps = fuse()
+        return maps
+
+    def _make_layers(self, cfgs):
+        self.height = int((len(cfgs) + 1) / 2)
+        self.add_module('downsample', nn.MaxPool2d(2, stride=2, ceil_mode=True))
+        for k, v in cfgs.items():
+            # build rsu block
+            self.add_module(k, RSU(v[0], *v[1]))
+            if v[2] > 0:
+                # build side layer
+                self.add_module(f'side{v[0][-1]}', nn.Conv2d(v[2], self.out_ch, 3, padding=1))
+        # build fuse layer
+        self.add_module('outconv', nn.Conv2d(int(self.height * self.out_ch), self.out_ch, 1))
+
+
+def U2NET_full():
+    full = {
+        # cfgs for building RSUs and sides
+        # {stage : [name, (height(L), in_ch, mid_ch, out_ch, dilated), side]}
+        'stage1': ['En_1', (7, 3, 32, 64), -1],
+        'stage2': ['En_2', (6, 64, 32, 128), -1],
+        'stage3': ['En_3', (5, 128, 64, 256), -1],
+        'stage4': ['En_4', (4, 256, 128, 512), -1],
+        'stage5': ['En_5', (4, 512, 256, 512, True), -1],
+        'stage6': ['En_6', (4, 512, 256, 512, True), 512],
+        'stage5d': ['De_5', (4, 1024, 256, 512, True), 512],
+        'stage4d': ['De_4', (4, 1024, 128, 256), 256],
+        'stage3d': ['De_3', (5, 512, 64, 128), 128],
+        'stage2d': ['De_2', (6, 256, 32, 64), 64],
+        'stage1d': ['De_1', (7, 128, 16, 64), 64],
+    }
+    return U2NET(cfgs=full, out_ch=1)
+
+
+def U2NET_lite():
+    lite = {
+        # cfgs for building RSUs and sides
+        # {stage : [name, (height(L), in_ch, mid_ch, out_ch, dilated), side]}
+        'stage1': ['En_1', (7, 3, 16, 64), -1],
+        'stage2': ['En_2', (6, 64, 16, 64), -1],
+        'stage3': ['En_3', (5, 64, 16, 64), -1],
+        'stage4': ['En_4', (4, 64, 16, 64), -1],
+        'stage5': ['En_5', (4, 64, 16, 64, True), -1],
+        'stage6': ['En_6', (4, 64, 16, 64, True), 64],
+        'stage5d': ['De_5', (4, 128, 16, 64, True), 64],
+        'stage4d': ['De_4', (4, 128, 16, 64), 64],
+        'stage3d': ['De_3', (5, 128, 16, 64), 64],
+        'stage2d': ['De_2', (6, 128, 16, 64), 64],
+        'stage1d': ['De_1', (7, 128, 16, 64), 64],
+    }
+    return U2NET(cfgs=lite, out_ch=1)

cog.yaml

@@ -0,0 +1,54 @@
+# Configuration for Cog ⚙️
+# Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md
+
+build:
+  # set to true if your model requires a GPU
+  gpu: true
+  cuda: "10.1"
+
+  # a list of ubuntu apt packages to install
+  system_packages:
+    #- "libopenmpi-dev"
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+    # - "cmake-"
+
+  # python version in the form '3.8' or '3.8.12'
+  python_version: "3.7.9"
+
+  # a list of packages in the format <package-name>==<version>
+  python_packages:
+  # cmake==3.21.2
+    # - "pip==21.2.2"
+    - "cmake==3.14.3"
+    - "torch==1.7.1"
+    - "torchvision==0.8.2"
+    - "numpy==1.19.2"
+    - "ipython==7.21.0"
+    - "Pillow==8.3.1"
+    - "svgwrite==1.4.1"
+    - "svgpathtools==1.4.1"
+    - "cssutils==2.3.0"
+    - "numba==0.55.1"
+    - "torch-tools==0.1.5"
+    - "visdom==0.1.8.9"
+    - "ftfy==6.1.1"
+    - "regex==2021.8.28"
+    - "tqdm==4.62.3"
+    - "scikit-image==0.18.3"
+    - "gdown==4.4.0"
+    - "wandb==0.12.0"
+    - "tensorflow-gpu==1.15.2"
+  
+  # commands run after the environment is setup
+  run:
+    # - /root/.pyenv/versions/3.7.9/bin/python3.7 -m pip install --upgrade pip
+    - export PYTHONPATH="/diffvg/build/lib.linux-x86_64-3.7"
+    - git clone https://github.com/BachiLi/diffvg && cd diffvg && git submodule update --init --recursive && CMAKE_PREFIX_PATH=$(pyenv prefix) DIFFVG_CUDA=1 python setup.py install
+    - pip install git+https://github.com/openai/CLIP.git --no-deps
+    - gdown "https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ" -O "/src/U2Net_/saved_models/"
+    - "echo env is ready!"
+    - "echo another command if needed"
+
+# predict.py defines how predictions are run on your model
+predict: "predict.py:Predictor"

config.py

@@ -0,0 +1,144 @@
+import argparse
+import os
+import random
+
+import numpy as np
+import pydiffvg
+import torch
+import wandb
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    # =================================
+    # ============ general ============
+    # =================================
+    parser.add_argument("target", help="target image path")
+    parser.add_argument("--output_dir", type=str,
+                        help="directory to save the output images and loss")
+    parser.add_argument("--path_svg", type=str, default="none",
+                        help="if you want to load an svg file and train from it")
+    parser.add_argument("--use_gpu", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--mask_object", type=int, default=0)
+    parser.add_argument("--fix_scale", type=int, default=0)
+    parser.add_argument("--display_logs", type=int, default=0)
+    parser.add_argument("--display", type=int, default=0)
+
+    # =================================
+    # ============ wandb ============
+    # =================================
+    parser.add_argument("--use_wandb", type=int, default=0)
+    parser.add_argument("--wandb_user", type=str, default="yael-vinker")
+    parser.add_argument("--wandb_name", type=str, default="test")
+    parser.add_argument("--wandb_project_name", type=str, default="none")
+
+    # =================================
+    # =========== training ============
+    # =================================
+    parser.add_argument("--num_iter", type=int, default=500,
+                        help="number of optimization iterations")
+    parser.add_argument("--num_stages", type=int, default=1,
+                        help="training stages, you can train x strokes, then freeze them and train another x strokes etc.")
+    parser.add_argument("--lr_scheduler", type=int, default=0)
+    parser.add_argument("--lr", type=float, default=1.0)
+    parser.add_argument("--color_lr", type=float, default=0.01)
+    parser.add_argument("--color_vars_threshold", type=float, default=0.0)
+    parser.add_argument("--batch_size", type=int, default=1,
+                        help="for optimization it's only one image")
+    parser.add_argument("--save_interval", type=int, default=10)
+    parser.add_argument("--eval_interval", type=int, default=10)
+    parser.add_argument("--image_scale", type=int, default=224)
+
+    # =================================
+    # ======== strokes params =========
+    # =================================
+    parser.add_argument("--num_paths", type=int,
+                        default=16, help="number of strokes")
+    parser.add_argument("--width", type=float,
+                        default=1.5, help="stroke width")
+    parser.add_argument("--control_points_per_seg", type=int, default=4)
+    parser.add_argument("--num_segments", type=int, default=1,
+                        help="number of segments for each stroke, each stroke is a bezier curve with 4 control points")
+    parser.add_argument("--attention_init", type=int, default=1,
+                        help="if True, use the attention heads of Dino model to set the location of the initial strokes")
+    parser.add_argument("--saliency_model", type=str, default="clip")
+    parser.add_argument("--saliency_clip_model", type=str, default="ViT-B/32")
+    parser.add_argument("--xdog_intersec", type=int, default=1)
+    parser.add_argument("--mask_object_attention", type=int, default=0)
+    parser.add_argument("--softmax_temp", type=float, default=0.3)
+
+    # =================================
+    # ============= loss ==============
+    # =================================
+    parser.add_argument("--percep_loss", type=str, default="none",
+                        help="the type of perceptual loss to be used (L2/LPIPS/none)")
+    parser.add_argument("--perceptual_weight", type=float, default=0,
+                        help="weight the perceptual loss")
+    parser.add_argument("--train_with_clip", type=int, default=0)
+    parser.add_argument("--clip_weight", type=float, default=0)
+    parser.add_argument("--start_clip", type=int, default=0)
+    parser.add_argument("--num_aug_clip", type=int, default=4)
+    parser.add_argument("--include_target_in_aug", type=int, default=0)
+    parser.add_argument("--augment_both", type=int, default=1,
+                        help="if you want to apply the affine augmentation to both the sketch and image")
+    parser.add_argument("--augemntations", type=str, default="affine",
+                        help="can be any combination of: 'affine_noise_eraserchunks_eraser_press'")
+    parser.add_argument("--noise_thresh", type=float, default=0.5)
+    parser.add_argument("--aug_scale_min", type=float, default=0.7)
+    parser.add_argument("--force_sparse", type=float, default=0,
+                        help="if True, use L1 regularization on stroke's opacity to encourage small number of strokes")
+    parser.add_argument("--clip_conv_loss", type=float, default=1)
+    parser.add_argument("--clip_conv_loss_type", type=str, default="L2")
+    parser.add_argument("--clip_conv_layer_weights",
+                        type=str, default="0,0,1.0,1.0,0")
+    parser.add_argument("--clip_model_name", type=str, default="RN101")
+    parser.add_argument("--clip_fc_loss_weight", type=float, default=0.1)
+    parser.add_argument("--clip_text_guide", type=float, default=0)
+    parser.add_argument("--text_target", type=str, default="none")
+
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    args.clip_conv_layer_weights = [
+        float(item) for item in args.clip_conv_layer_weights.split(',')]
+
+    args.output_dir = os.path.join(args.output_dir, args.wandb_name)
+    if not os.path.exists(args.output_dir):
+        os.mkdir(args.output_dir)
+
+    jpg_logs_dir = f"{args.output_dir}/jpg_logs"
+    svg_logs_dir = f"{args.output_dir}/svg_logs"
+    if not os.path.exists(jpg_logs_dir):
+        os.mkdir(jpg_logs_dir)
+    if not os.path.exists(svg_logs_dir):
+        os.mkdir(svg_logs_dir)
+
+    if args.use_wandb:
+        wandb.init(project=args.wandb_project_name, entity=args.wandb_user,
+                   config=args, name=args.wandb_name, id=wandb.util.generate_id())
+
+    if args.use_gpu:
+        args.device = torch.device("cuda" if (
+            torch.cuda.is_available() and torch.cuda.device_count() > 0) else "cpu")
+    else:
+        args.device = torch.device("cpu")
+    pydiffvg.set_use_gpu(torch.cuda.is_available() and args.use_gpu)
+    pydiffvg.set_device(args.device)
+    return args
+
+
+if __name__ == "__main__":
+    # for cog predict
+    args = parse_arguments()
+    final_config = vars(args)
+    np.save(f"{args.output_dir}/config_init.npy", final_config)

display_results.py

@@ -0,0 +1,94 @@
+import argparse
+import os
+import re
+
+import imageio
+import matplotlib.pyplot as plt
+import moviepy.editor as mvp
+import numpy as np
+import pydiffvg
+import torch
+from IPython.display import Image as Image_colab
+from IPython.display import display, SVG
+from PIL import Image
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--target_file", type=str,
+                    help="target image file, located in <target_images>")
+parser.add_argument("--num_strokes", type=int)
+args = parser.parse_args()
+
+
+def read_svg(path_svg, multiply=False):
+    device = torch.device("cuda" if (
+        torch.cuda.is_available() and torch.cuda.device_count() > 0) else "cpu")
+    canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(
+        path_svg)
+    if multiply:
+        canvas_width *= 2
+        canvas_height *= 2
+        for path in shapes:
+            path.points *= 2
+            path.stroke_width *= 2
+    _render = pydiffvg.RenderFunction.apply
+    scene_args = pydiffvg.RenderFunction.serialize_scene(
+        canvas_width, canvas_height, shapes, shape_groups)
+    img = _render(canvas_width,  # width
+                  canvas_height,  # height
+                  2,   # num_samples_x
+                  2,   # num_samples_y
+                  0,   # seed
+                  None,
+                  *scene_args)
+    img = img[:, :, 3:4] * img[:, :, :3] + \
+        torch.ones(img.shape[0], img.shape[1], 3,
+                   device=device) * (1 - img[:, :, 3:4])
+    img = img[:, :, :3]
+    return img
+
+
+abs_path = os.path.abspath(os.getcwd())
+
+result_path = f"{abs_path}/output_sketches/{os.path.splitext(args.target_file)[0]}"
+svg_files = os.listdir(result_path)
+svg_files = [f for f in svg_files if "best.svg" in f and f"{args.num_strokes}strokes" in f]
+svg_output_path = f"{result_path}/{svg_files[0]}"
+
+target_path = f"{svg_output_path[:-9]}/input.png"
+
+sketch_res = read_svg(svg_output_path, multiply=True).cpu().numpy()
+sketch_res = Image.fromarray((sketch_res * 255).astype('uint8'), 'RGB')
+
+input_im = Image.open(target_path).resize((224,224))
+display(input_im)
+display(SVG(svg_output_path))
+
+p = re.compile("_best")
+best_sketch_dir = ""
+for m in p.finditer(svg_files[0]):
+    best_sketch_dir += svg_files[0][0: m.start()]
+
+
+sketches = []
+cur_path = f"{result_path}/{best_sketch_dir}"
+sketch_res.save(f"{cur_path}/final_sketch.png")
+print(f"You can download the result sketch from {cur_path}/final_sketch.png")
+
+if not os.path.exists(f"{cur_path}/svg_to_png"):
+    os.mkdir(f"{cur_path}/svg_to_png")
+if os.path.exists(f"{cur_path}/config.npy"):
+    config = np.load(f"{cur_path}/config.npy", allow_pickle=True)[()]
+    inter = config["save_interval"]
+    loss_eval = np.array(config['loss_eval'])
+    inds = np.argsort(loss_eval)
+    intervals = list(range(0, (inds[0] + 1) * inter, inter))
+    for i_ in intervals:
+        path_svg = f"{cur_path}/svg_logs/svg_iter{i_}.svg"
+        sketch = read_svg(path_svg, multiply=True).cpu().numpy()
+        sketch = Image.fromarray((sketch * 255).astype('uint8'), 'RGB')
+        # print("{0}/iter_{1:04}.png".format(cur_path, int(i_)))
+        sketch.save("{0}/{1}/iter_{2:04}.png".format(cur_path, "svg_to_png", int(i_)))
+        sketches.append(sketch)
+    imageio.mimsave(f"{cur_path}/sketch.gif", sketches)
+
+print(cur_path)

models/loss.py

@@ -0,0 +1,463 @@
+
+import collections
+import CLIP_.clip as clip
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+
+
+class Loss(nn.Module):
+    def __init__(self, args):
+        super(Loss, self).__init__()
+        self.args = args
+        self.percep_loss = args.percep_loss
+
+        self.train_with_clip = args.train_with_clip
+        self.clip_weight = args.clip_weight
+        self.start_clip = args.start_clip
+
+        self.clip_conv_loss = args.clip_conv_loss
+        self.clip_fc_loss_weight = args.clip_fc_loss_weight
+        self.clip_text_guide = args.clip_text_guide
+
+        self.losses_to_apply = self.get_losses_to_apply()
+
+        self.loss_mapper = \
+            {
+                "clip": CLIPLoss(args),
+                "clip_conv_loss": CLIPConvLoss(args)
+            }
+
+    def get_losses_to_apply(self):
+        losses_to_apply = []
+        if self.percep_loss != "none":
+            losses_to_apply.append(self.percep_loss)
+        if self.train_with_clip and self.start_clip == 0:
+            losses_to_apply.append("clip")
+        if self.clip_conv_loss:
+            losses_to_apply.append("clip_conv_loss")
+        if self.clip_text_guide:
+            losses_to_apply.append("clip_text")
+        return losses_to_apply
+
+    def update_losses_to_apply(self, epoch):
+        if "clip" not in self.losses_to_apply:
+            if self.train_with_clip:
+                if epoch > self.start_clip:
+                    self.losses_to_apply.append("clip")
+
+    def forward(self, sketches, targets, color_parameters, renderer, epoch, points_optim=None, mode="train"):
+        loss = 0
+        self.update_losses_to_apply(epoch)
+
+        losses_dict = dict.fromkeys(
+            self.losses_to_apply, torch.tensor([0.0]).to(self.args.device))
+        loss_coeffs = dict.fromkeys(self.losses_to_apply, 1.0)
+        loss_coeffs["clip"] = self.clip_weight
+        loss_coeffs["clip_text"] = self.clip_text_guide
+
+        for loss_name in self.losses_to_apply:
+            if loss_name in ["clip_conv_loss"]:
+                conv_loss = self.loss_mapper[loss_name](
+                    sketches, targets, mode)
+                for layer in conv_loss.keys():
+                    losses_dict[layer] = conv_loss[layer]
+            elif loss_name == "l2":
+                losses_dict[loss_name] = self.loss_mapper[loss_name](
+                    sketches, targets).mean()
+            else:
+                losses_dict[loss_name] = self.loss_mapper[loss_name](
+                    sketches, targets, mode).mean()
+            # loss = loss + self.loss_mapper[loss_name](sketches, targets).mean() * loss_coeffs[loss_name]
+
+        for key in self.losses_to_apply:
+            # loss = loss + losses_dict[key] * loss_coeffs[key]
+            losses_dict[key] = losses_dict[key] * loss_coeffs[key]
+        # print(losses_dict)
+        return losses_dict
+
+
+class CLIPLoss(torch.nn.Module):
+    def __init__(self, args):
+        super(CLIPLoss, self).__init__()
+
+        self.args = args
+        self.model, clip_preprocess = clip.load(
+            'ViT-B/32', args.device, jit=False)
+        self.model.eval()
+        self.preprocess = transforms.Compose(
+            [clip_preprocess.transforms[-1]])  # clip normalisation
+        self.device = args.device
+        self.NUM_AUGS = args.num_aug_clip
+        augemntations = []
+        if "affine" in args.augemntations:
+            augemntations.append(transforms.RandomPerspective(
+                fill=0, p=1.0, distortion_scale=0.5))
+            augemntations.append(transforms.RandomResizedCrop(
+                224, scale=(0.8, 0.8), ratio=(1.0, 1.0)))
+        augemntations.append(
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)))
+        self.augment_trans = transforms.Compose(augemntations)
+
+        self.calc_target = True
+        self.include_target_in_aug = args.include_target_in_aug
+        self.counter = 0
+        self.augment_both = args.augment_both
+
+    def forward(self, sketches, targets, mode="train"):
+        if self.calc_target:
+            targets_ = self.preprocess(targets).to(self.device)
+            self.targets_features = self.model.encode_image(targets_).detach()
+            self.calc_target = False
+
+        if mode == "eval":
+            # for regular clip distance, no augmentations
+            with torch.no_grad():
+                sketches = self.preprocess(sketches).to(self.device)
+                sketches_features = self.model.encode_image(sketches)
+                return 1. - torch.cosine_similarity(sketches_features, self.targets_features)
+
+        loss_clip = 0
+        sketch_augs = []
+        img_augs = []
+        for n in range(self.NUM_AUGS):
+            augmented_pair = self.augment_trans(torch.cat([sketches, targets]))
+            sketch_augs.append(augmented_pair[0].unsqueeze(0))
+
+        sketch_batch = torch.cat(sketch_augs)
+        # sketch_utils.plot_batch(img_batch, sketch_batch, self.args, self.counter, use_wandb=False, title="fc_aug{}_iter{}_{}.jpg".format(1, self.counter, mode))
+        # if self.counter % 100 == 0:
+        # sketch_utils.plot_batch(img_batch, sketch_batch, self.args, self.counter, use_wandb=False, title="aug{}_iter{}_{}.jpg".format(1, self.counter, mode))
+
+        sketch_features = self.model.encode_image(sketch_batch)
+
+        for n in range(self.NUM_AUGS):
+            loss_clip += (1. - torch.cosine_similarity(
+                sketch_features[n:n+1], self.targets_features, dim=1))
+        self.counter += 1
+        return loss_clip
+        # return 1. - torch.cosine_similarity(sketches_features, self.targets_features)
+
+
+class LPIPS(torch.nn.Module):
+    def __init__(self, pretrained=True, normalize=True, pre_relu=True, device=None):
+        """
+        Args:
+            pre_relu(bool): if True, selects features **before** reLU activations
+        """
+        super(LPIPS, self).__init__()
+        # VGG using perceptually-learned weights (LPIPS metric)
+        self.normalize = normalize
+        self.pretrained = pretrained
+        augemntations = []
+        augemntations.append(transforms.RandomPerspective(
+            fill=0, p=1.0, distortion_scale=0.5))
+        augemntations.append(transforms.RandomResizedCrop(
+            224, scale=(0.8, 0.8), ratio=(1.0, 1.0)))
+        self.augment_trans = transforms.Compose(augemntations)
+        self.feature_extractor = LPIPS._FeatureExtractor(
+            pretrained, pre_relu).to(device)
+
+    def _l2_normalize_features(self, x, eps=1e-10):
+        nrm = torch.sqrt(torch.sum(x * x, dim=1, keepdim=True))
+        return x / (nrm + eps)
+
+    def forward(self, pred, target, mode="train"):
+        """Compare VGG features of two inputs."""
+
+        # Get VGG features
+
+        sketch_augs, img_augs = [pred], [target]
+        if mode == "train":
+            for n in range(4):
+                augmented_pair = self.augment_trans(torch.cat([pred, target]))
+                sketch_augs.append(augmented_pair[0].unsqueeze(0))
+                img_augs.append(augmented_pair[1].unsqueeze(0))
+
+        xs = torch.cat(sketch_augs, dim=0)
+        ys = torch.cat(img_augs, dim=0)
+
+        pred = self.feature_extractor(xs)
+        target = self.feature_extractor(ys)
+
+        # L2 normalize features
+        if self.normalize:
+            pred = [self._l2_normalize_features(f) for f in pred]
+            target = [self._l2_normalize_features(f) for f in target]
+
+        # TODO(mgharbi) Apply Richard's linear weights?
+
+        if self.normalize:
+            diffs = [torch.sum((p - t) ** 2, 1)
+                     for (p, t) in zip(pred, target)]
+        else:
+            # mean instead of sum to avoid super high range
+            diffs = [torch.mean((p - t) ** 2, 1)
+                     for (p, t) in zip(pred, target)]
+
+        # Spatial average
+        diffs = [diff.mean([1, 2]) for diff in diffs]
+
+        return sum(diffs)
+
+    class _FeatureExtractor(torch.nn.Module):
+        def __init__(self, pretrained, pre_relu):
+            super(LPIPS._FeatureExtractor, self).__init__()
+            vgg_pretrained = models.vgg16(pretrained=pretrained).features
+
+            self.breakpoints = [0, 4, 9, 16, 23, 30]
+            if pre_relu:
+                for i, _ in enumerate(self.breakpoints[1:]):
+                    self.breakpoints[i + 1] -= 1
+
+            # Split at the maxpools
+            for i, b in enumerate(self.breakpoints[:-1]):
+                ops = torch.nn.Sequential()
+                for idx in range(b, self.breakpoints[i + 1]):
+                    op = vgg_pretrained[idx]
+                    ops.add_module(str(idx), op)
+                # print(ops)
+                self.add_module("group{}".format(i), ops)
+
+            # No gradients
+            for p in self.parameters():
+                p.requires_grad = False
+
+            # Torchvision's normalization: <https://github.com/pytorch/examples/blob/42e5b996718797e45c46a25c55b031e6768f8440/imagenet/main.py#L89-L101>
+            self.register_buffer("shift", torch.Tensor(
+                [0.485, 0.456, 0.406]).view(1, 3, 1, 1))
+            self.register_buffer("scale", torch.Tensor(
+                [0.229, 0.224, 0.225]).view(1, 3, 1, 1))
+
+        def forward(self, x):
+            feats = []
+            x = (x - self.shift) / self.scale
+            for idx in range(len(self.breakpoints) - 1):
+                m = getattr(self, "group{}".format(idx))
+                x = m(x)
+                feats.append(x)
+            return feats
+
+
+class L2_(torch.nn.Module):
+    def __init__(self):
+        """
+        Args:
+            pre_relu(bool): if True, selects features **before** reLU activations
+        """
+        super(L2_, self).__init__()
+        # VGG using perceptually-learned weights (LPIPS metric)
+        augemntations = []
+        augemntations.append(transforms.RandomPerspective(
+            fill=0, p=1.0, distortion_scale=0.5))
+        augemntations.append(transforms.RandomResizedCrop(
+            224, scale=(0.8, 0.8), ratio=(1.0, 1.0)))
+        augemntations.append(
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)))
+        self.augment_trans = transforms.Compose(augemntations)
+        # LOG.warning("LPIPS is untested")
+
+    def forward(self, pred, target, mode="train"):
+        """Compare VGG features of two inputs."""
+
+        # Get VGG features
+
+        sketch_augs, img_augs = [pred], [target]
+        if mode == "train":
+            for n in range(4):
+                augmented_pair = self.augment_trans(torch.cat([pred, target]))
+                sketch_augs.append(augmented_pair[0].unsqueeze(0))
+                img_augs.append(augmented_pair[1].unsqueeze(0))
+
+        pred = torch.cat(sketch_augs, dim=0)
+        target = torch.cat(img_augs, dim=0)
+        diffs = [torch.square(p - t).mean() for (p, t) in zip(pred, target)]
+        return sum(diffs)
+
+
+class CLIPVisualEncoder(nn.Module):
+    def __init__(self, clip_model):
+        super().__init__()
+        self.clip_model = clip_model
+        self.featuremaps = None
+
+        for i in range(12):  # 12 resblocks in VIT visual transformer
+            self.clip_model.visual.transformer.resblocks[i].register_forward_hook(
+                self.make_hook(i))
+
+    def make_hook(self, name):
+        def hook(module, input, output):
+            if len(output.shape) == 3:
+                self.featuremaps[name] = output.permute(
+                    1, 0, 2)  # LND -> NLD bs, smth, 768
+            else:
+                self.featuremaps[name] = output
+
+        return hook
+
+    def forward(self, x):
+        self.featuremaps = collections.OrderedDict()
+        fc_features = self.clip_model.encode_image(x).float()
+        featuremaps = [self.featuremaps[k] for k in range(12)]
+
+        return fc_features, featuremaps
+
+
+def l2_layers(xs_conv_features, ys_conv_features, clip_model_name):
+    return [torch.square(x_conv - y_conv).mean() for x_conv, y_conv in
+            zip(xs_conv_features, ys_conv_features)]
+
+
+def l1_layers(xs_conv_features, ys_conv_features, clip_model_name):
+    return [torch.abs(x_conv - y_conv).mean() for x_conv, y_conv in
+            zip(xs_conv_features, ys_conv_features)]
+
+
+def cos_layers(xs_conv_features, ys_conv_features, clip_model_name):
+    if "RN" in clip_model_name:
+        return [torch.square(x_conv, y_conv, dim=1).mean() for x_conv, y_conv in
+                zip(xs_conv_features, ys_conv_features)]
+    return [(1 - torch.cosine_similarity(x_conv, y_conv, dim=1)).mean() for x_conv, y_conv in
+            zip(xs_conv_features, ys_conv_features)]
+
+
+class CLIPConvLoss(torch.nn.Module):
+    def __init__(self, args):
+        super(CLIPConvLoss, self).__init__()
+        self.clip_model_name = args.clip_model_name
+        assert self.clip_model_name in [
+            "RN50",
+            "RN101",
+            "RN50x4",
+            "RN50x16",
+            "ViT-B/32",
+            "ViT-B/16",
+        ]
+
+        self.clip_conv_loss_type = args.clip_conv_loss_type
+        self.clip_fc_loss_type = "Cos"  # args.clip_fc_loss_type
+        assert self.clip_conv_loss_type in [
+            "L2", "Cos", "L1",
+        ]
+        assert self.clip_fc_loss_type in [
+            "L2", "Cos", "L1",
+        ]
+
+        self.distance_metrics = \
+            {
+                "L2": l2_layers,
+                "L1": l1_layers,
+                "Cos": cos_layers
+            }
+
+        self.model, clip_preprocess = clip.load(
+            self.clip_model_name, args.device, jit=False)
+
+        if self.clip_model_name.startswith("ViT"):
+            self.visual_encoder = CLIPVisualEncoder(self.model)
+
+        else:
+            self.visual_model = self.model.visual
+            layers = list(self.model.visual.children())
+            init_layers = torch.nn.Sequential(*layers)[:8]
+            self.layer1 = layers[8]
+            self.layer2 = layers[9]
+            self.layer3 = layers[10]
+            self.layer4 = layers[11]
+            self.att_pool2d = layers[12]
+
+        self.args = args
+
+        self.img_size = clip_preprocess.transforms[1].size
+        self.model.eval()
+        self.target_transform = transforms.Compose([
+            transforms.ToTensor(),
+        ])  # clip normalisation
+        self.normalize_transform = transforms.Compose([
+            clip_preprocess.transforms[0],  # Resize
+            clip_preprocess.transforms[1],  # CenterCrop
+            clip_preprocess.transforms[-1],  # Normalize
+        ])
+
+        self.model.eval()
+        self.device = args.device
+        self.num_augs = self.args.num_aug_clip
+
+        augemntations = []
+        if "affine" in args.augemntations:
+            augemntations.append(transforms.RandomPerspective(
+                fill=0, p=1.0, distortion_scale=0.5))
+            augemntations.append(transforms.RandomResizedCrop(
+                224, scale=(0.8, 0.8), ratio=(1.0, 1.0)))
+        augemntations.append(
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)))
+        self.augment_trans = transforms.Compose(augemntations)
+
+        self.clip_fc_layer_dims = None  # self.args.clip_fc_layer_dims
+        self.clip_conv_layer_dims = None  # self.args.clip_conv_layer_dims
+        self.clip_fc_loss_weight = args.clip_fc_loss_weight
+        self.counter = 0
+
+    def forward(self, sketch, target, mode="train"):
+        """
+        Parameters
+        ----------
+        sketch: Torch Tensor [1, C, H, W]
+        target: Torch Tensor [1, C, H, W]
+        """
+        #         y = self.target_transform(target).to(self.args.device)
+        conv_loss_dict = {}
+        x = sketch.to(self.device)
+        y = target.to(self.device)
+        sketch_augs, img_augs = [self.normalize_transform(x)], [
+            self.normalize_transform(y)]
+        if mode == "train":
+            for n in range(self.num_augs):
+                augmented_pair = self.augment_trans(torch.cat([x, y]))
+                sketch_augs.append(augmented_pair[0].unsqueeze(0))
+                img_augs.append(augmented_pair[1].unsqueeze(0))
+
+        xs = torch.cat(sketch_augs, dim=0).to(self.device)
+        ys = torch.cat(img_augs, dim=0).to(self.device)
+
+        if self.clip_model_name.startswith("RN"):
+            xs_fc_features, xs_conv_features = self.forward_inspection_clip_resnet(
+                xs.contiguous())
+            ys_fc_features, ys_conv_features = self.forward_inspection_clip_resnet(
+                ys.detach())
+
+        else:
+            xs_fc_features, xs_conv_features = self.visual_encoder(xs)
+            ys_fc_features, ys_conv_features = self.visual_encoder(ys)
+
+        conv_loss = self.distance_metrics[self.clip_conv_loss_type](
+            xs_conv_features, ys_conv_features, self.clip_model_name)
+
+        for layer, w in enumerate(self.args.clip_conv_layer_weights):
+            if w:
+                conv_loss_dict[f"clip_conv_loss_layer{layer}"] = conv_loss[layer] * w
+
+        if self.clip_fc_loss_weight:
+            # fc distance is always cos
+            fc_loss = (1 - torch.cosine_similarity(xs_fc_features,
+                       ys_fc_features, dim=1)).mean()
+            conv_loss_dict["fc"] = fc_loss * self.clip_fc_loss_weight
+
+        self.counter += 1
+        return conv_loss_dict
+
+    def forward_inspection_clip_resnet(self, x):
+        def stem(m, x):
+            for conv, bn in [(m.conv1, m.bn1), (m.conv2, m.bn2), (m.conv3, m.bn3)]:
+                x = m.relu(bn(conv(x)))
+            x = m.avgpool(x)
+            return x
+        x = x.type(self.visual_model.conv1.weight.dtype)
+        x = stem(self.visual_model, x)
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+        y = self.att_pool2d(x4)
+        return y, [x, x1, x2, x3, x4]
+        

models/painter_params.py

@@ -0,0 +1,539 @@
+import random
+import CLIP_.clip as clip
+import numpy as np
+import pydiffvg
+import sketch_utils as utils
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from scipy.ndimage.filters import gaussian_filter
+from skimage.color import rgb2gray
+from skimage.filters import threshold_otsu
+from torchvision import transforms
+
+
+class Painter(torch.nn.Module):
+    def __init__(self, args,
+                num_strokes=4,
+                num_segments=4,
+                imsize=224,
+                device=None,
+                target_im=None,
+                mask=None):
+        super(Painter, self).__init__()
+
+        self.args = args
+        self.num_paths = num_strokes
+        self.num_segments = num_segments
+        self.width = args.width
+        self.control_points_per_seg = args.control_points_per_seg
+        self.opacity_optim = args.force_sparse
+        self.num_stages = args.num_stages
+        self.add_random_noise = "noise" in args.augemntations
+        self.noise_thresh = args.noise_thresh
+        self.softmax_temp = args.softmax_temp
+
+        self.shapes = []
+        self.shape_groups = []
+        self.device = device
+        self.canvas_width, self.canvas_height = imsize, imsize
+        self.points_vars = []
+        self.color_vars = []
+        self.color_vars_threshold = args.color_vars_threshold
+
+        self.path_svg = args.path_svg
+        self.strokes_per_stage = self.num_paths
+        self.optimize_flag = []
+
+        # attention related for strokes initialisation
+        self.attention_init = args.attention_init
+        self.target_path = args.target
+        self.saliency_model = args.saliency_model
+        self.xdog_intersec = args.xdog_intersec
+        self.mask_object = args.mask_object_attention
+        
+        self.text_target = args.text_target # for clip gradients
+        self.saliency_clip_model = args.saliency_clip_model
+        self.define_attention_input(target_im)
+        self.mask = mask
+        self.attention_map = self.set_attention_map() if self.attention_init else None
+        
+        self.thresh = self.set_attention_threshold_map() if self.attention_init else None
+        self.strokes_counter = 0 # counts the number of calls to "get_path"        
+        self.epoch = 0
+        self.final_epoch = args.num_iter - 1
+        
+
+    def init_image(self, stage=0):
+        if stage > 0:
+            # if multi stages training than add new strokes on existing ones
+            # don't optimize on previous strokes
+            self.optimize_flag = [False for i in range(len(self.shapes))]
+            for i in range(self.strokes_per_stage):
+                stroke_color = torch.tensor([0.0, 0.0, 0.0, 1.0])
+                path = self.get_path()
+                self.shapes.append(path)
+                path_group = pydiffvg.ShapeGroup(shape_ids = torch.tensor([len(self.shapes) - 1]),
+                                                    fill_color = None,
+                                                    stroke_color = stroke_color)
+                self.shape_groups.append(path_group)
+                self.optimize_flag.append(True)
+
+        else:
+            num_paths_exists = 0
+            if self.path_svg != "none":
+                self.canvas_width, self.canvas_height, self.shapes, self.shape_groups = utils.load_svg(self.path_svg)
+                # if you want to add more strokes to existing ones and optimize on all of them
+                num_paths_exists = len(self.shapes)
+
+            for i in range(num_paths_exists, self.num_paths):
+                stroke_color = torch.tensor([0.0, 0.0, 0.0, 1.0])
+                path = self.get_path()
+                self.shapes.append(path)
+                path_group = pydiffvg.ShapeGroup(shape_ids = torch.tensor([len(self.shapes) - 1]),
+                                                    fill_color = None,
+                                                    stroke_color = stroke_color)
+                self.shape_groups.append(path_group)        
+            self.optimize_flag = [True for i in range(len(self.shapes))]
+        
+        img = self.render_warp()
+        img = img[:, :, 3:4] * img[:, :, :3] + torch.ones(img.shape[0], img.shape[1], 3, device = self.device) * (1 - img[:, :, 3:4])
+        img = img[:, :, :3]
+        # Convert img from HWC to NCHW
+        img = img.unsqueeze(0)
+        img = img.permute(0, 3, 1, 2).to(self.device) # NHWC -> NCHW
+        return img
+        # utils.imwrite(img.cpu(), '{}/init.png'.format(args.output_dir), gamma=args.gamma, use_wandb=args.use_wandb, wandb_name="init")
+
+    def get_image(self):
+        img = self.render_warp()
+        opacity = img[:, :, 3:4]
+        img = opacity * img[:, :, :3] + torch.ones(img.shape[0], img.shape[1], 3, device = self.device) * (1 - opacity)
+        img = img[:, :, :3]
+        # Convert img from HWC to NCHW
+        img = img.unsqueeze(0)
+        img = img.permute(0, 3, 1, 2).to(self.device) # NHWC -> NCHW
+        return img
+
+    def get_path(self):
+        points = []
+        self.num_control_points = torch.zeros(self.num_segments, dtype = torch.int32) + (self.control_points_per_seg - 2)
+        p0 = self.inds_normalised[self.strokes_counter] if self.attention_init else (random.random(), random.random())
+        points.append(p0)
+
+        for j in range(self.num_segments):
+            radius = 0.05
+            for k in range(self.control_points_per_seg - 1):
+                p1 = (p0[0] + radius * (random.random() - 0.5), p0[1] + radius * (random.random() - 0.5))
+                points.append(p1)
+                p0 = p1
+        points = torch.tensor(points).to(self.device)
+        points[:, 0] *= self.canvas_width
+        points[:, 1] *= self.canvas_height
+        
+        path = pydiffvg.Path(num_control_points = self.num_control_points,
+                                points = points,
+                                stroke_width = torch.tensor(self.width),
+                                is_closed = False)
+        self.strokes_counter += 1
+        return path
+
+    def render_warp(self):
+        if self.opacity_optim:
+            for group in self.shape_groups:
+                group.stroke_color.data[:3].clamp_(0., 0.) # to force black stroke
+                group.stroke_color.data[-1].clamp_(0., 1.) # opacity
+                # group.stroke_color.data[-1] = (group.stroke_color.data[-1] >= self.color_vars_threshold).float()
+        _render = pydiffvg.RenderFunction.apply
+        # uncomment if you want to add random noise
+        if self.add_random_noise:
+            if random.random() > self.noise_thresh:
+                eps = 0.01 * min(self.canvas_width, self.canvas_height)
+                for path in self.shapes:
+                    path.points.data.add_(eps * torch.randn_like(path.points))
+        scene_args = pydiffvg.RenderFunction.serialize_scene(\
+            self.canvas_width, self.canvas_height, self.shapes, self.shape_groups)
+        img = _render(self.canvas_width, # width
+                    self.canvas_height, # height
+                    2,   # num_samples_x
+                    2,   # num_samples_y
+                    0,   # seed
+                    None,
+                    *scene_args)
+        return img
+    
+    def parameters(self):
+        self.points_vars = []
+        # storkes' location optimization
+        for i, path in enumerate(self.shapes):
+            if self.optimize_flag[i]:
+                path.points.requires_grad = True
+                self.points_vars.append(path.points)
+        return self.points_vars
+    
+    def get_points_parans(self):
+        return self.points_vars
+    
+    def set_color_parameters(self):
+        # for storkes' color optimization (opacity)
+        self.color_vars = []
+        for i, group in enumerate(self.shape_groups):
+            if self.optimize_flag[i]:
+                group.stroke_color.requires_grad = True
+                self.color_vars.append(group.stroke_color)
+        return self.color_vars
+
+    def get_color_parameters(self):
+        return self.color_vars
+        
+    def save_svg(self, output_dir, name):
+        pydiffvg.save_svg('{}/{}.svg'.format(output_dir, name), self.canvas_width, self.canvas_height, self.shapes, self.shape_groups)
+
+
+    def dino_attn(self):
+        patch_size=8 # dino hyperparameter
+        threshold=0.6
+
+        # for dino model
+        mean_imagenet = torch.Tensor([0.485, 0.456, 0.406])[None,:,None,None].to(self.device)
+        std_imagenet = torch.Tensor([0.229, 0.224, 0.225])[None,:,None,None].to(self.device)
+        totens = transforms.Compose([
+            transforms.Resize((self.canvas_height, self.canvas_width)),
+            transforms.ToTensor()
+            ])
+
+        dino_model = torch.hub.load('facebookresearch/dino:main', 'dino_vits8').eval().to(self.device)
+        
+        self.main_im = Image.open(self.target_path).convert("RGB")
+        main_im_tensor = totens(self.main_im).to(self.device)
+        img = (main_im_tensor.unsqueeze(0) - mean_imagenet) / std_imagenet
+        w_featmap = img.shape[-2] // patch_size
+        h_featmap = img.shape[-1] // patch_size
+        
+        with torch.no_grad():
+            attn = dino_model.get_last_selfattention(img).detach().cpu()[0]
+
+        nh = attn.shape[0]
+        attn = attn[:,0,1:].reshape(nh,-1)
+        val, idx = torch.sort(attn)
+        val /= torch.sum(val, dim=1, keepdim=True)
+        cumval = torch.cumsum(val, dim=1)
+        th_attn = cumval > (1 - threshold)
+        idx2 = torch.argsort(idx)
+        for head in range(nh):
+            th_attn[head] = th_attn[head][idx2[head]]
+        th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
+        th_attn = nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=patch_size, mode="nearest")[0].cpu()
+        
+        attn = attn.reshape(nh, w_featmap, h_featmap).float()
+        attn = nn.functional.interpolate(attn.unsqueeze(0), scale_factor=patch_size, mode="nearest")[0].cpu()
+        
+        return attn
+
+
+    def define_attention_input(self, target_im):
+        model, preprocess = clip.load(self.saliency_clip_model, device=self.device, jit=False)
+        model.eval().to(self.device)
+        data_transforms = transforms.Compose([
+                    preprocess.transforms[-1],
+                ])
+        self.image_input_attn_clip = data_transforms(target_im).to(self.device)
+        
+
+    def clip_attn(self):
+        model, preprocess = clip.load(self.saliency_clip_model, device=self.device, jit=False)
+        model.eval().to(self.device)
+        text_input = clip.tokenize([self.text_target]).to(self.device)
+
+        if "RN" in self.saliency_clip_model:
+            saliency_layer = "layer4"
+            attn_map = gradCAM(
+                model.visual,
+                self.image_input_attn_clip,
+                model.encode_text(text_input).float(),
+                getattr(model.visual, saliency_layer)
+            )
+            attn_map = attn_map.squeeze().detach().cpu().numpy()
+            attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min())
+
+        else:
+            # attn_map = interpret(self.image_input_attn_clip, text_input, model, device=self.device, index=0).astype(np.float32)
+            attn_map = interpret(self.image_input_attn_clip, text_input, model, device=self.device)
+            
+        del model
+        return attn_map
+
+    def set_attention_map(self):
+        assert self.saliency_model in ["dino", "clip"]
+        if self.saliency_model == "dino":
+            return self.dino_attn()
+        elif self.saliency_model == "clip":
+            return self.clip_attn()
+        
+
+    def softmax(self, x, tau=0.2):
+        e_x = np.exp(x / tau)
+        return e_x / e_x.sum() 
+
+    def set_inds_clip(self):
+        attn_map = (self.attention_map - self.attention_map.min()) / (self.attention_map.max() - self.attention_map.min())
+        if self.xdog_intersec:
+            xdog = XDoG_()
+            im_xdog = xdog(self.image_input_attn_clip[0].permute(1,2,0).cpu().numpy(), k=10)
+            intersec_map = (1 - im_xdog) * attn_map
+            attn_map = intersec_map
+            
+        attn_map_soft = np.copy(attn_map)
+        attn_map_soft[attn_map > 0] = self.softmax(attn_map[attn_map > 0], tau=self.softmax_temp)
+        
+        k = self.num_stages * self.num_paths
+        self.inds = np.random.choice(range(attn_map.flatten().shape[0]), size=k, replace=False, p=attn_map_soft.flatten())
+        self.inds = np.array(np.unravel_index(self.inds, attn_map.shape)).T
+    
+        self.inds_normalised = np.zeros(self.inds.shape)
+        self.inds_normalised[:, 0] =  self.inds[:, 1] / self.canvas_width
+        self.inds_normalised[:, 1] =  self.inds[:, 0] / self.canvas_height
+        self.inds_normalised = self.inds_normalised.tolist()
+        return attn_map_soft
+
+
+
+    def set_inds_dino(self):
+        k = max(3, (self.num_stages * self.num_paths) // 6 + 1) # sample top 3 three points from each attention head
+        num_heads = self.attention_map.shape[0]
+        self.inds = np.zeros((k * num_heads, 2))
+        # "thresh" is used for visualisaiton purposes only
+        thresh = torch.zeros(num_heads + 1, self.attention_map.shape[1], self.attention_map.shape[2])
+        softmax = nn.Softmax(dim=1)
+        for i in range(num_heads):
+            # replace "self.attention_map[i]" with "self.attention_map" to get the highest values among
+            # all heads. 
+            topk, indices = np.unique(self.attention_map[i].numpy(), return_index=True)
+            topk = topk[::-1][:k]
+            cur_attn_map = self.attention_map[i].numpy()
+            # prob function for uniform sampling
+            prob = cur_attn_map.flatten()
+            prob[prob > topk[-1]] = 1
+            prob[prob <= topk[-1]] = 0
+            prob = prob / prob.sum()
+            thresh[i] = torch.Tensor(prob.reshape(cur_attn_map.shape))
+
+            # choose k pixels from each head            
+            inds = np.random.choice(range(cur_attn_map.flatten().shape[0]), size=k, replace=False, p=prob)
+            inds = np.unravel_index(inds, cur_attn_map.shape)
+            self.inds[i * k: i * k + k, 0] = inds[0]
+            self.inds[i * k: i * k + k, 1] = inds[1]
+        
+        # for visualisaiton
+        sum_attn = self.attention_map.sum(0).numpy()
+        mask = np.zeros(sum_attn.shape)
+        mask[thresh[:-1].sum(0) > 0] = 1
+        sum_attn = sum_attn * mask
+        sum_attn = sum_attn / sum_attn.sum()
+        thresh[-1] = torch.Tensor(sum_attn)
+
+        # sample num_paths from the chosen pixels.
+        prob_sum = sum_attn[self.inds[:,0].astype(np.int), self.inds[:,1].astype(np.int)]
+        prob_sum = prob_sum / prob_sum.sum()
+        new_inds = []
+        for i in range(self.num_stages):
+            new_inds.extend(np.random.choice(range(self.inds.shape[0]), size=self.num_paths, replace=False, p=prob_sum))
+        self.inds = self.inds[new_inds]
+        print("self.inds",self.inds.shape)
+    
+        self.inds_normalised = np.zeros(self.inds.shape)
+        self.inds_normalised[:, 0] =  self.inds[:, 1] / self.canvas_width
+        self.inds_normalised[:, 1] =  self.inds[:, 0] / self.canvas_height
+        self.inds_normalised = self.inds_normalised.tolist()
+        return thresh
+
+    def set_attention_threshold_map(self):
+        assert self.saliency_model in ["dino", "clip"]
+        if self.saliency_model == "dino":
+            return self.set_inds_dino()
+        elif self.saliency_model == "clip":
+            return self.set_inds_clip()
+        
+
+    def get_attn(self):
+        return self.attention_map
+    
+    def get_thresh(self):
+        return self.thresh
+
+    def get_inds(self):
+        return self.inds
+    
+    def get_mask(self):
+        return self.mask
+
+    def set_random_noise(self, epoch):
+        if epoch % self.args.save_interval == 0:
+            self.add_random_noise = False
+        else:
+            self.add_random_noise = "noise" in self.args.augemntations
+
+class PainterOptimizer:
+    def __init__(self, args, renderer):
+        self.renderer = renderer
+        self.points_lr = args.lr
+        self.color_lr = args.color_lr
+        self.args = args
+        self.optim_color = args.force_sparse
+
+    def init_optimizers(self):
+        self.points_optim = torch.optim.Adam(self.renderer.parameters(), lr=self.points_lr)
+        if self.optim_color:
+            self.color_optim = torch.optim.Adam(self.renderer.set_color_parameters(), lr=self.color_lr)
+
+    def update_lr(self, counter):
+        new_lr = utils.get_epoch_lr(counter, self.args)
+        for param_group in self.points_optim.param_groups:
+            param_group["lr"] = new_lr
+    
+    def zero_grad_(self):
+        self.points_optim.zero_grad()
+        if self.optim_color:
+            self.color_optim.zero_grad()
+    
+    def step_(self):
+        self.points_optim.step()
+        if self.optim_color:
+            self.color_optim.step()
+    
+    def get_lr(self):
+        return self.points_optim.param_groups[0]['lr']
+
+
+class Hook:
+    """Attaches to a module and records its activations and gradients."""
+
+    def __init__(self, module: nn.Module):
+        self.data = None
+        self.hook = module.register_forward_hook(self.save_grad)
+        
+    def save_grad(self, module, input, output):
+        self.data = output
+        output.requires_grad_(True)
+        output.retain_grad()
+        
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        self.hook.remove()
+        
+    @property
+    def activation(self) -> torch.Tensor:
+        return self.data
+    
+    @property
+    def gradient(self) -> torch.Tensor:
+        return self.data.grad
+
+
+
+
+def interpret(image, texts, model, device):
+    images = image.repeat(1, 1, 1, 1)
+    res = model.encode_image(images)
+    model.zero_grad()
+    image_attn_blocks = list(dict(model.visual.transformer.resblocks.named_children()).values())
+    num_tokens = image_attn_blocks[0].attn_probs.shape[-1]
+    R = torch.eye(num_tokens, num_tokens, dtype=image_attn_blocks[0].attn_probs.dtype).to(device)
+    R = R.unsqueeze(0).expand(1, num_tokens, num_tokens)
+    cams = [] # there are 12 attention blocks
+    for i, blk in enumerate(image_attn_blocks):
+        cam = blk.attn_probs.detach() #attn_probs shape is 12, 50, 50
+        # each patch is 7x7 so we have 49 pixels + 1 for positional encoding
+        cam = cam.reshape(1, -1, cam.shape[-1], cam.shape[-1])
+        cam = cam.clamp(min=0)
+        cam = cam.clamp(min=0).mean(dim=1) # mean of the 12 something
+        cams.append(cam)  
+        R = R + torch.bmm(cam, R)
+              
+    cams_avg = torch.cat(cams) # 12, 50, 50
+    cams_avg = cams_avg[:, 0, 1:] # 12, 1, 49
+    image_relevance = cams_avg.mean(dim=0).unsqueeze(0)
+    image_relevance = image_relevance.reshape(1, 1, 7, 7)
+    image_relevance = torch.nn.functional.interpolate(image_relevance, size=224, mode='bicubic')
+    image_relevance = image_relevance.reshape(224, 224).data.cpu().numpy().astype(np.float32)
+    image_relevance = (image_relevance - image_relevance.min()) / (image_relevance.max() - image_relevance.min())
+    return image_relevance
+
+
+# Reference: https://arxiv.org/abs/1610.02391
+def gradCAM(
+    model: nn.Module,
+    input: torch.Tensor,
+    target: torch.Tensor,
+    layer: nn.Module
+) -> torch.Tensor:
+    # Zero out any gradients at the input.
+    if input.grad is not None:
+        input.grad.data.zero_()
+        
+    # Disable gradient settings.
+    requires_grad = {}
+    for name, param in model.named_parameters():
+        requires_grad[name] = param.requires_grad
+        param.requires_grad_(False)
+        
+    # Attach a hook to the model at the desired layer.
+    assert isinstance(layer, nn.Module)
+    with Hook(layer) as hook:        
+        # Do a forward and backward pass.
+        output = model(input)
+        output.backward(target)
+
+        grad = hook.gradient.float()
+        act = hook.activation.float()
+    
+        # Global average pool gradient across spatial dimension
+        # to obtain importance weights.
+        alpha = grad.mean(dim=(2, 3), keepdim=True)
+        # Weighted combination of activation maps over channel
+        # dimension.
+        gradcam = torch.sum(act * alpha, dim=1, keepdim=True)
+        # We only want neurons with positive influence so we
+        # clamp any negative ones.
+        gradcam = torch.clamp(gradcam, min=0)
+
+    # Resize gradcam to input resolution.
+    gradcam = F.interpolate(
+        gradcam,
+        input.shape[2:],
+        mode='bicubic',
+        align_corners=False)
+    
+    # Restore gradient settings.
+    for name, param in model.named_parameters():
+        param.requires_grad_(requires_grad[name])
+        
+    return gradcam    
+
+
+class XDoG_(object):
+    def __init__(self):
+        super(XDoG_, self).__init__()
+        self.gamma=0.98
+        self.phi=200
+        self.eps=-0.1
+        self.sigma=0.8
+        self.binarize=True
+        
+    def __call__(self, im, k=10):
+        if im.shape[2] == 3:
+            im = rgb2gray(im)
+        imf1 = gaussian_filter(im, self.sigma)
+        imf2 = gaussian_filter(im, self.sigma * k)
+        imdiff = imf1 - self.gamma * imf2
+        imdiff = (imdiff < self.eps) * 1.0  + (imdiff >= self.eps) * (1.0 + np.tanh(self.phi * imdiff))
+        imdiff -= imdiff.min()
+        imdiff /= imdiff.max()
+        if self.binarize:
+            th = threshold_otsu(imdiff)
+            imdiff = imdiff >= th
+        imdiff = imdiff.astype('float32')
+        return imdiff

painterly_rendering.py

@@ -0,0 +1,190 @@
+import warnings
+
+warnings.filterwarnings('ignore')
+warnings.simplefilter('ignore')
+
+import argparse
+import math
+import os
+import sys
+import time
+import traceback
+
+import numpy as np
+import PIL
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import wandb
+from PIL import Image
+from torchvision import models, transforms
+from tqdm.auto import tqdm, trange
+
+import config
+import sketch_utils as utils
+from models.loss import Loss
+from models.painter_params import Painter, PainterOptimizer
+from IPython.display import display, SVG
+
+
+def load_renderer(args, target_im=None, mask=None):
+    renderer = Painter(num_strokes=args.num_paths, args=args,
+                       num_segments=args.num_segments,
+                       imsize=args.image_scale,
+                       device=args.device,
+                       target_im=target_im,
+                       mask=mask)
+    renderer = renderer.to(args.device)
+    return renderer
+
+
+def get_target(args):
+    target = Image.open(args.target)
+    if target.mode == "RGBA":
+        # Create a white rgba background
+        new_image = Image.new("RGBA", target.size, "WHITE")
+        # Paste the image on the background.
+        new_image.paste(target, (0, 0), target)
+        target = new_image
+    target = target.convert("RGB")
+    masked_im, mask = utils.get_mask_u2net(args, target)
+    if args.mask_object:
+        target = masked_im
+    if args.fix_scale:
+        target = utils.fix_image_scale(target)
+
+    transforms_ = []
+    if target.size[0] != target.size[1]:
+        transforms_.append(transforms.Resize(
+            (args.image_scale, args.image_scale), interpolation=PIL.Image.BICUBIC))
+    else:
+        transforms_.append(transforms.Resize(
+            args.image_scale, interpolation=PIL.Image.BICUBIC))
+        transforms_.append(transforms.CenterCrop(args.image_scale))
+    transforms_.append(transforms.ToTensor())
+    data_transforms = transforms.Compose(transforms_)
+    target_ = data_transforms(target).unsqueeze(0).to(args.device)
+    return target_, mask
+
+
+def main(args):
+    loss_func = Loss(args)
+    inputs, mask = get_target(args)
+    utils.log_input(args.use_wandb, 0, inputs, args.output_dir)
+    renderer = load_renderer(args, inputs, mask)
+
+    optimizer = PainterOptimizer(args, renderer)
+    counter = 0
+    configs_to_save = {"loss_eval": []}
+    best_loss, best_fc_loss = 100, 100
+    best_iter, best_iter_fc = 0, 0
+    min_delta = 1e-5
+    terminate = False
+
+    renderer.set_random_noise(0)
+    img = renderer.init_image(stage=0)
+    optimizer.init_optimizers()
+    
+    # not using tdqm for jupyter demo
+    if args.display:
+        epoch_range = range(args.num_iter)
+    else:
+        epoch_range = tqdm(range(args.num_iter))
+
+    for epoch in epoch_range:
+        if not args.display:
+            epoch_range.refresh()
+        renderer.set_random_noise(epoch)
+        if args.lr_scheduler:
+            optimizer.update_lr(counter)
+
+        start = time.time()
+        optimizer.zero_grad_()
+        sketches = renderer.get_image().to(args.device)
+        losses_dict = loss_func(sketches, inputs.detach(
+        ), renderer.get_color_parameters(), renderer, counter, optimizer)
+        loss = sum(list(losses_dict.values()))
+        loss.backward()
+        optimizer.step_()
+        if epoch % args.save_interval == 0:
+            utils.plot_batch(inputs, sketches, f"{args.output_dir}/jpg_logs", counter,
+                             use_wandb=args.use_wandb, title=f"iter{epoch}.jpg")
+            renderer.save_svg(
+                f"{args.output_dir}/svg_logs", f"svg_iter{epoch}")
+        if epoch % args.eval_interval == 0:
+            with torch.no_grad():
+                losses_dict_eval = loss_func(sketches, inputs, renderer.get_color_parameters(
+                ), renderer.get_points_parans(), counter, optimizer, mode="eval")
+                loss_eval = sum(list(losses_dict_eval.values()))
+                configs_to_save["loss_eval"].append(loss_eval.item())
+                for k in losses_dict_eval.keys():
+                    if k not in configs_to_save.keys():
+                        configs_to_save[k] = []
+                    configs_to_save[k].append(losses_dict_eval[k].item())
+                if args.clip_fc_loss_weight:
+                    if losses_dict_eval["fc"].item() < best_fc_loss:
+                        best_fc_loss = losses_dict_eval["fc"].item(
+                        ) / args.clip_fc_loss_weight
+                        best_iter_fc = epoch
+                # print(
+                #     f"eval iter[{epoch}/{args.num_iter}] loss[{loss.item()}] time[{time.time() - start}]")
+
+                cur_delta = loss_eval.item() - best_loss
+                if abs(cur_delta) > min_delta:
+                    if cur_delta < 0:
+                        best_loss = loss_eval.item()
+                        best_iter = epoch
+                        terminate = False
+                        utils.plot_batch(
+                            inputs, sketches, args.output_dir, counter, use_wandb=args.use_wandb, title="best_iter.jpg")
+                        renderer.save_svg(args.output_dir, "best_iter")
+
+                if args.use_wandb:
+                    wandb.run.summary["best_loss"] = best_loss
+                    wandb.run.summary["best_loss_fc"] = best_fc_loss
+                    wandb_dict = {"delta": cur_delta,
+                                  "loss_eval": loss_eval.item()}
+                    for k in losses_dict_eval.keys():
+                        wandb_dict[k + "_eval"] = losses_dict_eval[k].item()
+                    wandb.log(wandb_dict, step=counter)
+
+                if abs(cur_delta) <= min_delta:
+                    if terminate:
+                        break
+                    terminate = True
+
+        if counter == 0 and args.attention_init:
+            utils.plot_atten(renderer.get_attn(), renderer.get_thresh(), inputs, renderer.get_inds(),
+                             args.use_wandb, "{}/{}.jpg".format(
+                                 args.output_dir, "attention_map"),
+                             args.saliency_model, args.display_logs)
+
+        if args.use_wandb:
+            wandb_dict = {"loss": loss.item(), "lr": optimizer.get_lr()}
+            for k in losses_dict.keys():
+                wandb_dict[k] = losses_dict[k].item()
+            wandb.log(wandb_dict, step=counter)
+
+        counter += 1
+
+    renderer.save_svg(args.output_dir, "final_svg")
+    path_svg = os.path.join(args.output_dir, "best_iter.svg")
+    utils.log_sketch_summary_final(
+        path_svg, args.use_wandb, args.device, best_iter, best_loss, "best total")
+
+    return configs_to_save
+
+if __name__ == "__main__":
+    args = config.parse_arguments()
+    final_config = vars(args)
+    try:
+        configs_to_save = main(args)
+    except BaseException as err:
+        print(f"Unexpected error occurred:\n {err}")
+        print(traceback.format_exc())
+        sys.exit(1)
+    for k in configs_to_save.keys():
+        final_config[k] = configs_to_save[k]
+    np.save(f"{args.output_dir}/config.npy", final_config)
+    if args.use_wandb:
+        wandb.finish()

predict.py

@@ -0,0 +1,317 @@
+# sudo cog push r8.im/yael-vinker/clipasso
+
+# Prediction interface for Cog ⚙️
+# https://github.com/replicate/cog/blob/main/docs/python.md
+
+import warnings
+
+warnings.filterwarnings('ignore')
+warnings.simplefilter('ignore')
+
+from cog import BasePredictor, Input, Path
+import subprocess as sp
+import os
+import re
+
+import imageio
+import matplotlib.pyplot as plt
+import numpy as np
+import pydiffvg
+import torch
+from PIL import Image
+import multiprocessing as mp
+from shutil import copyfile
+
+import argparse
+import math
+import sys
+import time
+import traceback
+
+import PIL
+import torch.nn as nn
+import torch.nn.functional as F
+import wandb
+from torchvision import models, transforms
+from tqdm import tqdm
+
+import config
+import sketch_utils as utils
+from models.loss import Loss
+from models.painter_params import Painter, PainterOptimizer
+
+
+class Predictor(BasePredictor):
+    def setup(self):
+        """Load the model into memory to make running multiple predictions efficient"""
+        self.num_iter = 2001
+        self.save_interval = 100
+        self.num_sketches = 3
+        self.use_gpu = True
+
+    def predict(
+        self,
+        target_image: Path = Input(description="Input image (square, without background)"),
+        num_strokes: int = Input(description="The number of strokes used to create the sketch, which determines the level of abstraction",default=16),
+        trials: int = Input(description="It is recommended to use 3 trials to recieve the best sketch, but it might be slower",default=3),
+        mask_object: int = Input(description="It is recommended to use images without a background, however, if your image contains a background, you can mask it out by using this flag with 1 as an argument",default=0),
+        fix_scale: int = Input(description="If your image is not squared, it might be cut off, it is recommended to use this flag with 1 as input to automatically fix the scale without cutting the image",default=0),
+    ) -> Path:
+        
+        self.num_sketches = trials
+        target_image_name = os.path.basename(str(target_image))
+
+        multiprocess = False
+        abs_path = os.path.abspath(os.getcwd())
+
+        target = str(target_image)
+        assert os.path.isfile(target), f"{target} does not exists!"
+
+        test_name = os.path.splitext(target_image_name)[0]
+        output_dir = f"{abs_path}/output_sketches/{test_name}/"
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        print("=" * 50)
+        print(f"Processing [{target_image_name}] ...")
+        print(f"Results will be saved to \n[{output_dir}] ...")
+        print("=" * 50)
+
+        if not torch.cuda.is_available():
+            self.use_gpu = False
+            print("CUDA is not configured with GPU, running with CPU instead.")
+            print("Note that this will be very slow, it is recommended to use colab.")
+        print(f"GPU: {self.use_gpu}")
+        seeds = list(range(0, self.num_sketches * 1000, 1000))
+
+        losses_all = {}
+        
+        for seed in seeds:
+            wandb_name = f"{test_name}_{num_strokes}strokes_seed{seed}"
+            sp.run(["python", "config.py", target,
+                        "--num_paths", str(num_strokes),
+                        "--output_dir", output_dir,
+                        "--wandb_name", wandb_name,
+                        "--num_iter", str(self.num_iter),
+                        "--save_interval", str(self.save_interval),
+                        "--seed", str(seed),
+                        "--use_gpu", str(int(self.use_gpu)),
+                        "--fix_scale", str(fix_scale),
+                        "--mask_object", str(mask_object),
+                        "--mask_object_attention", str(
+                            mask_object),
+                        "--display_logs", str(int(0))])
+            config_init = np.load(f"{output_dir}/{wandb_name}/config_init.npy", allow_pickle=True)[()]
+            args = Args(config_init)
+            args.cog_display = True
+            
+            final_config = vars(args)
+            try:
+                configs_to_save = main(args)
+            except BaseException as err:
+                print(f"Unexpected error occurred:\n {err}")
+                print(traceback.format_exc())
+                sys.exit(1)
+            for k in configs_to_save.keys():
+                final_config[k] = configs_to_save[k]
+            np.save(f"{args.output_dir}/config.npy", final_config)
+            if args.use_wandb:
+                wandb.finish()
+
+            config = np.load(f"{output_dir}/{wandb_name}/config.npy",
+                            allow_pickle=True)[()]
+            loss_eval = np.array(config['loss_eval'])
+            inds = np.argsort(loss_eval)
+            losses_all[wandb_name] = loss_eval[inds][0]
+            # return Path(f"{output_dir}/{wandb_name}/best_iter.svg") 
+
+
+        sorted_final = dict(sorted(losses_all.items(), key=lambda item: item[1]))
+        copyfile(f"{output_dir}/{list(sorted_final.keys())[0]}/best_iter.svg",
+                f"{output_dir}/{list(sorted_final.keys())[0]}_best.svg")
+        target_path = f"{abs_path}/target_images/{target_image_name}"
+        svg_files = os.listdir(output_dir)
+        svg_files = [f for f in svg_files if "best.svg" in f]
+        svg_output_path = f"{output_dir}/{svg_files[0]}"
+        sketch_res = read_svg(svg_output_path, multiply=True).cpu().numpy()
+        sketch_res = Image.fromarray((sketch_res * 255).astype('uint8'), 'RGB')
+        sketch_res.save(f"{abs_path}/output_sketches/sketch.png")
+        return Path(svg_output_path)
+
+
+class Args():
+    def __init__(self, config):
+        for k in config.keys():
+            setattr(self, k, config[k])
+        
+
+def load_renderer(args, target_im=None, mask=None):
+    renderer = Painter(num_strokes=args.num_paths, args=args,
+                       num_segments=args.num_segments,
+                       imsize=args.image_scale,
+                       device=args.device,
+                       target_im=target_im,
+                       mask=mask)
+    renderer = renderer.to(args.device)
+    return renderer
+
+
+def get_target(args):
+    target = Image.open(args.target)
+    if target.mode == "RGBA":
+        # Create a white rgba background
+        new_image = Image.new("RGBA", target.size, "WHITE")
+        # Paste the image on the background.
+        new_image.paste(target, (0, 0), target)
+        target = new_image
+    target = target.convert("RGB")
+    masked_im, mask = utils.get_mask_u2net(args, target)
+    if args.mask_object:
+        target = masked_im
+    if args.fix_scale:
+        target = utils.fix_image_scale(target)
+
+    transforms_ = []
+    if target.size[0] != target.size[1]:
+        transforms_.append(transforms.Resize(
+            (args.image_scale, args.image_scale), interpolation=PIL.Image.BICUBIC))
+    else:
+        transforms_.append(transforms.Resize(
+            args.image_scale, interpolation=PIL.Image.BICUBIC))
+        transforms_.append(transforms.CenterCrop(args.image_scale))
+    transforms_.append(transforms.ToTensor())
+    data_transforms = transforms.Compose(transforms_)
+    target_ = data_transforms(target).unsqueeze(0).to(args.device)
+    return target_, mask
+
+
+def main(args):
+    loss_func = Loss(args)
+    inputs, mask = get_target(args)
+    utils.log_input(args.use_wandb, 0, inputs, args.output_dir)
+    renderer = load_renderer(args, inputs, mask)
+
+    optimizer = PainterOptimizer(args, renderer)
+    counter = 0
+    configs_to_save = {"loss_eval": []}
+    best_loss, best_fc_loss = 100, 100
+    best_iter, best_iter_fc = 0, 0
+    min_delta = 1e-5
+    terminate = False
+
+    renderer.set_random_noise(0)
+    img = renderer.init_image(stage=0)
+    optimizer.init_optimizers()
+
+    for epoch in tqdm(range(args.num_iter)):
+        renderer.set_random_noise(epoch)
+        if args.lr_scheduler:
+            optimizer.update_lr(counter)
+
+        start = time.time()
+        optimizer.zero_grad_()
+        sketches = renderer.get_image().to(args.device)
+        losses_dict = loss_func(sketches, inputs.detach(
+        ), renderer.get_color_parameters(), renderer, counter, optimizer)
+        loss = sum(list(losses_dict.values()))
+        loss.backward()
+        optimizer.step_()
+        if epoch % args.save_interval == 0:
+            utils.plot_batch(inputs, sketches, f"{args.output_dir}/jpg_logs", counter,
+                             use_wandb=args.use_wandb, title=f"iter{epoch}.jpg")
+            renderer.save_svg(
+                f"{args.output_dir}/svg_logs", f"svg_iter{epoch}")
+            # if args.cog_display:
+            #     yield Path(f"{args.output_dir}/svg_logs/svg_iter{epoch}.svg")
+
+
+        if epoch % args.eval_interval == 0:
+            with torch.no_grad():
+                losses_dict_eval = loss_func(sketches, inputs, renderer.get_color_parameters(
+                ), renderer.get_points_parans(), counter, optimizer, mode="eval")
+                loss_eval = sum(list(losses_dict_eval.values()))
+                configs_to_save["loss_eval"].append(loss_eval.item())
+                for k in losses_dict_eval.keys():
+                    if k not in configs_to_save.keys():
+                        configs_to_save[k] = []
+                    configs_to_save[k].append(losses_dict_eval[k].item())
+                if args.clip_fc_loss_weight:
+                    if losses_dict_eval["fc"].item() < best_fc_loss:
+                        best_fc_loss = losses_dict_eval["fc"].item(
+                        ) / args.clip_fc_loss_weight
+                        best_iter_fc = epoch
+                # print(
+                #     f"eval iter[{epoch}/{args.num_iter}] loss[{loss.item()}] time[{time.time() - start}]")
+
+                cur_delta = loss_eval.item() - best_loss
+                if abs(cur_delta) > min_delta:
+                    if cur_delta < 0:
+                        best_loss = loss_eval.item()
+                        best_iter = epoch
+                        terminate = False
+                        utils.plot_batch(
+                            inputs, sketches, args.output_dir, counter, use_wandb=args.use_wandb, title="best_iter.jpg")
+                        renderer.save_svg(args.output_dir, "best_iter")
+
+                if args.use_wandb:
+                    wandb.run.summary["best_loss"] = best_loss
+                    wandb.run.summary["best_loss_fc"] = best_fc_loss
+                    wandb_dict = {"delta": cur_delta,
+                                  "loss_eval": loss_eval.item()}
+                    for k in losses_dict_eval.keys():
+                        wandb_dict[k + "_eval"] = losses_dict_eval[k].item()
+                    wandb.log(wandb_dict, step=counter)
+
+                if abs(cur_delta) <= min_delta:
+                    if terminate:
+                        break
+                    terminate = True
+
+        if counter == 0 and args.attention_init:
+            utils.plot_atten(renderer.get_attn(), renderer.get_thresh(), inputs, renderer.get_inds(),
+                             args.use_wandb, "{}/{}.jpg".format(
+                                 args.output_dir, "attention_map"),
+                             args.saliency_model, args.display_logs)
+
+        if args.use_wandb:
+            wandb_dict = {"loss": loss.item(), "lr": optimizer.get_lr()}
+            for k in losses_dict.keys():
+                wandb_dict[k] = losses_dict[k].item()
+            wandb.log(wandb_dict, step=counter)
+
+        counter += 1
+
+    renderer.save_svg(args.output_dir, "final_svg")
+    path_svg = os.path.join(args.output_dir, "best_iter.svg")
+    utils.log_sketch_summary_final(
+        path_svg, args.use_wandb, args.device, best_iter, best_loss, "best total")
+
+    return configs_to_save
+
+
+def read_svg(path_svg, multiply=False):
+    device = torch.device("cuda" if (
+        torch.cuda.is_available() and torch.cuda.device_count() > 0) else "cpu")
+    canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(
+        path_svg)
+    if multiply:
+        canvas_width *= 2
+        canvas_height *= 2
+        for path in shapes:
+            path.points *= 2
+            path.stroke_width *= 2
+    _render = pydiffvg.RenderFunction.apply
+    scene_args = pydiffvg.RenderFunction.serialize_scene(
+        canvas_width, canvas_height, shapes, shape_groups)
+    img = _render(canvas_width,  # width
+                  canvas_height,  # height
+                  2,   # num_samples_x
+                  2,   # num_samples_y
+                  0,   # seed
+                  None,
+                  *scene_args)
+    img = img[:, :, 3:4] * img[:, :, :3] + \
+        torch.ones(img.shape[0], img.shape[1], 3,
+                   device=device) * (1 - img[:, :, 3:4])
+    img = img[:, :, :3]
+    return img

requirements.txt

@@ -0,0 +1,53 @@
+ftfy==6.0.3
+cssutils==2.3.0
+gdown==4.4.0
+imageio==2.9.0
+imageio-ffmpeg==0.4.4
+importlib-metadata==4.6.4
+ipykernel==6.1.0
+ipython==7.26.0
+ipython-genutils==0.2.0
+json5==0.9.5
+jsonpatch==1.32
+jsonpointer==2.1
+jsonschema==3.2.0
+jupyter-client==6.1.12
+jupyter-core==4.7.1
+jupyter-server==1.10.2
+jupyterlab==3.1.6
+jupyterlab-pygments==0.1.2
+jupyterlab-server==2.7.0
+matplotlib==3.4.2
+matplotlib-inline==0.1.2
+moviepy==1.0.3
+notebook==6.4.3
+numba==0.53.1
+numpy==1.20.3
+nvidia-ml-py3==7.352.0
+opencv-python==4.5.3.56
+pandas==1.3.2
+pathtools==0.1.2
+Pillow==8.2.0
+pip==21.2.2
+plotly==5.2.1
+psutil==5.8.0
+ptyprocess==0.7.0
+pyaml==21.8.3
+regex==2021.11.10
+scikit-image==0.18.1
+scikit-learn==1.0.2
+scipy==1.6.2
+seaborn==0.11.2
+subprocess32==3.5.4
+svgpathtools==1.4.1
+svgwrite==1.4.1
+torch==1.7.1
+torch-tools==0.1.5
+torchfile==0.1.0
+torchvision==0.8.2
+tqdm==4.62.1
+visdom==0.1.8.9
+wandb==0.12.0
+webencodings==0.5.1
+websocket-client==0.57.0
+zipp==3.5.0

run_object_sketching.py

@@ -0,0 +1,158 @@
+import sys
+import warnings
+
+warnings.filterwarnings('ignore')
+warnings.simplefilter('ignore')
+
+import argparse
+import multiprocessing as mp
+import os
+import subprocess as sp
+from shutil import copyfile
+
+import numpy as np
+import torch
+from IPython.display import Image as Image_colab
+from IPython.display import display, SVG, clear_output
+from ipywidgets import IntSlider, Output, IntProgress, Button
+import time
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--target_file", type=str,
+                    help="target image file, located in <target_images>")
+parser.add_argument("--num_strokes", type=int, default=16,
+                    help="number of strokes used to generate the sketch, this defines the level of abstraction.")
+parser.add_argument("--num_iter", type=int, default=2001,
+                    help="number of iterations")
+parser.add_argument("--fix_scale", type=int, default=0,
+                    help="if the target image is not squared, it is recommended to fix the scale")
+parser.add_argument("--mask_object", type=int, default=0,
+                    help="if the target image contains background, it's better to mask it out")
+parser.add_argument("--num_sketches", type=int, default=3,
+                    help="it is recommended to draw 3 sketches and automatically chose the best one")
+parser.add_argument("--multiprocess", type=int, default=0,
+                    help="recommended to use multiprocess if your computer has enough memory")
+parser.add_argument('-colab', action='store_true')
+parser.add_argument('-cpu', action='store_true')
+parser.add_argument('-display', action='store_true')
+parser.add_argument('--gpunum', type=int, default=0)
+
+args = parser.parse_args()
+
+multiprocess = not args.colab and args.num_sketches > 1 and args.multiprocess
+
+abs_path = os.path.abspath(os.getcwd())
+
+target = f"{abs_path}/target_images/{args.target_file}"
+assert os.path.isfile(target), f"{target} does not exists!"
+
+if not os.path.isfile(f"{abs_path}/U2Net_/saved_models/u2net.pth"):
+    sp.run(["gdown", "https://drive.google.com/uc?id=1ao1ovG1Qtx4b7EoskHXmi2E9rp5CHLcZ",
+           "-O", "U2Net_/saved_models/"])
+
+test_name = os.path.splitext(args.target_file)[0]
+output_dir = f"{abs_path}/output_sketches/{test_name}/"
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+num_iter = args.num_iter
+save_interval = 10
+use_gpu = not args.cpu
+
+if not torch.cuda.is_available():
+    use_gpu = False
+    print("CUDA is not configured with GPU, running with CPU instead.")
+    print("Note that this will be very slow, it is recommended to use colab.")
+
+if args.colab:
+    print("=" * 50)
+    print(f"Processing [{args.target_file}] ...")
+    if args.colab or args.display:
+        img_ = Image_colab(target)
+        display(img_)
+        print(f"GPU: {use_gpu}, {torch.cuda.current_device()}")
+    print(f"Results will be saved to \n[{output_dir}] ...")
+    print("=" * 50)
+
+seeds = list(range(0, args.num_sketches * 1000, 1000))
+
+exit_codes = []
+manager = mp.Manager()
+losses_all = manager.dict()
+
+
+def run(seed, wandb_name):
+    exit_code = sp.run(["python", "painterly_rendering.py", target,
+                            "--num_paths", str(args.num_strokes),
+                            "--output_dir", output_dir,
+                            "--wandb_name", wandb_name,
+                            "--num_iter", str(num_iter),
+                            "--save_interval", str(save_interval),
+                            "--seed", str(seed),
+                            "--use_gpu", str(int(use_gpu)),
+                            "--fix_scale", str(args.fix_scale),
+                            "--mask_object", str(args.mask_object),
+                            "--mask_object_attention", str(
+                                args.mask_object),
+                            "--display_logs", str(int(args.colab)),
+                            "--display", str(int(args.display))])
+    if exit_code.returncode:
+        sys.exit(1)
+
+    config = np.load(f"{output_dir}/{wandb_name}/config.npy",
+                     allow_pickle=True)[()]
+    loss_eval = np.array(config['loss_eval'])
+    inds = np.argsort(loss_eval)
+    losses_all[wandb_name] = loss_eval[inds][0]
+ 
+    
+def display_(seed, wandb_name):
+    path_to_svg = f"{output_dir}/{wandb_name}/svg_logs/"
+    intervals_ = list(range(0, num_iter, save_interval))
+    filename = f"svg_iter0.svg"
+    display(IntSlider())
+    out = Output()
+    display(out)
+    for i in intervals_:
+        filename = f"svg_iter{i}.svg"
+        not_exist = True 
+        while not_exist:
+            not_exist = not os.path.isfile(f"{path_to_svg}/{filename}")
+            continue
+        with out:
+            clear_output()
+            print("")
+            display(IntProgress(
+                        value=i,
+                        min=0,
+                        max=num_iter,
+                        description='Processing:',
+                        bar_style='info', # 'success', 'info', 'warning', 'danger' or ''
+                        style={'bar_color': 'maroon'},
+                        orientation='horizontal'
+                    ))
+            display(SVG(f"{path_to_svg}/svg_iter{i}.svg"))
+
+    
+    
+if multiprocess:
+    ncpus = 10
+    P = mp.Pool(ncpus)  # Generate pool of workers
+
+for seed in seeds:
+    wandb_name = f"{test_name}_{args.num_strokes}strokes_seed{seed}"
+    if multiprocess:
+        P.apply_async(run, (seed, wandb_name))
+    else:
+        run(seed, wandb_name)
+
+if args.display:
+    time.sleep(10)
+    P.apply_async(display_, (0, f"{test_name}_{args.num_strokes}strokes_seed0"))
+
+if multiprocess:
+    P.close()
+    P.join()  # start processes
+sorted_final = dict(sorted(losses_all.items(), key=lambda item: item[1]))
+copyfile(f"{output_dir}/{list(sorted_final.keys())[0]}/best_iter.svg",
+         f"{output_dir}/{list(sorted_final.keys())[0]}_best.svg")

sketch_utils.py

@@ -0,0 +1,295 @@
+import os
+
+import imageio
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import pydiffvg
+import skimage
+import skimage.io
+import torch
+import wandb
+import PIL
+from PIL import Image
+from torchvision import transforms
+from torchvision.utils import make_grid
+from skimage.transform import resize
+
+from U2Net_.model import U2NET
+
+
+def imwrite(img, filename, gamma=2.2, normalize=False, use_wandb=False, wandb_name="", step=0, input_im=None):
+    directory = os.path.dirname(filename)
+    if directory != '' and not os.path.exists(directory):
+        os.makedirs(directory)
+
+    if not isinstance(img, np.ndarray):
+        img = img.data.numpy()
+    if normalize:
+        img_rng = np.max(img) - np.min(img)
+        if img_rng > 0:
+            img = (img - np.min(img)) / img_rng
+    img = np.clip(img, 0.0, 1.0)
+    if img.ndim == 2:
+        # repeat along the third dimension
+        img = np.expand_dims(img, 2)
+    img[:, :, :3] = np.power(img[:, :, :3], 1.0/gamma)
+    img = (img * 255).astype(np.uint8)
+
+    skimage.io.imsave(filename, img, check_contrast=False)
+    images = [wandb.Image(Image.fromarray(img), caption="output")]
+    if input_im is not None and step == 0:
+        images.append(wandb.Image(input_im, caption="input"))
+    if use_wandb:
+        wandb.log({wandb_name + "_": images}, step=step)
+
+
+def plot_batch(inputs, outputs, output_dir, step, use_wandb, title):
+    plt.figure()
+    plt.subplot(2, 1, 1)
+    grid = make_grid(inputs.clone().detach(), normalize=True, pad_value=2)
+    npgrid = grid.cpu().numpy()
+    plt.imshow(np.transpose(npgrid, (1, 2, 0)), interpolation='nearest')
+    plt.axis("off")
+    plt.title("inputs")
+
+    plt.subplot(2, 1, 2)
+    grid = make_grid(outputs, normalize=False, pad_value=2)
+    npgrid = grid.detach().cpu().numpy()
+    plt.imshow(np.transpose(npgrid, (1, 2, 0)), interpolation='nearest')
+    plt.axis("off")
+    plt.title("outputs")
+
+    plt.tight_layout()
+    if use_wandb:
+        wandb.log({"output": wandb.Image(plt)}, step=step)
+    plt.savefig("{}/{}".format(output_dir, title))
+    plt.close()
+
+
+def log_input(use_wandb, epoch, inputs, output_dir):
+    grid = make_grid(inputs.clone().detach(), normalize=True, pad_value=2)
+    npgrid = grid.cpu().numpy()
+    plt.imshow(np.transpose(npgrid, (1, 2, 0)), interpolation='nearest')
+    plt.axis("off")
+    plt.tight_layout()
+    if use_wandb:
+        wandb.log({"input": wandb.Image(plt)}, step=epoch)
+    plt.close()
+    input_ = inputs[0].cpu().clone().detach().permute(1, 2, 0).numpy()
+    input_ = (input_ - input_.min()) / (input_.max() - input_.min())
+    input_ = (input_ * 255).astype(np.uint8)
+    imageio.imwrite("{}/{}.png".format(output_dir, "input"), input_)
+
+
+def log_sketch_summary_final(path_svg, use_wandb, device, epoch, loss, title):
+    canvas_width, canvas_height, shapes, shape_groups = load_svg(path_svg)
+    _render = pydiffvg.RenderFunction.apply
+    scene_args = pydiffvg.RenderFunction.serialize_scene(
+        canvas_width, canvas_height, shapes, shape_groups)
+    img = _render(canvas_width,  # width
+                  canvas_height,  # height
+                  2,   # num_samples_x
+                  2,   # num_samples_y
+                  0,   # seed
+                  None,
+                  *scene_args)
+
+    img = img[:, :, 3:4] * img[:, :, :3] + \
+        torch.ones(img.shape[0], img.shape[1], 3,
+                   device=device) * (1 - img[:, :, 3:4])
+    img = img[:, :, :3]
+    plt.imshow(img.cpu().numpy())
+    plt.axis("off")
+    plt.title(f"{title} best res [{epoch}] [{loss}.]")
+    if use_wandb:
+        wandb.log({title: wandb.Image(plt)})
+    plt.close()
+
+
+def log_sketch_summary(sketch, title, use_wandb):
+    plt.figure()
+    grid = make_grid(sketch.clone().detach(), normalize=True, pad_value=2)
+    npgrid = grid.cpu().numpy()
+    plt.imshow(np.transpose(npgrid, (1, 2, 0)), interpolation='nearest')
+    plt.axis("off")
+    plt.title(title)
+    plt.tight_layout()
+    if use_wandb:
+        wandb.run.summary["best_loss_im"] = wandb.Image(plt)
+    plt.close()
+
+
+def load_svg(path_svg):
+    svg = os.path.join(path_svg)
+    canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(
+        svg)
+    return canvas_width, canvas_height, shapes, shape_groups
+
+
+def read_svg(path_svg, device, multiply=False):
+    canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(
+        path_svg)
+    if multiply:
+        canvas_width *= 2
+        canvas_height *= 2
+        for path in shapes:
+            path.points *= 2
+            path.stroke_width *= 2
+    _render = pydiffvg.RenderFunction.apply
+    scene_args = pydiffvg.RenderFunction.serialize_scene(
+        canvas_width, canvas_height, shapes, shape_groups)
+    img = _render(canvas_width,  # width
+                  canvas_height,  # height
+                  2,   # num_samples_x
+                  2,   # num_samples_y
+                  0,   # seed
+                  None,
+                  *scene_args)
+    img = img[:, :, 3:4] * img[:, :, :3] + \
+        torch.ones(img.shape[0], img.shape[1], 3,
+                   device=device) * (1 - img[:, :, 3:4])
+    img = img[:, :, :3]
+    return img
+
+
+def plot_attn_dino(attn, threshold_map, inputs, inds, use_wandb, output_path):
+    # currently supports one image (and not a batch)
+    plt.figure(figsize=(10, 5))
+
+    plt.subplot(2, attn.shape[0] + 2, 1)
+    main_im = make_grid(inputs, normalize=True, pad_value=2)
+    main_im = np.transpose(main_im.cpu().numpy(), (1, 2, 0))
+    plt.imshow(main_im, interpolation='nearest')
+    plt.scatter(inds[:, 1], inds[:, 0], s=10, c='red', marker='o')
+    plt.title("input im")
+    plt.axis("off")
+
+    plt.subplot(2, attn.shape[0] + 2, 2)
+    plt.imshow(attn.sum(0).numpy(), interpolation='nearest')
+    plt.title("atn map sum")
+    plt.axis("off")
+
+    plt.subplot(2, attn.shape[0] + 2, attn.shape[0] + 3)
+    plt.imshow(threshold_map[-1].numpy(), interpolation='nearest')
+    plt.title("prob sum")
+    plt.axis("off")
+
+    plt.subplot(2, attn.shape[0] + 2, attn.shape[0] + 4)
+    plt.imshow(threshold_map[:-1].sum(0).numpy(), interpolation='nearest')
+    plt.title("thresh sum")
+    plt.axis("off")
+
+    for i in range(attn.shape[0]):
+        plt.subplot(2, attn.shape[0] + 2, i + 3)
+        plt.imshow(attn[i].numpy())
+        plt.axis("off")
+        plt.subplot(2, attn.shape[0] + 2, attn.shape[0] + 1 + i + 4)
+        plt.imshow(threshold_map[i].numpy())
+        plt.axis("off")
+    plt.tight_layout()
+    if use_wandb:
+        wandb.log({"attention_map": wandb.Image(plt)})
+    plt.savefig(output_path)
+    plt.close()
+
+
+def plot_attn_clip(attn, threshold_map, inputs, inds, use_wandb, output_path, display_logs):
+    # currently supports one image (and not a batch)
+    plt.figure(figsize=(10, 5))
+
+    plt.subplot(1, 3, 1)
+    main_im = make_grid(inputs, normalize=True, pad_value=2)
+    main_im = np.transpose(main_im.cpu().numpy(), (1, 2, 0))
+    plt.imshow(main_im, interpolation='nearest')
+    plt.scatter(inds[:, 1], inds[:, 0], s=10, c='red', marker='o')
+    plt.title("input im")
+    plt.axis("off")
+
+    plt.subplot(1, 3, 2)
+    plt.imshow(attn, interpolation='nearest', vmin=0, vmax=1)
+    plt.title("atn map")
+    plt.axis("off")
+
+    plt.subplot(1, 3, 3)
+    threshold_map_ = (threshold_map - threshold_map.min()) / \
+        (threshold_map.max() - threshold_map.min())
+    plt.imshow(threshold_map_, interpolation='nearest', vmin=0, vmax=1)
+    plt.title("prob softmax")
+    plt.scatter(inds[:, 1], inds[:, 0], s=10, c='red', marker='o')
+    plt.axis("off")
+
+    plt.tight_layout()
+    if use_wandb:
+        wandb.log({"attention_map": wandb.Image(plt)})
+    plt.savefig(output_path)
+    plt.close()
+
+
+def plot_atten(attn, threshold_map, inputs, inds, use_wandb, output_path, saliency_model, display_logs):
+    if saliency_model == "dino":
+        plot_attn_dino(attn, threshold_map, inputs,
+                       inds, use_wandb, output_path)
+    elif saliency_model == "clip":
+        plot_attn_clip(attn, threshold_map, inputs, inds,
+                       use_wandb, output_path, display_logs)
+
+
+def fix_image_scale(im):
+    im_np = np.array(im) / 255
+    height, width = im_np.shape[0], im_np.shape[1]
+    max_len = max(height, width) + 20
+    new_background = np.ones((max_len, max_len, 3))
+    y, x = max_len // 2 - height // 2, max_len // 2 - width // 2
+    new_background[y: y + height, x: x + width] = im_np
+    new_background = (new_background / new_background.max()
+                      * 255).astype(np.uint8)
+    new_im = Image.fromarray(new_background)
+    return new_im
+
+
+def get_mask_u2net(args, pil_im):
+    w, h = pil_im.size[0], pil_im.size[1]
+    im_size = min(w, h)
+    data_transforms = transforms.Compose([
+        transforms.Resize(min(320, im_size), interpolation=PIL.Image.BICUBIC),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), std=(
+            0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+    input_im_trans = data_transforms(pil_im).unsqueeze(0).to(args.device)
+
+    model_dir = os.path.join("./U2Net_/saved_models/u2net.pth")
+    net = U2NET(3, 1)
+    if torch.cuda.is_available() and args.use_gpu:
+        net.load_state_dict(torch.load(model_dir))
+        net.to(args.device)
+    else:
+        net.load_state_dict(torch.load(model_dir, map_location='cpu'))
+    net.eval()
+    with torch.no_grad():
+        d1, d2, d3, d4, d5, d6, d7 = net(input_im_trans.detach())
+    pred = d1[:, 0, :, :]
+    pred = (pred - pred.min()) / (pred.max() - pred.min())
+    predict = pred
+    predict[predict < 0.5] = 0
+    predict[predict >= 0.5] = 1
+    mask = torch.cat([predict, predict, predict], axis=0).permute(1, 2, 0)
+    mask = mask.cpu().numpy()
+    mask = resize(mask, (h, w), anti_aliasing=False)
+    mask[mask < 0.5] = 0
+    mask[mask >= 0.5] = 1
+    
+    # predict_np = predict.clone().cpu().data.numpy()
+    im = Image.fromarray((mask[:, :, 0]*255).astype(np.uint8)).convert('RGB')
+    im.save(f"{args.output_dir}/mask.png")
+
+    im_np = np.array(pil_im)
+    im_np = im_np / im_np.max()
+    im_np = mask * im_np
+    im_np[mask == 0] = 1
+    im_final = (im_np / im_np.max() * 255).astype(np.uint8)
+    im_final = Image.fromarray(im_final)
+
+    return im_final, predict